diff --git a/db/migrate/20241106225114_add_agent_privileges_to_workflow.rb b/db/migrate/20241106225114_add_agent_privileges_to_workflow.rb new file mode 100644 index 0000000000000000000000000000000000000000..38a768a0e3ccf2e799889c7babc579e1e57bf054 --- /dev/null +++ b/db/migrate/20241106225114_add_agent_privileges_to_workflow.rb @@ -0,0 +1,13 @@ +# frozen_string_literal: true + +class AddAgentPrivilegesToWorkflow < Gitlab::Database::Migration[2.2] + milestone '17.6' + + READ_WRITE_FILES = 1 + READ_ONLY_GITLAB = 2 + + def change + add_column :duo_workflows_workflows, :agent_privileges, :smallint, array: true, + default: [READ_WRITE_FILES, READ_ONLY_GITLAB], null: false + end +end diff --git a/db/schema_migrations/20241106225114 b/db/schema_migrations/20241106225114 new file mode 100644 index 0000000000000000000000000000000000000000..c958f8049a9d1269f5a6cff86b933a7a9d1913ba --- /dev/null +++ b/db/schema_migrations/20241106225114 @@ -0,0 +1 @@ +3d5a17cd95b1c4c55ade68d396dca2fd1a6d7534df06a29b44fb5ca8f0abe002 \ No newline at end of file diff --git a/db/structure.sql b/db/structure.sql index 91c77e9f9353a56f1aaaba2f2d01cdcc996eab79..4c99e1fe459cf2db1815e233ee0ebabd42c56909 100644 --- a/db/structure.sql +++ b/db/structure.sql @@ -11413,6 +11413,7 @@ CREATE TABLE duo_workflows_workflows ( updated_at timestamp with time zone NOT NULL, status smallint DEFAULT 0 NOT NULL, goal text, + agent_privileges smallint[] DEFAULT '{1,2}'::smallint[] NOT NULL, CONSTRAINT check_5aedde451d CHECK ((char_length(goal) <= 4096)) ); diff --git a/ee/app/models/ai/duo_workflows/workflow.rb b/ee/app/models/ai/duo_workflows/workflow.rb index 124567b06bc10280da3ef96e5ace10646237aedf..6289f2eb62326857d567b2d4282e7c5525a0d384 100644 --- a/ee/app/models/ai/duo_workflows/workflow.rb +++ b/ee/app/models/ai/duo_workflows/workflow.rb @@ -13,10 +13,53 @@ class Workflow < ::ApplicationRecord validates :status, presence: true validates :goal, length: { maximum: 4096 } + validate :only_known_agent_priviliges + scope :for_user_with_id!, ->(user_id, id) { find_by!(user_id: user_id, id: id) } scope :for_user, ->(user_id) { where(user_id: user_id) } scope :for_project, ->(project) { where(project: project) } + class AgentPrivileges + READ_WRITE_FILES = 1 + READ_ONLY_GITLAB = 2 + READ_WRITE_GITLAB = 3 + RUN_COMMANDS = 4 + + ALL_PRIVILEGES = { + READ_WRITE_FILES => { + name: "read_write_files", + description: "Allow local filesystem read/write access" + }.freeze, + READ_ONLY_GITLAB => { + name: "read_only_gitlab", + description: "Allow read only access to GitLab APIs" + }.freeze, + READ_WRITE_GITLAB => { + name: "read_write_gitlab", + description: "Allow write access to GitLab APIs" + }.freeze, + RUN_COMMANDS => { + name: "run_commands", + description: "Allow running any commands" + }.freeze + }.freeze + + DEFAULT_PRIVILEGES = [ + READ_WRITE_FILES, + READ_ONLY_GITLAB + ].freeze + end + + def only_known_agent_priviliges + self.agent_privileges ||= AgentPrivileges::DEFAULT_PRIVILEGES + + agent_privileges.each do |privilege| + unless AgentPrivileges::ALL_PRIVILEGES.key?(privilege) + errors.add(:agent_privileges, "contains an invalid value #{privilege}") + end + end + end + state_machine :status, initial: :created do event :start do transition created: :running diff --git a/ee/lib/api/ai/duo_workflows/workflows.rb b/ee/lib/api/ai/duo_workflows/workflows.rb index 719c8bda3fa5bc3e2a74eb48bdf4e38ddad4bd46..3e6618512fd67d490268746af3c52a63cd8ef737 100644 --- a/ee/lib/api/ai/duo_workflows/workflows.rb +++ b/ee/lib/api/ai/duo_workflows/workflows.rb @@ -56,6 +56,8 @@ def create_workflow_params documentation: { example: '1' } optional :goal, type: String, desc: 'Goal of the workflow', documentation: { example: 'Fix pipeline for merge request 1 in project 1' } + optional :agent_privileges, type: [Integer], desc: 'The actions the agent is allowed to perform', + documentation: { example: [1] } end end @@ -128,10 +130,15 @@ def create_workflow_params present result[:workflow], with: ::API::Entities::Ai::DuoWorkflows::Workflow end - get '/:id' do - workflow = find_workflow!(params[:id]) - - present workflow, with: ::API::Entities::Ai::DuoWorkflows::Workflow + desc 'Get all possible agent privileges and descriptions' do + success code: 200 + failure [ + { code: 401, message: 'Unauthorized' } + ] + end + get 'agent_privileges' do + present ::Ai::DuoWorkflows::Workflow::AgentPrivileges, + with: ::API::Entities::Ai::DuoWorkflows::Workflow::AgentPrivileges end end end diff --git a/ee/lib/api/ai/duo_workflows/workflows_internal.rb b/ee/lib/api/ai/duo_workflows/workflows_internal.rb index 94dd13889aa186eb29b0820f154295f58af9cbb1..ffde986939794ba23e43a3d288dd47d5e5d687ef 100644 --- a/ee/lib/api/ai/duo_workflows/workflows_internal.rb +++ b/ee/lib/api/ai/duo_workflows/workflows_internal.rb @@ -44,6 +44,15 @@ def render_response(response) namespace :duo_workflows do namespace :workflows do namespace '/:id' do + params do + requires :id, type: Integer, desc: 'The ID of the workflow', documentation: { example: 1 } + end + get do + workflow = find_workflow!(params[:id]) + + present workflow, with: ::API::Entities::Ai::DuoWorkflows::Workflow + end + desc 'Updates the workflow status' do success code: 200 end diff --git a/ee/lib/api/entities/ai/duo_workflows/workflow.rb b/ee/lib/api/entities/ai/duo_workflows/workflow.rb index e9fdbd1c7f9b299881119833c25ba59779ac67c0..def6271caf1b63cdfd0fa2b67bc62b673842178b 100644 --- a/ee/lib/api/entities/ai/duo_workflows/workflow.rb +++ b/ee/lib/api/entities/ai/duo_workflows/workflow.rb @@ -6,6 +6,14 @@ module Ai module DuoWorkflows class Workflow < Grape::Entity expose :id + expose :agent_privileges + expose :agent_privileges_names + + def agent_privileges_names + object.agent_privileges.map do |privilege| + ::Ai::DuoWorkflows::Workflow::AgentPrivileges::ALL_PRIVILEGES[privilege][:name] + end + end end end end diff --git a/ee/lib/api/entities/ai/duo_workflows/workflow/agent_privileges.rb b/ee/lib/api/entities/ai/duo_workflows/workflow/agent_privileges.rb new file mode 100644 index 0000000000000000000000000000000000000000..1afe55cf1eade0f2373d6ed3d2bd6cce3b9d3be9 --- /dev/null +++ b/ee/lib/api/entities/ai/duo_workflows/workflow/agent_privileges.rb @@ -0,0 +1,26 @@ +# frozen_string_literal: true + +module API + module Entities + module Ai + module DuoWorkflows + class Workflow + class AgentPrivileges < Grape::Entity + expose :all_privileges + + def all_privileges + object::ALL_PRIVILEGES.map do |id, attributes| + { + id: id, + name: attributes[:name], + description: attributes[:description], + default_enabled: id.in?(object::DEFAULT_PRIVILEGES) + } + end + end + end + end + end + end + end +end diff --git a/ee/spec/factories/ai/duo_workflows/workflows.rb b/ee/spec/factories/ai/duo_workflows/workflows.rb index d82fe4bdf7c52d9e4e2c4c70d69536ae3a1e459d..aa67613d1cc8649b96a8cbf623d6bd108061aa97 100644 --- a/ee/spec/factories/ai/duo_workflows/workflows.rb +++ b/ee/spec/factories/ai/duo_workflows/workflows.rb @@ -5,5 +5,6 @@ project { association(:project) } user { association(:user, developer_of: project) } goal { "Fix pipeline" } + agent_privileges { [Ai::DuoWorkflows::Workflow::AgentPrivileges::READ_WRITE_FILES] } end end diff --git a/ee/spec/models/ai/duo_workflows/workflow_spec.rb b/ee/spec/models/ai/duo_workflows/workflow_spec.rb index 3dee3968a0b5e078cde5a6ceb8cd7cdd7b859c30..bcea7837ac0e37d2b289b6e936e86fae5e0a9e2a 100644 --- a/ee/spec/models/ai/duo_workflows/workflow_spec.rb +++ b/ee/spec/models/ai/duo_workflows/workflow_spec.rb @@ -37,6 +37,62 @@ describe 'validations' do it { is_expected.to validate_presence_of(:status) } it { is_expected.to validate_length_of(:goal).is_at_most(4096) } + + describe '#only_known_agent_priviliges' do + it 'is valid with a valid privilege' do + workflow = described_class.new(agent_privileges: [ + Ai::DuoWorkflows::Workflow::AgentPrivileges::READ_WRITE_FILES + ]) + expect(workflow).to be_valid + end + + it 'is invalid with an invalid privilege' do + workflow = described_class.new(agent_privileges: [999]) + expect(workflow).not_to be_valid + expect(workflow.errors[:agent_privileges]).to include("contains an invalid value 999") + end + end + end + + describe '#agent_privileges' do + it 'returns the privileges that are set' do + workflow = described_class.new(agent_privileges: [ + Ai::DuoWorkflows::Workflow::AgentPrivileges::READ_WRITE_FILES, + Ai::DuoWorkflows::Workflow::AgentPrivileges::READ_WRITE_GITLAB + ]) + + # Validation triggers setting the default + expect(workflow).to be_valid + + expect(workflow.agent_privileges).to match_array([ + described_class::AgentPrivileges::READ_WRITE_FILES, + described_class::AgentPrivileges::READ_WRITE_GITLAB + ]) + end + + it 'replaces with DEFAULT_PRIVILEGES when set to nil' do + workflow = described_class.new(agent_privileges: nil) + + # Validation triggers setting the default + expect(workflow).to be_valid + + expect(workflow.agent_privileges).to match_array([ + described_class::AgentPrivileges::READ_WRITE_FILES, + described_class::AgentPrivileges::READ_ONLY_GITLAB + ]) + end + + it 'replaces with DEFAULT_PRIVILEGES when not set' do + workflow = described_class.new + + # Validation triggers setting the default + expect(workflow).to be_valid + + expect(workflow.agent_privileges).to match_array([ + described_class::AgentPrivileges::READ_WRITE_FILES, + described_class::AgentPrivileges::READ_ONLY_GITLAB + ]) + end end describe 'state transitions' do diff --git a/ee/spec/requests/api/ai/duo_workflows/workflows_internal_spec.rb b/ee/spec/requests/api/ai/duo_workflows/workflows_internal_spec.rb index 84caa07704df710ff188cb05696dfca09d6ad09d..35b7131ff530e56c43231b7a8c8cc396a0b9c47c 100644 --- a/ee/spec/requests/api/ai/duo_workflows/workflows_internal_spec.rb +++ b/ee/spec/requests/api/ai/duo_workflows/workflows_internal_spec.rb @@ -10,7 +10,7 @@ let_it_be(:user) { create(:user, maintainer_of: project) } let_it_be(:workflow) { create(:duo_workflows_workflow, user: user, project: project) } let_it_be(:duo_workflow_service_url) { 'duo-workflow-service.example.com:50052' } - let_it_be(:oauth_token) { create(:oauth_access_token, user: user, scopes: [:ai_workflows]) } + let_it_be(:ai_workflows_oauth_token) { create(:oauth_access_token, user: user, scopes: [:ai_workflows]) } before do allow(::Gitlab::Llm::StageCheck).to receive(:available?).with(project, :duo_workflow).and_return(true) @@ -40,7 +40,7 @@ context 'when authenticated with a token that has the ai_workflows scope' do it 'is successful' do - post api(path, oauth_access_token: oauth_token), + post api(path, oauth_access_token: ai_workflows_oauth_token), params: params.merge(thread_ts: later_thread_ts, parent_ts: thread_ts) expect(response).to have_gitlab_http_status(:created) @@ -97,7 +97,7 @@ context 'when authenticated with a token that has the ai_workflows scope' do it 'is successful' do expect do - post api(path, oauth_access_token: oauth_token), params: params + post api(path, oauth_access_token: ai_workflows_oauth_token), params: params expect(response).to have_gitlab_http_status(:created) end.to change { workflow.events.count }.by(1) end @@ -186,7 +186,7 @@ context 'when authenticated with a token that has the ai_workflows scope' do it 'is successful' do - put api(path, oauth_access_token: oauth_token), params: params + put api(path, oauth_access_token: ai_workflows_oauth_token), params: params expect(response).to have_gitlab_http_status(:ok) end end @@ -229,6 +229,49 @@ end end + describe 'GET /ai/duo_workflows/workflows/:id' do + let(:path) { "/ai/duo_workflows/workflows/#{workflow.id}" } + + it 'returns the Ai::DuoWorkflows::Workflow' do + get api(path, user) + + expect(response).to have_gitlab_http_status(:ok) + expect(json_response['id']).to eq(workflow.id) + expect(json_response['agent_privileges']).to eq(workflow.agent_privileges) + expect(json_response['agent_privileges_names']).to eq(["read_write_files"]) + end + + context 'when authenticated with a token that has the ai_workflows scope' do + it 'returns the Ai::DuoWorkflows::Workflow' do + get api(path, oauth_access_token: ai_workflows_oauth_token) + + expect(response).to have_gitlab_http_status(:ok) + expect(json_response['id']).to eq(workflow.id) + end + end + + context 'when duo_features_enabled settings is turned off' do + before do + workflow.project.project_setting.update!(duo_features_enabled: false) + workflow.project.reload + end + + it 'returns forbidden' do + get api(path, user) + expect(response).to have_gitlab_http_status(:forbidden) + end + end + + context 'with a workflow belonging to a different user' do + let(:workflow) { create(:duo_workflows_workflow) } + + it 'returns 404' do + get api(path, user) + expect(response).to have_gitlab_http_status(:not_found) + end + end + end + describe 'PATCH /ai/duo_workflows/workflows/:id' do let(:path) { "/ai/duo_workflows/workflows/#{workflow.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 a54bf1da1973864fdae33a77feede54946ee622b..c352c721d2857e0b4d6f00d052ed307456082cd7 100644 --- a/ee/spec/requests/api/ai/duo_workflows/workflows_spec.rb +++ b/ee/spec/requests/api/ai/duo_workflows/workflows_spec.rb @@ -11,6 +11,7 @@ let_it_be(:workflow) { create(:duo_workflows_workflow, user: user, project: project) } let_it_be(:duo_workflow_service_url) { 'duo-workflow-service.example.com:50052' } let_it_be(:ai_workflows_oauth_token) { create(:oauth_access_token, user: user, scopes: [:ai_workflows]) } + let(:agent_privileges) { [::Ai::DuoWorkflows::Workflow::AgentPrivileges::READ_WRITE_FILES] } before do allow(::Gitlab::Llm::StageCheck).to receive(:available?).with(project, :duo_workflow).and_return(true) @@ -18,7 +19,7 @@ describe 'POST /ai/duo_workflows/workflows' do let(:path) { "/ai/duo_workflows/workflows" } - let(:params) { { project_id: project.id } } + let(:params) { { project_id: project.id, agent_privileges: agent_privileges } } context 'when success' do it 'creates the Ai::DuoWorkflows::Workflow' do @@ -27,6 +28,24 @@ expect(response).to have_gitlab_http_status(:created) end.to change { Ai::DuoWorkflows::Workflow.count }.by(1) expect(json_response['id']).to eq(Ai::DuoWorkflows::Workflow.last.id) + + created_workflow = Ai::DuoWorkflows::Workflow.last + + expect(created_workflow.agent_privileges).to eq(agent_privileges) + end + + context 'when agent_privileges is not provided' do + let(:params) { { project_id: project.id } } + + it 'creates a workflow with the default agent_privileges' do + post api(path, user), params: params + expect(response).to have_gitlab_http_status(:created) + + created_workflow = Ai::DuoWorkflows::Workflow.last + expect(created_workflow.agent_privileges).to match_array( + ::Ai::DuoWorkflows::Workflow::AgentPrivileges::DEFAULT_PRIVILEGES + ) + end end context 'when authenticated with a token that has the ai_workflows scope' do @@ -84,46 +103,6 @@ end end - describe 'GET /ai/duo_workflows/workflows/:id' do - let(:path) { "/ai/duo_workflows/workflows/#{workflow.id}" } - - it 'returns the Ai::DuoWorkflows::Workflow' do - get api(path, user) - - expect(response).to have_gitlab_http_status(:ok) - expect(json_response['id']).to eq(workflow.id) - end - - context 'when authenticated with a token that has the ai_workflows scope' do - it 'is forbidden' do - get api(path, oauth_access_token: ai_workflows_oauth_token) - - expect(response).to have_gitlab_http_status(:forbidden) - end - end - - context 'when duo_features_enabled settings is turned off' do - before do - workflow.project.project_setting.update!(duo_features_enabled: false) - workflow.project.reload - end - - it 'returns forbidden' do - get api(path, user) - expect(response).to have_gitlab_http_status(:forbidden) - end - end - - context 'with a workflow belonging to a different user' do - let(:workflow) { create(:duo_workflows_workflow) } - - it 'returns 404' do - get api(path, user) - expect(response).to have_gitlab_http_status(:not_found) - end - end - end - describe 'POST /ai/duo_workflows/direct_access' do let(:path) { '/ai/duo_workflows/direct_access' } @@ -234,4 +213,28 @@ end end end + + describe 'GET /ai/duo_workflows/workflows/agent_privileges' do + let(:path) { "/ai/duo_workflows/workflows/agent_privileges" } + + it 'returns a static set of privileges' do + get api(path, user) + + expect(response).to have_gitlab_http_status(:ok) + + expect(json_response['all_privileges'].count).to eq(4) + + privilege1 = json_response['all_privileges'][0] + expect(privilege1['id']).to eq(1) + expect(privilege1['name']).to eq('read_write_files') + expect(privilege1['description']).to eq('Allow local filesystem read/write access') + expect(privilege1['default_enabled']).to eq(true) + + privilege4 = json_response['all_privileges'][3] + expect(privilege4['id']).to eq(4) + expect(privilege4['name']).to eq('run_commands') + expect(privilege4['description']).to eq('Allow running any commands') + expect(privilege4['default_enabled']).to eq(false) + end + end end