diff --git a/db/docs/duo_workflows_leases.yml b/db/docs/duo_workflows_leases.yml new file mode 100644 index 0000000000000000000000000000000000000000..f5c3378d1bcc54bef64b08a957132011bed4bb81 --- /dev/null +++ b/db/docs/duo_workflows_leases.yml @@ -0,0 +1,14 @@ +--- +table_name: duo_workflows_leases +classes: +- Ai::DuoWorkflows::Lease +feature_categories: +- agent_foundations +description: +introduced_by_url: +milestone: '18.5' +gitlab_schema: gitlab_main_org +sharding_key: + project_id: projects + namespace_id: namespaces +table_size: small diff --git a/db/migrate/20251003044825_create_table_duo_workflows_leases.rb b/db/migrate/20251003044825_create_table_duo_workflows_leases.rb new file mode 100644 index 0000000000000000000000000000000000000000..2e08467828b70dd1b48a1661a8e4e15638c9deed --- /dev/null +++ b/db/migrate/20251003044825_create_table_duo_workflows_leases.rb @@ -0,0 +1,30 @@ +# frozen_string_literal: true + +class CreateTableDuoWorkflowsLeases < Gitlab::Database::Migration[2.3] + disable_ddl_transaction! + + milestone '18.5' + + def up + create_table :duo_workflows_leases, if_not_exists: true do |t| + t.bigint :workflow_id, null: false + t.timestamp :last_renewed_at + t.text :lease_id, null: false + t.bigint :project_id + t.bigint :namespace_id + + t.timestamps null: false + end + + add_concurrent_index :duo_workflows_leases, [:workflow_id], unique: true + + add_multi_column_not_null_constraint(:duo_workflows_leases, :namespace_id, :project_id, validate: true) + add_concurrent_foreign_key :duo_workflows_leases, :duo_workflows_workflows, column: :workflow_id, on_delete: :cascade + add_concurrent_foreign_key :duo_workflows_leases, :projects, column: :project_id, on_delete: :cascade + add_concurrent_foreign_key :duo_workflows_leases, :namespaces, column: :namespace_id, on_delete: :cascade + end + + def down + drop_table :duo_workflows_leases, if_exists: true + end +end diff --git a/db/schema_migrations/20251003044825 b/db/schema_migrations/20251003044825 new file mode 100644 index 0000000000000000000000000000000000000000..c442e44ad2b19abb5911610319cfe77d29275180 --- /dev/null +++ b/db/schema_migrations/20251003044825 @@ -0,0 +1 @@ +e2ec491a524b98ca41dd36c04595c14d4faefca2536052eda7100fb7fc928998 \ No newline at end of file diff --git a/db/structure.sql b/db/structure.sql index b5bdd311d8bfe26837cecb23a7ed051d3bf3e977..3e345bab890c1e51d12be531df35b49bccd85558 100644 --- a/db/structure.sql +++ b/db/structure.sql @@ -16345,6 +16345,27 @@ CREATE SEQUENCE duo_workflows_events_id_seq ALTER SEQUENCE duo_workflows_events_id_seq OWNED BY duo_workflows_events.id; +CREATE TABLE duo_workflows_leases ( + id bigint NOT NULL, + workflow_id bigint NOT NULL, + last_renewed_at timestamp without time zone, + lease_id text NOT NULL, + project_id bigint, + namespace_id bigint, + created_at timestamp(6) without time zone NOT NULL, + updated_at timestamp(6) without time zone NOT NULL, + CONSTRAINT check_8c86845a83 CHECK ((num_nonnulls(namespace_id, project_id) = 1)) +); + +CREATE SEQUENCE duo_workflows_leases_id_seq + START WITH 1 + INCREMENT BY 1 + NO MINVALUE + NO MAXVALUE + CACHE 1; + +ALTER SEQUENCE duo_workflows_leases_id_seq OWNED BY duo_workflows_leases.id; + CREATE TABLE duo_workflows_workflows ( id bigint NOT NULL, user_id bigint NOT NULL, @@ -30868,6 +30889,8 @@ ALTER TABLE ONLY duo_workflows_checkpoint_writes ALTER COLUMN id SET DEFAULT nex ALTER TABLE ONLY duo_workflows_events ALTER COLUMN id SET DEFAULT nextval('duo_workflows_events_id_seq'::regclass); +ALTER TABLE ONLY duo_workflows_leases ALTER COLUMN id SET DEFAULT nextval('duo_workflows_leases_id_seq'::regclass); + ALTER TABLE ONLY duo_workflows_workflows ALTER COLUMN id SET DEFAULT nextval('duo_workflows_workflows_id_seq'::regclass); ALTER TABLE ONLY duo_workflows_workloads ALTER COLUMN id SET DEFAULT nextval('duo_workflows_workloads_id_seq'::regclass); @@ -33741,6 +33764,9 @@ ALTER TABLE ONLY duo_workflows_checkpoint_writes ALTER TABLE ONLY duo_workflows_events ADD CONSTRAINT duo_workflows_events_pkey PRIMARY KEY (id); +ALTER TABLE ONLY duo_workflows_leases + ADD CONSTRAINT duo_workflows_leases_pkey PRIMARY KEY (id); + ALTER TABLE ONLY duo_workflows_workflows ADD CONSTRAINT duo_workflows_workflows_pkey PRIMARY KEY (id); @@ -39752,6 +39778,8 @@ CREATE INDEX index_duo_workflows_events_on_project_id ON duo_workflows_events US CREATE INDEX index_duo_workflows_events_on_workflow_id ON duo_workflows_events USING btree (workflow_id); +CREATE UNIQUE INDEX index_duo_workflows_leases_on_workflow_id ON duo_workflows_leases USING btree (workflow_id); + CREATE INDEX index_duo_workflows_workflows_on_ai_catalog_item_version_id ON duo_workflows_workflows USING btree (ai_catalog_item_version_id); CREATE INDEX index_duo_workflows_workflows_on_namespace_id ON duo_workflows_workflows USING btree (namespace_id); @@ -48187,6 +48215,9 @@ ALTER TABLE ONLY audit_events_google_cloud_logging_configurations ALTER TABLE ONLY releases ADD CONSTRAINT fk_47fe2a0596 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; +ALTER TABLE ONLY duo_workflows_leases + ADD CONSTRAINT fk_480d241192 FOREIGN KEY (workflow_id) REFERENCES duo_workflows_workflows(id) ON DELETE CASCADE; + ALTER TABLE ONLY operations_scopes ADD CONSTRAINT fk_4913f5d6a2 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; @@ -49375,6 +49406,9 @@ ALTER TABLE ONLY timelogs ALTER TABLE ONLY sbom_graph_paths ADD CONSTRAINT fk_c4c7d16f3e FOREIGN KEY (ancestor_id) REFERENCES sbom_occurrences(id) ON DELETE CASCADE; +ALTER TABLE ONLY duo_workflows_leases + ADD CONSTRAINT fk_c540134e48 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; + ALTER TABLE ONLY wiki_repository_states ADD CONSTRAINT fk_c558ca51b8 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; @@ -49468,6 +49502,9 @@ ALTER TABLE ONLY packages_debian_project_architectures ALTER TABLE ONLY incident_management_escalation_rules ADD CONSTRAINT fk_cdfc40b861 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; +ALTER TABLE ONLY duo_workflows_leases + ADD CONSTRAINT fk_ce7a3f897f FOREIGN KEY (namespace_id) REFERENCES namespaces(id) ON DELETE CASCADE; + ALTER TABLE ONLY packages_dependencies ADD CONSTRAINT fk_cea1124da7 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; diff --git a/ee/app/graphql/ee/types/mutation_type.rb b/ee/app/graphql/ee/types/mutation_type.rb index b66bb19e3a8b43241a9b6170d40341f998d0a3bb..e45928dc2d08a1e11b65001ac2c85d1029d6ec6f 100644 --- a/ee/app/graphql/ee/types/mutation_type.rb +++ b/ee/app/graphql/ee/types/mutation_type.rb @@ -350,6 +350,10 @@ def self.authorization_scopes mount_mutation ::Mutations::Ai::DuoWorkflows::DeleteWorkflow, experiment: { milestone: '18.1' } mount_mutation ::Mutations::Ai::DuoWorkflows::Create, experiment: { milestone: '18.1' }, scopes: [:api, :ai_features] + mount_mutation ::Mutations::Ai::DuoWorkflows::WorkflowWithLease, experiment: { milestone: '18.5' }, + scopes: [:api, :ai_workflows] + mount_mutation ::Mutations::Ai::DuoWorkflows::WorkflowLeaseRenew, experiment: { milestone: '18.5' }, + scopes: [:api, :ai_workflows] mount_mutation ::Mutations::Authz::LdapAdminRoleLinks::Create, experiment: { milestone: '17.11' } mount_mutation ::Mutations::Authz::LdapAdminRoleLinks::Destroy, experiment: { milestone: '18.0' } mount_mutation ::Mutations::Authz::AdminRoles::LdapSync, experiment: { milestone: '18.0' } diff --git a/ee/app/graphql/mutations/ai/duo_workflows/workflow_lease_renew.rb b/ee/app/graphql/mutations/ai/duo_workflows/workflow_lease_renew.rb new file mode 100644 index 0000000000000000000000000000000000000000..6fd46cf524abdf200b7326b1f35560a5b5e5780b --- /dev/null +++ b/ee/app/graphql/mutations/ai/duo_workflows/workflow_lease_renew.rb @@ -0,0 +1,52 @@ +# frozen_string_literal: true + +module Mutations + module Ai + module DuoWorkflows + class WorkflowLeaseRenew < BaseMutation + graphql_name 'AiDuoWorkflowLeaseRenew' + + # TODO: New permission for run_duo_workflow on specific workflow + authorize :read_duo_workflow + + def self.authorization_scopes + super + [:ai_workflows] + end + + argument :workflow_id, Types::GlobalIDType[::Ai::DuoWorkflows::Workflow], + required: true, + description: 'Flow ID of the flow to load and acquire a lease.' + + argument :lease_id, GraphQL::Types::String, + required: true, + description: 'Flow ID of the flow to load and acquire a lease.' + + field :last_renewed_at, GraphQL::Types::String, + null: true, + description: 'The value of the last time the lease was renewed. Will be roughly the current time since this lease was just renewed.' + + field :errors, [GraphQL::Types::String], + null: false, + description: 'Errors encountered during the creation process.' + + def resolve(workflow_id:, lease_id:) + workflow = authorized_find!(id: workflow_id) + last_renewed_at = Time.current + + updated = ::Ai::DuoWorkflows::Lease + .where(workflow_id: workflow.id, lease_id: lease_id) + .update_all(last_renewed_at: last_renewed_at) + + unless updated + return { errors: ['Lease was not updated. Likely this lease is no longer valid.'] } + end + + { + lease_id: last_renewed_at, + errors: [] + } + end + end + end + end +end diff --git a/ee/app/graphql/mutations/ai/duo_workflows/workflow_with_lease.rb b/ee/app/graphql/mutations/ai/duo_workflows/workflow_with_lease.rb new file mode 100644 index 0000000000000000000000000000000000000000..6c1f742e6090bfe4432e19bd7562f7d002cabd16 --- /dev/null +++ b/ee/app/graphql/mutations/ai/duo_workflows/workflow_with_lease.rb @@ -0,0 +1,53 @@ +# frozen_string_literal: true + +module Mutations + module Ai + module DuoWorkflows + class WorkflowWithLease < BaseMutation + graphql_name 'AiDuoWorkflowWithLease' + + # TODO: New permission for run_duo_workflow on specific workflow + authorize :read_duo_workflow + + def self.authorization_scopes + super + [:ai_workflows] + end + + argument :workflow_id, Types::GlobalIDType[::Ai::DuoWorkflows::Workflow], + required: true, + description: 'Flow ID of the flow to load and acquire a lease.' + + field :workflow, Types::Ai::DuoWorkflows::WorkflowType, + scopes: [:api, :ai_workflows], + null: true, + description: 'Created workflow.' + + field :lease_id, GraphQL::Types::String, + scopes: [:api, :ai_workflows], + null: true, + description: 'This lease ID must be sent back as a query paramater with all subsequent checkpoints, heartbeats and REST API calls. If a stale lease_id is provided in any of those subsequent API calls then they will fail. Additionally this mutation will fail if a lease already exists' + + def resolve(workflow_id:) + workflow = authorized_find!(id: workflow_id) + ::Ai::DuoWorkflows::WorkflowLeaseExpirationService.new.execute(workflow) + + lease = begin + ::Ai::DuoWorkflows::Lease.create(workflow_id: workflow.id, last_renewed_at: Time.current) + rescue ActiveRecord::RecordNotUnique + nil + end + + unless lease&.persisted? + return { errors: ['Failed to acquire lease'], workflow: nil } + end + + { + workflow: workflow, + lease_id: lease.lease_id, + errors: [] + } + end + end + end + end +end diff --git a/ee/app/graphql/resolvers/ai/duo_workflows/workflows_resolver.rb b/ee/app/graphql/resolvers/ai/duo_workflows/workflows_resolver.rb index b1367289b04c7af9a83c458343737dfd0d62d5bb..2f1cd8c6344171bebef6870c5afe90ea2389bb3f 100644 --- a/ee/app/graphql/resolvers/ai/duo_workflows/workflows_resolver.rb +++ b/ee/app/graphql/resolvers/ai/duo_workflows/workflows_resolver.rb @@ -49,7 +49,7 @@ def resolve(**args) return resolve_single_workflow(args[:workflow_id]) if args[:workflow_id].present? - ::Ai::DuoWorkflows::WorkflowsFinder.new( + workflows = ::Ai::DuoWorkflows::WorkflowsFinder.new( source: object, current_user: current_user, project_path: args[:project_path], @@ -60,6 +60,8 @@ def resolve(**args) search: args[:search], status_group: args[:status_group] ).results + + workflows end private @@ -72,12 +74,16 @@ def conflicting_type_filters?(args) end def resolve_single_workflow(workflow_id) + Gitlab::Graphql::Lazy.with_value(find_object(id: workflow_id)) do |workflow| if workflow.nil? raise_resource_not_available_error! "Workflow not found" elsif !Ability.allowed?(current_user, :read_duo_workflow, workflow) raise_resource_not_available_error! "You don't have permission to access this workflow" else + # TODO: Move to when we load all the workflows as well as single workflow? + # Any time a workflow is fetched we want to trigger expiration for it + ::Ai::DuoWorkflows::WorkflowLeaseExpirationWorker.perform_async(workflow.id) ::Ai::DuoWorkflows::Workflow.id_in([workflow.id]) end end diff --git a/ee/app/models/ai/duo_workflows/lease.rb b/ee/app/models/ai/duo_workflows/lease.rb new file mode 100644 index 0000000000000000000000000000000000000000..cd78e8aa010d05649239543e4219616950f1c0af --- /dev/null +++ b/ee/app/models/ai/duo_workflows/lease.rb @@ -0,0 +1,32 @@ +module Ai + module DuoWorkflows + class Lease < ::ApplicationRecord + self.table_name = :duo_workflows_leases + + EXPIRY = 60.seconds + + belongs_to :workflow, class_name: '::Ai::DuoWorkflows::Workflow' + belongs_to :project, class_name: '::Project', required: false + + before_save :ensure_lease_id + before_save :set_project_or_namespace + + def ensure_lease_id + self.lease_id ||= SecureRandom.uuid + end + + def set_project_or_namespace + self.project_id = workflow.project_id + self.namespace_id = workflow.namespace_id + end + + def renew! + update!(last_renewed_at: Time.current) + end + + def expired? + last_renewed_at < EXPIRY.ago + end + end + end +end diff --git a/ee/app/models/ai/duo_workflows/workflow.rb b/ee/app/models/ai/duo_workflows/workflow.rb index a6f92e4bb329fb7cfbbbd16ba962cc887b336673..2d2cfebeb487cc61b9d2cd98630094d2903843e7 100644 --- a/ee/app/models/ai/duo_workflows/workflow.rb +++ b/ee/app/models/ai/duo_workflows/workflow.rb @@ -7,6 +7,7 @@ class Workflow < ::ApplicationRecord include FromUnion include EachBatch include Sortable + include AfterCommitQueue self.table_name = :duo_workflows_workflows @@ -22,6 +23,8 @@ class Workflow < ::ApplicationRecord has_many :workloads, through: :workflows_workloads, disable_joins: true has_many :vulnerability_triggered_workflows, class_name: '::Vulnerabilities::TriggeredWorkflow' + has_one :current_lease, class_name: 'Ai::DuoWorkflows::Lease', foreign_key: :workflow_id + validates :status, presence: true validates :goal, length: { maximum: 16_384 } validates :image, length: { maximum: 2048 }, allow_blank: true @@ -308,6 +311,14 @@ def pre_approved_privileges_included_in_agent_privileges ] => ::Ai::DuoWorkflows::Workflow.target_status_for_event(:stop) end + after_transition any => [:paused, :finished, :failed, :stopped, :input_required, :plan_approval_required, :tool_call_approval_required] do |workflow, transition| + workflow.run_after_commit do + # The workflow gets marked as "input_required" before final checkpoint comes in. As such we need a 15s grace + # period to allow the final checkpoint to come in before we expire the lease. + ::Ai::DuoWorkflows::WorkflowLeaseExpirationWorker.perform_in(15.seconds, workflow.id) + end + end + state :created, value: 0 state :running, value: 1 state :paused, value: 2 diff --git a/ee/app/services/ai/duo_workflows/workflow_lease_expiration_service.rb b/ee/app/services/ai/duo_workflows/workflow_lease_expiration_service.rb new file mode 100644 index 0000000000000000000000000000000000000000..2fc94918140267235a186393a3fa106a0870052b --- /dev/null +++ b/ee/app/services/ai/duo_workflows/workflow_lease_expiration_service.rb @@ -0,0 +1,28 @@ +# frozen_string_literal: true + +module Ai + module DuoWorkflows + class WorkflowLeaseExpirationService + def execute(workflow, including_finished_flows: true) + if workflow.running? && workflow.current_lease&.expired? + ::Ai::DuoWorkflows::Workflow.transaction do + workflow.current_lease.destroy! + workflow.drop + + # Must reload otherwise caller gets stale current_lease + workflow.reload + end + elsif including_finished_flows && !workflow.running? && workflow.current_lease.present? + # If it's not running anymore the lease should be expired + ::Ai::DuoWorkflows::Workflow.transaction do + workflow.current_lease.destroy! + # No need to drop here as it's not running anyway + + # Must reload otherwise caller gets stale current_lease + workflow.reload + end + end + end + end + end +end diff --git a/ee/app/workers/ai/duo_workflows/workflow_lease_expiration_worker.rb b/ee/app/workers/ai/duo_workflows/workflow_lease_expiration_worker.rb new file mode 100644 index 0000000000000000000000000000000000000000..6eaea78a1886daec7256440de5d7a8a7ec4963b0 --- /dev/null +++ b/ee/app/workers/ai/duo_workflows/workflow_lease_expiration_worker.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +module Ai + module DuoWorkflows + class WorkflowLeaseExpirationWorker + include ::ApplicationWorker + + idempotent! + worker_resource_boundary :cpu + urgency :low + feature_category :agent_foundations + data_consistency :sticky + + def perform(workflow_id) + workflow = ::Ai::DuoWorkflows::Workflow.find(workflow_id) + ::Ai::DuoWorkflows::WorkflowLeaseExpirationService.new.execute(workflow) + end + 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 e3025ee44f8cc6d0504fea13fcb6172cb3a47abb..0e756b835f0bb0396469e75b094b8bea50f558f6 100644 --- a/ee/lib/api/ai/duo_workflows/workflows_internal.rb +++ b/ee/lib/api/ai/duo_workflows/workflows_internal.rb @@ -108,10 +108,30 @@ def render_response(response) optional :parent_ts, type: String, desc: 'The parent ts' requires :checkpoint, type: Hash, desc: "Checkpoint content" requires :metadata, type: Hash, desc: "Checkpoint metadata" + optional :lease_id, type: String, desc: 'The lease id. This will be renewed with each checkpoint. Checkpoints will be rejected if the lease_id is not current.' end post do workflow = find_workflow!(params[:id]) - checkpoint_params = declared_params(include_missing: false).except(:id) + + # TODO: Feature flag around all lease related logic + lease_id = params.delete(:lease_id) + + # Set including_finished_flows: false here because we don't want to expire flows that have just + # finished because they may still have some checkpoints coming in after they are marked as finished. + ::Ai::DuoWorkflows::WorkflowLeaseExpirationService.new.execute(workflow, including_finished_flows: false) + + if lease_id + lease = workflow.current_lease + if lease && lease.lease_id == lease_id + lease.renew! + else + # TODO: Special feature flag for failing checkpoints + bad_request!("Lease does not match existing lease for workflow") + end + end + + checkpoint_params = declared_params(include_missing: false).except(:id, :lease_id) + service = ::Ai::DuoWorkflows::CreateCheckpointService.new( workflow: workflow, params: checkpoint_params) result = service.execute diff --git a/ee/spec/factories/ai/duo_workflows/leases.rb b/ee/spec/factories/ai/duo_workflows/leases.rb new file mode 100644 index 0000000000000000000000000000000000000000..c41139c3c6636f277a2521daf6e5c1c77231d264 --- /dev/null +++ b/ee/spec/factories/ai/duo_workflows/leases.rb @@ -0,0 +1,14 @@ +# frozen_string_literal: true + +FactoryBot.define do + factory :duo_workflows_lease, class: 'Ai::DuoWorkflows::Lease' do + lease_id { SecureRandom.uuid } + last_renewed_at { Time.current } + workflow { association(:workflow) } + + after(:build) do |lease, _| + lease.project_id = lease.workflow.project_id + lease.namespace_id = lease.workflow.project_id + end + end +end 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 c06583111a0ccea6ccfd1b3c9d2fae3185a37735..714cbc0bbae4c5d48447513f7d03856a62ec7be6 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 @@ -44,7 +44,8 @@ let(:parent_ts) { Gitlab::Utils.uuid_v7 } let(:checkpoint) { { key: 'value' } } let(:metadata) { { key: 'value' } } - let(:params) { { thread_ts: thread_ts, checkpoint: checkpoint, parent_ts: parent_ts, metadata: metadata } } + let(:lease_id) { nil } + let(:params) { { thread_ts: thread_ts, checkpoint: checkpoint, parent_ts: parent_ts, metadata: metadata, lease_id: lease_id } } let(:path) { "/ai/duo_workflows/workflows/#{workflow.id}/checkpoints" } it 'allows creating multiple checkpoints for a workflow' do @@ -106,6 +107,36 @@ expect(response).to have_gitlab_http_status(:bad_request) expect(json_response['message']).to include("can't be blank") end + + context 'when a lease_id is provided' do + let(:lease) { create(:duo_workflows_lease, workflow: workflow) } + let(:lease_id) { lease.lease_id } + + it 'creates the checkpoint and updates the last_renewed_at' do + lease.update!(last_renewed_at: Time.current - 10.seconds) + checkpoints_count = workflow.checkpoints.count + + post api(path, user), params: params + expect(response).to have_gitlab_http_status(:created) + + expect(workflow.reload.checkpoints.count).to eq(checkpoints_count + 1) + expect(lease.reload.last_renewed_at).to be_within(1.second).of(Time.current) + end + + context 'when the lease_id is not current' do + let(:lease_id) { 'outdated_lease_id' } + + it 'returns a 400 response and does not create the checkpoint' do + checkpoints_count = workflow.checkpoints.count + + post api(path, user), params: params + + expect(response).to have_gitlab_http_status(:bad_request) + expect(json_response['message']).to include("Lease does not match existing lease for workflow") + expect(workflow.reload.checkpoints.count).to eq(checkpoints_count) + end + end + end end describe 'GET /ai/duo_workflows/workflows/:id/checkpoints' do diff --git a/ee/spec/requests/api/graphql/mutations/ai/duo_workflows/workflow_lease_renew_spec.rb b/ee/spec/requests/api/graphql/mutations/ai/duo_workflows/workflow_lease_renew_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..5be80117f26635df2bc82c5f04045ec2e0d88045 --- /dev/null +++ b/ee/spec/requests/api/graphql/mutations/ai/duo_workflows/workflow_lease_renew_spec.rb @@ -0,0 +1,75 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe 'AiDuoWorkflowLeaseRenew', :saas, feature_category: :agent_foundations do + include GraphqlHelpers + + let(:user) { create(:user) } + + # TODO: let_it_be was messed up. Need to figure out why. + let(:group) { create(:group, developers: [user], experiment_features_enabled: true) } + let(:project) { create(:project, group: group) } + + let!(:add_on_purchase) do + create(:gitlab_subscription_add_on_purchase, :duo_enterprise, :active, namespace: group) + end + + let!(:add_on_assignment) do + create(:gitlab_subscription_user_add_on_assignment, user: user, add_on_purchase: add_on_purchase) + end + + let(:workflow) { create(:duo_workflows_workflow, project: project, user: current_user) } + + let(:current_user) { user } + let(:input) do + { + 'workflow_id' => workflow.to_global_id, + 'lease_id' => lease_id + } + end + + let(:mutation) do + graphql_mutation(:ai_duo_workflow_lease_renew, input, + <<-QL + errors + lastRenewedAt + QL + ) + end + + let(:mutation_response) { graphql_mutation_response(:ai_duo_workflow_lease_renew) } + let(:lease) { ::Ai::DuoWorkflows::Lease.create!(workflow_id: workflow.id, lease_id: "the-lease-id") } + let(:lease_id) { lease.lease_id } + + subject(:request) { post_graphql_mutation(mutation, current_user: current_user) } + + before do + stub_licensed_features(ai_workflows: true, agentic_chat: true, ai_chat: true) + project.root_ancestor.update!(experiment_features_enabled: true) + end + + context 'when lease is current' do + it 'returns the workflow and the leaseId' do + post_graphql_mutation(mutation, current_user: current_user) + + expect(response).to have_gitlab_http_status(:success) + expect(json_response['errors']).not_to be_present + expect(mutation_response['errors']).not_to be_present + + expect(mutation_response['lastRenewedAt']).to be_present + end + end + + context 'when the lease is not valid' do + let(:lease_id) { 'invalid-lease-id' } + + it 'returns an error' do + post_graphql_mutation(mutation, current_user: current_user) + + expect(response).to have_gitlab_http_status(:success) + expect(json_response['errors']).not_to be_present + expect(mutation_response['errors']).to include('Lease was not updated. Likely this lease is no longer valid.') + end + end +end diff --git a/ee/spec/requests/api/graphql/mutations/ai/duo_workflows/workflow_with_lease_spec.rb b/ee/spec/requests/api/graphql/mutations/ai/duo_workflows/workflow_with_lease_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..fdb63acacb8be44f044add17f90e206cd48d8fb1 --- /dev/null +++ b/ee/spec/requests/api/graphql/mutations/ai/duo_workflows/workflow_with_lease_spec.rb @@ -0,0 +1,78 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe 'AiDuoWorkflowWithLease', :saas, feature_category: :agent_foundations do + include GraphqlHelpers + + let(:user) { create(:user) } + + # TODO: let_it_be was messed up. Need to figure out why. + let(:group) { create(:group, developers: [user], experiment_features_enabled: true) } + let(:project) { create(:project, group: group) } + + let!(:add_on_purchase) do + create(:gitlab_subscription_add_on_purchase, :duo_enterprise, :active, namespace: group) + end + + let!(:add_on_assignment) do + create(:gitlab_subscription_user_add_on_assignment, user: user, add_on_purchase: add_on_purchase) + end + + let(:workflow) { create(:duo_workflows_workflow, project: project, user: current_user) } + + let(:current_user) { user } + let(:input) do + { + 'workflow_id' => workflow.to_global_id + } + end + + let(:mutation) do + graphql_mutation(:ai_duo_workflow_with_lease, input, + <<-QL + errors + workflow { + id + } + leaseId + QL + ) + end + + let(:mutation_response) { graphql_mutation_response(:ai_duo_workflow_with_lease) } + + subject(:request) { post_graphql_mutation(mutation, current_user: current_user) } + + before do + stub_licensed_features(ai_workflows: true, agentic_chat: true, ai_chat: true) + project.root_ancestor.update!(experiment_features_enabled: true) + end + + context 'workflow is not already leased' do + it 'returns the workflow and the leaseId' do + post_graphql_mutation(mutation, current_user: current_user) + + expect(response).to have_gitlab_http_status(:success) + expect(json_response['errors']).not_to be_present + expect(mutation_response['errors']).not_to be_present + + expect(mutation_response['workflow']['id']).to eq("gid://gitlab/Ai::DuoWorkflows::Workflow/#{workflow.id}") + expect(mutation_response['leaseId']).to eq(workflow.current_lease.lease_id) + end + end + + context 'when workflow is already leased' do + before do + ::Ai::DuoWorkflows::Lease.create!(workflow: workflow, last_renewed_at: Time.current) + end + + it 'returns an error' do + post_graphql_mutation(mutation, current_user: current_user) + + expect(response).to have_gitlab_http_status(:success) + expect(json_response['errors']).not_to be_present + expect(mutation_response['errors']).to include('Failed to acquire lease') + end + end +end diff --git a/lib/api/helpers.rb b/lib/api/helpers.rb index 630deb64de6a4e3b961c84f3b90719f899e7fa4f..350b238733517eb37e1d475085a9564b3e3e61c0 100644 --- a/lib/api/helpers.rb +++ b/lib/api/helpers.rb @@ -661,7 +661,36 @@ def set_status_code_in_env(status) def handle_api_exception(exception) if report_exception?(exception) define_params_for_grape_middleware - Gitlab::ApplicationContext.push(user: current_user, remote_ip: request.ip) + # TODO: Open seperate MR about this. It raises + # + # ========================================================================================== + # Gitlab::Auth::InsufficientScopeError + # + # lib/gitlab/auth/auth_finders.rb, line 219 + # ----------------------------------------- + # + # ``` ruby + # 214 access_token.reload if reset_token # rubocop:disable Cop/ActiveRecordAssociationReload + # 215 + # 216 case AccessTokenValidationService.new(access_token, request: request).validate(scopes: scopes) + # 217 when AccessTokenValidationService::INSUFFICIENT_SCOPE + # 218 save_auth_failure_in_application_context(access_token, :insufficient_scope, scopes) if save_auth_context + # > 219 raise InsufficientScopeError, scopes + # 220 when AccessTokenValidationService::EXPIRED + # 221 save_auth_failure_in_application_context(access_token, :token_expired, scopes) if save_auth_context + # 222 raise ExpiredError + # 223 when AccessTokenValidationService::REVOKED + # 224 save_auth_failure_in_application_context(access_token, :token_revoked, scopes) if save_auth_context + # ``` + # + # App backtrace + # ------------- + # + # - lib/gitlab/auth/auth_finders.rb:219:in `validate_and_save_access_token!' + # - lib/api/helpers.rb:91:in `current_user' + # - lib/api/helpers.rb:664:in `handle_api_exception' + # - lib/api/api.rb:175:in `block in ' + #Gitlab::ApplicationContext.push(user: current_user, remote_ip: request.ip) Gitlab::ErrorTracking.track_exception(exception) end