From 8394c1d43f22f144ed52fc67958cd74854f89740 Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Mon, 18 Mar 2024 10:06:34 +0100 Subject: [PATCH 1/7] Return list of endpoints for code suggestions Returns additional 'completions' endpoint information. This can be used by clients to send requests directly to AI Gateway. --- doc/api/code_suggestions.md | 49 ++++++++++ .../code_suggestions_direct_completions.yml | 9 ++ ee/lib/api/code_suggestions.rb | 24 +++++ .../code_suggestions/connection_info.rb | 11 +++ .../entities/code_suggestions_access_token.rb | 13 --- ee/lib/code_suggestions/connection_info.rb | 37 ++++++++ ee/lib/ee/gitlab/application_rate_limiter.rb | 1 + .../code_suggestions_access_token_spec.rb | 17 ---- .../code_suggestions/connection_info_spec.rb | 47 +++++++++ ee/spec/requests/api/code_suggestions_spec.rb | 95 ++++++++++++++----- 10 files changed, 248 insertions(+), 55 deletions(-) create mode 100644 ee/config/feature_flags/wip/code_suggestions_direct_completions.yml create mode 100644 ee/lib/api/entities/code_suggestions/connection_info.rb delete mode 100644 ee/lib/api/entities/code_suggestions_access_token.rb create mode 100644 ee/lib/code_suggestions/connection_info.rb delete mode 100644 ee/spec/lib/api/entities/code_suggestions_access_token_spec.rb create mode 100644 ee/spec/lib/code_suggestions/connection_info_spec.rb diff --git a/doc/api/code_suggestions.md b/doc/api/code_suggestions.md index 4ce12887ad67b5..cb7c04fde4cc52 100644 --- a/doc/api/code_suggestions.md +++ b/doc/api/code_suggestions.md @@ -126,3 +126,52 @@ curl --request POST \ }' \ ``` + +## Fetch direct connection information + +DETAILS: +**Status:** Experiment + +> - [Introduced](https://gitlab.com/gitlab-org/gitlab/-/issues/452044) in GitLab 17.0 [with a flag](../administration/feature_flags.md) named `code_suggestions_direct_completions`. Disabled by default. This feature is an Experiment. + +FLAG: +On self-managed GitLab, by default this feature is not available. To make it available, an administrator can [enable the feature flag](../administration/feature_flags.md) named `code_suggestions_direct_completions`. +On GitLab.com and GitLab Dedicated, this feature is not available. +This feature is not ready for production use. + +```plaintext +POST /code_suggestions/connection_info +``` + +NOTE: +This endpoint rate-limits each user to 10 requests per 5-minute window. + +Returns user-specific connection details which can be used by IDEs/clients to send completion requests directly to AI Gateway. + +Example request: + +```shell +curl --request POST \ + --header "Authorization: Bearer " \ + --url "https://gitlab.example.com/api/v4/code_suggestions/connection_info" +``` + +Example response: + +```json +{ + "endpoints": { + "completions": { + "base_url": "http://0.0.0.0:5052", + "jwt_token": "a valid token", + "expires_at": 3600, + "headers": { + "X-Gitlab-Instance-Id": "292c3c7c-c5d5-48ec-b4bf-f00b724ce560", + "X-Gitlab-Realm": "saas", + "X-Gitlab-Global-User-Id": "Df0Jhs9xlbetQR8YoZCKDZJflhxO0ZBI8uoRzmpnd1w=", + "X-Gitlab-Host-Name": "gitlab.example.com" + } + } + } +} +``` diff --git a/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml b/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml new file mode 100644 index 00000000000000..7373665f79fa98 --- /dev/null +++ b/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml @@ -0,0 +1,9 @@ +--- +name: code_suggestions_direct_completions +feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/452044 +introduced_by_url: +rollout_issue_url: +milestone: '16.11' +group: group::code creation +type: wip +default_enabled: false diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index b1e09d1f1f254b..5fe39348d84a4f 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -114,6 +114,30 @@ def saas_headers end end + resources :connection_info do + desc 'Code suggestion connection information' do + success code: 200 + failure [ + { code: 401, message: 'Unauthorized' }, + { code: 404, message: 'Not found' } + ] + end + + post do + check_rate_limit!(:code_suggestions_connection_info, scope: current_user) do + Gitlab::InternalEvents.track_event( + 'code_suggestions_connection_info_rate_limit_exceeded', + user: current_user + ) + + render_api_error!({ error: _('This endpoint has been requested too many times. Try again later.') }, 429) + end + + info = ::CodeSuggestions::ConnectionInfo.new(current_user: current_user) + present info, with: ::API::Entities::CodeSuggestions::ConnectionInfo + end + end + resources :enabled do desc 'Code suggestions enabled for a project' do success code: 200 diff --git a/ee/lib/api/entities/code_suggestions/connection_info.rb b/ee/lib/api/entities/code_suggestions/connection_info.rb new file mode 100644 index 00000000000000..13e12a5bd5cb76 --- /dev/null +++ b/ee/lib/api/entities/code_suggestions/connection_info.rb @@ -0,0 +1,11 @@ +# frozen_string_literal: true + +module API + module Entities + module CodeSuggestions + class ConnectionInfo < Grape::Entity + expose :endpoints + end + end + end +end diff --git a/ee/lib/api/entities/code_suggestions_access_token.rb b/ee/lib/api/entities/code_suggestions_access_token.rb deleted file mode 100644 index b8a4c110c541b4..00000000000000 --- a/ee/lib/api/entities/code_suggestions_access_token.rb +++ /dev/null @@ -1,13 +0,0 @@ -# frozen_string_literal: true - -module API - module Entities - class CodeSuggestionsAccessToken < Grape::Entity - expose :encoded, as: :access_token, documentation: { type: 'string', example: 'eyJ0eXAi...' } - expose :expires_in, documentation: { type: 'integer', example: 3600 } do |token| - token.class::EXPIRES_IN - end - expose :issued_at, as: :created_at, documentation: { type: 'integer', example: 1684386897 } - end - end -end diff --git a/ee/lib/code_suggestions/connection_info.rb b/ee/lib/code_suggestions/connection_info.rb new file mode 100644 index 00000000000000..3c1a954a8678df --- /dev/null +++ b/ee/lib/code_suggestions/connection_info.rb @@ -0,0 +1,37 @@ +# frozen_string_literal: true + +module CodeSuggestions + class ConnectionInfo + include ::API::Helpers::CloudConnector + + def initialize(current_user:) + @current_user = current_user + end + + def endpoints + result = {} + + if Feature.enabled?(:code_suggestions_direct_completions, current_user) + result[:completions] = { + base_url: ::Gitlab::AiGateway.url, + # for development purposes we just return instance JWT, this should not be used in production + # until we generate a short-term token for user + # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429 + jwt_token: ::Gitlab::Llm::AiGateway::Client.access_token(scopes: [:code_suggestions]), + expires_at: ::Gitlab::CloudConnector::SelfIssuedToken::EXPIRES_IN, + headers: headers + } + end + + result + end + + private + + attr_reader :current_user, :params + + def headers + cloud_connector_headers(current_user).merge('X-Gitlab-Host-Name' => Gitlab.config.gitlab.host) + end + end +end diff --git a/ee/lib/ee/gitlab/application_rate_limiter.rb b/ee/lib/ee/gitlab/application_rate_limiter.rb index e049b8122a6988..d39c6a7303998e 100644 --- a/ee/lib/ee/gitlab/application_rate_limiter.rb +++ b/ee/lib/ee/gitlab/application_rate_limiter.rb @@ -22,6 +22,7 @@ def rate_limits }, credit_card_verification_check_for_reuse: { threshold: 10, interval: 1.hour }, code_suggestions_api_endpoint: { threshold: 60, interval: 1.minute }, + code_suggestions_connection_info: { threshold: 10, interval: 5.minutes }, soft_phone_verification_transactions_limit: { threshold: 16_000, interval: 1.day }, hard_phone_verification_transactions_limit: { threshold: 20_000, interval: 1.day } }).freeze diff --git a/ee/spec/lib/api/entities/code_suggestions_access_token_spec.rb b/ee/spec/lib/api/entities/code_suggestions_access_token_spec.rb deleted file mode 100644 index a04fcc22b7ccd3..00000000000000 --- a/ee/spec/lib/api/entities/code_suggestions_access_token_spec.rb +++ /dev/null @@ -1,17 +0,0 @@ -# frozen_string_literal: true - -require 'spec_helper' - -RSpec.describe API::Entities::CodeSuggestionsAccessToken, feature_category: :code_suggestions do - subject { described_class.new(token).as_json } - - let_it_be(:token) do - Gitlab::CloudConnector::SelfIssuedToken.new( - audience: 'gitlab-ai-gateway', subject: 'ABC-123', scopes: [:code_suggestions] - ) - end - - it 'exposes correct attributes' do - expect(subject.keys).to contain_exactly(:access_token, :expires_in, :created_at) - end -end diff --git a/ee/spec/lib/code_suggestions/connection_info_spec.rb b/ee/spec/lib/code_suggestions/connection_info_spec.rb new file mode 100644 index 00000000000000..8c318e58216452 --- /dev/null +++ b/ee/spec/lib/code_suggestions/connection_info_spec.rb @@ -0,0 +1,47 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe CodeSuggestions::ConnectionInfo, feature_category: :code_suggestions do + describe '#endpoints' do + let(:global_instance_id) { 'instance-ABC' } + let(:global_user_id) { 'user-ABC' } + let(:token) { 'some token' } + let(:user) { create(:user) } + + subject(:endpoints) { described_class.new(current_user: user).endpoints } + + before do + allow_next_instance_of(API::Helpers::GlobalIds::Generator) do |generator| + allow(generator).to receive(:generate).with(user).and_return([global_instance_id, global_user_id]) + end + + allow(Gitlab::Llm::AiGateway::Client).to receive(:access_token).and_return(token) + end + + it 'returns expected endpoints' do + expected_endpoints = { + completions: { + base_url: 'https://cloud.gitlab.com/ai', + expires_at: 3600, + jwt_token: token, + headers: { + 'X-Gitlab-Global-User-Id' => global_user_id, + 'X-Gitlab-Instance-Id' => global_instance_id, + 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, + 'X-Gitlab-Realm' => Gitlab::CloudConnector.gitlab_realm + } + } + } + expect(endpoints).to match(expected_endpoints) + end + + context 'when code_suggestions_direct_completions flag is disabled' do + before do + stub_feature_flags(code_suggestions_direct_completions: false) + end + + it { is_expected.to be_empty } + end + end +end diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index 7b71dd8d9eccb7..6ea40ea1a93f18 100644 --- a/ee/spec/requests/api/code_suggestions_spec.rb +++ b/ee/spec/requests/api/code_suggestions_spec.rb @@ -20,6 +20,9 @@ let(:headers) { {} } let(:access_code_suggestions) { true } let(:is_saas) { true } + let(:global_instance_id) { 'instance-ABC' } + let(:global_user_id) { 'user-ABC' } + let(:gitlab_realm) { 'saas' } before do allow(Gitlab).to receive(:com?).and_return(is_saas) @@ -35,6 +38,10 @@ stub_feature_flags(claude_3_code_generation_opus: false) stub_feature_flags(claude_3_code_generation_sonnet: false) stub_feature_flags(claude_3_code_generation_haiku: false) + + allow_next_instance_of(API::Helpers::GlobalIds::Generator) do |generator| + allow(generator).to receive(:generate).with(authorized_user).and_return([global_instance_id, global_user_id]) + end end shared_examples 'a response' do |case_name| @@ -129,10 +136,22 @@ end end + shared_examples_for 'rate limited and tracked endpoint' do |rate_limit_key:, event_name:| + it_behaves_like 'rate limited endpoint', rate_limit_key: rate_limit_key + + it 'tracks rate limit exceeded event' do + allow(Gitlab::ApplicationRateLimiter).to receive(:throttled_request?).and_return(true) + + request + + expect(Gitlab::InternalEvents) + .to have_received(:track_event) + .with(event_name, user: current_user) + end + end + describe 'POST /code_suggestions/completions' do let(:access_code_suggestions) { true } - let(:global_instance_id) { 'instance-ABC' } - let(:global_user_id) { 'user-ABC' } let(:prefix) do <<~PREFIX @@ -175,10 +194,6 @@ def is_even(n: int) -> before do allow(Gitlab::ApplicationRateLimiter).to receive(:threshold).and_return(0) - - allow_next_instance_of(API::Helpers::GlobalIds::Generator) do |generator| - allow(generator).to receive(:generate).with(authorized_user).and_return([global_instance_id, global_user_id]) - end end shared_examples 'code completions endpoint' do @@ -197,31 +212,18 @@ def is_even(n: int) -> context 'when user is logged in' do let(:current_user) { authorized_user } - it_behaves_like 'rate limited endpoint', rate_limit_key: :code_suggestions_api_endpoint do + it_behaves_like 'rate limited and tracked endpoint', + { rate_limit_key: :code_suggestions_api_endpoint, + event_name: 'code_suggestions_rate_limit_exceeded' } do def request post api('/code_suggestions/completions', current_user), headers: headers, params: body.to_json end - - context 'when rate limit is exceeded' do - it 'tracks a code_suggestions_rate_limit_exceeded event' do - allow(Gitlab::ApplicationRateLimiter).to receive(:throttled_request?).and_return(true) - - request - - expect(response).to have_gitlab_http_status(:too_many_requests) - # `error_message` defined in the shared example - expect(response.body).to eq({ message: { error: error_message } }.to_json) - expect(Gitlab::InternalEvents) - .to have_received(:track_event) - .with('code_suggestions_rate_limit_exceeded', user: current_user) - end - end end it 'delegates downstream service call to Workhorse with correct auth token' do post_api - expect(response.status).to be(200) + expect(response).to have_gitlab_http_status(:ok) expect(response.body).to eq("".to_json) command, params = workhorse_send_data expect(command).to eq('send-url') @@ -369,7 +371,6 @@ def request context 'when the instance is Gitlab.org_or_com' do let(:is_saas) { true } - let(:gitlab_realm) { 'saas' } let_it_be(:token) { 'generated-jwt' } let(:headers) do @@ -724,7 +725,7 @@ def get_user(session): it 'does not include additional headers, which are for SaaS only' do post_api - expect(response.status).to be(200) + expect(response).to have_gitlab_http_status(:ok) expect(response.body).to eq("".to_json) _, params = workhorse_send_data expect(params['Header']).not_to have_key('X-Gitlab-Saas-Namespace-Ids') @@ -749,6 +750,50 @@ def get_user(session): end end + describe 'POST /code_suggestions/connection_info' do + subject(:post_api) { post api('/code_suggestions/connection_info', current_user) } + + context 'when unauthorized' do + let(:current_user) { unauthorized_user } + + it_behaves_like 'an unauthorized response' + end + + context 'when authorized' do + let(:current_user) { authorized_user } + + it_behaves_like 'rate limited and tracked endpoint', + { rate_limit_key: :code_suggestions_connection_info, + event_name: 'code_suggestions_connection_info_rate_limit_exceeded' } do + def request + post api('/code_suggestions/connection_info', current_user) + end + end + + it 'returns connection info' do + post_api + + expected_response = { + 'endpoints' => { + 'completions' => { + 'base_url' => 'https://cloud.gitlab.com/ai', + 'expires_at' => 3600, + 'jwt_token' => an_instance_of(String), + 'headers' => { + 'X-Gitlab-Global-User-Id' => global_user_id, + 'X-Gitlab-Instance-Id' => global_instance_id, + 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, + 'X-Gitlab-Realm' => gitlab_realm + } + } + } + } + expect(response).to have_gitlab_http_status(:created) + expect(json_response).to match(expected_response) + end + end + end + context 'when checking if project has duo features enabled' do let_it_be(:enabled_project) { create(:project, :in_group, :private, :with_duo_features_enabled) } let_it_be(:disabled_project) { create(:project, :in_group, :with_duo_features_disabled) } -- GitLab From b11a43d78c65a27f9336f5f097d589f3f4c7cf43 Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Tue, 16 Apr 2024 13:22:25 +0200 Subject: [PATCH 2/7] Minor reference fixes --- .../wip/code_suggestions_direct_completions.yml | 4 ++-- ee/lib/api/code_suggestions.rb | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml b/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml index 7373665f79fa98..40c68bb4cd3b62 100644 --- a/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml +++ b/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml @@ -1,8 +1,8 @@ --- name: code_suggestions_direct_completions feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/452044 -introduced_by_url: -rollout_issue_url: +introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/147246 +rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/456443 milestone: '16.11' group: group::code creation type: wip diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index 5fe39348d84a4f..3a3645c48945e1 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -116,10 +116,11 @@ def saas_headers resources :connection_info do desc 'Code suggestion connection information' do - success code: 200 + success code: 201 failure [ { code: 401, message: 'Unauthorized' }, - { code: 404, message: 'Not found' } + { code: 404, message: 'Not found' }, + { code: 429, message: 'Too many requests' } ] end -- GitLab From 94d04970166205fcae08b089587c762c44ed8e3e Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Tue, 16 Apr 2024 13:36:08 +0200 Subject: [PATCH 3/7] Fix rubocop --- .rubocop_todo/rspec/named_subject.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.rubocop_todo/rspec/named_subject.yml b/.rubocop_todo/rspec/named_subject.yml index 789a831593f6a7..92a87a93252615 100644 --- a/.rubocop_todo/rspec/named_subject.yml +++ b/.rubocop_todo/rspec/named_subject.yml @@ -229,7 +229,6 @@ RSpec/NamedSubject: - 'ee/spec/lib/analytics/merge_request_metrics_refresh_spec.rb' - 'ee/spec/lib/analytics/productivity_analytics_request_params_spec.rb' - 'ee/spec/lib/analytics/refresh_comments_data_spec.rb' - - 'ee/spec/lib/api/entities/code_suggestions_access_token_spec.rb' - 'ee/spec/lib/api/entities/deployments/approval_spec.rb' - 'ee/spec/lib/api/entities/deployments/approval_summary_spec.rb' - 'ee/spec/lib/api/entities/epic_board_spec.rb' -- GitLab From 2927bd9eb42a72184d625740f5da898054b09096 Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Wed, 17 Apr 2024 09:43:42 +0200 Subject: [PATCH 4/7] Refactor response payload Instead of trying to return a generic payload with some compromises, we just return payload which fits our current needs. --- ee/lib/api/code_suggestions.rb | 29 ++++++++---- .../code_suggestions/connection_info.rb | 11 ----- ee/lib/code_suggestions/connection_info.rb | 37 --------------- ee/lib/ee/gitlab/application_rate_limiter.rb | 2 +- .../code_suggestions/connection_info_spec.rb | 47 ------------------- ee/spec/requests/api/code_suggestions_spec.rb | 44 ++++++++++------- 6 files changed, 48 insertions(+), 122 deletions(-) delete mode 100644 ee/lib/api/entities/code_suggestions/connection_info.rb delete mode 100644 ee/lib/code_suggestions/connection_info.rb delete mode 100644 ee/spec/lib/code_suggestions/connection_info_spec.rb diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index 3a3645c48945e1..d978a1fb332e4d 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -28,15 +28,18 @@ def model_gateway_headers(headers, gateway_token) telemetry_headers = headers.select { |k| /\Ax-gitlab-cs-/i.match?(k) } { - 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, 'X-Gitlab-Authentication-Type' => 'oidc', 'Authorization' => "Bearer #{gateway_token}", 'Content-Type' => 'application/json', 'User-Agent' => headers["User-Agent"] # Forward the User-Agent on to the model gateway - }.merge(telemetry_headers).merge(saas_headers).merge(cloud_connector_headers(current_user)) + }.merge(telemetry_headers).merge(saas_headers).merge(connector_headers) .transform_values { |v| Array(v) } end + def connector_headers + cloud_connector_headers(current_user).merge('X-Gitlab-Host-Name' => Gitlab.config.gitlab.host) + end + def saas_headers return {} unless Gitlab.com? @@ -114,8 +117,8 @@ def saas_headers end end - resources :connection_info do - desc 'Code suggestion connection information' do + resources :direct_access do + desc 'Connection details for accessing cloud connector directly' do success code: 201 failure [ { code: 401, message: 'Unauthorized' }, @@ -125,17 +128,27 @@ def saas_headers end post do - check_rate_limit!(:code_suggestions_connection_info, scope: current_user) do + check_rate_limit!(:code_suggestions_direct_access, scope: current_user) do Gitlab::InternalEvents.track_event( - 'code_suggestions_connection_info_rate_limit_exceeded', + 'code_suggestions_direct_access_rate_limit_exceeded', user: current_user ) render_api_error!({ error: _('This endpoint has been requested too many times. Try again later.') }, 429) end - info = ::CodeSuggestions::ConnectionInfo.new(current_user: current_user) - present info, with: ::API::Entities::CodeSuggestions::ConnectionInfo + not_found! unless Feature.enabled?(:code_suggestions_direct_completions, current_user) + + access = { + base_url: ::Gitlab::AiGateway.url, + # for development purposes we just return instance JWT, this should not be used in production + # until we generate a short-term token for user + # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429 + token: ::Gitlab::Llm::AiGateway::Client.access_token(scopes: [:code_suggestions]), + expires_at: ::Gitlab::CloudConnector::SelfIssuedToken::EXPIRES_IN + Time.now.to_i, + headers: connector_headers + } + present access, with: Grape::Presenters::Presenter end end diff --git a/ee/lib/api/entities/code_suggestions/connection_info.rb b/ee/lib/api/entities/code_suggestions/connection_info.rb deleted file mode 100644 index 13e12a5bd5cb76..00000000000000 --- a/ee/lib/api/entities/code_suggestions/connection_info.rb +++ /dev/null @@ -1,11 +0,0 @@ -# frozen_string_literal: true - -module API - module Entities - module CodeSuggestions - class ConnectionInfo < Grape::Entity - expose :endpoints - end - end - end -end diff --git a/ee/lib/code_suggestions/connection_info.rb b/ee/lib/code_suggestions/connection_info.rb deleted file mode 100644 index 3c1a954a8678df..00000000000000 --- a/ee/lib/code_suggestions/connection_info.rb +++ /dev/null @@ -1,37 +0,0 @@ -# frozen_string_literal: true - -module CodeSuggestions - class ConnectionInfo - include ::API::Helpers::CloudConnector - - def initialize(current_user:) - @current_user = current_user - end - - def endpoints - result = {} - - if Feature.enabled?(:code_suggestions_direct_completions, current_user) - result[:completions] = { - base_url: ::Gitlab::AiGateway.url, - # for development purposes we just return instance JWT, this should not be used in production - # until we generate a short-term token for user - # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429 - jwt_token: ::Gitlab::Llm::AiGateway::Client.access_token(scopes: [:code_suggestions]), - expires_at: ::Gitlab::CloudConnector::SelfIssuedToken::EXPIRES_IN, - headers: headers - } - end - - result - end - - private - - attr_reader :current_user, :params - - def headers - cloud_connector_headers(current_user).merge('X-Gitlab-Host-Name' => Gitlab.config.gitlab.host) - end - end -end diff --git a/ee/lib/ee/gitlab/application_rate_limiter.rb b/ee/lib/ee/gitlab/application_rate_limiter.rb index d39c6a7303998e..af0f291131b905 100644 --- a/ee/lib/ee/gitlab/application_rate_limiter.rb +++ b/ee/lib/ee/gitlab/application_rate_limiter.rb @@ -22,7 +22,7 @@ def rate_limits }, credit_card_verification_check_for_reuse: { threshold: 10, interval: 1.hour }, code_suggestions_api_endpoint: { threshold: 60, interval: 1.minute }, - code_suggestions_connection_info: { threshold: 10, interval: 5.minutes }, + code_suggestions_direct_access: { threshold: 10, interval: 5.minutes }, soft_phone_verification_transactions_limit: { threshold: 16_000, interval: 1.day }, hard_phone_verification_transactions_limit: { threshold: 20_000, interval: 1.day } }).freeze diff --git a/ee/spec/lib/code_suggestions/connection_info_spec.rb b/ee/spec/lib/code_suggestions/connection_info_spec.rb deleted file mode 100644 index 8c318e58216452..00000000000000 --- a/ee/spec/lib/code_suggestions/connection_info_spec.rb +++ /dev/null @@ -1,47 +0,0 @@ -# frozen_string_literal: true - -require 'spec_helper' - -RSpec.describe CodeSuggestions::ConnectionInfo, feature_category: :code_suggestions do - describe '#endpoints' do - let(:global_instance_id) { 'instance-ABC' } - let(:global_user_id) { 'user-ABC' } - let(:token) { 'some token' } - let(:user) { create(:user) } - - subject(:endpoints) { described_class.new(current_user: user).endpoints } - - before do - allow_next_instance_of(API::Helpers::GlobalIds::Generator) do |generator| - allow(generator).to receive(:generate).with(user).and_return([global_instance_id, global_user_id]) - end - - allow(Gitlab::Llm::AiGateway::Client).to receive(:access_token).and_return(token) - end - - it 'returns expected endpoints' do - expected_endpoints = { - completions: { - base_url: 'https://cloud.gitlab.com/ai', - expires_at: 3600, - jwt_token: token, - headers: { - 'X-Gitlab-Global-User-Id' => global_user_id, - 'X-Gitlab-Instance-Id' => global_instance_id, - 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, - 'X-Gitlab-Realm' => Gitlab::CloudConnector.gitlab_realm - } - } - } - expect(endpoints).to match(expected_endpoints) - end - - context 'when code_suggestions_direct_completions flag is disabled' do - before do - stub_feature_flags(code_suggestions_direct_completions: false) - end - - it { is_expected.to be_empty } - end - end -end diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index 6ea40ea1a93f18..3987d58579c210 100644 --- a/ee/spec/requests/api/code_suggestions_spec.rb +++ b/ee/spec/requests/api/code_suggestions_spec.rb @@ -750,8 +750,8 @@ def get_user(session): end end - describe 'POST /code_suggestions/connection_info' do - subject(:post_api) { post api('/code_suggestions/connection_info', current_user) } + describe 'POST /code_suggestions/direct_access' do + subject(:post_api) { post api('/code_suggestions/direct_access', current_user) } context 'when unauthorized' do let(:current_user) { unauthorized_user } @@ -763,34 +763,42 @@ def get_user(session): let(:current_user) { authorized_user } it_behaves_like 'rate limited and tracked endpoint', - { rate_limit_key: :code_suggestions_connection_info, - event_name: 'code_suggestions_connection_info_rate_limit_exceeded' } do + { rate_limit_key: :code_suggestions_direct_access, + event_name: 'code_suggestions_direct_access_rate_limit_exceeded' } do def request - post api('/code_suggestions/connection_info', current_user) + post api('/code_suggestions/direct_access', current_user) end end - it 'returns connection info' do + it 'returns direct access details', :freeze_time do post_api expected_response = { - 'endpoints' => { - 'completions' => { - 'base_url' => 'https://cloud.gitlab.com/ai', - 'expires_at' => 3600, - 'jwt_token' => an_instance_of(String), - 'headers' => { - 'X-Gitlab-Global-User-Id' => global_user_id, - 'X-Gitlab-Instance-Id' => global_instance_id, - 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, - 'X-Gitlab-Realm' => gitlab_realm - } - } + 'base_url' => 'https://cloud.gitlab.com/ai', + 'expires_at' => Time.now.to_i + 3600, + 'token' => an_instance_of(String), + 'headers' => { + 'X-Gitlab-Global-User-Id' => global_user_id, + 'X-Gitlab-Instance-Id' => global_instance_id, + 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, + 'X-Gitlab-Realm' => gitlab_realm } } expect(response).to have_gitlab_http_status(:created) expect(json_response).to match(expected_response) end + + context 'when code_suggestions_direct_completions flag is disabled' do + before do + stub_feature_flags(code_suggestions_direct_completions: false) + end + + it 'returns not_found' do + post_api + + expect(response).to have_gitlab_http_status(:not_found) + end + end end end -- GitLab From 63b58822802e7025a797c2de415f0e8b3fd56b6c Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Wed, 17 Apr 2024 14:46:22 +0200 Subject: [PATCH 5/7] Update API docs --- doc/api/code_suggestions.md | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/doc/api/code_suggestions.md b/doc/api/code_suggestions.md index cb7c04fde4cc52..dd41ef57887d31 100644 --- a/doc/api/code_suggestions.md +++ b/doc/api/code_suggestions.md @@ -160,18 +160,14 @@ Example response: ```json { - "endpoints": { - "completions": { - "base_url": "http://0.0.0.0:5052", - "jwt_token": "a valid token", - "expires_at": 3600, - "headers": { - "X-Gitlab-Instance-Id": "292c3c7c-c5d5-48ec-b4bf-f00b724ce560", - "X-Gitlab-Realm": "saas", - "X-Gitlab-Global-User-Id": "Df0Jhs9xlbetQR8YoZCKDZJflhxO0ZBI8uoRzmpnd1w=", - "X-Gitlab-Host-Name": "gitlab.example.com" - } - } + "base_url": "http://0.0.0.0:5052", + "token": "a valid token", + "expires_at": 1713343569, + "headers": { + "X-Gitlab-Instance-Id": "292c3c7c-c5d5-48ec-b4bf-f00b724ce560", + "X-Gitlab-Realm": "saas", + "X-Gitlab-Global-User-Id": "Df0Jhs9xlbetQR8YoZCKDZJflhxO0ZBI8uoRzmpnd1w=", + "X-Gitlab-Host-Name": "gitlab.example.com" } } ``` -- GitLab From 47d230390959ebe0c256e3ba3b492c62a1a94679 Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Thu, 18 Apr 2024 11:28:52 +0200 Subject: [PATCH 6/7] Fix expiration time --- .../code_suggestions_direct_completions.yml | 2 +- ee/lib/api/code_suggestions.rb | 12 +++++- ee/spec/requests/api/code_suggestions_spec.rb | 41 +++++++++++++------ 3 files changed, 41 insertions(+), 14 deletions(-) diff --git a/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml b/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml index 40c68bb4cd3b62..81e1cae452256a 100644 --- a/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml +++ b/ee/config/feature_flags/wip/code_suggestions_direct_completions.yml @@ -3,7 +3,7 @@ name: code_suggestions_direct_completions feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/452044 introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/147246 rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/456443 -milestone: '16.11' +milestone: '17.0' group: group::code creation type: wip default_enabled: false diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index d978a1fb332e4d..72ef91aff9c21d 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -50,6 +50,16 @@ def saas_headers .join(',') } end + + def token_expiration_time + # Because we temporarily use selfissued or instance JWT (not ready for production use) which doesn't expose + # expiration time, expiration time is taken directly from the token record. This helper method is temporary and + # should be removed with + # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429 + return ::CloudConnector::ServiceAccessToken.active.last&.expires_at&.to_i unless Gitlab.org_or_com? + + Time.now.to_i + ::Gitlab::CloudConnector::SelfIssuedToken::EXPIRES_IN + end end namespace 'code_suggestions' do @@ -145,7 +155,7 @@ def saas_headers # until we generate a short-term token for user # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429 token: ::Gitlab::Llm::AiGateway::Client.access_token(scopes: [:code_suggestions]), - expires_at: ::Gitlab::CloudConnector::SelfIssuedToken::EXPIRES_IN + Time.now.to_i, + expires_at: token_expiration_time, headers: connector_headers } present access, with: Grape::Presenters::Presenter diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index 3987d58579c210..9176dc4209038e 100644 --- a/ee/spec/requests/api/code_suggestions_spec.rb +++ b/ee/spec/requests/api/code_suggestions_spec.rb @@ -750,7 +750,7 @@ def get_user(session): end end - describe 'POST /code_suggestions/direct_access' do + describe 'POST /code_suggestions/direct_access', :freeze_time do subject(:post_api) { post api('/code_suggestions/direct_access', current_user) } context 'when unauthorized' do @@ -761,6 +761,20 @@ def get_user(session): context 'when authorized' do let(:current_user) { authorized_user } + let(:expected_expiration) { Time.now.to_i + 3600 } + let(:expected_response) do + { + 'base_url' => 'https://cloud.gitlab.com/ai', + 'expires_at' => expected_expiration, + 'token' => an_instance_of(String), + 'headers' => { + 'X-Gitlab-Global-User-Id' => global_user_id, + 'X-Gitlab-Instance-Id' => global_instance_id, + 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, + 'X-Gitlab-Realm' => gitlab_realm + } + } + end it_behaves_like 'rate limited and tracked endpoint', { rate_limit_key: :code_suggestions_direct_access, @@ -773,21 +787,24 @@ def request it 'returns direct access details', :freeze_time do post_api - expected_response = { - 'base_url' => 'https://cloud.gitlab.com/ai', - 'expires_at' => Time.now.to_i + 3600, - 'token' => an_instance_of(String), - 'headers' => { - 'X-Gitlab-Global-User-Id' => global_user_id, - 'X-Gitlab-Instance-Id' => global_instance_id, - 'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host, - 'X-Gitlab-Realm' => gitlab_realm - } - } expect(response).to have_gitlab_http_status(:created) expect(json_response).to match(expected_response) end + context 'when not SaaS' do + let_it_be(:active_token) { create(:service_access_token, :active) } + let(:is_saas) { false } + let(:expected_expiration) { active_token.expires_at.to_i } + let(:gitlab_realm) { 'self-managed' } + + it 'returns direct access details', :freeze_time do + post_api + + expect(response).to have_gitlab_http_status(:created) + expect(json_response).to match(expected_response) + end + end + context 'when code_suggestions_direct_completions flag is disabled' do before do stub_feature_flags(code_suggestions_direct_completions: false) -- GitLab From 86ff96d6ff83bc36b083dcbecf20fa25407a8359 Mon Sep 17 00:00:00 2001 From: Jan Provaznik Date: Thu, 18 Apr 2024 11:49:32 +0200 Subject: [PATCH 7/7] Address review feedback --- doc/api/code_suggestions.md | 4 ++-- ee/lib/api/code_suggestions.rb | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/doc/api/code_suggestions.md b/doc/api/code_suggestions.md index dd41ef57887d31..95fa6f505a7a05 100644 --- a/doc/api/code_suggestions.md +++ b/doc/api/code_suggestions.md @@ -140,7 +140,7 @@ On GitLab.com and GitLab Dedicated, this feature is not available. This feature is not ready for production use. ```plaintext -POST /code_suggestions/connection_info +POST /code_suggestions/direct_access ``` NOTE: @@ -153,7 +153,7 @@ Example request: ```shell curl --request POST \ --header "Authorization: Bearer " \ - --url "https://gitlab.example.com/api/v4/code_suggestions/connection_info" + --url "https://gitlab.example.com/api/v4/code_suggestions/direct_access" ``` Example response: diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index 72ef91aff9c21d..5405a9ad7d4e66 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -128,7 +128,7 @@ def token_expiration_time end resources :direct_access do - desc 'Connection details for accessing cloud connector directly' do + desc 'Connection details for accessing code suggestions directly' do success code: 201 failure [ { code: 401, message: 'Unauthorized' }, @@ -138,6 +138,8 @@ def token_expiration_time end post do + not_found! unless Feature.enabled?(:code_suggestions_direct_completions, current_user) + check_rate_limit!(:code_suggestions_direct_access, scope: current_user) do Gitlab::InternalEvents.track_event( 'code_suggestions_direct_access_rate_limit_exceeded', @@ -147,8 +149,6 @@ def token_expiration_time render_api_error!({ error: _('This endpoint has been requested too many times. Try again later.') }, 429) end - not_found! unless Feature.enabled?(:code_suggestions_direct_completions, current_user) - access = { base_url: ::Gitlab::AiGateway.url, # for development purposes we just return instance JWT, this should not be used in production -- GitLab