diff --git a/config/feature_flags/gitlab_com_derisk/require_resource_id.yml b/config/feature_flags/gitlab_com_derisk/require_resource_id.yml new file mode 100644 index 0000000000000000000000000000000000000000..8ad470624c8b198bd3a15d9e732fef5d84a19ff6 --- /dev/null +++ b/config/feature_flags/gitlab_com_derisk/require_resource_id.yml @@ -0,0 +1,9 @@ +--- +name: require_resource_id +feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/463046 +introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/154115 +rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/465363 +milestone: '17.1' +group: group::ai framework +type: gitlab_com_derisk +default_enabled: false diff --git a/ee/app/services/llm/base_service.rb b/ee/app/services/llm/base_service.rb index f19a294c0866ea9f95ab95425caeb91a823f219c..d5b69635fad36231726776a2f1c11e150430e250 100644 --- a/ee/app/services/llm/base_service.rb +++ b/ee/app/services/llm/base_service.rb @@ -3,6 +3,7 @@ module Llm class BaseService INVALID_MESSAGE = 'AI features are not enabled or resource is not permitted to be sent.' + MISSING_RESOURCE_ID_MESSAGE = 'ResourceId is required for slash command request.' def initialize(user, resource, options = {}) @user = user @@ -17,6 +18,11 @@ def execute return error(INVALID_MESSAGE) end + if Feature.enabled?(:require_resource_id, @group) && invalid_slash_command_request? + logger.info(message: "Returning from Service due to missing resource id. Associated group: #{@group}") + return error(MISSING_RESOURCE_ID_MESSAGE) + end + result = perform result.is_a?(ServiceResponse) ? result : success(ai_message: prompt_message) @@ -28,6 +34,11 @@ def valid? ai_integration_enabled? && user_can_send_to_ai? end + def invalid_slash_command_request? + true if contains_ai_action? && prompt_message.slash_command_prompt? && + !prompt_message.resource.present? + end + private attr_reader :user, :resource, :options, :logger @@ -46,6 +57,10 @@ def user_can_send_to_ai? user.any_group_with_ai_available? end + def contains_ai_action? + options.key?(:ai_action) + end + def prompt_message @prompt_message ||= build_prompt_message end diff --git a/ee/lib/gitlab/llm/ai_message.rb b/ee/lib/gitlab/llm/ai_message.rb index 3c9cc43d4f73b611d567fab5dfe74c852bc68695..5c4f48366dd9019b412c3c007aab29f1aff4bf46 100644 --- a/ee/lib/gitlab/llm/ai_message.rb +++ b/ee/lib/gitlab/llm/ai_message.rb @@ -16,6 +16,12 @@ class AiMessage :agent_version_id, :referer_url ].freeze + SLASH_COMMAND_TOOLS = [ + ::Gitlab::Llm::Chain::Tools::ExplainCode, + ::Gitlab::Llm::Chain::Tools::WriteTests, + ::Gitlab::Llm::Chain::Tools::RefactorCode + ].freeze + attr_accessor(*ATTRIBUTES_LIST) delegate :resource, to: :context @@ -73,6 +79,18 @@ def slash_command? content.to_s.match?(%r{\A/\w}) end + def slash_command_prompt? + false unless slash_command? + + command, _ = slash_command_and_input + + return false unless SLASH_COMMAND_TOOLS.find do |tool| + tool::Executor.slash_commands.has_key?(command) + end + + true + end + def slash_command_and_input return [] unless slash_command? diff --git a/ee/spec/services/llm/base_service_spec.rb b/ee/spec/services/llm/base_service_spec.rb index d4ec2e521c4461881106cedc7ac0c34da4301399..15e8c365b1bca6eda2e39930c93203334ca5fdae 100644 --- a/ee/spec/services/llm/base_service_spec.rb +++ b/ee/spec/services/llm/base_service_spec.rb @@ -18,6 +18,15 @@ end end + shared_examples 'returns a missing resource error' do + it 'returns a missing resource error' do + result = subject.execute + + expect(result).to be_error + expect(result.message).to eq(described_class::MISSING_RESOURCE_ID_MESSAGE) + end + end + shared_examples 'raises a NotImplementedError' do it 'raises a NotImplementedError' do expect { subject.execute }.to raise_error(NotImplementedError) @@ -42,6 +51,34 @@ def ai_action end end + shared_examples 'success when implemented with slash command' do + subject do + Class.new(described_class) do + def perform + schedule_completion_worker + end + end.new(user, resource, options) + end + + it_behaves_like 'schedules completion worker' do + let(:action_name) { "/explain def" } + end + end + + shared_examples 'success when implemented with invalid slash command' do + subject do + Class.new(described_class) do + def perform + schedule_completion_worker + end + end.new(user, resource, options) + end + + it_behaves_like 'schedules completion worker' do + let(:action_name) { "/where can credentials be set" } + end + end + shared_examples 'authorizing a resource' do let(:authorizer_response) { instance_double(Gitlab::Llm::Utils::Authorizer::Response, allowed?: allowed) } @@ -102,6 +139,10 @@ def ai_action end context 'when ai features are enabled' do + before do + stub_feature_flags(require_resource_id: false) + end + include_context 'with ai features enabled for group' it_behaves_like 'raises a NotImplementedError' @@ -123,6 +164,47 @@ def ai_action it_behaves_like 'success when implemented' end + + context 'when require_resource_id FF is enabled' do + context 'when resource is missing' do + let(:resource) { nil } + let(:options) { { ai_action: "/explain def" } } + + before do + stub_feature_flags(require_resource_id: true) + end + + it_behaves_like 'returns a missing resource error' + end + + context 'when non slash command request starts with a slash' do + let(:resource) { nil } + let(:options) { { ai_action: "/where can credentials be set" } } + + before do + stub_feature_flags(require_resource_id: true) + end + + it_behaves_like 'success when implemented with invalid slash command' + end + + context 'when non slash command request is received' do + let(:resource) { nil } + + before do + stub_feature_flags(require_resource_id: true) + end + + it_behaves_like 'success when implemented' + end + end + + context 'when resource is missing and require_resource_id FF is disabled, slash command request ' do + let(:resource) { nil } + let(:options) { { ai_action: "/explain def" } } + + it_behaves_like 'success when implemented with slash command' + end end end end