diff --git a/.rubocop_todo/rspec/named_subject.yml b/.rubocop_todo/rspec/named_subject.yml index 789a831593f6a7af86e3747607c58ef3556ba673..92a87a932526155294288f75fae7f08642acbae4 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' diff --git a/doc/api/code_suggestions.md b/doc/api/code_suggestions.md index 4ce12887ad67b54f10cc2707452cde907f6990ac..95fa6f505a7a0545f71b748073b733ec355ec3a7 100644 --- a/doc/api/code_suggestions.md +++ b/doc/api/code_suggestions.md @@ -126,3 +126,48 @@ 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/direct_access +``` + +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/direct_access" +``` + +Example response: + +```json +{ + "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" + } +} +``` 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 0000000000000000000000000000000000000000..81e1cae452256a191bfbeafe6074230fb7547314 --- /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: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/147246 +rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/456443 +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 b1e09d1f1f254b2a483d8be27ca3d51d7501dc46..5405a9ad7d4e669787a65e84da10b83207baaebb 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? @@ -47,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 @@ -114,6 +127,41 @@ def saas_headers end end + resources :direct_access do + desc 'Connection details for accessing code suggestions directly' do + success code: 201 + failure [ + { code: 401, message: 'Unauthorized' }, + { code: 404, message: 'Not found' }, + { code: 429, message: 'Too many requests' } + ] + 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', + user: current_user + ) + + render_api_error!({ error: _('This endpoint has been requested too many times. Try again later.') }, 429) + end + + 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: token_expiration_time, + headers: connector_headers + } + present access, with: Grape::Presenters::Presenter + 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_access_token.rb b/ee/lib/api/entities/code_suggestions_access_token.rb deleted file mode 100644 index b8a4c110c541b417d42e07b76518f7509cde472f..0000000000000000000000000000000000000000 --- 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/ee/gitlab/application_rate_limiter.rb b/ee/lib/ee/gitlab/application_rate_limiter.rb index e049b8122a69885198595f81b0755d9cde7516b1..af0f291131b90557b7fed0900536f9f020c78d27 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_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/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 a04fcc22b7ccd3bff181eb0b9eed7adbd8abf30f..0000000000000000000000000000000000000000 --- 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/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index 7b71dd8d9eccb7f9535b90b5e10f1325d02146ed..9176dc4209038e9d0a7dd12015dda675420898a0 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,75 @@ def get_user(session): end end + describe 'POST /code_suggestions/direct_access', :freeze_time do + subject(:post_api) { post api('/code_suggestions/direct_access', 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 } + 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, + event_name: 'code_suggestions_direct_access_rate_limit_exceeded' } do + def request + post api('/code_suggestions/direct_access', current_user) + end + end + + 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 + + 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) + end + + it 'returns not_found' do + post_api + + expect(response).to have_gitlab_http_status(:not_found) + end + 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) }