From 90449a0a052df0e093c7b54e1ed526678c2b1e2c Mon Sep 17 00:00:00 2001 From: Aleksei Lipniagov Date: Tue, 22 Jul 2025 14:18:11 +0200 Subject: [PATCH] Accept unit_primitive in AiGateway.headers Replace service signature arg with unit_primitive. --- .../fetch_model_definitions_service.rb | 6 +----- ee/lib/api/code_suggestions.rb | 8 +++++--- ee/lib/api/internal/ai/x_ray/scan.rb | 15 +++------------ ee/lib/gitlab/ai_gateway.rb | 14 +++++++++++--- ee/lib/gitlab/duo/chat/step_executor.rb | 7 +------ ee/lib/gitlab/llm/ai_gateway/client.rb | 2 +- .../llm/ai_gateway/code_suggestions_client.rb | 8 ++------ ee/lib/gitlab/llm/ai_gateway/docs_client.rb | 7 +------ ee/lib/gitlab/llm/anthropic/client.rb | 16 ++-------------- ee/lib/gitlab/llm/q_ai/client.rb | 11 +++-------- ee/lib/gitlab/llm/vertex_ai/configuration.rb | 6 +----- 11 files changed, 31 insertions(+), 69 deletions(-) diff --git a/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb b/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb index 6e58ce9462a8b3..45cfcd876ef8cd 100644 --- a/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb +++ b/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb @@ -56,7 +56,7 @@ def fetch_model_definitions def call_endpoint Gitlab::HTTP.get( endpoint, - headers: Gitlab::AiGateway.headers(user: user, service: service), + headers: Gitlab::AiGateway.headers(user: user, unit_primitive: :code_suggestions), timeout: DEFAULT_TIMEOUT, allow_local_requests: true ) @@ -79,10 +79,6 @@ def endpoint "#{base_url}/v1/#{endpoint_route}" end - - def service - ::CloudConnector::AvailableServices.find_by_name(:code_suggestions) - end end end end diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index 19f01f0b1fd992..8856c115ca5bca 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -38,10 +38,10 @@ def project(project_path) end end - def ai_gateway_headers(headers, service) + def ai_gateway_headers(headers) Gitlab::AiGateway.headers( user: current_user, - service: service, + unit_primitive: :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) } @@ -174,6 +174,8 @@ def model_prompt_cache_enabled?(project_path) unauthorized_with_origin_header! end + # TODO: do we still need this check? (and therefore obtaining the token there?) + # If yes, we will need to refactor it using `CloudConnector::Tokens.get`. token = service.access_token(current_user) unauthorized_with_origin_header! if token.nil? @@ -187,7 +189,7 @@ def model_prompt_cache_enabled?(project_path) Gitlab::Workhorse.send_url( task.endpoint, body: body, - headers: ai_gateway_headers(headers, service), + headers: ai_gateway_headers(headers), method: "POST", timeouts: { read: 55 } ) diff --git a/ee/lib/api/internal/ai/x_ray/scan.rb b/ee/lib/api/internal/ai/x_ray/scan.rb index fe67d1115f08aa..886f7b73749be4 100644 --- a/ee/lib/api/internal/ai/x_ray/scan.rb +++ b/ee/lib/api/internal/ai/x_ray/scan.rb @@ -27,12 +27,8 @@ def x_ray_available? ::GitlabSubscriptions::AddOnPurchase.exists_for_unit_primitive?(:complete_code, current_namespace) end - def code_suggestions_data - CloudConnector::AvailableServices.find_by_name(:code_suggestions) - end - - def model_gateway_headers(headers, code_suggestions_data) - Gitlab::AiGateway.headers(user: current_job.user, service: code_suggestions_data, + def model_gateway_headers(headers) + Gitlab::AiGateway.headers(user: current_job.user, unit_primitive: :code_suggestions, agent: headers["User-Agent"]) .merge(saas_headers) .transform_values { |v| Array(v) } @@ -50,11 +46,6 @@ def current_namespace current_job.namespace end strong_memoize_attr :current_namespace - - def ai_gateway_token - code_suggestions_data.access_token(current_namespace) - end - strong_memoize_attr :ai_gateway_token end namespace 'internal' do @@ -74,7 +65,7 @@ def ai_gateway_token Gitlab::Workhorse.send_url( File.join(::Gitlab::AiGateway.url, 'v1', 'x-ray', 'libraries'), body: params.except(:token, :id).to_json, - headers: model_gateway_headers(headers, code_suggestions_data), + headers: model_gateway_headers(headers), method: "POST" ) diff --git a/ee/lib/gitlab/ai_gateway.rb b/ee/lib/gitlab/ai_gateway.rb index d8e2a41501ee1d..5eda6258cbbe5d 100644 --- a/ee/lib/gitlab/ai_gateway.rb +++ b/ee/lib/gitlab/ai_gateway.rb @@ -55,17 +55,22 @@ def self.enabled_feature_flags Gitlab::SafeRequestStore.fetch(FEATURE_FLAG_CACHE_KEY) { [] } end - def self.headers(user:, service:, agent: nil, lsp_version: nil) + def self.headers(user:, unit_primitive:, agent: nil, lsp_version: nil) + token = ::CloudConnector::Tokens.get(unit_primitive: unit_primitive, resource: user) + { 'X-Gitlab-Authentication-Type' => 'oidc', - 'Authorization' => "Bearer #{service.access_token(user)}", + 'Authorization' => "Bearer #{token}", 'Content-Type' => 'application/json', 'X-Gitlab-Is-Team-Member' => (::Gitlab::Tracking::StandardContext.new.gitlab_team_member?(user&.id) || false).to_s, '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: user, service_name: unit_primitive)) + # Every Unit Primitive name that is currently passed into `AiGateway.headers` has corresponding service + # with the same name. It allows us to refactor token issuing in this method incrementally - for example, + # we can defer refactoring `User#allowed_to_use?` as it currently expects the service name, not unit primitive. .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] @@ -92,6 +97,9 @@ def self.headers(user:, service:, agent: nil, lsp_version: nil) end end + # Should be refactored to accept `unit_primitive:` (name) instead of `service_name:` after + # `5.2 Migrate access_token by feature or use case`(https://gitlab.com/groups/gitlab-org/-/epics/17999) is complete. + # We should also update `User#allowed_to_use` to accept Unit Primitive name there too. def self.public_headers(user:, service_name:) auth_response = user&.allowed_to_use(service_name) enablement_type = auth_response&.enablement_type || '' diff --git a/ee/lib/gitlab/duo/chat/step_executor.rb b/ee/lib/gitlab/duo/chat/step_executor.rb index 7cedd5acafb81a..a55a6bd69e9f72 100644 --- a/ee/lib/gitlab/duo/chat/step_executor.rb +++ b/ee/lib/gitlab/duo/chat/step_executor.rb @@ -81,7 +81,7 @@ def perform_agent_request(params) response = Gitlab::HTTP.post( "#{Gitlab::AiGateway.url}#{CHAT_V2_ENDPOINT}", - headers: Gitlab::AiGateway.headers(user: user, service: service), + headers: Gitlab::AiGateway.headers(user: user, unit_primitive: :duo_chat), body: params.to_json, timeout: DEFAULT_TIMEOUT, allow_local_requests: true, @@ -127,11 +127,6 @@ def perform_agent_request(params) raise ConnectionError, 'AI gateway not reachable' end - - def service - ::CloudConnector::AvailableServices.find_by_name(:duo_chat) - end - strong_memoize_attr :service end end end diff --git a/ee/lib/gitlab/llm/ai_gateway/client.rb b/ee/lib/gitlab/llm/ai_gateway/client.rb index 374fd76a96eceb..98ce590d3048b9 100644 --- a/ee/lib/gitlab/llm/ai_gateway/client.rb +++ b/ee/lib/gitlab/llm/ai_gateway/client.rb @@ -106,7 +106,7 @@ def perform_completion_request(url:, body:, timeout:, stream:) Gitlab::HTTP.post( url, - headers: Gitlab::AiGateway.headers(user: user, service: service), + headers: Gitlab::AiGateway.headers(user: user, unit_primitive: service.name), body: body.to_json, timeout: timeout, allow_local_requests: true, 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 8626ca48c0e5e2..a65df57d57483f 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: Gitlab::AiGateway.headers(user: user, unit_primitive: :code_suggestions), 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, - headers: Gitlab::AiGateway.headers(user: user, service: service), + headers: Gitlab::AiGateway.headers(user: user, unit_primitive: :code_suggestions), body: nil, timeout: DEFAULT_TIMEOUT, allow_local_requests: true, @@ -121,10 +121,6 @@ def success(pass_back = {}) pass_back end - def service - ::CloudConnector::AvailableServices.find_by_name(:code_suggestions) - end - def choice?(response) response['choices']&.first&.dig('text').present? 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..f2c8843558a3b2 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, unit_primitive: :duo_chat), body: request_body(query: query).to_json, timeout: timeout, allow_local_requests: true @@ -58,11 +58,6 @@ def feature_setting ::Ai::FeatureSetting.find_by_feature(:duo_chat) end - def service - ::CloudConnector::AvailableServices.find_by_name(:duo_chat) - end - strong_memoize_attr :service - def request_body(query:) { type: DEFAULT_TYPE, diff --git a/ee/lib/gitlab/llm/anthropic/client.rb b/ee/lib/gitlab/llm/anthropic/client.rb index 646d8855edc57b..545b4d11be722f 100644 --- a/ee/lib/gitlab/llm/anthropic/client.rb +++ b/ee/lib/gitlab/llm/anthropic/client.rb @@ -154,26 +154,14 @@ def url "#{Gitlab::AiGateway.url}/v1/proxy/anthropic" end - def service_name - :anthropic_proxy - end - - def service - ::CloudConnector::AvailableServices.find_by_name(service_name) - end - - def api_key - service.access_token(user) - end - - # We specificy the `anthropic-version` header to receive the stream word by word instead of the accumulated + # We specify the `anthropic-version` header to receive the stream word by word instead of the accumulated # response https://docs.anthropic.com/claude/reference/streaming. def request_headers { "Accept" => "application/json", 'anthropic-version' => '2023-06-01', 'X-Gitlab-Unit-Primitive' => unit_primitive - }.merge(Gitlab::AiGateway.headers(user: user, service: service)) + }.merge(Gitlab::AiGateway.headers(user: user, unit_primitive: unit_primitive)) end def request_body(prompt:, options: {}) diff --git a/ee/lib/gitlab/llm/q_ai/client.rb b/ee/lib/gitlab/llm/q_ai/client.rb index 69e57c637b2ce5..0cc07bdaa48b79 100644 --- a/ee/lib/gitlab/llm/q_ai/client.rb +++ b/ee/lib/gitlab/llm/q_ai/client.rb @@ -100,20 +100,15 @@ def url(path:) Gitlab::Utils.append_path(Gitlab::AiGateway.url, path) end - def service_name + def unit_primitive :amazon_q_integration end - def service - ::CloudConnector::AvailableServices.find_by_name(service_name) - end - def request_headers { "Accept" => "application/json", - # Note: In this case, the service is the same as the unit primitive name - 'X-Gitlab-Unit-Primitive' => service_name.to_s - }.merge(Gitlab::AiGateway.headers(user: user, service: service)) + 'X-Gitlab-Unit-Primitive' => unit_primitive.to_s + }.merge(Gitlab::AiGateway.headers(user: user, unit_primitive: unit_primitive)) end def with_response_logger diff --git a/ee/lib/gitlab/llm/vertex_ai/configuration.rb b/ee/lib/gitlab/llm/vertex_ai/configuration.rb index 87dfc4c4b86701..bbed3b01e54c66 100644 --- a/ee/lib/gitlab/llm/vertex_ai/configuration.rb +++ b/ee/lib/gitlab/llm/vertex_ai/configuration.rb @@ -30,16 +30,12 @@ def self.payload_parameters(params = {}) default_payload_parameters.merge(params) end - def service - ::CloudConnector::AvailableServices.find_by_name(:vertex_ai_proxy) - end - def headers { "Accept" => "application/json", "Host" => model_config.host, 'X-Gitlab-Unit-Primitive' => unit_primitive - }.merge(Gitlab::AiGateway.headers(user: user, service: service)) + }.merge(Gitlab::AiGateway.headers(user: user, unit_primitive: unit_primitive)) end private -- GitLab