diff --git a/ee/app/services/ai/duo_workflows/sanitize_job_variables_service.rb b/ee/app/services/ai/duo_workflows/sanitize_job_variables_service.rb new file mode 100644 index 0000000000000000000000000000000000000000..a12ec8e76902abed108b491c07439dfab35b2ee4 --- /dev/null +++ b/ee/app/services/ai/duo_workflows/sanitize_job_variables_service.rb @@ -0,0 +1,38 @@ +# frozen_string_literal: true + +module Ai + module DuoWorkflows + class SanitizeJobVariablesService + OVERRIDABLE_JOB_VARIABLES = { + GIT_DEPTH: 1 + }.freeze + + def initialize(dirty_variables) + @dirty_variables = (dirty_variables.presence || {}).symbolize_keys + end + + def execute + invalid_variables = dirty_variables.keys - OVERRIDABLE_JOB_VARIABLES.keys + + if invalid_variables.any? + return ServiceResponse.error( + message: <<~MESSAGE.squish, + Non-overridable variables detected: #{invalid_variables.to_sentence}. + Allowed variables: #{OVERRIDABLE_JOB_VARIABLES.keys.to_sentence}. + MESSAGE + reason: :unprocessable_entity) + end + + sanitized_variables = OVERRIDABLE_JOB_VARIABLES.each_with_object({}) do |(key, value), variables| + variables[key] = dirty_variables.fetch(key, value) + end + + ServiceResponse.success(payload: { variables: sanitized_variables }) + end + + private + + attr_reader :dirty_variables + end + end +end diff --git a/ee/app/services/ai/duo_workflows/start_workflow_service.rb b/ee/app/services/ai/duo_workflows/start_workflow_service.rb index 117e2003509841f056c8a4ca322a1065356d2832..cf017bcd37b22c77f42425f57a400bd0123b5e66 100644 --- a/ee/app/services/ai/duo_workflows/start_workflow_service.rb +++ b/ee/app/services/ai/duo_workflows/start_workflow_service.rb @@ -3,25 +3,19 @@ module Ai module DuoWorkflows class StartWorkflowService + include ::Gitlab::Utils::StrongMemoize + IMAGE = 'registry.gitlab.com/gitlab-org/duo-workflow/default-docker-image/workflow-generic-image:v0.0.4' def initialize(workflow:, params:) @workflow = workflow @current_user = workflow.user @params = params + @job_variables = (params.delete(:job_variables) || {}).symbolize_keys end def execute - unless @workflow.project_level? - return ServiceResponse.error( - message: 'Only project-level workflow is supported', - reason: :unprocessable_entity) - end - - unless @current_user.can?(:execute_duo_workflow_in_ci, @workflow) - return ServiceResponse.error(message: 'Can not execute workflow in CI', - reason: :feature_unavailable) - end + return validation if validation.error? workload_user = @current_user @@ -55,6 +49,33 @@ def execute private + def validation + sanitized_variables = ::Ai::DuoWorkflows::SanitizeJobVariablesService.new(@job_variables).execute + + if sanitized_variables.error? + return ServiceResponse.error( + message: sanitized_variables.message, + reason: :unprocessable_entity) + end + + unless @workflow.project_level? + return ServiceResponse.error( + message: 'Only project-level workflow is supported', + reason: :unprocessable_entity + ) + end + + unless @current_user.can?(:execute_duo_workflow_in_ci, @workflow) + return ServiceResponse.error( + message: 'Can not execute workflow in CI', + reason: :feature_unavailable + ) + end + + ServiceResponse.success + end + strong_memoize_attr :validation + def workload_definition ::Ci::Workloads::WorkloadDefinition.new do |d| d.image = @workflow.image.presence || configured_image || IMAGE @@ -86,6 +107,11 @@ def setup_script_commands duo_config.setup_script || [] end + def sanitized_variables + ::Ai::DuoWorkflows::SanitizeJobVariablesService.new(@job_variables).execute + end + strong_memoize_attr :validation + def variables { DUO_WORKFLOW_ADDITIONAL_CONTEXT_CONTENT: serialized_flow_additional_context, @@ -110,7 +136,7 @@ def variables DUO_WORKFLOW_METADATA: @params[:workflow_metadata], GITLAB_BASE_URL: Gitlab.config.gitlab.url, AGENT_PLATFORM_GITLAB_VERSION: Gitlab.version_info.to_s - } + }.merge(sanitized_variables.payload.fetch(:variables)) end def commands diff --git a/ee/lib/api/ai/duo_workflows/workflows.rb b/ee/lib/api/ai/duo_workflows/workflows.rb index 9174dc193bd0bd8983b27f5d4df9a0a990a6aad1..e5f22aa37e929db14c4be4ba01be8e0eb86b79fa 100644 --- a/ee/lib/api/ai/duo_workflows/workflows.rb +++ b/ee/lib/api/ai/duo_workflows/workflows.rb @@ -17,6 +17,8 @@ class Workflows < ::API::Base end helpers do + include Gitlab::Utils::StrongMemoize + def find_root_namespace # The IDE sends namespace data via the header, while the web agentic chat UI # sends it as a query param. @@ -78,7 +80,8 @@ def start_workflow_params(workflow_id) use_service_account: token_service.use_service_account?, source_branch: params[:source_branch], additional_context: params[:additional_context], - workflow_metadata: Gitlab::DuoWorkflow::Client.metadata(current_user).to_json + workflow_metadata: Gitlab::DuoWorkflow::Client.metadata(current_user).to_json, + job_variables: sanitized_variables.payload.fetch(:variables) } end @@ -121,8 +124,13 @@ def duo_workflow_list_tools end def create_workflow_params - wrkf_params = declared_params(include_missing: false).except(:start_workflow, :source_branch, - :additional_context) + wrkf_params = declared_params(include_missing: false).except( + :start_workflow, + :source_branch, + :additional_context, + :job_variables + ) + return wrkf_params unless wrkf_params[:ai_catalog_item_version_id] wrkf_params[:ai_catalog_item_version] = ::Ai::Catalog::ItemVersion @@ -130,6 +138,11 @@ def create_workflow_params wrkf_params end + def sanitized_variables + ::Ai::DuoWorkflows::SanitizeJobVariablesService.new(params[:job_variables]).execute + end + strong_memoize_attr :sanitized_variables + params :workflow_params do optional :project_id, type: String, desc: 'The ID or path of the workflow project', documentation: { example: '1' } @@ -179,6 +192,9 @@ def create_workflow_params example: '[{"Category": "agent_user_environment", "Content": "{\"merge_request_url\": ' \ 'https://gitlab.com/project/-/merge_requests/1\"}", "Metadata": "{}"}]' } + optional :job_variables, type: Hash, + desc: 'Overridable environment variables to pass to the workflow job.', + documentation: { example: '{"GIT_DEPTH": 30}' } exactly_one_of :project_id, :namespace_id end end @@ -317,6 +333,8 @@ def create_workflow_params 'https://gitlab.com/gitlab-org/gitlab/-/issues/566195', new_threshold: 105 ) + bad_request!(sanitized_variables.message) if sanitized_variables.error? + container = if params[:project_id] find_project!(params[:project_id]) elsif params[:namespace_id] diff --git a/ee/spec/requests/api/ai/duo_workflows/workflows_spec.rb b/ee/spec/requests/api/ai/duo_workflows/workflows_spec.rb index fc4975b28a10c8e805602054bbcad03b4594b6c0..0bad54378180e418d4a27955aaa7171efa853e73 100644 --- a/ee/spec/requests/api/ai/duo_workflows/workflows_spec.rb +++ b/ee/spec/requests/api/ai/duo_workflows/workflows_spec.rb @@ -573,6 +573,46 @@ expect(response).to have_gitlab_http_status(:bad_request) end end + + context 'when valid job_variables are provided' do + let(:job_variables) do + { 'GIT_DEPTH' => '15' } + end + + let(:params) do + super().merge(job_variables: job_variables) + end + + it 'passes job_variables to StartWorkflowService' do + expect(::Ai::DuoWorkflows::StartWorkflowService).to receive(:new).with( + workflow: anything, + params: hash_including(job_variables: { GIT_DEPTH: '15' }) + ).and_call_original + + post api(path, user), params: params + + expect(response).to have_gitlab_http_status(:created) + end + end + + context 'when invalid job_variables are provided' do + let(:job_variables) do + { 'GIT_STRATEGY' => 'fetch' } + end + + let(:params) do + super().merge(job_variables: job_variables) + end + + it 'returns an API error' do + post api(path, user), params: params + + expect(response).to have_gitlab_http_status(:bad_request) + expect(response.body).to include( + 'Non-overridable variables detected: GIT_STRATEGY. Allowed variables: GIT_DEPTH' + ) + end + end end end diff --git a/ee/spec/services/ai/duo_workflows/sanitize_job_variables_service_spec.rb b/ee/spec/services/ai/duo_workflows/sanitize_job_variables_service_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..047a655a33dafafba42a1a40b7ee7e3c9ca3a2d6 --- /dev/null +++ b/ee/spec/services/ai/duo_workflows/sanitize_job_variables_service_spec.rb @@ -0,0 +1,55 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Ai::DuoWorkflows::SanitizeJobVariablesService, feature_category: :agent_foundations do + subject(:result) { described_class.new(dirty_variables).execute } + + let(:variables) { result.payload.fetch(:variables) } + let(:default_values) { described_class::OVERRIDABLE_JOB_VARIABLES } + + context 'when dirty variables are invalid' do + let(:dirty_variables) { { GIT_STRATEGY: 'fetch' } } + + it 'returns an error' do + expect(result).to be_error + expect(result.message).to eq('Non-overridable variables detected: GIT_STRATEGY. Allowed variables: GIT_DEPTH.') + end + end + + context 'when dirty variables are valid' do + let(:dirty_variables) { { GIT_DEPTH: 15 } } + + it 'returns the sanitized variables' do + expect(result).to be_success + expect(variables).to eq(GIT_DEPTH: 15) + end + end + + context 'when dirty variables keys are string' do + let(:dirty_variables) { { 'GIT_DEPTH' => 15 } } + + it 'sanitize keys to symbols' do + expect(result).to be_success + expect(variables).to eq(GIT_DEPTH: 15) + end + end + + context 'when dirty variables are nil' do + let(:dirty_variables) { nil } + + it 'returns the default values' do + expect(result).to be_success + expect(variables).to eq(default_values) + end + end + + context 'when dirty variables are empty' do + let(:dirty_variables) { {} } + + it 'returns the default values' do + expect(result).to be_success + expect(variables).to eq(default_values) + end + end +end diff --git a/ee/spec/services/ai/duo_workflows/start_workflow_service_spec.rb b/ee/spec/services/ai/duo_workflows/start_workflow_service_spec.rb index 0290ee3a7fdf1bd24ba947074346caf04fd64977..d9621378a47df0a258b4d28592359c3470a69ba1 100644 --- a/ee/spec/services/ai/duo_workflows/start_workflow_service_spec.rb +++ b/ee/spec/services/ai/duo_workflows/start_workflow_service_spec.rb @@ -170,6 +170,13 @@ end end + shared_context 'with Duo enabled' do + before do + allow(::Gitlab::Llm::StageCheck).to receive(:available?).with(project, :duo_workflow).and_return(true) + allow(maintainer).to receive(:allowed_to_use?).and_return(true) + end + end + subject(:execute) { described_class.new(workflow: workflow, params: params).execute } context 'with workflow enablement checks' do @@ -787,4 +794,59 @@ end end end + + context 'without custom job variables', :aggregate_failures do + include_context 'with Duo enabled' + + it 'uses default variables' do + expect(Ci::Workloads::RunWorkloadService).to receive(:new).and_wrap_original do |method, **kwargs| + workload_definition = kwargs[:workload_definition] + variables = workload_definition.variables + + expect(variables[:GIT_DEPTH]).to eq(1) + + method.call(**kwargs) + end + + expect(execute).to be_success + end + end + + context 'with valid custom job variables' do + include_context 'with Duo enabled' + + let(:params) do + super().merge(job_variables: { 'GIT_DEPTH' => 15 }) + end + + it 'merges with default variables' do + expect(Ci::Workloads::RunWorkloadService).to receive(:new).and_wrap_original do |method, **kwargs| + workload_definition = kwargs[:workload_definition] + variables = workload_definition.variables + + expect(variables[:GIT_DEPTH]).to eq(15) + + method.call(**kwargs) + end + + expect(execute).to be_success + end + end + + context 'with invalid custom job variables' do + include_context 'with Duo enabled' + + let(:params) do + super().merge(job_variables: { 'GIT_DEPTH' => 15, 'GIT_STRATEGY' => 'fetch', 'FOO' => 'bar' }) + end + + it 'does not create a workflow' do + expect(execute).to be_error + expect(execute.reason).to eq(:unprocessable_entity) + expect(execute.message).to eq(<<~MESSAGE.squish) + Non-overridable variables detected: GIT_STRATEGY and FOO. + Allowed variables: GIT_DEPTH. + MESSAGE + end + end end