From a203d359bfba8c2e4bdc8d1141c3bd14b859cf8c Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Tue, 18 Aug 2026 01:40:11 -0500 Subject: [PATCH] Rig setup --- AGENTS.md | 6 +- app/src/ai/blocklist/action_model.rs | 4 +- app/src/ai/blocklist/action_model/execute.rs | 25 +- .../action_model/execute/run_agents.rs | 68 +- .../action_model/execute/run_agents_tests.rs | 14 +- .../action_model/execute/start_agent.rs | 18 +- .../action_model/execute/start_agent_tests.rs | 61 +- app/src/ai/blocklist/block/cli_controller.rs | 17 +- app/src/ai/blocklist/controller.rs | 628 +++++++++- .../controller/pending_response_streams.rs | 2 +- .../blocklist/controller/response_stream.rs | 4 + app/src/ai/blocklist/controller_tests.rs | 1006 ++++++++++++++++- .../inline_action/run_agents_card_view.rs | 75 +- .../run_agents_card_view_tests.rs | 29 +- crates/galaxy_agent_core/src/provider_run.rs | 320 ++++++ .../src/provider_run_tests.rs | 307 ++++- 16 files changed, 2495 insertions(+), 89 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f77f7900..71b9a6a7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -142,8 +142,9 @@ Key invariants: - Provider actions and results must correlate by `(conversation_id, run_id, epoch, call_id)`; stale or duplicate callbacks must not advance a run - Action status/result lookups and archived results are keyed by `(conversation_id, action_id)`; callers must supply the owning conversation and must not fall back to a global action-ID search - Action blocked/executing/finished events carry `conversation_id`; UI subscribers must match it, and CLI shell-control mutations must also match the active block's requested-command action ID -- Active provider runs must checkpoint before external work, persist without credentials, normalize unsafe restored states, and reconcile command state before continuing; cancellation intent stays checkpointed until the terminal outcome is projected and `finish_active_provider_run` performs cleanup -- Same-conversation direct-provider follow-ups queue behind the cancelling generation; the old run keeps the active slot until terminal projection and cleanup, then the next generation starts, and stale callbacks are ignored by stream identity +- Active provider runs must checkpoint before external work, persist without credentials, validate deserialized run invariants before normalization or runtime construction, normalize unsafe restored states, and reconcile command state before continuing; cancellation intent stays checkpointed until the terminal outcome is projected and `finish_active_provider_run` performs cleanup +- A restored `AwaitingModel` checkpoint has an uncertain remote outcome and must terminate as an explicit restore failure rather than replaying the call; known recoverable failures observed in-process retain the bounded model-retry lifecycle +- Same-conversation direct-provider follow-ups queue behind the cancelling generation; the old run keeps the active slot until terminal projection and cleanup, queued intent is checkpointed without credentials for restart recovery, and queued-only restore validates provider ownership and terminalizes the abandoned unprepared exchange from its persisted projection/stream identity before starting the successor; the next generation rebuilds provider history after cleanup so it includes the old generation's final committed output, and stale callbacks are ignored by stream identity - Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs` - Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose - Unknown or invalid tool calls receive one correlated synthetic error result and a visible `AgentOutput` message; the durable run owns any continuation @@ -156,6 +157,7 @@ Key invariants: - Loop prevention in `controller.rs` detects repeated tool failures (3+ identical) and injects a corrective instruction - Direct-provider long-running shell follow-ups retain CLI tasks as UI projections while command monitoring and completion stay owned by the same provider run - A direct-provider command completion is only queued when the terminal reports it; the CLI task remains active until the provider run applies that completion at a safe boundary and deactivates it +- Provider command ownership is resolved from the active slot or its durable snapshot by block/action identity; completion arriving during restore is persisted into that snapshot and must never fall back to the legacy assessment path - Direct-provider completed-command assessments are hidden, tool-free root-task turns whose output and hidden input survive CLI-task deactivation and restoration - ACP remains a separate session-owned runtime; direct-provider cleanup must not move ACP lifecycle into `ProviderRun` - Orchestrated child conversations are leaf workers: nested `RunAgents` and legacy `StartAgent` calls are rejected before autonomy or permission bypasses, and child requests do not advertise delegation tools diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index ab666d01..e7ef6c88 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -1733,7 +1733,7 @@ impl BlocklistAIActionModel { } self.provider_tool_executions.extend(refs); self.executor.update(ctx, |executor, ctx| { - executor.mark_restored_actions(&recovery_action_ids, ctx); + executor.mark_restored_actions(conversation_id, &recovery_action_ids, ctx); }); self.queue_actions(actions, conversation_id, ctx); Ok(()) @@ -1981,7 +1981,7 @@ impl BlocklistAIActionModel { return; }; for action in actions_to_cancel.drain(..).collect_vec() { - log::info!( + crate::ai::tool_diagnostics::tool_debug!( "Canceling pending action of type {:?} conversation_id={conversation_id:?} action_id={:?}, reason={:?}", AIAgentActionTypeDiscriminants::from(&action.action), action.id, diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 83c4580c..3feb67a2 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -316,7 +316,7 @@ pub struct BlocklistAIActionExecutor { wait_for_events_executor: ModelHandle, /// The actions currently executing asynchronously, scoped by conversation and action ID. async_executing_actions: AsyncExecutingActions, - restored_action_ids: HashSet, + restored_action_ids: HashSet, /// Reference to the terminal model for checking session sharing state. terminal_model: Arc>, @@ -428,12 +428,18 @@ impl BlocklistAIActionExecutor { pub fn mark_restored_actions( &mut self, + conversation_id: AIConversationId, action_ids: &HashSet, ctx: &mut ModelContext, ) { - self.restored_action_ids.extend(action_ids.iter().cloned()); + self.restored_action_ids.extend( + action_ids + .iter() + .cloned() + .map(|action_id| (conversation_id, action_id)), + ); self.run_agents_executor.update(ctx, |executor, _| { - executor.mark_recovery_actions(action_ids); + executor.mark_recovery_actions(conversation_id, action_ids); }); } @@ -753,7 +759,8 @@ impl BlocklistAIActionExecutor { action.id, std::mem::discriminant(&action.action) ); - let is_restored = self.restored_action_ids.remove(&action.id); + let action_key = (conversation_id, action.id.clone()); + let is_restored = self.restored_action_ids.remove(&action_key); let action_clone = action.clone(); let execution = match &action.action { AIAgentActionType::RequestCommandOutput { .. } @@ -1054,7 +1061,7 @@ impl BlocklistAIActionExecutor { .remove(conversation_id, action_id) { let action_kind = AIAgentActionTypeDiscriminants::from(&running.action.action); - log::info!( + crate::ai::tool_diagnostics::tool_debug!( "Canceling running async action of type {action_kind:?} action_id={action_id:?}, reason={reason:?}" ); if let Some(backtrace) = crate::ai::tool_diagnostics::capture_backtrace() { @@ -1075,11 +1082,11 @@ impl BlocklistAIActionExecutor { }); } else if matches!(running.action.action, AIAgentActionType::RunAgents(..)) { self.run_agents_executor.update(ctx, |executor, ctx| { - executor.cancel_execution(&running.action.id, ctx); + executor.cancel_execution(conversation_id, &running.action.id, ctx); }); } else if matches!(running.action.action, AIAgentActionType::StartAgent { .. }) { self.start_agent_executor.update(ctx, |executor, _| { - executor.cancel_execution(&running.action.id); + executor.cancel_execution(conversation_id, &running.action.id); }); } else if let AIAgentActionType::WaitForEvents { tool_call_id, .. } = &running.action.action @@ -1123,7 +1130,9 @@ impl BlocklistAIActionExecutor { } fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext) -> bool { - if self.restored_action_ids.contains(&input.action.id) + if self + .restored_action_ids + .contains(&(input.conversation_id, input.action.id.clone())) || cfg!(feature = "bedrock_smoke_test") { return true; diff --git a/app/src/ai/blocklist/action_model/execute/run_agents.rs b/app/src/ai/blocklist/action_model/execute/run_agents.rs index e555b951..eb727025 100644 --- a/app/src/ai/blocklist/action_model/execute/run_agents.rs +++ b/app/src/ai/blocklist/action_model/execute/run_agents.rs @@ -68,8 +68,8 @@ struct ExistingLaunchedAgent { } pub struct RunAgentsExecutor { - pending: HashMap, - recovery_action_ids: HashSet, + pending: HashMap<(AIConversationId, AIAgentActionId), PendingRunAgents>, + recovery_action_ids: HashSet<(AIConversationId, AIAgentActionId)>, launched_agents: HashMap>, start_agent_executor: ModelHandle, terminal_view_id: EntityId, @@ -78,10 +78,12 @@ pub struct RunAgentsExecutor { /// Lifecycle events for in-flight dispatches. pub enum RunAgentsExecutorEvent { SpawningStarted { + conversation_id: AIConversationId, action_id: AIAgentActionId, snapshot: RunAgentsSpawningSnapshot, }, SpawningFinished { + conversation_id: AIConversationId, action_id: AIAgentActionId, }, ChildConversationCreated { @@ -127,12 +129,26 @@ impl RunAgentsExecutor { } } - pub fn is_pending(&self, action_id: &AIAgentActionId) -> bool { - self.pending.contains_key(action_id) + pub fn is_pending( + &self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> bool { + self.pending + .contains_key(&(conversation_id, action_id.clone())) } - pub fn mark_recovery_actions(&mut self, action_ids: &HashSet) { - self.recovery_action_ids.extend(action_ids.iter().cloned()); + pub fn mark_recovery_actions( + &mut self, + conversation_id: AIConversationId, + action_ids: &HashSet, + ) { + self.recovery_action_ids.extend( + action_ids + .iter() + .cloned() + .map(|action_id| (conversation_id, action_id)), + ); } pub(crate) fn terminal_view_id(&self) -> EntityId { @@ -142,18 +158,21 @@ impl RunAgentsExecutor { /// Cancels the parent tool wait without cancelling independently-running children. pub(super) fn cancel_execution( &mut self, + conversation_id: AIConversationId, action_id: &AIAgentActionId, ctx: &mut ModelContext, ) { - self.recovery_action_ids.remove(action_id); + let action_key = (conversation_id, action_id.clone()); + self.recovery_action_ids.remove(&action_key); let detached_dispatches = self.start_agent_executor.update(ctx, |executor, _| { - executor.cancel_dispatches_for_action(action_id) + executor.cancel_dispatches_for_action(conversation_id, action_id) }); - log::info!( + crate::ai::tool_diagnostics::tool_debug!( "RunAgents cancellation detached {detached_dispatches} pending child dispatch(es) for action {action_id}" ); - if self.pending.remove(action_id).is_some() { + if self.pending.remove(&action_key).is_some() { ctx.emit(RunAgentsExecutorEvent::SpawningFinished { + conversation_id, action_id: action_id.clone(), }); } @@ -224,7 +243,8 @@ impl RunAgentsExecutor { ) -> async_channel::Receiver { let (sender, receiver) = async_channel::bounded(1); - if self.pending.contains_key(&action_id) { + let action_key = (parent_conversation_id, action_id.clone()); + if self.pending.contains_key(&action_key) { log::warn!("RunAgentsExecutor: dispatch reentered for {action_id:?}; rejecting"); #[cfg(not(target_family = "wasm"))] log_run_agents_event( @@ -265,7 +285,7 @@ impl RunAgentsExecutor { agent_count: request.agent_run_configs.len(), }; self.pending - .insert(action_id.clone(), PendingRunAgents::Publishing); + .insert(action_key, PendingRunAgents::Publishing); #[cfg(not(target_family = "wasm"))] log_run_agents_event( ctx, @@ -280,6 +300,7 @@ impl RunAgentsExecutor { }), ); ctx.emit(RunAgentsExecutorEvent::SpawningStarted { + conversation_id: parent_conversation_id, action_id: action_id.clone(), snapshot, }); @@ -294,7 +315,7 @@ impl RunAgentsExecutor { request }, move |me, request, ctx| { - if !me.is_pending(&action_id_for_wait) { + if !me.is_pending(parent_conversation_id, &action_id_for_wait) { return; } me.dispatch_children_for_prepared_request( @@ -320,7 +341,7 @@ impl RunAgentsExecutor { ctx: &mut ModelContext, ) -> async_channel::Receiver { let (sender, receiver) = async_channel::bounded(1); - if self.pending.contains_key(&action_id) { + if self.is_pending(parent_conversation_id, &action_id) { let _ = sender.try_send(RunAgentsResult::Cancelled); return receiver; } @@ -333,6 +354,7 @@ impl RunAgentsExecutor { agent_count: request.agent_run_configs.len(), }; ctx.emit(RunAgentsExecutorEvent::SpawningStarted { + conversation_id: parent_conversation_id, action_id: action_id.clone(), snapshot, }); @@ -356,8 +378,10 @@ impl RunAgentsExecutor { sender: async_channel::Sender, ctx: &mut ModelContext, ) { - self.pending - .insert(action_id.clone(), PendingRunAgents::Spawning); + self.pending.insert( + (parent_conversation_id, action_id.clone()), + PendingRunAgents::Spawning, + ); let parent_run_id = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&parent_conversation_id) .and_then(|c| c.run_id()); @@ -527,7 +551,7 @@ impl RunAgentsExecutor { resolved_slots }, move |me, resolved_slots, ctx| { - if !me.is_pending(&action_id_for_aggr) { + if !me.is_pending(parent_conversation_id_for_result, &action_id_for_aggr) { return; } let timed_out_request_ids = resolved_slots @@ -603,8 +627,12 @@ impl RunAgentsExecutor { execution_mode: launched_mode, agents, }; - me.pending.remove(&action_id_for_aggr); + me.pending.remove(&( + parent_conversation_id_for_result, + action_id_for_aggr.clone(), + )); ctx.emit(RunAgentsExecutorEvent::SpawningFinished { + conversation_id: parent_conversation_id_for_result, action_id: action_id_for_aggr, }); let _ = sender.try_send(result); @@ -624,7 +652,9 @@ impl RunAgentsExecutor { let mut request = request.clone(); let action_id = id.clone(); let parent_conversation_id = input.conversation_id; - let is_recovery = self.recovery_action_ids.remove(&action_id); + let is_recovery = self + .recovery_action_ids + .remove(&(parent_conversation_id, action_id.clone())); let recovery_children = if is_recovery { prepare_recovery_request_for_execution(&mut request, parent_conversation_id, ctx); diff --git a/app/src/ai/blocklist/action_model/execute/run_agents_tests.rs b/app/src/ai/blocklist/action_model/execute/run_agents_tests.rs index a4802f20..e8e38a11 100644 --- a/app/src/ai/blocklist/action_model/execute/run_agents_tests.rs +++ b/app/src/ai/blocklist/action_model/execute/run_agents_tests.rs @@ -386,7 +386,8 @@ fn recovered_run_agents_reattaches_existing_child_and_launches_only_missing_chil title: String::new(), }); state.executor.update(&mut app, |executor, _| { - executor.mark_recovery_actions(&HashSet::from([action.id.clone()])); + executor + .mark_recovery_actions(state.conversation_id, &HashSet::from([action.id.clone()])); }); let execution = state.executor.update(&mut app, |executor, ctx| { @@ -481,7 +482,8 @@ fn cancelling_recovered_run_agents_keeps_persisted_child_running() { }); let action = remote_run_agents_action("oz"); state.executor.update(&mut app, |executor, _| { - executor.mark_recovery_actions(&HashSet::from([action.id.clone()])); + executor + .mark_recovery_actions(state.conversation_id, &HashSet::from([action.id.clone()])); }); let execution = state.executor.update(&mut app, |executor, ctx| { executor @@ -503,7 +505,7 @@ fn cancelling_recovered_run_agents_keeps_persisted_child_running() { }; state.executor.update(&mut app, |executor, ctx| { - executor.cancel_execution(&action.id, ctx); + executor.cancel_execution(state.conversation_id, &action.id, ctx); }); let async_result = execute_future.await; let result = app.update(|ctx| on_complete(async_result, ctx)); @@ -1090,9 +1092,9 @@ fn cancel_during_plan_publication_does_not_dispatch_children() { // The action is awaiting plan publication, so it's pending but no children dispatched yet. assert!(matches!(execution, AnyActionExecution::Async { .. })); state.executor.update(&mut app, |executor, ctx| { - assert!(executor.is_pending(&action_id)); - executor.cancel_execution(&action_id, ctx); - assert!(!executor.is_pending(&action_id)); + assert!(executor.is_pending(state.conversation_id, &action_id)); + executor.cancel_execution(state.conversation_id, &action_id, ctx); + assert!(!executor.is_pending(state.conversation_id, &action_id)); }); // Finish publishing the plan, which resolves the wait the dispatch was blocked on. diff --git a/app/src/ai/blocklist/action_model/execute/start_agent.rs b/app/src/ai/blocklist/action_model/execute/start_agent.rs index 95a843eb..167e2534 100644 --- a/app/src/ai/blocklist/action_model/execute/start_agent.rs +++ b/app/src/ai/blocklist/action_model/execute/start_agent.rs @@ -788,12 +788,18 @@ impl StartAgentExecutor { self.pending.contains_key(&request_id) } - pub fn cancel_dispatches_for_action(&mut self, action_id: &AIAgentActionId) -> usize { + pub fn cancel_dispatches_for_action( + &mut self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> usize { let request_ids = self .pending .iter() .filter_map(|(request_id, pending)| { - (&pending.action_id == action_id).then_some(*request_id) + (pending.parent_conversation_id == conversation_id + && &pending.action_id == action_id) + .then_some(*request_id) }) .collect::>(); let detached_count = request_ids.len(); @@ -805,8 +811,12 @@ impl StartAgentExecutor { /// Cancels only the caller's pending tool wait. A child that was already created keeps /// running independently and remains available in conversation history. - pub(super) fn cancel_execution(&mut self, action_id: &AIAgentActionId) { - self.cancel_dispatches_for_action(action_id); + pub(super) fn cancel_execution( + &mut self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) { + self.cancel_dispatches_for_action(conversation_id, action_id); } pub(super) fn preprocess_action( diff --git a/app/src/ai/blocklist/action_model/execute/start_agent_tests.rs b/app/src/ai/blocklist/action_model/execute/start_agent_tests.rs index 765ae601..0d8c2543 100644 --- a/app/src/ai/blocklist/action_model/execute/start_agent_tests.rs +++ b/app/src/ai/blocklist/action_model/execute/start_agent_tests.rs @@ -986,7 +986,7 @@ fn cancelling_standalone_start_agent_drops_dispatch_but_keeps_child_running() { }); executor.update(&mut app, |executor, _| { - executor.cancel_execution(&action.id); + executor.cancel_execution(parent_conversation_id, &action.id); }); executor.read(&app, |executor, _| assert!(executor.pending.is_empty())); @@ -1006,6 +1006,65 @@ fn cancelling_standalone_start_agent_drops_dispatch_but_keeps_child_running() { }); } +#[test] +fn cancelling_duplicate_action_id_detaches_only_the_matching_conversation() { + App::test((), |mut app| async move { + initialize_history_persistence_for_tests(&mut app); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let executor = app.add_model(StartAgentExecutor::new); + let terminal_view_id = EntityId::new(); + let first_conversation = history_model.update(&mut app, |history, ctx| { + history.start_new_conversation(terminal_view_id, false, false, false, ctx) + }); + let second_conversation = history_model.update(&mut app, |history, ctx| { + history.start_new_conversation(terminal_view_id, false, false, false, ctx) + }); + let action = build_start_agent_action( + StartAgentVersion::V1, + StartAgentExecutionMode::local_with_defaults(), + ); + + let first = executor.update(&mut app, |executor, ctx| { + executor + .execute( + ExecuteActionInput { + action: &action, + conversation_id: first_conversation, + }, + ctx, + ) + .into() + }); + let second = executor.update(&mut app, |executor, ctx| { + executor + .execute( + ExecuteActionInput { + action: &action, + conversation_id: second_conversation, + }, + ctx, + ) + .into() + }); + assert!(matches!(first, AnyActionExecution::Async { .. })); + assert!(matches!(second, AnyActionExecution::Async { .. })); + + executor.update(&mut app, |executor, _| { + executor.cancel_execution(first_conversation, &action.id); + assert_eq!(executor.pending.len(), 1); + assert_eq!( + executor + .pending + .values() + .next() + .unwrap() + .parent_conversation_id, + second_conversation + ); + }); + }); +} + #[test] fn removing_direct_provider_child_resolves_pending_wait() { App::test((), |mut app| async move { diff --git a/app/src/ai/blocklist/block/cli_controller.rs b/app/src/ai/blocklist/block/cli_controller.rs index 24615ed9..59bb519b 100644 --- a/app/src/ai/blocklist/block/cli_controller.rs +++ b/app/src/ai/blocklist/block/cli_controller.rs @@ -356,16 +356,17 @@ impl CLISubagentController { drop(terminal_model); let provider_accepted_completion = completion.as_ref().is_some_and(|completion| { + let provider_completion = PendingProviderCommandCompletion::new( + completion.completed_command.block_id.clone(), + completion.initial_requested_command_action_id.clone(), + completion.completed_command.command.clone(), + completion.completed_command.grid_contents.clone(), + completion.exit_code, + ); me.controller.update(ctx, |controller, ctx| { - controller.accept_provider_command_completion( + controller.offer_provider_command_completion( completion.conversation_id, - PendingProviderCommandCompletion::new( - completion.completed_command.block_id.clone(), - completion.initial_requested_command_action_id.clone(), - completion.completed_command.command.clone(), - completion.completed_command.grid_contents.clone(), - completion.exit_code, - ), + provider_completion, ctx, ) }) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 3eef0c2f..21a77439 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -544,7 +544,7 @@ impl ActiveProviderRun { const ACTIVE_PROVIDER_RUN_SNAPSHOT_VERSION: u32 = 1; -#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] struct ProviderProjectionTarget { task_id: TaskId, exchange_id: AIAgentExchangeId, @@ -684,6 +684,40 @@ struct QueuedProviderRun { request_params: api::RequestParams, } +struct PreparedQueuedProviderRunRestoration { + snapshot: QueuedProviderRunSnapshot, + root_task_id: TaskId, + base_provider_config: crate::ai::provider::ProviderConfig, + cli_provider_config: crate::ai::provider::ProviderConfig, + request_params: api::RequestParams, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct QueuedProviderRunSnapshot { + run_id: ProviderRunId, + projection_target: ProviderProjectionTarget, + did_input_contain_user_query: bool, + supported_tools_override: Option>, +} + +const QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION: u32 = 2; + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct AbandonedProviderGenerationSnapshot { + run_id: ProviderRunId, + projection_target: ProviderProjectionTarget, + response_stream_id: String, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct QueuedProviderRunsOnlySnapshot { + version: u32, + active_run_id: ProviderRunId, + #[serde(default)] + abandoned_generation: Option, + queued_follow_ups: Vec, +} + #[derive(Clone)] struct ActiveProviderRunCheckpoint { run: ProviderRun, @@ -751,6 +785,8 @@ struct ActiveProviderRunSnapshot { pending_monitor_observation: Option, pending_command_completion: Option, monitor_prose_continuations: usize, + #[serde(default)] + queued_follow_ups: Vec, } impl ActiveProviderRunSnapshot { @@ -791,6 +827,7 @@ impl ActiveProviderRunSnapshot { pending_monitor_observation: slot.pending_monitor_observation.clone(), pending_command_completion: slot.pending_command_completion.clone(), monitor_prose_continuations: slot.monitor_prose_continuations, + queued_follow_ups: Vec::new(), }) } @@ -803,6 +840,10 @@ impl ActiveProviderRunSnapshot { snapshot.version )); } + snapshot + .run + .validate_restored_state() + .map_err(|error| error.to_string())?; Ok(snapshot) } @@ -997,6 +1038,10 @@ fn recoverable_run_agents_call_ids( fn normalize_restored_provider_snapshot( snapshot: &mut ActiveProviderRunSnapshot, ) -> Result<(), String> { + snapshot + .run + .validate_restored_state() + .map_err(|error| error.to_string())?; let recoverable_call_ids = recoverable_run_agents_call_ids(snapshot)?; let normalization = snapshot .run @@ -1090,6 +1135,19 @@ fn apply_restored_provider_command_evidence( Ok(()) } +fn merge_completion_offered_during_restore( + prepared: &mut ActiveProviderRunSnapshot, + latest: ActiveProviderRunSnapshot, +) { + if latest.run.id() != prepared.run.id() { + return; + } + prepared.pending_command_completion = latest.pending_command_completion; + if prepared.pending_command_completion.is_some() { + prepared.pending_monitor_observation = None; + } +} + fn restored_projection_was_initialized( has_output: bool, has_server_output_id: bool, @@ -1104,6 +1162,20 @@ fn restored_projection_was_initialized( } } +fn refresh_queued_provider_history( + request_params: &mut api::RequestParams, + conversation: &AIConversation, +) { + request_params.tasks = conversation.compute_active_tasks(); + request_params.root_task_id = Some(conversation.get_root_task_id().to_string()); + if conversation.is_child_agent_conversation() { + request_params.orchestration_enabled = false; + } + request_params.message_history = conversation.bedrock_message_history().to_vec(); + request_params.tool_result_archive = conversation.tool_result_archive().to_vec(); + request_params.progressive_summary = conversation.progressive_summary().map(str::to_owned); +} + fn provider_execution_matches_active_work( run_id: &ProviderRunId, active_work_id: Option<&ExternalWorkId>, @@ -4493,11 +4565,15 @@ impl BlocklistAIController { .all_inputs() .any(|input| input.is_passive_request()); - // Make sure there's no existing response stream for the conversation. If - // there is, something has gone wrong. - if self + // A same-conversation direct-provider follow-up is allowed to create its exchange while + // the cancelled generation is still terminalizing. Its provider run is queued below and + // cannot take the active slot until cleanup removes the old generation. Other overlapping + // streams remain invalid. + let has_in_flight_response = self .in_flight_response_streams - .has_active_stream_for_conversation(conversation_id, ctx) + .has_active_stream_for_conversation(conversation_id, ctx); + if has_in_flight_response + && !self.provider_generation_is_terminalizing_for_follow_up(conversation_id) { send_telemetry_from_ctx!( TelemetryEvent::AIInputNotSent { @@ -4764,6 +4840,9 @@ impl BlocklistAIController { cli_provider_config, request_params: request_params.clone(), }); + if let Err(error) = self.persist_active_provider_run(conversation_data.id, ctx) { + log::error!("Failed to persist queued provider follow-up: {error}"); + } } else { self.active_provider_runs.insert(conversation_data.id, slot); self.prepare_active_provider_run( @@ -4822,6 +4901,22 @@ impl BlocklistAIController { Ok((conversation_data.id, response_stream_id)) } + fn provider_generation_is_terminalizing_for_follow_up( + &self, + conversation_id: AIConversationId, + ) -> bool { + self.active_provider_runs + .get(&conversation_id) + .is_some_and(|slot| { + matches!( + slot.cancellation_reason, + Some(CancellationReason::FollowUpSubmitted { + is_for_same_conversation: true, + }) + ) && self.in_flight_response_streams.has_stream(&slot.stream_id) + }) + } + fn schedule_restored_provider_runs( &mut self, conversation_ids: &[AIConversationId], @@ -4875,6 +4970,79 @@ impl BlocklistAIController { self.restoring_provider_runs.remove(&conversation_id); return; }; + if let Ok(snapshot) = serde_json::from_str::(&snapshot_json) + { + if !matches!(snapshot.version, 1 | QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION) { + self.fail_restored_provider_run( + conversation_id, + format!( + "unsupported queued provider run snapshot version {}", + snapshot.version + ), + ctx, + ); + return; + } + let validation = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .ok_or_else(|| "queued provider conversation is missing".to_string()) + .and_then(|conversation| { + if conversation.agent_backend() != &AgentBackend::Provider { + return Err("queued provider run belongs to a non-provider conversation" + .to_string()); + } + Ok(()) + }); + if let Err(error) = validation { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + let Some(abandoned_generation) = snapshot.abandoned_generation else { + self.fail_restored_provider_run( + conversation_id, + "queued-only provider snapshot is missing abandoned generation identity" + .to_string(), + ctx, + ); + return; + }; + if abandoned_generation.run_id != snapshot.active_run_id { + self.fail_restored_provider_run( + conversation_id, + "queued provider abandoned generation identity mismatch".to_string(), + ctx, + ); + return; + } + if let Err(error) = validate_queued_provider_run_snapshots( + Some(&snapshot.active_run_id), + Some(&abandoned_generation.projection_target), + &snapshot.queued_follow_ups, + ) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + if let Err(error) = self.reconcile_abandoned_provider_generation( + conversation_id, + abandoned_generation, + ctx, + ) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + if let Err(error) = self.restore_queued_provider_follow_ups( + conversation_id, + snapshot.queued_follow_ups, + ctx, + ) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + self.restoring_provider_runs.remove(&conversation_id); + self.start_next_queued_provider_run(conversation_id, ctx); + return; + } let mut snapshot = match ActiveProviderRunSnapshot::parse(&snapshot_json) { Ok(snapshot) => snapshot, Err(error) => { @@ -4986,6 +5154,63 @@ impl BlocklistAIController { ); } + fn reconcile_abandoned_provider_generation( + &self, + conversation_id: AIConversationId, + abandoned: AbandonedProviderGenerationSnapshot, + ctx: &mut ModelContext, + ) -> Result<(), String> { + let stream_id = ResponseStreamId::from_persisted(abandoned.response_stream_id); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + let existing_target = history_model + .conversation(&conversation_id) + .ok_or_else(|| "queued provider conversation is missing".to_string())? + .provider_projection_target(&stream_id); + if existing_target.is_none() { + history_model + .rebind_provider_projection( + conversation_id, + &abandoned.projection_target.task_id, + abandoned.projection_target.exchange_id, + stream_id.clone(), + self.terminal_surface_id, + ctx, + ) + .map_err(|error| { + format!("failed to rebind abandoned provider projection: {error:?}") + })?; + } + let target = history_model + .conversation(&conversation_id) + .and_then(|conversation| conversation.provider_projection_target(&stream_id)) + .expect("abandoned provider projection was rebound"); + if target + != ( + abandoned.projection_target.task_id.clone(), + abandoned.projection_target.exchange_id, + ) + { + return Err( + "abandoned provider generation projection identity mismatch".to_string() + ); + } + history_model.mark_response_stream_cancelled( + &stream_id, + conversation_id, + self.terminal_surface_id, + CancellationReason::FollowUpSubmitted { + is_for_same_conversation: true, + }, + ctx, + ); + history_model + .conversation_mut(&conversation_id) + .expect("queued provider conversation was validated") + .cleanup_completed_response_stream(&stream_id); + Ok(()) + }) + } + fn reconcile_restored_provider_command( &self, conversation_id: AIConversationId, @@ -5037,10 +5262,12 @@ impl BlocklistAIController { || self.active_provider_runs.contains_key(&conversation_id) { self.restoring_provider_runs.remove(&conversation_id); + self.restoring_provider_command_completions + .remove(&conversation_id); return; } let PreparedRestoredProviderRun { - snapshot, + mut snapshot, profiles, projection_was_initialized, } = match result { @@ -5050,6 +5277,22 @@ impl BlocklistAIController { return; } }; + // Completion can arrive while provider runtimes are being rebuilt. Reload only that + // mailbox from the durable snapshot so the prepared run cannot overwrite the offer. + if let Some(latest_snapshot) = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .and_then(AIConversation::active_provider_run_json) + .and_then(|json| ActiveProviderRunSnapshot::parse(json).ok()) + { + merge_completion_offered_during_restore(&mut snapshot, latest_snapshot); + } + if let Some(completion) = self + .restoring_provider_command_completions + .remove(&conversation_id) + { + snapshot.pending_command_completion = Some(completion); + snapshot.pending_monitor_observation = None; + } let ActiveProviderRunSnapshot { version: _, run: provider_run, @@ -5069,6 +5312,7 @@ impl BlocklistAIController { pending_monitor_observation, pending_command_completion, monitor_prose_continuations, + queued_follow_ups, } = snapshot; let run_id = provider_run.id().clone(); let transcript = provider_run.transcript(); @@ -5175,6 +5419,12 @@ impl BlocklistAIController { monitor_prose_continuations, }, ); + if let Err(error) = + self.restore_queued_provider_follow_ups(conversation_id, queued_follow_ups, ctx) + { + self.fail_active_provider_run(conversation_id, error, ctx); + return; + } self.restoring_provider_runs.remove(&conversation_id); BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { history_model.update_conversation_status( @@ -5191,6 +5441,194 @@ impl BlocklistAIController { self.resume_restored_provider_run(conversation_id, ctx); } + fn restore_queued_provider_follow_ups( + &mut self, + conversation_id: AIConversationId, + snapshots: Vec, + ctx: &mut ModelContext, + ) -> Result<(), String> { + let (active_run_id, active_projection_target) = self + .active_provider_runs + .get(&conversation_id) + .map_or((None, None), |slot| { + (Some(&slot.run_id), Some(&slot.projection_target)) + }); + validate_queued_provider_run_snapshots( + active_run_id, + active_projection_target, + &snapshots, + )?; + let history_model = BlocklistAIHistoryModel::handle(ctx); + let mut prepared_runs = Vec::with_capacity(snapshots.len()); + for snapshot in snapshots { + let (request_input, conversation_data, message_history, tool_result_archive, summary) = { + let conversation = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .ok_or_else(|| "queued provider conversation is missing".to_string())?; + let exchange = conversation + .get_task(&snapshot.projection_target.task_id) + .and_then(|task| task.exchange(snapshot.projection_target.exchange_id)) + .ok_or_else(|| "queued provider projection exchange is missing".to_string())?; + let request_input = RequestInput { + conversation_id, + input_messages: HashMap::from([( + snapshot.projection_target.task_id.clone(), + exchange.input.clone(), + )]), + working_directory: exchange.working_directory.clone(), + model_id: exchange.model_id.clone(), + coding_model_id: exchange.coding_model_id.clone(), + cli_agent_model_id: exchange.cli_agent_model_id.clone(), + computer_use_model_id: exchange.computer_use_model_id.clone(), + shared_session_response_initiator: exchange.response_initiator.clone(), + request_start_ts: exchange.start_time, + supported_tools_override: snapshot + .supported_tools_override + .as_ref() + .map(|tools| { + tools + .iter() + .map(|tool| { + ToolType::try_from(*tool).map_err(|_| { + format!("queued provider tool type {tool} is invalid") + }) + }) + .collect::, _>>() + }) + .transpose()?, + }; + let conversation_data = api::ConversationData { + id: conversation_id, + tasks: conversation.compute_active_tasks(), + server_conversation_token: conversation.server_conversation_token().cloned(), + forked_from_conversation_token: conversation + .forked_from_server_conversation_token() + .cloned(), + ambient_agent_task_id: self.ambient_agent_task_id, + existing_suggestions: history_model + .as_ref(ctx) + .existing_suggestions_for_conversation(conversation_id) + .cloned(), + }; + ( + request_input, + conversation_data, + conversation.bedrock_message_history().to_vec(), + conversation.tool_result_archive().to_vec(), + conversation.progressive_summary().map(str::to_owned), + ) + }; + let mut request_params = api::RequestParams::new( + Some(self.terminal_surface_id), + SessionContext::from_session(self.active_session.as_ref(ctx), ctx), + &request_input, + conversation_data, + None, + ctx, + ); + request_params.message_history = message_history; + request_params.tool_result_archive = tool_result_archive; + request_params.progressive_summary = summary; + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .expect("queued provider conversation was validated"); + let root_task_id = conversation.get_root_task_id().clone(); + request_params.root_task_id = Some(root_task_id.to_string()); + if conversation.is_child_agent_conversation() { + request_params.orchestration_enabled = false; + } + let base_provider_config = + ResponseStream::resolve_provider_config(request_params.model.as_str(), ctx); + let cli_provider_config = ResponseStream::resolve_provider_config( + request_params.cli_agent_model.as_str(), + ctx, + ); + prepared_runs.push(PreparedQueuedProviderRunRestoration { + snapshot, + root_task_id, + base_provider_config, + cli_provider_config, + request_params, + }); + } + + // Do not rebind exchanges or register streams until every queued entry validates and its + // request can be rebuilt. A malformed later entry must not make an earlier one executable. + for prepared in prepared_runs { + let PreparedQueuedProviderRunRestoration { + snapshot, + root_task_id, + base_provider_config, + cli_provider_config, + request_params, + } = prepared; + let ai_identifiers = AIIdentifiers { + client_conversation_id: Some(conversation_id), + model_id: Some(request_params.model.clone()), + ..AIIdentifiers::default() + }; + let response_stream = ctx.add_model(|ctx| { + ResponseStream::new_provider_projection(request_params.clone(), ai_identifiers, ctx) + }); + let stream_id = response_stream.as_ref(ctx).id().clone(); + let response_stream_clone = response_stream.clone(); + let did_input_contain_user_query = snapshot.did_input_contain_user_query; + ctx.subscribe_to_model(&response_stream, move |me, _, event, ctx| { + let _ = me.handle_response_stream_event( + did_input_contain_user_query, + event, + &response_stream_clone, + ctx, + ); + }); + history_model + .update(ctx, |history_model, ctx| { + history_model.rebind_provider_projection( + conversation_id, + &snapshot.projection_target.task_id, + snapshot.projection_target.exchange_id, + stream_id.clone(), + self.terminal_surface_id, + ctx, + ) + }) + .map_err(|error| { + format!("failed to rebind queued provider projection: {error:?}") + })?; + self.in_flight_response_streams + .register_additional_stream(stream_id.clone(), response_stream.clone()); + self.queued_provider_runs + .entry(conversation_id) + .or_default() + .push_back(QueuedProviderRun { + slot: ActiveProviderRunSlot { + stream_id, + response_stream, + did_input_contain_user_query, + run_id: snapshot.run_id, + root_task_id, + projection_target: snapshot.projection_target, + run: None, + checkpoint: None, + 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, + pending_command_completion: None, + monitor_prose_continuations: 0, + }, + base_provider_config, + cli_provider_config, + request_params, + }); + } + Ok(()) + } + fn resume_restored_provider_run( &mut self, conversation_id: AIConversationId, @@ -5253,6 +5691,8 @@ impl BlocklistAIController { ctx: &mut ModelContext, ) { self.restoring_provider_runs.remove(&conversation_id); + self.restoring_provider_command_completions + .remove(&conversation_id); BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { history_model.update_conversation_status_with_error( self.terminal_surface_id, @@ -5382,7 +5822,42 @@ impl BlocklistAIController { .active_provider_runs .get(&conversation_id) .ok_or_else(|| "active provider run disappeared before persistence".to_string())?; - let snapshot = ActiveProviderRunSnapshot::from_slot(slot)?; + let queued_follow_ups = self.queued_provider_run_snapshots(conversation_id); + let mut snapshot = match ActiveProviderRunSnapshot::from_slot(slot) { + Ok(snapshot) => snapshot, + Err(error) if slot.run.is_none() && slot.checkpoint.is_none() => { + let persisted = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .and_then(AIConversation::active_provider_run_json); + if let Some(persisted) = persisted { + if let Ok(snapshot) = ActiveProviderRunSnapshot::parse(persisted) { + if snapshot.run.id() != &slot.run_id { + return Err( + "persisted provider run identity does not match active slot".into(), + ); + } + snapshot + } else { + return self.persist_queued_provider_runs_only( + conversation_id, + slot, + queued_follow_ups, + ctx, + ); + } + } else { + let _ = error; + return self.persist_queued_provider_runs_only( + conversation_id, + slot, + queued_follow_ups, + ctx, + ); + } + } + Err(error) => return Err(error), + }; + snapshot.queued_follow_ups = queued_follow_ups; let json = serde_json::to_string(&snapshot) .map_err(|error| format!("failed to serialize active provider run: {error}"))?; BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { @@ -5392,6 +5867,55 @@ impl BlocklistAIController { }) } + fn queued_provider_run_snapshots( + &self, + conversation_id: AIConversationId, + ) -> Vec { + self.queued_provider_runs + .get(&conversation_id) + .into_iter() + .flatten() + .map(|queued| QueuedProviderRunSnapshot { + run_id: queued.slot.run_id.clone(), + projection_target: queued.slot.projection_target.clone(), + did_input_contain_user_query: queued.slot.did_input_contain_user_query, + supported_tools_override: queued + .request_params + .supported_tools_override + .as_ref() + .map(|tools| tools.iter().map(|tool| *tool as i32).collect()), + }) + .collect() + } + + fn persist_queued_provider_runs_only( + &self, + conversation_id: AIConversationId, + active_slot: &ActiveProviderRunSlot, + queued_follow_ups: Vec, + ctx: &mut ModelContext, + ) -> Result<(), String> { + if queued_follow_ups.is_empty() { + return Err("provider run is not prepared".to_string()); + } + let json = serde_json::to_string(&QueuedProviderRunsOnlySnapshot { + version: QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION, + active_run_id: active_slot.run_id.clone(), + abandoned_generation: Some(AbandonedProviderGenerationSnapshot { + run_id: active_slot.run_id.clone(), + projection_target: active_slot.projection_target.clone(), + response_stream_id: active_slot.stream_id.as_str().to_owned(), + }), + queued_follow_ups, + }) + .map_err(|error| format!("failed to serialize queued provider runs: {error}"))?; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model + .persist_active_provider_run_json(conversation_id, Some(json), ctx) + .map_err(|error| format!("failed to persist queued provider runs: {error:?}")) + }) + } + fn clear_persisted_active_provider_run( &self, conversation_id: AIConversationId, @@ -6538,8 +7062,10 @@ impl BlocklistAIController { { return; } - if let Err(error) = self.clear_persisted_active_provider_run(conversation_id, ctx) { - log::error!("Failed to clear persisted provider run during cleanup: {error}"); + if !self.queued_provider_runs.contains_key(&conversation_id) { + if let Err(error) = self.clear_persisted_active_provider_run(conversation_id, ctx) { + log::error!("Failed to clear persisted provider run during cleanup: {error}"); + } } self.active_provider_runs.remove(&conversation_id); self.restoring_provider_runs.remove(&conversation_id); @@ -6566,6 +7092,11 @@ impl BlocklistAIController { conversation_id: AIConversationId, ctx: &mut ModelContext, ) { + // The old generation must be gone before its successor can own the conversation slot. + // This also makes delayed cleanup callbacks harmless: cleanup checks the stream identity. + if self.active_provider_runs.contains_key(&conversation_id) { + return; + } let next = self .queued_provider_runs .get_mut(&conversation_id) @@ -6580,6 +7111,12 @@ impl BlocklistAIController { let Some(next) = next else { return; }; + let mut request_params = next.request_params; + if let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + { + refresh_queued_provider_history(&mut request_params, conversation); + } let stream_id = next.slot.stream_id.clone(); self.active_provider_runs.insert(conversation_id, next.slot); self.prepare_active_provider_run( @@ -6587,7 +7124,7 @@ impl BlocklistAIController { stream_id, next.base_provider_config, next.cli_provider_config, - next.request_params, + request_params, ctx, ); } @@ -6689,6 +7226,19 @@ impl BlocklistAIController { self.active_provider_runs.contains_key(&conversation_id) } + pub(super) fn offer_provider_command_completion( + &mut self, + conversation_id: AIConversationId, + completion: PendingProviderCommandCompletion, + ctx: &mut ModelContext, + ) -> bool { + if self.active_provider_runs.contains_key(&conversation_id) { + self.accept_provider_command_completion(conversation_id, completion, ctx) + } else { + self.persist_restoring_provider_command_completion(conversation_id, completion, ctx) + } + } + pub(super) fn accept_provider_command_completion( &mut self, conversation_id: AIConversationId, @@ -6778,6 +7328,64 @@ impl BlocklistAIController { true } + pub(super) fn persist_restoring_provider_command_completion( + &mut self, + conversation_id: AIConversationId, + mut completion: PendingProviderCommandCompletion, + ctx: &mut ModelContext, + ) -> bool { + if self.active_provider_runs.contains_key(&conversation_id) { + return self.accept_provider_command_completion(conversation_id, completion, ctx); + } + let history_model = BlocklistAIHistoryModel::handle(ctx); + let Some(snapshot_json) = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .and_then(AIConversation::active_provider_run_json) + else { + return false; + }; + let Ok(mut snapshot) = ActiveProviderRunSnapshot::parse(snapshot_json) else { + return false; + }; + if !provider_command_completion_matches( + snapshot.run.id(), + &snapshot.command_action_refs, + snapshot.command_monitor.as_ref(), + &completion.block_id, + completion.initial_requested_command_action_id.as_ref(), + ) { + return false; + } + if completion.command.is_empty() { + completion.command = snapshot + .command_monitor + .as_ref() + .map(|monitor| monitor.command.clone()) + .unwrap_or_default(); + } + match self + .restoring_provider_command_completions + .get(&conversation_id) + { + Some(existing) => return existing == &completion, + None => { + self.restoring_provider_command_completions + .insert(conversation_id, completion.clone()); + } + } + snapshot.pending_monitor_observation = None; + snapshot.pending_command_completion = Some(completion); + match self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) { + Ok(()) => true, + Err(error) => { + log::error!("Failed to persist completion for restoring provider run: {error}"); + // The in-memory mailbox remains the exactly-once owner until restore installs it. + true + } + } + } + pub fn has_active_stream_for_conversation( &self, conversation_id: AIConversationId, diff --git a/app/src/ai/blocklist/controller/pending_response_streams.rs b/app/src/ai/blocklist/controller/pending_response_streams.rs index 63051220..68390a4a 100644 --- a/app/src/ai/blocklist/controller/pending_response_streams.rs +++ b/app/src/ai/blocklist/controller/pending_response_streams.rs @@ -148,7 +148,7 @@ impl PendingResponseStreams { false } else { for response_stream in streams_to_cancel.into_iter() { - log::info!( + crate::ai::tool_diagnostics::tool_debug!( "Canceling active stream for conversation_id={conversation_id:?}, \ reason={reason}" ); diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index 9c7686d2..5e8740fe 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -56,6 +56,10 @@ impl ResponseStreamId { &self.0 } + pub(crate) fn from_persisted(value: String) -> Self { + Self(value) + } + pub fn for_shared_session(init_event: &response_event::StreamInit) -> Self { // Make the stream ID unique per viewing by appending a local UUID // This prevents collisions when replaying the same conversation multiple times diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index be6bde8f..8cadc63e 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -274,6 +274,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + queued_follow_ups: Vec::new(), } } @@ -295,6 +296,687 @@ fn provider_snapshot_persists_cancellation_reason() { assert!(!restored.run.is_terminal()); } +#[test] +fn queued_provider_follow_up_snapshot_survives_restore_round_trip() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + snapshot + .queued_follow_ups + .push(super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("queued-run"), + projection_target: super::ProviderProjectionTarget { + task_id: snapshot.root_task_id.clone(), + exchange_id: AIAgentExchangeId::new(), + }, + did_input_contain_user_query: true, + supported_tools_override: None, + }); + + let json = serde_json::to_string(&snapshot).unwrap(); + let restored = super::ActiveProviderRunSnapshot::parse(&json).unwrap(); + + assert_eq!(restored.queued_follow_ups.len(), 1); + assert_eq!( + restored.queued_follow_ups[0].run_id, + ProviderRunId::new("queued-run") + ); + assert_eq!( + restored.queued_follow_ups[0].projection_target, + snapshot.queued_follow_ups[0].projection_target + ); +} + +#[test] +fn queued_snapshot_prevalidation_rejects_ambiguous_batch_identity() { + let active_run_id = ProviderRunId::new("active-run"); + let active_target = super::ProviderProjectionTarget { + task_id: TaskId::new("active-task".to_string()), + exchange_id: AIAgentExchangeId::new(), + }; + let queued_target = super::ProviderProjectionTarget { + task_id: TaskId::new("queued-task".to_string()), + exchange_id: AIAgentExchangeId::new(), + }; + let snapshot = |run_id: &str, projection_target| super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new(run_id), + projection_target, + did_input_contain_user_query: true, + supported_tools_override: None, + }; + + let empty_run_id = vec![snapshot("", queued_target.clone())]; + assert_eq!( + super::validate_queued_provider_run_snapshots( + Some(&active_run_id), + Some(&active_target), + &empty_run_id, + ) + .unwrap_err(), + "queued provider run ID must not be empty" + ); + + let active_run_reuse = vec![snapshot("active-run", queued_target.clone())]; + assert!(super::validate_queued_provider_run_snapshots( + Some(&active_run_id), + Some(&active_target), + &active_run_reuse, + ) + .unwrap_err() + .contains("active generation run ID")); + + let duplicate_runs = vec![ + snapshot("duplicate", queued_target.clone()), + snapshot( + "duplicate", + super::ProviderProjectionTarget { + task_id: TaskId::new("other-task".to_string()), + exchange_id: AIAgentExchangeId::new(), + }, + ), + ]; + assert!(super::validate_queued_provider_run_snapshots( + Some(&active_run_id), + Some(&active_target), + &duplicate_runs, + ) + .unwrap_err() + .contains("duplicate queued provider run ID")); + + let active_projection_reuse = vec![snapshot("queued", active_target.clone())]; + assert!(super::validate_queued_provider_run_snapshots( + Some(&active_run_id), + Some(&active_target), + &active_projection_reuse, + ) + .unwrap_err() + .contains("active generation projection target")); + + let duplicate_projections = vec![ + snapshot("first", queued_target.clone()), + snapshot("second", queued_target), + ]; + assert_eq!( + super::validate_queued_provider_run_snapshots( + Some(&active_run_id), + Some(&active_target), + &duplicate_projections, + ) + .unwrap_err(), + "duplicate queued provider projection target" + ); +} + +#[test] +fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(terminal.id(), false, false, false, ctx) + }); + let active_snapshot = provider_snapshot(conversation_id); + let queued_snapshot = super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("queued-unprepared"), + projection_target: super::ProviderProjectionTarget { + task_id: active_snapshot.root_task_id.clone(), + exchange_id: AIAgentExchangeId::new(), + }, + did_input_contain_user_query: true, + supported_tools_override: None, + }; + let active_stream_id = ResponseStreamId::new_for_test(); + let active_response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(active_stream_id.clone())); + let queued_stream_id = ResponseStreamId::new_for_test(); + let queued_response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(queued_stream_id.clone())); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_runs.insert( + conversation_id, + super::ActiveProviderRunSlot { + stream_id: active_stream_id.clone(), + response_stream: active_response_stream, + did_input_contain_user_query: true, + run_id: active_snapshot.run.id().clone(), + root_task_id: active_snapshot.root_task_id, + projection_target: active_snapshot.projection_target, + run: None, + checkpoint: None, + 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, + pending_command_completion: None, + monitor_prose_continuations: 0, + }, + ); + controller + .queued_provider_runs + .entry(conversation_id) + .or_default() + .push_back(super::QueuedProviderRun { + slot: super::ActiveProviderRunSlot { + stream_id: queued_stream_id, + response_stream: queued_response_stream, + did_input_contain_user_query: true, + run_id: queued_snapshot.run_id.clone(), + root_task_id: queued_snapshot.projection_target.task_id.clone(), + projection_target: queued_snapshot.projection_target.clone(), + run: None, + checkpoint: None, + 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, + pending_command_completion: None, + monitor_prose_continuations: 0, + }, + base_provider_config: crate::ai::provider::ProviderConfig::None, + cli_provider_config: crate::ai::provider::ProviderConfig::None, + request_params: crate::ai::agent::api::RequestParams::new_for_test(), + }); + + controller + .persist_active_provider_run(conversation_id, ctx) + .unwrap(); + }); + + let persisted = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap() + .active_provider_run_json() + .unwrap(); + let persisted: super::QueuedProviderRunsOnlySnapshot = + serde_json::from_str(persisted).unwrap(); + assert_eq!(persisted.queued_follow_ups.len(), 1); + let abandoned = persisted + .abandoned_generation + .expect("unprepared active generation identity should be durable"); + assert_eq!(abandoned.run_id, active_snapshot.run.id().clone()); + assert_eq!(abandoned.response_stream_id, active_stream_id.as_str()); + assert_eq!( + persisted.queued_follow_ups[0].run_id, + queued_snapshot.run_id + ); + }); + }); +} + +#[test] +fn queued_only_provider_restore_rejects_acp_conversation_before_starting_successor() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal.id(), false, false, false, ctx); + assert!(history + .conversation_mut(&conversation_id) + .unwrap() + .set_agent_backend_if_no_output(AgentBackend::Acp( + AcpConversationData::default(), + ))); + conversation_id + }); + let snapshot = super::QueuedProviderRunsOnlySnapshot { + version: super::QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION, + active_run_id: ProviderRunId::new("abandoned-acp-run"), + abandoned_generation: None, + queued_follow_ups: vec![super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("must-not-start"), + projection_target: super::ProviderProjectionTarget { + task_id: BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap() + .get_root_task_id() + .clone(), + exchange_id: AIAgentExchangeId::new(), + }, + did_input_contain_user_query: true, + supported_tools_override: None, + }], + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + history + .persist_active_provider_run_json( + conversation_id, + Some(serde_json::to_string(&snapshot).unwrap()), + ctx, + ) + .unwrap(); + }); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.restoring_provider_runs.insert(conversation_id); + controller.restore_active_provider_run(conversation_id, ctx); + assert!(!controller + .active_provider_runs + .contains_key(&conversation_id)); + assert!(!controller + .queued_provider_runs + .contains_key(&conversation_id)); + }); + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert_eq!(conversation.status(), &ConversationStatus::Error); + assert!(conversation.active_provider_run_json().is_none()); + }); + }); +} + +#[test] +fn queued_only_v1_without_abandoned_identity_rejects_successor() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + history.start_new_conversation(terminal.id(), false, false, false, ctx) + }); + let root_task_id = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap() + .get_root_task_id() + .clone(); + let legacy = super::QueuedProviderRunsOnlySnapshot { + version: 1, + active_run_id: ProviderRunId::new("legacy-abandoned-run"), + abandoned_generation: None, + queued_follow_ups: vec![super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("must-not-start"), + projection_target: super::ProviderProjectionTarget { + task_id: root_task_id, + exchange_id: AIAgentExchangeId::new(), + }, + did_input_contain_user_query: true, + supported_tools_override: None, + }], + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + history + .persist_active_provider_run_json( + conversation_id, + Some(serde_json::to_string(&legacy).unwrap()), + ctx, + ) + .unwrap(); + }); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.restoring_provider_runs.insert(conversation_id); + controller.restore_active_provider_run(conversation_id, ctx); + assert!(!controller + .active_provider_runs + .contains_key(&conversation_id)); + assert!(!controller + .queued_provider_runs + .contains_key(&conversation_id)); + }); + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert_eq!(conversation.status(), &ConversationStatus::Error); + assert!(conversation.active_provider_run_json().is_none()); + }); + }); +} + +#[test] +fn queued_only_restore_cancels_abandoned_exchange_before_starting_successor() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let terminal_id = terminal.id(); + let (conversation_id, abandoned_stream_id, abandoned_target, queued_target) = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal_id, false, false, false, ctx); + let task_id = history + .conversation(&conversation_id) + .unwrap() + .get_root_task_id() + .clone(); + let abandoned_stream_id = ResponseStreamId::new_for_test(); + let add_exchange = |history: &mut BlocklistAIHistoryModel, + stream_id: ResponseStreamId, + ctx: &mut warpui::ModelContext< + BlocklistAIHistoryModel, + >| { + history + .update_conversation_for_new_request_input( + RequestInput { + conversation_id, + input_messages: HashMap::from([(task_id.clone(), vec![])]), + working_directory: None, + model_id: LLMId::from("test-model"), + coding_model_id: LLMId::from("test-model"), + cli_agent_model_id: LLMId::from("test-model"), + computer_use_model_id: LLMId::from("test-model"), + shared_session_response_initiator: None, + request_start_ts: Local::now(), + supported_tools_override: None, + }, + stream_id, + terminal_id, + ctx, + ) + .unwrap(); + }; + add_exchange(history, abandoned_stream_id.clone(), ctx); + let abandoned_target = history + .conversation(&conversation_id) + .unwrap() + .provider_projection_target(&abandoned_stream_id) + .unwrap(); + let queued_stream_id = ResponseStreamId::new_for_test(); + add_exchange(history, queued_stream_id.clone(), ctx); + let queued_target = history + .conversation(&conversation_id) + .unwrap() + .provider_projection_target(&queued_stream_id) + .unwrap(); + history + .conversation_mut(&conversation_id) + .unwrap() + .cleanup_completed_response_stream(&abandoned_stream_id); + ( + conversation_id, + abandoned_stream_id, + abandoned_target, + queued_target, + ) + }); + let abandoned_exchange_id = abandoned_target.1; + let active_run_id = ProviderRunId::new("abandoned-unprepared-run"); + let snapshot = super::QueuedProviderRunsOnlySnapshot { + version: super::QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION, + active_run_id: active_run_id.clone(), + abandoned_generation: Some(super::AbandonedProviderGenerationSnapshot { + run_id: active_run_id, + projection_target: super::ProviderProjectionTarget { + task_id: abandoned_target.0, + exchange_id: abandoned_exchange_id, + }, + response_stream_id: abandoned_stream_id.as_str().to_owned(), + }), + queued_follow_ups: vec![super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("successor-run"), + projection_target: super::ProviderProjectionTarget { + task_id: queued_target.0, + exchange_id: queued_target.1, + }, + did_input_contain_user_query: true, + supported_tools_override: None, + }], + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + history + .persist_active_provider_run_json( + conversation_id, + Some(serde_json::to_string(&snapshot).unwrap()), + ctx, + ) + .unwrap(); + }); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.restoring_provider_runs.insert(conversation_id); + controller.restore_active_provider_run(conversation_id, ctx); + assert_eq!( + controller.active_provider_runs[&conversation_id].run_id, + ProviderRunId::new("successor-run") + ); + }); + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert!(!conversation.is_processing_response_stream(&abandoned_stream_id)); + assert!(conversation + .exchange_with_id(abandoned_exchange_id) + .unwrap() + .output_status + .is_cancelled()); + }); + }); +} + +#[test] +fn queued_provider_follow_up_refreshes_late_committed_history() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(terminal.id(), false, false, false, ctx) + }); + let late_message = ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::Text("late old-generation output".to_owned()), + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, _| { + history_model + .conversation_mut(&conversation_id) + .unwrap() + .append_to_bedrock_history(vec![late_message.clone()]); + }); + let mut params = crate::ai::agent::api::RequestParams::new_for_test(); + params.tasks.clear(); + params.root_task_id = Some("stale-root".to_owned()); + params.message_history = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("stale history".to_owned()), + }]; + + super::refresh_queued_provider_history( + &mut params, + BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(), + ); + + assert_eq!(params.message_history, vec![late_message]); + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert_eq!(params.tasks, conversation.compute_active_tasks()); + assert_eq!( + params.root_task_id, + Some(conversation.get_root_task_id().to_string()) + ); + }); + }); +} + +#[test] +fn restored_queued_child_follow_up_keeps_orchestration_disabled() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let terminal_id = terminal.id(); + let (child_id, projection_target) = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let parent_id = + history.start_new_conversation(terminal_id, false, false, false, ctx); + let child_id = history.start_new_child_conversation( + terminal_id, + "child".to_owned(), + parent_id, + None, + ctx, + ); + let task_id = history + .conversation(&child_id) + .unwrap() + .get_root_task_id() + .clone(); + let stream_id = ResponseStreamId::new_for_test(); + history + .update_conversation_for_new_request_input( + RequestInput { + conversation_id: child_id, + input_messages: HashMap::from([(task_id.clone(), vec![])]), + working_directory: None, + model_id: LLMId::from("test-model"), + coding_model_id: LLMId::from("test-model"), + cli_agent_model_id: LLMId::from("test-model"), + computer_use_model_id: LLMId::from("test-model"), + shared_session_response_initiator: None, + request_start_ts: Local::now(), + supported_tools_override: None, + }, + stream_id.clone(), + terminal_id, + ctx, + ) + .unwrap(); + let (task_id, exchange_id) = history + .conversation(&child_id) + .unwrap() + .provider_projection_target(&stream_id) + .unwrap(); + ( + child_id, + super::ProviderProjectionTarget { + task_id, + exchange_id, + }, + ) + }); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller + .restore_queued_provider_follow_ups( + child_id, + vec![super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("restored-child-follow-up"), + projection_target, + did_input_contain_user_query: true, + supported_tools_override: None, + }], + ctx, + ) + .unwrap(); + assert!( + !controller.queued_provider_runs[&child_id][0] + .request_params + .orchestration_enabled + ); + }); + }); + }); +} + +#[test] +fn malformed_queued_restoration_mutates_none_of_the_batch() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let terminal_id = terminal.id(); + let (conversation_id, original_stream_id, valid_target) = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal_id, false, false, false, ctx); + let task_id = history + .conversation(&conversation_id) + .unwrap() + .get_root_task_id() + .clone(); + let stream_id = ResponseStreamId::new_for_test(); + history + .update_conversation_for_new_request_input( + RequestInput { + conversation_id, + input_messages: HashMap::from([(task_id, vec![])]), + working_directory: None, + model_id: LLMId::from("test-model"), + coding_model_id: LLMId::from("test-model"), + cli_agent_model_id: LLMId::from("test-model"), + computer_use_model_id: LLMId::from("test-model"), + shared_session_response_initiator: None, + request_start_ts: Local::now(), + supported_tools_override: None, + }, + stream_id.clone(), + terminal_id, + ctx, + ) + .unwrap(); + let (task_id, exchange_id) = history + .conversation(&conversation_id) + .unwrap() + .provider_projection_target(&stream_id) + .unwrap(); + ( + conversation_id, + stream_id, + super::ProviderProjectionTarget { + task_id, + exchange_id, + }, + ) + }); + let malformed_target = super::ProviderProjectionTarget { + task_id: valid_target.task_id.clone(), + exchange_id: AIAgentExchangeId::new(), + }; + + terminal.ai_controller().update(ctx, |controller, ctx| { + assert!(controller + .restore_queued_provider_follow_ups( + conversation_id, + vec![ + super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("valid-first"), + projection_target: valid_target, + did_input_contain_user_query: true, + supported_tools_override: None, + }, + super::QueuedProviderRunSnapshot { + run_id: ProviderRunId::new("malformed-second"), + projection_target: malformed_target, + did_input_contain_user_query: true, + supported_tools_override: None, + }, + ], + ctx, + ) + .is_err()); + assert!(!controller + .queued_provider_runs + .contains_key(&conversation_id)); + }); + assert!(BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap() + .is_processing_response_stream(&original_stream_id)); + }); + }); +} + #[test] fn cancelling_provider_startup_keeps_slot_and_durable_cancellation() { App::test((), |mut app| async move { @@ -447,6 +1129,12 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() monitor_prose_continuations: 0, }, ); + controller + .in_flight_response_streams + .register_additional_stream(old_stream_id.clone(), old_response_stream.clone()); + assert!( + controller.provider_generation_is_terminalizing_for_follow_up(conversation_id) + ); controller .queued_provider_runs .entry(conversation_id) @@ -454,7 +1142,7 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() .push_back(super::QueuedProviderRun { slot: super::ActiveProviderRunSlot { stream_id: new_stream_id.clone(), - response_stream: new_response_stream, + response_stream: new_response_stream.clone(), did_input_contain_user_query: true, run_id: new_snapshot.run.id().clone(), root_task_id: new_snapshot.root_task_id, @@ -475,11 +1163,27 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() cli_provider_config: crate::ai::provider::ProviderConfig::None, request_params: crate::ai::agent::api::RequestParams::new_for_test(), }); + controller + .in_flight_response_streams + .register_additional_stream(new_stream_id.clone(), new_response_stream); assert_eq!( controller.active_provider_runs[&conversation_id].stream_id, old_stream_id ); + assert!(controller + .in_flight_response_streams + .has_stream(&old_stream_id)); + assert!(controller + .in_flight_response_streams + .has_stream(&new_stream_id)); + controller.start_next_queued_provider_run(conversation_id, ctx); + assert_eq!( + controller.active_provider_runs[&conversation_id].stream_id, + old_stream_id + ); + assert_eq!(controller.queued_provider_runs[&conversation_id].len(), 1); + controller.cleanup_active_provider_run( conversation_id, &old_stream_id, @@ -490,6 +1194,12 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() controller.active_provider_runs[&conversation_id].stream_id, new_stream_id ); + assert!(!controller + .in_flight_response_streams + .has_stream(&old_stream_id)); + assert!(controller + .in_flight_response_streams + .has_stream(&new_stream_id)); assert!(!controller .queued_provider_runs .contains_key(&conversation_id)); @@ -510,6 +1220,64 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() }); } +#[test] +fn non_follow_up_provider_cancellation_does_not_admit_an_overlapping_generation() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let terminal_surface_id = terminal.id(); + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation( + terminal_surface_id, + false, + false, + false, + ctx, + ) + }); + let snapshot = provider_snapshot(conversation_id); + let stream_id = ResponseStreamId::new_for_test(); + let response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(stream_id.clone())); + + terminal.ai_controller().update(ctx, |controller, _| { + controller + .in_flight_response_streams + .register_additional_stream(stream_id.clone(), response_stream.clone()); + controller.active_provider_runs.insert( + conversation_id, + super::ActiveProviderRunSlot { + stream_id, + response_stream, + did_input_contain_user_query: true, + run_id: snapshot.run.id().clone(), + root_task_id: snapshot.root_task_id, + projection_target: snapshot.projection_target, + run: None, + checkpoint: None, + turn_control: None, + cancellation_reason: Some(CancellationReason::ManuallyCancelled), + committed_provider_batch: None, + finished_provider_batch: None, + command_action_refs: HashMap::new(), + command_monitor: None, + pending_monitor_observation: None, + pending_command_completion: None, + monitor_prose_continuations: 0, + }, + ); + + assert!( + !controller.provider_generation_is_terminalizing_for_follow_up(conversation_id) + ); + }); + }); + }); +} + #[test] fn cancelled_provider_command_detaches_running_process_to_user() { App::test((), |mut app| async move { @@ -750,6 +1518,30 @@ fn provider_snapshot_parse_and_validation_reject_corrupt_restore_identity() { ); } +#[test] +fn restored_provider_snapshot_validates_run_before_normalization() { + let conversation_id = AIConversationId::new(); + let snapshot = provider_snapshot(conversation_id); + let mut value = serde_json::to_value(&snapshot).unwrap(); + value["run"]["state"] = serde_json::json!({ + "AwaitingTools": { + "batch": { + "work_id": {"run_id": snapshot.run.id().as_str(), "epoch": 0}, + "calls": [] + } + } + }); + let parse_error = super::ActiveProviderRunSnapshot::parse(&value.to_string()).unwrap_err(); + assert!(parse_error.contains("invalid restored provider run: pending tool batch is empty")); + let mut corrupted: super::ActiveProviderRunSnapshot = serde_json::from_value(value).unwrap(); + let before = serde_json::to_value(&corrupted.run).unwrap(); + + let error = super::normalize_restored_provider_snapshot(&mut corrupted).unwrap_err(); + + assert!(error.contains("invalid restored provider run: pending tool batch is empty")); + assert_eq!(serde_json::to_value(&corrupted.run).unwrap(), before); +} + #[test] fn restored_committed_command_requires_durable_terminal_owner() { let conversation_id = AIConversationId::new(); @@ -772,6 +1564,50 @@ fn restored_committed_command_requires_durable_terminal_owner() { ); } +#[test] +fn crash_after_model_acceptance_before_snapshot_checkpoint_terminates_without_replay() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let Some(ProviderRunStep::CallModel(dispatched_call)) = snapshot.run.next_step().unwrap() + else { + panic!("expected provider model call"); + }; + + let persisted_at_dispatch_boundary = serde_json::to_string(&snapshot).unwrap(); + snapshot + .run + .accept_model_turn( + &dispatched_call.work_id, + CompletedModelTurn { + assistant_content: vec![ContentPart::Text( + "accepted but not checkpointed".to_owned(), + )], + tool_calls: Vec::new(), + usage: Usage::default(), + stop_reason: StopReason::Completed, + advertised_tools: BTreeSet::new(), + }, + ) + .unwrap(); + assert_eq!(snapshot.run.model_turns(), 1); + + let mut restored = super::ActiveProviderRunSnapshot::parse(&persisted_at_dispatch_boundary) + .expect("dispatch-boundary snapshot should deserialize"); + super::normalize_restored_provider_snapshot(&mut restored).unwrap(); + + let ProviderRunState::Failed { failure } = restored.run.state() else { + panic!("uncertain restored model dispatch must terminate"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::Restore); + assert!(failure.message.contains("outcome is unknown")); + assert_eq!(restored.run.active_work_id(), None); + assert_eq!(dispatched_call.work_id.epoch, RunEpoch::new(0)); + assert!(matches!( + restored.run.next_step().unwrap(), + Some(ProviderRunStep::Done(ProviderRunOutcome::Failed(_))) + )); +} + #[test] fn restore_normalization_removes_interrupted_command_correlation() { let conversation_id = AIConversationId::new(); @@ -981,6 +1817,174 @@ fn restored_completed_command_rebuilds_exact_completion_evidence() { assert_eq!(completion.exit_code, 17); } +#[test] +fn completion_offered_during_restore_stays_provider_owned_and_durable() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(terminal.id(), false, false, false, ctx) + }); + let mut snapshot = provider_snapshot(conversation_id); + let (action_id, block_id, _) = + attach_snapshot_command_monitor(&mut snapshot, conversation_id); + let json = serde_json::to_string(&snapshot).unwrap(); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model + .persist_active_provider_run_json(conversation_id, Some(json), ctx) + .unwrap(); + }); + + terminal.ai_controller().update(ctx, |controller, ctx| { + let completion = super::PendingProviderCommandCompletion::new( + block_id.clone(), + Some(action_id.clone()), + "sleep 10".to_owned(), + "done before restore".to_owned(), + 0, + ); + assert!(controller.offer_provider_command_completion( + conversation_id, + completion.clone(), + ctx, + )); + assert!(controller.offer_provider_command_completion( + conversation_id, + completion, + ctx, + )); + assert!(!controller.offer_provider_command_completion( + conversation_id, + super::PendingProviderCommandCompletion::new( + block_id.clone(), + Some(action_id.clone()), + "sleep 10".to_owned(), + "conflicting duplicate".to_owned(), + 0, + ), + ctx, + )); + assert_eq!( + controller.restoring_provider_command_completions[&conversation_id].output, + "done before restore" + ); + }); + + let restored = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .and_then(|conversation| conversation.active_provider_run_json()) + .and_then(|json| super::ActiveProviderRunSnapshot::parse(json).ok()) + .unwrap(); + assert!(restored.pending_monitor_observation.is_none()); + let completion = restored.pending_command_completion.unwrap(); + assert_eq!(completion.block_id, block_id); + assert_eq!( + completion.initial_requested_command_action_id, + Some(action_id) + ); + assert_eq!(completion.output, "done before restore"); + }); + }); +} + +#[test] +fn prepared_restore_merges_a_later_durable_completion_once() { + let conversation_id = AIConversationId::new(); + let mut prepared = provider_snapshot(conversation_id); + let (action_id, block_id, cli_task_id) = + attach_snapshot_command_monitor(&mut prepared, conversation_id); + prepared.pending_monitor_observation = Some(super::PendingProviderMonitorObservation { + block_id: block_id.clone(), + cli_task_id, + }); + let mut latest = prepared.clone(); + latest.pending_command_completion = Some(super::PendingProviderCommandCompletion::new( + block_id.clone(), + Some(action_id.clone()), + "sleep 10".to_owned(), + "done during runtime preparation".to_owned(), + 0, + )); + latest.pending_monitor_observation = None; + + super::merge_completion_offered_during_restore(&mut prepared, latest); + + assert!(prepared.pending_monitor_observation.is_none()); + let completion = prepared.pending_command_completion.unwrap(); + assert_eq!(completion.block_id, block_id); + assert_eq!( + completion.initial_requested_command_action_id, + Some(action_id) + ); + assert_eq!(completion.output, "done during runtime preparation"); +} + +#[test] +fn prepared_restore_ignores_completion_from_a_different_run() { + let conversation_id = AIConversationId::new(); + let mut prepared = provider_snapshot(conversation_id); + let (_, block_id, cli_task_id) = + attach_snapshot_command_monitor(&mut prepared, conversation_id); + prepared.pending_monitor_observation = Some(super::PendingProviderMonitorObservation { + block_id: block_id.clone(), + cli_task_id, + }); + let mut other = provider_snapshot(conversation_id); + other.run = ProviderRun::new( + "different-restored-run", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Finish the task".to_owned()), + }], + crate::ai::runtime::BASE_PROVIDER_PROFILE, + ProviderRunLimits::default(), + ); + other.pending_command_completion = Some(super::PendingProviderCommandCompletion::new( + block_id, + None, + "other".to_owned(), + "stale".to_owned(), + 0, + )); + + super::merge_completion_offered_during_restore(&mut prepared, other); + + assert!(prepared.pending_monitor_observation.is_some()); + assert!(prepared.pending_command_completion.is_none()); +} + +#[test] +fn restoring_provider_ownership_requires_exact_block_and_action() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let (action_id, block_id, _) = attach_snapshot_command_monitor(&mut snapshot, conversation_id); + + assert!(super::provider_command_completion_matches( + snapshot.run.id(), + &snapshot.command_action_refs, + snapshot.command_monitor.as_ref(), + &block_id, + Some(&action_id), + )); + assert!(!super::provider_command_completion_matches( + snapshot.run.id(), + &snapshot.command_action_refs, + snapshot.command_monitor.as_ref(), + &BlockId::new(), + Some(&action_id), + )); + assert!(!super::provider_command_completion_matches( + snapshot.run.id(), + &snapshot.command_action_refs, + snapshot.command_monitor.as_ref(), + &block_id, + Some(&AIAgentActionId::from("legacy-action".to_owned())), + )); +} + #[test] fn restored_missing_command_block_becomes_interrupted_completion_evidence() { let conversation_id = AIConversationId::new(); diff --git a/app/src/ai/blocklist/inline_action/run_agents_card_view.rs b/app/src/ai/blocklist/inline_action/run_agents_card_view.rs index b41ad9a2..6872de85 100644 --- a/app/src/ai/blocklist/inline_action/run_agents_card_view.rs +++ b/app/src/ai/blocklist/inline_action/run_agents_card_view.rs @@ -296,6 +296,31 @@ fn mark_run_agents_child_removed( true } +fn run_agents_event_matches_card( + event: &RunAgentsExecutorEvent, + conversation_id: Option, + action_id: &AIAgentActionId, +) -> bool { + let (event_conversation_id, event_action_id) = match event { + RunAgentsExecutorEvent::SpawningStarted { + conversation_id, + action_id, + .. + } + | RunAgentsExecutorEvent::SpawningFinished { + conversation_id, + action_id, + } => (*conversation_id, action_id), + RunAgentsExecutorEvent::ChildConversationCreated { + action_id, + parent_conversation_id, + .. + } => (*parent_conversation_id, action_id), + }; + + Some(event_conversation_id) == conversation_id && event_action_id == action_id +} + pub struct RunAgentsCardView { action_id: AIAgentActionId, state: RunAgentsEditState, @@ -441,32 +466,34 @@ impl RunAgentsCardView { }); let action_id_for_subscription = action_id.clone(); - ctx.subscribe_to_model(&run_agents_executor, move |me, _, event, ctx| match event { - RunAgentsExecutorEvent::SpawningStarted { - action_id, - snapshot, - } if action_id == &action_id_for_subscription => { - me.spawning = Some(*snapshot); - ctx.notify(); + let conversation_id_for_subscription = block_model.conversation_id(ctx); + ctx.subscribe_to_model(&run_agents_executor, move |me, _, event, ctx| { + if !run_agents_event_matches_card( + event, + conversation_id_for_subscription, + &action_id_for_subscription, + ) { + return; } - RunAgentsExecutorEvent::SpawningFinished { action_id } - if action_id == &action_id_for_subscription => - { - me.spawning = None; - ctx.notify(); + + match event { + RunAgentsExecutorEvent::SpawningStarted { snapshot, .. } => { + me.spawning = Some(*snapshot); + ctx.notify(); + } + RunAgentsExecutorEvent::SpawningFinished { .. } => { + me.spawning = None; + ctx.notify(); + } + RunAgentsExecutorEvent::ChildConversationCreated { + agent_name, + child_conversation_id, + .. + } => { + me.link_child_conversation(agent_name, *child_conversation_id); + ctx.notify(); + } } - RunAgentsExecutorEvent::ChildConversationCreated { - action_id, - agent_name, - child_conversation_id, - .. - } if action_id == &action_id_for_subscription => { - me.link_child_conversation(agent_name, *child_conversation_id); - ctx.notify(); - } - RunAgentsExecutorEvent::SpawningStarted { .. } - | RunAgentsExecutorEvent::SpawningFinished { .. } - | RunAgentsExecutorEvent::ChildConversationCreated { .. } => {} }); let history_model = BlocklistAIHistoryModel::handle(ctx); diff --git a/app/src/ai/blocklist/inline_action/run_agents_card_view_tests.rs b/app/src/ai/blocklist/inline_action/run_agents_card_view_tests.rs index f4ab1159..85aff73c 100644 --- a/app/src/ai/blocklist/inline_action/run_agents_card_view_tests.rs +++ b/app/src/ai/blocklist/inline_action/run_agents_card_view_tests.rs @@ -10,9 +10,11 @@ use warp_util::local_or_remote_path::LocalOrRemotePath; use super::{ has_run_agents_child, link_run_agents_child, mark_run_agents_child_removed, - sync_run_agents_children, RunAgentsChildState, RunAgentsEditState, + run_agents_event_matches_card, sync_run_agents_children, RunAgentsChildState, + RunAgentsEditState, RunAgentsExecutorEvent, }; use crate::ai::agent::conversation::AIConversationId; +use crate::ai::agent::AIAgentActionId; use crate::ai::blocklist::inline_action::orchestration_controls::OrchestrationEditState; fn make_request(harness: &str, mode: RunAgentsExecutionMode) -> RunAgentsRequest { @@ -305,6 +307,31 @@ fn live_child_links_and_removal_survive_streaming_config_sync() { )); } +#[test] +fn child_created_with_duplicate_action_id_only_matches_parent_conversation() { + let card_conversation_id = AIConversationId::new(); + let other_conversation_id = AIConversationId::new(); + let child_conversation_id = AIConversationId::new(); + let duplicate_action_id = AIAgentActionId::from("duplicate-action".to_string()); + let event = RunAgentsExecutorEvent::ChildConversationCreated { + action_id: duplicate_action_id.clone(), + agent_name: "child".to_string(), + parent_conversation_id: other_conversation_id, + child_conversation_id, + }; + + assert!(!run_agents_event_matches_card( + &event, + Some(card_conversation_id), + &duplicate_action_id, + )); + assert!(run_agents_event_matches_card( + &event, + Some(other_conversation_id), + &duplicate_action_id, + )); +} + mod format_terminal_state_tests { use super::super::{format_terminal_state, StatusKind}; use super::*; diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 9652fba4..33f702a6 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -302,6 +302,7 @@ pub enum ModelFailureDisposition { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct ProviderRunRestoreNormalization { + pub interrupted_model_call: bool, pub permission_call_ids_reset: Vec, pub interrupted_call_ids: Vec, pub recovery_call_ids: Vec, @@ -350,6 +351,9 @@ pub enum ProviderRunProtocolError { InvalidDriverObservation { message: String, }, + InvalidRestoredState { + message: String, + }, EpochExhausted, Terminal, } @@ -409,6 +413,9 @@ impl fmt::Display for ProviderRunProtocolError { Self::InvalidDriverObservation { message } => { write!(f, "invalid driver observation: {message}") } + Self::InvalidRestoredState { message } => { + write!(f, "invalid restored provider run: {message}") + } Self::EpochExhausted => f.write_str("provider run epoch is exhausted"), Self::Terminal => f.write_str("provider run is already terminal"), } @@ -502,6 +509,226 @@ impl ProviderRun { ) } + /// Validates persisted state before restore normalization can mutate it or external work can + /// be reconstructed from it. + pub fn validate_restored_state(&self) -> Result<(), ProviderRunProtocolError> { + let invalid = |message: String| ProviderRunProtocolError::InvalidRestoredState { message }; + + if self.id.as_str().is_empty() { + return Err(invalid("run ID must not be empty".to_string())); + } + if self.profile.as_str().is_empty() { + return Err(invalid("request profile must not be empty".to_string())); + } + if self.limits.max_model_turns == 0 { + return Err(invalid("model-turn limit must be at least one".to_string())); + } + if self.model_turns > self.limits.max_model_turns { + return Err(invalid(format!( + "model-turn counter {} exceeds limit {}", + self.model_turns, self.limits.max_model_turns + ))); + } + if self.epoch.get() < u64::from(self.model_turns) { + return Err(invalid(format!( + "epoch {} is behind model-turn counter {}", + self.epoch.get(), + self.model_turns + ))); + } + let retry_slots = u64::from(self.model_turns) + .saturating_add(1) + .saturating_mul(u64::from(self.limits.max_model_retries_per_turn)); + if u64::from(self.model_retries) > retry_slots { + return Err(invalid(format!( + "model-retry counter {} exceeds maximum possible {}", + self.model_retries, retry_slots + ))); + } + + if let Some(work_id) = self.active_work_id() { + self.validate_restored_work_id(work_id)?; + } + + match &self.state { + ProviderRunState::ReadyToCallModel => {} + ProviderRunState::AwaitingModel { call } => { + if self.model_turns >= self.limits.max_model_turns { + return Err(invalid( + "awaiting a model call after reaching the model-turn limit".to_string(), + )); + } + if call.retry_attempt > self.limits.max_model_retries_per_turn { + return Err(invalid(format!( + "pending retry attempt {} exceeds per-turn limit {}", + call.retry_attempt, self.limits.max_model_retries_per_turn + ))); + } + if call.retry_attempt > self.model_retries { + return Err(invalid(format!( + "pending retry attempt {} exceeds total retry counter {}", + call.retry_attempt, self.model_retries + ))); + } + if (call.retry_attempt == 0) != call.last_error.is_none() { + return Err(invalid( + "pending retry error does not match its retry attempt".to_string(), + )); + } + if call + .last_error + .as_ref() + .is_some_and(|error| !error.recoverable) + { + return Err(invalid( + "pending retry retains a non-recoverable model error".to_string(), + )); + } + } + ProviderRunState::ResolvingModel { turn } => { + self.validate_post_model_phase()?; + validate_model_turn(turn).map_err(|error| invalid(error.to_string()))?; + if self.transcript.last() != Some(&assistant_message(turn)) { + return Err(invalid( + "resolving model turn does not own the latest transcript message" + .to_string(), + )); + } + if !usage_contains(&self.usage, &turn.usage) { + return Err(invalid( + "aggregate usage does not include the resolving model turn".to_string(), + )); + } + } + ProviderRunState::AwaitingTools { batch } => { + self.validate_post_model_phase()?; + if batch.calls.is_empty() { + return Err(invalid("pending tool batch is empty".to_string())); + } + let mut call_ids = HashSet::new(); + for pending in &batch.calls { + if pending.call.id.is_empty() { + return Err(invalid("pending tool call ID is empty".to_string())); + } + if !call_ids.insert(pending.call.id.as_str()) { + return Err(invalid(format!( + "duplicate pending tool call ID '{}'", + pending.call.id + ))); + } + validate_pending_tool_state(pending).map_err(invalid)?; + } + let transcript_calls = self + .transcript + .last() + .map(tool_calls_from_message) + .unwrap_or_default(); + if transcript_calls + != batch + .calls + .iter() + .map(|pending| pending.call.clone()) + .collect::>() + { + return Err(invalid( + "pending tool batch does not match the latest assistant message" + .to_string(), + )); + } + } + ProviderRunState::AwaitingDriver { .. } => { + self.validate_post_model_phase()?; + let Some(message) = self.transcript.last() else { + return Err(invalid( + "driver wait is missing its assistant transcript message".to_string(), + )); + }; + if message.role != MessageRole::Assistant { + return Err(invalid( + "driver wait does not follow an assistant transcript message".to_string(), + )); + } + if !tool_calls_from_message(message).is_empty() { + return Err(invalid( + "driver wait follows an uncommitted assistant tool call".to_string(), + )); + } + } + ProviderRunState::Done { .. } => { + if self.model_turns == 0 { + return Err(invalid( + "completed run has no completed model turn".to_string(), + )); + } + } + ProviderRunState::Failed { failure } => match failure.kind { + ProviderRunFailureKind::ModelCall + if !failure + .source + .as_ref() + .is_some_and(|source| !source.recoverable) => + { + return Err(invalid( + "model-call failure lacks a non-recoverable source".to_string(), + )); + } + ProviderRunFailureKind::RetryLimitExceeded + if !failure + .source + .as_ref() + .is_some_and(|source| source.recoverable) => + { + return Err(invalid( + "retry-limit failure lacks a recoverable source".to_string(), + )); + } + ProviderRunFailureKind::TurnLimitExceeded + if self.model_turns < self.limits.max_model_turns => + { + return Err(invalid( + "turn-limit failure occurred before reaching the limit".to_string(), + )); + } + ProviderRunFailureKind::ModelCall + | ProviderRunFailureKind::RetryLimitExceeded + | ProviderRunFailureKind::TurnLimitExceeded + | ProviderRunFailureKind::Protocol + | ProviderRunFailureKind::Projection + | ProviderRunFailureKind::Restore + | ProviderRunFailureKind::ExternalWork => {} + }, + ProviderRunState::Cancelled { .. } => {} + } + Ok(()) + } + + fn validate_restored_work_id( + &self, + work_id: &ExternalWorkId, + ) -> Result<(), ProviderRunProtocolError> { + validate_work_id(&self.current_work_id(), work_id).map_err(|_| { + ProviderRunProtocolError::InvalidRestoredState { + message: format!( + "active work identity {}:{} does not match run {}:{}", + work_id.run_id.as_str(), + work_id.epoch.get(), + self.id.as_str(), + self.epoch.get() + ), + } + }) + } + + fn validate_post_model_phase(&self) -> Result<(), ProviderRunProtocolError> { + if self.model_turns == 0 { + Err(ProviderRunProtocolError::InvalidRestoredState { + message: format!("{:?} phase has no completed model turn", self.state.phase()), + }) + } else { + Ok(()) + } + } + pub fn normalize_after_restore( &mut self, ) -> Result { @@ -512,6 +739,21 @@ impl ProviderRun { &mut self, recoverable_call_ids: &HashSet, ) -> Result { + if matches!(self.state, ProviderRunState::AwaitingModel { .. }) { + self.state = ProviderRunState::Failed { + failure: ProviderRunFailure { + kind: ProviderRunFailureKind::Restore, + message: "The model call was interrupted by application restart after dispatch may have begun. Its outcome is unknown, so it was not replayed to avoid duplicate billing or output." + .to_string(), + source: None, + }, + }; + return Ok(ProviderRunRestoreNormalization { + interrupted_model_call: true, + ..ProviderRunRestoreNormalization::default() + }); + } + let ProviderRunState::AwaitingTools { batch } = &mut self.state else { return Ok(ProviderRunRestoreNormalization::default()); }; @@ -1220,6 +1462,84 @@ fn validate_work_id( } } +fn validate_pending_tool_state(pending: &PendingToolCall) -> Result<(), String> { + match &pending.state { + PendingToolCallState::PermissionPending { request } => { + if request.id.is_empty() { + return Err(format!( + "permission request for '{}' has an empty request ID", + pending.call.id + )); + } + if request.call_id != pending.call.id { + return Err(format!( + "permission request for '{}' belongs to call '{}'", + pending.call.id, request.call_id + )); + } + } + PendingToolCallState::Approved { + request_id, + decision, + } => { + if request_id.is_empty() { + return Err(format!( + "approved tool call '{}' has an empty request ID", + pending.call.id + )); + } + if matches!(decision, PermissionDecision::Denied { .. }) { + return Err(format!( + "approved tool call '{}' contains a denied decision", + pending.call.id + )); + } + } + PendingToolCallState::Resolved { result } if result.call_id != pending.call.id => { + return Err(format!( + "resolved result for '{}' belongs to call '{}'", + pending.call.id, result.call_id + )); + } + PendingToolCallState::Proposed + | PendingToolCallState::Executing + | PendingToolCallState::RecoveryPending + | PendingToolCallState::Resolved { .. } => {} + } + Ok(()) +} + +fn tool_calls_from_message(message: &ConversationMessage) -> Vec { + let MessageContent::MultiPart(parts) = &message.content else { + return Vec::new(); + }; + parts + .iter() + .filter_map(|part| match part { + ContentPart::ToolUse { + tool_use_id, + name, + input, + } => Some(ToolCall { + id: tool_use_id.clone(), + name: name.clone(), + arguments: input.clone(), + }), + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::ToolResult { .. } + | ContentPart::Image { .. } => None, + }) + .collect() +} + +fn usage_contains(total: &Usage, part: &Usage) -> bool { + total.input_tokens >= part.input_tokens + && total.output_tokens >= part.output_tokens + && total.cached_input_tokens >= part.cached_input_tokens + && total.cache_creation_input_tokens >= part.cache_creation_input_tokens +} + fn validate_model_turn(turn: &CompletedModelTurn) -> Result<(), ProviderRunProtocolError> { for part in &turn.assistant_content { match part { diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index 6d991f83..3af81e11 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -88,6 +88,36 @@ fn assert_serialization_round_trip(run: &ProviderRun) { assert_eq!(&restored, run); } +fn mutate_run_json(run: &ProviderRun, mutate: impl FnOnce(&mut serde_json::Value)) -> ProviderRun { + let mut value = serde_json::to_value(run).unwrap(); + mutate(&mut value); + serde_json::from_value(value).unwrap() +} + +fn restored_state_error(run: &ProviderRun) -> String { + let ProviderRunProtocolError::InvalidRestoredState { message } = + run.validate_restored_state().unwrap_err() + else { + panic!("expected restored-state validation error"); + }; + message +} + +fn awaiting_tool_run() -> ProviderRun { + let mut run = run(); + accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("first", "read_files"), + tool_call("second", "grep"), + ], + &["read_files", "grep"], + ), + ); + run +} + #[test] fn next_step_reemits_identical_pending_model_work() { let mut run = run(); @@ -740,7 +770,228 @@ fn every_nonterminal_phase_round_trips_through_json() { } #[test] -fn restore_normalization_preserves_safe_nonterminal_states() { +fn restored_state_validation_accepts_valid_snapshots_in_every_phase() { + let ready = run(); + let mut awaiting_model = ready.clone(); + let call = next_model_call(&mut awaiting_model); + let mut resolving = awaiting_model.clone(); + resolving + .accept_model_turn(&call.work_id, text_turn("done")) + .unwrap(); + let mut awaiting_driver = resolving.clone(); + assert_eq!(awaiting_driver.next_step().unwrap(), None); + let awaiting_tools = awaiting_tool_run(); + let mut done = awaiting_driver.clone(); + let work_id = awaiting_driver.active_work_id().unwrap().clone(); + done.complete(&work_id).unwrap(); + let mut failed = ready.clone(); + failed + .fail(ProviderRunFailureKind::ExternalWork, "failed") + .unwrap(); + let mut cancelled = ready; + cancelled.cancel("cancelled").unwrap(); + + for candidate in [ + awaiting_model, + resolving, + awaiting_tools, + awaiting_driver, + done, + failed, + cancelled, + ] { + candidate.validate_restored_state().unwrap(); + } +} + +#[test] +fn restored_state_validation_rejects_active_work_run_and_epoch_mismatches() { + let mut awaiting_model = run(); + next_model_call(&mut awaiting_model); + let wrong_run = mutate_run_json(&awaiting_model, |value| { + value["state"]["AwaitingModel"]["call"]["work_id"]["run_id"] = json!("other-run"); + }); + assert!(restored_state_error(&wrong_run).contains("active work identity other-run:0")); + + let wrong_epoch = mutate_run_json(&awaiting_model, |value| { + value["state"]["AwaitingModel"]["call"]["work_id"]["epoch"] = json!(9); + }); + assert!(restored_state_error(&wrong_epoch).contains("active work identity run-1:9")); +} + +#[test] +fn restored_state_validation_rejects_duplicate_tool_call_ids() { + let corrupted = mutate_run_json(&awaiting_tool_run(), |value| { + value["state"]["AwaitingTools"]["batch"]["calls"][1]["call"]["id"] = json!("first"); + }); + assert_eq!( + restored_state_error(&corrupted), + "duplicate pending tool call ID 'first'" + ); +} + +#[test] +fn restored_state_validation_rejects_result_and_permission_call_ownership() { + let mut resolved = awaiting_tool_run(); + let work_id = resolved.active_work_id().unwrap().clone(); + resolved + .complete_tool(&work_id, successful_result("first", "done")) + .unwrap(); + let wrong_result = mutate_run_json(&resolved, |value| { + value["state"]["AwaitingTools"]["batch"]["calls"][0]["state"]["Resolved"]["result"]["call_id"] = + json!("second"); + }); + assert_eq!( + restored_state_error(&wrong_result), + "resolved result for 'first' belongs to call 'second'" + ); + + let mut permission = awaiting_tool_run(); + let work_id = permission.active_work_id().unwrap().clone(); + permission + .request_tool_permission( + &work_id, + PermissionRequest { + id: "request-1".to_string(), + call_id: "first".to_string(), + kind: PermissionKind::Read, + reason: None, + }, + ) + .unwrap(); + let wrong_permission = mutate_run_json(&permission, |value| { + value["state"]["AwaitingTools"]["batch"]["calls"][0]["state"]["PermissionPending"]["request"] + ["call_id"] = json!("second"); + }); + assert_eq!( + restored_state_error(&wrong_permission), + "permission request for 'first' belongs to call 'second'" + ); +} + +#[test] +fn restored_state_validation_rejects_phase_specific_corruption() { + let empty_batch = mutate_run_json(&awaiting_tool_run(), |value| { + value["state"]["AwaitingTools"]["batch"]["calls"] = json!([]); + }); + assert_eq!( + restored_state_error(&empty_batch), + "pending tool batch is empty" + ); + + let mismatched_batch = mutate_run_json(&awaiting_tool_run(), |value| { + value["state"]["AwaitingTools"]["batch"]["calls"][0]["call"]["name"] = + json!("different_tool"); + }); + assert!(restored_state_error(&mismatched_batch).contains("latest assistant message")); + + let resolving_without_turn = mutate_run_json(&awaiting_tool_run(), |value| { + value["model_turns"] = json!(0); + }); + assert!(restored_state_error(&resolving_without_turn).contains("has no completed model turn")); + + let mut resolving = run(); + let call = next_model_call(&mut resolving); + resolving + .accept_model_turn(&call.work_id, text_turn("done")) + .unwrap(); + let invalid_turn = mutate_run_json(&resolving, |value| { + value["state"]["ResolvingModel"]["turn"]["assistant_content"] = json!([{ + "ToolUse": { + "tool_use_id": "injected", + "name": "read_files", + "input": {} + } + }]); + }); + assert!(restored_state_error(&invalid_turn).contains("assistant_content")); + + let mut awaiting_driver = resolving; + assert_eq!(awaiting_driver.next_step().unwrap(), None); + let wrong_driver_owner = mutate_run_json(&awaiting_driver, |value| { + let last = value["transcript"] + .as_array_mut() + .unwrap() + .last_mut() + .unwrap(); + last["role"] = json!("User"); + }); + assert!(restored_state_error(&wrong_driver_owner).contains("does not follow an assistant")); +} + +#[test] +fn restored_state_validation_rejects_retry_counter_and_terminal_corruption() { + let zero_limit = mutate_run_json(&run(), |value| { + value["limits"]["max_model_turns"] = json!(0); + }); + assert_eq!( + restored_state_error(&zero_limit), + "model-turn limit must be at least one" + ); + + let excessive_retries = mutate_run_json(&run(), |value| { + value["model_retries"] = json!(3); + }); + assert!(restored_state_error(&excessive_retries).contains("model-retry counter")); + + let mut awaiting_model = run(); + next_model_call(&mut awaiting_model); + let inconsistent_retry = mutate_run_json(&awaiting_model, |value| { + value["state"]["AwaitingModel"]["call"]["retry_attempt"] = json!(1); + }); + assert!(restored_state_error(&inconsistent_retry).contains("total retry counter")); + + let completed_without_turn = mutate_run_json(&run(), |value| { + value["state"] = json!({"Done": {"completion": {"stop_reason": "Completed"}}}); + }); + assert_eq!( + restored_state_error(&completed_without_turn), + "completed run has no completed model turn" + ); + + let early_turn_limit = mutate_run_json(&run(), |value| { + value["state"] = json!({ + "Failed": {"failure": { + "kind": "TurnLimitExceeded", + "message": "bad", + "source": null + }} + }); + }); + assert_eq!( + restored_state_error(&early_turn_limit), + "turn-limit failure occurred before reaching the limit" + ); + + let missing_retry_source = mutate_run_json(&run(), |value| { + value["state"] = json!({ + "Failed": {"failure": { + "kind": "RetryLimitExceeded", + "message": "bad", + "source": null + }} + }); + }); + assert_eq!( + restored_state_error(&missing_retry_source), + "retry-limit failure lacks a recoverable source" + ); +} + +#[test] +fn restored_state_validation_rejects_empty_run_identity_and_profile() { + let empty_run = mutate_run_json(&run(), |value| value["id"] = json!("")); + assert_eq!(restored_state_error(&empty_run), "run ID must not be empty"); + + let empty_profile = mutate_run_json(&run(), |value| value["profile"] = json!("")); + assert_eq!( + restored_state_error(&empty_profile), + "request profile must not be empty" + ); +} + +#[test] +fn restore_normalization_preserves_model_work_outside_the_uncertain_dispatch_boundary() { let ready = run(); let mut awaiting_model = ready.clone(); let call = next_model_call(&mut awaiting_model); @@ -752,7 +1003,7 @@ fn restore_normalization_preserves_safe_nonterminal_states() { let mut awaiting_driver = resolving.clone(); assert_eq!(awaiting_driver.next_step().unwrap(), None); - for mut candidate in [ready, awaiting_model, resolving, awaiting_driver] { + for mut candidate in [ready, resolving, awaiting_driver] { let before = candidate.clone(); assert_eq!( candidate.normalize_after_restore().unwrap(), @@ -762,6 +1013,58 @@ fn restore_normalization_preserves_safe_nonterminal_states() { } } +#[test] +fn crash_after_model_acceptance_before_checkpoint_does_not_replay_the_persisted_call() { + let mut live_run = run(); + let call = next_model_call(&mut live_run); + let serialized_at_dispatch_boundary = serde_json::to_string(&live_run).unwrap(); + + // Simulate remote acceptance followed by a crash before the accepted turn is checkpointed. + live_run + .accept_model_turn(&call.work_id, text_turn("accepted but not checkpointed")) + .unwrap(); + assert_eq!(live_run.model_turns(), 1); + + let mut restored: ProviderRun = serde_json::from_str(&serialized_at_dispatch_boundary).unwrap(); + + let normalization = restored.normalize_after_restore().unwrap(); + + assert!(normalization.interrupted_model_call); + let ProviderRunState::Failed { failure } = restored.state() else { + panic!("uncertain model work must become terminal on restore"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::Restore); + assert!(failure.message.contains("outcome is unknown")); + assert!(failure.message.contains("not replayed")); + assert_eq!(restored.active_work_id(), None); + assert!(matches!( + restored.next_step().unwrap(), + Some(ProviderRunStep::Done(ProviderRunOutcome::Failed(_))) + )); + + // The persisted dispatch identity remains useful for diagnostics but can never be called again. + assert_eq!(call.work_id.epoch, RunEpoch::new(0)); +} + +#[test] +fn restore_after_model_acceptance_keeps_the_committed_turn_without_replaying() { + let mut run = run(); + let call = next_model_call(&mut run); + run.accept_model_turn(&call.work_id, text_turn("accepted output")) + .unwrap(); + let serialized_after_acceptance = serde_json::to_string(&run).unwrap(); + let mut restored: ProviderRun = serde_json::from_str(&serialized_after_acceptance).unwrap(); + + assert_eq!( + restored.normalize_after_restore().unwrap(), + ProviderRunRestoreNormalization::default() + ); + assert_eq!(restored.state().phase(), ProviderRunPhase::ResolvingModel); + assert_eq!(restored.model_turns(), 1); + assert_eq!(restored.next_step().unwrap(), None); + assert_eq!(restored.state().phase(), ProviderRunPhase::AwaitingDriver); +} + #[test] fn restore_normalization_reproposes_permissions_and_interrupts_unsafe_tools() { let mut run = run();