Skip to content

Migrate DocsClient to CloudConnector::Tokens

What does this MR do and why?

Migrate DocsClient to the updated Token interface.

This is a refactoring and should not result in any user-facing changes.

📚 Built on top of Migrate Code Suggestions to CloudConnector::Tokens (!200349 - merged) because the mentioned MR introduces an updated .headers() signature that includes ai_feature: argument. We need to supply ai_feature: separately because it does not match Unit Primitive name in our case. As soon as the mentioned MR lands master, the diff will be normalized.

How to test locally

Pull this branch.
Make sure you run GitLab as gitlab.com/SaaS (GITLAB_SIMULATE_SAAS=1).
Restart GitLab.
In rails console, enable the FF: Feature.enable(:documentation_search_new_tokens_path)
Open Duo Chat.
Ask Chat to search through GitLab docs.
I used this prompt to invoke that:

explain me how gitlab ci works in 3 sentences. search gitlab docs.

First of all, you should get proper response.
Second, in tail -f log/llm.log you will see relevant log records:

{
  "severity": "INFO",
  "time": "2025-08-12T09:38:38.984Z",
  "correlation_id": "01K2ETAM7S334KZ9QRM7QBZ2XE",
  "meta.caller_id": "Llm::CompletionWorker",
  "meta.feature_category": "ai_abstraction_layer",
  "meta.organization_id": 1,
  "meta.remote_ip": "127.0.0.1",
  "meta.http_router_rule_action": "classify",
  "meta.http_router_rule_type": "SESSION_PREFIX",
  "meta.user": "root",
  "meta.user_id": 1,
  "meta.client_id": "user/1",
  "meta.root_caller_id": "GraphqlController#execute",
  "options": {},
  "message": "Searching docs from AI Gateway",
  "class": "Gitlab::Llm::AiGateway::DocsClient",
  "ai_event_name": "performing_request",
  "ai_component": "duo_chat"
}

Finally, you should also see relevant logs in your AI GW (if running locally):

127.0.0.1:61533 - "POST /v2/chat/agent HTTP/1.1" 200  
127.0.0.1:61537 - "POST /v1/search/gitlab-docs HTTP/1.1" 200   
127.0.0.1:61539 - "POST /v1/prompts/chat/documentation_search HTTP/1.1" 200

I also debug-printed the result of TokenIssuer.new() (under if use_new_token_path_for?) -> with that I confirmed that we indeed enter the new issuing path. By inspecting the token and verifying I am getting the correct response from AI GW, I confirmed that the functionality remained unaffected

The user-facing behavior will be similar for master and this branch (in any FF state) because it is a refactor.
You can also check similar steps in Self-Managed to make sure Docs Search (and Chat) works as before.

References

Related to #560757 (closed)

Edited by Aleksei Lipniagov

Merge request reports

Loading