diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 85edc94d..cf5d5695 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -642,6 +642,7 @@ struct ActiveProviderRunSlot { turn_control: Option, cancellation_reason: Option, committed_provider_batch: Option, + finished_provider_batch: Option, command_action_refs: HashMap, command_monitor: Option, pending_monitor_observation: Option, @@ -706,6 +707,8 @@ struct ActiveProviderRunSnapshot { did_input_contain_user_query: bool, persistence_offset: usize, committed_provider_batch: Option, + #[serde(default)] + finished_provider_batch: Option, command_action_refs: HashMap, command_monitor: Option, pending_monitor_observation: Option, @@ -744,6 +747,7 @@ impl ActiveProviderRunSnapshot { did_input_contain_user_query: slot.did_input_contain_user_query, persistence_offset: checkpoint.persistence_offset, committed_provider_batch: slot.committed_provider_batch.clone(), + finished_provider_batch: slot.finished_provider_batch.clone(), command_action_refs: slot.command_action_refs.clone(), command_monitor: slot.command_monitor.clone(), pending_monitor_observation: slot.pending_monitor_observation.clone(), @@ -813,6 +817,13 @@ impl ActiveProviderRunSnapshot { { return Err("committed provider batch belongs to a different run".to_string()); } + if self + .finished_provider_batch + .as_ref() + .is_some_and(|work_id| &work_id.run_id != run_id) + { + return Err("finished provider batch belongs to a different run".to_string()); + } for (action_id, execution_ref) in &self.command_action_refs { if execution_ref.conversation_id != conversation_id || &execution_ref.run_id != run_id @@ -891,6 +902,34 @@ enum ProviderFinishedActionDisposition { Resume, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ProviderBatchSignal { + BatchCommitted, + ActionsFinished, +} + +fn record_provider_batch_signal( + committed_work_id: &mut Option, + finished_work_id: &mut Option, + work_id: &ExternalWorkId, + signal: ProviderBatchSignal, +) -> bool { + if committed_work_id + .as_ref() + .is_some_and(|recorded_work_id| recorded_work_id != work_id) + || finished_work_id + .as_ref() + .is_some_and(|recorded_work_id| recorded_work_id != work_id) + { + return false; + } + match signal { + ProviderBatchSignal::BatchCommitted => *committed_work_id = Some(work_id.clone()), + ProviderBatchSignal::ActionsFinished => *finished_work_id = Some(work_id.clone()), + } + committed_work_id.as_ref() == Some(work_id) && finished_work_id.as_ref() == Some(work_id) +} + fn normalize_restored_provider_snapshot( snapshot: &mut ActiveProviderRunSnapshot, ) -> Result<(), String> { @@ -923,6 +962,7 @@ fn normalize_restored_provider_snapshot( } snapshot.committed_provider_batch = None; } + snapshot.finished_provider_batch = None; Ok(()) } @@ -4342,6 +4382,7 @@ impl BlocklistAIController { turn_control: None, cancellation_reason: None, committed_provider_batch: None, + finished_provider_batch: None, command_action_refs: HashMap::new(), command_monitor: None, pending_monitor_observation: None, @@ -4629,6 +4670,7 @@ impl BlocklistAIController { did_input_contain_user_query, persistence_offset, committed_provider_batch, + finished_provider_batch, command_action_refs, command_monitor, pending_monitor_observation, @@ -4721,6 +4763,7 @@ impl BlocklistAIController { turn_control: None, cancellation_reason: None, committed_provider_batch, + finished_provider_batch, command_action_refs, command_monitor, pending_monitor_observation, @@ -5281,6 +5324,8 @@ impl BlocklistAIController { }; match block { ProviderRunBlock::Tools(batch) => { + slot.committed_provider_batch = None; + slot.finished_provider_batch = None; slot.run = Some(run); if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { self.fail_active_provider_run( @@ -5445,10 +5490,17 @@ impl BlocklistAIController { ) { return; } + let mut should_resume = false; let should_drive = match run.coordinator.apply_tool_lifecycle(execution_ref, event) { Ok(ProviderToolLifecycleOutcome::Pending) => false, Ok(ProviderToolLifecycleOutcome::BatchCommitted) => { - slot.committed_provider_batch = Some(execution_ref.work_id()); + let work_id = execution_ref.work_id(); + should_resume = record_provider_batch_signal( + &mut slot.committed_provider_batch, + &mut slot.finished_provider_batch, + &work_id, + ProviderBatchSignal::BatchCommitted, + ); false } Err(error) => { @@ -5471,7 +5523,9 @@ impl BlocklistAIController { ); return; } - if should_drive { + if should_resume { + self.handle_provider_actions_finished(conversation_id, execution_ref, ctx); + } else if should_drive { self.drive_active_provider_run(conversation_id, ctx); } } @@ -5655,11 +5709,31 @@ impl BlocklistAIController { )) }) .unwrap_or(ProviderFinishedActionDisposition::Ignore); + let work_id = execution_ref.work_id(); + let should_resume = if disposition == ProviderFinishedActionDisposition::Ignore { + false + } else { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + record_provider_batch_signal( + &mut slot.committed_provider_batch, + &mut slot.finished_provider_batch, + &work_id, + ProviderBatchSignal::ActionsFinished, + ) + }; if disposition == ProviderFinishedActionDisposition::AwaitBatchCommit { + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist finished provider action phase: {error}"), + ctx, + ); + } return; } - let work_id = execution_ref.work_id(); - let results = (disposition == ProviderFinishedActionDisposition::Resume).then(|| { + let results = should_resume.then(|| { self.action_model .as_ref(ctx) .provider_finished_action_results(conversation_id, &work_id) @@ -5674,7 +5748,7 @@ impl BlocklistAIController { self.fail_active_provider_run(conversation_id, error, ctx); return; } - if disposition != ProviderFinishedActionDisposition::Resume { + if !should_resume { return; } @@ -5682,6 +5756,7 @@ impl BlocklistAIController { return; }; slot.committed_provider_batch = None; + slot.finished_provider_batch = None; let Some(mut run) = slot.run.take() else { return; }; diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index a1b45f2c..3732bb88 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -141,6 +141,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider did_input_contain_user_query: true, persistence_offset: 0, committed_provider_batch: None, + finished_provider_batch: None, command_action_refs: HashMap::new(), command_monitor: None, pending_monitor_observation: None, @@ -599,6 +600,67 @@ fn provider_finished_action_only_resumes_its_committed_batch() { ); } +#[test] +fn provider_batch_resumes_after_both_signals_in_either_order() { + let work_id = ExternalWorkId { + run_id: ProviderRunId::new("current"), + epoch: RunEpoch::new(3), + }; + for signals in [ + [ + super::ProviderBatchSignal::ActionsFinished, + super::ProviderBatchSignal::BatchCommitted, + ], + [ + super::ProviderBatchSignal::BatchCommitted, + super::ProviderBatchSignal::ActionsFinished, + ], + ] { + let mut committed = None; + let mut finished = None; + assert!(!super::record_provider_batch_signal( + &mut committed, + &mut finished, + &work_id, + signals[0], + )); + assert!(super::record_provider_batch_signal( + &mut committed, + &mut finished, + &work_id, + signals[1], + )); + } +} + +#[test] +fn provider_batch_signals_do_not_cross_work_ids() { + let current = ExternalWorkId { + run_id: ProviderRunId::new("current"), + epoch: RunEpoch::new(3), + }; + let stale = ExternalWorkId { + run_id: current.run_id.clone(), + epoch: RunEpoch::new(2), + }; + let mut committed = None; + let mut finished = None; + assert!(!super::record_provider_batch_signal( + &mut committed, + &mut finished, + ¤t, + super::ProviderBatchSignal::BatchCommitted, + )); + assert!(!super::record_provider_batch_signal( + &mut committed, + &mut finished, + &stale, + super::ProviderBatchSignal::ActionsFinished, + )); + assert_eq!(committed, Some(current)); + assert_eq!(finished, None); +} + #[test] fn provider_boundary_prioritizes_completion_and_waits_for_committed_results() { assert_eq!(