diff --git a/ee/lib/tasks/gitlab/duo_workflow/duo_workflow.rake b/ee/lib/tasks/gitlab/duo_workflow/duo_workflow.rake index 957d000f72bae99835b08d1b89e7c10cfa96e7dd..b6995c5287d8bb52707b2a9b4584f6344839f458 100644 --- a/ee/lib/tasks/gitlab/duo_workflow/duo_workflow.rake +++ b/ee/lib/tasks/gitlab/duo_workflow/duo_workflow.rake @@ -30,7 +30,7 @@ namespace :gitlab do raise 'This task cannot be run in production' if Rails.env.production? # rubocop:disable Rake/TopLevelMethodDefinition -- Instance metods withing task scope do not leak - def create_checkpoint_data(workflow_goal, workflow_status, step = 0) + def create_checkpoint_data(workflow_goal, workflow_status, step = 0, previous_checkpoint_thread_ts = nil) # rubocop:enable Rake/TopLevelMethodDefinition timestamp = Time.current.iso8601(6) @@ -132,7 +132,7 @@ namespace :gitlab do } { - parent_ts: step > 0 ? Gitlab::Utils.uuid_v7 : nil, + parent_ts: step > 0 ? previous_checkpoint_thread_ts : nil, thread_ts: Gitlab::Utils.uuid_v7, checkpoint: checkpoint_data, metadata: metadata @@ -401,8 +401,9 @@ namespace :gitlab do end # Create workflow checkpoints + previous_thread_ts = nil checkpoint_configs.each_with_index do |config, checkpoint_index| - checkpoint_data = create_checkpoint_data(goal, config[:workflow_status], checkpoint_index) + checkpoint_data = create_checkpoint_data(goal, config[:workflow_status], checkpoint_index, previous_thread_ts) Ai::DuoWorkflows::Checkpoint.create!( workflow: workflow, @@ -413,6 +414,8 @@ namespace :gitlab do metadata: checkpoint_data[:metadata] ) + # Store the current thread_ts for the next checkpoint's parent_ts + previous_thread_ts = checkpoint_data[:thread_ts] checkpoints_created += 1 end diff --git a/ee/spec/tasks/gitlab/duo_workflow/duo_workflow_spec.rb b/ee/spec/tasks/gitlab/duo_workflow/duo_workflow_spec.rb index 0cedd6981b7a4dee34769041d6938b2eb72d3013..1416672d39603ad21db22f4674a631cd67011c8e 100644 --- a/ee/spec/tasks/gitlab/duo_workflow/duo_workflow_spec.rb +++ b/ee/spec/tasks/gitlab/duo_workflow/duo_workflow_spec.rb @@ -185,6 +185,42 @@ expect(first_message).to include('Starting workflow with goal:') expect(first_message).to include(workflow.goal.split(' (Workflow #').first) end + + it 'creates checkpoints with correct parent_ts relationships' do + run_rake_task('gitlab:duo_workflow:populate', '5', '2') + + workflows = Ai::DuoWorkflows::Workflow.all + workflows.each do |workflow| + checkpoints = workflow.checkpoints.order(:created_at) + + # First checkpoint should have nil parent_ts + expect(checkpoints.first.parent_ts).to be_nil + + # Subsequent checkpoints should have parent_ts referencing the previous checkpoint's thread_ts + if checkpoints.count > 1 + checkpoints[1..-1].each_with_index do |checkpoint, index| + previous_checkpoint = checkpoints[index] + expect(checkpoint.parent_ts).to eq(previous_checkpoint.thread_ts) + end + end + end + end + + it 'ensures parent_ts references actual existing checkpoint thread_ts values' do + run_rake_task('gitlab:duo_workflow:populate', '3', '1') + + workflows = Ai::DuoWorkflows::Workflow.all + all_thread_ts_values = Ai::DuoWorkflows::Checkpoint.pluck(:thread_ts) + + workflows.each do |workflow| + workflow.checkpoints.each do |checkpoint| + if checkpoint.parent_ts.present? + # parent_ts should reference an existing thread_ts + expect(all_thread_ts_values).to include(checkpoint.parent_ts) + end + end + end + end end it 'uses default values when no arguments provided' do