From 56e3b51d48a746501887e61f6bf85de3217a7acd Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Mon, 17 Aug 2026 18:19:37 -0500 Subject: [PATCH] Lots of changes... not done yet. --- AGENTS.md | 9 +- app/src/ai/agent/conversation.rs | 3 +- app/src/ai/agent/mod.rs | 42 +- app/src/ai/agent_sdk/driver/output.rs | 8 + app/src/ai/blocklist/action_model.rs | 288 ++++--- app/src/ai/blocklist/action_model/execute.rs | 266 ++++-- .../action_model/execute/call_mcp_tool.rs | 20 +- .../action_model/execute/file_glob.rs | 4 +- .../ai/blocklist/action_model/execute/grep.rs | 4 +- .../action_model/execute/read_files.rs | 4 +- .../execute/request_file_edits.rs | 8 +- .../action_model/execute/run_agents.rs | 126 ++- .../action_model/execute/run_agents_tests.rs | 76 +- .../action_model/execute/shell_command.rs | 191 +++-- .../execute/shell_command_tests.rs | 208 ++++- .../action_model/execute/start_agent.rs | 103 ++- .../action_model/execute/start_agent_tests.rs | 130 ++- app/src/ai/blocklist/action_model_tests.rs | 125 +++ app/src/ai/blocklist/block.rs | 93 ++- app/src/ai/blocklist/block/cli.rs | 2 +- app/src/ai/blocklist/block/cli_controller.rs | 170 +++- app/src/ai/blocklist/block/model/helper.rs | 6 +- app/src/ai/blocklist/block/status_bar.rs | 25 +- app/src/ai/blocklist/block/view_impl.rs | 1 + .../block/view_impl/orchestration.rs | 10 +- .../ai/blocklist/block/view_impl/output.rs | 90 +- app/src/ai/blocklist/controller.rs | 608 ++++++++++---- .../controller/pending_response_streams.rs | 14 +- .../ai/blocklist/controller/shared_session.rs | 2 +- app/src/ai/blocklist/controller_tests.rs | 463 ++++++++++- .../inline_action/ask_user_question_view.rs | 9 +- .../blocklist/inline_action/code_diff_view.rs | 18 +- .../inline_action/requested_command.rs | 37 +- .../inline_action/run_agents_card_view.rs | 106 ++- .../run_agents_card_view_tests.rs | 34 + .../ai/blocklist/inline_action/tool_pane.rs | 3 +- app/src/ai/mod.rs | 1 + app/src/ai/runtime/event_translator.rs | 111 ++- app/src/ai/runtime/event_translator_tests.rs | 173 +++- .../ai/runtime/provider_run_coordinator.rs | 411 ++++++--- .../runtime/provider_run_coordinator_tests.rs | 141 +++- app/src/ai/runtime/rig_request.rs | 10 + app/src/ai/runtime/rig_request_tests.rs | 22 + app/src/ai/runtime/rig_tool.rs | 778 +++++++++++------- app/src/ai/runtime/rig_tool_tests.rs | 241 ++++++ app/src/ai/tool_diagnostics.rs | 46 ++ app/src/terminal/view.rs | 49 +- app/src/terminal/view/context_menu.rs | 6 +- app/src/terminal/view/load_ai_conversation.rs | 14 +- crates/ai/src/agent/action_result/convert.rs | 27 + .../src/agent/action_result/convert_tests.rs | 19 + crates/ai/src/agent/action_result/mod.rs | 61 +- .../ai/src/agent/action_result/mod_tests.rs | 87 +- crates/galaxy_agent_core/src/provider_run.rs | 6 +- .../src/provider_run_tests.rs | 83 ++ 55 files changed, 4494 insertions(+), 1098 deletions(-) create mode 100644 app/src/ai/tool_diagnostics.rs diff --git a/AGENTS.md b/AGENTS.md index e8bc4128..f77f7900 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -40,6 +40,7 @@ Environment variables: - Set `GALAXY_BEDROCK_DIAGNOSTICS=1` to enable Bedrock diagnostic output, including: - `Error_.txt` snapshot files written to the repository root on request/stream failures (includes the serialized Bedrock context window, tool definitions, protobuf request debug payload, captured Bedrock diagnostic lines, and log tails) - Per-event Bedrock diagnostic logs written to `bedrock-diagnostics.log` in the active Warp log directory +- Set `GALAXY_TOOL_DIAGNOSTICS=1` to enable verbose local tool queue/execution debug logs and cancellation backtraces. These diagnostics are disabled during routine operation. ### AI Provider Architecture @@ -139,7 +140,10 @@ Key invariants: - Direct-provider `RequestFileEdits` views must register from streaming output before provider-run completion; preprocessing results must survive delayed view registration, and `NotReady` retries must remain automatic rather than emitting a synthetic user permission decision - A clean direct-provider `ProviderRunOutcome::Completed` explicitly finalizes the conversation as `Success` after terminal output projection, even if earlier turns added tool actions; child-completion waits rely on that status - Provider actions and results must correlate by `(conversation_id, run_id, epoch, call_id)`; stale or duplicate callbacks must not advance a run -- Active provider runs must checkpoint before external work, persist without credentials, normalize unsafe restored states, and reconcile command state before continuing +- 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 - 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 @@ -151,10 +155,13 @@ Key invariants: - Progressive summaries are prepended to provider requests as a user/assistant pair; background summarization remains independent of the active provider run - 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 - 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 - Direct-provider `RunAgents` remains pending until every local child reaches `Success`, `Error`, or `Cancelled`, or is removed/deleted; recoverable `Blocked`, `TransientError`, and `WaitingForEvents` states remain pending, and the hosted 30-second startup timeout must not apply to these completion waits +- `StartAgentWaitPolicy` is selected from child execution mode, not parent `run_id`: local children wait for completion and only remote/hosted children use startup acknowledgement +- Hosted `RunAgents` startup timeouts detach the exact `StartAgentRequestId`; late launch callbacks must not register the child after the timeout result, while children linked before cancellation remain independently running ### Platform Setup - `./script/bootstrap` - Platform-specific setup plus common agent skill installation from `skills-lock.json`; prompts for project/global when an install or update is needed unless a target flag or environment override is provided. diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 938def3e..08d437c8 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -1930,7 +1930,8 @@ impl AIConversation { ) -> String { let mut result = Vec::new(); for exchange in self.all_exchanges() { - let formatted_exchange = exchange.format_for_copy(action_model); + let formatted_exchange = + exchange.format_for_copy_for_conversation(action_model, Some(self.id())); if !formatted_exchange.is_empty() { result.push(formatted_exchange); } diff --git a/app/src/ai/agent/mod.rs b/app/src/ai/agent/mod.rs index d27a0ea0..9f7c2d5b 100644 --- a/app/src/ai/agent/mod.rs +++ b/app/src/ai/agent/mod.rs @@ -581,6 +581,14 @@ impl AIAgentOutput { pub fn format_for_copy( &self, action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>, + ) -> String { + self.format_for_copy_for_conversation(action_model, None) + } + + pub fn format_for_copy_for_conversation( + &self, + action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>, + conversation_id: Option, ) -> String { let mut result = Vec::new(); let mut last_was_action = false; @@ -612,8 +620,12 @@ impl AIAgentOutput { } AIAgentOutputMessageType::Action(action) => { // Include action results from the action model if available - if let Some(action_model) = action_model { - if let Some(action_result) = action_model.get_action_result(&action.id) { + if let (Some(action_model), Some(conversation_id)) = + (action_model, conversation_id) + { + if let Some(action_result) = + action_model.get_action_result(conversation_id, &action.id) + { result.push(format!("{}", MarkdownActionResult(&action_result.result))); // Add an extra newline after tool call results for readability result.push(String::new()); @@ -1222,6 +1234,9 @@ impl<'a> std::fmt::Display for MarkdownActionResult<'a> { RequestCommandOutputResult::CancelledBeforeExecution => { write!(f, "\n_Command cancelled_") } + RequestCommandOutputResult::ExecutionError { command, message } => { + write!(f, "\n_Command `{command}` was not executed: {message}_") + } RequestCommandOutputResult::Denylisted { command } => { write!( f, @@ -3225,9 +3240,19 @@ impl AIAgentExchange { pub fn format_output_for_copy( &self, action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>, + ) -> String { + self.format_output_for_copy_for_conversation(action_model, None) + } + + pub fn format_output_for_copy_for_conversation( + &self, + action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>, + conversation_id: Option, ) -> String { match self.output_status.output() { - Some(output) => output.get().format_for_copy(action_model), + Some(output) => output + .get() + .format_for_copy_for_conversation(action_model, conversation_id), None => String::new(), } } @@ -3238,9 +3263,18 @@ impl AIAgentExchange { pub fn format_for_copy( &self, action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>, + ) -> String { + self.format_for_copy_for_conversation(action_model, None) + } + + pub fn format_for_copy_for_conversation( + &self, + action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>, + conversation_id: Option, ) -> String { let input_text = self.format_input_for_copy(); - let output_text = self.format_output_for_copy(action_model); + let output_text = + self.format_output_for_copy_for_conversation(action_model, conversation_id); let has_user_input = !input_text.is_empty(); let has_agent_output = !output_text.is_empty(); diff --git a/app/src/ai/agent_sdk/driver/output.rs b/app/src/ai/agent_sdk/driver/output.rs index 58688e3f..a40ce9ad 100644 --- a/app/src/ai/agent_sdk/driver/output.rs +++ b/app/src/ai/agent_sdk/driver/output.rs @@ -61,6 +61,9 @@ pub mod text { RequestCommandOutputResult::CancelledBeforeExecution => { writeln!(w, "{CANCELLED_MESSAGE}") } + RequestCommandOutputResult::ExecutionError { command, message } => { + writeln!(w, "Command `{command}` was not executed: {message}") + } RequestCommandOutputResult::Denylisted { .. } => { writeln!( w, @@ -829,6 +832,11 @@ pub mod json { RequestCommandOutputResult::CancelledBeforeExecution => { Some(JsonMessage::ToolCanceled) } + RequestCommandOutputResult::ExecutionError { message, .. } => { + Some(JsonMessage::ToolError { + error: Cow::Borrowed(message), + }) + } RequestCommandOutputResult::Denylisted { .. } => Some(JsonMessage::ToolError { error: Cow::Borrowed( "Command was not allowed to run due to presence on denylist", diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 445c6009..ab666d01 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -31,7 +31,7 @@ pub use execute::{ ReadFileContextResult, RequestFileEditsExecutor, RequestFileEditsFormatKind, RequestFileEditsTelemetryEvent, RunAgentsExecutor, RunAgentsExecutorEvent, RunAgentsSpawningSnapshot, ShellCommandExecutor, ShellCommandExecutorEvent, StartAgentExecutor, - StartAgentExecutorEvent, StartAgentRequest, StartAgentRequestId, + StartAgentExecutorEvent, StartAgentRequest, StartAgentRequestId, StartAgentWaitPolicy, }; use futures::future::{join_all, BoxFuture}; use galaxy_agent_core::{ @@ -355,6 +355,45 @@ fn sort_action_results_by_order( results.sort_by_key(|result| action_order.get(&result.id).copied().unwrap_or(usize::MAX)); } +fn action_result_for_conversation<'a>( + finished_action_results: &'a HashMap>>, + provider_finished_action_results: &'a HashMap< + (AIConversationId, ExternalWorkId), + Vec>, + >, + past_action_results: &'a HashMap<(AIConversationId, AIAgentActionId), Arc>, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, +) -> Option<&'a Arc> { + finished_action_results + .get(&conversation_id) + .into_iter() + .chain(provider_finished_action_results.iter().filter_map( + |((result_conversation_id, _), results)| { + (*result_conversation_id == conversation_id).then_some(results) + }, + )) + .flat_map(|results| results.iter()) + .find(|result| &result.id == action_id) + .or_else(|| past_action_results.get(&(conversation_id, action_id.clone()))) +} + +fn pending_action_status( + pending_actions: &HashMap>, + running_actions: &HashMap, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + is_view_only: bool, +) -> Option { + let actions = pending_actions.get(&conversation_id)?; + let index = actions.iter().position(|action| &action.id == action_id)?; + if index == 0 && !is_view_only && !running_actions.contains_key(&conversation_id) { + Some(AIActionStatus::Blocked) + } else { + Some(AIActionStatus::Queued) + } +} + fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bool) -> ToolResult { let status = if permission_denied { ToolResultStatus::Denied @@ -587,6 +626,9 @@ fn action_result_failure_summary(result: &AIAgentActionResultType) -> Option { Some(error.clone()) } + AIAgentActionResultType::RequestCommandOutput( + RequestCommandOutputResult::ExecutionError { message, .. }, + ) => Some(message.clone()), AIAgentActionResultType::RequestCommandOutput( RequestCommandOutputResult::Completed { .. } | RequestCommandOutputResult::CancelledBeforeExecution @@ -734,7 +776,7 @@ pub struct BlocklistAIActionModel { HashMap<(AIConversationId, AIAgentActionId), ProviderToolExecutionRef>, /// Past actions and their corresponding statuses from previous AI exchanges. - past_action_results: HashMap>, + past_action_results: HashMap<(AIConversationId, AIAgentActionId), Arc>, /// The ID of the terminal view this controller is associated with. terminal_view_id: EntityId, @@ -774,6 +816,7 @@ impl BlocklistAIActionModel { let execution_ref = me.provider_tool_execution_ref(*conversation_id, action_id); ctx.emit(BlocklistAIActionEvent::ExecutingAction { action_id: action_id.clone(), + conversation_id: *conversation_id, execution_ref: execution_ref.clone(), }); ctx.emit(BlocklistAIActionEvent::ToolLifecycle { @@ -863,6 +906,7 @@ impl BlocklistAIActionModel { ); ctx.emit(BlocklistAIActionEvent::ExecutingAction { action_id: action_id.clone(), + conversation_id, execution_ref: self.provider_tool_execution_ref(conversation_id, action_id), }); } @@ -988,8 +1032,8 @@ impl BlocklistAIActionModel { .get(&conversation_id) .map(|q| q.len()) .unwrap_or(0); - log::info!( - "[tool-debug] try_to_execute_available_actions: conversation={:?}, pending_count={}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_available_actions: conversation={:?}, pending_count={}", conversation_id, pending_count ); @@ -1000,14 +1044,14 @@ impl BlocklistAIActionModel { .and_then(|queue| queue.front()) .cloned() else { - log::info!( - "[tool-debug] try_to_execute_available_actions: no more pending actions" + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_available_actions: no more pending actions" ); return; }; - log::info!( - "[tool-debug] try_to_execute_available_actions: trying action id={:?}, type={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_available_actions: trying action id={:?}, type={:?}", front_action.id, std::mem::discriminant(&front_action.action) ); @@ -1019,8 +1063,8 @@ impl BlocklistAIActionModel { current_phase, ctx, ) { - log::info!( - "[tool-debug] try_to_execute_available_actions: cannot start in current phase {:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_available_actions: cannot start in current phase {:?}", current_phase ); return; @@ -1033,12 +1077,12 @@ impl BlocklistAIActionModel { ActionExecutionInitiator::Automatic, ctx, ) else { - log::info!("[tool-debug] try_to_execute_available_actions: start_pending_action_by_id returned None (blocked)"); + crate::ai::tool_diagnostics::tool_debug!("try_to_execute_available_actions: start_pending_action_by_id returned None (blocked)"); return; }; - log::info!( - "[tool-debug] try_to_execute_available_actions: action started, result={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_available_actions: action started, result={:?}", std::mem::discriminant(&result) ); @@ -1048,7 +1092,9 @@ impl BlocklistAIActionModel { phase: RunningActionPhase::Serial } ) { - log::info!("[tool-debug] try_to_execute_available_actions: serial async action, stopping loop"); + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_available_actions: serial async action, stopping loop" + ); return; } } @@ -1081,14 +1127,6 @@ impl BlocklistAIActionModel { } } - /// Returns all pending actions for all conversations. - pub fn get_pending_actions(&self) -> Vec<&AIAgentAction> { - self.pending_actions - .values() - .flat_map(|queue| queue.iter()) - .collect() - } - /// Returns all pending actions for a specific conversation. pub fn get_pending_actions_for_conversation( &self, @@ -1106,11 +1144,15 @@ impl BlocklistAIActionModel { self.blocked_action_for_conversation(&conversation_id) } - /// Returns a pending action by its ID, searching across all conversations. - pub fn get_pending_action_by_id(&self, action_id: &AIAgentActionId) -> Option<&AIAgentAction> { + /// Returns a pending action by its ID within the given conversation. + pub fn get_pending_action_by_id( + &self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> Option<&AIAgentAction> { self.pending_actions - .values() - .flat_map(|queue| queue.iter()) + .get(&conversation_id)? + .iter() .find(|action| &action.id == action_id) } @@ -1142,7 +1184,11 @@ impl BlocklistAIActionModel { self.running_actions .get(&conversation_id) .and_then(RunningActions::first_action_id) - .and_then(|action_id| self.executor.as_ref(app).async_executing_action(action_id)) + .and_then(|action_id| { + self.executor + .as_ref(app) + .async_executing_action(conversation_id, action_id) + }) } /// Returns whether there is a pending or running action for the active conversation. @@ -1194,53 +1240,58 @@ impl BlocklistAIActionModel { self.finished_action_results.get(&conversation_id) } - /// Returns the `AIActionStatus` for the action corresponding to the given `id`, if any. - pub fn get_action_status(&self, id: &AIAgentActionId) -> Option { - for (conversation_id, pending_actions_for_conversation) in &self.pending_actions { - for (index, action) in pending_actions_for_conversation.iter().enumerate() { - if &action.id != id { - continue; - } - - if index == 0 - && !self.is_view_only - && !self.running_actions.contains_key(conversation_id) - { - return Some(AIActionStatus::Blocked); - } - - return Some(AIActionStatus::Queued); - } + /// Returns the status for an action within the given conversation. + pub fn get_action_status( + &self, + conversation_id: AIConversationId, + id: &AIAgentActionId, + ) -> Option { + if let Some(status) = pending_action_status( + &self.pending_actions, + &self.running_actions, + conversation_id, + id, + self.is_view_only, + ) { + return Some(status); } self.running_actions - .values() - .find(|running| running.contains(id)) + .get(&conversation_id) + .filter(|running| running.contains(id)) .map(|_| AIActionStatus::RunningAsync) .or_else(|| { - self.get_action_result(id) + self.get_action_result(conversation_id, id) .map(|result| AIActionStatus::Finished(result.clone())) }) .or_else(|| { self.pending_preprocessed_actions - .values() - .any(|preprocessing| preprocessing.contains(id)) + .get(&conversation_id) + .is_some_and(|preprocessing| preprocessing.contains(id)) .then_some(AIActionStatus::Preprocessing) }) } - pub fn get_action_result(&self, id: &AIAgentActionId) -> Option<&Arc> { - // Search through all conversations' finished action results - self.finished_action_results - .values() - .chain(self.provider_finished_action_results.values()) - .flat_map(|results| results.iter()) - .find(|result| &result.id == id) - .or_else(|| self.past_action_results.get(id)) + pub fn get_action_result( + &self, + conversation_id: AIConversationId, + id: &AIAgentActionId, + ) -> Option<&Arc> { + action_result_for_conversation( + &self.finished_action_results, + &self.provider_finished_action_results, + &self.past_action_results, + conversation_id, + id, + ) } /// Bulk restore action results from a list of exchanges (used when loading conversations from tasks) - pub fn restore_action_results_from_exchanges(&mut self, exchanges: Vec<&AIAgentExchange>) { + pub fn restore_action_results_from_exchanges( + &mut self, + conversation_id: AIConversationId, + exchanges: Vec<&AIAgentExchange>, + ) { for exchange in exchanges.iter() { for input in &exchange.input { if let AIAgentInput::ActionResult { result, .. } = input { @@ -1257,7 +1308,7 @@ impl BlocklistAIActionModel { ); } self.past_action_results - .insert(result_id, Arc::new(result_to_insert)); + .insert((conversation_id, result_id), Arc::new(result_to_insert)); } } } @@ -1267,18 +1318,16 @@ impl BlocklistAIActionModel { /// from the confirmation card. pub fn execute_run_agents( &mut self, + conversation_id: AIConversationId, action_id: &AIAgentActionId, request: ai::agent::action::RunAgentsRequest, ctx: &mut ModelContext, ) { - let mut found = None; - for (conv_id, queue) in self.pending_actions.iter_mut() { - if let Some(action) = queue.iter_mut().find(|action| &action.id == action_id) { - found = Some((*conv_id, action)); - break; - } - } - let Some((conversation_id, action)) = found else { + let Some(action) = self + .pending_actions + .get_mut(&conversation_id) + .and_then(|queue| queue.iter_mut().find(|action| &action.id == action_id)) + else { log::warn!( "BlocklistAIActionModel::execute_run_agents: no pending action for {action_id:?}" ); @@ -1299,20 +1348,19 @@ impl BlocklistAIActionModel { /// the time the action becomes blocked on user confirmation. pub fn deny_run_agents( &mut self, + conversation_id: AIConversationId, action_id: &AIAgentActionId, reason: String, ctx: &mut ModelContext, ) { - let mut found: Option<(AIConversationId, AIAgentAction)> = None; - for (conv_id, queue) in self.pending_actions.iter_mut() { - if let Some(idx) = queue.iter().position(|a| &a.id == action_id) { - if let Some(action) = queue.remove(idx) { - found = Some((*conv_id, action)); - } - break; - } - } - let Some((conversation_id, action)) = found else { + let Some(action) = self + .pending_actions + .get_mut(&conversation_id) + .and_then(|queue| { + let index = queue.iter().position(|action| &action.id == action_id)?; + queue.remove(index) + }) + else { log::warn!( "BlocklistAIActionModel::deny_run_agents: no pending action for {action_id:?}" ); @@ -1451,6 +1499,7 @@ impl BlocklistAIActionModel { let execution_ref = self.provider_tool_execution_ref(conversation_id, &action.id); ctx.emit(BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { action_id: action.id.clone(), + conversation_id, execution_ref: execution_ref.clone(), }); ctx.emit(BlocklistAIActionEvent::ToolLifecycle { @@ -1698,14 +1747,14 @@ impl BlocklistAIActionModel { conversation_id: AIConversationId, ctx: &mut ModelContext, ) { - log::info!( - "[tool-debug] queue_actions: queuing {} actions for conversation {:?}", + crate::ai::tool_diagnostics::tool_debug!( + "queue_actions: queuing {} actions for conversation {:?}", actions.len(), conversation_id ); for (i, action) in actions.iter().enumerate() { - log::info!( - "[tool-debug] queue_actions: [{}] id={:?}, type={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "queue_actions: [{}] id={:?}, type={:?}", i, action.id, std::mem::discriminant(&action.action) @@ -1856,7 +1905,7 @@ impl BlocklistAIActionModel { reason: CancellationReason, ctx: &mut ModelContext, ) { - let status = self.get_action_status(action_id); + let status = self.get_action_status(conversation_id, action_id); let permission_denied = is_permission_denial(reason, status.as_ref()); if self .running_actions @@ -1864,7 +1913,7 @@ impl BlocklistAIActionModel { .is_some_and(|running| running.contains(action_id)) { self.executor.update(ctx, |executor, ctx| { - executor.cancel_running_async_action(action_id, Some(reason), ctx) + executor.cancel_running_async_action(conversation_id, action_id, Some(reason), ctx) }); } else { let Some(pending_actions_for_conversation) = @@ -1933,12 +1982,14 @@ impl BlocklistAIActionModel { }; for action in actions_to_cancel.drain(..).collect_vec() { log::info!( - "Canceling pending action of type {:?} conversation_id={conversation_id:?} action_id={:?}, reason={:?}, backtrace=\n{}", + "Canceling pending action of type {:?} conversation_id={conversation_id:?} action_id={:?}, reason={:?}", AIAgentActionTypeDiscriminants::from(&action.action), action.id, - reason, - std::backtrace::Backtrace::force_capture() + reason ); + if let Some(backtrace) = crate::ai::tool_diagnostics::capture_backtrace() { + log::debug!("Pending action cancellation backtrace:\n{backtrace}"); + } self.cancel_pending_action(conversation_id, action, reason, false, ctx); } } @@ -2053,7 +2104,7 @@ impl BlocklistAIActionModel { for result in finished_action_results.iter() { self.past_action_results - .insert(result.id.clone(), result.clone()); + .insert((conversation_id, result.id.clone()), result.clone()); } finished_action_results .into_iter() @@ -2091,7 +2142,8 @@ impl BlocklistAIActionModel { .remove(&(conversation_id, work_id.clone())) .unwrap_or_default(); for result in results { - self.past_action_results.insert(result.id.clone(), result); + self.past_action_results + .insert((conversation_id, result.id.clone()), result); } } @@ -2123,33 +2175,28 @@ impl BlocklistAIActionModel { /// respective functions. pub fn handle_requested_command_accepted( &mut self, + conversation_id: AIConversationId, action_id: &AIAgentActionId, command: String, ctx: &mut ModelContext, ) { - // Search through all pending conversations to find the action and conversation ID - let mut found_conversation_id = None; - for (conversation_id, pending_actions_for_conversation) in self.pending_actions.iter_mut() { - if let Some(action) = pending_actions_for_conversation - .iter_mut() - .find(|action| action.id == *action_id) - { - if let AIAgentActionType::RequestCommandOutput { - command: original_command, - .. - } = &mut action.action - { - *original_command = command; - found_conversation_id = Some(*conversation_id); - break; - } - } - } - - let Some(conversation_id) = found_conversation_id else { + let Some(action) = self + .pending_actions + .get_mut(&conversation_id) + .and_then(|actions| actions.iter_mut().find(|action| action.id == *action_id)) + else { log::warn!("Ignoring acceptance for non-pending requested command: {action_id:?}"); return; }; + let AIAgentActionType::RequestCommandOutput { + command: original_command, + .. + } = &mut action.action + else { + log::warn!("Ignoring acceptance for non-command action: {action_id:?}"); + return; + }; + *original_command = command; self.execute_action(action_id, conversation_id, ctx); } @@ -2161,8 +2208,8 @@ impl BlocklistAIActionModel { cancellation_reason: Option, ctx: &mut ModelContext, ) { - log::info!( - "[tool-debug] handle_action_result: action_id={:?}, result_type={:?}, cancellation={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "handle_action_result: action_id={:?}, result_type={:?}, cancellation={:?}", action_result.id, std::mem::discriminant(&action_result.result), cancellation_reason @@ -2453,11 +2500,13 @@ pub enum BlocklistAIActionEvent { /// Emitted when the action with the given ID requires user confirmation to execute. ActionBlockedOnUserConfirmation { action_id: AIAgentActionId, + conversation_id: AIConversationId, execution_ref: Option, }, /// Emitted when the action with the given ID begins execution. ExecutingAction { action_id: AIAgentActionId, + conversation_id: AIConversationId, execution_ref: Option, }, /// Emitted when the action with the given ID has finished. @@ -2496,6 +2545,25 @@ impl BlocklistAIActionEvent { BlocklistAIActionEvent::InsertCodeReviewComments { action_id, .. } => action_id, } } + + pub fn conversation_id(&self) -> Option { + match self { + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { + conversation_id, .. + } + | BlocklistAIActionEvent::ExecutingAction { + conversation_id, .. + } + | BlocklistAIActionEvent::FinishedAction { + conversation_id, .. + } => Some(*conversation_id), + BlocklistAIActionEvent::QueuedAction { .. } + | BlocklistAIActionEvent::ToolLifecycle { .. } + | BlocklistAIActionEvent::InitProject(_) + | BlocklistAIActionEvent::ToggleCodeReview(_) + | BlocklistAIActionEvent::InsertCodeReviewComments { .. } => None, + } + } } impl Entity for BlocklistAIActionModel { diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index c0050ee8..83c4580c 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -74,6 +74,7 @@ use serde::{Deserialize, Serialize}; pub use shell_command::{ShellCommandExecutor, ShellCommandExecutorEvent}; pub use start_agent::{ StartAgentExecutor, StartAgentExecutorEvent, StartAgentRequest, StartAgentRequestId, + StartAgentWaitPolicy, }; pub use suggest_new_conversation::NewConversationDecision; use suggest_new_conversation::SuggestNewConversationExecutor; @@ -245,9 +246,36 @@ pub(super) enum TryExecuteResult { #[derive(Clone)] struct AsyncExecutingAction { action: AIAgentAction, - /// The conversation this action belongs to so cancellation and follow-up scheduling remain - /// scoped even when several conversations have async actions in flight. - conversation_id: AIConversationId, +} + +type AsyncExecutingActionKey = (AIConversationId, AIAgentActionId); + +#[derive(Default)] +struct AsyncExecutingActions( + std::collections::HashMap, +); + +impl AsyncExecutingActions { + fn insert(&mut self, conversation_id: AIConversationId, running: AsyncExecutingAction) { + self.0 + .insert((conversation_id, running.action.id.clone()), running); + } + + fn get( + &self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> Option<&AsyncExecutingAction> { + self.0.get(&(conversation_id, action_id.clone())) + } + + fn remove( + &mut self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> Option { + self.0.remove(&(conversation_id, action_id.clone())) + } } impl AsyncExecutingAction { @@ -286,10 +314,8 @@ pub struct BlocklistAIActionExecutor { send_message_executor: ModelHandle, ask_user_question_executor: ModelHandle, wait_for_events_executor: ModelHandle, - /// The actions currently executing asynchronously, keyed by action ID. - /// 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, + /// The actions currently executing asynchronously, scoped by conversation and action ID. + async_executing_actions: AsyncExecutingActions, restored_action_ids: HashSet, /// Reference to the terminal model for checking session sharing state. @@ -390,9 +416,13 @@ impl BlocklistAIActionExecutor { } } - pub fn async_executing_action(&self, action_id: &AIAgentActionId) -> Option<&AIAgentAction> { + pub fn async_executing_action( + &self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> Option<&AIAgentAction> { self.async_executing_actions - .get(action_id) + .get(conversation_id, action_id) .map(|running| &running.action) } @@ -408,13 +438,16 @@ impl BlocklistAIActionExecutor { } 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 - && matches!( - running.action.action, - AIAgentActionType::AskUserQuestion { .. } - ) - }) + self.async_executing_actions + .0 + .iter() + .any(|((running_conversation_id, _), running)| { + *running_conversation_id == conversation_id + && matches!( + running.action.action, + AIAgentActionType::AskUserQuestion { .. } + ) + }) } /// Returns the action_id of any running WaitForEvents action for the @@ -424,10 +457,9 @@ impl BlocklistAIActionExecutor { &self, conversation_id: AIConversationId, ) -> Option { - self.async_executing_actions - .iter() - .find_map(|(action_id, running)| { - if running.conversation_id == conversation_id + self.async_executing_actions.0.iter().find_map( + |((running_conversation_id, action_id), running)| { + if *running_conversation_id == conversation_id && matches!( running.action.action, AIAgentActionType::WaitForEvents { .. } @@ -437,7 +469,8 @@ impl BlocklistAIActionExecutor { } else { None } - }) + }, + ) } pub fn shell_command_executor(&self) -> &ModelHandle { @@ -642,8 +675,8 @@ impl BlocklistAIActionExecutor { is_user_initiated: bool, ctx: &mut ModelContext, ) -> TryExecuteResult { - log::info!( - "[tool-debug] try_to_execute_action: action_id={:?}, type={:?}, is_user_initiated={}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: action_id={:?}, type={:?}, is_user_initiated={}", action.id, std::mem::discriminant(&action.action), is_user_initiated @@ -651,7 +684,9 @@ impl BlocklistAIActionExecutor { // We should never actually execute actions in view-only mode. if self.is_shared_session_viewer() { - log::info!("[tool-debug] try_to_execute_action: BLOCKED - shared session viewer mode"); + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: BLOCKED - shared session viewer mode" + ); return TryExecuteResult::NotExecuted { reason: NotExecutedReason::WaitingOnSharer, action: Box::new(action), @@ -664,8 +699,8 @@ impl BlocklistAIActionExecutor { }; let can_auto_execute = self.should_autoexecute(input, ctx); let is_agent_autonomous = AppExecutionMode::as_ref(ctx).is_autonomous(); - log::info!( - "[tool-debug] try_to_execute_action: can_auto_execute={}, is_agent_autonomous={}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: can_auto_execute={}, is_agent_autonomous={}", can_auto_execute, is_agent_autonomous ); @@ -677,8 +712,8 @@ impl BlocklistAIActionExecutor { || can_auto_execute || (is_agent_autonomous && action.action.is_request_command_output())); if needs_confirmation { - log::info!( - "[tool-debug] try_to_execute_action: NEEDS CONFIRMATION - action_id={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: NEEDS CONFIRMATION - action_id={:?}", action.id ); return TryExecuteResult::NotExecuted { @@ -713,8 +748,8 @@ impl BlocklistAIActionExecutor { } } - log::info!( - "[tool-debug] try_to_execute_action: EXECUTING action_id={:?}, type={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: EXECUTING action_id={:?}, type={:?}", action.id, std::mem::discriminant(&action.action) ); @@ -870,8 +905,8 @@ impl BlocklistAIActionExecutor { }; let action_id = action_clone.id.clone(); - log::info!( - "[tool-debug] try_to_execute_action: execution result type={:?} for action_id={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: execution result type={:?} for action_id={:?}", match &execution { AnyActionExecution::NotReady => "NotReady", AnyActionExecution::InvalidAction => "InvalidAction", @@ -882,8 +917,8 @@ impl BlocklistAIActionExecutor { ); match execution { AnyActionExecution::NotReady => { - log::info!( - "[tool-debug] try_to_execute_action: NOT READY - action_id={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: NOT READY - action_id={:?}", action_id ); TryExecuteResult::NotExecuted { @@ -893,7 +928,7 @@ impl BlocklistAIActionExecutor { } AnyActionExecution::InvalidAction => { log::error!( - "[tool-debug] try_to_execute_action: INVALID ACTION - action_id={:?}", + "try_to_execute_action: invalid action, action_id={:?}", action_id ); debug_assert!(false, "Tried to execute AIAgentAction with wrong executor."); @@ -907,10 +942,9 @@ impl BlocklistAIActionExecutor { on_complete, } => { self.async_executing_actions.insert( - action_id.clone(), + conversation_id, AsyncExecutingAction { action: action_clone, - conversation_id, }, ); if !is_restored { @@ -919,15 +953,21 @@ impl BlocklistAIActionExecutor { conversation_id, }); } - log::info!("[tool-debug] try_to_execute_action: spawning ASYNC execution for action_id={:?}", action_id); + crate::ai::tool_diagnostics::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 { - log::warn!("[tool-debug] try_to_execute_action: async action completed but not found in executing map, action_id={:?}", action_id); + let Some(running) = me + .async_executing_actions + .remove(conversation_id, &action_id) + else { + log::warn!("try_to_execute_action: async action completed but not found in executing map, conversation_id={conversation_id}, action_id={action_id:?}"); return; }; let result = on_complete(result, ctx); - log::info!( - "[tool-debug] try_to_execute_action: ASYNC action COMPLETED action_id={:?}, result_type={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "try_to_execute_action: ASYNC action COMPLETED action_id={:?}, result_type={:?}", action_id, std::mem::discriminant(&result) ); @@ -937,7 +977,7 @@ impl BlocklistAIActionExecutor { task_id: running.action.task_id, result, }), - conversation_id: running.conversation_id, + conversation_id, cancellation_reason: None, }); }); @@ -981,6 +1021,7 @@ impl BlocklistAIActionExecutor { pub fn cancel_running_async_action( &mut self, + conversation_id: AIConversationId, action_id: &AIAgentActionId, reason: Option, ctx: &mut ModelContext, @@ -989,13 +1030,42 @@ impl BlocklistAIActionExecutor { if self.is_shared_session_viewer() { return; } - if let Some(running) = self.async_executing_actions.remove(action_id) { + if self + .async_executing_actions + .get(conversation_id, action_id) + .is_some_and(|running| { + matches!( + running.action.action, + AIAgentActionType::RequestCommandOutput { .. } + ) + }) + { + let termination_requested = self.shell_command_executor.update(ctx, |executor, ctx| { + executor.cancel_execution(action_id, ctx) + }); + if termination_requested { + // Keep the action in flight until block completion proves the process stopped. + // Its normal async completion will report the actual terminal exit status. + return; + } + } + if let Some(running) = self + .async_executing_actions + .remove(conversation_id, action_id) + { let action_kind = AIAgentActionTypeDiscriminants::from(&running.action.action); log::info!( - "Canceling running async action of type {action_kind:?} action_id={action_id:?}, reason={reason:?}, backtrace=\n{}", - std::backtrace::Backtrace::force_capture() + "Canceling running async action of type {action_kind:?} action_id={action_id:?}, reason={reason:?}" ); - if running.is_shell_command_action() { + if let Some(backtrace) = crate::ai::tool_diagnostics::capture_backtrace() { + log::debug!("Running action cancellation backtrace:\n{backtrace}"); + } + if running.is_shell_command_action() + && !matches!( + running.action.action, + AIAgentActionType::RequestCommandOutput { .. } + ) + { self.shell_command_executor.update(ctx, |executor, ctx| { executor.cancel_execution(&running.action.id, ctx); }); @@ -1007,6 +1077,10 @@ impl BlocklistAIActionExecutor { self.run_agents_executor.update(ctx, |executor, ctx| { executor.cancel_execution(&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); + }); } else if let AIAgentActionType::WaitForEvents { tool_call_id, .. } = &running.action.action { @@ -1023,7 +1097,7 @@ impl BlocklistAIActionExecutor { task_id: running.action.task_id, result: running.action.action.cancelled_result(), }), - conversation_id: running.conversation_id, + conversation_id, cancellation_reason: reason, }); } @@ -1037,13 +1111,14 @@ impl BlocklistAIActionExecutor { ) { let action_ids = self .async_executing_actions + .0 .iter() - .filter_map(|(action_id, running)| { - (running.conversation_id == conversation_id).then_some(action_id.clone()) + .filter_map(|((running_conversation_id, action_id), _)| { + (*running_conversation_id == conversation_id).then_some(action_id.clone()) }) .collect::>(); for action_id in action_ids { - self.cancel_running_async_action(&action_id, reason, ctx); + self.cancel_running_async_action(conversation_id, &action_id, reason, ctx); } } @@ -1493,6 +1568,95 @@ async fn read_file_as_binary(file_path: &std::path::Path) -> Result, Fil async_fs::read(file_path).await.map_err(FileLoadError::from) } +#[cfg(test)] +mod async_executing_action_tests { + use super::*; + use crate::ai::agent::task::TaskId; + + fn action(id: &str, task_id: &str) -> AIAgentAction { + AIAgentAction { + id: AIAgentActionId::from(id.to_owned()), + action: AIAgentActionType::InitProject, + task_id: TaskId::new(task_id.to_owned()), + requires_result: true, + tool_name: Some("init_project".to_owned()), + } + } + + #[test] + fn duplicate_action_ids_can_execute_concurrently_in_different_conversations() { + let first_conversation = AIConversationId::new(); + let second_conversation = AIConversationId::new(); + let duplicate_id = AIAgentActionId::from("duplicate".to_owned()); + let mut running = AsyncExecutingActions::default(); + + running.insert( + first_conversation, + AsyncExecutingAction { + action: action("duplicate", "first-task"), + }, + ); + running.insert( + second_conversation, + AsyncExecutingAction { + action: action("duplicate", "second-task"), + }, + ); + + assert_eq!(running.0.len(), 2); + assert_eq!( + running + .get(first_conversation, &duplicate_id) + .unwrap() + .action + .task_id, + TaskId::new("first-task".to_owned()) + ); + assert_eq!( + running + .get(second_conversation, &duplicate_id) + .unwrap() + .action + .task_id, + TaskId::new("second-task".to_owned()) + ); + } + + #[test] + fn duplicate_action_completion_and_cancellation_remove_only_the_matching_conversation() { + let first_conversation = AIConversationId::new(); + let second_conversation = AIConversationId::new(); + let duplicate_id = AIAgentActionId::from("duplicate".to_owned()); + let mut running = AsyncExecutingActions::default(); + running.insert( + first_conversation, + AsyncExecutingAction { + action: action("duplicate", "first-task"), + }, + ); + running.insert( + second_conversation, + AsyncExecutingAction { + action: action("duplicate", "second-task"), + }, + ); + + let completed = running.remove(first_conversation, &duplicate_id).unwrap(); + assert_eq!( + completed.action.task_id, + TaskId::new("first-task".to_owned()) + ); + assert!(running.get(second_conversation, &duplicate_id).is_some()); + + let cancelled = running.remove(second_conversation, &duplicate_id).unwrap(); + assert_eq!( + cancelled.action.task_id, + TaskId::new("second-task".to_owned()) + ); + assert!(running.0.is_empty()); + } +} + #[cfg(all(test, feature = "local_fs"))] #[path = "execute_tests.rs"] mod tests; diff --git a/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs b/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs index 7b224695..dd98f9db 100644 --- a/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs +++ b/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs @@ -85,7 +85,7 @@ impl CallMCPToolExecutor { #[cfg(not(target_family = "wasm"))] { - log::info!("[tool-debug] CallMCPToolExecutor::execute called"); + crate::ai::tool_diagnostics::tool_debug!("CallMCPToolExecutor::execute called"); let server_output_id = get_server_output_id(input.conversation_id, ctx); let AIAgentAction { action: @@ -97,21 +97,21 @@ impl CallMCPToolExecutor { .. } = input.action else { - log::error!("[tool-debug] CallMCPToolExecutor::execute: action type mismatch!"); + log::error!("CallMCPToolExecutor::execute: action type mismatch"); return ActionExecution::InvalidAction; }; let name_owned = name.to_owned(); let name_clone = name_owned.clone(); - log::info!( - "[tool-debug] CallMCPToolExecutor: tool_name={}, server_id={:?}, input={}", + crate::ai::tool_diagnostics::tool_debug!( + "CallMCPToolExecutor: tool_name={}, server_id={:?}, input={}", name, server_id, serde_json::to_string(input).unwrap_or_else(|_| "".to_string()) ); let serde_json::Value::Object(mut arguments) = input.clone() else { - log::error!("[tool-debug] CallMCPToolExecutor: input is not an object!"); + log::error!("CallMCPToolExecutor: input is not an object"); return ActionExecution::Sync(AIAgentActionResultType::CallMCPTool( CallMCPToolResult::Error("MCP server tool input not an object".to_owned()), )); @@ -143,15 +143,15 @@ impl CallMCPToolExecutor { let Some(reconnecting_peer) = templatable_peer else { log::error!( - "[tool-debug] CallMCPToolExecutor: MCP server for tool '{}' NOT FOUND", + "CallMCPToolExecutor: MCP server for tool '{}' not found", name_owned ); return ActionExecution::Sync(AIAgentActionResultType::CallMCPTool( CallMCPToolResult::Error("MCP server for tool not found".to_owned()), )); }; - log::info!( - "[tool-debug] CallMCPToolExecutor: found MCP server peer for tool '{}'", + crate::ai::tool_diagnostics::tool_debug!( + "CallMCPToolExecutor: found MCP server peer for tool '{}'", name_owned ); @@ -314,8 +314,8 @@ fn handle_call_tool_result( tool_name: String, ctx: &galaxyui::AppContext, ) -> AIAgentActionResultType { - log::info!( - "[tool-debug] handle_call_tool_result: tool_name={}, is_ok={}", + crate::ai::tool_diagnostics::tool_debug!( + "handle_call_tool_result: tool_name={}, is_ok={}", tool_name, res.is_ok() ); diff --git a/app/src/ai/blocklist/action_model/execute/file_glob.rs b/app/src/ai/blocklist/action_model/execute/file_glob.rs index 079f6749..e201da08 100644 --- a/app/src/ai/blocklist/action_model/execute/file_glob.rs +++ b/app/src/ai/blocklist/action_model/execute/file_glob.rs @@ -108,8 +108,8 @@ impl FileGlobExecutor { else { return ActionExecution::InvalidAction; }; - log::info!( - "[tool-debug] FileGlobExecutor::execute: patterns={:?}, path={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "FileGlobExecutor::execute: patterns={:?}, path={:?}", patterns, path ); diff --git a/app/src/ai/blocklist/action_model/execute/grep.rs b/app/src/ai/blocklist/action_model/execute/grep.rs index 0cd14f07..2d1b4b5a 100644 --- a/app/src/ai/blocklist/action_model/execute/grep.rs +++ b/app/src/ai/blocklist/action_model/execute/grep.rs @@ -237,8 +237,8 @@ impl GrepExecutor { else { return ActionExecution::InvalidAction; }; - log::info!( - "[tool-debug] GrepExecutor::execute: queries={:?}, path={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "GrepExecutor::execute: queries={:?}, path={:?}", queries, path ); diff --git a/app/src/ai/blocklist/action_model/execute/read_files.rs b/app/src/ai/blocklist/action_model/execute/read_files.rs index 8dc917ab..afa5a5c7 100644 --- a/app/src/ai/blocklist/action_model/execute/read_files.rs +++ b/app/src/ai/blocklist/action_model/execute/read_files.rs @@ -91,8 +91,8 @@ impl ReadFilesExecutor { else { return ActionExecution::InvalidAction; }; - log::info!( - "[tool-debug] ReadFilesExecutor::execute: {} files requested", + crate::ai::tool_diagnostics::tool_debug!( + "ReadFilesExecutor::execute: {} files requested", locations.len() ); diff --git a/app/src/ai/blocklist/action_model/execute/request_file_edits.rs b/app/src/ai/blocklist/action_model/execute/request_file_edits.rs index e40c101a..1be790c7 100644 --- a/app/src/ai/blocklist/action_model/execute/request_file_edits.rs +++ b/app/src/ai/blocklist/action_model/execute/request_file_edits.rs @@ -173,14 +173,14 @@ impl RequestFileEditsExecutor { else { return ActionExecution::InvalidAction; }; - log::info!( - "[tool-debug] RequestFileEditsExecutor::execute: action_id={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "RequestFileEditsExecutor::execute: action_id={:?}", id ); let Some(diff_view) = self.diff_views.get(id) else { - log::warn!( - "[tool-debug] RequestFileEditsExecutor: no diff view found for action_id={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "RequestFileEditsExecutor: no diff view found for action_id={:?}", id ); return ActionExecution::NotReady; 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 4d8a8054..e555b951 100644 --- a/app/src/ai/blocklist/action_model/execute/run_agents.rs +++ b/app/src/ai/blocklist/action_model/execute/run_agents.rs @@ -21,7 +21,7 @@ use warpui::{Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; use super::start_agent::{ StartAgentDispatch, StartAgentExecutor, StartAgentExecutorEvent, StartAgentOutcome, - StartAgentWaitPolicy, + StartAgentRequestId, StartAgentWaitPolicy, }; use super::{ child_agent_delegation_denial_reason, ActionExecution, AnyActionExecution, ExecuteActionInput, @@ -146,9 +146,12 @@ impl RunAgentsExecutor { ctx: &mut ModelContext, ) { self.recovery_action_ids.remove(action_id); - self.start_agent_executor.update(ctx, |executor, _| { - executor.cancel_dispatches_for_action(action_id); + let detached_dispatches = self.start_agent_executor.update(ctx, |executor, _| { + executor.cancel_dispatches_for_action(action_id) }); + log::info!( + "RunAgents cancellation detached {detached_dispatches} pending child dispatch(es) for action {action_id}" + ); if self.pending.remove(action_id).is_some() { ctx.emit(RunAgentsExecutorEvent::SpawningFinished { action_id: action_id.clone(), @@ -163,6 +166,22 @@ impl RunAgentsExecutor { ) { for agent in agents { let RunAgentsAgentOutcomeKind::Launched { agent_id } = &agent.kind else { + let RunAgentsAgentOutcomeKind::Completed { agent_id, .. } = &agent.kind else { + continue; + }; + let Some(normalized_name) = normalize_agent_name(&agent.name) else { + continue; + }; + self.launched_agents + .entry(conversation_id) + .or_default() + .insert( + normalized_name, + ExistingLaunchedAgent { + name: agent.name.clone(), + agent_id: agent_id.clone(), + }, + ); continue; }; let Some(normalized_name) = normalize_agent_name(&agent.name) else { @@ -372,6 +391,10 @@ impl RunAgentsExecutor { ); let mut slots: Vec = Vec::with_capacity(agent_run_configs.len()); + let wait_policy = match &run_execution_mode { + RunAgentsExecutionMode::Local => StartAgentWaitPolicy::Completion, + RunAgentsExecutionMode::Remote { .. } => StartAgentWaitPolicy::Startup, + }; for cfg in &agent_run_configs { let normalized_name = normalize_agent_name(&cfg.name) .expect("validated RunAgents requests have non-empty agent names"); @@ -382,7 +405,7 @@ impl RunAgentsExecutor { cfg.name.clone(), parent_conversation_id, child_conversation_id, - parent_run_id.clone(), + wait_policy, exec_ctx, ) }); @@ -485,9 +508,9 @@ impl RunAgentsExecutor { ctx.spawn( async move { - let outcomes = join_all(slots.into_iter().map(resolve_child_slot)).await; + let resolved_slots = join_all(slots.into_iter().map(resolve_child_slot)).await; #[cfg(not(target_family = "wasm"))] - for (slot_index, kind) in outcomes.iter().enumerate() { + for (slot_index, resolved) in resolved_slots.iter().enumerate() { log::info!( "RunAgents child launch outcome action_id={} parent_conversation_id={} \ agent_name={} slot_index={} outcome={}", @@ -498,21 +521,32 @@ impl RunAgentsExecutor { .map(String::as_str) .unwrap_or(""), slot_index, - run_agents_agent_outcome_kind_label(kind) + run_agents_agent_outcome_kind_label(&resolved.outcome) ); } - outcomes + resolved_slots }, - move |me, outcomes, ctx| { + move |me, resolved_slots, ctx| { if !me.is_pending(&action_id_for_aggr) { return; } + let timed_out_request_ids = resolved_slots + .iter() + .filter_map(|resolved| resolved.timed_out_request_id) + .collect::>(); + if !timed_out_request_ids.is_empty() { + me.start_agent_executor.update(ctx, |executor, _| { + for request_id in timed_out_request_ids { + executor.detach_dispatch(request_id); + } + }); + } let agents: Vec = agent_run_configs_for_result .iter() - .zip(outcomes) - .map(|(cfg, kind)| RunAgentsAgentOutcome { + .zip(resolved_slots) + .map(|(cfg, resolved)| RunAgentsAgentOutcome { name: cfg.name.clone(), - kind, + kind: resolved.outcome, }) .collect(); me.record_launched_agents(parent_conversation_id_for_result, &agents); @@ -526,7 +560,7 @@ impl RunAgentsExecutor { "action_id": action_id_for_aggr.to_string(), "parent_conversation_id": parent_conversation_id_for_result.to_string(), "agent_count": agents.len(), - "launched_count": agents.iter().filter(|agent| matches!(agent.kind, RunAgentsAgentOutcomeKind::Launched { .. })).count(), + "launched_count": agents.iter().filter(|agent| matches!(agent.kind, RunAgentsAgentOutcomeKind::Launched { .. } | RunAgentsAgentOutcomeKind::Completed { .. })).count(), "failed_count": agents.iter().filter(|agent| matches!(agent.kind, RunAgentsAgentOutcomeKind::Failed { .. })).count(), "agents": agents .iter() @@ -536,6 +570,12 @@ impl RunAgentsExecutor { "status": "launched", "agent_id": agent_id.as_str(), }), + RunAgentsAgentOutcomeKind::Completed { agent_id, output } => serde_json::json!({ + "name": agent.name.as_str(), + "status": "completed", + "agent_id": agent_id.as_str(), + "output": output, + }), RunAgentsAgentOutcomeKind::Failed { error } => serde_json::json!({ "name": agent.name.as_str(), "status": "failed", @@ -723,6 +763,7 @@ fn start_agent_execution_mode_label(mode: &StartAgentExecutionMode) -> &'static fn run_agents_agent_outcome_kind_label(kind: &RunAgentsAgentOutcomeKind) -> &'static str { match kind { RunAgentsAgentOutcomeKind::Launched { .. } => "launched", + RunAgentsAgentOutcomeKind::Completed { .. } => "completed", RunAgentsAgentOutcomeKind::Failed { .. } => "failed", } } @@ -732,18 +773,30 @@ enum ChildSlot { Pending(StartAgentDispatch), } -async fn resolve_child_slot(slot: ChildSlot) -> RunAgentsAgentOutcomeKind { +#[derive(Debug)] +struct ResolvedChildSlot { + outcome: RunAgentsAgentOutcomeKind, + timed_out_request_id: Option, +} + +async fn resolve_child_slot(slot: ChildSlot) -> ResolvedChildSlot { resolve_child_slot_with_timeout(slot, SPAWN_TIMEOUT).await } async fn resolve_child_slot_with_timeout( slot: ChildSlot, spawn_timeout: Duration, -) -> RunAgentsAgentOutcomeKind { +) -> ResolvedChildSlot { let dispatch = match slot { - ChildSlot::Failed(error) => return RunAgentsAgentOutcomeKind::Failed { error }, + ChildSlot::Failed(error) => { + return ResolvedChildSlot { + outcome: RunAgentsAgentOutcomeKind::Failed { error }, + timed_out_request_id: None, + }; + } ChildSlot::Pending(dispatch) => dispatch, }; + let request_id = dispatch.request_id; let outcome = match dispatch.wait_policy { StartAgentWaitPolicy::Completion => dispatch.receiver.recv().await.ok(), @@ -754,31 +807,41 @@ async fn resolve_child_slot_with_timeout( { futures::future::Either::Left((outcome, _)) => outcome.ok(), futures::future::Either::Right((_, _)) => { + dispatch.mark_detached(); log::warn!( "Agent spawn timed out after {} seconds", spawn_timeout.as_secs() ); - return RunAgentsAgentOutcomeKind::Failed { - error: format!( - "Agent failed to start within {} seconds. \ - The harness binary may not be installed.", - spawn_timeout.as_secs() - ), + return ResolvedChildSlot { + outcome: RunAgentsAgentOutcomeKind::Failed { + error: format!( + "Agent failed to start within {} seconds. \ + The harness binary may not be installed.", + spawn_timeout.as_secs() + ), + }, + timed_out_request_id: Some(request_id), }; } } } }; - match outcome { - Some(StartAgentOutcome::Started { agent_id }) - | Some(StartAgentOutcome::Completed { agent_id, .. }) => { + let outcome = match outcome { + Some(StartAgentOutcome::Started { agent_id }) => { RunAgentsAgentOutcomeKind::Launched { agent_id } } + Some(StartAgentOutcome::Completed { agent_id, output }) => { + RunAgentsAgentOutcomeKind::Completed { agent_id, output } + } Some(StartAgentOutcome::Error(error)) => RunAgentsAgentOutcomeKind::Failed { error }, None => RunAgentsAgentOutcomeKind::Failed { error: "Child agent was cancelled before completion".to_string(), }, + }; + ResolvedChildSlot { + outcome, + timed_out_request_id: None, } } @@ -967,6 +1030,19 @@ fn existing_launched_agents_for_conversation( }; for agent in agents { let RunAgentsAgentOutcomeKind::Launched { agent_id } = &agent.kind else { + let RunAgentsAgentOutcomeKind::Completed { agent_id, .. } = &agent.kind + else { + continue; + }; + let Some(normalized_name) = normalize_agent_name(&agent.name) else { + continue; + }; + existing_agents.entry(normalized_name).or_insert_with(|| { + ExistingLaunchedAgent { + name: agent.name.clone(), + agent_id: agent_id.clone(), + } + }); continue; }; let Some(normalized_name) = normalize_agent_name(&agent.name) else { 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 39a32f1b..a4802f20 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 @@ -1,4 +1,6 @@ use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; use ai::agent::action::{RunAgentsAgentRunConfig, RunAgentsExecutionMode, RunAgentsRequest}; use ai::agent::orchestration_config::{ @@ -525,12 +527,16 @@ fn completion_slots_are_polled_concurrently_and_preserve_request_order() { let (second_sender, second_receiver) = async_channel::bounded(1); let slots = vec![ ChildSlot::Pending(StartAgentDispatch { + request_id: StartAgentRequestId::from_raw_for_test(1), receiver: first_receiver, wait_policy: StartAgentWaitPolicy::Completion, + detached: Arc::new(AtomicBool::new(false)), }), ChildSlot::Pending(StartAgentDispatch { + request_id: StartAgentRequestId::from_raw_for_test(2), receiver: second_receiver, wait_policy: StartAgentWaitPolicy::Completion, + detached: Arc::new(AtomicBool::new(false)), }), ChildSlot::Failed("prelaunch failure".to_string()), ]; @@ -557,15 +563,16 @@ fn completion_slots_are_polled_concurrently_and_preserve_request_order() { let outcomes = outcomes.await; assert!(matches!( - &outcomes[0], + &outcomes[0].outcome, RunAgentsAgentOutcomeKind::Failed { error } if error == "first failed" )); assert!(matches!( - &outcomes[1], - RunAgentsAgentOutcomeKind::Launched { agent_id } if agent_id == "second-agent" + &outcomes[1].outcome, + RunAgentsAgentOutcomeKind::Completed { agent_id, output } + if agent_id == "second-agent" && output == "done" )); assert!(matches!( - &outcomes[2], + &outcomes[2].outcome, RunAgentsAgentOutcomeKind::Failed { error } if error == "prelaunch failure" )); }); @@ -577,8 +584,10 @@ fn completion_wait_ignores_spawn_timeout() { let (sender, receiver) = async_channel::bounded(1); let completion = Box::pin(resolve_child_slot_with_timeout( ChildSlot::Pending(StartAgentDispatch { + request_id: StartAgentRequestId::from_raw_for_test(1), receiver, wait_policy: StartAgentWaitPolicy::Completion, + detached: Arc::new(AtomicBool::new(false)), }), Duration::from_millis(1), )); @@ -598,8 +607,9 @@ fn completion_wait_ignores_spawn_timeout() { .unwrap(); assert!(matches!( - completion.await, - RunAgentsAgentOutcomeKind::Launched { agent_id } if agent_id == "child-agent" + completion.await.outcome, + RunAgentsAgentOutcomeKind::Completed { agent_id, output } + if agent_id == "child-agent" && output == "done" )); }); } @@ -610,21 +620,73 @@ fn startup_wait_retains_spawn_timeout() { let (_sender, receiver) = async_channel::bounded(1); let outcome = resolve_child_slot_with_timeout( ChildSlot::Pending(StartAgentDispatch { + request_id: StartAgentRequestId::from_raw_for_test(1), receiver, wait_policy: StartAgentWaitPolicy::Startup, + detached: Arc::new(AtomicBool::new(false)), }), Duration::from_millis(1), ) .await; + assert!(outcome.timed_out_request_id.is_some()); + assert_eq!( + outcome.timed_out_request_id, + Some(StartAgentRequestId::from_raw_for_test(1)) + ); assert!(matches!( - outcome, + outcome.outcome, RunAgentsAgentOutcomeKind::Failed { error } if error.contains("Agent failed to start within") )); }); } +#[test] +fn startup_timeout_detaches_exact_pending_request() { + App::test((), |mut app| async move { + let state = initialize_run_agents_test(&mut app, ExecutionMode::App); + let start_agent_executor = state.start_agent_executor; + let parent_conversation_id = state.conversation_id; + let dispatch = start_agent_executor.update(&mut app, |executor, ctx| { + executor.dispatch( + AIAgentActionId::from("run-agents-timeout".to_string()), + "child".to_string(), + "work".to_string(), + StartAgentExecutionMode::Remote { + environment_id: "environment".to_string(), + skill_references: Vec::new(), + model_id: "model".to_string(), + computer_use_enabled: false, + worker_host: String::new(), + harness_type: "oz".to_string(), + title: String::new(), + auth_secret_name: None, + }, + None, + parent_conversation_id, + Some("parent-run".to_string()), + ctx, + ) + }); + let request_id = dispatch.request_id; + + let resolved = + resolve_child_slot_with_timeout(ChildSlot::Pending(dispatch), Duration::from_millis(1)) + .await; + let timed_out_request_id = resolved + .timed_out_request_id + .expect("startup timeout should expose request identity"); + start_agent_executor.update(&mut app, |executor, _| { + assert!(executor.detach_dispatch(timed_out_request_id)); + }); + + start_agent_executor.read(&app, |executor, _| { + assert!(!executor.has_pending_dispatch_for_test(request_id)); + }); + }); +} + fn initialize_run_agents_test(app: &mut App, mode: ExecutionMode) -> RunAgentsTestState { initialize_settings_for_tests_with_mode(app, mode, false); let global_resource_handles = GlobalResourceHandles::mock(app); diff --git a/app/src/ai/blocklist/action_model/execute/shell_command.rs b/app/src/ai/blocklist/action_model/execute/shell_command.rs index 084ebb49..1a312726 100644 --- a/app/src/ai/blocklist/action_model/execute/shell_command.rs +++ b/app/src/ai/blocklist/action_model/execute/shell_command.rs @@ -13,7 +13,6 @@ use galaxy_core::execution_mode::AppExecutionMode; use galaxy_util::path::ShellFamily; use galaxyui::r#async::{Spawnable, Timer}; use galaxyui::{Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; -use itertools::Itertools; use parking_lot::FairMutex; use super::{ActionExecution, AnyActionExecution, ExecuteActionInput, PreprocessActionInput}; @@ -37,11 +36,11 @@ use crate::{send_telemetry_from_ctx, TelemetryEvent}; pub struct ShellCommandExecutor { active_session: ModelHandle, - block_finished_senders: HashMap>, + block_finished_senders: HashMap>>, /// Senders used by `Check now` and the automatic monitor watchdog to force a long-running /// shell command's pending poll future to resolve immediately with a fresh snapshot, /// bypassing the agent-set timeout. - force_refresh_senders: HashMap>, + force_refresh_senders: HashMap>>, terminal_model: Arc>, terminal_view_id: EntityId, /// Sender to notify when user hands control back to agent after TransferShellCommandControlToUser. @@ -80,24 +79,39 @@ impl ShellCommandExecutor { event: &ModelEvent, _ctx: &mut ModelContext, ) { - // We wait for precmd for the block _after_ the requested command's block so that - // downstream checks for current working directory are fresh. The precmd hook is when - // the shell relays current working directory to warp. - if let ModelEvent::BlockMetadataReceived(BlockMetadataReceivedEvent { .. }) = event { + // Precmd provides fresh CWD metadata, while BlockCompleted is definitive completion + // evidence for shells that never deliver a subsequent precmd. + if matches!( + event, + ModelEvent::BlockMetadataReceived(BlockMetadataReceivedEvent { .. }) + | ModelEvent::BlockCompleted(_) + ) { let model = self.terminal_model.lock(); - let block_finished_senders = self.block_finished_senders.drain().collect_vec(); - for (block_selector, block_finished_tx) in block_finished_senders.into_iter() { - if let Some(block) = block_selector.get_block(&model) { - if block.is_command_finished() { + let block_finished_senders = self.block_finished_senders.drain().collect::>(); + for (block_selector, block_finished_txs) in block_finished_senders { + let completed_block = block_selector.get_block(&model).filter(|block| { + block.is_command_finished() + && match event { + ModelEvent::BlockCompleted(completed) => { + block.id() == &completed.block_id + } + ModelEvent::BlockMetadataReceived(_) => true, + _ => false, + } + }); + if completed_block.is_some() { + for block_finished_tx in block_finished_txs { if let Err(e) = block_finished_tx.send(()) { log::warn!( "Failed to notify block completion for running requested command: {e:?}" ) } - } else { - self.block_finished_senders - .insert(block_selector, block_finished_tx); } + } else { + // The requested-command association may not exist yet. Keep all waiters until + // this selector resolves and its block actually completes, or it is cancelled. + self.block_finished_senders + .insert(block_selector, block_finished_txs); } } } @@ -195,8 +209,8 @@ impl ShellCommandExecutor { input: ExecuteActionInput, ctx: &mut ModelContext, ) -> impl Into { - log::info!( - "[tool-debug] ShellCommandExecutor::execute: action_type={:?}", + crate::ai::tool_diagnostics::tool_debug!( + "ShellCommandExecutor::execute: action_type={:?}", std::mem::discriminant(&input.action.action) ); let model = self.terminal_model.lock(); @@ -204,12 +218,6 @@ impl ShellCommandExecutor { // Determine the action we want to take based on the input. let action_id = input.action.id.clone(); - let command = model - .block_list() - .active_block() - .command_with_secrets_unobfuscated(false) - .clone(); - let handle = ctx.handle(); match &input.action.action { AIAgentActionType::RequestCommandOutput { @@ -222,18 +230,13 @@ impl ShellCommandExecutor { .active_block() .is_active_and_long_running() { - // Another command is still running (e.g. stuck in a pager). Return an error - // result so the model receives feedback and can adapt. Using Completed with a - // non-zero exit code ensures a follow-up request is triggered. - return ActionExecution::Sync(AIAgentActionResultType::RequestCommandOutput( - RequestCommandOutputResult::Completed { - command: command.clone(), - block_id: model.block_list().active_block().id().clone(), - output: "Error: Cannot execute command because another command is still running in the terminal.".to_string(), - exit_code: ExitCode::from(1), - start_ts: None, - completed_ts: None, - }, + let running_command = model + .block_list() + .active_block() + .command_with_secrets_unobfuscated(false); + return ActionExecution::Sync(terminal_busy_execution_error( + command, + &running_command, )); } // If another conversation has taken over the agent view since this command @@ -275,8 +278,7 @@ impl ShellCommandExecutor { // Remove the senders from the maps. if let Some(handle) = handle.upgrade(ctx) { handle.update(ctx, |me, _| { - me.block_finished_senders.remove(&block_selector); - me.force_refresh_senders.remove(&block_selector); + me.prune_closed_senders(&block_selector); }); } @@ -339,8 +341,7 @@ impl ShellCommandExecutor { // Remove the senders from the maps. if let Some(handle) = handle.upgrade(ctx) { handle.update(ctx, |me, _| { - me.block_finished_senders.remove(&block_selector); - me.force_refresh_senders.remove(&block_selector); + me.prune_closed_senders(&block_selector); }); } @@ -371,6 +372,7 @@ impl ShellCommandExecutor { }, )); } + let command = block.command_with_secrets_unobfuscated(false); drop(model); let block_selector = BlockSelector::Id(block_id.clone()); @@ -380,8 +382,7 @@ impl ShellCommandExecutor { // Remove the senders from the maps. if let Some(handle) = handle.upgrade(ctx) { handle.update(ctx, |me, _| { - me.block_finished_senders.remove(&block_selector); - me.force_refresh_senders.remove(&block_selector); + me.prune_closed_senders(&block_selector); }); } @@ -419,7 +420,9 @@ impl ShellCommandExecutor { // Set up a future to also wait for block completion. let (block_finished_tx, block_finished_rx) = oneshot::channel(); self.block_finished_senders - .insert(block_selector.clone(), block_finished_tx); + .entry(block_selector.clone()) + .or_default() + .push(block_finished_tx); // Build the future that captures terminal model and block data. let transfer_future = { @@ -491,7 +494,7 @@ impl ShellCommandExecutor { // Clean up. if let Some(handle) = handle.upgrade(ctx) { handle.update(ctx, |me, _| { - me.block_finished_senders.remove(&block_selector); + me.prune_closed_senders(&block_selector); me.control_handback_sender = None; }); } @@ -520,13 +523,17 @@ impl ShellCommandExecutor { // Create a channel to notify us when we receive block metadata. let (block_metadata_received_tx, block_metadata_received_rx) = oneshot::channel(); self.block_finished_senders - .insert(block_selector.clone(), block_metadata_received_tx); + .entry(block_selector.clone()) + .or_default() + .push(block_metadata_received_tx); // Create a channel so `Check now` or the automatic monitor watchdog can short-circuit // the timeout and deliver the agent a fresh snapshot immediately. let (force_refresh_tx, force_refresh_rx) = oneshot::channel(); self.force_refresh_senders - .insert(block_selector.clone(), force_refresh_tx); + .entry(block_selector.clone()) + .or_default() + .push(force_refresh_tx); // Create a future that resolves when we should send a result to the agent. let terminal_model = self.terminal_model.clone(); @@ -600,7 +607,12 @@ impl ShellCommandExecutor { completed_ts: block.completed_ts().cloned(), } } else { - let grid_contents = if model.is_alt_screen_active() { + let selected_block_owns_alt_screen = selected_block_owns_alt_screen( + model.is_alt_screen_active(), + model.active_block_id(), + block.id(), + ); + let grid_contents = if selected_block_owns_alt_screen { formatted_terminal_contents_for_input( model.alt_screen().grid_handler(), None, @@ -618,7 +630,7 @@ impl ShellCommandExecutor { block_id: block.id().clone(), grid_contents, cursor: CURSOR_MARKER, - is_alt_screen_active: model.is_alt_screen_active(), + is_alt_screen_active: selected_block_owns_alt_screen, is_preempted, } } @@ -630,23 +642,50 @@ impl ShellCommandExecutor { } } - pub(super) fn cancel_execution(&mut self, id: &AIAgentActionId, _ctx: &mut ModelContext) { + pub(super) fn cancel_execution( + &mut self, + id: &AIAgentActionId, + ctx: &mut ModelContext, + ) -> bool { let terminal_model = self.terminal_model.lock(); - let active_block = terminal_model.block_list().active_block(); - if !active_block.is_active_and_long_running() { - return; - } - - let selector = if active_block - .requested_command_action_id() - .is_some_and(|requested_command_id| requested_command_id == id) - { - BlockSelector::RequestedCommandId(id.clone()) + let requested_selector = BlockSelector::RequestedCommandId(id.clone()); + let requested_block_is_running = requested_selector + .get_block(&terminal_model) + .is_some_and(|block| block.is_active_and_long_running() && !block.finished()); + let selector = if requested_block_is_running { + requested_selector } else { - BlockSelector::Id(active_block.id().clone()) + BlockSelector::Id(terminal_model.active_block_id().clone()) }; - self.block_finished_senders.remove(&selector); - self.force_refresh_senders.remove(&selector); + // Cancelling the wait future alone would report cancellation while the process keeps + // running. Terminate the exact requested command before resolving the action as cancelled. + if requested_block_is_running { + ctx.emit(ShellCommandExecutorEvent::CancelExecution { + action_id: id.clone(), + }); + } + if !requested_block_is_running { + self.block_finished_senders.remove(&selector); + self.force_refresh_senders.remove(&selector); + } + requested_block_is_running + } + + fn prune_closed_senders(&mut self, selector: &BlockSelector) { + Self::prune_closed_sender_group(&mut self.block_finished_senders, selector); + Self::prune_closed_sender_group(&mut self.force_refresh_senders, selector); + } + + fn prune_closed_sender_group( + senders: &mut HashMap>>, + selector: &BlockSelector, + ) { + if let Some(selector_senders) = senders.get_mut(selector) { + selector_senders.retain(|sender| !sender.is_canceled()); + if selector_senders.is_empty() { + senders.remove(selector); + } + } } /// Force any in-flight poll for the given long-running command block to resolve @@ -657,9 +696,8 @@ impl ShellCommandExecutor { /// control to the user). Returns whether a matching poll was successfully refreshed. pub fn force_refresh_block(&mut self, block_id: &BlockId) -> bool { let terminal_model = self.terminal_model.lock(); - // Find a sender whose selector resolves to this block. In practice there is at - // most one: a given block can have at most one in-flight `action_result_future` - // at a time. + // Find every pending poll whose selector resolves to this block. Multiple provider polls + // may legitimately wait on the same command and must be refreshed together. let matching_selector = self .force_refresh_senders .keys() @@ -674,8 +712,12 @@ impl ShellCommandExecutor { drop(terminal_model); if let Some(selector) = matching_selector { - if let Some(sender) = self.force_refresh_senders.remove(&selector) { - return sender.send(()).is_ok(); + if let Some(senders) = self.force_refresh_senders.remove(&selector) { + let mut refreshed = false; + for sender in senders { + refreshed |= sender.send(()).is_ok(); + } + return refreshed; } } false @@ -714,6 +756,21 @@ fn command_for_execution( } } +fn terminal_busy_execution_error(command: &str, running_command: &str) -> AIAgentActionResultType { + AIAgentActionResultType::RequestCommandOutput(RequestCommandOutputResult::ExecutionError { + command: command.to_string(), + message: format!("terminal is busy running command '{running_command}'"), + }) +} + +fn selected_block_owns_alt_screen( + is_alt_screen_active: bool, + active_block_id: &BlockId, + selected_block_id: &BlockId, +) -> bool { + is_alt_screen_active && active_block_id == selected_block_id +} + #[derive(Debug, Clone, Hash, PartialEq, Eq)] enum BlockSelector { Id(BlockId), @@ -919,7 +976,9 @@ pub enum ShellCommandExecutorEvent { input: Bytes, mode: AIAgentPtyWriteMode, }, - CancelExecution, + CancelExecution { + action_id: AIAgentActionId, + }, /// Emitted when the agent requests to transfer control of a long-running command to the user. TransferControlToUser { action_id: AIAgentActionId, diff --git a/app/src/ai/blocklist/action_model/execute/shell_command_tests.rs b/app/src/ai/blocklist/action_model/execute/shell_command_tests.rs index 8530111c..c1439600 100644 --- a/app/src/ai/blocklist/action_model/execute/shell_command_tests.rs +++ b/app/src/ai/blocklist/action_model/execute/shell_command_tests.rs @@ -1,19 +1,29 @@ use std::sync::Arc; +use std::task::Poll; use async_channel::unbounded; use futures::channel::oneshot; +use futures::{pin_mut, poll}; use parking_lot::FairMutex; use warpui::{App, EntityId}; -use super::{command_for_execution, ActionResult, BlockSelector, ShellCommandExecutor}; -use crate::ai::agent::ShellCommandDelay; -use crate::terminal::event::{BlockMetadataReceivedEvent, BlockWorkingDirectoryUpdatedEvent}; +use super::{ + command_for_execution, selected_block_owns_alt_screen, terminal_busy_execution_error, + ActionResult, BlockSelector, ShellCommandExecutor, +}; +use crate::ai::agent::{ + AIAgentActionId, AIAgentActionResultType, RequestCommandOutputResult, ShellCommandDelay, +}; +use crate::terminal::event::{ + BlockCompletedEvent, BlockMetadataReceivedEvent, BlockType, BlockWorkingDirectoryUpdatedEvent, +}; use crate::terminal::model::block::{BlockId, BlockMetadata}; use crate::terminal::model::session::active_session::ActiveSession; use crate::terminal::model::session::Sessions; use crate::terminal::model::terminal_model::{BlockIndex, TerminalModel}; use crate::terminal::model_events::{ModelEvent, ModelEventDispatcher}; use crate::terminal::shell::ShellType; +use crate::AIConversationId; #[test] fn wait_commands_disable_implicit_posix_pagers_without_masking_exit_status() { @@ -31,6 +41,41 @@ fn wait_commands_disable_implicit_posix_pagers_without_masking_exit_status() { ); } +#[test] +fn terminal_busy_is_an_execution_error_for_the_unstarted_command() { + let result = terminal_busy_execution_error("cargo test", "sleep 120"); + + assert!(matches!( + result, + AIAgentActionResultType::RequestCommandOutput( + RequestCommandOutputResult::ExecutionError { command, message } + ) if command == "cargo test" + && message == "terminal is busy running command 'sleep 120'" + )); +} + +#[test] +fn targeted_poll_uses_alt_screen_only_for_its_owning_block() { + let active_block_id = BlockId::new(); + let selected_block_id = BlockId::new(); + + assert!(!selected_block_owns_alt_screen( + true, + &active_block_id, + &selected_block_id + )); + assert!(selected_block_owns_alt_screen( + true, + &active_block_id, + &active_block_id + )); + assert!(!selected_block_owns_alt_screen( + false, + &active_block_id, + &active_block_id + )); +} + /// Locks in the contract that `ShellCommandExecutor`'s requested-command finish /// detector reacts only to `BlockMetadataReceived` (precmd) and not to /// `BlockWorkingDirectoryUpdated` (OSC 7). The detector relies on @@ -63,7 +108,7 @@ fn block_working_directory_updated_does_not_drain_finish_senders() { let selector = BlockSelector::Id(block_id); let (tx, _rx) = oneshot::channel::<()>(); executor.update(&mut app, |executor, _ctx| { - executor.block_finished_senders.insert(selector, tx); + executor.block_finished_senders.insert(selector, vec![tx]); }); assert_eq!( app.read(|ctx| executor.as_ref(ctx).block_finished_senders.len()), @@ -88,8 +133,7 @@ fn block_working_directory_updated_does_not_drain_finish_senders() { that map is reserved for precmd (BlockMetadataReceived)" ); - // Precmd event — the senders map should be drained (and since the - // block isn't in the terminal model, the sender is dropped). + // An unrelated precmd cannot resolve this selector, so its waiter must survive. model_event_dispatcher.update(&mut app, |_dispatcher, ctx| { ctx.emit(ModelEvent::BlockMetadataReceived( BlockMetadataReceivedEvent { @@ -102,8 +146,8 @@ fn block_working_directory_updated_does_not_drain_finish_senders() { }); assert_eq!( app.read(|ctx| executor.as_ref(ctx).block_finished_senders.len()), - 0, - "BlockMetadataReceived should drain the finish senders" + 1, + "BlockMetadataReceived must retain unresolved finish senders" ); }); } @@ -138,7 +182,7 @@ fn force_refresh_block_reports_and_resolves_matching_poll() { executor.update(&mut app, |executor, _| { executor .force_refresh_senders - .insert(BlockSelector::Id(block_id.clone()), tx); + .insert(BlockSelector::Id(block_id.clone()), vec![tx]); assert!(executor.force_refresh_block(&block_id)); assert!(!executor.force_refresh_block(&block_id)); }); @@ -149,7 +193,7 @@ fn force_refresh_block_reports_and_resolves_matching_poll() { executor.update(&mut app, |executor, _| { executor .force_refresh_senders - .insert(BlockSelector::Id(block_id.clone()), tx); + .insert(BlockSelector::Id(block_id.clone()), vec![tx]); }); terminal_model.lock().finish_block(); assert!(executor.update(&mut app, |executor, _| { @@ -158,6 +202,150 @@ fn force_refresh_block_reports_and_resolves_matching_poll() { }); } +#[test] +fn requested_command_waiter_survives_early_metadata_and_resolves_after_association() { + App::test((), |mut app| async move { + let terminal_view_id = EntityId::new(); + let sessions = app.add_model(|_| Sessions::new_for_test()); + let (_model_events_tx, model_events_rx) = unbounded(); + let model_event_dispatcher = + app.add_model(|ctx| ModelEventDispatcher::new(model_events_rx, sessions.clone(), ctx)); + let active_session = app.add_model(|ctx| { + ActiveSession::new(sessions.clone(), model_event_dispatcher.clone(), ctx) + }); + let terminal_model = Arc::new(FairMutex::new(TerminalModel::mock(None, None))); + let executor = app.add_model(|ctx| { + ShellCommandExecutor::new( + active_session, + terminal_model.clone(), + &model_event_dispatcher, + terminal_view_id, + ctx, + ) + }); + let action_id = AIAgentActionId::from("requested-command".to_string()); + let result_future = executor.update(&mut app, |executor, _| { + executor.action_result_future( + BlockSelector::RequestedCommandId(action_id.clone()), + Some(ShellCommandDelay::OnCompletion), + ) + }); + pin_mut!(result_future); + + model_event_dispatcher.update(&mut app, |_dispatcher, ctx| { + ctx.emit(ModelEvent::BlockMetadataReceived( + BlockMetadataReceivedEvent { + block_metadata: BlockMetadata::new(None, Some("/tmp/early".to_string())), + block_index: BlockIndex::zero(), + is_after_in_band_command: false, + is_done_bootstrapping: true, + }, + )); + }); + assert!(matches!(poll!(&mut result_future), Poll::Pending)); + + terminal_model + .lock() + .simulate_long_running_block("printf done", "done"); + let block_id = terminal_model.lock().active_block_id().clone(); + terminal_model + .lock() + .block_list_mut() + .active_block_mut() + .set_agent_interaction_mode_for_requested_command( + action_id, + None, + AIConversationId::new(), + ); + terminal_model.lock().finish_block(); + model_event_dispatcher.update(&mut app, |_dispatcher, ctx| { + ctx.emit(ModelEvent::BlockCompleted(block_completed_event( + block_id.clone(), + ))); + }); + + assert!(matches!( + result_future.await, + ActionResult::CommandFinished { + block_id: result_block_id, + .. + } if result_block_id == block_id + )); + }); +} + +#[test] +fn duplicate_completion_polls_for_same_block_both_resolve_on_block_completed() { + App::test((), |mut app| async move { + let terminal_view_id = EntityId::new(); + let sessions = app.add_model(|_| Sessions::new_for_test()); + let (_model_events_tx, model_events_rx) = unbounded(); + let model_event_dispatcher = + app.add_model(|ctx| ModelEventDispatcher::new(model_events_rx, sessions.clone(), ctx)); + let active_session = app.add_model(|ctx| { + ActiveSession::new(sessions.clone(), model_event_dispatcher.clone(), ctx) + }); + let terminal_model = Arc::new(FairMutex::new(TerminalModel::mock(None, None))); + terminal_model + .lock() + .simulate_long_running_block("sleep 1", "finished"); + let block_id = terminal_model.lock().active_block_id().clone(); + let executor = app.add_model(|ctx| { + ShellCommandExecutor::new( + active_session, + terminal_model.clone(), + &model_event_dispatcher, + terminal_view_id, + ctx, + ) + }); + + let first = executor.update(&mut app, |executor, _| { + executor.action_result_future( + BlockSelector::Id(block_id.clone()), + Some(ShellCommandDelay::OnCompletion), + ) + }); + let second = executor.update(&mut app, |executor, _| { + executor.action_result_future( + BlockSelector::Id(block_id.clone()), + Some(ShellCommandDelay::OnCompletion), + ) + }); + pin_mut!(first); + pin_mut!(second); + assert!(matches!(poll!(&mut first), Poll::Pending)); + assert!(matches!(poll!(&mut second), Poll::Pending)); + + terminal_model.lock().finish_block(); + model_event_dispatcher.update(&mut app, |_dispatcher, ctx| { + ctx.emit(ModelEvent::BlockCompleted(block_completed_event( + block_id.clone(), + ))); + }); + + let first_result = first.await; + let second_result = second.await; + assert!(matches!(first_result, ActionResult::CommandFinished { .. })); + assert!(matches!( + second_result, + ActionResult::CommandFinished { .. } + )); + }); +} + +fn block_completed_event(block_id: BlockId) -> BlockCompletedEvent { + BlockCompletedEvent { + block_latency_data: None, + block_type: BlockType::Restored, + num_secrets_obfuscated: 0, + block_index: BlockIndex::zero(), + block_id, + session_id: None, + restored_block_was_local: None, + } +} + #[test] fn force_refresh_wakes_on_completion_poll_with_preempted_snapshot() { App::test((), |mut 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 5c882498..95a843eb 100644 --- a/app/src/ai/blocklist/action_model/execute/start_agent.rs +++ b/app/src/ai/blocklist/action_model/execute/start_agent.rs @@ -1,4 +1,6 @@ use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; use futures::future::BoxFuture; use futures::FutureExt; @@ -42,9 +44,24 @@ pub enum StartAgentWaitPolicy { Completion, } +fn wait_policy_for_execution_mode(mode: &StartAgentExecutionMode) -> StartAgentWaitPolicy { + match mode { + StartAgentExecutionMode::Local { .. } => StartAgentWaitPolicy::Completion, + StartAgentExecutionMode::Remote { .. } => StartAgentWaitPolicy::Startup, + } +} + pub struct StartAgentDispatch { + pub request_id: StartAgentRequestId, pub receiver: async_channel::Receiver, pub wait_policy: StartAgentWaitPolicy, + pub(super) detached: Arc, +} + +impl StartAgentDispatch { + pub(super) fn mark_detached(&self) { + self.detached.store(true, Ordering::Release); + } } fn invalid_local_child_harness_error(harness_type: &str) -> String { @@ -139,6 +156,7 @@ struct PendingStartAgent { /// Set once the child conversation is synchronously created. child_conversation_id: Option, sender: async_channel::Sender, + detached: Arc, /// Direct Bedrock/OpenAI parents do not have a server run id or an /// orchestration event stream. Keep the tool call open until their local /// child finishes, then return the child's output inline. @@ -176,9 +194,17 @@ impl StartAgentExecutor { ctx: &mut ModelContext, ) { let child_link_event = { - let Some(pending) = self.pending.get_mut(&request_id) else { + let Some(pending) = self.pending.get(&request_id) else { return; }; + if pending.detached.load(Ordering::Acquire) { + self.pending.remove(&request_id); + return; + } + let pending = self + .pending + .get_mut(&request_id) + .expect("pending request was checked above"); pending.child_conversation_id = Some(child_conversation_id); if let Some(agent_name) = pending.run_agents_child_name.clone() { Some(StartAgentExecutorEvent::RunAgentsChildConversationCreated { @@ -595,16 +621,13 @@ impl StartAgentExecutor { } }; - // In local mode (no parent_run_id), block until the child finishes - // so the parent model receives the child's output as the tool result. - let wait_policy = if parent_run_id.is_none() { - StartAgentWaitPolicy::Completion - } else { - StartAgentWaitPolicy::Startup - }; + // Local children return their completed work; remote children acknowledge startup and + // continue through the hosted orchestration lifecycle. + let wait_policy = wait_policy_for_execution_mode(&execution_mode); let (sender, receiver) = async_channel::bounded(1); let request_id = self.next_request_id(); + let detached = Arc::new(AtomicBool::new(false)); self.pending.insert( request_id, PendingStartAgent { @@ -613,6 +636,7 @@ impl StartAgentExecutor { parent_conversation_id, child_conversation_id: None, sender, + detached, wait_policy, }, ); @@ -667,24 +691,23 @@ impl StartAgentExecutor { parent_run_id: Option, ctx: &mut ModelContext, ) -> StartAgentDispatch { - let wait_policy = if parent_run_id.is_none() { - StartAgentWaitPolicy::Completion - } else { - StartAgentWaitPolicy::Startup - }; + let wait_policy = wait_policy_for_execution_mode(&execution_mode); let (sender, receiver) = async_channel::bounded(1); + let request_id = self.next_request_id(); + let detached = Arc::new(AtomicBool::new(false)); if let Some(error) = child_agent_delegation_denial_reason(parent_conversation_id, ctx) { let _ = sender.try_send(StartAgentOutcome::Error(error)); return StartAgentDispatch { + request_id, receiver, wait_policy, + detached, }; } let (prompt, execution_mode) = normalize_legacy_local_child_harness_command(prompt, execution_mode); let prompt = compose_leaf_agent_prompt(&prompt); - let request_id = self.next_request_id(); self.pending.insert( request_id, PendingStartAgent { @@ -693,6 +716,7 @@ impl StartAgentExecutor { parent_conversation_id, child_conversation_id: None, sender, + detached: detached.clone(), wait_policy, }, ); @@ -708,8 +732,10 @@ impl StartAgentExecutor { }, ))); StartAgentDispatch { + request_id, receiver, wait_policy, + detached, } } @@ -719,16 +745,12 @@ impl StartAgentExecutor { name: String, parent_conversation_id: AIConversationId, child_conversation_id: AIConversationId, - parent_run_id: Option, + wait_policy: StartAgentWaitPolicy, 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(); + let detached = Arc::new(AtomicBool::new(false)); self.pending.insert( request_id, PendingStartAgent { @@ -737,19 +759,54 @@ impl StartAgentExecutor { parent_conversation_id, child_conversation_id: Some(child_conversation_id), sender, + detached: detached.clone(), wait_policy, }, ); self.record_child_conversation(request_id, child_conversation_id, ctx); StartAgentDispatch { + request_id, receiver, wait_policy, + detached, } } - pub fn cancel_dispatches_for_action(&mut self, action_id: &AIAgentActionId) { - self.pending - .retain(|_, pending| &pending.action_id != action_id); + /// Detaches one exact dispatch. If its launch callback is already queued, + /// the shared marker prevents that callback from linking a late child. + pub fn detach_dispatch(&mut self, request_id: StartAgentRequestId) -> bool { + let Some(pending) = self.pending.remove(&request_id) else { + return false; + }; + pending.detached.store(true, Ordering::Release); + true + } + + /// Test-only lookup for request ownership without exposing executor internals. + #[cfg(test)] + pub fn has_pending_dispatch_for_test(&self, request_id: StartAgentRequestId) -> bool { + self.pending.contains_key(&request_id) + } + + pub fn cancel_dispatches_for_action(&mut self, action_id: &AIAgentActionId) -> usize { + let request_ids = self + .pending + .iter() + .filter_map(|(request_id, pending)| { + (&pending.action_id == action_id).then_some(*request_id) + }) + .collect::>(); + let detached_count = request_ids.len(); + for request_id in request_ids { + self.detach_dispatch(request_id); + } + detached_count + } + + /// 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 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 312ec533..765ae601 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 @@ -257,6 +257,71 @@ fn dispatch_denies_child_conversation_defense_in_depth() { }); } +#[test] +fn local_execution_waits_for_completion() { + assert_eq!( + wait_policy_for_execution_mode(&StartAgentExecutionMode::local_with_defaults()), + StartAgentWaitPolicy::Completion + ); +} + +#[test] +fn detach_dispatch_rejects_late_child_callback() { + 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 parent_conversation_id = history_model.update(&mut app, |history, ctx| { + history.start_new_conversation(terminal_view_id, false, false, false, ctx) + }); + let dispatch = executor.update(&mut app, |executor, ctx| { + executor.dispatch( + AIAgentActionId::from("run-agents".to_string()), + "child".to_string(), + "work".to_string(), + StartAgentExecutionMode::Remote { + environment_id: "environment".to_string(), + skill_references: Vec::new(), + model_id: "model".to_string(), + computer_use_enabled: false, + worker_host: String::new(), + harness_type: "oz".to_string(), + title: String::new(), + auth_secret_name: None, + }, + None, + parent_conversation_id, + Some(PARENT_RUN_ID.to_string()), + ctx, + ) + }); + assert!(executor.update(&mut app, |executor, _| { + executor.detach_dispatch(dispatch.request_id) + })); + + let child_conversation_id = history_model.update(&mut app, |history, ctx| { + history.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + parent_conversation_id, + None, + ctx, + ) + }); + history_model.update(&mut app, |history, ctx| { + history.record_new_conversation_request_complete( + dispatch.request_id, + child_conversation_id, + ctx, + ); + }); + + executor.read(&app, |executor, _| assert!(executor.pending.is_empty())); + assert!(dispatch.receiver.try_recv().is_err()); + }); +} + #[test] fn legacy_local_codex_command_prompt_normalizes_to_local_harness() { let (prompt, execution_mode) = normalize_legacy_local_child_harness_command( @@ -878,6 +943,69 @@ fn direct_provider_error_preserves_child_for_inspection() { }); } +#[test] +fn cancelling_standalone_start_agent_drops_dispatch_but_keeps_child_running() { + 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 parent_conversation_id = 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 execution = executor.update(&mut app, |executor, ctx| { + executor + .execute( + ExecuteActionInput { + action: &action, + conversation_id: parent_conversation_id, + }, + ctx, + ) + .into() + }); + let child_conversation_id = history_model.update(&mut app, |history, ctx| { + history.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + parent_conversation_id, + None, + ctx, + ) + }); + history_model.update(&mut app, |history, ctx| { + history.record_new_conversation_request_complete( + FIRST_REQUEST_ID, + child_conversation_id, + ctx, + ); + }); + + executor.update(&mut app, |executor, _| { + executor.cancel_execution(&action.id); + }); + + executor.read(&app, |executor, _| assert!(executor.pending.is_empty())); + history_model.read(&app, |history, _| { + assert_eq!( + history + .conversation(&child_conversation_id) + .expect("child should remain in history") + .status(), + &ConversationStatus::InProgress + ); + }); + let AnyActionExecution::Async { execute_future, .. } = execution else { + panic!("expected async StartAgent execution"); + }; + let _ = execute_future.await; + }); +} + #[test] fn removing_direct_provider_child_resolves_pending_wait() { App::test((), |mut app| async move { @@ -975,7 +1103,7 @@ fn reattach_reuses_persisted_child_without_launching_another_agent() { "child".to_string(), parent_conversation_id, child_conversation_id, - None, + StartAgentWaitPolicy::Completion, ctx, ) }); diff --git a/app/src/ai/blocklist/action_model_tests.rs b/app/src/ai/blocklist/action_model_tests.rs index c36fb85d..c6973daa 100644 --- a/app/src/ai/blocklist/action_model_tests.rs +++ b/app/src/ai/blocklist/action_model_tests.rs @@ -346,3 +346,128 @@ fn only_rejecting_a_blocked_action_is_a_permission_denial() { Some(&AIActionStatus::Blocked), )); } + +#[test] +fn duplicate_action_ids_resolve_only_within_the_requested_conversation() { + let first_conversation = AIConversationId::new(); + let second_conversation = AIConversationId::new(); + let duplicate_id = AIAgentActionId::from("duplicate".to_string()); + let first_result = make_action_result("duplicate"); + let mut second_result = action_result("duplicate", AIAgentActionResultType::InitProject); + second_result.task_id = TaskId::new("second-task".to_string()); + let second_result = Arc::new(second_result); + let finished_results = HashMap::from([(first_conversation, vec![first_result.clone()])]); + let provider_results = HashMap::new(); + let archive = HashMap::from([ + ( + (first_conversation, duplicate_id.clone()), + first_result.clone(), + ), + ( + (second_conversation, duplicate_id.clone()), + second_result.clone(), + ), + ]); + + assert!(Arc::ptr_eq( + action_result_for_conversation( + &finished_results, + &provider_results, + &archive, + first_conversation, + &duplicate_id, + ) + .unwrap(), + &first_result, + )); + assert!(Arc::ptr_eq( + action_result_for_conversation( + &finished_results, + &provider_results, + &archive, + second_conversation, + &duplicate_id, + ) + .unwrap(), + &second_result, + )); + assert!(action_result_for_conversation( + &finished_results, + &provider_results, + &archive, + AIConversationId::new(), + &duplicate_id, + ) + .is_none()); +} + +#[test] +fn cancellation_permission_inference_uses_the_matching_conversation_status() { + let blocked_conversation = AIConversationId::new(); + let queued_conversation = AIConversationId::new(); + let duplicate_id = AIAgentActionId::from("duplicate".to_string()); + let pending_actions = HashMap::from([ + (blocked_conversation, VecDeque::from([action("duplicate")])), + ( + queued_conversation, + VecDeque::from([action("first"), action("duplicate")]), + ), + ]); + let running_actions = HashMap::new(); + + let blocked_status = pending_action_status( + &pending_actions, + &running_actions, + blocked_conversation, + &duplicate_id, + false, + ); + let queued_status = pending_action_status( + &pending_actions, + &running_actions, + queued_conversation, + &duplicate_id, + false, + ); + + assert!(is_permission_denial( + CancellationReason::ManuallyCancelled, + blocked_status.as_ref(), + )); + assert!(!is_permission_denial( + CancellationReason::ManuallyCancelled, + queued_status.as_ref(), + )); +} + +#[test] +fn action_lifecycle_events_disambiguate_duplicate_ids_by_conversation() { + let first_conversation = AIConversationId::new(); + let second_conversation = AIConversationId::new(); + let duplicate_id = AIAgentActionId::from("duplicate".to_owned()); + let events = [ + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { + action_id: duplicate_id.clone(), + conversation_id: first_conversation, + execution_ref: None, + }, + BlocklistAIActionEvent::ExecutingAction { + action_id: duplicate_id.clone(), + conversation_id: second_conversation, + execution_ref: None, + }, + BlocklistAIActionEvent::FinishedAction { + action_id: duplicate_id.clone(), + conversation_id: first_conversation, + cancellation_reason: None, + execution_ref: None, + }, + ]; + + assert_eq!(events[0].conversation_id(), Some(first_conversation)); + assert_eq!(events[1].conversation_id(), Some(second_conversation)); + assert_eq!(events[2].conversation_id(), Some(first_conversation)); + assert!(events + .iter() + .all(|event| event.action_id() == &duplicate_id)); +} diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index ba84d860..f2e40dbd 100644 --- a/app/src/ai/blocklist/block.rs +++ b/app/src/ai/blocklist/block.rs @@ -3578,7 +3578,10 @@ impl AIBlock { } // Set the state based on the action status from the action model - let action_status = self.action_model.as_ref(ctx).get_action_status(action_id); + let action_status = self + .action_model + .as_ref(ctx) + .get_action_status(self.client_ids.conversation_id, action_id); let is_reverted = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&self.client_ids.conversation_id) @@ -3673,6 +3676,7 @@ impl AIBlock { RequestedCommandViewEvent::Accepted => { self.action_model.update(ctx, |action_model, ctx| { action_model.handle_requested_command_accepted( + self.client_ids.conversation_id, action_id, view.as_ref(ctx).command_text().to_string(), ctx, @@ -3691,7 +3695,10 @@ impl AIBlock { RequestedCommandViewEvent::UpdatedExpansionState { is_expanded } => { // We only care about expansion state updates when the command // is running or finished (i.e. when it has a block). - let action_status = self.action_model.as_ref(ctx).get_action_status(action_id); + let action_status = self + .action_model + .as_ref(ctx) + .get_action_status(self.client_ids.conversation_id, action_id); let has_finished_command_block = { let terminal_model = self.terminal_model.lock(); terminal_model @@ -3890,7 +3897,7 @@ impl AIBlock { if self .action_model .as_ref(ctx) - .get_action_status(action_id) + .get_action_status(self.client_ids.conversation_id, action_id) .is_some_and(|status| status.is_blocked()) { ctx.focus(&view); @@ -4274,7 +4281,10 @@ impl AIBlock { // but it's not incorrect to populate if it is, and we rely on this for // for restored conversations because action model events don't re-fire // after the view is created. - let action_status = self.action_model.as_ref(ctx).get_action_status(action_id); + let action_status = self + .action_model + .as_ref(ctx) + .get_action_status(self.client_ids.conversation_id, action_id); if let Some(view) = self.search_codebase_view.get(action_id) { let files = if let Some(AIActionStatus::Finished(ref result)) = action_status { if let AIAgentActionResultType::SearchCodebase(SearchCodebaseResult::Success { @@ -4708,7 +4718,11 @@ impl AIBlock { pub fn is_blocked_on_user_confirmation(&self, app: &AppContext) -> bool { self.requested_action_ids .iter() - .filter_map(|id| self.action_model.as_ref(app).get_action_status(id)) + .filter_map(|id| { + self.action_model + .as_ref(app) + .get_action_status(self.client_ids.conversation_id, id) + }) .any(|status| status.is_blocked()) } @@ -4734,7 +4748,12 @@ impl AIBlock { ctx.subscribe_to_model(action_model, |me, action_model, event, ctx| { let action_id = event.action_id(); - if me.is_finished() || !me.requested_action_ids.contains(action_id) { + if event + .conversation_id() + .is_some_and(|conversation_id| conversation_id != me.client_ids.conversation_id) + || me.is_finished() + || !me.requested_action_ids.contains(action_id) + { // Technically, this subscription should be unregistered after `is_finished` is // set to true, but it seems that the callback is called once more after the `unsubscribe_to_model` // call, so early return here if this is errantly being called. @@ -4828,7 +4847,7 @@ impl AIBlock { { let should_collapse = action_model .as_ref(ctx) - .get_action_result(action_id) + .get_action_result(me.client_ids.conversation_id, action_id) .is_none_or(|result| match &result.result { AIAgentActionResultType::RequestCommandOutput( RequestCommandOutputResult::Completed { exit_code, .. }, @@ -4843,7 +4862,9 @@ impl AIBlock { } if let Some(view) = me.search_codebase_view.get(action_id) { - let new_status = action_model.as_ref(ctx).get_action_status(action_id); + let new_status = action_model + .as_ref(ctx) + .get_action_status(me.client_ids.conversation_id, action_id); view.update(ctx, |view, ctx| { view.update_status(new_status); ctx.notify(); @@ -4852,7 +4873,9 @@ impl AIBlock { // Create subagent panel state for finished StartAgent actions if let Some(AIActionStatus::Finished(result)) = - action_model.as_ref(ctx).get_action_status(action_id) + action_model + .as_ref(ctx) + .get_action_status(me.client_ids.conversation_id, action_id) { if let AIAgentActionResultType::StartAgent( crate::ai::agent::StartAgentResult::Success { agent_id, .. }, @@ -4874,7 +4897,11 @@ impl AIBlock { let action_statuses = me .requested_action_ids .iter() - .filter_map(|id| action_model.as_ref(ctx).get_action_status(id)) + .filter_map(|id| { + action_model + .as_ref(ctx) + .get_action_status(me.client_ids.conversation_id, id) + }) .collect_vec(); // Detecting links on SearchCodebase tool call outputs @@ -4907,7 +4934,9 @@ impl AIBlock { view.update_render_read_file_args( &me.find_state, files.clone(), - action_model.as_ref(ctx).get_action_status(action_id), + action_model + .as_ref(ctx) + .get_action_status(me.client_ids.conversation_id, action_id), ); ctx.notify(); }) @@ -4917,7 +4946,9 @@ impl AIBlock { // Open the AI document pane when documents are created or edited if let Some(action_result) = - action_model.as_ref(ctx).get_action_result(action_id) + action_model + .as_ref(ctx) + .get_action_result(me.client_ids.conversation_id, action_id) { match &action_result.result { AIAgentActionResultType::CreateDocuments( @@ -5677,7 +5708,9 @@ impl AIBlock { /// This hides their keybindings in the UI and makes them less interactive. pub fn ignore_passive_actions(&mut self, ctx: &mut ViewContext) { self.action_model.update(ctx, |action_model, ctx| { - for action in action_model.get_pending_actions() { + for action in + action_model.get_pending_actions_for_conversation(&self.client_ids.conversation_id) + { if let Some(edit) = self.requested_edits.get(&action.id) { edit.view.update(ctx, |view, ctx| view.dismiss(ctx)); } else if let Some(suggested_prompt) = self.unit_tests_suggestions.get(&action.id) { @@ -5730,7 +5763,12 @@ impl AIBlock { .view .update(ctx, |view, ctx| view.commit_and_get_command_text(ctx)); self.action_model.update(ctx, |action_model, ctx| { - action_model.handle_requested_command_accepted(&action_id, command_text, ctx); + action_model.handle_requested_command_accepted( + self.client_ids.conversation_id, + &action_id, + command_text, + ctx, + ); }); ctx.notify(); } @@ -5758,12 +5796,11 @@ impl AIBlock { /// Finds the undismissed passive code diff across all pending actions. /// This is needed because passive code diffs are NOT added to the active conversation by default, when they first appear. pub(crate) fn find_undismissed_code_diff(&self, app: &AppContext) -> Option<&RequestedEdit> { - let all_pending_actions = self.action_model.as_ref(app).get_pending_actions(); - // Find any RequestFileEdits action that has a corresponding passive code diff view. // Note that we only expect a maximum of 1 passive code diff to be undismissed at any given time. - all_pending_actions - .iter() + self.action_model + .as_ref(app) + .get_pending_actions_for_conversation(&self.client_ids.conversation_id) .find_map(|action| match &action.action { AIAgentActionType::RequestFileEdits { file_edits: _, @@ -5803,7 +5840,10 @@ impl AIBlock { .is_none_or(|output| { output.get().actions().last().is_none_or(|action| { let is_streaming = self.model.status(app).is_streaming(); - let status = self.action_model.as_ref(app).get_action_status(&action.id); + let status = self + .action_model + .as_ref(app) + .get_action_status(self.client_ids.conversation_id, &action.id); is_streaming || status.is_some_and(|status| status.is_running()) }) }) @@ -5830,7 +5870,7 @@ impl AIBlock { .any(|(action_id, requested_command)| { self.action_model .as_ref(app) - .get_action_status(action_id) + .get_action_status(self.client_ids.conversation_id, action_id) .is_some_and(|status| status.is_running()) && requested_command.view.as_ref(app).is_header_expanded() }) @@ -5930,7 +5970,10 @@ impl AIBlock { return String::new(); }; let output = output.get(); - output.format_for_copy(Some(self.action_model.as_ref(app))) + output.format_for_copy_for_conversation( + Some(self.action_model.as_ref(app)), + Some(self.client_ids.conversation_id), + ) } /// Gets AI output text for copying from the preceding user query until the next user query @@ -5985,8 +6028,10 @@ impl AIBlock { // Collect all AI outputs from start_idx to end_idx (exclusive) let mut combined_result = Vec::new(); for exchange in exchanges.iter().take(end_idx).skip(start_idx) { - let formatted_output = - exchange.format_output_for_copy(Some(self.action_model.as_ref(app))); + let formatted_output = exchange.format_output_for_copy_for_conversation( + Some(self.action_model.as_ref(app)), + Some(self.client_ids.conversation_id), + ); if !formatted_output.is_empty() { combined_result.push(formatted_output); } @@ -7158,7 +7203,7 @@ impl TypedActionView for AIBlock { let Some(result) = self .action_model .as_ref(ctx) - .get_action_result(action_id) + .get_action_result(self.client_ids.conversation_id, action_id) .map(Arc::clone) else { continue; diff --git a/app/src/ai/blocklist/block/cli.rs b/app/src/ai/blocklist/block/cli.rs index 6847df21..00148595 100644 --- a/app/src/ai/blocklist/block/cli.rs +++ b/app/src/ai/blocklist/block/cli.rs @@ -1170,7 +1170,7 @@ impl View for CLISubagentView { let is_cancelled = self .action_model .as_ref(app) - .get_action_status(&action.id) + .get_action_status(self.conversation_id, &action.id) .is_some_and(|status| status.is_cancelled()); if blocked_action.is_none() && !is_cancelled && !should_hide_responses { if let Some(rendered_action) = render_action(action.action.clone(), app) diff --git a/app/src/ai/blocklist/block/cli_controller.rs b/app/src/ai/blocklist/block/cli_controller.rs index d42579ef..24615ed9 100644 --- a/app/src/ai/blocklist/block/cli_controller.rs +++ b/app/src/ai/blocklist/block/cli_controller.rs @@ -41,6 +41,7 @@ pub enum UserTakeOverReason { #[derive(Debug, Clone, Default)] struct ActiveCLISubagentState { + initial_requested_command_conversation_id: Option, initial_requested_command_action_id: Option, task_id: Option, last_snapshot_at: Option, @@ -171,9 +172,21 @@ impl CLISubagentController { }); ctx.subscribe_to_model(action_model, |me, _, event, ctx| match event { - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { .. } => { + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { + action_id, + conversation_id, + .. + } => { let mut terminal_model = me.terminal_model.lock(); let active_block = terminal_model.block_list_mut().active_block_mut(); + if !matches_active_requested_command( + *conversation_id, + action_id, + active_block.ai_conversation_id(), + active_block.requested_command_action_id(), + ) { + return; + } active_block.update_is_agent_blocked(true); let action_id = active_block.requested_command_action_id().cloned(); @@ -183,9 +196,21 @@ impl CLISubagentController { agent_has_control: active_block.is_agent_in_control(), }); } - BlocklistAIActionEvent::ExecutingAction { .. } => { + BlocklistAIActionEvent::ExecutingAction { + action_id, + conversation_id, + .. + } => { let mut terminal_model = me.terminal_model.lock(); let active_block = terminal_model.block_list_mut().active_block_mut(); + if !matches_active_requested_command( + *conversation_id, + action_id, + active_block.ai_conversation_id(), + active_block.requested_command_action_id(), + ) { + return; + } active_block.update_is_agent_blocked(false); let action_id = active_block.requested_command_action_id().cloned(); @@ -197,12 +222,13 @@ impl CLISubagentController { } BlocklistAIActionEvent::FinishedAction { action_id: finished_action_id, + conversation_id, .. } => { let action_result = me .action_model .as_ref(ctx) - .get_action_result(finished_action_id); + .get_action_result(*conversation_id, finished_action_id); let initial_command_finished_without_snapshot = action_result.is_some_and(|result| { matches!( @@ -222,14 +248,22 @@ impl CLISubagentController { .cloned(); let mut terminal_model = me.terminal_model.lock(); let active_block = terminal_model.block_list_mut().active_block_mut(); - active_block.update_is_agent_blocked(false); + if matches_active_requested_command( + *conversation_id, + finished_action_id, + active_block.ai_conversation_id(), + active_block.requested_command_action_id(), + ) { + active_block.update_is_agent_blocked(false); - let active_command_action_id = active_block.requested_command_action_id().cloned(); - ctx.emit(CLISubagentEvent::UpdatedControl { - block_id: active_block.id().clone(), - requested_command_action_id: active_command_action_id, - agent_has_control: active_block.is_agent_in_control(), - }); + let active_command_action_id = + active_block.requested_command_action_id().cloned(); + ctx.emit(CLISubagentEvent::UpdatedControl { + block_id: active_block.id().clone(), + requested_command_action_id: active_command_action_id, + agent_has_control: active_block.is_agent_in_control(), + }); + } // Updates the last snapshot timestamp for the active block after the agent has read the block output. if let Some(snapshot_block_id) = snapshot_block_id { @@ -244,18 +278,17 @@ impl CLISubagentController { if initial_command_finished_without_snapshot { me.active_subagents_by_block.retain(|_, state| { state.task_id.is_some() - || state.initial_requested_command_action_id.as_ref() - != Some(finished_action_id) + || !matches_requested_command_identity( + *conversation_id, + finished_action_id, + state.initial_requested_command_conversation_id, + state.initial_requested_command_action_id.as_ref(), + ) }); } + drop(terminal_model); if let Some(block_id) = command_finished_block_id { - if let Some(completion) = me - .active_subagents_by_block - .get_mut(&block_id) - .and_then(|state| state.completion.as_mut()) - { - completion.final_turn_started = true; - } + me.advance_completed_subagent(&block_id, ctx); } } _ => (), @@ -322,7 +355,7 @@ impl CLISubagentController { }; drop(terminal_model); - let provider_consumed_completion = completion.as_ref().is_some_and(|completion| { + let provider_accepted_completion = completion.as_ref().is_some_and(|completion| { me.controller.update(ctx, |controller, ctx| { controller.accept_provider_command_completion( completion.conversation_id, @@ -344,13 +377,9 @@ impl CLISubagentController { if has_last_snapshot { ctx.emit(CLISubagentEvent::UpdatedLastSnapshot); } - if provider_consumed_completion { - me.finish_subagent( - &block_id, - conversation_id, - requested_command_action_id, - ctx, - ); + if provider_accepted_completion { + // The provider controller owns deactivation after it applies the queued + // completion at a safe run boundary. return; } if !me.active_subagents_by_block.contains_key(&block_id) { @@ -628,11 +657,18 @@ impl CLISubagentController { /// /// The placeholder lets command completion and action-result events arrive in either order /// without losing the completion that a subsequently-created CLI monitor needs. - pub fn track_requested_command(&mut self, block_id: &BlockId, action_id: &AIAgentActionId) { - self.active_subagents_by_block + pub fn track_requested_command( + &mut self, + block_id: &BlockId, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) { + let state = self + .active_subagents_by_block .entry(block_id.clone()) - .or_default() - .initial_requested_command_action_id = Some(action_id.clone()); + .or_default(); + state.initial_requested_command_conversation_id = Some(conversation_id); + state.initial_requested_command_action_id = Some(action_id.clone()); } /// Force the currently in-flight poll for the given long-running command block to @@ -879,6 +915,10 @@ impl CLISubagentController { requested_command_action_id: action_id.clone(), agent_has_control, }); + self.active_subagents_by_block + .entry(block_id.clone()) + .or_default() + .initial_requested_command_conversation_id = Some(conversation_id); self.active_subagents_by_block .entry(block_id.clone()) .or_default() @@ -1027,6 +1067,7 @@ fn command_finished_block_id(result: &AIAgentActionResultType) -> Option<&BlockI AIAgentActionResultType::RequestCommandOutput( RequestCommandOutputResult::LongRunningCommandSnapshot { .. } | RequestCommandOutputResult::CancelledBeforeExecution + | RequestCommandOutputResult::ExecutionError { .. } | RequestCommandOutputResult::Denylisted { .. }, ) | AIAgentActionResultType::WriteToLongRunningShellCommand( @@ -1085,6 +1126,26 @@ fn should_nudge_monitor_turn(last_exchange_has_action: bool, monitor_nudge_sent: !last_exchange_has_action && !monitor_nudge_sent } +fn matches_active_requested_command( + event_conversation_id: AIConversationId, + event_action_id: &AIAgentActionId, + active_conversation_id: Option, + active_requested_command_id: Option<&AIAgentActionId>, +) -> bool { + active_conversation_id == Some(event_conversation_id) + && active_requested_command_id == Some(event_action_id) +} + +fn matches_requested_command_identity( + event_conversation_id: AIConversationId, + event_action_id: &AIAgentActionId, + requested_command_conversation_id: Option, + requested_command_action_id: Option<&AIAgentActionId>, +) -> bool { + requested_command_conversation_id == Some(event_conversation_id) + && requested_command_action_id == Some(event_action_id) +} + #[cfg(test)] mod tests { use super::*; @@ -1121,4 +1182,51 @@ mod tests { assert!(!should_nudge_monitor_turn(false, true)); assert!(!should_nudge_monitor_turn(true, false)); } + + #[test] + fn shell_control_event_must_match_conversation_and_requested_command() { + let active_conversation_id = AIConversationId::new(); + let other_conversation_id = AIConversationId::new(); + let active_action_id = AIAgentActionId::from("same-action".to_owned()); + let other_action_id = AIAgentActionId::from("other-action".to_owned()); + + assert!(matches_active_requested_command( + active_conversation_id, + &active_action_id, + Some(active_conversation_id), + Some(&active_action_id), + )); + assert!(!matches_active_requested_command( + other_conversation_id, + &active_action_id, + Some(active_conversation_id), + Some(&active_action_id), + )); + assert!(!matches_active_requested_command( + active_conversation_id, + &other_action_id, + Some(active_conversation_id), + Some(&active_action_id), + )); + } + + #[test] + fn requested_command_identity_rejects_duplicate_id_from_another_conversation() { + let active_conversation_id = AIConversationId::new(); + let other_conversation_id = AIConversationId::new(); + let duplicate_action_id = AIAgentActionId::from("duplicate-action".to_owned()); + + assert!(matches_requested_command_identity( + active_conversation_id, + &duplicate_action_id, + Some(active_conversation_id), + Some(&duplicate_action_id), + )); + assert!(!matches_requested_command_identity( + other_conversation_id, + &duplicate_action_id, + Some(active_conversation_id), + Some(&duplicate_action_id), + )); + } } diff --git a/app/src/ai/blocklist/block/model/helper.rs b/app/src/ai/blocklist/block/model/helper.rs index 7e355988..3d6344eb 100644 --- a/app/src/ai/blocklist/block/model/helper.rs +++ b/app/src/ai/blocklist/block/model/helper.rs @@ -149,7 +149,11 @@ impl AIBlockModelHelper for T { let output = output.get(); output.messages.iter().find_map(|message| { if let AIAgentOutputMessageType::Action(action) = &message.message { - if let Some(status) = action_model.as_ref(app).get_action_status(&action.id) { + if let Some(status) = self.conversation_id(app).and_then(|conversation_id| { + action_model + .as_ref(app) + .get_action_status(conversation_id, &action.id) + }) { return status.is_blocked().then_some(action.clone()); } } diff --git a/app/src/ai/blocklist/block/status_bar.rs b/app/src/ai/blocklist/block/status_bar.rs index a157f276..50dc9615 100644 --- a/app/src/ai/blocklist/block/status_bar.rs +++ b/app/src/ai/blocklist/block/status_bar.rs @@ -328,10 +328,27 @@ impl BlocklistAIStatusBar { }, ); - ctx.subscribe_to_model(&action_model, |_, _, event, ctx| match event { - BlocklistAIActionEvent::ExecutingAction { .. } - | BlocklistAIActionEvent::FinishedAction { .. } => ctx.notify(), - _ => (), + ctx.subscribe_to_model(&action_model, |me, _, event, ctx| match event { + BlocklistAIActionEvent::ExecutingAction { + conversation_id, .. + } + | BlocklistAIActionEvent::FinishedAction { + conversation_id, .. + } if me + .active_exchange_model + .as_ref() + .is_some_and(|model| model.conversation_id(ctx) == Some(*conversation_id)) => + { + ctx.notify(); + } + BlocklistAIActionEvent::QueuedAction { .. } + | BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { .. } + | BlocklistAIActionEvent::ExecutingAction { .. } + | BlocklistAIActionEvent::FinishedAction { .. } + | BlocklistAIActionEvent::ToolLifecycle { .. } + | BlocklistAIActionEvent::InitProject(_) + | BlocklistAIActionEvent::ToggleCodeReview(_) + | BlocklistAIActionEvent::InsertCodeReviewComments { .. } => {} }); ctx.subscribe_to_model(model_event_dispatcher, |me, _, event, ctx| match event { ModelEvent::AfterBlockStarted { block_id, .. } => { diff --git a/app/src/ai/blocklist/block/view_impl.rs b/app/src/ai/blocklist/block/view_impl.rs index 376f3738..6f2fa430 100644 --- a/app/src/ai/blocklist/block/view_impl.rs +++ b/app/src/ai/blocklist/block/view_impl.rs @@ -1079,6 +1079,7 @@ impl View for AIBlock { contents.add_child(output::render( output::Props { + conversation_id: self.client_ids.conversation_id, model: self.model.as_ref(), state_handles: &self.state_handles, action_buttons: &self.action_buttons, diff --git a/app/src/ai/blocklist/block/view_impl/orchestration.rs b/app/src/ai/blocklist/block/view_impl/orchestration.rs index 8ba59df6..d4b21c60 100644 --- a/app/src/ai/blocklist/block/view_impl/orchestration.rs +++ b/app/src/ai/blocklist/block/view_impl/orchestration.rs @@ -420,7 +420,10 @@ pub(super) fn render_send_message( ) -> Box { let appearance = Appearance::as_ref(app); let theme = appearance.theme(); - let status = props.action_model.as_ref(app).get_action_status(action_id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, action_id); let orchestrator_agent_id = props .model .conversation(app) @@ -564,7 +567,10 @@ pub(super) fn render_start_agent( ) -> Box { let appearance = Appearance::as_ref(app); let theme = appearance.theme(); - let status = props.action_model.as_ref(app).get_action_status(action_id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, action_id); if let Some(AIActionStatus::Finished(result)) = &status { let AIAgentActionResultType::StartAgent(result) = &result.result else { diff --git a/app/src/ai/blocklist/block/view_impl/output.rs b/app/src/ai/blocklist/block/view_impl/output.rs index 18ceaadf..a900ec8d 100644 --- a/app/src/ai/blocklist/block/view_impl/output.rs +++ b/app/src/ai/blocklist/block/view_impl/output.rs @@ -56,6 +56,7 @@ use super::{ }; use crate::ai::agent::api::ServerConversationToken; use crate::ai::agent::comment::ReviewComment; +use crate::ai::agent::conversation::AIConversationId; use crate::ai::agent::icons::{self, gray_stop_icon, yellow_stop_icon}; use crate::ai::agent::task::TaskId; use crate::ai::agent::{ @@ -143,6 +144,7 @@ fn should_render_requested_edit(action_status: Option<&AIActionStatus>) -> bool /// Data required to render the AI block output component. #[derive(Copy, Clone)] pub(crate) struct Props<'a> { + pub(crate) conversation_id: AIConversationId, pub(crate) model: &'a dyn AIBlockModel, pub(super) state_handles: &'a AIBlockStateHandles, pub(super) action_buttons: &'a HashMap, @@ -436,7 +438,7 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { let is_action_done = props .action_model .as_ref(app) - .get_action_status(id) + .get_action_status(props.conversation_id, id) .as_ref() .is_some_and(|status| status.is_done()); if !is_action_done { @@ -476,7 +478,7 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { let agent_action_results = props .action_model .as_ref(app) - .get_action_result(id) + .get_action_result(props.conversation_id, id) .map(|action_result| action_result.as_ref()); // checks if the read file action result is completed and successful. @@ -565,8 +567,10 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { id, .. }) => { - let action_status = - props.action_model.as_ref(app).get_action_status(id); + let action_status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); if should_render_requested_edit(action_status.as_ref()) { if let Some(requested_edit) = props.requested_edits.get(id) { @@ -656,7 +660,7 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { let is_action_done = props .action_model .as_ref(app) - .get_action_status(id) + .get_action_status(props.conversation_id, id) .as_ref() .is_some_and(|status| status.is_done()); if !is_action_done { @@ -1374,7 +1378,13 @@ fn render_runtime_activity( } } - Some(render_tool_pane_shell(content.finish(), false, false, app)) + Some(render_tool_pane_shell( + content.finish(), + false, + is_expanded, + false, + app, + )) } fn should_render_stopped_output(props: Props, app: &AppContext) -> bool { @@ -1473,7 +1483,10 @@ fn render_search_codebase( id: &AIAgentActionId, app: &AppContext, ) -> Option> { - let status = props.action_model.as_ref(app).get_action_status(id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); let appearance = Appearance::as_ref(app); let theme = appearance.theme(); @@ -1974,7 +1987,10 @@ fn render_read_files( parsed_skill: Option<&ai::skills::ParsedSkill>, action_index: usize, ) -> Box { - let status = props.action_model.as_ref(app).get_action_status(id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); let appearance = Appearance::as_ref(app); let formatted_files = render_read_files_text(props.into(), file_names, app, appearance, action_index); @@ -2091,7 +2107,10 @@ fn maybe_render_edit_document( id: &AIAgentActionId, app: &AppContext, ) -> Option> { - let status = props.action_model.as_ref(app).get_action_status(id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); // Document operations are always auto-executed for now if status.as_ref().is_some_and(|status| status.is_blocked()) { @@ -2101,7 +2120,7 @@ fn maybe_render_edit_document( let agent_action_results = props .action_model .as_ref(app) - .get_action_result(id) + .get_action_result(props.conversation_id, id) .map(|action_result| action_result.as_ref()); let Some(AIAgentActionResult { @@ -2128,7 +2147,10 @@ fn maybe_render_create_document( id: &AIAgentActionId, app: &AppContext, ) -> Option> { - let status = props.action_model.as_ref(app).get_action_status(id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); // Document operations are always auto-executed for now if status.as_ref().is_some_and(|status| status.is_blocked()) { @@ -2138,7 +2160,7 @@ fn maybe_render_create_document( let agent_action_results = props .action_model .as_ref(app) - .get_action_result(id) + .get_action_result(props.conversation_id, id) .map(|action_result| action_result.as_ref()); let Some(AIAgentActionResult { @@ -2441,7 +2463,7 @@ fn render_suggest_new_conversation( let status = props .action_model .as_ref(app) - .get_action_status(action_id) + .get_action_status(props.conversation_id, action_id) .unwrap_or(AIActionStatus::Finished(Arc::new(AIAgentActionResult { result: AIAgentActionResultType::SuggestNewConversation( SuggestNewConversationResult::Cancelled, @@ -2549,7 +2571,10 @@ fn create_formatted_text_for_grep( let appearance = Appearance::as_ref(app); let theme = appearance.theme(); - let action_status = props.action_model.as_ref(app).get_action_status(id); + let action_status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); let is_cancelled = action_status .as_ref() .is_some_and(|status| status.is_cancelled()); @@ -2653,7 +2678,10 @@ fn create_formatted_text_for_file_glob( let appearance = Appearance::as_ref(app); let theme = appearance.theme(); - let action_status = props.action_model.as_ref(app).get_action_status(id); + let action_status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, id); let is_cancelled = action_status .as_ref() .is_some_and(|status| status.is_cancelled()); @@ -2754,7 +2782,10 @@ fn render_file_retrieval_tool( app: &AppContext, ) -> Box { let appearance = Appearance::as_ref(app); - let status = props.action_model.as_ref(app).get_action_status(action_id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, action_id); let mut config = RenderableAction::new_with_formatted_text(tool_formatted_text, app); @@ -2871,7 +2902,10 @@ fn render_read_mcp_resource( app: &AppContext, ) -> Box { let appearance = Appearance::as_ref(app); - let status = props.action_model.as_ref(app).get_action_status(action_id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, action_id); let mut renderable_action = RenderableAction::new(name, app); @@ -2948,11 +2982,14 @@ fn render_upload_artifact( app: &AppContext, ) -> Box { let appearance = Appearance::as_ref(app); - let status = props.action_model.as_ref(app).get_action_status(action_id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, action_id); let result = props .action_model .as_ref(app) - .get_action_result(action_id) + .get_action_result(props.conversation_id, action_id) .and_then(|result| match &result.result { AIAgentActionResultType::UploadArtifact(upload_result) => Some(upload_result), _ => None, @@ -3011,7 +3048,7 @@ fn render_use_computer( let has_screenshot = props .action_model .as_ref(app) - .get_action_result(action_id) + .get_action_result(props.conversation_id, action_id) .is_some_and(|result| { matches!( &result.result, @@ -3057,7 +3094,10 @@ fn render_request_computer_use( app: &AppContext, ) -> Box { let appearance = Appearance::as_ref(app); - let status = props.action_model.as_ref(app).get_action_status(action_id); + let status = props + .action_model + .as_ref(app) + .get_action_status(props.conversation_id, action_id); let mut renderable_action = RenderableAction::new(&request.task_summary, app); @@ -3638,7 +3678,13 @@ pub fn action_icon( app: &AppContext, ) -> galaxyui::elements::Icon { let appearance = Appearance::as_ref(app); - let status = action_model.as_ref(app).get_action_status(action_id); + let status = ai_block_model + .conversation_id(app) + .and_then(|conversation_id| { + action_model + .as_ref(app) + .get_action_status(conversation_id, action_id) + }); match status { Some(status) => match status { AIActionStatus::Preprocessing => icons::gray_circle_icon(appearance), diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 04df5d8f..3eef0c2f 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -10,7 +10,7 @@ mod pending_response_streams; pub mod response_stream; pub(super) mod shared_session; mod slash_command; -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; #[cfg(not(target_family = "wasm"))] use std::path::PathBuf; use std::sync::Arc; @@ -23,7 +23,8 @@ use futures::channel::oneshot; use galaxy_agent_core::{ turn_control, ExternalWorkId, PendingToolBatch, PendingToolCallState, ProviderRun, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, - ToolLoopGuard, TurnCommand, TurnCommandSender, TurnRequest, + StopReason, ToolLoopGuard, ToolResult, ToolResultStatus, TurnCommand, TurnCommandSender, + TurnRequest, }; use galaxy_core::assertions::safe_assert; use input_context::{input_context_for_request, parse_context_attachments}; @@ -630,6 +631,32 @@ enum ProviderCommandResult { }, } +fn convert_provider_tool_batch( + action_context: &ProviderActionContext, + batch: &PendingToolBatch, +) -> (Vec<(AIAgentAction, bool)>, Vec) { + let mut actions = Vec::new(); + let mut invalid_results = Vec::new(); + for pending in batch + .calls + .iter() + .filter(|pending| pending.state.result().is_none()) + { + match action_context.action_from_tool_call(&pending.call) { + Ok(action) => actions.push(( + action, + matches!(pending.state, PendingToolCallState::RecoveryPending), + )), + Err(message) => invalid_results.push(ToolResult { + call_id: pending.call.id.clone(), + content: format!("Invalid {} tool input: {message}", pending.call.name), + status: ToolResultStatus::Error, + }), + } + } + (actions, invalid_results) +} + struct ActiveProviderRunSlot { stream_id: ResponseStreamId, response_stream: ModelHandle, @@ -650,6 +677,13 @@ struct ActiveProviderRunSlot { monitor_prose_continuations: usize, } +struct QueuedProviderRun { + slot: ActiveProviderRunSlot, + base_provider_config: crate::ai::provider::ProviderConfig, + cli_provider_config: crate::ai::provider::ProviderConfig, + request_params: api::RequestParams, +} + #[derive(Clone)] struct ActiveProviderRunCheckpoint { run: ProviderRun, @@ -663,6 +697,7 @@ struct ActiveProviderRunCheckpoint { struct PreparedRestoredProviderRun { snapshot: ActiveProviderRunSnapshot, profiles: BTreeMap, + projection_was_initialized: bool, } impl ActiveProviderRunCheckpoint { @@ -706,6 +741,8 @@ struct ActiveProviderRunSnapshot { root_task_id: TaskId, did_input_contain_user_query: bool, persistence_offset: usize, + #[serde(default)] + cancellation_reason: Option, committed_provider_batch: Option, #[serde(default)] finished_provider_batch: Option, @@ -746,6 +783,7 @@ impl ActiveProviderRunSnapshot { root_task_id: slot.root_task_id.clone(), did_input_contain_user_query: slot.did_input_contain_user_query, persistence_offset: checkpoint.persistence_offset, + cancellation_reason: slot.cancellation_reason, committed_provider_batch: slot.committed_provider_batch.clone(), finished_provider_batch: slot.finished_provider_batch.clone(), command_action_refs: slot.command_action_refs.clone(), @@ -996,11 +1034,25 @@ fn normalize_restored_provider_snapshot( fn apply_restored_provider_command_evidence( conversation_id: AIConversationId, snapshot: &mut ActiveProviderRunSnapshot, - evidence: RestoredProviderCommandEvidence, + evidence: Option, ) -> Result<(), String> { let Some(monitor) = snapshot.command_monitor.as_ref() else { return Ok(()); }; + let Some(evidence) = evidence else { + snapshot.pending_monitor_observation = None; + snapshot.pending_command_completion = Some(PendingProviderCommandCompletion { + block_id: monitor.block_id.clone(), + initial_requested_command_action_id: Some( + monitor.initial_requested_command_action_id.clone(), + ), + command: monitor.command.clone(), + output: "The monitored command was interrupted while Galaxy was offline; its terminal block is no longer available." + .to_owned(), + exit_code: 130, + }); + return Ok(()); + }; if evidence.conversation_id != Some(conversation_id) || evidence.requested_command_action_id.as_ref() != Some(&monitor.initial_requested_command_action_id) @@ -1038,6 +1090,20 @@ fn apply_restored_provider_command_evidence( Ok(()) } +fn restored_projection_was_initialized( + has_output: bool, + has_server_output_id: bool, + has_added_messages: bool, +) -> Result { + match (has_output, has_server_output_id, has_added_messages) { + (false, false, false) => Ok(false), + (true, true, _) => Ok(true), + (false, true, _) | (false, false, true) | (true, false, _) => { + Err("restored provider projection exchange is partially initialized".to_owned()) + } + } +} + fn provider_execution_matches_active_work( run_id: &ProviderRunId, active_work_id: Option<&ExternalWorkId>, @@ -1147,6 +1213,7 @@ fn classify_provider_command_result( command: Some(command.clone()), }), RequestCommandOutputResult::CancelledBeforeExecution + | RequestCommandOutputResult::ExecutionError { .. } | RequestCommandOutputResult::Denylisted { .. } => None, }, AIAgentActionResultType::WriteToLongRunningShellCommand(result) => match result { @@ -1430,7 +1497,7 @@ fn provider_llm_lifecycle(projection: &ProviderRunProjection) -> Option { - return None + return None; } }; Some(lifecycle) @@ -1535,8 +1602,11 @@ fn provider_run_terminal_remote_log_record( } enum ProviderDriveMessage { - Response(warp_multi_agent_api::ResponseEvent), - Lifecycle(ProviderLlmLifecycle), + Projection { + lifecycle: Option, + events: Vec, + acknowledgement: oneshot::Sender>, + }, Checkpoint { checkpoint: ActiveProviderRunCheckpoint, acknowledgement: oneshot::Sender>, @@ -1559,6 +1629,7 @@ pub struct BlocklistAIController { in_flight_response_streams: PendingResponseStreams, active_provider_runs: HashMap, + queued_provider_runs: HashMap>, restoring_provider_runs: HashSet, /// The ID of the terminal surface this controller is associated with. @@ -2048,6 +2119,7 @@ impl BlocklistAIController { terminal_model, in_flight_response_streams: PendingResponseStreams::new(), active_provider_runs: HashMap::new(), + queued_provider_runs: HashMap::new(), restoring_provider_runs: HashSet::new(), terminal_surface_id, should_refresh_available_llms_on_stream_finish: false, @@ -2606,6 +2678,10 @@ impl BlocklistAIController { if self .in_flight_response_streams .has_active_stream_for_conversation(conversation_id, ctx) + && !self + .active_provider_runs + .get(&conversation_id) + .is_some_and(|slot| slot.cancellation_reason.is_some()) || self .action_model .as_ref(ctx) @@ -4563,7 +4639,7 @@ impl BlocklistAIController { .all_inputs() .any(|input| input.is_user_query()); ctx.subscribe_to_model(&response_stream, move |me, _, event, ctx| { - me.handle_response_stream_event( + let _ = me.handle_response_stream_event( input_contains_user_query, event, &response_stream_clone, @@ -4625,15 +4701,24 @@ impl BlocklistAIController { } else { None }; - self.in_flight_response_streams.register_new_stream( - response_stream_id.clone(), - conversation_data.id, - response_stream.clone(), - CancellationReason::FollowUpSubmitted { - is_for_same_conversation: true, - }, - ctx, - ); + if provider_configs.is_some() + && self + .active_provider_runs + .contains_key(&conversation_data.id) + { + self.in_flight_response_streams + .register_additional_stream(response_stream_id.clone(), response_stream.clone()); + } else { + self.in_flight_response_streams.register_new_stream( + response_stream_id.clone(), + conversation_data.id, + response_stream.clone(), + CancellationReason::FollowUpSubmitted { + is_for_same_conversation: true, + }, + ctx, + ); + } if let Some((base_provider_config, cli_provider_config)) = provider_configs { let provider_run_id = ProviderRunId::new(format!( "{}:{}", @@ -4646,37 +4731,50 @@ impl BlocklistAIController { .expect("conversation exists while starting provider run") .get_root_task_id() .clone(); - self.active_provider_runs.insert( - conversation_data.id, - ActiveProviderRunSlot { - stream_id: response_stream_id.clone(), - response_stream, - did_input_contain_user_query: input_contains_user_query, - run_id: provider_run_id, - root_task_id, - projection_target: provider_projection_target - .expect("provider projection target was validated"), - 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, - }, - ); - self.prepare_active_provider_run( - conversation_data.id, - response_stream_id.clone(), - base_provider_config, - cli_provider_config, - request_params.clone(), - ctx, - ); + let slot = ActiveProviderRunSlot { + stream_id: response_stream_id.clone(), + response_stream, + did_input_contain_user_query: input_contains_user_query, + run_id: provider_run_id, + root_task_id, + projection_target: provider_projection_target + .expect("provider projection target was validated"), + 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, + }; + if self + .active_provider_runs + .contains_key(&conversation_data.id) + { + self.queued_provider_runs + .entry(conversation_data.id) + .or_default() + .push_back(QueuedProviderRun { + slot, + base_provider_config, + cli_provider_config, + request_params: request_params.clone(), + }); + } else { + self.active_provider_runs.insert(conversation_data.id, slot); + self.prepare_active_provider_run( + conversation_data.id, + response_stream_id.clone(), + base_provider_config, + cli_provider_config, + request_params.clone(), + ctx, + ); + } } // Skip the context reset for a fired queued-prompt row (`is_queued_prompt`): its @@ -4806,21 +4904,29 @@ impl BlocklistAIController { let Some(task) = conversation.get_task(&snapshot.projection_target.task_id) else { return Err("restored provider projection task is missing".to_string()); }; - if !task + let Some(exchange) = task .exchanges() - .any(|exchange| exchange.id == snapshot.projection_target.exchange_id) - { + .find(|exchange| exchange.id == snapshot.projection_target.exchange_id) + else { return Err( "restored provider projection exchange is missing from its task" .to_string(), ); - } - Ok(()) + }; + let output = exchange.output_status.output(); + restored_projection_was_initialized( + output.is_some(), + output.is_some_and(|output| output.get().server_output_id.is_some()), + !exchange.added_message_ids.is_empty(), + ) }); - if let Err(error) = history_validation { - self.fail_restored_provider_run(conversation_id, error, ctx); - return; - } + let projection_was_initialized = match history_validation { + Ok(initialized) => initialized, + Err(error) => { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + }; if let Err(error) = normalize_restored_provider_snapshot(&mut snapshot) { self.fail_restored_provider_run(conversation_id, error, ctx); @@ -4868,7 +4974,11 @@ impl BlocklistAIController { ProviderRunProfile::new(cli_runtime, cli_monitor_request.clone()), ); } - Ok::<_, anyhow::Error>(PreparedRestoredProviderRun { snapshot, profiles }) + Ok::<_, anyhow::Error>(PreparedRestoredProviderRun { + snapshot, + profiles, + projection_was_initialized, + }) }, move |me, result, ctx| { me.handle_prepared_restored_provider_run(conversation_id, result, ctx); @@ -4886,19 +4996,18 @@ impl BlocklistAIController { }; let evidence = { let terminal_model = self.terminal_model.lock(); - let block = terminal_model + terminal_model .block_list() .block_with_id(&monitor.block_id) - .ok_or_else(|| "restored provider command block is missing".to_string())?; - RestoredProviderCommandEvidence { - conversation_id: block.ai_conversation_id(), - requested_command_action_id: block.requested_command_action_id().cloned(), - cli_task_id: block.cli_subagent_task_id().cloned(), - command: block.command_to_string(), - state: block.state(), - output: block.output_to_string(), - exit_code: block.exit_code().value(), - } + .map(|block| RestoredProviderCommandEvidence { + conversation_id: block.ai_conversation_id(), + requested_command_action_id: block.requested_command_action_id().cloned(), + cli_task_id: block.cli_subagent_task_id().cloned(), + command: block.command_to_string(), + state: block.state(), + output: block.output_to_string(), + exit_code: block.exit_code().value(), + }) }; apply_restored_provider_command_evidence(conversation_id, snapshot, evidence) } @@ -4930,7 +5039,11 @@ impl BlocklistAIController { self.restoring_provider_runs.remove(&conversation_id); return; } - let PreparedRestoredProviderRun { snapshot, profiles } = match result { + let PreparedRestoredProviderRun { + snapshot, + profiles, + projection_was_initialized, + } = match result { Ok(prepared) => prepared, Err(error) => { self.fail_restored_provider_run(conversation_id, error.to_string(), ctx); @@ -4948,6 +5061,7 @@ impl BlocklistAIController { root_task_id, did_input_contain_user_query, persistence_offset, + cancellation_reason, committed_provider_batch, finished_provider_batch, command_action_refs, @@ -4960,13 +5074,21 @@ impl BlocklistAIController { let transcript = provider_run.transcript(); let offset = persistence_offset.min(transcript.len()); let messages_sent = Arc::new(std::sync::Mutex::new(transcript[offset..].to_vec())); - let coordinator = match ProviderRunCoordinator::new(provider_run, profiles) { + let mut coordinator = match ProviderRunCoordinator::new(provider_run, profiles) { Ok(coordinator) => coordinator, Err(error) => { self.fail_restored_provider_run(conversation_id, error.to_string(), ctx); return; } }; + if let Some(reason) = cancellation_reason { + if !coordinator.run().is_terminal() { + if let Err(error) = coordinator.run_mut().cancel(reason.to_string()) { + self.fail_restored_provider_run(conversation_id, error.to_string(), ctx); + return; + } + } + } let model = LLMId::from(response_config.model_id.as_str()); let ai_identifiers = AIIdentifiers { client_conversation_id: Some(conversation_id), @@ -4984,7 +5106,7 @@ impl BlocklistAIController { let stream_id = response_stream.as_ref(ctx).id().clone(); let response_stream_clone = response_stream.clone(); ctx.subscribe_to_model(&response_stream, move |me, _, event, ctx| { - me.handle_response_stream_event( + let _ = me.handle_response_stream_event( did_input_contain_user_query, event, &response_stream_clone, @@ -5032,7 +5154,10 @@ impl BlocklistAIController { projection_target, run: Some(ActiveProviderRun { coordinator, - projector: ProviderRunResponseProjector::restored(response_config.clone()), + projector: ProviderRunResponseProjector::restored( + response_config.clone(), + projection_was_initialized, + ), response_config, action_context, messages_sent, @@ -5040,7 +5165,7 @@ impl BlocklistAIController { }), checkpoint: None, turn_control: None, - cancellation_reason: None, + cancellation_reason, committed_provider_batch, finished_provider_batch, command_action_refs, @@ -5320,26 +5445,29 @@ impl BlocklistAIController { let checkpoint_sender = sender.clone(); let result = run .coordinator - .drive_until_blocked_with_checkpoint( + .drive_until_blocked_with_acknowledgements( turn_control, |projection| { - if let Some(lifecycle) = provider_llm_lifecycle(&projection) { + let lifecycle = provider_llm_lifecycle(&projection); + let events = run.projector.project(projection); + let projection_sender = projection_sender.clone(); + Box::pin(async move { + let events = events?; + let (acknowledgement, receiver) = oneshot::channel(); projection_sender - .try_send(ProviderDriveMessage::Lifecycle(lifecycle)) + .send(ProviderDriveMessage::Projection { + lifecycle, + events, + acknowledgement, + }) + .await .map_err(|_| { - "provider lifecycle projection receiver was closed" - .to_string() + "provider projection receiver was closed".to_string() })?; - } - for event in run.projector.project(projection)? { - projection_sender - .try_send(ProviderDriveMessage::Response(event)) - .map_err(|_| { - "provider response projection receiver was closed" - .to_string() - })?; - } - Ok(()) + receiver.await.map_err(|_| { + "provider projection acknowledgement was dropped".to_string() + })? + }) }, move |provider_run| { let checkpoint_sender = checkpoint_sender.clone(); @@ -5385,29 +5513,40 @@ impl BlocklistAIController { return; } match message { - ProviderDriveMessage::Response(event) => { + ProviderDriveMessage::Projection { + lifecycle, + events, + acknowledgement, + } => { let response_stream = slot.response_stream.clone(); let did_input_contain_user_query = slot.did_input_contain_user_query; - let event = ResponseStream::projected_event(event); - self.handle_response_stream_event( - did_input_contain_user_query, - &event, - &response_stream, - ctx, - ); - } - ProviderDriveMessage::Lifecycle(lifecycle) => { + let mut result = Ok(()); + for event in events { + let event = ResponseStream::projected_event(event); + if let Err(error) = self.handle_response_stream_event( + did_input_contain_user_query, + &event, + &response_stream, + ctx, + ) { + result = Err(error); + break; + } + } #[cfg(not(target_family = "wasm"))] - remote_logging::log_model_event( - ctx, - provider_llm_lifecycle_remote_log_record( - conversation_id, - stream_id, - &lifecycle, - ), - ); + if let Some(lifecycle) = lifecycle.as_ref() { + remote_logging::log_model_event( + ctx, + provider_llm_lifecycle_remote_log_record( + conversation_id, + stream_id, + lifecycle, + ), + ); + } #[cfg(target_family = "wasm")] let _ = lifecycle; + let _ = acknowledgement.send(result); } ProviderDriveMessage::Checkpoint { checkpoint, @@ -5691,35 +5830,60 @@ impl BlocklistAIController { batch: PendingToolBatch, ctx: &mut ModelContext, ) { - let conversion = self + let Some(run) = self .active_provider_runs - .get(&conversation_id) - .and_then(|slot| slot.run.as_ref()) - .map(|run| { - batch - .calls - .iter() - .filter(|pending| pending.state.result().is_none()) - .map(|pending| { - run.action_context - .action_from_tool_call(&pending.call) - .map(|action| { - ( - action, - matches!(pending.state, PendingToolCallState::RecoveryPending), - ) - }) - }) - .collect::, _>>() - }); - let converted_actions = match conversion { - Some(Ok(actions)) => actions, - Some(Err(message)) => { - self.fail_active_provider_run(conversation_id, message, ctx); + .get_mut(&conversation_id) + .and_then(|slot| slot.run.as_mut()) + else { + return; + }; + let (converted_actions, invalid_results) = + convert_provider_tool_batch(&run.action_context, &batch); + for result in &invalid_results { + if let Err(error) = run + .coordinator + .run_mut() + .complete_tool(&batch.work_id, result.clone()) + { + self.fail_active_provider_run( + conversation_id, + format!("failed to record invalid provider tool input: {error}"), + ctx, + ); return; } - None => return, - }; + } + if converted_actions.is_empty() { + if let Err(error) = run.coordinator.run_mut().commit_tool_batch(&batch.work_id) { + self.fail_active_provider_run( + conversation_id, + format!("failed to commit invalid provider tool batch: {error}"), + ctx, + ); + return; + } + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist invalid provider tool results: {error}"), + ctx, + ); + return; + } + self.drive_active_provider_run(conversation_id, ctx); + return; + } + let mut executable_batch = batch.clone(); + for pending in &mut executable_batch.calls { + if let Some(result) = invalid_results + .iter() + .find(|result| result.call_id == pending.call.id) + { + pending.state = PendingToolCallState::Resolved { + result: result.clone(), + }; + } + } let stream_id = self.active_provider_runs[&conversation_id] .stream_id .clone(); @@ -5806,7 +5970,7 @@ impl BlocklistAIController { actions, recovery_action_ids, conversation_id, - &batch, + &executable_batch, ctx, ) }); @@ -6038,6 +6202,33 @@ impl BlocklistAIController { } } + fn detach_cancelled_provider_command( + &mut self, + conversation_id: AIConversationId, + block_id: &BlockId, + ctx: &mut ModelContext, + ) { + let detached = { + let mut terminal_model = self.terminal_model.lock(); + let active_block = terminal_model.block_list_mut().active_block_mut(); + if active_block.id() == block_id + && active_block.ai_conversation_id() == Some(conversation_id) + && active_block.is_active_and_long_running() + { + active_block.set_user_control_with_stop_reason(); + true + } else { + false + } + }; + self.deactivate_provider_cli_task(conversation_id, block_id, ctx); + if !detached { + log::warn!( + "Could not detach cancelled provider command for conversation {conversation_id:?} block {block_id:?}" + ); + } + } + fn handle_provider_actions_finished( &mut self, conversation_id: AIConversationId, @@ -6219,7 +6410,7 @@ impl BlocklistAIController { let did_input_contain_user_query = slot.did_input_contain_user_query; for event in events { let event = ResponseStream::projected_event(event); - self.handle_response_stream_event( + let _ = self.handle_response_stream_event( did_input_contain_user_query, &event, &response_stream, @@ -6237,9 +6428,35 @@ impl BlocklistAIController { ), ); match outcome { - ProviderRunOutcome::Completed(_) => { - self.finalize_completed_provider_conversation(conversation_id, ctx); - } + ProviderRunOutcome::Completed(completion) => match completion.stop_reason { + StopReason::Completed => { + self.finalize_completed_provider_conversation(conversation_id, ctx); + } + StopReason::Cancelled => { + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_conversation_status( + self.terminal_surface_id, + conversation_id, + ConversationStatus::Cancelled, + ctx, + ); + }); + } + StopReason::MaxTokens + | StopReason::ContextWindowExceeded + | StopReason::Refusal + | StopReason::ToolLoopLimit + | StopReason::Other(_) => { + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_conversation_status( + self.terminal_surface_id, + conversation_id, + ConversationStatus::Error, + ctx, + ); + }); + } + }, // Failed outcomes are finalized by the projected InternalError event. ProviderRunOutcome::Failed(_) => {} ProviderRunOutcome::Cancelled { .. } => { @@ -6314,6 +6531,13 @@ impl BlocklistAIController { response_stream: &ModelHandle, ctx: &mut ModelContext, ) { + if !self + .active_provider_runs + .get(&conversation_id) + .is_some_and(|slot| &slot.stream_id == stream_id) + { + 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}"); } @@ -6334,6 +6558,38 @@ impl BlocklistAIController { request_usage_model.refresh_request_usage_async(ctx); }); self.maybe_refresh_ai_overages(ctx); + self.start_next_queued_provider_run(conversation_id, ctx); + } + + fn start_next_queued_provider_run( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let next = self + .queued_provider_runs + .get_mut(&conversation_id) + .and_then(VecDeque::pop_front); + if self + .queued_provider_runs + .get(&conversation_id) + .is_some_and(VecDeque::is_empty) + { + self.queued_provider_runs.remove(&conversation_id); + } + let Some(next) = next else { + return; + }; + let stream_id = next.slot.stream_id.clone(); + self.active_provider_runs.insert(conversation_id, next.slot); + self.prepare_active_provider_run( + conversation_id, + stream_id, + next.base_provider_config, + next.cli_provider_config, + next.request_params, + ctx, + ); } fn cancel_active_provider_run( @@ -6342,53 +6598,59 @@ impl BlocklistAIController { reason: CancellationReason, ctx: &mut ModelContext, ) -> bool { - let Some(mut slot) = self.active_provider_runs.remove(&conversation_id) else { + let cancellation_outcome = reason.conversation_outcome(); + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { return false; }; slot.cancellation_reason = Some(reason); + let command_block_id = if matches!(cancellation_outcome, CancellationOutcome::Cancelled) { + let monitor = slot.command_monitor.take(); + if let Some(monitor) = &monitor { + slot.command_action_refs + .remove(&monitor.initial_requested_command_action_id); + } + slot.pending_monitor_observation = None; + slot.pending_command_completion = None; + monitor.map(|monitor| monitor.block_id) + } else { + None + }; if let Some(turn_control) = &slot.turn_control { let _ = turn_control.try_send(TurnCommand::Cancel); } - if let Some(mut run) = slot.run.take() { + let should_drive = if let Some(run) = slot.run.as_mut() { if !run.coordinator.run().is_terminal() { let _ = run.coordinator.run_mut().cancel(reason.to_string()); } - if let Ok(mut messages_sent) = run.messages_sent.lock() { - let transcript = run.coordinator.run().transcript(); - let offset = run.persistence_offset.min(transcript.len()); - *messages_sent = transcript[offset..].to_vec(); - } + true + } else { + false + }; + + // Keep the terminal run and its slot durable until the normal driver path projects the + // cancellation and finalizes it through `finish_active_provider_run`. + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + log::error!("Failed to persist provider cancellation: {error}"); } self.action_model.update(ctx, |action_model, ctx| { action_model.cancel_all_pending_actions(conversation_id, Some(reason), ctx); }); - let cancellation_outcome = reason.conversation_outcome(); if FeatureFlag::AgentSharedSessions.is_enabled() && !matches!(cancellation_outcome, CancellationOutcome::KeepInProgress) { self.send_cancellation_to_viewers(ctx); } - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { - history_model.mark_response_stream_cancelled( - &slot.stream_id, - conversation_id, - self.terminal_surface_id, - reason, - ctx, - ); - }); if matches!(cancellation_outcome, CancellationOutcome::Cancelled) { self.set_input_mode_for_cancellation(ctx); + if let Some(block_id) = command_block_id { + self.detach_cancelled_provider_command(conversation_id, &block_id, ctx); + } + } + if should_drive { + self.drive_active_provider_run(conversation_id, ctx); } - - self.cleanup_active_provider_run( - conversation_id, - &slot.stream_id, - &slot.response_stream, - ctx, - ); true } @@ -6535,7 +6797,7 @@ impl BlocklistAIController { ) { let stream_clone = stream.clone(); ctx.subscribe_to_model(&stream, move |me, _, event, ctx| { - me.handle_response_stream_event(false, event, &stream_clone, ctx); + let _ = me.handle_response_stream_event(false, event, &stream_clone, ctx); }); self.in_flight_response_streams.register_new_stream( stream_id, @@ -6739,7 +7001,7 @@ impl BlocklistAIController { event: &ResponseStreamEvent, response_stream: &ModelHandle, ctx: &mut ModelContext, - ) { + ) -> Result<(), String> { let stream_id = response_stream.as_ref(ctx).id().clone(); match event { @@ -6749,14 +7011,16 @@ impl BlocklistAIController { .conversation_for_response_stream(&stream_id) else { log::warn!("Could not find conversation for response stream: {stream_id:?}"); - return; + return Err(format!( + "could not find conversation for response stream {stream_id:?}" + )); }; let Some(event) = event.consume() else { debug_assert!( false, "This model should only have a single subscriber that takes ownership over the event." ); - return; + return Err("response stream event was already consumed".to_string()); }; let history_model = BlocklistAIHistoryModel::handle(ctx); match event { @@ -6794,7 +7058,7 @@ impl BlocklistAIController { } } let Some(event) = event.r#type else { - return; + return Err("response event did not contain a type".to_string()); }; match event { warp_multi_agent_api::response_event::Type::Init(init_event) => { @@ -6909,6 +7173,9 @@ impl BlocklistAIController { log::error!( "Failed to apply client actions to conversation: {e:?}" ); + return Err(format!( + "failed to apply provider client actions: {e:?}" + )); } } } @@ -6954,7 +7221,9 @@ impl BlocklistAIController { log::warn!( "Could not find conversation for response stream: {stream_id:?}" ); - return; + return Err(format!( + "could not find conversation for response stream {stream_id:?}" + )); }; id } @@ -6980,7 +7249,7 @@ impl BlocklistAIController { }) else { log::warn!("Conversation not found."); - return; + return Err("conversation not found for completed response stream".to_string()); }; #[cfg(not(target_family = "wasm"))] if let Some(metadata) = response_stream.as_ref(ctx).acp_session_metadata() { @@ -7036,7 +7305,7 @@ impl BlocklistAIController { for new_exchange_id in new_exchange_ids { let Some(exchange) = exchanges.exchange_with_id(new_exchange_id) else { log::warn!("Exchange not found."); - return; + return Err("exchange not found for completed response stream".to_string()); }; was_passive_request |= exchange.has_passive_request(); is_any_exchange_unfinished |= !exchange.output_status.is_finished(); @@ -7328,6 +7597,7 @@ impl BlocklistAIController { self.maybe_refresh_ai_overages(ctx); } } + Ok(()) } /// Sets the terminal input state after an AI request is cancelled. diff --git a/app/src/ai/blocklist/controller/pending_response_streams.rs b/app/src/ai/blocklist/controller/pending_response_streams.rs index 41b70450..63051220 100644 --- a/app/src/ai/blocklist/controller/pending_response_streams.rs +++ b/app/src/ai/blocklist/controller/pending_response_streams.rs @@ -91,6 +91,14 @@ impl PendingResponseStreams { self.streams.insert(stream_id, stream); } + pub fn register_additional_stream( + &mut self, + stream_id: ResponseStreamId, + stream: ModelHandle, + ) { + self.streams.insert(stream_id, stream); + } + pub fn cleanup_stream(&mut self, stream_id: &ResponseStreamId) { self.streams.remove(stream_id); } @@ -142,9 +150,11 @@ impl PendingResponseStreams { for response_stream in streams_to_cancel.into_iter() { log::info!( "Canceling active stream for conversation_id={conversation_id:?}, \ - reason={reason}, backtrace=\n{}", - std::backtrace::Backtrace::force_capture() + reason={reason}" ); + if let Some(backtrace) = crate::ai::tool_diagnostics::capture_backtrace() { + log::debug!("Active stream cancellation backtrace:\n{backtrace}"); + } response_stream.update(ctx, |stream, ctx| { stream.cancel(reason, conversation_id, ctx) }); diff --git a/app/src/ai/blocklist/controller/shared_session.rs b/app/src/ai/blocklist/controller/shared_session.rs index 5bf079e6..00f842ab 100644 --- a/app/src/ai/blocklist/controller/shared_session.rs +++ b/app/src/ai/blocklist/controller/shared_session.rs @@ -354,7 +354,7 @@ impl BlocklistAIController { if self .action_model .as_ref(ctx) - .get_action_result(&result.id) + .get_action_result(conversation_id, &result.id) .is_none() { self.action_model.update(ctx, |action_model, ctx| { diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 456c4d37..be6bde8f 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -25,6 +25,7 @@ use crate::ai::agent::{ WriteToLongRunningShellCommandResult, }; use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::ai::blocklist::action_model::StartAgentWaitPolicy; use crate::ai::blocklist::{ BlocklistAIHistoryEvent, BlocklistAIHistoryModel, PendingAttachment, PendingFile, RequestInput, ResponseStream, ResponseStreamId, StartAgentExecutor, @@ -265,6 +266,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider root_task_id: task_id, did_input_contain_user_query: true, persistence_offset: 0, + cancellation_reason: None, committed_provider_batch: None, finished_provider_batch: None, command_action_refs: HashMap::new(), @@ -275,6 +277,278 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider } } +#[test] +fn provider_snapshot_persists_cancellation_reason() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + snapshot.cancellation_reason = Some(CancellationReason::ManuallyCancelled); + + let restored = super::ActiveProviderRunSnapshot::parse( + &serde_json::to_string(&snapshot).expect("cancellation snapshot should serialize"), + ) + .expect("cancellation snapshot should parse"); + + assert_eq!( + restored.cancellation_reason, + Some(CancellationReason::ManuallyCancelled) + ); + assert!(!restored.run.is_terminal()); +} + +#[test] +fn cancelling_provider_startup_keeps_slot_and_durable_cancellation() { + 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())); + let checkpoint = super::ActiveProviderRunCheckpoint { + run: snapshot.run.clone(), + base_request: snapshot.base_request.clone(), + cli_monitor_request: snapshot.cli_monitor_request.clone(), + response_config: snapshot.response_config.clone(), + action_context: snapshot.action_context.clone(), + persistence_offset: snapshot.persistence_offset, + }; + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_runs.insert( + conversation_id, + super::ActiveProviderRunSlot { + stream_id, + response_stream, + did_input_contain_user_query: snapshot.did_input_contain_user_query, + run_id: snapshot.run.id().clone(), + root_task_id: snapshot.root_task_id, + projection_target: snapshot.projection_target, + run: None, + checkpoint: Some(checkpoint), + 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, + }, + ); + + assert!(controller.cancel_active_provider_run( + conversation_id, + CancellationReason::ManuallyCancelled, + ctx, + )); + assert_eq!( + controller + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.cancellation_reason), + Some(CancellationReason::ManuallyCancelled) + ); + }); + + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .expect("cancelled provider conversation should remain durable"); + let persisted = super::ActiveProviderRunSnapshot::parse( + conversation + .active_provider_run_json() + .expect("cancelled provider run should remain checkpointed"), + ) + .expect("persisted cancellation should parse"); + assert_eq!( + persisted.cancellation_reason, + Some(CancellationReason::ManuallyCancelled) + ); + assert_eq!(conversation.status(), &ConversationStatus::InProgress); + }); + }); +} + +#[test] +fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() { + 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 mut old_snapshot = provider_snapshot(conversation_id); + start_snapshot_tool(&mut old_snapshot, "old-tool"); + assert!(matches!( + old_snapshot.run.state(), + ProviderRunState::AwaitingTools { .. } + )); + + let old_stream_id = ResponseStreamId::new_for_test(); + let old_response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(old_stream_id.clone())); + let new_stream_id = ResponseStreamId::new_for_test(); + let new_response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(new_stream_id.clone())); + let new_snapshot = provider_snapshot(conversation_id); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_runs.insert( + conversation_id, + super::ActiveProviderRunSlot { + stream_id: old_stream_id.clone(), + response_stream: old_response_stream.clone(), + did_input_contain_user_query: true, + run_id: old_snapshot.run.id().clone(), + root_task_id: old_snapshot.root_task_id.clone(), + projection_target: old_snapshot.projection_target.clone(), + run: None, + checkpoint: Some(super::ActiveProviderRunCheckpoint { + run: old_snapshot.run.clone(), + base_request: old_snapshot.base_request.clone(), + cli_monitor_request: old_snapshot.cli_monitor_request.clone(), + response_config: old_snapshot.response_config.clone(), + action_context: old_snapshot.action_context.clone(), + persistence_offset: old_snapshot.persistence_offset, + }), + turn_control: None, + cancellation_reason: Some(CancellationReason::FollowUpSubmitted { + is_for_same_conversation: true, + }), + 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: new_stream_id.clone(), + response_stream: new_response_stream, + did_input_contain_user_query: true, + run_id: new_snapshot.run.id().clone(), + root_task_id: new_snapshot.root_task_id, + projection_target: new_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: crate::ai::provider::ProviderConfig::None, + cli_provider_config: crate::ai::provider::ProviderConfig::None, + request_params: crate::ai::agent::api::RequestParams::new_for_test(), + }); + + assert_eq!( + controller.active_provider_runs[&conversation_id].stream_id, + old_stream_id + ); + controller.cleanup_active_provider_run( + conversation_id, + &old_stream_id, + &old_response_stream, + ctx, + ); + assert_eq!( + controller.active_provider_runs[&conversation_id].stream_id, + new_stream_id + ); + assert!(!controller + .queued_provider_runs + .contains_key(&conversation_id)); + + // A delayed callback from the old generation cannot remove its replacement. + controller.cleanup_active_provider_run( + conversation_id, + &old_stream_id, + &old_response_stream, + ctx, + ); + assert_eq!( + controller.active_provider_runs[&conversation_id].stream_id, + new_stream_id + ); + }); + }); + }); +} + +#[test] +fn cancelled_provider_command_detaches_running_process_to_user() { + 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 = AIConversationId::new(); + let task_id = TaskId::new("provider-command-monitor".to_owned()); + let block_id = { + let mut terminal_model = terminal.model.lock(); + terminal_model.simulate_long_running_block("sleep 100", "running"); + let active_block = terminal_model.block_list_mut().active_block_mut(); + active_block.set_is_agent_tagged_in(true); + active_block + .set_agent_interaction_mode_for_agent_monitored_command( + &task_id, + conversation_id, + ) + .expect("command should become agent monitored"); + active_block.id().clone() + }; + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.detach_cancelled_provider_command(conversation_id, &block_id, ctx); + }); + + let terminal_model = terminal.model.lock(); + let active_block = terminal_model.block_list().active_block(); + assert!(!active_block.is_agent_in_control()); + assert!(active_block + .long_running_control_state() + .and_then(|state| state.user_take_over_reason()) + .is_some_and(|reason| reason.is_stop())); + assert!(active_block.is_active_and_long_running()); + }); + }); +} + fn start_snapshot_tool( snapshot: &mut super::ActiveProviderRunSnapshot, call_id: &str, @@ -305,6 +579,105 @@ fn start_snapshot_tool( batch.work_id } +#[test] +fn malformed_provider_tool_inputs_become_correlated_errors_without_dropping_valid_calls() { + 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![], + tool_calls: vec![ + ToolCall { + id: "bad-read".to_owned(), + name: "read_files".to_owned(), + arguments: serde_json::json!({"files": "not-an-array"}), + }, + ToolCall { + id: "good-grep".to_owned(), + name: "grep".to_owned(), + arguments: serde_json::json!({"queries": ["ProviderRun"]}), + }, + ], + usage: Usage::default(), + stop_reason: StopReason::Completed, + advertised_tools: BTreeSet::from(["grep".to_owned(), "read_files".to_owned()]), + }, + ) + .unwrap(); + let Some(ProviderRunStep::DispatchTools(batch)) = snapshot.run.next_step().unwrap() else { + panic!("expected provider tool batch"); + }; + + let (actions, errors) = super::convert_provider_tool_batch(&snapshot.action_context, &batch); + + assert_eq!(actions.len(), 1); + assert_eq!(actions[0].0.id.to_string(), "good-grep"); + assert_eq!(errors.len(), 1); + assert_eq!(errors[0].call_id, "bad-read"); + assert_eq!(errors[0].status, galaxy_agent_core::ToolResultStatus::Error); + assert!(errors[0].content.contains("expected an array")); +} + +#[test] +fn malformed_provider_tool_error_can_be_committed_and_run_continues() { + 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![], + tool_calls: vec![ToolCall { + id: "bad-read".to_owned(), + name: "read_files".to_owned(), + arguments: serde_json::json!({}), + }], + usage: Usage::default(), + stop_reason: StopReason::Completed, + advertised_tools: BTreeSet::from(["read_files".to_owned()]), + }, + ) + .unwrap(); + let Some(ProviderRunStep::DispatchTools(batch)) = snapshot.run.next_step().unwrap() else { + panic!("expected provider tool batch"); + }; + let (actions, errors) = super::convert_provider_tool_batch(&snapshot.action_context, &batch); + assert!(actions.is_empty()); + + snapshot + .run + .complete_tool(&batch.work_id, errors[0].clone()) + .unwrap(); + snapshot.run.commit_tool_batch(&batch.work_id).unwrap(); + + assert!(matches!( + snapshot.run.state(), + ProviderRunState::ReadyToCallModel + )); + let MessageContent::MultiPart(parts) = &snapshot.run.transcript().last().unwrap().content + else { + panic!("expected correlated tool result"); + }; + assert!(matches!( + &parts[0], + ContentPart::ToolResult { + tool_use_id, + is_error: true, + .. + } if tool_use_id == "bad-read" + )); +} + fn attach_snapshot_command_monitor( snapshot: &mut super::ActiveProviderRunSnapshot, conversation_id: AIConversationId, @@ -548,7 +921,7 @@ fn restored_active_command_rebuilds_monitor_observation() { super::apply_restored_provider_command_evidence( conversation_id, &mut snapshot, - super::RestoredProviderCommandEvidence { + Some(super::RestoredProviderCommandEvidence { conversation_id: Some(conversation_id), requested_command_action_id: Some(action_id), cli_task_id: Some(cli_task_id.clone()), @@ -556,7 +929,7 @@ fn restored_active_command_rebuilds_monitor_observation() { state: BlockState::Executing, output: "running".to_owned(), exit_code: 0, - }, + }), ) .unwrap(); @@ -582,7 +955,7 @@ fn restored_completed_command_rebuilds_exact_completion_evidence() { super::apply_restored_provider_command_evidence( conversation_id, &mut snapshot, - super::RestoredProviderCommandEvidence { + Some(super::RestoredProviderCommandEvidence { conversation_id: Some(conversation_id), requested_command_action_id: Some(action_id.clone()), cli_task_id: Some(cli_task_id), @@ -590,7 +963,7 @@ fn restored_completed_command_rebuilds_exact_completion_evidence() { state: BlockState::DoneWithExecution, output: "done".to_owned(), exit_code: 17, - }, + }), ) .unwrap(); @@ -608,6 +981,78 @@ fn restored_completed_command_rebuilds_exact_completion_evidence() { assert_eq!(completion.exit_code, 17); } +#[test] +fn restored_missing_command_block_becomes_interrupted_completion_evidence() { + 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); + snapshot.pending_monitor_observation = Some(super::PendingProviderMonitorObservation { + block_id: block_id.clone(), + cli_task_id: TaskId::new("stale-cli-task".to_owned()), + }); + + super::apply_restored_provider_command_evidence(conversation_id, &mut snapshot, None).unwrap(); + + assert!(snapshot.pending_monitor_observation.is_none()); + let completion = snapshot + .pending_command_completion + .expect("missing terminal block should become interrupted-command evidence"); + assert_eq!(completion.block_id, block_id); + assert_eq!( + completion.initial_requested_command_action_id, + Some(action_id) + ); + assert_eq!(completion.command, "sleep 10"); + assert_eq!(completion.exit_code, 130); + assert!(completion.output.contains("interrupted")); +} + +#[test] +fn restored_evidence_is_ignored_without_a_command_monitor() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let evidence = super::RestoredProviderCommandEvidence { + conversation_id: Some(conversation_id), + requested_command_action_id: None, + cli_task_id: None, + command: "sleep 10".to_owned(), + state: BlockState::Executing, + output: "running".to_owned(), + exit_code: 0, + }; + + super::apply_restored_provider_command_evidence(conversation_id, &mut snapshot, Some(evidence)) + .unwrap(); + assert!(snapshot.pending_monitor_observation.is_none()); + assert!(snapshot.pending_command_completion.is_none()); +} + +#[test] +fn restored_projection_accepts_empty_or_complete_and_rejects_partial_state() { + assert_eq!( + super::restored_projection_was_initialized(false, false, false).unwrap(), + false + ); + assert_eq!( + super::restored_projection_was_initialized(true, true, false).unwrap(), + true + ); + assert_eq!( + super::restored_projection_was_initialized(true, true, true).unwrap(), + true + ); + for state in [ + (false, false, true), + (false, true, false), + (true, false, false), + ] { + assert_eq!( + super::restored_projection_was_initialized(state.0, state.1, state.2).unwrap_err(), + "restored provider projection exchange is partially initialized" + ); + } +} + #[test] fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { let conversation_id = AIConversationId::new(); @@ -619,7 +1064,7 @@ fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { super::apply_restored_provider_command_evidence( conversation_id, &mut snapshot, - super::RestoredProviderCommandEvidence { + Some(super::RestoredProviderCommandEvidence { conversation_id: Some(AIConversationId::new()), requested_command_action_id: Some(action_id.clone()), cli_task_id: Some(cli_task_id.clone()), @@ -627,7 +1072,7 @@ fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { state: BlockState::Executing, output: String::new(), exit_code: 0, - }, + }), ) .unwrap_err(), "restored provider command block identity does not match" @@ -636,7 +1081,7 @@ fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { super::apply_restored_provider_command_evidence( conversation_id, &mut snapshot, - super::RestoredProviderCommandEvidence { + Some(super::RestoredProviderCommandEvidence { conversation_id: Some(conversation_id), requested_command_action_id: Some(action_id), cli_task_id: Some(cli_task_id), @@ -644,7 +1089,7 @@ fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { state: BlockState::Background, output: String::new(), exit_code: 0, - }, + }), ) .unwrap_err(), "restored provider command block has an invalid state" @@ -2042,7 +2487,7 @@ fn completed_provider_run_with_prior_action_resolves_child_completion_wait() { "child".to_owned(), parent_conversation_id, child_conversation_id, - None, + StartAgentWaitPolicy::Completion, ctx, ) }); diff --git a/app/src/ai/blocklist/inline_action/ask_user_question_view.rs b/app/src/ai/blocklist/inline_action/ask_user_question_view.rs index 38a6193f..4502bd7b 100644 --- a/app/src/ai/blocklist/inline_action/ask_user_question_view.rs +++ b/app/src/ai/blocklist/inline_action/ask_user_question_view.rs @@ -782,7 +782,11 @@ impl AskUserQuestionView { }; ctx.subscribe_to_model(&action_model, |me, _, event, ctx| { - if event.action_id() != me.action_id() { + if event.action_id() != me.action_id() + || event + .conversation_id() + .is_some_and(|conversation_id| conversation_id != me.conversation_id) + { return; } @@ -879,7 +883,8 @@ impl AskUserQuestionView { /// conversations still render deterministically. fn action_status(&self, app: &AppContext) -> Option { let action_model = self.action_model.as_ref(app); - if let Some(status) = action_model.get_action_status(self.action_id()) { + if let Some(status) = action_model.get_action_status(self.conversation_id, self.action_id()) + { return Some(status); } diff --git a/app/src/ai/blocklist/inline_action/code_diff_view.rs b/app/src/ai/blocklist/inline_action/code_diff_view.rs index 438d7a08..a4ee73d6 100644 --- a/app/src/ai/blocklist/inline_action/code_diff_view.rs +++ b/app/src/ai/blocklist/inline_action/code_diff_view.rs @@ -695,12 +695,26 @@ impl CodeDiffView { session_platform, ctx, ); + let action_id = (*action_id).clone(); ctx.subscribe_to_model( &action_model, move |me, action_model, event, ctx| match event { - BlocklistAIActionEvent::FinishedAction { action_id, .. } if !me.is_complete() => { - match action_model.as_ref(ctx).get_action_status(&me.action_id) { + BlocklistAIActionEvent::FinishedAction { + action_id: event_action_id, + conversation_id: event_conversation_id, + .. + } if !me.is_complete() + && *event_action_id == me.action_id + && me.identifiers.client_conversation_id == Some(*event_conversation_id) => + { + let Some(conversation_id) = me.identifiers.client_conversation_id else { + return; + }; + match action_model + .as_ref(ctx) + .get_action_status(conversation_id, &me.action_id) + { Some(AIActionStatus::Blocked) => { me.state = CodeDiffState::WaitingForUser; ctx.notify(); diff --git a/app/src/ai/blocklist/inline_action/requested_command.rs b/app/src/ai/blocklist/inline_action/requested_command.rs index 0dd3ef87..ecd11a23 100644 --- a/app/src/ai/blocklist/inline_action/requested_command.rs +++ b/app/src/ai/blocklist/inline_action/requested_command.rs @@ -412,7 +412,7 @@ impl RequestedCommandView { let is_finished = action_model .as_ref(ctx) - .get_action_result(&action_id) + .get_action_result(client_ids.conversation_id, &action_id) .is_some(); if !is_finished { @@ -424,16 +424,24 @@ impl RequestedCommandView { ctx.notify(); } BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { - action_id, .. - } if *action_id == me.action_id => { + action_id, + conversation_id, + .. + } if *conversation_id == me.client_ids.conversation_id + && *action_id == me.action_id => + { if me.action_type.is_requested_command() { me.ensure_editor(ctx); } me.set_is_header_expanded(true, ctx); ctx.notify(); } - BlocklistAIActionEvent::ExecutingAction { action_id, .. } - if *action_id == me.action_id => + BlocklistAIActionEvent::ExecutingAction { + action_id, + conversation_id, + .. + } if *conversation_id == me.client_ids.conversation_id + && *action_id == me.action_id => { // For shared-session viewers, sync the command text from the action when it starts executing. if me.action_model.as_ref(ctx).is_view_only() { @@ -467,11 +475,15 @@ impl RequestedCommandView { } ctx.notify(); } - BlocklistAIActionEvent::FinishedAction { action_id, .. } => { + BlocklistAIActionEvent::FinishedAction { + action_id, + conversation_id, + .. + } if *conversation_id == me.client_ids.conversation_id => { let Some(action_result) = me .action_model .as_ref(ctx) - .get_action_result(action_id) + .get_action_result(me.client_ids.conversation_id, action_id) .cloned() else { log::info!("Got finished action event without result: {action_id}."); @@ -724,7 +736,7 @@ impl RequestedCommandView { fn is_waiting_for_user_confirmation(&self, app: &AppContext) -> bool { self.action_model .as_ref(app) - .get_action_status(&self.action_id) + .get_action_status(self.client_ids.conversation_id, &self.action_id) .is_some_and(|status| status.is_blocked()) } @@ -750,7 +762,9 @@ impl RequestedCommandView { let Some(mouse_state_handle) = self.citation_state_handles.get(copied_citation).cloned() else { - log::warn!("Tried to retrieve mouse state handle for citation, but no mouse state handle exists."); + log::warn!( + "Tried to retrieve mouse state handle for citation, but no mouse state handle exists." + ); return None; }; render_citation( @@ -1108,7 +1122,7 @@ impl RequestedCommandView { let action_status = self .action_model .as_ref(app) - .get_action_status(&self.action_id); + .get_action_status(self.client_ids.conversation_id, &self.action_id); let mut title: Cow<'static, str>; let mut font_override = None; @@ -1457,7 +1471,7 @@ impl View for RequestedCommandView { let action_status = self .action_model .as_ref(app) - .get_action_status(&self.action_id); + .get_action_status(self.client_ids.conversation_id, &self.action_id); let is_last_output_message_in_output = self .block_model @@ -1635,6 +1649,7 @@ impl View for RequestedCommandView { let container = render_tool_pane_shell( content.finish(), has_highlighted_border, + self.is_header_expanded, should_remove_bottom_margin, app, ); 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 62c5a2a2..b41ad9a2 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 @@ -489,29 +489,36 @@ impl RunAgentsCardView { // Re-render when this action finishes or becomes blocked. let action_id_for_action_events = action_id.clone(); - ctx.subscribe_to_model(&action_model, move |me, _, event, ctx| match event { - BlocklistAIActionEvent::FinishedAction { action_id, .. } - if action_id == &action_id_for_action_events => - { - ctx.notify(); + ctx.subscribe_to_model(&action_model, move |me, _, event, ctx| { + if event.conversation_id().is_some_and(|conversation_id| { + me.block_model.conversation_id(ctx) != Some(conversation_id) + }) { + return; } - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { action_id, .. } - if action_id == &action_id_for_action_events => - { - // Normal case: streaming is complete and the action is - // ready for user confirmation. Re-render so the card - // transitions from the "Configuring agents..." placeholder - // to the full confirmation UI. - resolve_interactive_defaults(&mut me.state, &*me.block_model, ctx); - oc::repopulate_all_pickers(&mut me.state.orch, &me.handles.pickers, ctx); - me.refresh_accept_button_state(ctx); - me.maybe_auto_open_create_modal(ctx); - if let Some(conversation_id) = me.block_model.conversation_id(ctx) { - me.emit_orchestration_entered_once(conversation_id, ctx); + match event { + BlocklistAIActionEvent::FinishedAction { action_id, .. } + if action_id == &action_id_for_action_events => + { + ctx.notify(); } - ctx.notify(); + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { action_id, .. } + if action_id == &action_id_for_action_events => + { + // Normal case: streaming is complete and the action is + // ready for user confirmation. Re-render so the card + // transitions from the "Configuring agents..." placeholder + // to the full confirmation UI. + resolve_interactive_defaults(&mut me.state, &*me.block_model, ctx); + oc::repopulate_all_pickers(&mut me.state.orch, &me.handles.pickers, ctx); + me.refresh_accept_button_state(ctx); + me.maybe_auto_open_create_modal(ctx); + if let Some(conversation_id) = me.block_model.conversation_id(ctx) { + me.emit_orchestration_entered_once(conversation_id, ctx); + } + ctx.notify(); + } + _ => {} } - _ => {} }); // Repopulate the model picker when available Warp LLMs change. @@ -713,8 +720,11 @@ impl RunAgentsCardView { let request = self.state.to_request(); self.emit_decision(RunAgentsCardDecision::Accept, ctx); let action_id = self.action_id.clone(); + let Some(conversation_id) = self.block_model.conversation_id(ctx) else { + return; + }; self.action_model.update(ctx, |action_model, action_ctx| { - action_model.execute_run_agents(&action_id, request, action_ctx); + action_model.execute_run_agents(conversation_id, &action_id, request, action_ctx); }); } @@ -806,10 +816,13 @@ impl RunAgentsCardView { if self.block_model.is_restored() { return; } + let Some(conversation_id) = self.block_model.conversation_id(ctx) else { + return; + }; if matches!( self.action_model .as_ref(ctx) - .get_action_status(&self.action_id), + .get_action_status(conversation_id, &self.action_id), Some(AIActionStatus::Finished(_)) | Some(AIActionStatus::RunningAsync) ) { return; @@ -1093,9 +1106,13 @@ impl View for RunAgentsCardView { fn render(&self, app: &AppContext) -> Box { let appearance = Appearance::as_ref(app); let status = self - .action_model - .as_ref(app) - .get_action_status(&self.action_id); + .block_model + .conversation_id(app) + .and_then(|conversation_id| { + self.action_model + .as_ref(app) + .get_action_status(conversation_id, &self.action_id) + }); if let Some(AIActionStatus::Finished(result)) = &status { if let AIAgentActionResultType::RunAgents(orchestrate_result) = &result.result { @@ -1208,8 +1225,16 @@ impl TypedActionView for RunAgentsCardView { RunAgentsCardViewAction::AcceptWithoutOrchestration => { self.emit_decision(RunAgentsCardDecision::AcceptWithoutOrchestration, ctx); let action_id = self.action_id.clone(); + let Some(conversation_id) = self.block_model.conversation_id(ctx) else { + return; + }; self.action_model.update(ctx, |action_model, action_ctx| { - action_model.deny_run_agents(&action_id, String::new(), action_ctx); + action_model.deny_run_agents( + conversation_id, + &action_id, + String::new(), + action_ctx, + ); }); } RunAgentsCardViewAction::ToggleAcceptMenu => { @@ -1537,14 +1562,31 @@ pub(crate) fn format_terminal_state(result: &RunAgentsResult) -> (String, Status .iter() .filter(|a| matches!(a.kind, RunAgentsAgentOutcomeKind::Launched { .. })) .count(); - if launched == total { + let completed = agents + .iter() + .filter(|a| matches!(a.kind, RunAgentsAgentOutcomeKind::Completed { .. })) + .count(); + let successful = launched + completed; + if completed > 0 && completed == total { + let label = if total == 1 { + "Completed 1 agent".to_string() + } else { + format!("Completed {total} agents") + }; + (label, StatusKind::Success) + } else if launched == 0 && completed > 0 { + ( + format!("Completed {completed} of {total} agents"), + StatusKind::Mixed, + ) + } else if successful == total { let label = if total == 1 { "Spawned 1 agent".to_string() } else { format!("Spawned {total} agents") }; (label, StatusKind::Success) - } else if launched == 0 { + } else if successful == 0 { // Every child failed to launch: surface a terminal failure // rather than the in-progress-looking mixed state. let label = if total == 1 { @@ -1555,7 +1597,7 @@ pub(crate) fn format_terminal_state(result: &RunAgentsResult) -> (String, Status (label, StatusKind::Failure) } else { ( - format!("Spawned {launched} of {total} agents"), + format!("Spawned {successful} of {total} agents"), StatusKind::Mixed, ) } @@ -1713,7 +1755,8 @@ fn render_run_agents_child_row( app: &AppContext, ) -> Box { let outcome_conversation_id = outcome.and_then(|outcome| match &outcome.kind { - RunAgentsAgentOutcomeKind::Launched { agent_id } => { + RunAgentsAgentOutcomeKind::Launched { agent_id } + | RunAgentsAgentOutcomeKind::Completed { agent_id, .. } => { conversation_id_for_agent_id(agent_id, app) } RunAgentsAgentOutcomeKind::Failed { .. } => None, @@ -1757,6 +1800,9 @@ fn render_run_agents_child_row( RunAgentsAgentOutcomeKind::Launched { .. } => { (ConversationStatus::Success, "Started".to_string()) } + RunAgentsAgentOutcomeKind::Completed { .. } => { + (ConversationStatus::Success, "Completed".to_string()) + } RunAgentsAgentOutcomeKind::Failed { error } => ( ConversationStatus::Error, if error.trim().is_empty() { 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 5946d81f..f4ab1159 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 @@ -327,6 +327,16 @@ mod format_terminal_state_tests { } } + fn completed(name: &str, agent_id: &str) -> RunAgentsAgentOutcome { + RunAgentsAgentOutcome { + name: name.to_string(), + kind: RunAgentsAgentOutcomeKind::Completed { + agent_id: agent_id.to_string(), + output: format!("{name} output"), + }, + } + } + fn launched_result(agents: Vec) -> RunAgentsResult { RunAgentsResult::Launched { model_id: "auto".to_string(), @@ -368,6 +378,30 @@ mod format_terminal_state_tests { assert!(matches!(kind, StatusKind::Mixed)); } + #[test] + fn all_completed_uses_completed_label_and_success_status() { + let result = launched_result(vec![ + completed("a", "a-1"), + completed("b", "a-2"), + completed("c", "a-3"), + ]); + let (label, kind) = format_terminal_state(&result); + assert_eq!(label, "Completed 3 agents"); + assert!(matches!(kind, StatusKind::Success)); + } + + #[test] + fn mixed_completed_and_failed_uses_completed_label_and_mixed_status() { + let result = launched_result(vec![ + completed("a", "a-1"), + failed("b", "boom"), + completed("c", "a-3"), + ]); + let (label, kind) = format_terminal_state(&result); + assert_eq!(label, "Completed 2 of 3 agents"); + assert!(matches!(kind, StatusKind::Mixed)); + } + #[test] fn all_failed_uses_failure_status_not_mixed() { let result = launched_result(vec![ diff --git a/app/src/ai/blocklist/inline_action/tool_pane.rs b/app/src/ai/blocklist/inline_action/tool_pane.rs index 135dc3a0..86a58c37 100644 --- a/app/src/ai/blocklist/inline_action/tool_pane.rs +++ b/app/src/ai/blocklist/inline_action/tool_pane.rs @@ -14,6 +14,7 @@ use crate::ai::blocklist::block::view_impl::{ pub(crate) fn render_tool_pane_shell( content: Box, has_highlighted_border: bool, + spans_conversation_width: bool, should_remove_bottom_margin: bool, app: &AppContext, ) -> Box { @@ -25,7 +26,7 @@ pub(crate) fn render_tool_pane_shell( }; Container::new(content) - .with_margin_left(if has_highlighted_border { + .with_margin_left(if has_highlighted_border || spans_conversation_width { CONTENT_HORIZONTAL_PADDING } else { CONTENT_HORIZONTAL_PADDING + icon_size(app) + 16. diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index c6bd6e20..682ad3bf 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -61,6 +61,7 @@ pub mod request_usage_model; pub(crate) mod restored_conversations; pub(crate) mod runtime; pub(crate) mod skills; +pub(crate) mod tool_diagnostics; pub(crate) mod voice; pub use agent_tips::*; use galaxyui::AppContext; diff --git a/app/src/ai/runtime/event_translator.rs b/app/src/ai/runtime/event_translator.rs index e85b26f8..948d4e42 100644 --- a/app/src/ai/runtime/event_translator.rs +++ b/app/src/ai/runtime/event_translator.rs @@ -62,9 +62,12 @@ impl ProviderRunResponseProjector { } } - pub(crate) fn restored(config: RuntimeResponseConfig) -> Self { + pub(crate) fn restored( + config: RuntimeResponseConfig, + projection_was_initialized: bool, + ) -> Self { Self { - translator: RuntimeResponseTranslator::restored(config), + translator: RuntimeResponseTranslator::restored(config, projection_was_initialized), has_started_model_turn: false, finished: false, } @@ -88,9 +91,11 @@ impl ProviderRunResponseProjector { }) } ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event), + ProviderRunProjection::ModelRetry { .. } => { + Ok(self.translator.discard_failed_turn_output()) + } ProviderRunProjection::ModelTurnRequested { .. } | ProviderRunProjection::ModelTurnFinished { .. } - | ProviderRunProjection::ModelRetry { .. } | ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()), } } @@ -130,8 +135,15 @@ impl RuntimeResponseTranslator { Self::with_initialization(config, false) } - pub(crate) fn restored(config: RuntimeResponseConfig) -> Self { - Self::with_initialization(config, true) + pub(crate) fn restored( + mut config: RuntimeResponseConfig, + projection_was_initialized: bool, + ) -> Self { + // The task and exchange already exist in restored history. If its output was never + // initialized, replay only the stream Init rather than duplicating task/input messages. + config.needs_create_task = false; + config.user_query = None; + Self::with_initialization(config, projection_was_initialized) } fn with_initialization(config: RuntimeResponseConfig, initialized: bool) -> Self { @@ -143,7 +155,7 @@ impl RuntimeResponseTranslator { reasoning_message_id: None, activity_message_ids: HashMap::new(), activities: HashMap::new(), - has_visible_output: initialized, + has_visible_output: false, usage: Usage::default(), context_usage: None, } @@ -163,9 +175,7 @@ impl RuntimeResponseTranslator { } AgentEvent::ReasoningCompleted { text, .. } => { self.initialize(&mut events); - if self.reasoning_message_id.is_none() && !text.is_empty() { - self.add_or_append_reasoning(&text, &mut events); - } + self.complete_reasoning(&text, &mut events); } AgentEvent::RuntimeActivityUpdated { activity } => { if self.config.capabilities.host_tool_execution { @@ -244,6 +254,25 @@ impl RuntimeResponseTranslator { self.reasoning_message_id = None; } + fn discard_failed_turn_output(&mut self) -> Vec { + let mut events = Vec::new(); + if let Some(message_id) = self.text_message_id.take() { + events.push(build_replace_text_message( + &self.config.task_id, + &message_id, + "", + )); + } + if let Some(message_id) = self.reasoning_message_id.take() { + events.push(build_replace_reasoning_message( + &self.config.task_id, + &message_id, + "", + )); + } + events + } + pub(crate) fn set_task_id(&mut self, task_id: impl Into) { let task_id = task_id.into(); if self.config.task_id == task_id { @@ -315,6 +344,18 @@ impl RuntimeResponseTranslator { } } + fn complete_reasoning(&mut self, text: &str, events: &mut Vec) { + if let Some(message_id) = &self.reasoning_message_id { + events.push(build_replace_reasoning_message( + &self.config.task_id, + message_id, + text, + )); + } else if !text.is_empty() { + self.add_or_append_reasoning(text, events); + } + } + fn upsert_runtime_activity( &mut self, activity: RuntimeActivity, @@ -441,6 +482,58 @@ fn build_reasoning_message( runtime_client_action(action) } +fn build_replace_reasoning_message(task_id: &str, message_id: &str, text: &str) -> ResponseEvent { + let message = api::Message { + id: message_id.to_owned(), + task_id: task_id.to_owned(), + request_id: String::new(), + timestamp: None, + server_message_data: String::new(), + citations: Vec::new(), + fetched_memories: Vec::new(), + message: Some(api::message::Message::AgentReasoning( + api::message::AgentReasoning { + reasoning: text.to_owned(), + finished_duration: None, + }, + )), + }; + runtime_client_action(api::client_action::Action::UpdateTaskMessage( + api::client_action::UpdateTaskMessage { + task_id: task_id.to_owned(), + message: Some(message), + mask: Some(prost_types::FieldMask { + paths: vec!["agent_reasoning.reasoning".to_owned()], + }), + }, + )) +} + +fn build_replace_text_message(task_id: &str, message_id: &str, text: &str) -> ResponseEvent { + runtime_client_action(api::client_action::Action::UpdateTaskMessage( + api::client_action::UpdateTaskMessage { + task_id: task_id.to_owned(), + message: Some(api::Message { + id: message_id.to_owned(), + task_id: task_id.to_owned(), + request_id: String::new(), + timestamp: None, + server_message_data: String::new(), + citations: Vec::new(), + fetched_memories: Vec::new(), + message: Some(api::message::Message::AgentOutput( + api::message::AgentOutput { + text: text.to_owned(), + }, + )), + }), + mask: Some(prost_types::FieldMask { + paths: vec!["agent_output.text".to_owned()], + }), + }, + )) +} + fn runtime_activity_fallback_text(activity: &RuntimeActivity) -> String { let title = &activity.title; let status = activity.status.as_ref().map(|status| match status { diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index 82d7ee04..5f9cef2d 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -45,7 +45,7 @@ fn restored_provider_projection_skips_stream_initialization() { capabilities: RuntimeCapabilities::provider(), empty_output_message: None, }; - let mut projector = ProviderRunResponseProjector::restored(config); + let mut projector = ProviderRunResponseProjector::restored(config, true); let work_id = galaxy_agent_core::ExternalWorkId { run_id: galaxy_agent_core::ProviderRunId::new("run"), epoch: galaxy_agent_core::RunEpoch::new(2), @@ -82,6 +82,56 @@ fn restored_provider_projection_skips_stream_initialization() { )); } +#[test] +fn restored_uninitialized_projection_replays_init_before_live_delta() { + let config = RuntimeResponseConfig { + task_id: "task".to_owned(), + conversation_id: "conversation".to_owned(), + needs_create_task: true, + user_query: Some("already persisted".to_owned()), + model_id: "model".to_owned(), + max_context_tokens: Some(1_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + }; + let mut projector = ProviderRunResponseProjector::restored(config, false); + let work_id = galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(2), + }; + + let started = projector + .project(ProviderRunProjection::ModelTurnStarted { + work_id: work_id.clone(), + profile: galaxy_agent_core::ProviderRequestProfile::new("base"), + runtime_id: "runtime".to_owned(), + model_id: "model".to_owned(), + runtime_request_id: "request".to_owned(), + retry_attempt: 0, + elapsed_ms: 1, + }) + .unwrap(); + let delta = projector + .project(ProviderRunProjection::ModelEvent { + work_id, + event: AgentEvent::TextDelta { + text: "continued".to_owned(), + }, + }) + .unwrap(); + + assert_eq!(started.len(), 1); + assert!(matches!( + started[0].r#type, + Some(response_event::Type::Init(_)) + )); + assert_eq!(delta.len(), 1); + assert!(matches!( + delta[0].r#type, + Some(response_event::Type::ClientActions(_)) + )); +} + #[test] fn provider_followup_turn_starts_a_distinct_text_message() { let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig { @@ -261,6 +311,127 @@ fn reasoning_uses_the_native_reasoning_message_contract() { )); } +#[test] +fn reasoning_completed_authoritatively_replaces_streamed_reasoning() { + let mut translator = provider_translator(); + let streamed = translator + .translate(AgentEvent::ReasoningDelta { + text: "draft reasoning".to_owned(), + }) + .expect("reasoning delta"); + let completed = translator + .translate(AgentEvent::ReasoningCompleted { + text: "authoritative reasoning".to_owned(), + signature: Some("signature".to_owned()), + }) + .expect("reasoning completion"); + + let Some(response_event::Type::ClientActions(streamed_actions)) = &streamed[1].r#type else { + panic!("expected streamed reasoning action"); + }; + let Some(client_action::Action::AddMessagesToTask(add)) = &streamed_actions.actions[0].action + else { + panic!("expected streamed reasoning message"); + }; + let message_id = add.messages[0].id.clone(); + let Some(response_event::Type::ClientActions(completed_actions)) = &completed[0].r#type else { + panic!("expected completed reasoning action"); + }; + let Some(client_action::Action::UpdateTaskMessage(update)) = + &completed_actions.actions[0].action + else { + panic!("completed reasoning must replace the streamed value"); + }; + let message = update.message.as_ref().expect("replacement message"); + assert_eq!(message.id, message_id); + assert!(matches!( + &message.message, + Some(message::Message::AgentReasoning(reasoning)) + if reasoning.reasoning == "authoritative reasoning" + )); + assert_eq!( + update.mask.as_ref().expect("replacement mask").paths, + ["agent_reasoning.reasoning"] + ); +} + +#[test] +fn provider_retry_clears_failed_attempt_output_before_new_messages() { + let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig { + task_id: "task".to_owned(), + conversation_id: "conversation".to_owned(), + needs_create_task: false, + user_query: None, + model_id: "model".to_owned(), + max_context_tokens: Some(1_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + }); + let work_id = galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(1), + }; + projector + .project(ProviderRunProjection::ModelTurnStarted { + work_id: work_id.clone(), + profile: galaxy_agent_core::ProviderRequestProfile::new("base"), + runtime_id: "runtime".to_owned(), + model_id: "model".to_owned(), + runtime_request_id: "request".to_owned(), + retry_attempt: 0, + elapsed_ms: 1, + }) + .unwrap(); + for event in [ + AgentEvent::TextDelta { + text: "failed text".to_owned(), + }, + AgentEvent::ReasoningDelta { + text: "failed reasoning".to_owned(), + }, + ] { + projector + .project(ProviderRunProjection::ModelEvent { + work_id: work_id.clone(), + event, + }) + .unwrap(); + } + let retry = projector + .project(ProviderRunProjection::ModelRetry { + work_id, + profile: galaxy_agent_core::ProviderRequestProfile::new("base"), + runtime_id: "runtime".to_owned(), + model_id: "model".to_owned(), + retry_attempt: 1, + elapsed_ms: 2, + error: galaxy_agent_core::AgentError::new( + galaxy_agent_core::AgentErrorKind::Transport, + "retry", + ), + }) + .unwrap(); + + assert_eq!(retry.len(), 2); + for event in retry { + let Some(response_event::Type::ClientActions(actions)) = event.r#type else { + panic!("cleanup must use a client action"); + }; + let Some(client_action::Action::UpdateTaskMessage(update)) = &actions.actions[0].action + else { + panic!("cleanup must replace failed output"); + }; + let message = update.message.as_ref().expect("cleanup message"); + match &message.message { + Some(message::Message::AgentOutput(output)) => assert!(output.text.is_empty()), + Some(message::Message::AgentReasoning(reasoning)) => { + assert!(reasoning.reasoning.is_empty()) + } + _ => panic!("cleanup must target visible text or reasoning"), + } + } +} + #[test] fn session_activity_updates_the_same_structured_message() { let mut translator = session_translator(); diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index aa836066..f7b4856d 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -5,16 +5,17 @@ use std::sync::Arc; use std::time::Duration; use futures::future::BoxFuture; -use futures::StreamExt; +use futures::{FutureExt as _, StreamExt}; use galaxy_agent_core::{ - AgentError, AgentErrorKind, AgentEvent, AgentRuntime, CompletedModelTurn, ContentPart, - ExternalWorkId, ModelFailureDisposition, PendingToolBatch, ProviderModelCall, + turn_control, AgentError, AgentErrorKind, AgentEvent, AgentRuntime, CompletedModelTurn, + ContentPart, ExternalWorkId, ModelFailureDisposition, PendingToolBatch, ProviderModelCall, ProviderRequestProfile, ProviderRun, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState, - ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnControl, TurnRequest, Usage, + ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnCommand, TurnControl, + TurnRequest, Usage, }; use instant::Instant; -use warpui::r#async::FutureExt as _; +use warpui::r#async::Timer; use crate::ai::agent::conversation::AIConversationId; @@ -321,31 +322,55 @@ impl ProviderRunCoordinator { &mut self, control: TurnControl, mut project: F, - mut checkpoint: C, + checkpoint: C, ) -> Result where F: FnMut(ProviderRunProjection) -> Result<(), String>, C: FnMut(ProviderRun) -> BoxFuture<'static, Result<(), String>>, + { + self.drive_until_blocked_with_acknowledgements( + control, + move |projection| { + let result = project(projection); + Box::pin(async move { result }) + }, + checkpoint, + ) + .await + } + + pub(crate) async fn drive_until_blocked_with_acknowledgements( + &mut self, + control: TurnControl, + mut project: F, + mut checkpoint: C, + ) -> Result + where + F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>, + C: FnMut(ProviderRun) -> BoxFuture<'static, Result<(), String>>, { loop { match self.run.next_step()? { Some(ProviderRunStep::CallModel(call)) => { if !self.checkpoint_or_fail(&mut checkpoint).await? { - continue; + return self.terminal_block(); } - self.drive_model_call(call, control.clone(), &mut project) + self.drive_model_call_acknowledged(call, control.clone(), &mut project) .await?; } Some(ProviderRunStep::DispatchTools(batch)) => { if !self.checkpoint_or_fail(&mut checkpoint).await? { - continue; + return self.terminal_block(); } - if !self.project_or_fail( - ProviderRunProjection::ToolBatchReady { - batch: batch.clone(), - }, - &mut project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ToolBatchReady { + batch: batch.clone(), + }, + &mut project, + ) + .await? + { continue; } if batch.is_complete() { @@ -356,7 +381,7 @@ impl ProviderRunCoordinator { } Some(ProviderRunStep::Done(outcome)) => { if !self.checkpoint_or_fail(&mut checkpoint).await? { - continue; + return self.terminal_block(); } return Ok(ProviderRunBlock::Done(outcome)); } @@ -368,7 +393,7 @@ impl ProviderRunCoordinator { let work_id = work_id.clone(); let stop_reason = stop_reason.clone(); if !self.checkpoint_or_fail(&mut checkpoint).await? { - continue; + return self.terminal_block(); } return Ok(ProviderRunBlock::AwaitingDriver { work_id, @@ -413,14 +438,26 @@ impl ProviderRunCoordinator { } } - async fn drive_model_call( + fn terminal_block(&mut self) -> Result { + match self.run.next_step()? { + Some(ProviderRunStep::Done(outcome)) => Ok(ProviderRunBlock::Done(outcome)), + Some(ProviderRunStep::CallModel(_) | ProviderRunStep::DispatchTools(_)) | None => Err( + ProviderRunCoordinatorError::Core(ProviderRunProtocolError::UnexpectedState { + expected: ProviderRunPhase::Failed, + actual: self.run.state().phase(), + }), + ), + } + } + + async fn drive_model_call_acknowledged( &mut self, call: ProviderModelCall, control: TurnControl, project: &mut F, ) -> Result<(), ProviderRunCoordinatorError> where - F: FnMut(ProviderRunProjection) -> Result<(), String>, + F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>, { let Some(profile) = self.profiles.get(call.profile.as_str()).cloned() else { self.run.fail( @@ -440,39 +477,66 @@ impl ProviderRunCoordinator { .collect::>(); let runtime_id = profile.runtime.descriptor().id.clone(); let model_id = profile.request.model.as_str().to_string(); - if !self.project_or_fail( - ProviderRunProjection::ModelTurnRequested { - work_id: call.work_id.clone(), - profile: call.profile.clone(), - runtime_id: runtime_id.clone(), - model_id: model_id.clone(), - retry_attempt: call.retry_attempt, - }, - project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelTurnRequested { + work_id: call.work_id.clone(), + profile: call.profile.clone(), + runtime_id: runtime_id.clone(), + model_id: model_id.clone(), + retry_attempt: call.retry_attempt, + }, + project, + ) + .await? + { return Ok(()); } let request = request_for_model_call(profile.request.clone(), &call); let started_at = Instant::now(); - let stream = match profile - .runtime - .start_turn(request, control) - .with_timeout(self.model_start_timeout) - .await - { - Ok(Ok(stream)) => stream, - Ok(Err(error)) => { - self.handle_model_failure(&call, &profile, started_at, error, project)?; + let (attempt_sender, attempt_control) = turn_control(); + let start_future = profile.runtime.start_turn(request, attempt_control).fuse(); + let timeout = futures::FutureExt::fuse(Timer::after(self.model_start_timeout)); + futures::pin_mut!(start_future, timeout); + let mut control_open = true; + let start_result = loop { + let command = if control_open { + futures::future::Either::Left(control.receive()) + } else { + futures::future::Either::Right(futures::future::pending()) + } + .fuse(); + futures::pin_mut!(command); + futures::select_biased! { + result = start_future => break Some(result), + _ = timeout => { + let _ = attempt_sender.send(TurnCommand::Cancel).await; + break None; + } + command = command => match command { + Ok(command) => { + let _ = attempt_sender.send(command).await; + } + Err(_) => control_open = false, + }, + } + }; + let stream = match start_result { + Some(Ok(stream)) => stream, + Some(Err(error)) => { + self.handle_model_failure_acknowledged(&call, &profile, started_at, error, project) + .await?; return Ok(()); } - Err(_) => { - self.handle_model_failure( + None => { + self.handle_model_failure_acknowledged( &call, &profile, started_at, provider_timeout_error("start", self.model_start_timeout), project, - )?; + ) + .await?; return Ok(()); } }; @@ -480,157 +544,230 @@ impl ProviderRunCoordinator { let mut buffer = ModelTurnBuffer::default(); loop { - let event = match stream - .next() - .with_timeout(self.model_event_idle_timeout) - .await - { - Ok(Some(Ok(event))) => event, - Ok(Some(Err(error))) => { - self.handle_model_failure(&call, &profile, started_at, error, project)?; + let next_event = stream.next().fuse(); + let timeout = futures::FutureExt::fuse(Timer::after(self.model_event_idle_timeout)); + futures::pin_mut!(next_event, timeout); + let event_result = loop { + let command = if control_open { + futures::future::Either::Left(control.receive()) + } else { + futures::future::Either::Right(futures::future::pending()) + } + .fuse(); + futures::pin_mut!(command); + futures::select_biased! { + event = next_event => break Some(event), + _ = timeout => { + let _ = attempt_sender.send(TurnCommand::Cancel).await; + break None; + } + command = command => match command { + Ok(command) => { + let _ = attempt_sender.send(command).await; + } + Err(_) => control_open = false, + }, + } + }; + let event = match event_result { + Some(Some(Ok(event))) => event, + Some(Some(Err(error))) => { + self.handle_model_failure_acknowledged( + &call, &profile, started_at, error, project, + ) + .await?; return Ok(()); } - Ok(None) => break, - Err(_) => { - self.handle_model_failure( + Some(None) => break, + None => { + self.handle_model_failure_acknowledged( &call, &profile, started_at, provider_timeout_error("event", self.model_event_idle_timeout), project, - )?; + ) + .await?; return Ok(()); } }; match event { AgentEvent::TurnStarted { runtime_request_id } => { if buffer.started { - self.handle_model_failure( + self.handle_model_failure_acknowledged( &call, &profile, started_at, protocol_error("provider emitted more than one TurnStarted event"), project, - )?; + ) + .await?; return Ok(()); } if runtime_request_id.is_empty() { - self.handle_model_failure( + self.handle_model_failure_acknowledged( &call, &profile, started_at, protocol_error("provider emitted an empty runtime request ID"), project, - )?; + ) + .await?; return Ok(()); } buffer.started = true; - if !self.project_or_fail( - ProviderRunProjection::ModelTurnStarted { - work_id: call.work_id.clone(), - profile: call.profile.clone(), - runtime_id: runtime_id.clone(), - model_id: model_id.clone(), - runtime_request_id, - retry_attempt: call.retry_attempt, - elapsed_ms: elapsed_millis(started_at), - }, - project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelTurnStarted { + work_id: call.work_id.clone(), + profile: call.profile.clone(), + runtime_id: runtime_id.clone(), + model_id: model_id.clone(), + runtime_request_id, + retry_attempt: call.retry_attempt, + elapsed_ms: elapsed_millis(started_at), + }, + project, + ) + .await? + { return Ok(()); } } AgentEvent::TextDelta { text } => { - if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? { + if !self + .ensure_model_started_acknowledged( + &call, &profile, started_at, &buffer, project, + ) + .await? + { return Ok(()); } buffer.text.push_str(&text); - if !self.project_or_fail( - ProviderRunProjection::ModelEvent { - work_id: call.work_id.clone(), - event: AgentEvent::TextDelta { text }, - }, - project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::TextDelta { text }, + }, + project, + ) + .await? + { return Ok(()); } } AgentEvent::ReasoningDelta { text } => { - if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? { + if !self + .ensure_model_started_acknowledged( + &call, &profile, started_at, &buffer, project, + ) + .await? + { return Ok(()); } buffer.reasoning.push_str(&text); - if !self.project_or_fail( - ProviderRunProjection::ModelEvent { - work_id: call.work_id.clone(), - event: AgentEvent::ReasoningDelta { text }, - }, - project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::ReasoningDelta { text }, + }, + project, + ) + .await? + { return Ok(()); } } AgentEvent::ReasoningCompleted { text, signature } => { - if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? { + if !self + .ensure_model_started_acknowledged( + &call, &profile, started_at, &buffer, project, + ) + .await? + { return Ok(()); } - if !text.is_empty() { - buffer.reasoning.clone_from(&text); - } + buffer.reasoning.clone_from(&text); buffer.reasoning_signature.clone_from(&signature); - if !self.project_or_fail( - ProviderRunProjection::ModelEvent { - work_id: call.work_id.clone(), - event: AgentEvent::ReasoningCompleted { text, signature }, - }, - project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::ReasoningCompleted { text, signature }, + }, + project, + ) + .await? + { return Ok(()); } } AgentEvent::Tool { event: ToolEvent::Proposed { call: tool_call }, } => { - if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? { + if !self + .ensure_model_started_acknowledged( + &call, &profile, started_at, &buffer, project, + ) + .await? + { return Ok(()); } buffer.tool_calls.push(tool_call); } AgentEvent::UsageUpdated { usage } => { - if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? { + if !self + .ensure_model_started_acknowledged( + &call, &profile, started_at, &buffer, project, + ) + .await? + { return Ok(()); } buffer.usage.clone_from(&usage); let cumulative_usage = combined_usage(self.run.usage(), &usage); - if !self.project_or_fail( - ProviderRunProjection::ModelEvent { - work_id: call.work_id.clone(), - event: AgentEvent::UsageUpdated { - usage: cumulative_usage, + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::UsageUpdated { + usage: cumulative_usage, + }, }, - }, - project, - )? { + project, + ) + .await? + { return Ok(()); } } AgentEvent::TurnStopped { reason } => { - if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? { + if !self + .ensure_model_started_acknowledged( + &call, &profile, started_at, &buffer, project, + ) + .await? + { return Ok(()); } - if !self.project_or_fail( - ProviderRunProjection::ModelTurnFinished { - work_id: call.work_id.clone(), - profile: call.profile.clone(), - runtime_id: runtime_id.clone(), - model_id: model_id.clone(), - stop_reason: reason.clone(), - retry_attempt: call.retry_attempt, - elapsed_ms: elapsed_millis(started_at), - tool_call_count: buffer.tool_calls.len(), - }, - project, - )? { + if !self + .project_or_fail_acknowledged( + ProviderRunProjection::ModelTurnFinished { + work_id: call.work_id.clone(), + profile: call.profile.clone(), + runtime_id: runtime_id.clone(), + model_id: model_id.clone(), + stop_reason: reason.clone(), + retry_attempt: call.retry_attempt, + elapsed_ms: elapsed_millis(started_at), + tool_call_count: buffer.tool_calls.len(), + }, + project, + ) + .await? + { return Ok(()); } if reason == StopReason::Cancelled { @@ -657,7 +794,7 @@ impl ProviderRunCoordinator { | AgentEvent::ContextUsageUpdated { .. } | AgentEvent::UserInputAccepted { .. } | AgentEvent::RuntimeNotice { .. } => { - self.handle_model_failure( + self.handle_model_failure_acknowledged( &call, &profile, started_at, @@ -665,23 +802,25 @@ impl ProviderRunCoordinator { "direct-provider transport emitted a non-model lifecycle event", ), project, - )?; + ) + .await?; return Ok(()); } } } - self.handle_model_failure( + self.handle_model_failure_acknowledged( &call, &profile, started_at, protocol_error("provider stream ended before TurnStopped"), project, - )?; + ) + .await?; Ok(()) } - fn ensure_model_started( + async fn ensure_model_started_acknowledged( &mut self, call: &ProviderModelCall, profile: &ProviderRunProfile, @@ -690,22 +829,23 @@ impl ProviderRunCoordinator { project: &mut F, ) -> Result where - F: FnMut(ProviderRunProjection) -> Result<(), String>, + F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>, { if buffer.started { return Ok(true); } - self.handle_model_failure( + self.handle_model_failure_acknowledged( call, profile, started_at, protocol_error("provider emitted model output before TurnStarted"), project, - )?; + ) + .await?; Ok(false) } - fn handle_model_failure( + async fn handle_model_failure_acknowledged( &mut self, call: &ProviderModelCall, profile: &ProviderRunProfile, @@ -714,7 +854,7 @@ impl ProviderRunCoordinator { project: &mut F, ) -> Result<(), ProviderRunCoordinatorError> where - F: FnMut(ProviderRunProjection) -> Result<(), String>, + F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>, { let disposition = self .run @@ -737,7 +877,7 @@ impl ProviderRunCoordinator { )); } }; - self.project_or_fail( + self.project_or_fail_acknowledged( ProviderRunProjection::ModelRetry { work_id: call.work_id.clone(), profile: call.profile.clone(), @@ -748,20 +888,21 @@ impl ProviderRunCoordinator { error, }, project, - )?; + ) + .await?; } Ok(()) } - fn project_or_fail( + async fn project_or_fail_acknowledged( &mut self, event: ProviderRunProjection, project: &mut F, ) -> Result where - F: FnMut(ProviderRunProjection) -> Result<(), String>, + F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>, { - match project(event) { + match project(event).await { Ok(()) => Ok(true), Err(message) => { self.run.fail( diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index 177363f5..6992a587 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -79,6 +79,7 @@ struct StallingRuntime { first_attempt_stall: FirstAttemptStall, attempts: AtomicUsize, requests: Mutex>, + controls: Mutex>, } impl StallingRuntime { @@ -93,12 +94,22 @@ impl StallingRuntime { first_attempt_stall, attempts: AtomicUsize::new(0), requests: Mutex::new(Vec::new()), + controls: Mutex::new(Vec::new()), } } fn requests(&self) -> Vec { self.requests.lock().unwrap().clone() } + + fn attempt_commands(&self) -> Vec> { + self.controls + .lock() + .unwrap() + .iter() + .map(|control| control.try_receive().ok()) + .collect() + } } #[async_trait] @@ -110,9 +121,10 @@ impl AgentRuntime for StallingRuntime { async fn start_turn( &self, request: TurnRequest, - _control: TurnControl, + control: TurnControl, ) -> Result { self.requests.lock().unwrap().push(request); + self.controls.lock().unwrap().push(control); let attempt = self.attempts.fetch_add(1, Ordering::Relaxed); if attempt == 0 { match self.first_attempt_stall { @@ -778,6 +790,64 @@ async fn recoverable_start_failure_retries_the_same_work_identity() { assert_eq!(retry.1, 1); } +#[tokio::test] +async fn model_progress_waits_for_each_projection_acknowledgement() { + let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let mut coordinator = coordinator(runtime.clone()); + let (acknowledgement_sender, acknowledgement_receiver) = async_channel::unbounded(); + let (observed_sender, observed_receiver) = async_channel::unbounded(); + let (_sender, control) = turn_control(); + + let drive = async { + coordinator + .drive_until_blocked_with_acknowledgements( + control, + move |projection| { + let acknowledgement_receiver = acknowledgement_receiver.clone(); + let observed_sender = observed_sender.clone(); + Box::pin(async move { + observed_sender.send(projection).await.unwrap(); + acknowledgement_receiver.recv().await.unwrap() + }) + }, + |_| Box::pin(async { Ok(()) }), + ) + .await + .unwrap() + }; + let driver = async { + let first = observed_receiver.recv().await.unwrap(); + assert!(matches!( + first, + ProviderRunProjection::ModelTurnRequested { .. } + )); + assert_eq!(runtime.requests().len(), 0); + assert!(observed_receiver.try_recv().is_err()); + acknowledgement_sender.send(Ok(())).await.unwrap(); + + let second = observed_receiver.recv().await.unwrap(); + assert!(matches!( + second, + ProviderRunProjection::ModelTurnStarted { .. } + )); + assert_eq!(runtime.requests().len(), 1); + assert!(observed_receiver.try_recv().is_err()); + acknowledgement_sender.send(Ok(())).await.unwrap(); + + loop { + let projection = observed_receiver.recv().await.unwrap(); + let finished = matches!(projection, ProviderRunProjection::ModelTurnFinished { .. }); + acknowledgement_sender.send(Ok(())).await.unwrap(); + if finished { + break; + } + } + }; + + let (block, ()) = futures::join!(drive, driver); + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); +} + fn assert_single_retry_lifecycle( projections: &[ProviderRunProjection], expected_timeout_stage: &str, @@ -868,6 +938,10 @@ async fn model_start_timeout_retries_the_same_work_identity() { assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); assert_eq!(runtime.requests().len(), 2); + assert_eq!( + runtime.attempt_commands(), + vec![Some(TurnCommand::Cancel), None] + ); assert_eq!(coordinator.run().model_retries(), 1); assert_single_retry_lifecycle(&projections, "start timed out", false); } @@ -887,10 +961,43 @@ async fn model_event_idle_timeout_retries_the_same_work_identity() { assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); assert_eq!(runtime.requests().len(), 2); + assert_eq!( + runtime.attempt_commands(), + vec![Some(TurnCommand::Cancel), None] + ); assert_eq!(coordinator.run().model_retries(), 1); assert_single_retry_lifecycle(&projections, "event timed out", true); } +#[tokio::test] +async fn persistent_checkpoint_failure_terminates_without_redrive() { + let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let mut coordinator = coordinator(runtime.clone()); + let checkpoint_attempts = Arc::new(AtomicUsize::new(0)); + let attempts = checkpoint_attempts.clone(); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked_with_checkpoint( + control, + |_| Ok(()), + move |_| { + attempts.fetch_add(1, Ordering::Relaxed); + Box::pin(async { Err("database unavailable".to_string()) }) + }, + ) + .await + .unwrap(); + + let ProviderRunBlock::Done(ProviderRunOutcome::Failed(failure)) = block else { + panic!("checkpoint failure must terminate the drive"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::ExternalWork); + assert!(failure.message.contains("database unavailable")); + assert_eq!(checkpoint_attempts.load(Ordering::Relaxed), 1); + assert!(runtime.requests().is_empty()); +} + #[tokio::test] async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() { let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn()])); @@ -953,6 +1060,38 @@ async fn stream_without_terminal_event_fails_instead_of_committing_partial_outpu assert_eq!(coordinator.run().transcript().len(), 1); } +#[tokio::test] +async fn reasoning_completed_authoritatively_replaces_canonical_reasoning() { + let runtime = Arc::new(ScriptedRuntime::new(vec![Ok(vec![ + started("request-reasoning"), + Ok(AgentEvent::ReasoningDelta { + text: "draft".to_string(), + }), + Ok(AgentEvent::ReasoningCompleted { + text: "final".to_string(), + signature: Some("signature".to_string()), + }), + stopped(StopReason::Completed), + ])])); + let mut coordinator = coordinator(runtime); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + let MessageContent::MultiPart(parts) = &coordinator.run().transcript()[1].content else { + panic!("expected canonical reasoning content"); + }; + assert!(matches!( + parts.as_slice(), + [ContentPart::Reasoning { text, signature }] + if text == "final" && signature.as_deref() == Some("signature") + )); +} + #[tokio::test] async fn provider_cancellation_does_not_commit_partial_assistant_content() { let expected_transcript = request().messages; diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index 53fcf58e..bc14036f 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -835,6 +835,16 @@ fn build_system_prompt( .join(", "), ); prompt.push_str(".\nNever invent tool names or parameters. Check command exit codes and tool error results before claiming success.\n"); + if tools.iter().any(|tool| tool.name == "run_shell_command") { + let has_file_tools = tools + .iter() + .any(|tool| matches!(tool.name.as_str(), "file_glob" | "grep" | "read_files")); + if has_file_tools { + prompt.push_str( + "Prefer `file_glob`, `grep`, and `read_files` for file discovery, content search, and file reading when they are available. Reserve `run_shell_command` for operations those specialized tools cannot perform; do not use shell `find`, `grep`, `rg`, `cat`, `head`, or `tail` as substitutes.\n", + ); + } + } if tools.iter().any(|tool| tool.name == "create_plan") { prompt.push_str( "Plan document creation is available through `create_plan`. When the user asks you to create a plan for review, research first as needed, then call `create_plan`; do not merely return the plan as prose or claim that no plan-creation tool is available. If the user asks to review the plan before implementation, creating the document and presenting it for review is the requested outcome; do not implement it until they approve.\n", diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index 5e61c3be..7644e94e 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -127,6 +127,28 @@ fn builds_a_rig_turn_directly_from_galaxy_request_state() { )); } +#[test] +fn system_prompt_prefers_specialized_file_tools_over_shell_substitutes() { + let mut params = RequestParams::new_for_test(); + params.input = vec![user_query("Find every Rust file containing ProviderRun")]; + + let prepared = prepare_rig_turn( + &config(), + params, + vec![ + ToolType::RunShellCommand, + ToolType::FileGlob, + ToolType::Grep, + ToolType::ReadFiles, + ], + Vec::new(), + ); + let prompt = prepared.request.system_prompt.expect("system prompt"); + + assert!(prompt.contains("Prefer `file_glob`, `grep`, and `read_files`")); + assert!(prompt.contains("do not use shell `find`, `grep`, `rg`, `cat`, `head`, or `tail`")); +} + #[test] fn normal_turn_advertises_plan_creation_and_corrects_false_unavailability_claims() { let mut params = RequestParams::new_for_test(); diff --git a/app/src/ai/runtime/rig_tool.rs b/app/src/ai/runtime/rig_tool.rs index 37c1a31b..cb898d02 100644 --- a/app/src/ai/runtime/rig_tool.rs +++ b/app/src/ai/runtime/rig_tool.rs @@ -25,6 +25,9 @@ pub(crate) fn action_from_tool_call( mcp_tool_aliases: &HashMap, ) -> Result { let input = &call.arguments; + if !input.is_object() { + return Err(format!("invalid {} input: expected an object", call.name)); + } let action = if let Some(target) = mcp_tool_aliases.get(&call.name) { AIAgentActionType::CallMCPTool { server_id: target.server_id, @@ -33,197 +36,220 @@ pub(crate) fn action_from_tool_call( } } else { match call.name.as_str() { - "run_shell_command" => AIAgentActionType::RequestCommandOutput { - command: string(input, "command"), - is_read_only: Some(boolean(input, "is_read_only")), - is_risky: Some(boolean(input, "is_risky")), - wait_until_completion: boolean(input, "wait_until_complete"), - uses_pager: Some(boolean(input, "uses_pager")), - rationale: None, - citations: Vec::new(), - }, - "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest { - locations: input - .get("files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(file_location) - .collect(), - }), - "apply_file_diffs" => AIAgentActionType::RequestFileEdits { - file_edits: file_edits(input), - title: nonempty_string(input, "summary"), - }, - "grep" => AIAgentActionType::Grep { - queries: strings(input, "queries"), - path: string(input, "path"), - }, - "file_glob" => AIAgentActionType::FileGlob { - patterns: strings(input, "patterns"), - path: nonempty_string(input, "path"), - }, - "search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest { - query: string(input, "query"), - partial_paths: nonempty_strings(input, "path_filters"), - codebase_path: nonempty_string(input, "path"), - }), - "write_to_long_running_shell_command" => { - AIAgentActionType::WriteToLongRunningShellCommand { - block_id: string(input, "command_id").into(), - input: string(input, "input").into_bytes().into(), - mode: match input.get("mode").and_then(serde_json::Value::as_str) { - Some("line") => AIAgentPtyWriteMode::Line, - Some("block") => AIAgentPtyWriteMode::Block, - Some("raw") | Some(_) | None => AIAgentPtyWriteMode::Raw, - }, + "run_shell_command" => AIAgentActionType::RequestCommandOutput { + command: required_nonempty_string(input, "command")?, + is_read_only: Some(optional_boolean(input, "is_read_only")?.unwrap_or(false)), + is_risky: Some(optional_boolean(input, "is_risky")?.unwrap_or(false)), + wait_until_completion: optional_boolean(input, "wait_until_complete")? + .unwrap_or(false), + uses_pager: Some(optional_boolean(input, "uses_pager")?.unwrap_or(false)), + rationale: None, + citations: Vec::new(), + }, + "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest { + locations: required_array(input, "files")? + .iter() + .enumerate() + .map(|(index, file)| file_location(file, index)) + .collect::>()?, + }), + "apply_file_diffs" => AIAgentActionType::RequestFileEdits { + file_edits: file_edits(input)?, + title: Some(required_string(input, "summary")?), + }, + "grep" => AIAgentActionType::Grep { + queries: required_strings(input, "queries")?, + path: optional_string(input, "path")?.unwrap_or_default(), + }, + "file_glob" => AIAgentActionType::FileGlob { + patterns: required_strings(input, "patterns")?, + path: optional_string(input, "path")?.filter(|path| !path.is_empty()), + }, + "search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest { + query: required_string(input, "query")?, + partial_paths: optional_strings(input, "path_filters")? + .filter(|paths| !paths.is_empty()), + codebase_path: optional_string(input, "path")?.filter(|path| !path.is_empty()), + }), + "write_to_long_running_shell_command" => { + AIAgentActionType::WriteToLongRunningShellCommand { + block_id: required_nonempty_string(input, "command_id")?.into(), + input: required_string(input, "input")?.into_bytes().into(), + mode: match optional_string(input, "mode")?.as_deref() { + Some("line") => AIAgentPtyWriteMode::Line, + Some("block") => AIAgentPtyWriteMode::Block, + Some("raw") | None => AIAgentPtyWriteMode::Raw, + Some(mode) => { + return Err(format!( + "invalid field \"mode\": expected \"raw\", \"line\", or \"block\", got {mode:?}" + )); + } + }, + } } - } - "interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand { - block_id: string(input, "command_id").into(), - input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(), - mode: AIAgentPtyWriteMode::Raw, - }, - "read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput { - block_id: string(input, "command_id").into(), - delay: Some(ShellCommandDelay::Duration(Duration::from_secs( - input - .get("wait_seconds") - .and_then(serde_json::Value::as_u64) - .unwrap_or(2) - .min(crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS), - ))), - }, - "read_mcp_resource" => AIAgentActionType::ReadMCPResource { - server_id: uuid(input, "server_id"), - name: String::new(), - uri: nonempty_string(input, "uri"), - }, - "read_plan" | "read_documents" | "read_notebook" => { - AIAgentActionType::ReadDocuments(ReadDocumentsRequest { - document_ids: strings(input, "document_ids") - .into_iter() - .filter_map(|id| AIDocumentId::try_from(id).ok()) - .collect(), - }) - } - "create_plan" | "create_documents" | "create_notebook" => { - AIAgentActionType::CreateDocuments(CreateDocumentsRequest { - documents: input - .get("documents") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|document| { - Some(DocumentToCreate { - title: document.get("title")?.as_str()?.to_string(), - content: document.get("content")?.as_str()?.to_string(), - }) - }) - .collect(), - }) - } - "edit_plan" | "edit_documents" | "edit_notebook" => { - AIAgentActionType::EditDocuments(EditDocumentsRequest { - diffs: input - .get("diffs") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|diff| { - Some(DocumentDiff { - document_id: AIDocumentId::try_from(diff.get("document_id")?.as_str()?) - .ok()?, - search: string(diff, "search"), - replace: string(diff, "replace"), - }) - }) - .collect(), - }) - } - "run_agents" => AIAgentActionType::RunAgents(RunAgentsRequest { - summary: string(input, "summary"), - base_prompt: string(input, "base_prompt"), - skills: skill_references(input, skill_path_origin), - model_id: string(input, "model_id"), - harness_type: string(input, "harness_type"), - execution_mode: run_agents_execution_mode(input), - agent_run_configs: input - .get("agent_run_configs") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .map(|config| RunAgentsAgentRunConfig { - name: string(config, "name"), - prompt: string(config, "prompt"), - title: string(config, "title"), - }) - .collect(), - plan_id: string(input, "plan_id"), - harness_auth_secret_name: None, - }), - "start_agent" => AIAgentActionType::StartAgent { - version: StartAgentVersion::V1, - name: string(input, "name"), - prompt: string(input, "prompt"), - execution_mode: StartAgentExecutionMode::local_with_defaults(), - lifecycle_subscription: None, - }, - "send_message_to_agent" => AIAgentActionType::SendMessageToAgent { - addresses: vec![string(input, "agent_id")], - subject: String::new(), - message: string(input, "message"), - }, - "ask_user_question" => AIAgentActionType::AskUserQuestion { - questions: vec![AskUserQuestionItem { - question_id: Uuid::new_v4().to_string(), - question: string(input, "question"), - question_type: AskUserQuestionType::MultipleChoice { - is_multiselect: false, - options: strings(input, "options") + "interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand { + block_id: required_nonempty_string(input, "command_id")?.into(), + input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(), + mode: AIAgentPtyWriteMode::Raw, + }, + "read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput { + block_id: required_nonempty_string(input, "command_id")?.into(), + delay: Some(ShellCommandDelay::Duration(Duration::from_secs( + optional_bounded_u64( + input, + "wait_seconds", + crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS, + )? + .unwrap_or(2), + ))), + }, + "read_mcp_resource" => AIAgentActionType::ReadMCPResource { + server_id: Some(required_uuid(input, "server_id")?), + name: String::new(), + uri: Some(required_string(input, "uri")?), + }, + "read_plan" | "read_documents" | "read_notebook" => { + AIAgentActionType::ReadDocuments(ReadDocumentsRequest { + document_ids: required_strings(input, "document_ids")? .into_iter() - .enumerate() - .map(|(index, label)| AskUserQuestionOption { - label, - recommended: index == 0, + .map(|id| { + AIDocumentId::try_from(id.clone()).map_err(|_| { + format!("invalid document_ids entry: {id:?} is not a document ID") + }) }) - .collect(), - supports_other: true, - }, - }], - }, - "read_skill" => { - let skill = string(input, "skill"); - let skill = match input - .get("reference_type") - .and_then(serde_json::Value::as_str) - { - Some("bundled") => SkillReference::BundledSkillId(skill), - Some("path") | Some(_) | None => SkillReference::Path( - skill_path_origin - .location_for_path(skill) - .map_err(|error| error.to_string())?, - ), - }; - AIAgentActionType::ReadSkill(ReadSkillRequest { skill }) - } - "fetch_conversation" => AIAgentActionType::FetchConversation { - conversation_id: string(input, "conversation_id"), - }, - name if name.starts_with("mcp__") => { - let mut parts = name.splitn(3, "__"); - let _prefix = parts.next(); - let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok()); - let name = parts - .next() - .unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name)) - .to_string(); - AIAgentActionType::CallMCPTool { - server_id, - name, - input: input.clone(), + .collect::>()?, + }) + } + "create_plan" | "create_documents" | "create_notebook" => { + AIAgentActionType::CreateDocuments(CreateDocumentsRequest { + documents: required_array(input, "documents")? + .iter() + .enumerate() + .map(|(index, document)| { + require_object(document, &format!("documents[{index}]"))?; + Ok(DocumentToCreate { + title: required_string(document, "title")?, + content: required_string(document, "content")?, + }) + }) + .collect::>()?, + }) + } + "edit_plan" | "edit_documents" | "edit_notebook" => { + AIAgentActionType::EditDocuments(EditDocumentsRequest { + diffs: required_array(input, "diffs")? + .iter() + .enumerate() + .map(|(index, diff)| { + require_object(diff, &format!("diffs[{index}]"))?; + let document_id = required_string(diff, "document_id")?; + Ok(DocumentDiff { + document_id: AIDocumentId::try_from(document_id.clone()) + .map_err(|_| format!("invalid document_id: {document_id:?}"))?, + search: required_string(diff, "search")?, + replace: required_string(diff, "replace")?, + }) + }) + .collect::>()?, + }) + } + "run_agents" => AIAgentActionType::RunAgents(RunAgentsRequest { + summary: required_nonempty_string(input, "summary")?, + base_prompt: optional_string(input, "base_prompt")?.unwrap_or_default(), + skills: skill_references(input, skill_path_origin)?, + model_id: optional_string(input, "model_id")?.unwrap_or_default(), + harness_type: optional_string(input, "harness_type")?.unwrap_or_default(), + execution_mode: run_agents_execution_mode(input)?, + agent_run_configs: nonempty_required_array(input, "agent_run_configs")? + .iter() + .enumerate() + .map(|(index, config)| { + require_object(config, &format!("agent_run_configs[{index}]"))?; + Ok(RunAgentsAgentRunConfig { + name: required_nonempty_string(config, "name")?, + prompt: required_nonempty_string(config, "prompt")?, + title: optional_string(config, "title")?.unwrap_or_default(), + }) + }) + .collect::>()?, + plan_id: optional_string(input, "plan_id")?.unwrap_or_default(), + harness_auth_secret_name: None, + }), + "start_agent" => AIAgentActionType::StartAgent { + version: StartAgentVersion::V1, + name: required_nonempty_string(input, "name")?, + prompt: required_nonempty_string(input, "prompt")?, + execution_mode: StartAgentExecutionMode::local_with_defaults(), + lifecycle_subscription: None, + }, + "send_message_to_agent" => AIAgentActionType::SendMessageToAgent { + addresses: vec![required_string(input, "agent_id")?], + subject: String::new(), + message: required_string(input, "message")?, + }, + "transfer_shell_command_control_to_user" => { + AIAgentActionType::TransferShellCommandControlToUser { + reason: required_nonempty_string(input, "reason")?, + } + } + "wait_for_events" => AIAgentActionType::WaitForEvents { + tool_call_id: call.id.clone(), + idle_timeout_seconds: optional_nonnegative_i32(input, "idle_timeout_seconds")? + .unwrap_or(0), + }, + "ask_user_question" => AIAgentActionType::AskUserQuestion { + questions: vec![AskUserQuestionItem { + question_id: Uuid::new_v4().to_string(), + question: required_string(input, "question")?, + question_type: AskUserQuestionType::MultipleChoice { + is_multiselect: false, + options: optional_strings(input, "options")? + .unwrap_or_default() + .into_iter() + .enumerate() + .map(|(index, label)| AskUserQuestionOption { + label, + recommended: index == 0, + }) + .collect(), + supports_other: true, + }, + }], + }, + "read_skill" => { + let skill = required_string(input, "skill")?; + let skill = match required_string(input, "reference_type")?.as_str() { + "bundled" => SkillReference::BundledSkillId(skill), + "path" => SkillReference::Path( + skill_path_origin + .location_for_path(skill) + .map_err(|error| error.to_string())?, + ), + reference_type => { + return Err(format!( + "invalid reference_type: expected \"path\" or \"bundled\", got {reference_type:?}" + )); + } + }; + AIAgentActionType::ReadSkill(ReadSkillRequest { skill }) + } + "fetch_conversation" => AIAgentActionType::FetchConversation { + conversation_id: required_string(input, "conversation_id")?, + }, + name if name.starts_with("mcp__") => { + let mut parts = name.splitn(3, "__"); + let _prefix = parts.next(); + let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok()); + let name = parts + .next() + .unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name)) + .to_string(); + AIAgentActionType::CallMCPTool { + server_id, + name, + input: input.clone(), + } } - } name => return Err(format!("unsupported Rig tool proposal: {name}")), } }; @@ -242,155 +268,309 @@ pub(crate) fn action_from_tool_call( }) } -fn string(input: &serde_json::Value, key: &str) -> String { +fn require_object(input: &serde_json::Value, field: &str) -> Result<(), String> { input - .get(key) - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string() + .is_object() + .then_some(()) + .ok_or_else(|| format!("invalid {field}: expected an object")) } -fn nonempty_string(input: &serde_json::Value, key: &str) -> Option { - let value = string(input, key); - (!value.is_empty()).then_some(value) -} - -fn boolean(input: &serde_json::Value, key: &str) -> bool { +fn required_string(input: &serde_json::Value, key: &str) -> Result { input .get(key) - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) -} - -fn strings(input: &serde_json::Value, key: &str) -> Vec { - input - .get(key) - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(serde_json::Value::as_str) + .ok_or_else(|| format!("missing required field {key:?}"))? + .as_str() .map(ToOwned::to_owned) + .ok_or_else(|| format!("invalid field {key:?}: expected a string")) +} + +fn required_nonempty_string(input: &serde_json::Value, key: &str) -> Result { + let value = required_string(input, key)?; + if value.trim().is_empty() { + return Err(format!( + "invalid field {key:?}: expected a non-empty string" + )); + } + Ok(value) +} + +fn optional_string(input: &serde_json::Value, key: &str) -> Result, String> { + input + .get(key) + .map(|value| { + value + .as_str() + .map(ToOwned::to_owned) + .ok_or_else(|| format!("invalid field {key:?}: expected a string")) + }) + .transpose() +} + +fn required_array<'a>( + input: &'a serde_json::Value, + key: &str, +) -> Result<&'a Vec, String> { + input + .get(key) + .ok_or_else(|| format!("missing required field {key:?}"))? + .as_array() + .ok_or_else(|| format!("invalid field {key:?}: expected an array")) +} + +fn nonempty_required_array<'a>( + input: &'a serde_json::Value, + key: &str, +) -> Result<&'a Vec, String> { + let values = required_array(input, key)?; + if values.is_empty() { + return Err(format!("invalid field {key:?}: expected at least one item")); + } + Ok(values) +} + +fn required_strings(input: &serde_json::Value, key: &str) -> Result, String> { + strings_from_array(required_array(input, key)?, key) +} + +fn optional_strings(input: &serde_json::Value, key: &str) -> Result>, String> { + input + .get(key) + .map(|value| { + let values = value + .as_array() + .ok_or_else(|| format!("invalid field {key:?}: expected an array"))?; + strings_from_array(values, key) + }) + .transpose() +} + +fn strings_from_array(values: &[serde_json::Value], key: &str) -> Result, String> { + values + .iter() + .enumerate() + .map(|(index, value)| { + value + .as_str() + .map(ToOwned::to_owned) + .ok_or_else(|| format!("invalid {key}[{index}]: expected a string")) + }) .collect() } -fn nonempty_strings(input: &serde_json::Value, key: &str) -> Option> { - let values = strings(input, key); - (!values.is_empty()).then_some(values) +fn required_uuid(input: &serde_json::Value, key: &str) -> Result { + let value = required_string(input, key)?; + Uuid::parse_str(&value).map_err(|_| format!("invalid field {key:?}: expected a UUID")) } -fn uuid(input: &serde_json::Value, key: &str) -> Option { +fn optional_boolean(input: &serde_json::Value, key: &str) -> Result, String> { input .get(key) - .and_then(serde_json::Value::as_str) - .and_then(|value| Uuid::parse_str(value).ok()) + .map(|value| { + value + .as_bool() + .ok_or_else(|| format!("invalid field {key:?}: expected a boolean")) + }) + .transpose() +} + +fn optional_bounded_u64( + input: &serde_json::Value, + key: &str, + maximum: u64, +) -> Result, String> { + input + .get(key) + .map(|value| { + let value = value + .as_u64() + .ok_or_else(|| format!("invalid field {key:?}: expected a non-negative integer"))?; + if value > maximum { + return Err(format!( + "invalid field {key:?}: expected an integer no greater than {maximum}" + )); + } + Ok(value) + }) + .transpose() +} + +fn optional_nonnegative_i32(input: &serde_json::Value, key: &str) -> Result, String> { + input + .get(key) + .map(|value| { + value + .as_i64() + .and_then(|value| i32::try_from(value).ok()) + .filter(|value| *value >= 0) + .ok_or_else(|| { + format!("invalid field {key:?}: expected a non-negative 32-bit integer") + }) + }) + .transpose() } fn skill_references( input: &serde_json::Value, skill_path_origin: &SkillPathOrigin, -) -> Vec { - input - .get("skills") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|skill| { - let reference = string(skill, "skill"); - if reference.is_empty() { - return None; - } - match skill - .get("reference_type") - .and_then(serde_json::Value::as_str) - { - Some("bundled") => Some(SkillReference::BundledSkillId(reference)), - Some("path") | Some(_) | None => skill_path_origin +) -> Result, String> { + let Some(skills) = optional_array(input, "skills")? else { + return Ok(Vec::new()); + }; + skills + .iter() + .enumerate() + .map(|(index, skill)| { + require_object(skill, &format!("skills[{index}]"))?; + let reference = required_string(skill, "skill")?; + match required_string(skill, "reference_type")?.as_str() { + "bundled" => Ok(SkillReference::BundledSkillId(reference)), + "path" => skill_path_origin .location_for_path(reference) - .ok() - .map(SkillReference::Path), + .map(SkillReference::Path) + .map_err(|error| error.to_string()), + reference_type => Err(format!( + "invalid skills[{index}].reference_type: expected \"path\" or \"bundled\", got {reference_type:?}" + )), } }) .collect() } -fn run_agents_execution_mode(input: &serde_json::Value) -> RunAgentsExecutionMode { +fn run_agents_execution_mode(input: &serde_json::Value) -> Result { let Some(execution_mode) = input.get("execution_mode") else { - return RunAgentsExecutionMode::Local; + return Ok(RunAgentsExecutionMode::Local); }; - let mode_type = execution_mode - .get("type") - .and_then(serde_json::Value::as_str) - .or_else(|| execution_mode.as_str()); - match mode_type { - Some("remote") => RunAgentsExecutionMode::Remote { - environment_id: string(execution_mode, "environment_id"), - worker_host: string(execution_mode, "worker_host"), - computer_use_enabled: boolean(execution_mode, "computer_use_enabled"), - }, - Some("local") | Some(_) | None => RunAgentsExecutionMode::Local, + require_object(execution_mode, "execution_mode")?; + match optional_string(execution_mode, "type")?.as_deref() { + Some("remote") => Ok(RunAgentsExecutionMode::Remote { + environment_id: optional_string(execution_mode, "environment_id")?.unwrap_or_default(), + worker_host: optional_string(execution_mode, "worker_host")?.unwrap_or_default(), + computer_use_enabled: optional_boolean(execution_mode, "computer_use_enabled")? + .unwrap_or(false), + }), + Some("local") | None => { + optional_string(execution_mode, "environment_id")?; + optional_string(execution_mode, "worker_host")?; + optional_boolean(execution_mode, "computer_use_enabled")?; + Ok(RunAgentsExecutionMode::Local) + } + Some(mode_type) => Err(format!( + "invalid execution_mode.type: expected \"local\" or \"remote\", got {mode_type:?}" + )), } } -fn file_location(file: &serde_json::Value) -> Option { +fn file_location(file: &serde_json::Value, file_index: usize) -> Result { if let Some(name) = file.as_str() { - return Some(FileLocations { + return Ok(FileLocations { name: name.to_string(), lines: Vec::new(), }); } - let name = file - .get("path") - .or_else(|| file.get("name"))? - .as_str()? - .to_string(); - let lines = file - .get("line_ranges") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|range| { - let start = usize::try_from(range.get("start")?.as_u64()?).ok()?; - let end = usize::try_from(range.get("end")?.as_u64()?).ok()?; - (start > 0 && end >= start).then_some(start..end) - }) - .collect(); - Some(FileLocations { name, lines }) + require_object(file, &format!("files[{file_index}]"))?; + let name = required_string(file, "path")?; + let lines = match file.get("line_ranges") { + None => Vec::new(), + Some(value) => value + .as_array() + .ok_or_else(|| format!("invalid files[{file_index}].line_ranges: expected an array"))? + .iter() + .enumerate() + .map(|(range_index, range)| { + require_object( + range, + &format!("files[{file_index}].line_ranges[{range_index}]"), + )?; + let start = required_line_number(range, "start", file_index, range_index)?; + let inclusive_end = required_line_number(range, "end", file_index, range_index)?; + if inclusive_end < start { + return Err(format!( + "invalid files[{file_index}].line_ranges[{range_index}]: end must be greater than or equal to start" + )); + } + let exclusive_end = inclusive_end.checked_add(1).ok_or_else(|| { + format!( + "invalid files[{file_index}].line_ranges[{range_index}].end: inclusive end is too large" + ) + })?; + Ok(start..exclusive_end) + }) + .collect::>()?, + }; + Ok(FileLocations { name, lines }) } -fn file_edits(input: &serde_json::Value) -> Vec { - let diffs = input - .get("diffs") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .map(|diff| { - FileEdit::Edit(ParsedDiff::StrReplaceEdit { - file: nonempty_string(diff, "file_path"), - search: nonempty_string(diff, "search"), - replace: nonempty_string(diff, "replace"), - }) - }); - let creates = input - .get("new_files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .map(|file| FileEdit::Create { - file: nonempty_string(file, "file_path"), - content: nonempty_string(file, "content"), - }); - let deletes = input - .get("deleted_files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .map(|file| FileEdit::Delete { - file: file +fn required_line_number( + range: &serde_json::Value, + key: &str, + file_index: usize, + range_index: usize, +) -> Result { + let value = range + .get(key) + .ok_or_else(|| format!("missing required field {key:?}"))? + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0) + .ok_or_else(|| { + format!( + "invalid files[{file_index}].line_ranges[{range_index}].{key}: expected a positive integer" + ) + })?; + Ok(value) +} + +fn file_edits(input: &serde_json::Value) -> Result, String> { + let mut edits = Vec::new(); + if let Some(diffs) = optional_array(input, "diffs")? { + for (index, diff) in diffs.iter().enumerate() { + require_object(diff, &format!("diffs[{index}]"))?; + edits.push(FileEdit::Edit(ParsedDiff::StrReplaceEdit { + file: Some(required_string(diff, "file_path")?), + search: Some(required_string(diff, "search")?), + replace: Some(required_string(diff, "replace")?), + })); + } + } + if let Some(files) = optional_array(input, "new_files")? { + for (index, file) in files.iter().enumerate() { + require_object(file, &format!("new_files[{index}]"))?; + edits.push(FileEdit::Create { + file: Some(required_string(file, "file_path")?), + content: Some(required_string(file, "content")?), + }); + } + } + if let Some(files) = optional_array(input, "deleted_files")? { + for (index, file) in files.iter().enumerate() { + let path = file .as_str() - .map(ToOwned::to_owned) - .or_else(|| nonempty_string(file, "file_path")), - }); - diffs.chain(creates).chain(deletes).collect() + .ok_or_else(|| format!("invalid deleted_files[{index}]: expected a string"))?; + edits.push(FileEdit::Delete { + file: Some(path.to_owned()), + }); + } + } + if edits.is_empty() { + return Err( + "invalid file edits: expected at least one diff, new file, or deleted file".to_string(), + ); + } + Ok(edits) +} + +fn optional_array<'a>( + input: &'a serde_json::Value, + key: &str, +) -> Result>, String> { + input + .get(key) + .map(|value| { + value + .as_array() + .ok_or_else(|| format!("invalid field {key:?}: expected an array")) + }) + .transpose() } #[cfg(test)] diff --git a/app/src/ai/runtime/rig_tool_tests.rs b/app/src/ai/runtime/rig_tool_tests.rs index 8bb17192..8445ef84 100644 --- a/app/src/ai/runtime/rig_tool_tests.rs +++ b/app/src/ai/runtime/rig_tool_tests.rs @@ -46,6 +46,26 @@ fn shell_calls_become_domain_actions_without_a_proto_round_trip() { )); } +#[test] +fn transfer_control_calls_become_domain_actions() { + let action = action_from_tool_call( + "task-1", + &call( + "transfer_shell_command_control_to_user", + serde_json::json!({"reason": "The command needs interactive input"}), + ), + &SkillPathOrigin::Local, + &HashMap::new(), + ) + .unwrap(); + + assert!(matches!( + action.action, + AIAgentActionType::TransferShellCommandControlToUser { reason } + if reason == "The command needs interactive input" + )); +} + #[test] fn create_plan_calls_become_document_actions() { let action = action_from_tool_call( @@ -72,6 +92,227 @@ fn create_plan_calls_become_document_actions() { assert_eq!(request.documents[0].content, "# Implementation plan"); } +#[test] +fn read_files_converts_advertised_inclusive_ranges_to_half_open_ranges() { + let action = action_from_tool_call( + "task-1", + &call( + "read_files", + serde_json::json!({ + "files": [{ + "path": "/tmp/example.rs", + "line_ranges": [ + {"start": 1, "end": 1}, + {"start": 10, "end": 25} + ] + }] + }), + ), + &SkillPathOrigin::Local, + &HashMap::new(), + ) + .unwrap(); + + let AIAgentActionType::ReadFiles(request) = action.action else { + panic!("expected read-files action"); + }; + assert_eq!(request.locations[0].lines, vec![1..2, 10..26]); +} + +#[test] +fn known_tools_reject_malformed_required_inputs() { + let cases = [ + ("read_files", serde_json::json!({}), "files"), + ( + "read_files", + serde_json::json!({"files": "not-an-array"}), + "expected an array", + ), + ( + "read_files", + serde_json::json!({"files": [{"path": "/tmp/a", "line_ranges": [{"start": 0, "end": 1}]}]}), + "positive integer", + ), + ( + "read_files", + serde_json::json!({"files": [{"path": "/tmp/a", "line_ranges": [{"start": 3, "end": 2}]}]}), + "greater than or equal", + ), + ( + "read_files", + serde_json::json!({"files": [{"path": "/tmp/a", "line_ranges": [{"start": 1, "end": u64::MAX}]}]}), + "inclusive end is too large", + ), + ( + "grep", + serde_json::json!({"queries": ["ok", 7]}), + "queries[1]", + ), + ( + "file_glob", + serde_json::json!({"patterns": false}), + "expected an array", + ), + ( + "search_codebase", + serde_json::json!({"query": 42}), + "expected a string", + ), + ( + "apply_file_diffs", + serde_json::json!({"summary": "edit", "diffs": [{"file_path": "/tmp/a", "search": "x"}]}), + "replace", + ), + ( + "apply_file_diffs", + serde_json::json!({"summary": "Nothing to do"}), + "at least one diff", + ), + ( + "run_shell_command", + serde_json::json!({"command": 42}), + "expected a string", + ), + ( + "run_shell_command", + serde_json::json!({"command": " "}), + "non-empty string", + ), + ( + "run_shell_command", + serde_json::json!({"command": "pwd", "is_read_only": "yes"}), + "expected a boolean", + ), + ( + "write_to_long_running_shell_command", + serde_json::json!({"command_id": "command-1", "input": "yes", "mode": "words"}), + "mode", + ), + ( + "interrupt_shell_command", + serde_json::json!({}), + "command_id", + ), + ( + "read_shell_command_output", + serde_json::json!({"command_id": 12}), + "expected a string", + ), + ( + "read_shell_command_output", + serde_json::json!({"command_id": "command-1", "wait_seconds": 11}), + "no greater than", + ), + ( + "run_agents", + serde_json::json!({"summary": "Investigate", "agent_run_configs": []}), + "at least one item", + ), + ( + "run_agents", + serde_json::json!({"summary": "Investigate", "agent_run_configs": [{"name": "one"}]}), + "prompt", + ), + ( + "run_agents", + serde_json::json!({"summary": "Investigate", "agent_run_configs": [{"name": "one", "prompt": "Inspect"}], "execution_mode": {"type": "other"}}), + "execution_mode.type", + ), + ( + "run_agents", + serde_json::json!({"summary": "Investigate", "agent_run_configs": [{"name": "one", "prompt": "Inspect"}], "skills": [{"skill": "test", "reference_type": "other"}]}), + "skills[0].reference_type", + ), + ( + "start_agent", + serde_json::json!({"name": "worker"}), + "prompt", + ), + ( + "transfer_shell_command_control_to_user", + serde_json::json!({"reason": false}), + "expected a string", + ), + ( + "wait_for_events", + serde_json::json!({"idle_timeout_seconds": -1}), + "non-negative", + ), + ( + "create_plan", + serde_json::json!({"documents": [{"title": "Plan"}]}), + "content", + ), + ( + "read_skill", + serde_json::json!({"skill": "/tmp/SKILL.md", "reference_type": "other"}), + "reference_type", + ), + ( + "fetch_conversation", + serde_json::json!({"conversation_id": null}), + "expected a string", + ), + ]; + + for (name, arguments, expected_error) in cases { + let error = action_from_tool_call( + "task-1", + &call(name, arguments), + &SkillPathOrigin::Local, + &HashMap::new(), + ) + .unwrap_err(); + assert!( + error.contains(expected_error), + "{name} error {error:?} did not contain {expected_error:?}" + ); + } +} + +#[test] +fn known_tools_preserve_legitimate_optional_defaults() { + let cases = [ + ("grep", serde_json::json!({"queries": ["needle"]})), + ("file_glob", serde_json::json!({"patterns": ["**/*.rs"]})), + ( + "ask_user_question", + serde_json::json!({"question": "Continue?"}), + ), + ( + "apply_file_diffs", + serde_json::json!({"summary": "Create file", "new_files": [{"file_path": "/tmp/new", "content": ""}]}), + ), + ("run_shell_command", serde_json::json!({"command": "pwd"})), + ( + "write_to_long_running_shell_command", + serde_json::json!({"command_id": "command-1", "input": ""}), + ), + ( + "read_shell_command_output", + serde_json::json!({"command_id": "command-1"}), + ), + ( + "run_agents", + serde_json::json!({ + "summary": "Investigate", + "agent_run_configs": [{"name": "worker", "prompt": "Inspect"}] + }), + ), + ("wait_for_events", serde_json::json!({})), + ]; + + for (name, arguments) in cases { + action_from_tool_call( + "task-1", + &call(name, arguments), + &SkillPathOrigin::Local, + &HashMap::new(), + ) + .unwrap_or_else(|error| panic!("{name} rejected optional defaults: {error}")); + } +} + #[test] fn edit_calls_preserve_file_edits_in_the_domain_model() { let action = action_from_tool_call( diff --git a/app/src/ai/tool_diagnostics.rs b/app/src/ai/tool_diagnostics.rs new file mode 100644 index 00000000..6e29706e --- /dev/null +++ b/app/src/ai/tool_diagnostics.rs @@ -0,0 +1,46 @@ +use std::backtrace::Backtrace; +use std::ffi::OsStr; +use std::sync::OnceLock; + +const ENV_VAR: &str = "GALAXY_TOOL_DIAGNOSTICS"; + +fn env_value_is_enabled(value: Option<&OsStr>) -> bool { + value + .and_then(OsStr::to_str) + .is_some_and(|value| value == "1" || value.eq_ignore_ascii_case("true")) +} + +pub(crate) fn is_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| env_value_is_enabled(std::env::var_os(ENV_VAR).as_deref())) +} + +pub(crate) fn capture_backtrace() -> Option { + is_enabled().then(Backtrace::force_capture) +} + +macro_rules! tool_debug { + ($($arg:tt)*) => { + if $crate::ai::tool_diagnostics::is_enabled() { + log::debug!("[tool-debug] {}", format_args!($($arg)*)); + } + }; +} + +pub(crate) use tool_debug; + +#[cfg(test)] +mod tests { + use std::ffi::OsStr; + + use super::env_value_is_enabled; + + #[test] + fn diagnostic_env_accepts_only_explicit_true_values() { + assert!(env_value_is_enabled(Some(OsStr::new("1")))); + assert!(env_value_is_enabled(Some(OsStr::new("TRUE")))); + assert!(!env_value_is_enabled(Some(OsStr::new("0")))); + assert!(!env_value_is_enabled(Some(OsStr::new("yes")))); + assert!(!env_value_is_enabled(None)); + } +} diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 2b475762..f9f870b4 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -7254,8 +7254,22 @@ impl TerminalView { event: &BlocklistAIActionEvent, ctx: &mut ViewContext, ) { + let event_matches_active_conversation = || { + let Some(event_conversation_id) = event.conversation_id() else { + return true; + }; + self.model + .lock() + .block_list() + .active_block() + .ai_conversation_id() + == Some(event_conversation_id) + }; match event { BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { .. } => { + if !event_matches_active_conversation() { + return; + } let is_agent_in_control = self .model .lock() @@ -7267,14 +7281,20 @@ impl TerminalView { } } BlocklistAIActionEvent::ExecutingAction { .. } => { - self.redetermine_terminal_focus(ctx); + if event_matches_active_conversation() { + self.redetermine_terminal_focus(ctx); + } ctx.notify(); } - BlocklistAIActionEvent::FinishedAction { action_id, .. } => { + BlocklistAIActionEvent::FinishedAction { + action_id, + conversation_id, + .. + } => { // Refresh git line changes when files are potentially updated by an action let action_result = action_model .as_ref(ctx) - .get_action_result(action_id) + .get_action_result(*conversation_id, action_id) .cloned(); let maybe_modified_files = action_result @@ -7644,7 +7664,11 @@ impl TerminalView { drop(model); self.cli_subagent_controller.update(ctx, |controller, _| { - controller.track_requested_command(&block_id, action_id); + controller.track_requested_command( + &block_id, + parent_conversation_id, + action_id, + ); }); ctx.emit(Event::ExecuteCommand(ExecuteCommandEvent { @@ -7672,10 +7696,23 @@ impl TerminalView { ShellCommandExecutorEvent::WriteToPty { input, mode } => { self.write_agent_bytes_to_pty(input.to_vec(), mode, ctx); } - ShellCommandExecutorEvent::CancelExecution => { + ShellCommandExecutorEvent::CancelExecution { action_id } => { // We need to manually invoke ctrl-c to terminate the running command because the // user's ctrl-c was directed to the AIBlock instead of the command's shell block. - self.ctrl_c(ctx); + let is_exact_active_command = self + .model + .lock() + .block_list() + .active_block() + .requested_command_action_id() + .is_some_and(|requested_id| requested_id == action_id); + if is_exact_active_command { + self.ctrl_c(ctx); + } else { + log::warn!( + "Refusing to interrupt active terminal command for stale requested action {action_id}" + ); + } } ShellCommandExecutorEvent::TransferControlToUser { reason, .. } => { // Transfer control of the long-running command to the user. diff --git a/app/src/terminal/view/context_menu.rs b/app/src/terminal/view/context_menu.rs index cd2f2b8d..292b6099 100644 --- a/app/src/terminal/view/context_menu.rs +++ b/app/src/terminal/view/context_menu.rs @@ -178,8 +178,10 @@ impl TerminalView { let mut result = Vec::new(); for exchange in conversation.root_task_exchanges() { - let formatted_exchange = - exchange.format_for_copy(Some(self.ai_action_model.as_ref(ctx))); + let formatted_exchange = exchange.format_for_copy_for_conversation( + Some(self.ai_action_model.as_ref(ctx)), + Some(conversation_id), + ); if !formatted_exchange.is_empty() { result.push(formatted_exchange); } diff --git a/app/src/terminal/view/load_ai_conversation.rs b/app/src/terminal/view/load_ai_conversation.rs index 6b46cd7a..9ab4a6d3 100644 --- a/app/src/terminal/view/load_ai_conversation.rs +++ b/app/src/terminal/view/load_ai_conversation.rs @@ -364,7 +364,9 @@ impl TerminalView { }) => { if let Some(result) = self.ai_action_model.read(ctx, |action_model, _| { - action_model.get_action_result(&action.id).cloned() + action_model + .get_action_result(conversation_id, &action.id) + .cloned() }) { if let AIAgentActionResultType::CreateDocuments( @@ -402,7 +404,9 @@ impl TerminalView { AIAgentActionType::EditDocuments { .. } => { if let Some(result) = self.ai_action_model.read(ctx, |action_model, _| { - action_model.get_action_result(&action.id).cloned() + action_model + .get_action_result(conversation_id, &action.id) + .cloned() }) { if let AIAgentActionResultType::EditDocuments( @@ -449,8 +453,10 @@ impl TerminalView { for conversation in &conversations { self.ai_action_model.update(ctx, |action_model, _ctx| { - action_model - .restore_action_results_from_exchanges(exchanges_for_blocklist(conversation)); + action_model.restore_action_results_from_exchanges( + conversation.id(), + exchanges_for_blocklist(conversation), + ); }); } diff --git a/crates/ai/src/agent/action_result/convert.rs b/crates/ai/src/agent/action_result/convert.rs index d495805e..29238e71 100644 --- a/crates/ai/src/agent/action_result/convert.rs +++ b/crates/ai/src/agent/action_result/convert.rs @@ -94,6 +94,25 @@ impl TryFrom for api::request::input::tool_call_resu }, ), ), + RequestCommandOutputResult::ExecutionError { command, message } => Ok( + api::request::input::tool_call_result::Result::RunShellCommand( + #[allow(deprecated)] + api::RunShellCommandResult { + command, + output: Default::default(), + exit_code: Default::default(), + result: Some(api::run_shell_command_result::Result::CommandFinished( + api::ShellCommandFinished { + command_id: String::new(), + output: format!("Command was not executed: {message}"), + exit_code: 1, + start_ts: None, + finish_ts: None, + }, + )), + }, + ), + ), RequestCommandOutputResult::Denylisted { command } => { #[allow(deprecated)] @@ -1551,6 +1570,14 @@ impl From for api::run_agents_result::AgentOutcome { api::run_agents_result::LaunchedAgent { agent_id }, ) } + // The legacy wire schema has no completed-child shape. Preserve the child identity; + // direct-provider history retains the richer local result and output. + RunAgentsAgentOutcomeKind::Completed { + agent_id, + output: _, + } => api::run_agents_result::agent_outcome::Result::Launched( + api::run_agents_result::LaunchedAgent { agent_id }, + ), RunAgentsAgentOutcomeKind::Failed { error } => { api::run_agents_result::agent_outcome::Result::Failed( api::run_agents_result::FailedAgent { error }, diff --git a/crates/ai/src/agent/action_result/convert_tests.rs b/crates/ai/src/agent/action_result/convert_tests.rs index 6adcbdda..d24ef45c 100644 --- a/crates/ai/src/agent/action_result/convert_tests.rs +++ b/crates/ai/src/agent/action_result/convert_tests.rs @@ -28,3 +28,22 @@ fn ask_user_question_skipped_by_auto_approve_converts_to_skipped_answers() { Some(AskUserQuestionAnswer::Skipped(())) )); } + +#[test] +fn completed_run_agents_child_converts_to_legacy_launched_wire_outcome() { + let outcome = RunAgentsAgentOutcome { + name: "research".to_string(), + kind: RunAgentsAgentOutcomeKind::Completed { + agent_id: "child-1".to_string(), + output: "local output".to_string(), + }, + }; + + let converted = api::run_agents_result::AgentOutcome::from(outcome); + assert!(matches!( + converted.result, + Some(api::run_agents_result::agent_outcome::Result::Launched( + api::run_agents_result::LaunchedAgent { agent_id } + )) if agent_id == "child-1" + )); +} diff --git a/crates/ai/src/agent/action_result/mod.rs b/crates/ai/src/agent/action_result/mod.rs index 59dc39fd..be4d799f 100644 --- a/crates/ai/src/agent/action_result/mod.rs +++ b/crates/ai/src/agent/action_result/mod.rs @@ -165,6 +165,7 @@ impl AIAgentActionResultType { None, ), RequestCommandOutputResult::CancelledBeforeExecution + | RequestCommandOutputResult::ExecutionError { .. } | RequestCommandOutputResult::Denylisted { .. } => result.to_string(), }, Self::WriteToLongRunningShellCommand(result) => match result { @@ -418,6 +419,8 @@ pub enum RequestCommandOutputResult { /// A running command canceled via ctrl-c /// would have Completed result with exit code 130. CancelledBeforeExecution, + /// The command could not start because the terminal was unavailable for execution. + ExecutionError { command: String, message: String }, /// The command was denied because it was present on the denylist. Denylisted { command: String }, } @@ -427,14 +430,16 @@ impl RequestCommandOutputResult { match self { Self::Completed { exit_code, .. } => exit_code.was_successful(), Self::LongRunningCommandSnapshot { .. } => true, - Self::CancelledBeforeExecution | Self::Denylisted { .. } => false, + Self::CancelledBeforeExecution + | Self::ExecutionError { .. } + | Self::Denylisted { .. } => false, } } pub fn failed(&self) -> bool { match self { Self::Completed { exit_code, .. } => !exit_code.was_successful(), - Self::Denylisted { .. } => true, + Self::ExecutionError { .. } | Self::Denylisted { .. } => true, Self::CancelledBeforeExecution | Self::LongRunningCommandSnapshot { .. } => false, } } @@ -444,6 +449,7 @@ impl RequestCommandOutputResult { match self { Self::Completed { command, .. } | Self::LongRunningCommandSnapshot { command, .. } + | Self::ExecutionError { command, .. } | Self::Denylisted { command } => command.clone(), Self::CancelledBeforeExecution => "cancelled".to_string(), } @@ -473,6 +479,9 @@ impl Display for RequestCommandOutputResult { RequestCommandOutputResult::CancelledBeforeExecution => { write!(f, "Command output cancelled") } + RequestCommandOutputResult::ExecutionError { command, message } => { + write!(f, "Command '{command}' could not be executed: {message}") + } RequestCommandOutputResult::Denylisted { .. } => { write!(f, "Command output was on denylist") } @@ -1042,7 +1051,9 @@ impl AIAgentActionResultType { | TransferShellCommandControlToUserResult::CommandFinished { .. }, ) => true, Self::AskUserQuestion(AskUserQuestionResult::Success { .. }) => true, - Self::RunAgents(RunAgentsResult::Launched { .. }) => true, + Self::RunAgents(RunAgentsResult::Launched { agents, .. }) => agents + .iter() + .any(|agent| !matches!(agent.kind, RunAgentsAgentOutcomeKind::Failed { .. })), Self::WaitForEvents(WaitForEventsResult::Completed) => true, _ => false, } @@ -1076,6 +1087,12 @@ impl AIAgentActionResultType { | Self::RunAgents(RunAgentsResult::Failure { .. } | RunAgentsResult::Denied { .. }) => { true } + Self::RunAgents(RunAgentsResult::Launched { agents, .. }) => { + !agents.is_empty() + && agents + .iter() + .all(|agent| matches!(agent.kind, RunAgentsAgentOutcomeKind::Failed { .. })) + } _ => false, } } @@ -1627,6 +1644,7 @@ pub struct RunAgentsAgentOutcome { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub enum RunAgentsAgentOutcomeKind { Launched { agent_id: String }, + Completed { agent_id: String, output: String }, Failed { error: String }, } @@ -1654,6 +1672,17 @@ impl RunAgentsResult { "computer_use_enabled": computer_use_enabled, }), }; + let children_completed = agents.iter().all(|agent| { + matches!( + agent.kind, + RunAgentsAgentOutcomeKind::Completed { .. } + | RunAgentsAgentOutcomeKind::Failed { .. } + ) + }); + let all_failed = !agents.is_empty() + && agents.iter().all(|agent| { + matches!(agent.kind, RunAgentsAgentOutcomeKind::Failed { .. }) + }); let agents = agents .iter() .map(|agent| match &agent.kind { @@ -1662,6 +1691,14 @@ impl RunAgentsResult { "status": "launched", "agent_id": agent_id, }), + RunAgentsAgentOutcomeKind::Completed { agent_id, output } => { + serde_json::json!({ + "name": agent.name, + "status": "completed", + "agent_id": agent_id, + "output": output, + }) + } RunAgentsAgentOutcomeKind::Failed { error } => serde_json::json!({ "name": agent.name, "status": "failed", @@ -1670,9 +1707,13 @@ impl RunAgentsResult { }) .collect::>(); serde_json::json!({ - "status": "launched", - "completion_state": "children_running", - "instruction": "Child agents have only been launched, not completed. Wait for child-agent updates before repeating this work or reporting final results.", + "status": if all_failed { "failure" } else { "launched" }, + "completion_state": if children_completed { "children_completed" } else { "children_running" }, + "instruction": if children_completed { + "Child agents reached terminal states. Use their structured outputs and errors to complete the task." + } else { + "Child agents have only been launched, not completed. Wait for child-agent updates before repeating this work or reporting final results." + }, "model_id": model_id, "harness_type": harness_type, "execution_mode": execution_mode, @@ -1701,7 +1742,13 @@ impl Display for RunAgentsResult { RunAgentsResult::Launched { agents, .. } => { let launched = agents .iter() - .filter(|a| matches!(a.kind, RunAgentsAgentOutcomeKind::Launched { .. })) + .filter(|a| { + matches!( + a.kind, + RunAgentsAgentOutcomeKind::Launched { .. } + | RunAgentsAgentOutcomeKind::Completed { .. } + ) + }) .count(); write!( f, diff --git a/crates/ai/src/agent/action_result/mod_tests.rs b/crates/ai/src/agent/action_result/mod_tests.rs index 3fe3909d..b661f522 100644 --- a/crates/ai/src/agent/action_result/mod_tests.rs +++ b/crates/ai/src/agent/action_result/mod_tests.rs @@ -1,8 +1,22 @@ use super::{ - AIAgentActionResultType, RunAgentsAgentOutcome, RunAgentsAgentOutcomeKind, - RunAgentsLaunchedExecutionMode, RunAgentsResult, StartAgentResult, StartAgentVersion, + AIAgentActionResultType, RequestCommandOutputResult, RunAgentsAgentOutcome, + RunAgentsAgentOutcomeKind, RunAgentsLaunchedExecutionMode, RunAgentsResult, StartAgentResult, + StartAgentVersion, }; +#[test] +fn shell_execution_error_is_failed_but_not_cancelled() { + let result = + AIAgentActionResultType::RequestCommandOutput(RequestCommandOutputResult::ExecutionError { + command: "cargo test".to_string(), + message: "terminal is busy".to_string(), + }); + + assert!(result.is_failed()); + assert!(!result.is_cancelled()); + assert!(result.model_content().contains("terminal is busy")); +} + #[test] fn deserializes_legacy_start_agent_success_without_version_as_v1() { let result: StartAgentResult = @@ -134,3 +148,72 @@ fn run_agents_model_content_serializes_terminal_non_launch_outcomes() { assert_eq!(content, expected); } } + +#[test] +fn completed_local_run_agents_preserves_outputs_and_terminal_state() { + let result = AIAgentActionResultType::RunAgents(RunAgentsResult::Launched { + model_id: "model".to_string(), + harness_type: "codex".to_string(), + execution_mode: RunAgentsLaunchedExecutionMode::Local, + agents: vec![RunAgentsAgentOutcome { + name: "research".to_string(), + kind: RunAgentsAgentOutcomeKind::Completed { + agent_id: "child-1".to_string(), + output: "Found the root cause".to_string(), + }, + }], + }); + + let content: serde_json::Value = serde_json::from_str(&result.model_content()).unwrap(); + assert_eq!(content["status"], "launched"); + assert_eq!(content["completion_state"], "children_completed"); + assert_eq!(content["agents"][0]["status"], "completed"); + assert_eq!(content["agents"][0]["output"], "Found the root cause"); + assert!(result.is_successful()); + assert!(!result.is_failed()); +} + +#[test] +fn all_failed_run_agents_is_failure_but_mixed_batch_is_successful() { + let failed = |name: &str, error: &str| RunAgentsAgentOutcome { + name: name.to_string(), + kind: RunAgentsAgentOutcomeKind::Failed { + error: error.to_string(), + }, + }; + let result = AIAgentActionResultType::RunAgents(RunAgentsResult::Launched { + model_id: "model".to_string(), + harness_type: "codex".to_string(), + execution_mode: RunAgentsLaunchedExecutionMode::Local, + agents: vec![failed("one", "first error"), failed("two", "second error")], + }); + + let content: serde_json::Value = serde_json::from_str(&result.model_content()).unwrap(); + assert_eq!(content["status"], "failure"); + assert_eq!(content["agents"][0]["error"], "first error"); + assert_eq!(content["agents"][1]["error"], "second error"); + assert!(result.is_failed()); + assert!(!result.is_successful()); + + let mixed = AIAgentActionResultType::RunAgents(RunAgentsResult::Launched { + model_id: "model".to_string(), + harness_type: "codex".to_string(), + execution_mode: RunAgentsLaunchedExecutionMode::Local, + agents: vec![ + failed("one", "first error"), + RunAgentsAgentOutcome { + name: "two".to_string(), + kind: RunAgentsAgentOutcomeKind::Completed { + agent_id: "child-2".to_string(), + output: "useful output".to_string(), + }, + }, + ], + }); + let mixed_content: serde_json::Value = serde_json::from_str(&mixed.model_content()).unwrap(); + assert_eq!(mixed_content["status"], "launched"); + assert_eq!(mixed_content["agents"][0]["error"], "first error"); + assert_eq!(mixed_content["agents"][1]["output"], "useful output"); + assert!(mixed.is_successful()); + assert!(!mixed.is_failed()); +} diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index d824fb60..9652fba4 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -820,7 +820,8 @@ impl ProviderRun { call_id: call.call.id.clone(), }) } - PendingToolCallState::PermissionPending { .. } | PendingToolCallState::Executing => { + PendingToolCallState::Executing => Ok(()), + PendingToolCallState::PermissionPending { .. } => { Err(invalid_tool_transition(call, "tool start")) } } @@ -840,6 +841,9 @@ impl ProviderRun { call.state = PendingToolCallState::Resolved { result }; Ok(()) } + PendingToolCallState::Resolved { + result: completed_result, + } if completed_result == &result => Ok(()), PendingToolCallState::Resolved { .. } => { Err(ProviderRunProtocolError::DuplicateToolUpdate { call_id: call.call.id.clone(), diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index 0f482cad..6d991f83 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -358,6 +358,89 @@ fn parallel_tool_results_commit_atomically_in_original_call_order() { assert_eq!(ids, vec!["first", "second"]); } +#[test] +fn duplicate_tool_start_is_idempotent() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ); + + run.start_tool(&batch.work_id, "read").unwrap(); + let started = run.clone(); + run.start_tool(&batch.work_id, "read").unwrap(); + + assert_eq!(run, started); +} + +#[test] +fn identical_tool_completion_is_idempotent() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ); + let result = successful_result("read", "contents"); + + run.complete_tool(&batch.work_id, result.clone()).unwrap(); + let completed = run.clone(); + run.complete_tool(&batch.work_id, result).unwrap(); + + assert_eq!(run, completed); +} + +#[test] +fn conflicting_tool_completion_is_rejected_without_mutation() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ); + run.complete_tool(&batch.work_id, successful_result("read", "contents")) + .unwrap(); + let completed = run.clone(); + + assert_eq!( + run.complete_tool(&batch.work_id, successful_result("read", "different")) + .unwrap_err(), + ProviderRunProtocolError::DuplicateToolUpdate { + call_id: "read".to_string(), + } + ); + assert_eq!(run, completed); +} + +#[test] +fn duplicate_tool_callbacks_do_not_prevent_eventual_batch_completion() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("first", "read_files"), + tool_call("second", "grep"), + ], + &["read_files", "grep"], + ), + ); + let first_result = successful_result("first", "one"); + + run.start_tool(&batch.work_id, "first").unwrap(); + run.start_tool(&batch.work_id, "first").unwrap(); + run.complete_tool(&batch.work_id, first_result.clone()) + .unwrap(); + run.complete_tool(&batch.work_id, first_result).unwrap(); + run.complete_tool(&batch.work_id, successful_result("second", "two")) + .unwrap(); + + let ProviderRunState::AwaitingTools { batch: completed } = run.state() else { + panic!("expected completed tool batch"); + }; + assert!(completed.is_complete()); + run.commit_tool_batch(&batch.work_id).unwrap(); + assert_eq!(run.state().phase(), ProviderRunPhase::ReadyToCallModel); +} + #[test] fn exact_batch_submission_rejects_missing_duplicate_and_unknown_results_without_mutation() { let mut run = run();