diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index cd0f78c6..938def3e 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -2008,21 +2008,25 @@ impl AIConversation { .sum() } - pub fn contains_action(&self, action_id: &AIAgentActionId) -> bool { - self.task_store.tasks().any(|task| { - task.exchanges() - .any(|exchange| { - let Some(output) = exchange.output_status.output() - else { - return false; - }; - output.get().messages.iter().any(|step| { - matches!(step, AIAgentOutputMessage{ message: AIAgentOutputMessageType::Action(AIAgentAction { id, .. }), .. } if id == action_id) + pub fn action(&self, action_id: &AIAgentActionId) -> Option { + self.task_store.tasks().find_map(|task| { + task.exchanges().find_map(|exchange| { + let output = exchange.output_status.output()?; + output.get().messages.iter().find_map(|step| match step { + AIAgentOutputMessage { + message: AIAgentOutputMessageType::Action(action), + .. + } if &action.id == action_id => Some(action.clone()), + AIAgentOutputMessage { .. } => None, }) }) }) } + pub fn contains_action(&self, action_id: &AIAgentActionId) -> bool { + self.action(action_id).is_some() + } + /// Returns the exchange ID that contains the given action ID, if any. pub fn exchange_id_for_action(&self, action_id: &AIAgentActionId) -> Option { for task in self.task_store.tasks() { diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 65a80a72..b4f1d7fd 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -1546,6 +1546,7 @@ impl BlocklistAIActionModel { pub(super) fn queue_provider_actions( &mut self, actions: Vec, + recovery_action_ids: HashSet, conversation_id: AIConversationId, batch: &PendingToolBatch, ctx: &mut ModelContext, @@ -1562,6 +1563,9 @@ impl BlocklistAIActionModel { } } self.provider_tool_executions.extend(refs); + self.executor.update(ctx, |executor, ctx| { + executor.mark_restored_actions(&recovery_action_ids, ctx); + }); self.queue_actions(actions, conversation_id, ctx); Ok(()) } diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index e4deffe1..76b0c8e1 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -24,6 +24,7 @@ pub(super) mod use_computer; pub(super) mod wait_for_events; use std::any::Any; +use std::collections::HashSet; use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; @@ -295,6 +296,7 @@ pub struct BlocklistAIActionExecutor { /// We track them per action rather than as a single slot so multiple actions from the same /// parallel phase can complete independently. async_executing_actions: std::collections::HashMap, + restored_action_ids: HashSet, /// Reference to the terminal model for checking session sharing state. terminal_model: Arc>, @@ -382,6 +384,7 @@ impl BlocklistAIActionExecutor { use_computer_executor, request_computer_use_executor, async_executing_actions: Default::default(), + restored_action_ids: Default::default(), terminal_model, read_skill_executor, fetch_conversation_executor, @@ -399,6 +402,17 @@ impl BlocklistAIActionExecutor { .map(|running| &running.action) } + pub fn mark_restored_actions( + &mut self, + action_ids: &HashSet, + ctx: &mut ModelContext, + ) { + self.restored_action_ids.extend(action_ids.iter().cloned()); + self.run_agents_executor.update(ctx, |executor, _| { + executor.mark_recovery_actions(action_ids); + }); + } + pub(super) fn has_running_ask_user_question(&self, conversation_id: AIConversationId) -> bool { self.async_executing_actions.values().any(|running| { running.conversation_id == conversation_id @@ -710,6 +724,7 @@ impl BlocklistAIActionExecutor { action.id, std::mem::discriminant(&action.action) ); + let is_restored = self.restored_action_ids.remove(&action.id); let action_clone = action.clone(); let execution = match &action.action { AIAgentActionType::RequestCommandOutput { .. } @@ -904,10 +919,12 @@ impl BlocklistAIActionExecutor { conversation_id, }, ); - ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { - action_id: action_id.clone(), - conversation_id, - }); + if !is_restored { + ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { + action_id: action_id.clone(), + conversation_id, + }); + } log::info!("[tool-debug] try_to_execute_action: spawning ASYNC execution for action_id={:?}", action_id); ctx.spawn(execute_future, move |me, result, ctx| { let Some(running) = me.async_executing_actions.remove(&action_id) else { @@ -933,10 +950,12 @@ impl BlocklistAIActionExecutor { TryExecuteResult::ExecutedAsync } AnyActionExecution::Sync(action_result) => { - ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { - action_id: action_id.clone(), - conversation_id, - }); + if !is_restored { + ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { + action_id: action_id.clone(), + conversation_id, + }); + } ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { result: Arc::new(AIAgentActionResult { id: action_id, @@ -1035,7 +1054,9 @@ impl BlocklistAIActionExecutor { } fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext) -> bool { - if cfg!(feature = "bedrock_smoke_test") { + if self.restored_action_ids.contains(&input.action.id) + || cfg!(feature = "bedrock_smoke_test") + { return true; } match input.action.action { 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 d31f09aa..4d8a8054 100644 --- a/app/src/ai/blocklist/action_model/execute/run_agents.rs +++ b/app/src/ai/blocklist/action_model/execute/run_agents.rs @@ -69,6 +69,7 @@ struct ExistingLaunchedAgent { pub struct RunAgentsExecutor { pending: HashMap, + recovery_action_ids: HashSet, launched_agents: HashMap>, start_agent_executor: ModelHandle, terminal_view_id: EntityId, @@ -119,6 +120,7 @@ impl RunAgentsExecutor { }); Self { pending: HashMap::new(), + recovery_action_ids: HashSet::new(), launched_agents: HashMap::new(), start_agent_executor, terminal_view_id, @@ -129,21 +131,25 @@ impl RunAgentsExecutor { self.pending.contains_key(action_id) } + pub fn mark_recovery_actions(&mut self, action_ids: &HashSet) { + self.recovery_action_ids.extend(action_ids.iter().cloned()); + } + pub(crate) fn terminal_view_id(&self) -> EntityId { self.terminal_view_id } - /// Cancels a pending run so publication completion cannot fan out children. + /// Cancels the parent tool wait without cancelling independently-running children. pub(super) fn cancel_execution( &mut self, action_id: &AIAgentActionId, ctx: &mut ModelContext, ) { - if matches!( - self.pending.get(action_id), - Some(PendingRunAgents::Publishing) - ) { - self.pending.remove(action_id); + self.recovery_action_ids.remove(action_id); + self.start_agent_executor.update(ctx, |executor, _| { + executor.cancel_dispatches_for_action(action_id); + }); + if self.pending.remove(action_id).is_some() { ctx.emit(RunAgentsExecutorEvent::SpawningFinished { action_id: action_id.clone(), }); @@ -276,6 +282,7 @@ impl RunAgentsExecutor { action_id_for_wait.clone(), request, parent_conversation_id, + HashMap::new(), sender, ctx, ) @@ -285,11 +292,48 @@ impl RunAgentsExecutor { receiver } + fn dispatch_recovered_run_agents( + &mut self, + action_id: AIAgentActionId, + request: RunAgentsRequest, + parent_conversation_id: AIConversationId, + recovery_children: HashMap, + ctx: &mut ModelContext, + ) -> async_channel::Receiver { + let (sender, receiver) = async_channel::bounded(1); + if self.pending.contains_key(&action_id) { + let _ = sender.try_send(RunAgentsResult::Cancelled); + return receiver; + } + if let Err(error) = validate_request(&request) { + let _ = sender.try_send(RunAgentsResult::Failure { error }); + return receiver; + } + + let snapshot = RunAgentsSpawningSnapshot { + agent_count: request.agent_run_configs.len(), + }; + ctx.emit(RunAgentsExecutorEvent::SpawningStarted { + action_id: action_id.clone(), + snapshot, + }); + self.dispatch_children_for_prepared_request( + action_id, + request, + parent_conversation_id, + recovery_children, + sender, + ctx, + ); + receiver + } + fn dispatch_children_for_prepared_request( &mut self, action_id: AIAgentActionId, request: RunAgentsRequest, parent_conversation_id: AIConversationId, + mut recovery_children: HashMap, sender: async_channel::Sender, ctx: &mut ModelContext, ) { @@ -329,6 +373,23 @@ impl RunAgentsExecutor { let mut slots: Vec = Vec::with_capacity(agent_run_configs.len()); for cfg in &agent_run_configs { + let normalized_name = normalize_agent_name(&cfg.name) + .expect("validated RunAgents requests have non-empty agent names"); + if let Some(child_conversation_id) = recovery_children.remove(&normalized_name) { + let dispatch = self.start_agent_executor.update(ctx, |executor, exec_ctx| { + executor.reattach( + action_id.clone(), + cfg.name.clone(), + parent_conversation_id, + child_conversation_id, + parent_run_id.clone(), + exec_ctx, + ) + }); + slots.push(ChildSlot::Pending(dispatch)); + continue; + } + let prompt = compose_run_agents_child_prompt(&base_prompt, &cfg.prompt); let mode = match run_agents_to_start_agent_mode( &run_execution_mode, @@ -443,6 +504,9 @@ impl RunAgentsExecutor { outcomes }, move |me, outcomes, ctx| { + if !me.is_pending(&action_id_for_aggr) { + return; + } let agents: Vec = agent_run_configs_for_result .iter() .zip(outcomes) @@ -520,32 +584,56 @@ impl RunAgentsExecutor { let mut request = request.clone(); let action_id = id.clone(); let parent_conversation_id = input.conversation_id; - if let Some(reason) = prepare_request_for_execution( - &mut request, - parent_conversation_id, - self.terminal_view_id, - &self.launched_agents, - ctx, - ) { - #[cfg(not(target_family = "wasm"))] - log_run_agents_event( - ctx, - RemoteLogLevel::Warn, - "RunAgents execution denied", - serde_json::json!({ - "event": "run_agents_execution_denied", - "action_id": action_id.to_string(), - "parent_conversation_id": parent_conversation_id.to_string(), - "reason": remote_logging::sanitize_error(&reason), - }), - ); - return ActionExecution::Sync(AIAgentActionResultType::RunAgents( - RunAgentsResult::Denied { reason }, - )); - } + let is_recovery = self.recovery_action_ids.remove(&action_id); - let receiver = - self.dispatch_prepared_run_agents(action_id, request, parent_conversation_id, ctx); + let recovery_children = if is_recovery { + prepare_recovery_request_for_execution(&mut request, parent_conversation_id, ctx); + match recovery_children_by_name(parent_conversation_id, ctx) { + Ok(children) => children, + Err(error) => { + return ActionExecution::Sync(AIAgentActionResultType::RunAgents( + RunAgentsResult::Failure { error }, + )); + } + } + } else { + if let Some(reason) = prepare_request_for_execution( + &mut request, + parent_conversation_id, + self.terminal_view_id, + &self.launched_agents, + ctx, + ) { + #[cfg(not(target_family = "wasm"))] + log_run_agents_event( + ctx, + RemoteLogLevel::Warn, + "RunAgents execution denied", + serde_json::json!({ + "event": "run_agents_execution_denied", + "action_id": action_id.to_string(), + "parent_conversation_id": parent_conversation_id.to_string(), + "reason": remote_logging::sanitize_error(&reason), + }), + ); + return ActionExecution::Sync(AIAgentActionResultType::RunAgents( + RunAgentsResult::Denied { reason }, + )); + } + HashMap::new() + }; + + let receiver = if is_recovery { + self.dispatch_recovered_run_agents( + action_id, + request, + parent_conversation_id, + recovery_children, + ctx, + ) + } else { + self.dispatch_prepared_run_agents(action_id, request, parent_conversation_id, ctx) + }; ActionExecution::new_async( async move { receiver.recv().await }, @@ -771,6 +859,42 @@ fn prepare_request_for_execution( None } +fn prepare_recovery_request_for_execution( + request: &mut RunAgentsRequest, + parent_conversation_id: AIConversationId, + ctx: &ModelContext, +) { + normalize_request_for_local_execution(request); + resolve_request_from_approved_config(request, parent_conversation_id, ctx); + populate_default_auth_secret_for_execution(request, ctx); +} + +fn recovery_children_by_name( + parent_conversation_id: AIConversationId, + ctx: &ModelContext, +) -> Result, String> { + let mut children_by_name = HashMap::new(); + for conversation in + BlocklistAIHistoryModel::as_ref(ctx).child_conversations_of(parent_conversation_id) + { + let Some(name) = conversation.agent_name() else { + continue; + }; + let Some(normalized_name) = normalize_agent_name(name) else { + continue; + }; + if children_by_name + .insert(normalized_name.clone(), conversation.id()) + .is_some() + { + return Err(format!( + "Cannot recover child agent '{name}': multiple persisted child conversations have the same name." + )); + } + } + Ok(children_by_name) +} + fn duplicate_launched_agents_reason( request: &RunAgentsRequest, parent_conversation_id: AIConversationId, 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 b2e541af..39a32f1b 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 @@ -358,6 +358,166 @@ fn validate_request_rejects_remote_dispatch() { ); } +#[test] +fn recovered_run_agents_reattaches_existing_child_and_launches_only_missing_child() { + App::test((), |mut app| async move { + let state = initialize_run_agents_test(&mut app, ExecutionMode::Sdk); + let terminal_view_id = EntityId::new(); + let history = BlocklistAIHistoryModel::handle(&app); + let existing_child_id = history.update(&mut app, |history, ctx| { + history.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + state.conversation_id, + None, + ctx, + ) + }); + let captured = subscribe_to_start_agent_requests(&mut app, &state.start_agent_executor); + let mut action = remote_run_agents_action("oz"); + let AIAgentActionType::RunAgents(request) = &mut action.action else { + panic!("expected run_agents action"); + }; + request.agent_run_configs.push(RunAgentsAgentRunConfig { + name: "missing-child".to_string(), + prompt: "Do separate work".to_string(), + title: String::new(), + }); + state.executor.update(&mut app, |executor, _| { + executor.mark_recovery_actions(&HashSet::from([action.id.clone()])); + }); + + let execution = state.executor.update(&mut app, |executor, ctx| { + executor + .execute( + ExecuteActionInput { + action: &action, + conversation_id: state.conversation_id, + }, + ctx, + ) + .into() + }); + let AnyActionExecution::Async { + execute_future, + on_complete, + } = execution + else { + panic!("expected async recovery execution"); + }; + let missing_request = captured.read(&app, |captured, _| { + assert_eq!(captured.0.len(), 1); + assert_eq!(captured.0[0].name, "missing-child"); + captured.0[0].clone() + }); + + history.update(&mut app, |history, ctx| { + history.update_conversation_status( + terminal_view_id, + existing_child_id, + crate::ai::agent::conversation::ConversationStatus::Success, + ctx, + ); + }); + let missing_child_id = history.update(&mut app, |history, ctx| { + history.start_new_child_conversation( + terminal_view_id, + "missing-child".to_string(), + state.conversation_id, + None, + ctx, + ) + }); + history.update(&mut app, |history, ctx| { + history.record_new_conversation_request_complete( + missing_request.id, + missing_child_id, + ctx, + ); + history.update_conversation_status( + terminal_view_id, + missing_child_id, + crate::ai::agent::conversation::ConversationStatus::Success, + ctx, + ); + }); + + let async_result = execute_future.await; + let result = app.update(|ctx| on_complete(async_result, ctx)); + let AIAgentActionResultType::RunAgents(RunAgentsResult::Launched { agents, .. }) = result + else { + panic!("expected recovered RunAgents result"); + }; + assert_eq!(agents.len(), 2); + assert!(matches!( + &agents[0].kind, + RunAgentsAgentOutcomeKind::Launched { agent_id } + if agent_id == &existing_child_id.to_string() + )); + assert!(matches!( + &agents[1].kind, + RunAgentsAgentOutcomeKind::Launched { agent_id } + if agent_id == &missing_child_id.to_string() + )); + }); +} + +#[test] +fn cancelling_recovered_run_agents_keeps_persisted_child_running() { + App::test((), |mut app| async move { + let state = initialize_run_agents_test(&mut app, ExecutionMode::Sdk); + let terminal_view_id = EntityId::new(); + let history = BlocklistAIHistoryModel::handle(&app); + let child_id = history.update(&mut app, |history, ctx| { + history.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + state.conversation_id, + None, + ctx, + ) + }); + let action = remote_run_agents_action("oz"); + state.executor.update(&mut app, |executor, _| { + executor.mark_recovery_actions(&HashSet::from([action.id.clone()])); + }); + let execution = state.executor.update(&mut app, |executor, ctx| { + executor + .execute( + ExecuteActionInput { + action: &action, + conversation_id: state.conversation_id, + }, + ctx, + ) + .into() + }); + let AnyActionExecution::Async { + execute_future, + on_complete, + } = execution + else { + panic!("expected async recovery execution"); + }; + + state.executor.update(&mut app, |executor, ctx| { + executor.cancel_execution(&action.id, ctx); + }); + let async_result = execute_future.await; + let result = app.update(|ctx| on_complete(async_result, ctx)); + assert!(matches!( + result, + AIAgentActionResultType::RunAgents(RunAgentsResult::Cancelled) + )); + history.read(&app, |history, _| { + assert!(matches!( + history.conversation(&child_id).map(|child| child.status()), + Some(crate::ai::agent::conversation::ConversationStatus::InProgress) + )); + }); + }); +} + #[test] fn completion_slots_are_polled_concurrently_and_preserve_request_order() { App::test((), |_app| async move { 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 4aa82346..5c882498 100644 --- a/app/src/ai/blocklist/action_model/execute/start_agent.rs +++ b/app/src/ai/blocklist/action_model/execute/start_agent.rs @@ -713,6 +713,45 @@ impl StartAgentExecutor { } } + pub fn reattach( + &mut self, + action_id: AIAgentActionId, + name: String, + parent_conversation_id: AIConversationId, + child_conversation_id: AIConversationId, + parent_run_id: Option, + ctx: &mut ModelContext, + ) -> StartAgentDispatch { + let wait_policy = if parent_run_id.is_none() { + StartAgentWaitPolicy::Completion + } else { + StartAgentWaitPolicy::Startup + }; + let (sender, receiver) = async_channel::bounded(1); + let request_id = self.next_request_id(); + self.pending.insert( + request_id, + PendingStartAgent { + action_id, + run_agents_child_name: Some(name), + parent_conversation_id, + child_conversation_id: Some(child_conversation_id), + sender, + wait_policy, + }, + ); + self.record_child_conversation(request_id, child_conversation_id, ctx); + StartAgentDispatch { + receiver, + wait_policy, + } + } + + pub fn cancel_dispatches_for_action(&mut self, action_id: &AIAgentActionId) { + self.pending + .retain(|_, pending| &pending.action_id != action_id); + } + pub(super) fn preprocess_action( &mut self, _action: PreprocessActionInput, 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 d0bb5103..312ec533 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 @@ -946,6 +946,79 @@ fn run_agents_dispatch_publishes_only_run_agents_child_link() { }); } +#[test] +fn reattach_reuses_persisted_child_without_launching_another_agent() { + App::test((), |mut app| async move { + initialize_history_persistence_for_tests(&mut app); + let terminal_view_id = EntityId::new(); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let executor = app.add_model(StartAgentExecutor::new); + let captured_prompts = capture_start_agent_prompts(&mut app, &executor); + let captured_links = capture_run_agents_child_links(&mut app, &executor); + let parent_conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, false, ctx) + }); + let child_conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + parent_conversation_id, + None, + ctx, + ) + }); + let action_id = AIAgentActionId::from("run-agents-action".to_string()); + + let dispatch = executor.update(&mut app, |executor, ctx| { + executor.reattach( + action_id.clone(), + "child".to_string(), + parent_conversation_id, + child_conversation_id, + None, + ctx, + ) + }); + + assert_eq!(dispatch.wait_policy, StartAgentWaitPolicy::Completion); + assert!(matches!( + dispatch.receiver.try_recv(), + Err(async_channel::TryRecvError::Empty) + )); + captured_prompts.read(&app, |captured, _| { + assert!(captured.0.is_empty()); + }); + captured_links.read(&app, |captured, _| { + assert_eq!( + captured.0, + vec![( + action_id, + "child".to_string(), + parent_conversation_id, + child_conversation_id, + )] + ); + }); + + history_model.update(&mut app, |history_model, ctx| { + history_model.update_conversation_status( + terminal_view_id, + child_conversation_id, + ConversationStatus::Success, + ctx, + ); + }); + assert!(matches!( + dispatch.receiver.try_recv(), + Ok(StartAgentOutcome::Completed { agent_id, .. }) + if agent_id == child_conversation_id.to_string() + )); + executor.read(&app, |executor, _| { + assert!(executor.pending.is_empty()); + }); + }); +} + #[test] fn execute_waits_for_direct_provider_child_and_returns_its_output() { App::test((), |mut app| async move { diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 5854389e..0006c920 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -21,9 +21,9 @@ use anyhow::anyhow; use chrono::{DateTime, Local}; use futures::channel::oneshot; use galaxy_agent_core::{ - turn_control, ExternalWorkId, PendingToolBatch, ProviderRun, ProviderRunFailureKind, - ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, ToolLoopGuard, - TurnCommand, TurnCommandSender, TurnRequest, + turn_control, ExternalWorkId, PendingToolBatch, PendingToolCallState, ProviderRun, + ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, + ToolLoopGuard, TurnCommand, TurnCommandSender, TurnRequest, }; use galaxy_core::assertions::safe_assert; use input_context::{input_context_for_request, parse_context_attachments}; @@ -930,12 +930,39 @@ fn record_provider_batch_signal( committed_work_id.as_ref() == Some(work_id) && finished_work_id.as_ref() == Some(work_id) } +fn recoverable_run_agents_call_ids( + snapshot: &ActiveProviderRunSnapshot, +) -> Result, String> { + let ProviderRunState::AwaitingTools { batch } = snapshot.run.state() else { + return Ok(HashSet::new()); + }; + batch + .calls + .iter() + .filter(|pending| { + matches!( + pending.state, + PendingToolCallState::Executing | PendingToolCallState::RecoveryPending + ) + }) + .try_fold(HashSet::new(), |mut call_ids, pending| { + let action = snapshot + .action_context + .action_from_tool_call(&pending.call)?; + if matches!(action.action, AIAgentActionType::RunAgents(_)) { + call_ids.insert(pending.call.id.clone()); + } + Ok(call_ids) + }) +} + fn normalize_restored_provider_snapshot( snapshot: &mut ActiveProviderRunSnapshot, ) -> Result<(), String> { + let recoverable_call_ids = recoverable_run_agents_call_ids(snapshot)?; let normalization = snapshot .run - .normalize_after_restore() + .normalize_after_restore_with_recoverable_calls(&recoverable_call_ids) .map_err(|error| error.to_string())?; let interrupted_call_ids = normalization .interrupted_call_ids @@ -5404,10 +5431,19 @@ impl BlocklistAIController { .calls .iter() .filter(|pending| pending.state.result().is_none()) - .map(|pending| run.action_context.action_from_tool_call(&pending.call)) + .map(|pending| { + run.action_context + .action_from_tool_call(&pending.call) + .map(|action| { + ( + action, + matches!(pending.state, PendingToolCallState::RecoveryPending), + ) + }) + }) .collect::, _>>() }); - let actions = match conversion { + let converted_actions = match conversion { Some(Ok(actions)) => actions, Some(Err(message)) => { self.fail_active_provider_run(conversation_id, message, ctx); @@ -5418,25 +5454,58 @@ impl BlocklistAIController { let stream_id = self.active_provider_runs[&conversation_id] .stream_id .clone(); - for action in &actions { - let apply_result = - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { - history_model.apply_domain_tool_proposal( - &stream_id, + let mut recovery_action_ids = HashSet::new(); + let mut actions = Vec::with_capacity(converted_actions.len()); + for (mut action, is_recovery) in converted_actions { + if is_recovery { + let restored_action = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .and_then(|conversation| conversation.action(&action.id)); + let Some(restored_action) = restored_action else { + self.fail_active_provider_run( conversation_id, - self.terminal_surface_id, - action.clone(), + format!( + "restored RunAgents action {} is missing from conversation history", + action.id + ), ctx, - ) - }); - if let Err(error) = apply_result { - self.fail_active_provider_run( - conversation_id, - format!("failed to attach provider tool proposal: {error:?}"), - ctx, - ); - return; + ); + return; + }; + if !matches!(restored_action.action, AIAgentActionType::RunAgents(_)) { + self.fail_active_provider_run( + conversation_id, + format!( + "restored provider action {} no longer matches RunAgents history", + action.id + ), + ctx, + ); + return; + } + recovery_action_ids.insert(action.id.clone()); + action = restored_action; + } else { + let apply_result = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.apply_domain_tool_proposal( + &stream_id, + conversation_id, + self.terminal_surface_id, + action.clone(), + ctx, + ) + }); + if let Err(error) = apply_result { + self.fail_active_provider_run( + conversation_id, + format!("failed to attach provider tool proposal: {error:?}"), + ctx, + ); + return; + } } + actions.push(action); } if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) { slot.command_action_refs.extend( @@ -5464,7 +5533,13 @@ impl BlocklistAIController { return; } let queue_result = self.action_model.update(ctx, |action_model, ctx| { - action_model.queue_provider_actions(actions, conversation_id, &batch, ctx) + action_model.queue_provider_actions( + actions, + recovery_action_ids, + conversation_id, + &batch, + ctx, + ) }); if let Err(error) = queue_result { self.fail_active_provider_run(conversation_id, error.to_string(), ctx); diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 8bc1f39a..ee8d12b2 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -312,6 +312,60 @@ fn restore_normalization_removes_interrupted_command_correlation() { )); } +#[test] +fn restore_normalization_preserves_executing_run_agents_for_recovery() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let Some(ProviderRunStep::CallModel(call)) = snapshot.run.next_step().unwrap() else { + panic!("expected provider model call"); + }; + snapshot + .run + .accept_model_turn( + &call.work_id, + CompletedModelTurn { + assistant_content: vec![ContentPart::Text("I will run child agents.".to_owned())], + tool_calls: vec![ToolCall { + id: "run-agents-call".to_owned(), + name: "run_agents".to_owned(), + arguments: serde_json::json!({ + "summary": "Run child agents", + "base_prompt": "Shared instructions", + "agent_run_configs": [{ + "name": "child", + "prompt": "Do work", + "title": "Child", + }], + }), + }], + usage: Usage::default(), + stop_reason: StopReason::Completed, + advertised_tools: BTreeSet::from(["run_agents".to_owned()]), + }, + ) + .unwrap(); + let Some(ProviderRunStep::DispatchTools(batch)) = snapshot.run.next_step().unwrap() else { + panic!("expected provider tool batch"); + }; + snapshot + .run + .start_tool(&batch.work_id, "run-agents-call") + .unwrap(); + + let recoverable = super::recoverable_run_agents_call_ids(&snapshot).unwrap(); + assert_eq!(recoverable.len(), 1); + assert!(recoverable.contains("run-agents-call")); + super::normalize_restored_provider_snapshot(&mut snapshot).unwrap(); + + let ProviderRunState::AwaitingTools { batch } = snapshot.run.state() else { + panic!("recovered RunAgents call should keep the provider batch pending"); + }; + assert!(matches!( + batch.calls[0].state, + galaxy_agent_core::PendingToolCallState::RecoveryPending + )); +} + #[test] fn restore_normalization_reproposes_permission_without_losing_correlation() { let conversation_id = AIConversationId::new(); diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 340da19c..d824fb60 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -139,6 +139,8 @@ pub enum PendingToolCallState { decision: PermissionDecision, }, Executing, + /// External work survived a restart and must be reattached before it can complete. + RecoveryPending, Resolved { result: ToolResult, }, @@ -151,7 +153,8 @@ impl PendingToolCallState { Self::Proposed | Self::PermissionPending { .. } | Self::Approved { .. } - | Self::Executing => None, + | Self::Executing + | Self::RecoveryPending => None, } } @@ -161,6 +164,7 @@ impl PendingToolCallState { Self::PermissionPending { .. } => "permission_pending", Self::Approved { .. } => "approved", Self::Executing => "executing", + Self::RecoveryPending => "recovery_pending", Self::Resolved { .. } => "resolved", } } @@ -300,6 +304,7 @@ pub enum ModelFailureDisposition { pub struct ProviderRunRestoreNormalization { pub permission_call_ids_reset: Vec, pub interrupted_call_ids: Vec, + pub recovery_call_ids: Vec, pub committed_tool_batch: bool, } @@ -499,6 +504,13 @@ impl ProviderRun { pub fn normalize_after_restore( &mut self, + ) -> Result { + self.normalize_after_restore_with_recoverable_calls(&HashSet::new()) + } + + pub fn normalize_after_restore_with_recoverable_calls( + &mut self, + recoverable_call_ids: &HashSet, ) -> Result { let ProviderRunState::AwaitingTools { batch } = &mut self.state else { return Ok(ProviderRunRestoreNormalization::default()); @@ -513,6 +525,19 @@ impl ProviderRun { .push(pending.call.id.clone()); pending.state = PendingToolCallState::Proposed; } + PendingToolCallState::Executing + if recoverable_call_ids.contains(&pending.call.id) => + { + normalization + .recovery_call_ids + .push(pending.call.id.clone()); + pending.state = PendingToolCallState::RecoveryPending; + } + PendingToolCallState::RecoveryPending => { + normalization + .recovery_call_ids + .push(pending.call.id.clone()); + } PendingToolCallState::Approved { .. } | PendingToolCallState::Executing => { normalization .interrupted_call_ids @@ -719,7 +744,8 @@ impl ProviderRun { } PendingToolCallState::PermissionPending { .. } | PendingToolCallState::Approved { .. } - | PendingToolCallState::Executing => { + | PendingToolCallState::Executing + | PendingToolCallState::RecoveryPending => { Err(invalid_tool_transition(call, "permission request")) } } @@ -742,7 +768,8 @@ impl ProviderRun { } PendingToolCallState::Proposed | PendingToolCallState::Approved { .. } - | PendingToolCallState::Executing => { + | PendingToolCallState::Executing + | PendingToolCallState::RecoveryPending => { return Err(invalid_tool_transition(call, "permission resolution")); } }; @@ -782,7 +809,9 @@ impl ProviderRun { ) -> Result<(), ProviderRunProtocolError> { let call = self.pending_tool_call_mut(work_id, call_id)?; match &call.state { - PendingToolCallState::Proposed | PendingToolCallState::Approved { .. } => { + PendingToolCallState::Proposed + | PendingToolCallState::Approved { .. } + | PendingToolCallState::RecoveryPending => { call.state = PendingToolCallState::Executing; Ok(()) } @@ -806,7 +835,8 @@ impl ProviderRun { match &call.state { PendingToolCallState::Proposed | PendingToolCallState::Approved { .. } - | PendingToolCallState::Executing => { + | PendingToolCallState::Executing + | PendingToolCallState::RecoveryPending => { call.state = PendingToolCallState::Resolved { result }; Ok(()) } @@ -837,7 +867,8 @@ impl ProviderRun { PendingToolCallState::Proposed | PendingToolCallState::PermissionPending { .. } | PendingToolCallState::Approved { .. } - | PendingToolCallState::Executing => { + | PendingToolCallState::Executing + | PendingToolCallState::RecoveryPending => { call.state = PendingToolCallState::Resolved { result: ToolResult { call_id: call_id.to_string(), diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index d2be435c..0f482cad 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -758,6 +758,73 @@ fn restore_normalization_reproposes_permissions_and_interrupts_unsafe_tools() { ); } +#[test] +fn restore_normalization_preserves_selected_executing_tools_for_recovery() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("run-agents", "run_agents"), + tool_call("read", "read_files"), + ], + &["run_agents", "read_files"], + ), + ); + run.start_tool(&batch.work_id, "run-agents").unwrap(); + run.start_tool(&batch.work_id, "read").unwrap(); + + let normalization = run + .normalize_after_restore_with_recoverable_calls(&HashSet::from(["run-agents".to_string()])) + .unwrap(); + + assert_eq!(normalization.recovery_call_ids, vec!["run-agents"]); + assert_eq!(normalization.interrupted_call_ids, vec!["read"]); + let ProviderRunState::AwaitingTools { batch } = run.state() else { + panic!("recovered tool batch should remain pending"); + }; + assert!(matches!( + batch.calls[0].state, + PendingToolCallState::RecoveryPending + )); + assert_eq!( + batch.calls[1].state.result().unwrap().status, + ToolResultStatus::Error + ); +} + +#[test] +fn recovery_pending_tool_survives_another_restore_and_completes_once() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("run-agents", "run_agents")], &["run_agents"]), + ); + run.start_tool(&batch.work_id, "run-agents").unwrap(); + run.normalize_after_restore_with_recoverable_calls(&HashSet::from(["run-agents".to_string()])) + .unwrap(); + let serialized = serde_json::to_string(&run).unwrap(); + let mut restored: ProviderRun = serde_json::from_str(&serialized).unwrap(); + + let normalization = restored.normalize_after_restore().unwrap(); + assert_eq!(normalization.recovery_call_ids, vec!["run-agents"]); + restored + .complete_tool( + &batch.work_id, + successful_result("run-agents", "children completed"), + ) + .unwrap(); + assert!(matches!( + restored.complete_tool( + &batch.work_id, + successful_result("run-agents", "duplicate completion"), + ), + Err(ProviderRunProtocolError::DuplicateToolUpdate { .. }) + )); + restored.commit_tool_batch(&batch.work_id).unwrap(); + assert_eq!(restored.state().phase(), ProviderRunPhase::ReadyToCallModel); +} + #[test] fn restore_normalization_commits_a_fully_resolved_batch() { let mut run = run();