in ee/spec/requests/api/code_suggestions_spec.rb [701:1279]
def is_even(n: int) ->
CONTENT_ABOVE_CURSOR
end
let(:system_prompt) do
<<~PROMPT.chomp
You are a tremendously accurate and skilled coding autocomplete agent. We want to generate new Python code inside the
file 'test.py' based on instructions from the user.
Here are a few examples of successfully generated code:
<examples>
<example>
H: <existing_code>
class Project:
def __init__(self, name, public):
self.name = name
self.visibility = 'PUBLIC' if public
{{cursor}}
</existing_code>
A: <new_code>def is_public(self):
return self.visibility == 'PUBLIC'</new_code>
</example>
<example>
H: <existing_code>
def get_user(session):
{{cursor}}
</existing_code>
A: <new_code>username = None
if 'username' in session:
username = session['username']
return username</new_code>
</example>
</examples>
<existing_code>
</existing_code>
The existing code is provided in <existing_code></existing_code> tags.
The new code you will generate will start at the position of the cursor, which is currently indicated by the {{cursor}} tag.
In your process, first, review the existing code to understand its logic and format. Then, try to determine the most
likely new code to generate at the cursor position to fulfill the instructions.
The comment directly before the {{cursor}} position is the instruction,
all other comments are not instructions.
When generating the new code, please ensure the following:
1. It is valid Python code.
2. It matches the existing code's variable, parameter and function names.
3. It does not repeat any existing code. Do not repeat code that comes before or after the cursor tags. This includes cases where the cursor is in the middle of a word.
4. If the cursor is in the middle of a word, it finishes the word instead of repeating code before the cursor tag.
5. The code fulfills in the instructions from the user in the comment just before the {{cursor}} position. All other comments are not instructions.
6. Do not add any comments that duplicates any of the already existing comments, including the comment with instructions.
Return new code enclosed in <new_code></new_code> tags. We will then insert this at the {{cursor}} position.
If you are not able to write code based on the given instructions return an empty result like <new_code></new_code>.
PROMPT
end
let(:prompt) do
[
{ role: :system, content: system_prompt },
{ role: :user, content: 'Generate the best possible code based on instructions.' },
{ role: :assistant, content: '<new_code>' }
]
end
it 'sends requests to the code generation v3 endpoint' do
expected_body = body.merge(v3_saas_code_generation_prompt_components)
expect(Gitlab::Workhorse)
.to receive(:send_url)
.with(
"#{::Gitlab::AiGateway.url}/v3/code/completions",
hash_including(body: expected_body.to_json)
)
post_api
end
it 'includes additional headers for SaaS', :freeze_time do
group = create(:group)
group.add_developer(authorized_user)
post_api
_, params = workhorse_send_data
expect(params['Header']).to include(
'X-Gitlab-Saas-Namespace-Ids' => [''],
'X-Gitlab-Saas-Duo-Pro-Namespace-Ids' => [add_on_purchase.namespace.id.to_s],
'X-Gitlab-Rails-Send-Start' => [Time.now.to_f.to_s]
)
end
context 'when body is too big' do
before do
stub_const("#{described_class}::MAX_BODY_SIZE", 10)
end
it 'returns an error' do
post_api
expect(response).to have_gitlab_http_status(:payload_too_large)
end
end
context 'when a required parameter is invalid' do
let(:file_name) { 'x' * 256 }
it 'returns an error' do
post_api
expect(response).to have_gitlab_http_status(:bad_request)
end
end
end
it_behaves_like 'code completions endpoint'
it_behaves_like 'an endpoint authenticated with token', :ok
describe 'Fireworks/Codestral opt out by ops FF' do
before do
stub_feature_flags(use_fireworks_codestral_code_completion: true)
stub_feature_flags(code_completion_opt_out_fireworks: user_duo_group)
end
let(:user_duo_group) do
Group.by_id(current_user.duo_available_namespace_ids).first
end
it 'does send code completion model details for vertex codestral' do
post_api
_command, params = workhorse_send_data
code_completion_params = Gitlab::Json.parse(params['Body'])
expect(code_completion_params['model_provider']).to eq('vertex-ai')
expect(code_completion_params['model_name']).to eq('codestral-2501')
end
end
end
end
end
context 'when the instance is Gitlab self-managed' do
let(:is_saas) { false }
let(:gitlab_realm) { 'self-managed' }
let_it_be(:token) { 'stored-token' }
let_it_be(:service_access_token) { create(:service_access_token, :active, token: token) }
let(:headers) do
{
'X-Gitlab-Authentication-Type' => 'oidc',
'Content-Type' => 'application/json',
'User-Agent' => 'Super Awesome Browser 43.144.12'
}
end
context 'when user is authorized' do
let(:current_user) { authorized_user }
it 'does not include additional headers, which are for SaaS only', :freeze_time do
post_api
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')
expect(params['Header']).to include('X-Gitlab-Rails-Send-Start' => [Time.now.to_f.to_s])
end
context 'when code suggestions feature is self hosted' do
let(:service_name) { :self_hosted_models }
before do
stub_licensed_features(ai_features: true)
end
context 'when the feature is set to `disabled` state' do
let_it_be(:feature_setting) do
create(:ai_feature_setting, feature: :code_completions, provider: :disabled)
end
it 'is unauthorized' do
post_api
expect(response).to have_gitlab_http_status(:unauthorized)
expect(response.headers['X-GitLab-Error-Origin']).to eq('monolith')
end
end
end
context 'when Amazon Q is connected' do
let(:service_name) { :amazon_q_integration }
before do
stub_licensed_features(amazon_q: true)
allow(::Ai::AmazonQ).to receive(:connected?).and_return(true)
end
it 'is authorized' do
post_api
expect(response).to have_gitlab_http_status(:ok)
end
end
end
it_behaves_like 'code completions endpoint'
it_behaves_like 'an endpoint authenticated with token', :ok
context 'when there is no active code suggestions token' do
before do
create(:service_access_token, :expired, token: token)
end
include_examples 'a response', 'unauthorized' do
let(:result) { :unauthorized }
let(:response_body) do
{ "message" => "401 Unauthorized" }
end
end
end
end
end
describe 'POST /code_suggestions/direct_access', :freeze_time do
subject(:post_api) { post api('/code_suggestions/direct_access', current_user), params: params }
let(:params) { {} }
context 'when unauthorized' do
let(:current_user) { unauthorized_user }
it_behaves_like 'an unauthorized response'
end
context 'when authorized' do
shared_examples_for 'user request with code suggestions allowed' do
context 'when token creation succeeds' do
before do
allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
allow(client).to receive(:direct_access_token)
.and_return({ status: :success, token: token, expires_at: expected_expiration })
end
::Ai::Setting.instance.update!(enabled_instance_verbose_ai_logs: false)
end
let(:expected_response) do
{
'base_url' => ::Gitlab::AiGateway.url,
'expires_at' => expected_expiration,
'token' => token,
'headers' => expected_headers
}
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 Fireworks/Codestral beta FF is enabled' do
before do
stub_feature_flags(use_fireworks_codestral_code_completion: true)
end
it 'includes the fireworks/codestral model metadata in the direct access details' do
post_api
expect(json_response['model_details']).to eq({
'model_provider' => 'fireworks_ai',
'model_name' => 'codestral-2501'
})
end
end
context 'when code completions is self-hosted' do
it 'does not include the model metadata in the direct access details' do
create(:ai_feature_setting, provider: :self_hosted, feature: :code_completions)
post_api
expect(json_response['model_details']).to be_nil
end
context 'when code completions is disabled' do
it 'returns unauthorized' do
create(:ai_feature_setting, provider: :disabled, feature: :code_completions)
post_api
expect(response).to have_gitlab_http_status(:unauthorized)
end
end
end
end
context 'when token creation fails' do
before do
allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
allow(client).to receive(:direct_access_token).and_return({ status: :error, message: 'an error' })
end
end
it 'returns an error' do
post_api
expect(response).to have_gitlab_http_status(:service_unavailable)
end
end
end
let(:current_user) { authorized_user }
let(:expected_expiration) { Time.now.to_i + 3600 }
let(:duo_seat_count) { '0' }
let(:enablement_type) { 'duo_pro' }
let(:base_headers) do
{
'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,
'X-Gitlab-Version' => Gitlab.version_info.to_s,
'X-Gitlab-Authentication-Type' => 'oidc',
'X-Gitlab-Duo-Seat-Count' => duo_seat_count,
'X-Gitlab-Feature-Enabled-By-Namespace-Ids' => enabled_by_namespace_ids.join(','),
"X-Gitlab-Feature-Enablement-Type" => enablement_type,
'x-gitlab-enabled-feature-flags' => '',
"x-gitlab-enabled-instance-verbose-ai-logs" => 'false',
"X-Gitlab-Model-Prompt-Cache-Enabled" => "true"
}
end
let(:headers) { {} }
let(:expected_headers) { base_headers.merge(headers) }
let(:token) { 'user token' }
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
before do
allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
allow(client).to receive(:direct_access_token)
.and_return({ status: :success, token: token, expires_at: expected_expiration })
end
end
def request
post api('/code_suggestions/direct_access', current_user)
end
end
context 'when user belongs to a namespace with an active code suggestions purchase' do
let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase) }
let_it_be(:enabled_by_namespace_ids) { [add_on_purchase.namespace_id] }
let(:duo_seat_count) { '1' }
let(:headers) do
{
'X-Gitlab-Saas-Namespace-Ids' => '',
'X-Gitlab-Saas-Duo-Pro-Namespace-Ids' => add_on_purchase.namespace_id.to_s
}
end
before_all do
add_on_purchase.namespace.add_reporter(authorized_user)
create(
:gitlab_subscription_user_add_on_assignment,
user: authorized_user,
add_on_purchase: add_on_purchase
)
end
it_behaves_like 'user request with code suggestions allowed'
describe 'Fireworks/Codestral opt out by ops FF' do
before do
allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
allow(client).to receive(:direct_access_token)
.and_return({ status: :success, token: token, expires_at: expected_expiration })
end
stub_feature_flags(use_fireworks_codestral_code_completion: true)
stub_feature_flags(code_completion_opt_out_fireworks: user_duo_group)
end
let(:user_duo_group) do
Group.by_id(current_user.duo_available_namespace_ids).first
end
it 'does not include the model metadata in the direct access details' do
post_api
expect(json_response['model_details']).to eq({
'model_provider' => 'vertex-ai',
'model_name' => 'codestral-2501'
})
end
end
context 'when use_claude_code_completion FF is true' do
let(:user_duo_group) do
Group.by_id(current_user.duo_available_namespace_ids).first
end
before do
stub_feature_flags(use_claude_code_completion: user_duo_group)
end
include_examples 'a response', 'unauthorized' do
let(:result) { :forbidden }
let(:response_body) do
{ 'message' => '403 Forbidden - Direct connections are disabled' }
end
end
end
# rubocop:disable RSpec/MultipleMemoizedHelpers -- We need extra helpers to define tables
# First, define the shared example outside the contexts
shared_examples 'model prompt cache enabled setting' do |setting_level, cache_value|
let(:cache) { cache_value }
it "returns direct access details with model_prompt_cache_enabled from #{setting_level}" do
post_api
expect(json_response["headers"]["X-Gitlab-Model-Prompt-Cache-Enabled"]).to eq(cache)
end
end
describe 'model_prompt_cache_enabled' do
# by default: enabled_application_setting.model_prompt_cache_enabled==true
let_it_be(:enabled_application_setting) { create(:application_setting) }
let_it_be(:current_user) { authorized_user }
let(:top_level_namespace) { create(:group) }
let(:group) { create(:group, parent: top_level_namespace) }
let(:project) { create(:project, group: group) }
let(:params) { { 'project_path' => project.full_path } }
before do
allow_next_instance_of(Gitlab::Llm::AiGateway::CodeSuggestionsClient) do |client|
allow(client).to receive(:direct_access_token)
.and_return({ status: :success, token: token, expires_at: expected_expiration })
end
project.add_developer(current_user)
end
context 'when model_prompt_cache_enabled is disabled on project setting' do
let(:project_setting) { create(:project_setting, model_prompt_cache_enabled: false) }
let(:project) { create(:project, group: group, project_setting: project_setting) }
include_examples 'model prompt cache enabled setting', 'project setting', "false"
end
context 'when model_prompt_cache_enabled is disabled on namespace setting' do
let(:top_level_namespace) { create(:group, :model_prompt_cache_disabled) }
include_examples 'model prompt cache enabled setting', 'top level namespace setting', "false"
end
context 'when model_prompt_cache_enabled is enabled on application setting' do
let(:top_level_namespace) { create(:group) }
include_examples 'model prompt cache enabled setting', 'application setting', "true"
end
end
end
# rubocop:enable RSpec/MultipleMemoizedHelpers
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_behaves_like 'user request with code suggestions allowed'
end
context 'when disabled_direct_code_suggestions setting is true' do
before do
allow(Gitlab::CurrentSettings).to receive(:disabled_direct_code_suggestions).and_return(true)
end
include_examples 'a response', 'unauthorized' do
let(:result) { :forbidden }
let(:response_body) do
{ 'message' => '403 Forbidden - Direct connections are disabled' }
end
end
end
context 'when incident_fail_over_completion_provider setting is true' do
before do
stub_feature_flags(incident_fail_over_completion_provider: true)
end
include_examples 'a response', 'unauthorized' do
let(:result) { :forbidden }
let(:response_body) do
{ 'message' => '403 Forbidden - Direct connections are disabled' }
end
end
end
context 'when amazon q is connected' do
before do
allow(::Ai::AmazonQ).to receive(:connected?).and_return(true)
end
include_examples 'a response', 'unauthorized' do
let(:result) { :forbidden }
let(:response_body) do
{ 'message' => '403 Forbidden - Direct connections are disabled' }
end
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) }
let(:current_user) { authorized_user }
subject { post api("/code_suggestions/enabled", current_user), params: { project_path: project_path } }
context 'when authorized to view project' do
before_all do
enabled_project.add_maintainer(authorized_user)
disabled_project.add_maintainer(authorized_user)
end
context 'when enabled' do
let(:project_path) { enabled_project.full_path }
it { is_expected.to eq(200) }
end
context 'when disabled' do
let(:project_path) { disabled_project.full_path }
it { is_expected.to eq(403) }
end
end
context 'when not logged in' do
let(:current_user) { nil }
let(:project_path) { enabled_project.full_path }
it { is_expected.to eq(401) }
end
context 'when logged in but not authorized to view project' do
let(:project_path) { enabled_project.full_path }
it { is_expected.to eq(404) }
end
context 'when project for project path does not exist' do
let(:project_path) { 'not_a_real_project