From a1e07ecb2c8d8fe60252ff7382b94e27269eb0d4 Mon Sep 17 00:00:00 2001
From: Matthias Kaeppler <mkaeppler@gitlab.com>
Date: Mon, 4 Aug 2025 12:05:31 +0200
Subject: [PATCH 1/2] Migrate Code Suggestions to CloudConnector::Tokens

This moves CS onto the new Tokens factory.
This is behind a FF and not user-facing.
---
 .../code_suggestions_new_tokens_path.yml      | 10 +++++
 ee/lib/api/code_suggestions.rb                | 19 +++++----
 ee/lib/api/helpers/duo_workflow_helpers.rb    |  5 ++-
 ee/lib/cloud_connector/tokens.rb              | 19 ++++-----
 ee/lib/code_suggestions/model_details/base.rb | 11 ++++-
 .../model_details/code_completion.rb          |  5 ++-
 ee/lib/code_suggestions/tasks/base.rb         |  2 +-
 .../code_suggestions/tasks/code_generation.rb |  1 +
 ee/lib/gitlab/ai_gateway.rb                   |  8 ++--
 .../llm/ai_gateway/code_suggestions_client.rb | 10 +++--
 ee/spec/lib/cloud_connector/tokens_spec.rb    | 30 ++++++++++++++
 .../model_details/base_spec.rb                | 23 +++++++++--
 .../lib/code_suggestions/tasks/base_spec.rb   |  3 +-
 .../tasks/code_completion_spec.rb             |  3 ++
 .../tasks/code_generation_spec.rb             |  3 ++
 ee/spec/lib/gitlab/ai_gateway_spec.rb         | 19 ++++++---
 .../code_suggestions_client_spec.rb           | 41 ++++++++-----------
 ee/spec/requests/api/code_suggestions_spec.rb | 19 ++++++---
 .../code_suggestions/task_shared_examples.rb  |  4 ++
 19 files changed, 166 insertions(+), 69 deletions(-)
 create mode 100644 ee/config/feature_flags/gitlab_com_derisk/code_suggestions_new_tokens_path.yml

diff --git a/ee/config/feature_flags/gitlab_com_derisk/code_suggestions_new_tokens_path.yml b/ee/config/feature_flags/gitlab_com_derisk/code_suggestions_new_tokens_path.yml
new file mode 100644
index 00000000000000..59a46f5d88a551
--- /dev/null
+++ b/ee/config/feature_flags/gitlab_com_derisk/code_suggestions_new_tokens_path.yml
@@ -0,0 +1,10 @@
+---
+name: code_suggestions_new_tokens_path
+description: Use CloudConnector::Tokens for Code Suggestions, replacing a legacy implementation.
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/559256
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/200349
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/560317
+milestone: '18.3'
+group: group::authentication
+type: gitlab_com_derisk
+default_enabled: false
diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb
index c887d2779ba215..65a67c5ef3cb55 100644
--- a/ee/lib/api/code_suggestions.rb
+++ b/ee/lib/api/code_suggestions.rb
@@ -42,13 +42,14 @@ def ai_gateway_headers(headers, service)
         Gitlab::AiGateway.headers(
           user: current_user,
           service: service,
+          ai_feature_name: :code_suggestions,
           agent: headers["User-Agent"],
           lsp_version: headers["X-Gitlab-Language-Server-Version"]
         ).merge(saas_headers).merge(model_config_headers).transform_values { |v| Array(v) }
       end
 
-      def ai_gateway_public_headers(service_name)
-        Gitlab::AiGateway.public_headers(user: current_user, service_name: service_name)
+      def ai_gateway_public_headers(ai_feature_name, service_name)
+        Gitlab::AiGateway.public_headers(current_user, ai_feature_name, service_name)
           .merge(saas_headers)
           .merge(model_config_headers)
           .merge('X-Gitlab-Authentication-Type' => 'oidc')
@@ -175,10 +176,10 @@ def model_prompt_cache_enabled?(project_path)
 
           unauthorized_with_origin_header! if task.feature_disabled?
 
-          service = CloudConnector::AvailableServices.find_by_name(task.feature_name)
+          service = CloudConnector::AvailableServices.find_by_name(task.unit_primitive_name)
 
           unless current_user.allowed_to_use?(:code_suggestions,
-            service_name: task.feature_name,
+            service_name: task.unit_primitive_name,
             licensed_feature: task.licensed_feature
           )
             unauthorized_with_origin_header!
@@ -249,7 +250,10 @@ def model_prompt_cache_enabled?(project_path)
             # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429
             token: token[:token],
             expires_at: token[:expires_at],
-            headers: ai_gateway_public_headers(completion_model_details.feature_name)
+            headers: ai_gateway_public_headers(
+              completion_model_details.feature_name,
+              completion_model_details.unit_primitive_name
+            )
           }.tap do |a|
             a[:model_details] = details_hash unless details_hash.blank?
           end
@@ -310,8 +314,9 @@ def model_prompt_cache_enabled?(project_path)
           end
 
           aigw_headers = Gitlab::AiGateway.public_headers(
-            user: current_user,
-            service_name: completion_model_details.feature_name
+            current_user,
+            completion_model_details.feature_name,
+            completion_model_details.unit_primitive_name
           )
 
           details = {
diff --git a/ee/lib/api/helpers/duo_workflow_helpers.rb b/ee/lib/api/helpers/duo_workflow_helpers.rb
index 8763548879e76d..6ba24e3f7312b0 100644
--- a/ee/lib/api/helpers/duo_workflow_helpers.rb
+++ b/ee/lib/api/helpers/duo_workflow_helpers.rb
@@ -6,7 +6,10 @@ module DuoWorkflowHelpers
       def push_ai_gateway_headers
         push_feature_flags
 
-        Gitlab::AiGateway.public_headers(user: current_user, service_name: :duo_workflow).each do |name, value|
+        Gitlab::AiGateway.public_headers(
+          current_user,
+          :duo_workflow,
+          :duo_workflow_execute_workflow).each do |name, value|
           header(name, value)
         end
       end
diff --git a/ee/lib/cloud_connector/tokens.rb b/ee/lib/cloud_connector/tokens.rb
index 5e8acc04e90379..df13c1bc2a52f6 100644
--- a/ee/lib/cloud_connector/tokens.rb
+++ b/ee/lib/cloud_connector/tokens.rb
@@ -53,20 +53,15 @@ def fetch_active_add_ons(resource)
       GitlabSubscriptions::AddOnPurchase.for_active_add_ons(add_on_names, resource).uniq_add_on_names
     end
 
-    def use_new_token_path_for?(unit_primitive, _user_or_namespace)
+    def use_new_token_path_for?(unit_primitive, actor)
       return true if ROLLED_OUT_UNIT_PRIMITIVES.include?(unit_primitive)
 
-      # Add a feature flag temporary logic here. Rollout first less critical unit_primitives
-      # Once feature flag is rolled out, this will be removed.
-      #
-      # Example:
-      #   case unit_primitive
-      #   when :troubleshooting_job
-      #     Feature.enabled?(:use_cloud_connector_tokens_for_troubleshooting_job, _user_or_namespace)
-      #   else
-      #     false
-      #   end
-      false
+      case unit_primitive
+      when :complete_code, :generate_code
+        Feature.enabled?(:code_suggestions_new_tokens_path, actor)
+      else
+        false
+      end
     end
 
     def use_self_signed_token?(unit_primitive)
diff --git a/ee/lib/code_suggestions/model_details/base.rb b/ee/lib/code_suggestions/model_details/base.rb
index cb44affbd69a71..028506f4f50800 100644
--- a/ee/lib/code_suggestions/model_details/base.rb
+++ b/ee/lib/code_suggestions/model_details/base.rb
@@ -6,9 +6,10 @@ class Base
       include Gitlab::Utils::StrongMemoize
       include ::Ai::ModelSelection::SelectionApplicable
 
-      def initialize(current_user:, feature_setting_name:, root_namespace: nil)
+      def initialize(current_user:, feature_setting_name:, unit_primitive_name:, root_namespace: nil)
         @current_user = current_user
         @feature_setting_name = feature_setting_name
+        @unit_primitive_name = unit_primitive_name
         @root_namespace = root_namespace
       end
 
@@ -28,6 +29,14 @@ def feature_name
         :code_suggestions
       end
 
+      def unit_primitive_name
+        # We don't need to override this for SHM because this already happens
+        # in AvailableServices.find_by_name.
+        return :amazon_q_integration if ::Ai::AmazonQ.connected?
+
+        @unit_primitive_name
+      end
+
       def licensed_feature
         return :amazon_q if ::Ai::AmazonQ.connected?
 
diff --git a/ee/lib/code_suggestions/model_details/code_completion.rb b/ee/lib/code_suggestions/model_details/code_completion.rb
index c2dc83c7ebeb20..6147b9ff6fcdd5 100644
--- a/ee/lib/code_suggestions/model_details/code_completion.rb
+++ b/ee/lib/code_suggestions/model_details/code_completion.rb
@@ -8,7 +8,10 @@ class CodeCompletion < Base
       FEATURE_SETTING_NAME = 'code_completions'
 
       def initialize(current_user:, root_namespace: nil)
-        super(current_user: current_user, feature_setting_name: FEATURE_SETTING_NAME, root_namespace: root_namespace)
+        super(
+          current_user: current_user, feature_setting_name: FEATURE_SETTING_NAME,
+          unit_primitive_name: :complete_code, root_namespace: root_namespace
+        )
       end
 
       # Returns model details for using direct connection in the IDE.
diff --git a/ee/lib/code_suggestions/tasks/base.rb b/ee/lib/code_suggestions/tasks/base.rb
index ba56a7bd4ab75d..d95b2845629f32 100644
--- a/ee/lib/code_suggestions/tasks/base.rb
+++ b/ee/lib/code_suggestions/tasks/base.rb
@@ -6,7 +6,7 @@ class Base
       AI_GATEWAY_CONTENT_SIZE = 100_000
 
       delegate :base_url, :self_hosted?, :feature_setting, :feature_name, :feature_disabled?, :licensed_feature,
-        :namespace_feature_setting?, :vendored?, :duo_context_not_found?, to: :model_details
+        :namespace_feature_setting?, :vendored?, :duo_context_not_found?, :unit_primitive_name, to: :model_details
 
       def initialize(current_user:, params: {}, unsafe_passthrough_params: {}, client: nil)
         @params = params
diff --git a/ee/lib/code_suggestions/tasks/code_generation.rb b/ee/lib/code_suggestions/tasks/code_generation.rb
index 6f4e5da825822d..a529c3d86772f5 100644
--- a/ee/lib/code_suggestions/tasks/code_generation.rb
+++ b/ee/lib/code_suggestions/tasks/code_generation.rb
@@ -21,6 +21,7 @@ def model_details
         @model_details ||= CodeSuggestions::ModelDetails::Base.new(
           current_user: current_user,
           feature_setting_name: :code_generations,
+          unit_primitive_name: :generate_code,
           root_namespace: params[:project]&.root_ancestor
         )
       end
diff --git a/ee/lib/gitlab/ai_gateway.rb b/ee/lib/gitlab/ai_gateway.rb
index 0fd72a060f36e4..806bfb47c68bbd 100644
--- a/ee/lib/gitlab/ai_gateway.rb
+++ b/ee/lib/gitlab/ai_gateway.rb
@@ -80,7 +80,7 @@ def self.expanded_ai_logging_on_self_managed?(name)
       name.to_sym == :expanded_ai_logging && self_managed_instance
     end
 
-    def self.headers(user:, service:, agent: nil, lsp_version: nil)
+    def self.headers(user:, service:, ai_feature_name: service.name, agent: nil, lsp_version: nil)
       {
         'X-Gitlab-Authentication-Type' => 'oidc',
         'Authorization' => "Bearer #{cloud_connector_token(service, user)}",
@@ -90,7 +90,7 @@ def self.headers(user:, service:, agent: nil, lsp_version: nil)
         'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id,
         # Forward the request time to the model gateway to calculate latency
         'X-Gitlab-Rails-Send-Start' => Time.now.to_f.to_s
-      }.merge(public_headers(user: user, service_name: service.name))
+      }.merge(public_headers(user, ai_feature_name, service.name))
         .tap do |result|
           result['User-Agent'] = agent if agent # Forward the User-Agent on to the model gateway
           if current_context[:x_gitlab_client_type]
@@ -117,8 +117,8 @@ def self.headers(user:, service:, agent: nil, lsp_version: nil)
         end
     end
 
-    def self.public_headers(user:, service_name:)
-      auth_response = user&.allowed_to_use(service_name)
+    def self.public_headers(user, ai_feature_name, service_name)
+      auth_response = user&.allowed_to_use(ai_feature_name, service_name: service_name)
       enablement_type = auth_response&.enablement_type || ''
       namespace_ids = auth_response&.namespace_ids || []
 
diff --git a/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb b/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
index 0fabd4d8209980..babc0706216d1b 100644
--- a/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
@@ -64,7 +64,7 @@ def test_model_connection(self_hosted_model)
         def call_endpoint(endpoint, body)
           Gitlab::HTTP.post(
             endpoint,
-            headers: Gitlab::AiGateway.headers(user: user, service: service),
+            headers: ai_gateway_headers,
             body: body,
             timeout: COMPLETION_CHECK_TIMEOUT,
             allow_local_requests: true
@@ -79,7 +79,7 @@ def direct_access_token
 
           response = Gitlab::HTTP.post(
             Gitlab::AiGateway.access_token_url(code_completions_feature_setting),
-            headers: Gitlab::AiGateway.headers(user: user, service: service),
+            headers: ai_gateway_headers,
             body: nil,
             timeout: DEFAULT_TIMEOUT,
             allow_local_requests: true,
@@ -109,6 +109,10 @@ def direct_access_token
 
         attr_reader :user
 
+        def ai_gateway_headers
+          Gitlab::AiGateway.headers(user: user, service: service, ai_feature_name: task.feature_name)
+        end
+
         # We only need to look at the code completion feature setting for self-hosted models.
         # Namespace level model switching record for code completions
         # (::Ai::ModelSelection::NamespaceFeatureSetting) need not be looked at because
@@ -130,7 +134,7 @@ def success(pass_back = {})
         end
 
         def service
-          ::CloudConnector::AvailableServices.find_by_name(:code_suggestions)
+          ::CloudConnector::AvailableServices.find_by_name(task.unit_primitive_name)
         end
 
         def choice?(response)
diff --git a/ee/spec/lib/cloud_connector/tokens_spec.rb b/ee/spec/lib/cloud_connector/tokens_spec.rb
index cabc5f9dea7aec..a33f28f574f10c 100644
--- a/ee/spec/lib/cloud_connector/tokens_spec.rb
+++ b/ee/spec/lib/cloud_connector/tokens_spec.rb
@@ -71,6 +71,36 @@
         it_behaves_like 'uses self-signed path'
       end
 
+      context 'with code suggestions unit primitives' do
+        context 'with generate_code' do
+          let(:unit_primitive) { :generate_code }
+
+          it_behaves_like 'uses self-signed path'
+
+          context 'with FF disabled' do
+            before do
+              stub_feature_flags(code_suggestions_new_tokens_path: false)
+            end
+
+            it_behaves_like 'uses AvailableServices legacy path'
+          end
+        end
+
+        context 'with complete_code' do
+          let(:unit_primitive) { :complete_code }
+
+          it_behaves_like 'uses self-signed path'
+
+          context 'with FF disabled' do
+            before do
+              stub_feature_flags(code_suggestions_new_tokens_path: false)
+            end
+
+            it_behaves_like 'uses AvailableServices legacy path'
+          end
+        end
+      end
+
       context 'with unknown unit primitive' do
         let(:unit_primitive) { :not_rolled_out }
 
diff --git a/ee/spec/lib/code_suggestions/model_details/base_spec.rb b/ee/spec/lib/code_suggestions/model_details/base_spec.rb
index 8ee95285e5f51a..1c796491423357 100644
--- a/ee/spec/lib/code_suggestions/model_details/base_spec.rb
+++ b/ee/spec/lib/code_suggestions/model_details/base_spec.rb
@@ -4,12 +4,14 @@
 
 RSpec.describe CodeSuggestions::ModelDetails::Base, feature_category: :code_suggestions do
   let_it_be(:feature_setting_name) { 'code_generations' }
+  let_it_be(:unit_primitive_name) { 'generate_code' }
   let_it_be(:user) { create(:user) }
 
   let(:root_namespace) { nil }
 
-  let(:model_details) do
-    described_class.new(current_user: user, feature_setting_name: feature_setting_name, root_namespace: root_namespace)
+  subject(:model_details) do
+    described_class.new(current_user: user, feature_setting_name: feature_setting_name,
+      unit_primitive_name: unit_primitive_name, root_namespace: root_namespace)
   end
 
   shared_context 'with a default duo namespace assigned' do
@@ -272,15 +274,30 @@
     end
   end
 
+  describe '#unit_primitive_name' do
+    it 'matches the initializer argument' do
+      expect(model_details.unit_primitive_name).to eq(unit_primitive_name)
+    end
+  end
+
   context 'when Amazon Q is connected' do
     let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase, :duo_amazon_q) }
 
-    it 'returns correct feature name and licensed feature' do
+    before do
       stub_licensed_features(amazon_q: true)
       Ai::Setting.instance.update!(amazon_q_ready: true)
+      allow(::Ai::AmazonQ).to receive(:connected?).and_return(true)
+    end
 
+    it 'returns correct feature name and licensed feature' do
       expect(model_details.feature_name).to eq(:amazon_q_integration)
       expect(model_details.licensed_feature).to eq(:amazon_q)
     end
+
+    describe '#unit_primitive_name' do
+      it 'is amazon_q_integration' do
+        expect(model_details.unit_primitive_name).to eq(:amazon_q_integration)
+      end
+    end
   end
 end
diff --git a/ee/spec/lib/code_suggestions/tasks/base_spec.rb b/ee/spec/lib/code_suggestions/tasks/base_spec.rb
index 162f6c48785bf9..4aacdbe231face 100644
--- a/ee/spec/lib/code_suggestions/tasks/base_spec.rb
+++ b/ee/spec/lib/code_suggestions/tasks/base_spec.rb
@@ -12,7 +12,8 @@
         def model_details
           @model_details ||= CodeSuggestions::ModelDetails::Base.new(
             current_user: current_user,
-            feature_setting_name: :code_generations
+            feature_setting_name: :code_generations,
+            unit_primitive_name: :generate_code
           )
         end
       end
diff --git a/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb b/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb
index d3c8f80e6a79c6..ceca60d146f0ce 100644
--- a/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb
+++ b/ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb
@@ -7,6 +7,8 @@
 
   let(:endpoint_path) { 'v2/code/completions' }
 
+  let(:expected_unit_primitive_name) { :complete_code }
+
   let(:current_file) do
     {
       'file_name' => 'test.py',
@@ -511,6 +513,7 @@
 
   describe 'when amazon q is connected' do
     let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase, :duo_amazon_q) }
+    let(:expected_unit_primitive_name) { :amazon_q_integration }
 
     before do
       stub_licensed_features(amazon_q: true)
diff --git a/ee/spec/lib/code_suggestions/tasks/code_generation_spec.rb b/ee/spec/lib/code_suggestions/tasks/code_generation_spec.rb
index 540e95c53d4ac9..f0a5313facc673 100644
--- a/ee/spec/lib/code_suggestions/tasks/code_generation_spec.rb
+++ b/ee/spec/lib/code_suggestions/tasks/code_generation_spec.rb
@@ -16,6 +16,8 @@
     }.with_indifferent_access
   end
 
+  let(:expected_unit_primitive_name) { :generate_code }
+
   let(:expected_current_file) do
     { current_file: { file_name: 'test.py', content_above_cursor: 'sor', content_below_cursor: 'som' } }
   end
@@ -285,6 +287,7 @@
 
   context 'when amazon q is connected' do
     let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase, :duo_amazon_q) }
+    let(:expected_unit_primitive_name) { :amazon_q_integration }
 
     let(:unsafe_params) do
       {
diff --git a/ee/spec/lib/gitlab/ai_gateway_spec.rb b/ee/spec/lib/gitlab/ai_gateway_spec.rb
index 03f866d216bfd5..37227102956b95 100644
--- a/ee/spec/lib/gitlab/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/ai_gateway_spec.rb
@@ -185,7 +185,8 @@
   describe '.headers', :request_store do
     let(:user) { build(:user, id: 1) }
     let(:token) { 'instance token' }
-    let(:service_name) { :test }
+    let(:ai_feature) { :test_feature }
+    let(:service_name) { :test_service }
     let(:service) { instance_double(CloudConnector::BaseAvailableServiceData, name: service_name) }
     let(:agent) { nil }
     let(:lsp_version) { nil }
@@ -222,7 +223,12 @@
       }.merge(cloud_connector_headers)
     end
 
-    subject(:headers) { described_class.headers(user: user, service: service, agent: agent, lsp_version: lsp_version) }
+    subject(:headers) do
+      described_class.headers(
+        user: user, service: service, ai_feature_name: ai_feature,
+        agent: agent, lsp_version: lsp_version
+      )
+    end
 
     before do
       allow_next_instance_of(::Ai::Setting) do |setting|
@@ -232,7 +238,7 @@
       allow(::CloudConnector::Tokens).to receive(:get)
         .with(unit_primitive: service_name, resource: user)
         .and_return(token)
-      allow(user).to receive(:allowed_to_use).with(service_name).and_return(auth_response)
+      allow(user).to receive(:allowed_to_use).with(ai_feature, service_name: service_name).and_return(auth_response)
       allow(::CloudConnector).to(
         receive(:ai_headers).with(user, namespace_ids: namespace_ids).and_return(cloud_connector_headers)
       )
@@ -311,11 +317,12 @@
 
   describe '.public_headers' do
     let(:user) { build(:user, id: 1) }
-    let(:service_name) { :test }
+    let(:ai_feature) { :test_feature }
+    let(:service_name) { :test_service }
     let(:enabled_feature_flags) { %w[feature_a feature_b] }
     let(:ai_headers) { { 'x-gitlab-feature-enabled-by-namespace-ids' => '' } }
 
-    subject(:public_headers) { described_class.public_headers(user: user, service_name: service_name) }
+    subject(:public_headers) { described_class.public_headers(user, ai_feature, service_name) }
 
     before do
       allow_next_instance_of(::Ai::Setting) do |setting|
@@ -323,7 +330,7 @@
       end
 
       allow(user).to receive(:allowed_to_use)
-        .with(service_name)
+        .with(ai_feature, service_name: service_name)
         .and_return(auth_response)
 
       allow(described_class).to receive(:enabled_feature_flags)
diff --git a/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb b/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
index 931262bb5be238..fe0f22ab8f434b 100644
--- a/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
@@ -4,10 +4,12 @@
 
 RSpec.describe Gitlab::Llm::AiGateway::CodeSuggestionsClient, feature_category: :code_suggestions do
   let_it_be(:user) { create(:user) }
-  let_it_be(:instance_token) { create(:service_access_token, :active) }
-  let(:service) { instance_double(CloudConnector::BaseAvailableServiceData, name: :code_suggestions) }
+
+  let(:unit_primitive) { :complete_code }
+  let(:service) { instance_double(CloudConnector::BaseAvailableServiceData, name: unit_primitive) }
   let(:enabled_by_namespace_ids) { [1, 2] }
   let(:enablement_type) { 'add_on' }
+  let(:ai_gateway_headers) { { 'header' => 'value' } }
   let(:auth_response) do
     instance_double(Ai::UserAuthorizable::Response,
       namespace_ids: enabled_by_namespace_ids, enablement_type: enablement_type)
@@ -19,18 +21,22 @@
 
   let(:self_hosted_auth_endpoint_url) { "#{Gitlab::AiGateway.self_hosted_url}#{Gitlab::AiGateway::ACCESS_TOKEN_PATH}" }
 
+  let(:expected_ai_feature) { :code_suggestions }
+
   let(:body) { { choices: [{ text: "puts \"Hello World!\"\nend", index: 0, finish_reason: "length" }] } }
   let(:code) { 200 }
 
   before do
-    allow(CloudConnector::AvailableServices).to receive(:find_by_name).and_return(service)
-    allow(service).to receive(:access_token).and_return(instance_token&.token)
+    allow(CloudConnector::AvailableServices).to receive(:find_by_name).with(unit_primitive).and_return(service)
     allow(user).to receive(:allowed_to_use).and_return(auth_response)
     allow(Gitlab::AiGateway).to receive_messages(
       self_hosted_url: 'http://local-aigw:5052',
       cloud_connector_url: 'https://cloud-connector.gitlab.com',
       cloud_connector_auth_url: 'https://cloud-connector.gitlab.com/auth'
     )
+    allow(Gitlab::AiGateway).to receive(:headers).with(
+      user: user, service: service, ai_feature_name: expected_ai_feature
+    ).and_return(ai_gateway_headers)
   end
 
   shared_examples "error response" do |message|
@@ -47,14 +53,13 @@
     end
   end
 
-  shared_context 'with tests requests' do |expected_service_name|
+  shared_context 'with tests requests' do
     before do
       stub_request(:post, /#{Gitlab::AiGateway.url}/)
         .to_return(status: code, body: body.to_json, headers: { "Content-Type" => "application/json" })
     end
 
     it 'returns nil if there is no error' do
-      expect(::CloudConnector::AvailableServices).to receive(:find_by_name).with(expected_service_name)
       expect(result).to be_nil
     end
 
@@ -76,7 +81,7 @@
   describe "#test_completion" do
     subject(:result) { described_class.new(user).test_completion }
 
-    include_examples 'with tests requests', :code_suggestions do
+    include_examples 'with tests requests' do
       include_examples 'with completions'
     end
 
@@ -142,7 +147,7 @@
         "your AI Gateway URL is configured correctly."
     end
 
-    include_examples 'with tests requests', :code_suggestions
+    include_examples 'with tests requests'
   end
 
   describe '#direct_access_token', :with_cloud_connector do
@@ -154,20 +159,6 @@
     let(:response_body) { expected_response.to_json }
     let(:http_status) { 200 }
     let(:client) { described_class.new(user) }
-    let(:expected_request_headers) do
-      {
-        'X-Gitlab-Instance-Id' => Gitlab::GlobalAnonymousId.instance_id,
-        'X-Gitlab-Global-User-Id' => Gitlab::GlobalAnonymousId.user_id(user),
-        'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host,
-        'X-Gitlab-Realm' => ::CloudConnector::GITLAB_REALM_SELF_MANAGED,
-        'X-Gitlab-Authentication-Type' => 'oidc',
-        'Authorization' => "Bearer #{instance_token.token}",
-        "X-Gitlab-Feature-Enabled-By-Namespace-Ids" => [enabled_by_namespace_ids.join(',')],
-        'X-Gitlab-Feature-Enablement-Type' => enablement_type,
-        'Content-Type' => 'application/json',
-        'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id
-      }
-    end
 
     let(:auth_url) { self_hosted_auth_endpoint_url }
 
@@ -177,7 +168,7 @@
       stub_request(:post, auth_url)
         .with(
           body: nil,
-          headers: expected_request_headers
+          headers: ai_gateway_headers
         )
         .to_return(
           status: http_status,
@@ -199,6 +190,8 @@
     end
 
     context 'when code_completions is self-hosted' do
+      let(:expected_ai_feature) { :self_hosted_models }
+
       before do
         create(:ai_feature_setting, :code_completions, provider: :self_hosted)
       end
@@ -251,7 +244,7 @@
         stub_request(:post, auth_url)
           .with(
             body: nil,
-            headers: expected_request_headers
+            headers: ai_gateway_headers
           )
           .to_return(
             status: http_status,
diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb
index e5fba3c658a170..5e2a424ad99d80 100644
--- a/ee/spec/requests/api/code_suggestions_spec.rb
+++ b/ee/spec/requests/api/code_suggestions_spec.rb
@@ -29,7 +29,7 @@
   let(:global_instance_id) { 'instance-ABC' }
   let(:global_user_id) { 'user-ABC' }
   let(:gitlab_realm) { 'saas' }
-  let(:service_name) { :code_suggestions }
+  let(:unit_primitive_name) { :complete_code }
   let(:service) { instance_double('::CloudConnector::SelfSigned::AvailableServiceData') }
 
   subject(:default_namespace_example_state) do
@@ -53,8 +53,13 @@
     allow(Gitlab::GlobalAnonymousId).to receive(:user_id).and_return(global_user_id)
     allow(Gitlab::GlobalAnonymousId).to receive(:instance_id).and_return(global_instance_id)
 
-    allow(::CloudConnector::AvailableServices).to receive(:find_by_name).with(service_name).and_return(service)
-    allow(service).to receive_messages(access_token: token, name: service_name, add_on_names: ['code_suggestions'])
+    allow(::CloudConnector::AvailableServices).to receive(:find_by_name).with(unit_primitive_name).and_return(service)
+    allow(service).to receive_messages(
+      access_token: token, name: unit_primitive_name, add_on_names: ['code_suggestions']
+    )
+    allow(::CloudConnector::Tokens).to receive(:get)
+      .with(unit_primitive: unit_primitive_name, resource: authorized_user)
+      .and_return(token)
 
     purchases = class_double(GitlabSubscriptions::AddOnPurchase)
     mock_purchase = instance_double(GitlabSubscriptions::AddOnPurchase, normalized_add_on_name: 'duo_pro')
@@ -474,6 +479,7 @@ def request
 
           context 'with generation intent' do
             let(:additional_params) { { intent: 'generation' } }
+            let(:unit_primitive_name) { :generate_code }
 
             it 'passes generation intent into TaskFactory.new' do
               expect(::CodeSuggestions::TaskFactory).to receive(:new)
@@ -507,6 +513,7 @@ def request
 
         context 'when passing generation_type parameter' do
           let(:additional_params) { { generation_type: :small_file } }
+          let(:unit_primitive_name) { :generate_code }
 
           it 'passes generation_type into TaskFactory.new' do
             expect(::CodeSuggestions::TaskFactory).to receive(:new)
@@ -568,6 +575,7 @@ def request
 
         context 'when passing user_instruction parameter' do
           let(:additional_params) { { user_instruction: 'Generate tests for this file' } }
+          let(:unit_primitive_name) { :generate_code }
 
           it 'passes user_instruction into TaskFactory.new' do
             expect(::CodeSuggestions::TaskFactory).to receive(:new)
@@ -742,6 +750,7 @@ def request
           end
 
           context 'when the task is code generation' do
+            let(:unit_primitive_name) { :generate_code }
             let(:content_above_cursor) do
               <<~CONTENT_ABOVE_CURSOR
                 def is_even(n: int) ->
@@ -952,7 +961,7 @@ def get_user(session):
 
         context 'when code suggestions feature is self hosted' do
           let(:top_level_namespace) { nil }
-          let(:service_name) { :self_hosted_models }
+          let(:unit_primitive_name) { :self_hosted_models }
 
           before do
             stub_licensed_features(ai_features: true)
@@ -973,7 +982,7 @@ def get_user(session):
         end
 
         context 'when Amazon Q is connected' do
-          let(:service_name) { :amazon_q_integration }
+          let(:unit_primitive_name) { :amazon_q_integration }
 
           before do
             stub_licensed_features(amazon_q: true)
diff --git a/ee/spec/support/shared_examples/lib/code_suggestions/task_shared_examples.rb b/ee/spec/support/shared_examples/lib/code_suggestions/task_shared_examples.rb
index abf3b33a11dfcf..a08cc46b334453 100644
--- a/ee/spec/support/shared_examples/lib/code_suggestions/task_shared_examples.rb
+++ b/ee/spec/support/shared_examples/lib/code_suggestions/task_shared_examples.rb
@@ -16,6 +16,10 @@
     expect(task.feature_name).to eq expected_feature_name
   end
 
+  it 'has correct unit_primitive_name' do
+    expect(task.unit_primitive_name).to eq expected_unit_primitive_name
+  end
+
   it 'is not disabled' do
     expect(task.feature_disabled?).to eq false
   end
-- 
GitLab


From 5658c77115971a4abd1041762715bbcf526c3d47 Mon Sep 17 00:00:00 2001
From: Aleksei Lipniagov <alipniagov@gitlab.com>
Date: Tue, 12 Aug 2025 11:25:06 +0200
Subject: [PATCH 2/2] Migrate DocsClient to new Tokens interface

Supply proper Unit Primitive.
---
 .../documentation_search_new_tokens_path.yml         | 10 ++++++++++
 doc/api/openapi/openapi_v2.yaml                      | 12 +++++++++---
 ee/lib/cloud_connector/tokens.rb                     |  2 ++
 ee/lib/gitlab/llm/ai_gateway/docs_client.rb          |  4 ++--
 .../lib/gitlab/llm/ai_gateway/docs_client_spec.rb    |  6 ++----
 5 files changed, 25 insertions(+), 9 deletions(-)
 create mode 100644 config/feature_flags/gitlab_com_derisk/documentation_search_new_tokens_path.yml

diff --git a/config/feature_flags/gitlab_com_derisk/documentation_search_new_tokens_path.yml b/config/feature_flags/gitlab_com_derisk/documentation_search_new_tokens_path.yml
new file mode 100644
index 00000000000000..7451fea4f5fc9e
--- /dev/null
+++ b/config/feature_flags/gitlab_com_derisk/documentation_search_new_tokens_path.yml
@@ -0,0 +1,10 @@
+---
+name: documentation_search_new_tokens_path
+description: Use CloudConnector::Tokens for Documentation Search, replacing a legacy implementation.
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/560757
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/200952
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/560757
+milestone: '18.4'
+group: group::authentication
+type: gitlab_com_derisk
+default_enabled: false
diff --git a/doc/api/openapi/openapi_v2.yaml b/doc/api/openapi/openapi_v2.yaml
index 2f2731c4017d4d..c6a7a4a320721f 100644
--- a/doc/api/openapi/openapi_v2.yaml
+++ b/doc/api/openapi/openapi_v2.yaml
@@ -51123,9 +51123,15 @@ definitions:
     type: object
     properties:
       total:
-        type: integer
-        format: int32
-        example: 3363
+        type: object
+        example:
+          time: 0.42
+          count: 2
+          success: 2
+          failed: 0
+          skipped: 0
+          error: 0
+          suite_error:
       test_suites:
         "$ref": "#/definitions/TestSuiteSummaryEntity"
     description: TestReportSummaryEntity model
diff --git a/ee/lib/cloud_connector/tokens.rb b/ee/lib/cloud_connector/tokens.rb
index df13c1bc2a52f6..7669780daff19c 100644
--- a/ee/lib/cloud_connector/tokens.rb
+++ b/ee/lib/cloud_connector/tokens.rb
@@ -59,6 +59,8 @@ def use_new_token_path_for?(unit_primitive, actor)
       case unit_primitive
       when :complete_code, :generate_code
         Feature.enabled?(:code_suggestions_new_tokens_path, actor)
+      when :documentation_search
+        Feature.enabled?(:documentation_search_new_tokens_path, actor)
       else
         false
       end
diff --git a/ee/lib/gitlab/llm/ai_gateway/docs_client.rb b/ee/lib/gitlab/llm/ai_gateway/docs_client.rb
index 5cc7fdd0e2928d..dc8ccfda77a0ef 100644
--- a/ee/lib/gitlab/llm/ai_gateway/docs_client.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/docs_client.rb
@@ -35,7 +35,7 @@ def perform_search_request(query:, options:)
 
           response = Gitlab::HTTP.post(
             "#{base_url}/v1/search/gitlab-docs",
-            headers: Gitlab::AiGateway.headers(user: user, service: service),
+            headers: Gitlab::AiGateway.headers(user: user, service: service, ai_feature_name: :duo_chat),
             body: request_body(query: query).to_json,
             timeout: timeout,
             allow_local_requests: true
@@ -59,7 +59,7 @@ def feature_setting
         end
 
         def service
-          ::CloudConnector::AvailableServices.find_by_name(:duo_chat)
+          ::CloudConnector::AvailableServices.find_by_name(:documentation_search)
         end
         strong_memoize_attr :service
 
diff --git a/ee/spec/lib/gitlab/llm/ai_gateway/docs_client_spec.rb b/ee/spec/lib/gitlab/llm/ai_gateway/docs_client_spec.rb
index f34c9ed859643a..22625f7ac8327c 100644
--- a/ee/spec/lib/gitlab/llm/ai_gateway/docs_client_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_gateway/docs_client_spec.rb
@@ -17,7 +17,6 @@
       namespace_ids: enabled_by_namespace_ids, enablement_type: enablement_type)
   end
 
-  let(:expected_feature_name) { :duo_chat }
   let(:expected_request_headers) { { 'header' => 'value' } }
 
   let(:default_body_params) do
@@ -47,9 +46,9 @@
 
   before do
     service = instance_double(CloudConnector::BaseAvailableServiceData)
-    allow(::CloudConnector::AvailableServices).to receive(:find_by_name).with(expected_feature_name).and_return(service)
+    allow(::CloudConnector::AvailableServices).to receive(:find_by_name).with(:documentation_search).and_return(service)
     allow(Gitlab::AiGateway).to receive(:headers)
-      .with(user: user, service: service)
+      .with(user: user, service: service, ai_feature_name: :duo_chat)
       .and_return(expected_request_headers)
     allow(Gitlab::AiGateway).to receive_messages(
       self_hosted_url: self_hosted_url,
@@ -86,7 +85,6 @@
 
     context 'when duo chat is self-hosted' do
       let_it_be(:feature_setting) { create(:ai_feature_setting, feature: :duo_chat) }
-      let(:expected_feature_name) { :duo_chat }
 
       it 'returns response for duo_chat service' do
         expect(Gitlab::HTTP).to receive(:post).with(
-- 
GitLab