From 4bb7282cd6e9bcbc73405f0e5e43419b6095a010 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 8 Sep 2026 17:04:03 +0200 Subject: [PATCH 01/15] refactor(core): make conversation checkpoints preserve session-owned state --- crates/agent_core/src/persistence.rs | 7 +- crates/agent_core/src/runtime.rs | 912 ++++++++---------- crates/agent_core/src/runtime/tests.rs | 305 ++++++ crates/agent_core/src/tree.rs | 111 +++ .../src/agent/checkpoint_tests.rs | 175 ++++ .../src/agent/persistence.rs | 2 + crates/code_assistant_core/src/agent/tests.rs | 102 +- crates/code_assistant_core/src/persistence.rs | 120 +++ .../src/session/manager.rs | 391 ++++++-- 9 files changed, 1461 insertions(+), 664 deletions(-) create mode 100644 crates/agent_core/src/runtime/tests.rs create mode 100644 crates/code_assistant_core/src/agent/checkpoint_tests.rs diff --git a/crates/agent_core/src/persistence.rs b/crates/agent_core/src/persistence.rs index 632b2458..e98bb3d0 100644 --- a/crates/agent_core/src/persistence.rs +++ b/crates/agent_core/src/persistence.rs @@ -1,5 +1,6 @@ //! Core-shaped persistence: the loop saves what it owns — the conversation -//! tree, the linearized history, the tool executions, and the id counters. +//! tree, its derived linear history, the tool executions, and the id counters. +//! Prompt-only repairs and context-recovery projections are never checkpointed. //! Application-level fields travel separately through the extension state //! and are assembled into the application's storage format by its adapter. @@ -15,7 +16,9 @@ pub struct AgentSnapshot { pub message_nodes: BTreeMap, pub active_path: ConversationPath, pub next_node_id: NodeId, - /// Linearized message history (derived from `active_path`). + /// Canonical linear history derived from `active_path`, retained for API + /// compatibility. Never the rendered/repaired LLM prompt. A supplied tree + /// takes precedence over this field on restore. pub messages: Vec, pub tool_executions: Vec, pub next_request_id: u64, diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index 454562d8..27d0b5e2 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -1,10 +1,13 @@ //! The agent loop. Application behavior plugs in through the hook traits in //! [`crate::hooks`]; application state travels type-erased in `extensions`. +#[cfg(test)] +mod tests; + use crate::dialect::ToolDialect; use crate::hooks::{ContextSnapshot, HookRegistry, LoopCtx, RecoveryAction, ToolServicesProvider}; use crate::persistence::{AgentSnapshot, SnapshotPersistence}; -use crate::tree::{ConversationPath, MessageNode, NodeId}; +use crate::tree::{Conversation, ConversationPath, MessageNode, NodeId}; use crate::types::{ToolExecution, ToolRequest, text_summary_from_blocks, to_tool_definitions}; use crate::ui::{AgentActivity, AgentUi, AgentUiEvent, DisplayFragment, HiddenTools, UIError}; use anyhow::Result; @@ -60,6 +63,13 @@ enum LoopFlow { GetUserInput, } +/// Purely derived request adjustments; canonical state is never edited here. +#[derive(Default)] +struct PromptProjection { + omitted_nodes: std::collections::HashSet, + tool_results: HashMap, +} + pub struct AgentRuntime { hooks: HookRegistry, /// Application-specific loop state, exposed to the hooks type-erased via @@ -81,21 +91,9 @@ pub struct AgentRuntime { permission_handler: Option>, permissions: ToolPermissions, - // ======================================================================== - // Branching: Tree-based message storage - // ======================================================================== - /// All message nodes in the session (tree structure) - message_nodes: BTreeMap, - /// The currently active path through the tree - active_path: ConversationPath, - /// Counter for generating unique node IDs - next_node_id: NodeId, - - // ======================================================================== - // Legacy: Linearized message history (derived from active_path) - // ======================================================================== - /// Store all messages exchanged (kept in sync with active_path) - message_history: Vec, + conversation: Conversation, + /// Run-local LLM projection. Never included in a checkpoint. + prompt_projection: PromptProjection, // Store the history of tool executions tool_executions: Vec, @@ -162,12 +160,8 @@ impl AgentRuntime { services_provider, permission_handler, permissions, - // Branching tree structure - message_nodes: BTreeMap::new(), - active_path: Vec::new(), - next_node_id: 1, - // Linearized message history - message_history: Vec::new(), + conversation: Conversation::default(), + prompt_projection: PromptProjection::default(), tool_executions: Vec::new(), cached_system_prompts: HashMap::new(), next_request_id: 1, // Start from 1 @@ -216,10 +210,9 @@ impl AgentRuntime { next_node_id: NodeId, messages: Vec, ) { - self.message_nodes = message_nodes; - self.active_path = active_path; - self.next_node_id = next_node_id; - self.message_history = messages; + self.conversation = + Conversation::restore(message_nodes, active_path, next_node_id, messages); + self.prompt_projection = PromptProjection::default(); } /// Restore the tool execution records from persisted state. @@ -259,7 +252,7 @@ impl AgentRuntime { /// Get a reference to the message history pub fn message_history(&self) -> &[Message] { - &self.message_history + self.conversation.history() } /// Get and clear the pending message from shared state @@ -290,16 +283,16 @@ impl AgentRuntime { fn save_state(&mut self) -> Result<()> { trace!( "saving {} messages to persistence (tree nodes: {})", - self.message_history.len(), - self.message_nodes.len() + self.conversation.history().len(), + self.conversation.nodes().len() ); let snapshot = AgentSnapshot { session_id: self.session_id.clone(), - message_nodes: self.message_nodes.clone(), - active_path: self.active_path.clone(), - next_node_id: self.next_node_id, - messages: self.message_history.clone(), + message_nodes: self.conversation.nodes().clone(), + active_path: self.conversation.path().clone(), + next_node_id: self.conversation.next_id(), + messages: self.conversation.history().to_vec(), tool_executions: self.tool_executions.clone(), next_request_id: self.next_request_id, }; @@ -311,35 +304,19 @@ impl AgentRuntime { /// The returned ID is guaranteed to be used by the next `append_message` call /// (or `append_message_with_node_id`). pub fn reserve_node_id(&mut self) -> NodeId { - let id = self.next_node_id; - self.next_node_id += 1; - id + self.conversation.reserve_id() } /// Adds a message to the history using a pre-allocated node_id. /// Use `reserve_node_id()` to obtain the ID before streaming starts, /// then call this after streaming completes. pub fn append_message_with_node_id(&mut self, message: Message, node_id: NodeId) -> Result<()> { - let parent_id = self.active_path.last().copied(); - - let node = MessageNode { - id: node_id, - message: message.clone(), - parent_id, - created_at: std::time::SystemTime::now(), - extension: None, - }; - - self.message_nodes.insert(node_id, node); - self.active_path.push(node_id); + self.conversation.append(message.clone(), node_id); for observer in &self.hooks.observers { observer.on_message(self.session_id.as_deref(), &message); } - // Also add to linearized history - self.message_history.push(message); - self.save_state()?; Ok(()) } @@ -433,22 +410,14 @@ impl AgentRuntime { .extract_tool_requests_from_response(&llm_response, request_id) .await?; - // 4. If we have a truncated response different from the original, update the last message + // Persist the parser's corrected response through the same tree + // mutation boundary used by format-on-save. if !truncated_response.content.is_empty() - && !self.message_history.is_empty() && truncated_response.content != llm_response.content { - // Replace the last message with the truncated version - if let Some(last_msg) = self.message_history.last_mut() - && last_msg.role == MessageRole::Assistant - { - last_msg.content = - MessageContent::Structured(truncated_response.content.clone()); - last_msg.usage = Some(truncated_response.usage.clone()); - } + self.correct_last_assistant_response(&truncated_response)?; } - // 5. Act based on the flow instruction match flow { LoopFlow::GetUserInput => { // In on-demand mode, we don't wait for user input @@ -481,61 +450,22 @@ impl AgentRuntime { } } - /// Drop dangling assistant tool requests (no following tool result) - /// from a freshly restored history. - pub fn normalize_loaded_message_history(&mut self) { - if self.message_history.is_empty() { - return; - } - - let dialect = self.dialect.clone(); - let mut removed = 0usize; - - while let Some(last_assistant_idx) = self - .message_history - .iter() - .rposition(|message| message.role == MessageRole::Assistant) - { - let last_assistant = &self.message_history[last_assistant_idx]; - - if !dialect.message_contains_invocation(last_assistant, self.registry.as_ref()) { - break; - } - - let has_tool_result_after = self.message_history[last_assistant_idx + 1..] - .iter() - .any(Self::is_user_tool_result_message); - - if has_tool_result_after { - break; - } - - let message = self.message_history.remove(last_assistant_idx); - debug!( - "Removing dangling assistant tool request (request_id={:?}) from history", - message.request_id - ); - removed += 1; - } - - if removed > 0 { - debug!( - "Normalized message history by dropping {removed} dangling tool request message(s)" - ); - } - } - - fn is_user_tool_result_message(message: &Message) -> bool { - if message.role != MessageRole::User { - return false; - } - - match &message.content { - MessageContent::Structured(blocks) => blocks - .iter() - .any(|block| matches!(block, ContentBlock::ToolResult { .. })), - MessageContent::Text(text) => text.trim().is_empty(), + /// Compatibility entry point. Restores no longer delete incomplete tool + /// calls: the tree/cache retain evidence, and prompt rendering supplies + /// missing outcomes without guessing that a user cancelled the operation. + pub fn normalize_loaded_message_history(&mut self) {} + + fn correct_last_assistant_response(&mut self, response: &llm::LLMResponse) -> Result<()> { + if let Some(id) = self.conversation.path().last().copied() { + self.conversation.edit_message(id, |message| { + if message.role == MessageRole::Assistant { + message.content = MessageContent::Structured(response.content.clone()); + message.usage = Some(response.usage.clone()); + } + }); + self.save_state()?; } + Ok(()) } /// Parses tool requests from the LLM response and returns a truncated response. @@ -911,70 +841,29 @@ impl AgentRuntime { /// Convert ToolResult blocks to Text blocks for custom tool-syntax mode fn convert_tool_results_to_text(&self, messages: Vec) -> Vec { - // Create a fresh ResourcesTracker for rendering - let mut resources_tracker = ResourcesTracker::new(); - - // First, build a map of tool_use_id to rendered output - let mut tool_outputs = std::collections::HashMap::new(); - - // Process tool executions in reverse chronological order (newest first) - for execution in self.tool_executions.iter().rev() { - let tool_use_id = &execution.tool_request.id; - let rendered_output = execution.result.as_render().render(&mut resources_tracker); - tool_outputs.insert(tool_use_id.clone(), rendered_output); - } - - // Process each message + // Inputs are already rendered, including recovery overrides. Rendering + // executions a second time here would undo the projection for XML/caret. messages .into_iter() - .map(|msg| { - match &msg.content { - MessageContent::Structured(blocks) => { - // Check if there are any ToolResult blocks that need conversion - let has_tool_results = blocks - .iter() - .any(|block| matches!(block, ContentBlock::ToolResult { .. })); - - if !has_tool_results { - // No conversion needed - return msg; - } - - // Convert all blocks to Text - let mut text_content = String::new(); - - for block in blocks { - match block { - ContentBlock::ToolResult { tool_use_id, .. } => { - // Get the dynamically rendered content for this tool result - if let Some(rendered_output) = tool_outputs.get(tool_use_id) { - // Add the rendered tool output from actual tool execution - text_content.push_str(rendered_output); - text_content.push_str("\n\n"); - } - } - ContentBlock::Text { text, .. } => { - // For existing Text blocks, keep as is - text_content.push_str(text); - text_content.push_str("\n\n"); - } - _ => {} // Ignore other block types + .map(|mut message| { + if let MessageContent::Structured(blocks) = &message.content + && blocks + .iter() + .any(|b| matches!(b, ContentBlock::ToolResult { .. })) + { + let text: Vec<_> = blocks + .iter() + .filter_map(|block| match block { + ContentBlock::ToolResult { content, .. } => { + Some(content.text_content().to_string()) } - } - - // Create a new message with Text content - Message { - role: msg.role, - content: MessageContent::Text(text_content.trim().to_string()), - volatile: msg.volatile, - request_id: msg.request_id, - usage: msg.usage.clone(), - ..Default::default() - } - } - // For non-structured content, keep as is - _ => msg, + ContentBlock::Text { text, .. } => Some(text.clone()), + _ => None, + }) + .collect(); + message.content = MessageContent::Text(text.join("\n\n").trim().to_string()); } + message }) .collect() } @@ -982,19 +871,22 @@ impl AgentRuntime { /// Runs the iteration hooks over the rendered messages right before they /// are sent to the LLM (e.g. to inject system reminders). pub fn shape_request_messages(&mut self, mut messages: Vec) -> Vec { - let ctx = LoopCtx { - tool_executions: &mut self.tool_executions, - message_nodes: &mut self.message_nodes, - active_path: &self.active_path, - session_id: self.session_id.as_deref(), - registry: self.registry.as_ref(), - extensions: self.extensions.as_mut(), - }; - for hook in &self.hooks.iteration_hooks { - if let Err(e) = hook.shape_request(&mut messages, &ctx) { - warn!("Iteration hook failed to shape the request: {}", e); - } - } + self.conversation + .with_nodes_mut(|message_nodes, active_path| { + let ctx = LoopCtx { + tool_executions: &mut self.tool_executions, + message_nodes, + active_path, + session_id: self.session_id.as_deref(), + registry: self.registry.as_ref(), + extensions: self.extensions.as_mut(), + }; + for hook in &self.hooks.iteration_hooks { + if let Err(e) = hook.shape_request(&mut messages, &ctx) { + warn!("Iteration hook failed to shape the request: {}", e); + } + } + }); messages } @@ -1247,15 +1139,31 @@ impl AgentRuntime { } fn active_messages(&self) -> &[Message] { - if self.message_history.is_empty() { - return &[]; - } - let start = self - .message_history + let history = self.conversation.history(); + let start = history .iter() .rposition(|message| message.is_compaction_summary) .unwrap_or(0); - &self.message_history[start..] + &history[start..] + } + + fn prompt_messages(&self) -> Vec { + let path = self.conversation.path(); + let start = path + .iter() + .rposition(|id| { + self.conversation + .nodes() + .get(id) + .is_some_and(|node| node.message.is_compaction_summary) + }) + .unwrap_or(0); + path[start..] + .iter() + .filter(|id| !self.prompt_projection.omitted_nodes.contains(id)) + .filter_map(|id| self.conversation.nodes().get(id)) + .map(|node| node.message.clone()) + .collect() } fn context_usage_ratio(&mut self) -> Result> { @@ -1298,11 +1206,9 @@ impl AgentRuntime { Ok(self.hooks.compaction.should_compact(&snapshot)) } - /// Shrinks the conversation after the provider rejected the prompt as too long. - /// Replaces large tool results with error placeholders when possible — the next - /// render of the message history then produces a much smaller prompt. If nothing - /// is large enough to replace, drops the last assistant+tool-result exchange and - /// forces context compaction as a last resort. + /// Shrinks only the request projection after an oversized-prompt rejection. + /// Large results become prompt placeholders; otherwise the last exchange is + /// omitted from the compaction request. Canonical evidence is never removed. async fn recover_from_oversized_prompt(&mut self) -> Result<()> { warn!("Prompt too long error detected, replacing large tool results with error messages"); let replaced = self.replace_large_tool_results(); @@ -1313,7 +1219,8 @@ impl AgentRuntime { self.drop_last_tool_exchange(); return self.perform_compaction().await; } - // Notify the UI that these tools switched from success → error + // Keep the existing transient UI notification for a rejected output; + // it is not a change to the persisted execution's actual outcome. for (tool_id, error_message) in &replaced { let _ = self .send_ui(AgentUiEvent::UpdateToolStatus { @@ -1362,9 +1269,8 @@ impl AgentRuntime { .await; } - /// Replace the largest tool execution results **from the most recent turn** - /// with [`PromptTooLongError`] placeholders so that the next LLM request has a - /// chance to succeed. + /// Project the largest results from the most recent turn as small error + /// placeholders for the next request. Original execution records survive. /// /// Returns a vec of `(tool_id, error_message)` for each replaced result, /// empty if nothing was replaced. The caller is responsible for sending @@ -1375,7 +1281,7 @@ impl AgentRuntime { // Collect tool_use_ids from the last user message that contains ToolResult // blocks — these are the results from the most recent turn. let current_turn_ids: std::collections::HashSet = self - .message_history + .prompt_messages() .iter() .rev() .find_map(|msg| { @@ -1410,7 +1316,12 @@ impl AgentRuntime { let mut sizes: Vec<(usize, usize)> = Vec::new(); // (index, byte_size) let mut tracker = ResourcesTracker::new(); for (i, exec) in self.tool_executions.iter().enumerate() { - if !current_turn_ids.contains(&exec.tool_request.id) { + if !current_turn_ids.contains(&exec.tool_request.id) + || self + .prompt_projection + .tool_results + .contains_key(&exec.tool_request.id) + { continue; } let rendered = exec.result.as_render().render(&mut tracker); @@ -1438,92 +1349,40 @@ impl AgentRuntime { ); let error = PromptTooLongError::new(&tool_name, byte_size); let error_message = error.error_message.clone(); - self.tool_executions[idx].result = Box::new(error); + self.prompt_projection + .tool_results + .insert(tool_id.clone(), error_message.clone()); replaced.push((tool_id, error_message)); } - // Also update the corresponding ToolResult content blocks in message history - // so the is_error flag is set correctly - if !replaced.is_empty() { - let replaced_ids: std::collections::HashSet<&str> = - replaced.iter().map(|(id, _)| id.as_str()).collect(); - - for msg in &mut self.message_history { - if let MessageContent::Structured(blocks) = &mut msg.content { - for block in blocks { - if let ContentBlock::ToolResult { - tool_use_id, - is_error, - .. - } = block - && replaced_ids.contains(tool_use_id.as_str()) - { - *is_error = Some(true); - } - } - } - } - } - replaced } - /// Drop the last assistant → tool-result message pair from history. - /// Also removes the corresponding `tool_executions` entries. - /// Used as a last-resort fallback before forcing compaction when the prompt - /// is too long but no individual tool result is large enough to replace. + /// Omit the last tool exchange from the prompt only, as the fallback + /// before compaction. Canonical messages and execution evidence survive. fn drop_last_tool_exchange(&mut self) { - // Walk backwards to find the last user message with ToolResult blocks - // and the assistant message immediately before it. - let mut tool_result_idx = None; - for i in (0..self.message_history.len()).rev() { - let msg = &self.message_history[i]; - if msg.role == MessageRole::User - && let MessageContent::Structured(blocks) = &msg.content - && blocks - .iter() - .any(|b| matches!(b, ContentBlock::ToolResult { .. })) - { - tool_result_idx = Some(i); - break; - } - } - - let Some(tr_idx) = tool_result_idx else { - return; - }; - - // Collect the tool_use_ids we're about to drop so we can clean up - // tool_executions too. - let mut dropped_ids: std::collections::HashSet = std::collections::HashSet::new(); - if let MessageContent::Structured(blocks) = &self.message_history[tr_idx].content { - for block in blocks { - if let ContentBlock::ToolResult { tool_use_id, .. } = block { - dropped_ids.insert(tool_use_id.clone()); - } - } - } - - // Remove the tool-result user message - self.message_history.remove(tr_idx); - - // If the message right before it was the assistant message with the - // corresponding ToolUse blocks, remove that too. - if tr_idx > 0 { - let prev = &self.message_history[tr_idx - 1]; - if prev.role == MessageRole::Assistant { - self.message_history.remove(tr_idx - 1); - } + let visible: Vec<_> = self + .conversation + .path() + .iter() + .copied() + .filter(|id| !self.prompt_projection.omitted_nodes.contains(id)) + .collect(); + let Some(index) = visible.iter().rposition(|id| { + self.conversation.nodes().get(id).is_some_and(|node| { + node.message.role == MessageRole::User + && matches!(&node.message.content, MessageContent::Structured(blocks) + if blocks.iter().any(|block| matches!(block, ContentBlock::ToolResult { .. }))) + }) + }) else { return; }; + self.prompt_projection.omitted_nodes.insert(visible[index]); + if index > 0 + && self.conversation.nodes()[&visible[index - 1]].message.role == MessageRole::Assistant + { + self.prompt_projection + .omitted_nodes + .insert(visible[index - 1]); } - - // Remove corresponding tool executions - self.tool_executions - .retain(|e| !dropped_ids.contains(&e.tool_request.id)); - - debug!( - "Dropped last tool exchange ({} tool result(s)) from history", - dropped_ids.len() - ); } async fn perform_compaction(&mut self) -> Result<()> { @@ -1582,209 +1441,142 @@ impl AgentRuntime { Ok(()) } - /// Prepare messages for LLM request, dynamically rendering tool outputs. - /// - /// This function also handles cancelled tool executions: if an assistant message - /// contains `ToolUse` blocks but there's no corresponding `ToolResult` in the - /// following user message (or no following user message at all), we generate - /// a synthetic "user cancelled" `ToolResult` to satisfy the API requirement that - /// every `tool_use` must have a corresponding `tool_result`. - pub fn render_tool_results_in_messages(&self) -> Vec { - // Start with a clean slate - let mut messages = Vec::new(); - - // Create a fresh ResourcesTracker for this rendering pass - let mut resources_tracker = ResourcesTracker::new(); - - // First, collect all tool executions and build a map from tool_use_id to rendered output - let mut tool_outputs = std::collections::HashMap::new(); - // Collect image data from tools that produce visual output - let mut tool_images: std::collections::HashMap> = - std::collections::HashMap::new(); + fn prompt_tool_use_ids(&self, message: &Message) -> Vec { + if message.role != MessageRole::Assistant { + return Vec::new(); + } + let content = match &message.content { + MessageContent::Structured(blocks) => blocks.clone(), + MessageContent::Text(text) => vec![ContentBlock::new_text(text.clone())], + }; + let response = llm::LLMResponse { + content, + usage: llm::Usage::zero(), + rate_limit_info: None, + }; + self.dialect + .extract_requests( + &response, + message.request_id.unwrap_or(0), + 0, + self.registry.as_ref(), + ) + .map(|(requests, _)| requests.into_iter().map(|request| request.id).collect()) + .unwrap_or_default() + } - // Process tool executions in reverse chronological order (newest first) - // so newer tool calls take precedence in resource conflicts + /// Render the run-local LLM projection, never modifying conversation or + /// tool evidence. Missing results are repaired by id in the immediate + /// follow-up message; absent evidence means an unknown outcome, not cancel. + pub fn render_tool_results_in_messages(&self) -> Vec { + let mut messages = self.prompt_messages(); + let mut tracker = ResourcesTracker::new(); + let mut outputs = HashMap::new(); + // Only render executions visible in this prompt. Inactive branches and + // omitted exchanges must not claim resources in the render tracker. + let visible_ids: std::collections::HashSet<_> = messages + .iter() + .flat_map(|message| { + let mut ids = self.prompt_tool_use_ids(message); + if let MessageContent::Structured(blocks) = &message.content { + ids.extend(blocks.iter().filter_map(|block| match block { + ContentBlock::ToolResult { tool_use_id, .. } => Some(tool_use_id.clone()), + _ => None, + })); + } + ids + }) + .collect(); for execution in self.tool_executions.iter().rev() { - let tool_use_id = &execution.tool_request.id; - let rendered_output = execution.result.as_render().render(&mut resources_tracker); - tool_outputs.insert(tool_use_id.clone(), rendered_output); - - // Collect any image data from the tool output - let images = execution.result.render_images(); - if !images.is_empty() { - tool_images.insert(tool_use_id.clone(), images); + let id = &execution.tool_request.id; + if !visible_ids.contains(id) || outputs.contains_key(id) { + continue; } + let (content, is_error) = + if let Some(replacement) = self.prompt_projection.tool_results.get(id) { + (ToolResultContent::text(replacement.clone()), true) + } else { + let text = execution.result.as_render().render(&mut tracker); + let images = execution + .result + .render_images() + .into_iter() + .map(|image| ToolResultImage { + media_type: image.media_type, + base64_data: image.base64_data, + }) + .collect(); + ( + ToolResultContent::with_images(text, images), + !execution.result.is_success(), + ) + }; + outputs.insert(id.clone(), (content, is_error)); } - // Build a set of all tool_use_ids that have corresponding tool_results in the message history - let mut tool_ids_with_results: std::collections::HashSet = - std::collections::HashSet::new(); - - for msg in self.active_messages() { - if let MessageContent::Structured(blocks) = &msg.content { - for block in blocks { - if let ContentBlock::ToolResult { tool_use_id, .. } = block { - tool_ids_with_results.insert(tool_use_id.clone()); - } + // Repair only the projection, including partially recorded batches. + let mut index = 0; + while index < messages.len() { + let missing: Vec<_> = if messages[index].role == MessageRole::Assistant { + let ids = self.prompt_tool_use_ids(&messages[index]); + ids.into_iter().filter(|id| { + !messages.get(index + 1).is_some_and(|next| { + next.role == MessageRole::User && matches!(&next.content, MessageContent::Structured(blocks) + if blocks.iter().any(|block| matches!(block, ContentBlock::ToolResult { tool_use_id, .. } if tool_use_id == id))) + }) + }).map(|id| ContentBlock::ToolResult { + content: outputs.get(&id).map(|(content, _)| content.clone()).unwrap_or_else(|| { + ToolResultContent::text("Tool result is missing; execution outcome is unknown. Verify the state before retrying any side effects.") + }), + is_error: Some(outputs.get(&id).map(|(_, error)| *error).unwrap_or(true)), + tool_use_id: id, + start_time: None, + end_time: None, + }).collect() + } else { + Vec::new() + }; + if !missing.is_empty() { + if let Some(next) = messages.get_mut(index + 1) + && next.role == MessageRole::User + && let MessageContent::Structured(blocks) = &mut next.content + && blocks + .iter() + .any(|block| matches!(block, ContentBlock::ToolResult { .. })) + { + blocks.extend(missing); + } else { + messages.insert(index + 1, Message::new_user_content(missing)); } } + index += 1; } - // Now rebuild the message history, replacing tool outputs with our dynamically rendered versions - let active_msgs: Vec<_> = self.active_messages().to_vec(); - for (idx, msg) in active_msgs.iter().enumerate() { - match &msg.content { + for message in &mut messages { + match &mut message.content { MessageContent::Structured(blocks) => { - if msg.role == MessageRole::Assistant { - // Check for ToolUse blocks that need synthetic ToolResults - let tool_use_ids: Vec = blocks - .iter() - .filter_map(|block| { - if let ContentBlock::ToolUse { id, .. } = block { - Some(id.clone()) - } else { - None - } - }) - .collect(); - - // Find tool_use_ids without corresponding tool_results - let missing_results: Vec<&String> = tool_use_ids - .iter() - .filter(|id| !tool_ids_with_results.contains(*id)) - .collect(); - - if !missing_results.is_empty() { - // We need to add the assistant message first, then add a synthetic - // user message with cancelled tool results - messages.push(msg.clone()); - - // Generate synthetic ToolResult blocks for cancelled tools - let cancelled_blocks: Vec = missing_results - .iter() - .map(|tool_id| { - debug!( - "Generating synthetic 'cancelled' tool result for tool_use_id: {}", - tool_id - ); - - ContentBlock::ToolResult { - tool_use_id: (*tool_id).clone(), - content: ToolResultContent::text( - "Tool execution was cancelled by user.", - ), - is_error: Some(true), - start_time: None, - end_time: None, - } - }) - .collect(); - - // Check if the next message is already a user message with tool results - // In that case, we need to merge the cancelled results - let next_msg = active_msgs.get(idx + 1); - let should_create_new_message = match next_msg { - Some(next) if next.role == MessageRole::User => { - // Check if this user message has tool results - match &next.content { - MessageContent::Structured(next_blocks) => !next_blocks - .iter() - .any(|b| matches!(b, ContentBlock::ToolResult { .. })), - _ => true, - } - } - _ => true, - }; - - if should_create_new_message { - // Insert a new user message with the cancelled tool results - let cancelled_msg = - Message::new_user_content(cancelled_blocks.clone()); - messages.push(cancelled_msg); - } - // If next message already has tool results, we'll handle merging when we process it - continue; - } - } - - // Look for ToolResult blocks and update with rendered output. - // When a tool produces images, they are embedded inside the - // ToolResultContent so Anthropic receives them in the - // `tool_result.content` array (per the API spec). - let mut new_blocks = Vec::new(); - let mut need_update = false; - for block in blocks { - match block { - ContentBlock::ToolResult { - tool_use_id, - is_error, - start_time, - end_time, - .. - } => { - // If we have an execution result for this tool use, use it - if let Some(output) = tool_outputs.get(tool_use_id) { - // Build content with optional images - let content = if let Some(images) = tool_images.get(tool_use_id) - { - ToolResultContent::with_images( - output.clone(), - images - .iter() - .map(|img| ToolResultImage { - media_type: img.media_type.clone(), - base64_data: img.base64_data.clone(), - }) - .collect(), - ) - } else { - ToolResultContent::text(output.clone()) - }; - - new_blocks.push(ContentBlock::ToolResult { - tool_use_id: tool_use_id.clone(), - content, - is_error: *is_error, - start_time: *start_time, - end_time: *end_time, - }); - - need_update = true; - } else { - // Keep the original block - new_blocks.push(block.clone()); - } - } - _ => { - // Keep other blocks as is - new_blocks.push(block.clone()); + if let ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + .. + } = block + && let Some((output, error)) = outputs.get(tool_use_id) + { + *content = output.clone(); + if *error { + *is_error = Some(true); } } } - - if need_update { - let mut updated = msg.clone(); - updated.content = MessageContent::Structured(new_blocks); - messages.push(updated); - } else { - // No changes needed, use original message - messages.push(msg.clone()); - } } - MessageContent::Text(text) => { - if msg.is_compaction_summary { - let mut updated = msg.clone(); - updated.content = - MessageContent::Text(Self::format_compaction_summary_for_prompt(text)); - messages.push(updated); - } else { - messages.push(msg.clone()); - } + MessageContent::Text(text) if message.is_compaction_summary => { + *text = Self::format_compaction_summary_for_prompt(text); } + _ => {} } } - messages } @@ -1792,35 +1584,41 @@ impl AgentRuntime { /// Gives the registered interceptors a chance to handle the request /// before the standard dispatch. Returns `Some(result)` when one did. fn intercept_tool(&mut self, tool_request: &ToolRequest) -> Option> { - let mut ctx = LoopCtx { - tool_executions: &mut self.tool_executions, - message_nodes: &mut self.message_nodes, - active_path: &self.active_path, - session_id: self.session_id.as_deref(), - registry: self.registry.as_ref(), - extensions: self.extensions.as_mut(), - }; - for interceptor in &self.hooks.interceptors { - if let Some(result) = interceptor.try_intercept(tool_request, &mut ctx) { - return Some(result); - } - } - None + self.conversation + .with_nodes_mut(|message_nodes, active_path| { + let mut ctx = LoopCtx { + tool_executions: &mut self.tool_executions, + message_nodes, + active_path, + session_id: self.session_id.as_deref(), + registry: self.registry.as_ref(), + extensions: self.extensions.as_mut(), + }; + for interceptor in &self.hooks.interceptors { + if let Some(result) = interceptor.try_intercept(tool_request, &mut ctx) { + return Some(result); + } + } + None + }) } /// Notifies the registered interceptors that a tool executed successfully. fn after_tool_success(&mut self, tool_request: &ToolRequest) { - let mut ctx = LoopCtx { - tool_executions: &mut self.tool_executions, - message_nodes: &mut self.message_nodes, - active_path: &self.active_path, - session_id: self.session_id.as_deref(), - registry: self.registry.as_ref(), - extensions: self.extensions.as_mut(), - }; - for interceptor in &self.hooks.interceptors { - interceptor.after_tool_success(tool_request, &mut ctx); - } + self.conversation + .with_nodes_mut(|message_nodes, active_path| { + let mut ctx = LoopCtx { + tool_executions: &mut self.tool_executions, + message_nodes, + active_path, + session_id: self.session_id.as_deref(), + registry: self.registry.as_ref(), + extensions: self.extensions.as_mut(), + }; + for interceptor in &self.hooks.interceptors { + interceptor.after_tool_success(tool_request, &mut ctx); + } + }); } /// A tool may rewrite its own input while executing (e.g. format-on-save). @@ -2125,59 +1923,115 @@ impl AgentRuntime { Ok(()) } - /// Update message history to reflect formatted tool parameters + /// Persist formatted inputs through the conversation mutation boundary. fn update_message_history_with_formatted_tool( &mut self, updated_request: &ToolRequest, ) -> Result<()> { let dialect = self.dialect.clone(); let registry = self.registry.clone(); - // Find the most recent assistant message that contains the tool call - for message in self.message_history.iter_mut().rev() { - if message.role == MessageRole::Assistant { - match &mut message.content { - MessageContent::Structured(blocks) => { - // Look for the ToolUse block with matching ID - for block in blocks { - if let ContentBlock::ToolUse { - id, name, input, .. - } = block - && *id == updated_request.id - && *name == updated_request.name - { - *input = updated_request.input.clone(); - debug!("Updated tool call {} in message history", id); - return Ok(()); - } + let Some(id) = self + .conversation + .path() + .iter() + .rev() + .find(|id| { + self.conversation + .nodes() + .get(id) + .is_some_and(|node| node.message.role == MessageRole::Assistant) + }) + .copied() + else { + return Ok(()); + }; + let mut updated = false; + self.conversation.edit_message(id, |message| { + let request_id = message.request_id.unwrap_or(0); + match &mut message.content { + MessageContent::Structured(blocks) => { + for block in blocks.iter_mut() { + if let ContentBlock::ToolUse { + id, name, input, .. + } = block + && id == &updated_request.id + && name == &updated_request.name + { + *input = updated_request.input.clone(); + updated = true; + return; } } - MessageContent::Text(text) => { - // For text content, we need to update the tool call in the text - // This is more complex and depends on the tool syntax - if let Ok(updated_text) = Self::update_tool_call_in_text_static( - text, + if !dialect.uses_native_tools() { + updated = Self::update_tool_call_in_text_blocks( + blocks, updated_request, + request_id, dialect.as_ref(), registry.as_ref(), - ) { - *text = updated_text; - debug!("Updated tool call {} in text message", updated_request.id); - return Ok(()); - } + ); + } + } + MessageContent::Text(text) => { + if let Ok(replacement) = Self::update_tool_call_in_text_static( + text, + updated_request, + dialect.as_ref(), + registry.as_ref(), + ) { + *text = replacement; + updated = true; } } - // Only check the most recent assistant message - break; } + }); + if updated { + self.save_state()?; + } else { + warn!("Could not find tool call {} to update", updated_request.id); } - - warn!( - "Could not find tool call {} to update in message history", - updated_request.id - ); Ok(()) } + fn update_tool_call_in_text_blocks( + blocks: &mut [ContentBlock], + request: &ToolRequest, + request_id: u64, + dialect: &dyn ToolDialect, + registry: &ToolRegistry, + ) -> bool { + // XML/caret offsets are local to the Text block that was parsed. + // Reparse that block to find the id and current offsets: an earlier + // formatted call may have changed its length. Never rewrite preambles + // or thinking blocks merely because their offsets happen to fit. + for block in blocks { + if let ContentBlock::Text { text, .. } = block { + let response = llm::LLMResponse { + content: vec![ContentBlock::new_text(text.clone())], + usage: llm::Usage::zero(), + rate_limit_info: None, + }; + if let Ok((requests, _)) = + dialect.extract_requests(&response, request_id, 0, registry) + && let Some(current) = requests + .iter() + .find(|current| current.id == request.id && current.name == request.name) + { + let mut corrected = request.clone(); + corrected.start_offset = current.start_offset; + corrected.end_offset = current.end_offset; + if let Ok(replacement) = + Self::update_tool_call_in_text_static(text, &corrected, dialect, registry) + { + *text = replacement; + return true; + } + } + } + } + false + } + /// Static helper to update tool call in text (to avoid borrowing issues) pub fn update_tool_call_in_text_static( text: &str, diff --git a/crates/agent_core/src/runtime/tests.rs b/crates/agent_core/src/runtime/tests.rs new file mode 100644 index 00000000..8bd5e5b8 --- /dev/null +++ b/crates/agent_core/src/runtime/tests.rs @@ -0,0 +1,305 @@ +use super::*; +use crate::hooks::*; +use serde_json::json; + +struct Stub; +#[async_trait::async_trait] +impl LLMProvider for Stub { + async fn send_message( + &mut self, + _: LLMRequest, + _: Option<&StreamingCallback>, + ) -> Result { + anyhow::bail!("unexpected LLM call") + } +} +#[async_trait::async_trait] +impl AgentUi for Stub { + async fn send_event(&self, _: AgentUiEvent) -> Result<(), UIError> { + Ok(()) + } + fn display_fragment(&self, _: &DisplayFragment) -> Result<(), UIError> { + Ok(()) + } + fn should_streaming_continue(&self) -> bool { + true + } + fn notify_rate_limit(&self, _: u64) {} + fn clear_rate_limit(&self) {} +} +impl ToolServicesProvider for Stub { + fn begin(&self, _: &mut (dyn Any + Send), _: &str) -> Box { + Box::new(()) + } + fn end(&self, _: &mut (dyn Any + Send), _: Box) {} + fn detached(&self, _: &str) -> Box { + Box::new(()) + } +} +impl ToolDispatchPolicy for Stub { + fn parallel_indices(&self, _: &[ToolRequest]) -> Vec { + vec![] + } +} +impl CompactionPolicy for Stub { + fn context_limit(&self, _: &(dyn Any + Send)) -> Result> { + Ok(None) + } + fn should_compact(&self, _: &ContextSnapshot) -> bool { + false + } + fn compaction_prompt(&self) -> &str { + "summarize" + } +} +impl RecoveryPolicy for Stub { + fn classify(&self, _: &anyhow::Error, _: u32) -> RecoveryAction { + RecoveryAction::Fail + } +} +impl SystemPromptProvider for Stub { + fn build(&self, _: &PromptCtx) -> String { + String::new() + } +} +#[derive(Clone, Default)] +struct Capture(Arc>>); +impl SnapshotPersistence for Capture { + fn save(&mut self, snapshot: AgentSnapshot, _: &(dyn Any + Send)) -> Result<()> { + *self.0.lock().unwrap() = Some(snapshot); + Ok(()) + } +} +fn runtime() -> (AgentRuntime, Capture) { + let capture = Capture::default(); + let runtime = AgentRuntime::new(AgentRuntimeComponents { + llm_provider: Box::new(Stub), + dialect: Arc::new(crate::native::NativeDialect), + ui: Arc::new(Stub), + registry: Arc::new(ToolRegistry::new()), + tool_capability: String::new(), + excluded_tool_capabilities: vec![], + stream_hidden_tools: Arc::new(|_| false), + command_executor: Arc::new(command_executor::DefaultCommandExecutor), + permission_handler: None, + permissions: Default::default(), + services_provider: Arc::new(Stub), + state_persistence: Box::new(capture.clone()), + hooks: HookRegistry { + interceptors: vec![], + iteration_hooks: vec![], + observers: vec![], + dispatch: Box::new(Stub), + compaction: Box::new(Stub), + recovery: Box::new(Stub), + system_prompt: Box::new(Stub), + }, + extensions: Box::new(()), + }); + (runtime, capture) +} +fn call(id: &str) -> ContentBlock { + ContentBlock::new_tool_use(id, "write_file", json!({"content": "unformatted"})) +} +fn result(id: &str) -> ContentBlock { + ContentBlock::ToolResult { + tool_use_id: id.into(), + content: ToolResultContent::text("evidence"), + is_error: None, + start_time: None, + end_time: None, + } +} +#[test] +fn checkpoint_legacy_history_is_imported_only_without_a_tree() { + let (mut agent, saved) = runtime(); + agent.restore_conversation( + BTreeMap::new(), + Vec::new(), + 1, + vec![Message::new_user("legacy")], + ); + agent.append_message(Message::new_assistant("new")).unwrap(); + let mut snapshot = saved.0.lock().unwrap().take().unwrap(); + assert_eq!(snapshot.message_nodes.len(), 2); + assert_eq!(snapshot.message_nodes[&2].parent_id, Some(1)); + // A nonempty tree with an intentionally empty active path is authoritative + // too: neither reactivate a branch nor import stale linear messages. + snapshot.active_path.clear(); + let restored = reload(snapshot); + assert!(restored.message_history().is_empty()); + assert_eq!(restored.conversation.nodes().len(), 2); +} + +#[test] +fn checkpoint_hook_message_corrections_rebuild_cache_even_on_early_return() { + struct Correction; + impl ToolInterceptor for Correction { + fn try_intercept(&self, _: &ToolRequest, ctx: &mut LoopCtx) -> Option> { + ctx.message_nodes.get_mut(&1).unwrap().message = Message::new_user("corrected by hook"); + Some(Ok(true)) + } + } + let (mut agent, saved) = runtime(); + agent.append_message(Message::new_user("before")).unwrap(); + agent.hooks.interceptors.push(Box::new(Correction)); + assert!( + agent + .intercept_tool(&ToolRequest::from(&call("a"))) + .unwrap() + .unwrap() + ); + agent.save_state().unwrap(); + let snapshot = saved.0.lock().unwrap().take().unwrap(); + assert_eq!( + serde_json::to_value(&snapshot.message_nodes[&1].message).unwrap(), + serde_json::to_value(&snapshot.messages[0]).unwrap() + ); + assert!( + matches!(&snapshot.messages[0].content, MessageContent::Text(text) if text == "corrected by hook") + ); +} + +fn reload(snapshot: AgentSnapshot) -> AgentRuntime { + let (mut restored, _) = runtime(); + // Exercise the serialized tree, not an in-memory alias of its messages. + let nodes = + serde_json::from_value(serde_json::to_value(snapshot.message_nodes).unwrap()).unwrap(); + restored.restore_conversation( + nodes, + snapshot.active_path, + snapshot.next_node_id, + snapshot.messages, + ); + restored.set_tool_executions(snapshot.tool_executions); + restored.normalize_loaded_message_history(); + restored +} + +#[test] +fn checkpoint_tree_wins_over_stale_linear_history() { + let (mut agent, saved) = runtime(); + agent + .append_message(Message::new_user("canonical")) + .unwrap(); + let mut snapshot = saved.0.lock().unwrap().take().unwrap(); + snapshot.messages = vec![Message::new_user("stale")]; + let restored = reload(snapshot); + assert!( + matches!(&restored.message_history()[0].content, MessageContent::Text(text) if text == "canonical") + ); +} + +#[test] +fn checkpoint_formatted_input_survives_roundtrip() { + let (mut agent, saved) = runtime(); + agent + .append_message(Message::new_assistant_content(vec![call("a")])) + .unwrap(); + let mut request = ToolRequest::from(&call("a")); + request.input = json!({"content": "formatted"}); + agent + .update_message_history_with_formatted_tool(&request) + .unwrap(); + agent + .append_message(Message::new_user_content(vec![result("a")])) + .unwrap(); + let snapshot = saved.0.lock().unwrap().take().unwrap(); + assert_eq!( + serde_json::to_value(&snapshot.message_nodes[&1].message.content).unwrap(), + serde_json::to_value(&snapshot.messages[0].content).unwrap() + ); + let restored = reload(snapshot); + assert!(matches!(&restored.message_history()[0].content, + MessageContent::Structured(blocks) if matches!(&blocks[0], ContentBlock::ToolUse { input, .. } if input == &request.input))); +} + +#[test] +fn checkpoint_dangling_calls_survive_reload_with_unknown_prompt_outcome() { + let (mut agent, saved) = runtime(); + agent + .append_message(Message::new_assistant_content(vec![call("a"), call("b")])) + .unwrap(); + // Partial result: the missing result must be merged into this prompt message. + agent + .append_message(Message::new_user_content(vec![result("a")])) + .unwrap(); + let mut restored = reload(saved.0.lock().unwrap().take().unwrap()); + let before = serde_json::to_value(restored.message_history()).unwrap(); + let prompt = restored.render_tool_results_in_messages(); + let MessageContent::Structured(blocks) = &prompt[1].content else { + panic!("results") + }; + assert_eq!( + blocks.len(), + 2, + "partial tool results must be repaired by id" + ); + assert!(blocks.iter().any(|block| matches!(block, ContentBlock::ToolResult { tool_use_id, content, .. } + if tool_use_id == "b" && content.contains("unknown") && !content.contains("cancelled by user")))); + restored.normalize_loaded_message_history(); + assert_eq!( + before, + serde_json::to_value(restored.message_history()).unwrap() + ); +} + +#[test] +fn checkpoint_dangling_tail_is_not_deleted() { + let (mut agent, saved) = runtime(); + agent.append_message(Message::new_user("task")).unwrap(); + agent + .append_message(Message::new_assistant_content(vec![call("a")])) + .unwrap(); + let restored = reload(saved.0.lock().unwrap().take().unwrap()); + assert_eq!(restored.message_history().len(), 2); + assert_eq!(restored.render_tool_results_in_messages().len(), 3); +} + +#[test] +fn checkpoint_recovery_keeps_canonical_messages_and_tool_evidence() { + let (mut agent, saved) = runtime(); + agent + .append_message(Message::new_assistant_content(vec![call("a")])) + .unwrap(); + agent + .append_message(Message::new_user_content(vec![result("a")])) + .unwrap(); + // A serializable large output suffices to exercise size-based recovery. + agent.set_tool_executions(vec![ToolExecution::create_parse_error( + "a".into(), + "x".repeat(60 * 1024), + )]); + let before = serde_json::to_value(agent.message_history()).unwrap(); + let evidence = agent.tool_executions[0].serialize().unwrap(); + assert_eq!(agent.replace_large_tool_results().len(), 1); + let projected = agent.render_tool_results_in_messages(); + assert!(serde_json::to_string(&projected).unwrap().len() < 10 * 1024); + // XML/caret conversion must not re-render the original large execution. + let text_projection = agent.convert_tool_results_to_text(projected); + assert!(serde_json::to_string(&text_projection).unwrap().len() < 10 * 1024); + assert!(agent.replace_large_tool_results().is_empty()); + assert_eq!( + before, + serde_json::to_value(agent.message_history()).unwrap() + ); + assert_eq!( + serde_json::to_value(&evidence).unwrap(), + serde_json::to_value(agent.tool_executions[0].serialize().unwrap()).unwrap() + ); + agent.drop_last_tool_exchange(); + assert!(agent.render_tool_results_in_messages().is_empty()); + agent.save_state().unwrap(); + let restored = reload(saved.0.lock().unwrap().take().unwrap()); + assert_eq!( + before, + serde_json::to_value(restored.message_history()).unwrap() + ); + assert_eq!(restored.tool_executions.len(), 1); + assert!( + serde_json::to_string(&restored.render_tool_results_in_messages()) + .unwrap() + .len() + > 50 * 1024 + ); +} diff --git a/crates/agent_core/src/tree.rs b/crates/agent_core/src/tree.rs index 0c95df1b..be97da9d 100644 --- a/crates/agent_core/src/tree.rs +++ b/crates/agent_core/src/tree.rs @@ -36,3 +36,114 @@ pub struct MessageNode { )] pub extension: Option, } +/// Runtime-owned conversation. Only the tree is writable; the linear history +/// is a derived cache, never a prompt recovery workspace or restore authority. +/// Kept crate-private so the persisted/public tree representation stays stable. +pub(crate) struct Conversation { + nodes: std::collections::BTreeMap, + path: ConversationPath, + next_id: NodeId, + history: Vec, +} + +impl Default for Conversation { + fn default() -> Self { + Self::restore(Default::default(), Vec::new(), 1, Vec::new()) + } +} + +impl Conversation { + pub(crate) fn restore( + nodes: std::collections::BTreeMap, + path: ConversationPath, + next_id: NodeId, + legacy_messages: Vec, + ) -> Self { + let mut conversation = Self { + next_id: next_id.max(nodes.keys().next_back().copied().unwrap_or(0) + 1), + nodes, + path, + history: Vec::new(), + }; + if conversation.nodes.is_empty() { + conversation.path.clear(); + for message in legacy_messages { + let id = conversation.reserve_id(); + conversation.append(message, id); + } + } + conversation.rebuild_history(); + conversation + } + + pub(crate) fn nodes(&self) -> &std::collections::BTreeMap { + &self.nodes + } + + pub(crate) fn path(&self) -> &ConversationPath { + &self.path + } + + pub(crate) fn next_id(&self) -> NodeId { + self.next_id + } + + pub(crate) fn history(&self) -> &[Message] { + &self.history + } + + pub(crate) fn reserve_id(&mut self) -> NodeId { + let id = self.next_id; + self.next_id += 1; + id + } + + pub(crate) fn append(&mut self, message: Message, id: NodeId) { + assert!( + !self.nodes.contains_key(&id), + "message node id already exists" + ); + self.next_id = self.next_id.max(id + 1); + self.nodes.insert( + id, + MessageNode { + id, + message, + parent_id: self.path.last().copied(), + created_at: SystemTime::now(), + extension: None, + }, + ); + self.path.push(id); + self.rebuild_history(); + } + + /// Persistent correction of one active-path message (content, usage, etc.). + /// Node identity, parent links, extensions and inactive branches survive. + pub(crate) fn edit_message(&mut self, id: NodeId, edit: impl FnOnce(&mut Message)) { + if let Some(node) = self.nodes.get_mut(&id) { + edit(&mut node.message); + self.rebuild_history(); + } + } + + /// Compatibility boundary for existing hooks that take mutable tree nodes. + /// Re-derive history once after the hook batch, including early results. + pub(crate) fn with_nodes_mut( + &mut self, + edit: impl FnOnce(&mut std::collections::BTreeMap, &ConversationPath) -> T, + ) -> T { + let result = edit(&mut self.nodes, &self.path); + self.rebuild_history(); + result + } + + fn rebuild_history(&mut self) { + self.history = self + .path + .iter() + .filter_map(|id| self.nodes.get(id)) + .map(|node| node.message.clone()) + .collect(); + } +} diff --git a/crates/code_assistant_core/src/agent/checkpoint_tests.rs b/crates/code_assistant_core/src/agent/checkpoint_tests.rs new file mode 100644 index 00000000..32bef056 --- /dev/null +++ b/crates/code_assistant_core/src/agent/checkpoint_tests.rs @@ -0,0 +1,175 @@ +use super::*; +use crate::agent::persistence::AgentStatePersistence; +use crate::persistence::{ChatSession, FileSessionPersistence, MessageNode}; +use std::sync::Mutex; + +#[derive(Clone, Default)] +struct Capture(Arc>>); +impl AgentStatePersistence for Capture { + fn save_agent_state(&mut self, state: SessionState) -> Result<()> { + *self.0.lock().unwrap() = Some(state); + Ok(()) + } +} + +/// Deterministic format-on-save without depending on an installed formatter. +struct FormattingTool; +#[async_trait::async_trait] +impl tools_core::Tool for FormattingTool { + type Input = serde_json::Value; + type Output = agent_core::types::ParseError; + fn spec(&self) -> tools_core::ToolSpec { + crate::tools::test_registry() + .get("write_file") + .unwrap() + .spec() + } + async fn execute<'a>( + &self, + _: &mut tools_core::ToolContext<'a>, + input: &mut Self::Input, + ) -> Result { + input["content"] = "formatted content\n".into(); + Ok(agent_core::types::ParseError::new("test output".into())) + } +} + +async fn formatted_roundtrip(syntax: ToolSyntax) -> Result<()> { + let dir = tempdir()?; + let mut registry = tools_core::ToolRegistry::new(); + registry.register(Box::new(FormattingTool)); + let registry = Arc::new(registry); + let dialect = crate::tool_dialects::dialect_for(syntax); + let request = crate::tools::ToolRequest { + id: "call".into(), + name: "write_file".into(), + input: serde_json::json!({"project":"test", "path":"test.txt", "content":"unformatted"}), + start_offset: None, + end_offset: None, + }; + let content = if syntax == ToolSyntax::Native { + vec![ContentBlock::new_tool_use( + &request.id, + &request.name, + request.input.clone(), + )] + } else { + vec![ + // Offsets belong to the tool-containing block, not this preamble. + ContentBlock::new_text("Preamble remains intact.\n"), + ContentBlock::new_text(format!( + "{}\nTRAILING TEXT MUST BE TRUNCATED", + dialect.format_tool_request(&request, ®istry)? + )), + ] + }; + let mock_llm = MockLLMProvider::new(vec![ + Ok(create_test_response_text("done")), + Ok(LLMResponse { + content, + usage: Usage::zero(), + rate_limit_info: None, + }), + ]); + let captured = Capture::default(); + let components = AgentComponents { + llm_provider: Box::new(mock_llm), + project_manager: Arc::new(MockProjectManager::new()), + command_executor: Arc::new(create_command_executor_mock()), + ui: Arc::new(MockUI::default()), + state_persistence: Box::new(captured.clone()), + permission_handler: None, + permissions: Default::default(), + tool_registry: registry.clone(), + sub_agent_runner: None, + wakeups: None, + pty_sessions: None, + browser_sessions: None, + terminal_interrupts: None, + session_source: None, + hooks_factory: None, + }; + let config = SessionConfig { + tool_syntax: syntax, + ..Default::default() + }; + let mut agent = Agent::new(components, config.clone()); + let mut initial = SessionState::from_messages( + "checkpoint", + "test", + vec![Message::new_user("write")], + config.clone(), + ); + initial.message_nodes.insert( + 99, + MessageNode { + id: 99, + message: Message::new_assistant("inactive branch"), + parent_id: Some(1), + created_at: std::time::SystemTime::now(), + extension: Some(serde_json::json!({"branch-data": true})), + }, + ); + initial.next_node_id = 100; + agent.load_from_session_state(initial).await?; + agent.run_single_iteration().await?; + let state = captured.0.lock().unwrap().take().unwrap(); + let mut session = ChatSession::new_empty("checkpoint".into(), "test".into(), config, None); + session.message_nodes = state.message_nodes; + session.active_path = state.active_path; + session.next_node_id = state.next_node_id; + session.tool_executions = state + .tool_executions + .iter() + .map(|e| e.serialize()) + .collect::>()?; + let mut persistence = FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()); + persistence.save_chat_session(&session)?; + let loaded = persistence.load_chat_session("checkpoint")?.unwrap(); + assert_eq!( + loaded.message_nodes[&99].extension, + Some(serde_json::json!({"branch-data": true})) + ); + assert!(!loaded.active_path.contains(&99)); + let canonical = loaded.get_active_messages_cloned(); + assert_eq!( + serde_json::to_value(&canonical)?, + serde_json::to_value(&state.messages)? + ); + let tool_message = &canonical[1]; + let MessageContent::Structured(blocks) = &tool_message.content else { + panic!("tool message") + }; + let response = LLMResponse { + content: blocks.clone(), + usage: Usage::zero(), + rate_limit_info: None, + }; + let (parsed, _) = + dialect.extract_requests(&response, tool_message.request_id.unwrap(), 0, ®istry)?; + assert_eq!(parsed.len(), 1); + // XML's multiline delimiters include surrounding newlines in parsed input. + assert_eq!( + parsed[0].input["content"].as_str().unwrap().trim(), + "formatted content" + ); + let serialized = serde_json::to_string(tool_message)?; + if syntax != ToolSyntax::Native { + assert!(serialized.contains("Preamble remains intact.")); + assert!(!serialized.contains("TRAILING TEXT MUST BE TRUNCATED")); + } + Ok(()) +} + +#[tokio::test] +async fn checkpoint_native_formatted_roundtrip() -> Result<()> { + formatted_roundtrip(ToolSyntax::Native).await +} +#[tokio::test] +async fn checkpoint_xml_formatted_roundtrip() -> Result<()> { + formatted_roundtrip(ToolSyntax::Xml).await +} +#[tokio::test] +async fn checkpoint_caret_formatted_roundtrip() -> Result<()> { + formatted_roundtrip(ToolSyntax::Caret).await +} diff --git a/crates/code_assistant_core/src/agent/persistence.rs b/crates/code_assistant_core/src/agent/persistence.rs index bc3727db..4aab4dd1 100644 --- a/crates/code_assistant_core/src/agent/persistence.rs +++ b/crates/code_assistant_core/src/agent/persistence.rs @@ -58,6 +58,8 @@ impl SnapshotPersistence for SessionStateAdapter { tool_executions: snapshot.tool_executions, plan: state.plan.clone(), active_skills: state.active_skills.clone(), + // Compatibility restore inputs only. The session manager owns the + // persistent settings and never applies a run's config on save. config: state.session_config.clone(), next_request_id: Some(snapshot.next_request_id), model_config: state.model_config.clone(), diff --git a/crates/code_assistant_core/src/agent/tests.rs b/crates/code_assistant_core/src/agent/tests.rs index 4c0c23f5..5fc7c32b 100644 --- a/crates/code_assistant_core/src/agent/tests.rs +++ b/crates/code_assistant_core/src/agent/tests.rs @@ -1,3 +1,6 @@ +#[path = "checkpoint_tests.rs"] +mod checkpoint_tests; + use super::*; use crate::agent::persistence::MockStatePersistence; use crate::mocks::MockLLMProvider; @@ -1596,7 +1599,11 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { agent.load_from_session_state(session_state).await?; let history = agent.message_history_for_tests(); - assert_eq!(history.len(), 1); + assert_eq!(history.len(), 2); + assert!(matches!(history[1].role, MessageRole::Assistant)); + let prompt = agent.render_tool_results_in_messages(); + assert_eq!(prompt.len(), 3); + assert!(serde_json::to_string(&prompt[2])?.contains("unknown")); assert!(matches!(history[0].role, MessageRole::User)); Ok(()) @@ -1660,14 +1667,16 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() agent.load_from_session_state(session_state).await?; let history = agent.message_history_for_tests(); - assert_eq!(history.len(), 2); + assert_eq!(history.len(), 3); assert!(matches!(history[0].role, MessageRole::User)); - assert!(matches!(history[1].role, MessageRole::User)); + assert!(matches!(history[1].role, MessageRole::Assistant)); + assert!(matches!(history[2].role, MessageRole::User)); + assert_eq!(agent.render_tool_results_in_messages().len(), 4); match &history[0].content { MessageContent::Text(content) => assert_eq!(content, "Please inspect the project."), _ => panic!("Expected initial user message to be preserved"), } - match &history[1].content { + match &history[2].content { MessageContent::Text(content) => assert_eq!(content, "Also check the contributing guide."), _ => panic!("Expected follow-up user message to be preserved"), } @@ -1726,7 +1735,11 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { agent.load_from_session_state(session_state).await?; let history = agent.message_history_for_tests(); - assert_eq!(history.len(), 1); + assert_eq!(history.len(), 2); + assert!(matches!(history[1].role, MessageRole::Assistant)); + let prompt = agent.render_tool_results_in_messages(); + assert_eq!(prompt.len(), 3); + assert!(serde_json::to_string(&prompt[2])?.contains("unknown")); assert!(matches!(history[0].role, MessageRole::User)); Ok(()) @@ -1788,12 +1801,10 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() } #[tokio::test] -async fn test_render_tool_results_generates_cancelled_results_for_missing_executions() -> Result<()> -{ +async fn test_render_tool_results_generates_unknown_results_for_missing_executions() -> Result<()> { // This test verifies that when an assistant message contains ToolUse blocks - // but there's no corresponding ToolResult in the message history (because the - // user cancelled the tool execution), we generate synthetic "user cancelled" - // ToolResult blocks to satisfy the API requirement. + // but no corresponding ToolResult, the prompt supplies an unknown outcome + // rather than claiming cancellation or silently repeating side effects. let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { @@ -1825,7 +1836,7 @@ async fn test_render_tool_results_generates_cancelled_results_for_missing_execut // Simulate a scenario where: // 1. User asks a question // 2. Assistant responds with a tool call - // 3. User cancels the tool execution (no ToolResult message added) + // 3. User interrupts the tool execution (no ToolResult message added) // 4. User asks a follow-up question // Add user message @@ -1847,7 +1858,7 @@ async fn test_render_tool_results_generates_cancelled_results_for_missing_execut .with_request_id(1), )?; - // Note: We do NOT add a ToolResult message - simulating user cancellation + // Note: We do NOT add a ToolResult message - simulating user interruption // Add another user message (user continues the conversation) agent.append_message(Message::new_user("Never mind, let's do something else."))?; @@ -1858,19 +1869,19 @@ async fn test_render_tool_results_generates_cancelled_results_for_missing_execut // We should have: // 1. Original user message // 2. Assistant message with tool call - // 3. Synthetic user message with cancelled tool result + // 3. Synthetic user message with unknown tool result // 4. Follow-up user message assert_eq!( rendered_messages.len(), 4, - "Expected 4 messages: user, assistant, cancelled tool result, follow-up user" + "Expected 4 messages: user, assistant, unknown tool result, follow-up user" ); - // Verify the synthetic cancelled tool result was inserted - let cancelled_message = &rendered_messages[2]; - assert_eq!(cancelled_message.role, MessageRole::User); + // Verify the synthetic unknown tool result was inserted + let unknown_message = &rendered_messages[2]; + assert_eq!(unknown_message.role, MessageRole::User); - if let MessageContent::Structured(blocks) = &cancelled_message.content { + if let MessageContent::Structured(blocks) = &unknown_message.content { assert_eq!(blocks.len(), 1); if let ContentBlock::ToolResult { tool_use_id, @@ -1880,13 +1891,13 @@ async fn test_render_tool_results_generates_cancelled_results_for_missing_execut } = &blocks[0] { assert_eq!(tool_use_id, "tool-1-1"); - assert!(content.contains("cancelled")); + assert!(content.contains("unknown")); assert!(is_error.unwrap_or(false)); } else { panic!("Expected ToolResult block"); } } else { - panic!("Expected Structured content for cancelled tool result"); + panic!("Expected Structured content for unknown tool result"); } // Verify the follow-up user message is still present @@ -1904,7 +1915,7 @@ async fn test_render_tool_results_generates_cancelled_results_for_missing_execut #[tokio::test] async fn test_render_tool_results_preserves_existing_tool_results() -> Result<()> { // This test verifies that when tool results already exist, we don't add - // synthetic cancelled results for them. + // synthetic unknown results for them. let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { @@ -1966,14 +1977,14 @@ async fn test_render_tool_results_preserves_existing_tool_results() -> Result<() // Now call render_tool_results_in_messages let rendered_messages = agent.render_tool_results_in_messages(); - // We should have exactly 3 messages - no synthetic cancelled results added + // We should have exactly 3 messages - no synthetic unknown results added assert_eq!( rendered_messages.len(), 3, "Expected 3 messages: user, assistant, tool result" ); - // Verify the tool result is the original one (not a cancelled one) + // Verify the tool result is the original one (not a unknown one) let result_message = &rendered_messages[2]; if let MessageContent::Structured(blocks) = &result_message.content && let ContentBlock::ToolResult { @@ -1984,7 +1995,7 @@ async fn test_render_tool_results_preserves_existing_tool_results() -> Result<() } = &blocks[0] { assert_eq!(tool_use_id, "tool-1-1"); - // Content should be the original, not "cancelled" + // Content should be the original, not "unknown" assert!(content.contains("File contents") || content.is_empty()); // Should not be marked as error assert!(!is_error.unwrap_or(false)); @@ -1994,8 +2005,8 @@ async fn test_render_tool_results_preserves_existing_tool_results() -> Result<() } #[tokio::test] -async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<()> { - // This test verifies that multiple cancelled tool calls are all handled correctly. +async fn test_render_tool_results_handles_multiple_unknown_tools() -> Result<()> { + // This test verifies that multiple unknown tool calls are all handled correctly. let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { @@ -2027,7 +2038,7 @@ async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<( // Add user message agent.append_message(Message::new_user("Check the project."))?; - // Add assistant message with multiple tool calls (all cancelled) + // Add assistant message with multiple tool calls (all unknown) agent.append_message( Message::new_assistant_content(vec![ ContentBlock::new_text("I'll check multiple things."), @@ -2051,7 +2062,7 @@ async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<( .with_request_id(1), )?; - // No tool results added - both cancelled + // No tool results added - both unknown // Now call render_tool_results_in_messages let rendered_messages = agent.render_tool_results_in_messages(); @@ -2059,21 +2070,21 @@ async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<( // We should have: // 1. Original user message // 2. Assistant message with tool calls - // 3. Synthetic user message with both cancelled tool results + // 3. Synthetic user message with both unknown tool results assert_eq!( rendered_messages.len(), 3, - "Expected 3 messages: user, assistant, cancelled tool results" + "Expected 3 messages: user, assistant, unknown tool results" ); - // Verify the synthetic cancelled results - let cancelled_message = &rendered_messages[2]; - assert_eq!(cancelled_message.role, MessageRole::User); + // Verify the synthetic unknown results + let unknown_message = &rendered_messages[2]; + assert_eq!(unknown_message.role, MessageRole::User); - if let MessageContent::Structured(blocks) = &cancelled_message.content { - assert_eq!(blocks.len(), 2, "Should have 2 cancelled tool results"); + if let MessageContent::Structured(blocks) = &unknown_message.content { + assert_eq!(blocks.len(), 2, "Should have 2 unknown tool results"); - // Check first cancelled result + // Check first unknown result if let ContentBlock::ToolResult { tool_use_id, content, @@ -2082,13 +2093,13 @@ async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<( } = &blocks[0] { assert_eq!(tool_use_id, "tool-1-1"); - assert!(content.contains("cancelled")); + assert!(content.contains("unknown")); assert!(is_error.unwrap_or(false)); } else { - panic!("Expected ToolResult block for first cancelled tool"); + panic!("Expected ToolResult block for first unknown tool"); } - // Check second cancelled result + // Check second unknown result if let ContentBlock::ToolResult { tool_use_id, content, @@ -2097,13 +2108,13 @@ async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<( } = &blocks[1] { assert_eq!(tool_use_id, "tool-1-2"); - assert!(content.contains("cancelled")); + assert!(content.contains("unknown")); assert!(is_error.unwrap_or(false)); } else { - panic!("Expected ToolResult block for second cancelled tool"); + panic!("Expected ToolResult block for second unknown tool"); } } else { - panic!("Expected Structured content for cancelled tool results"); + panic!("Expected Structured content for unknown tool results"); } Ok(()) @@ -2360,8 +2371,7 @@ async fn test_prompt_too_long_fallback_drops_exchange_and_compacts() -> Result<( "Expected compaction summary in message history" ); - // The dropped exchange (assistant tool_use + user tool_result) should no longer - // be in the message history + // Compaction omits the exchange from the prompt, not from canonical history. let has_tool_result = agent.message_history_for_tests().iter().any(|msg| { if let MessageContent::Structured(blocks) = &msg.content { blocks.iter().any(|b| { @@ -2378,8 +2388,8 @@ async fn test_prompt_too_long_fallback_drops_exchange_and_compacts() -> Result<( } }); assert!( - !has_tool_result, - "Expected the dropped tool result to be removed from message history" + has_tool_result, + "The prompt projection must retain canonical tool evidence" ); // Verify UI received compaction divider diff --git a/crates/code_assistant_core/src/persistence.rs b/crates/code_assistant_core/src/persistence.rs index 2e0b1722..8f8767bb 100644 --- a/crates/code_assistant_core/src/persistence.rs +++ b/crates/code_assistant_core/src/persistence.rs @@ -734,7 +734,42 @@ impl FileSessionPersistence { self.ensure_chats_dir() } + fn entry_lock_path(&self, session_id: &str) -> Result { + Ok(self + .ensure_chats_dir()? + .join(format!("{session_id}.entry.lock"))) + } + + /// Update an existing session under a cross-process, per-entry lock. + /// The closure sees the latest on-disk entry; an error leaves it unchanged. + /// Lock order is entry -> metadata. This is separate from the long-lived + /// agent lock so settings can still change during a run. Do not re-enter + /// persistence from the closure. Lock files must never be unlinked. + pub fn update_entry( + &mut self, + session_id: &str, + update: impl FnOnce(&mut ChatSession) -> Result<()>, + ) -> Result { + let _lock = lock_exclusive(&self.entry_lock_path(session_id)?)?; + let mut session = self + .load_chat_session(session_id)? + .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; + update(&mut session)?; + anyhow::ensure!(session.id == session_id, "Cannot change session identity"); + session.ensure_config()?; + self.save_chat_session_unlocked(&session)?; + Ok(session) + } + + /// Full replacement, retained for creation and legacy callers. The lock + /// serializes writes but cannot make a stale supplied snapshot current; + /// read-modify-write callers must use `update_entry` instead. pub fn save_chat_session(&mut self, session: &ChatSession) -> Result<()> { + let _lock = lock_exclusive(&self.entry_lock_path(&session.id)?)?; + self.save_chat_session_unlocked(session) + } + + fn save_chat_session_unlocked(&mut self, session: &ChatSession) -> Result<()> { let mut session = session.clone(); session.ensure_config()?; @@ -852,6 +887,7 @@ impl FileSessionPersistence { } pub fn delete_chat_session(&mut self, session_id: &str) -> Result<()> { + let _entry_lock = lock_exclusive(&self.entry_lock_path(session_id)?)?; // Remove the session file let session_path = self.chat_file_path(session_id)?; if session_path.exists() { @@ -1316,6 +1352,90 @@ mod tests { use base64::Engine as _; use tempfile::tempdir; + #[test] + fn checkpoint_update_entry_serializes_independent_persistence_instances() { + let dir = tempdir().unwrap(); + let mut persistence = FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()); + persistence + .save_chat_session(&ChatSession::new_empty( + "shared".into(), + "shared".into(), + SessionConfig::default(), + None, + )) + .unwrap(); + let start = Arc::new(std::sync::Barrier::new(4)); + std::thread::scope(|scope| { + for _ in 0..4 { + let root = dir.path().to_path_buf(); + let start = start.clone(); + scope.spawn(move || { + let mut persistence = FileSessionPersistence::new_with_root_dir(root); + start.wait(); + for _ in 0..10 { + persistence + .update_entry("shared", |session| { + let previous = session.next_request_id; + std::thread::sleep(std::time::Duration::from_millis(1)); + session.next_request_id = previous + 1; + session.add_message(Message::new_user("concurrent append")); + Ok(()) + }) + .unwrap(); + } + }); + } + }); + let saved = persistence.load_chat_session("shared").unwrap().unwrap(); + assert_eq!(saved.next_request_id, 41); + assert_eq!(saved.message_count(), 40); + assert_eq!(saved.get_active_messages().len(), 40); + assert_eq!( + persistence + .get_chat_session_metadata("shared") + .unwrap() + .unwrap() + .message_count, + 40 + ); + } + + #[test] + fn checkpoint_update_entry_error_does_not_write_or_create() { + let dir = tempdir().unwrap(); + let mut persistence = FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()); + assert!(persistence.update_entry("missing", |_| Ok(())).is_err()); + assert!(persistence.load_chat_session("missing").unwrap().is_none()); + persistence + .save_chat_session(&ChatSession::new_empty( + "existing".into(), + "original".into(), + SessionConfig::default(), + None, + )) + .unwrap(); + let before = std::fs::read(persistence.chat_file_path("existing").unwrap()).unwrap(); + assert!( + persistence + .update_entry("existing", |session| { + session.name = "not committed".into(); + anyhow::bail!("abort update") + }) + .is_err() + ); + assert_eq!( + before, + std::fs::read(persistence.chat_file_path("existing").unwrap()).unwrap() + ); + // The failed transaction also released its lock. + persistence + .update_entry("existing", |session| { + session.name = "committed".into(); + Ok(()) + }) + .unwrap(); + } + fn oversized_png_base64(width: u32, height: u32) -> String { let img = image::DynamicImage::ImageRgba8(image::RgbaImage::from_pixel( width, diff --git a/crates/code_assistant_core/src/session/manager.rs b/crates/code_assistant_core/src/session/manager.rs index 7c2da331..4191f22f 100644 --- a/crates/code_assistant_core/src/session/manager.rs +++ b/crates/code_assistant_core/src/session/manager.rs @@ -547,7 +547,14 @@ impl SessionManager { let session_instance = self.active_sessions.get(&session_id).unwrap(); session_instance.session.clone() }; - self.persistence.save_chat_session(&session_snapshot)?; + self.persistence.update_entry(&session_id, |session| { + // Backfill only if a concurrent settings update has not already + // selected a model. Never replace its conversation or config. + if session.model_config.is_none() { + session.model_config = session_snapshot.model_config; + } + Ok(()) + })?; } Ok(snapshot) @@ -771,6 +778,23 @@ impl SessionManager { .collect() } + fn initialize_session_project( + &mut self, + session_id: &str, + project_manager: &dyn ProjectManager, + ) -> Result<()> { + self.persistence.update_entry(session_id, |session| { + if session.config.initial_project.is_empty() + && let Some(path) = session.config.effective_project_path() + { + session.config.initial_project = + project_manager.add_temporary_project(path.clone())?; + } + Ok(()) + })?; + Ok(()) + } + /// Start an agent for a session (message must already be added via add_user_message) /// This is the key method - agents run on-demand for specific messages /// @@ -851,6 +875,10 @@ impl SessionManager { ) })?; + // Project grouping is session-owned initialization, not a side effect + // of a later agent checkpoint carrying its stale run configuration. + self.initialize_session_project(session_id, project_manager.as_ref())?; + // Prepare session - need to scope the mutable borrow carefully let ( session_config, @@ -1422,33 +1450,30 @@ impl SessionManager { session_id: &str, model_config: Option, ) -> Result { - let mut session = self - .persistence - .load_chat_session(session_id)? - .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; - - if let Some(new_config) = model_config.as_ref() { - let check = Self::check_model_switch_for_session(&session, new_config)?; - if !check.allowed { - anyhow::bail!(check.error_message()); + let use_diff_blocks = self.resolve_use_diff_blocks(model_config.as_ref()); + let session = self.persistence.update_entry(session_id, |session| { + if let Some(new_config) = model_config.as_ref() { + let check = Self::check_model_switch_for_session(session, new_config)?; + if !check.allowed { + anyhow::bail!(check.error_message()); + } + } else if session.message_count() > 0 || !session.tool_executions.is_empty() { + anyhow::bail!( + "Cannot clear the model for a session after a conversation has started." + ); } - } else if session.message_count() > 0 || !session.tool_executions.is_empty() { - anyhow::bail!("Cannot clear the model for a session after a conversation has started."); - } - + if session.message_count() == 0 && session.tool_executions.is_empty() { + session.config.use_diff_blocks = use_diff_blocks; + } + session.model_config = model_config.clone(); + Ok(()) + })?; let is_empty = session.message_count() == 0 && session.tool_executions.is_empty(); - if is_empty { - session.config.use_diff_blocks = self.resolve_use_diff_blocks(model_config.as_ref()); - } - let agent_running = self .active_sessions .get(session_id) .is_some_and(|instance| !instance.get_activity_state().is_terminal()); - session.model_config = model_config.clone(); - self.persistence.save_chat_session(&session)?; - if let Some(instance) = self.active_sessions.get_mut(session_id) { instance.session.model_config = model_config; if is_empty { @@ -1470,13 +1495,10 @@ impl SessionManager { session_id: &str, policy: SandboxPolicy, ) -> Result<()> { - let mut session = self - .persistence - .load_chat_session(session_id)? - .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; - - session.config.sandbox_policy = policy.clone(); - self.persistence.save_chat_session(&session)?; + self.persistence.update_entry(session_id, |session| { + session.config.sandbox_policy = policy.clone(); + Ok(()) + })?; if let Some(instance) = self.active_sessions.get_mut(session_id) { instance.session.config.sandbox_policy = policy; @@ -1492,13 +1514,10 @@ impl SessionManager { session_id: &str, tier: tools_core::PermissionTier, ) -> Result<()> { - let mut session = self - .persistence - .load_chat_session(session_id)? - .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; - - session.config.permission_tier = tier; - self.persistence.save_chat_session(&session)?; + self.persistence.update_entry(session_id, |session| { + session.config.permission_tier = tier; + Ok(()) + })?; if let Some(instance) = self.active_sessions.get_mut(session_id) { instance.session.config.permission_tier = tier; @@ -1520,13 +1539,10 @@ impl SessionManager { session_id: &str, disabled: Vec, ) -> Result> { - let mut session = self - .persistence - .load_chat_session(session_id)? - .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; - - session.config.disabled_mcp_servers = disabled.clone(); - self.persistence.save_chat_session(&session)?; + let session = self.persistence.update_entry(session_id, |session| { + session.config.disabled_mcp_servers = disabled.clone(); + Ok(()) + })?; if let Some(instance) = self.active_sessions.get_mut(session_id) { instance.session.config.disabled_mcp_servers = disabled.clone(); @@ -1588,14 +1604,11 @@ impl SessionManager { worktree_path: Option, branch: Option, ) -> Result<()> { - let mut session = self - .persistence - .load_chat_session(session_id)? - .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; - - session.config.worktree_path = worktree_path.clone(); - session.config.branch = branch.clone(); - self.persistence.save_chat_session(&session)?; + self.persistence.update_entry(session_id, |session| { + session.config.worktree_path = worktree_path.clone(); + session.config.branch = branch.clone(); + Ok(()) + })?; if let Some(instance) = self.active_sessions.get_mut(session_id) { instance.session.config.worktree_path = worktree_path.clone(); @@ -1690,51 +1703,47 @@ impl SessionManager { .save_chat_session(&session_instance.session) } - /// Save agent state to a specific session + /// Save only run-owned conversation state. `state.config` and + /// `state.model_config` are restore/run inputs, never checkpoint writes. pub fn save_session_state(&mut self, state: SessionState) -> Result<()> { - let mut session = self - .persistence - .load_chat_session(&state.session_id)? - .ok_or_else(|| anyhow::anyhow!("Session not found: {}", state.session_id))?; - - // Preserve session-level settings that can be changed outside the - // running agent. A model switch while an agent is running should take - // effect on the next iteration; the old agent must not overwrite it - // when it saves its captured state. - let persisted_model_config = session.model_config.clone(); - let persisted_use_diff_blocks = session.config.use_diff_blocks; - - // Update session with current state - session.name = state.name; - - // Update tree structure. - session.message_nodes = state.message_nodes; - session.active_path = state.active_path; - session.next_node_id = state.next_node_id; - - // Clear legacy messages (tree is now authoritative) - session.messages.clear(); - - session.tool_executions = state + let session_id = state.session_id.clone(); + let executions = state .tool_executions .into_iter() .map(|te| te.serialize()) .collect::>>()?; - session.plan = state.plan; - session.active_skills = state.active_skills; - session.config = state.config; - session.config.use_diff_blocks = persisted_use_diff_blocks; - session.model_config = persisted_model_config; - session.next_request_id = state.next_request_id.unwrap_or(0); - session.updated_at = SystemTime::now(); - - self.persistence.save_chat_session(&session)?; + let session = self.persistence.update_entry(&session_id, |session| { + session.name = state.name; + // Retain branches not carried by this run. Active-path corrections + // replace nodes by id; concurrent conversation writers still require + // the existing single-agent/branch guards, not just this disk lock. + session.message_nodes.extend(state.message_nodes); + session.active_path = state.active_path; + session.next_node_id = session.next_node_id.max(state.next_node_id); + session.messages.clear(); + for execution in executions { + if let Some(existing) = session + .tool_executions + .iter_mut() + .find(|existing| existing.tool_request.id == execution.tool_request.id) + { + *existing = execution; + } else { + session.tool_executions.push(execution); + } + } + session.plan = state.plan; + session.active_skills = state.active_skills; + if let Some(next_id) = state.next_request_id { + session.next_request_id = session.next_request_id.max(next_id); + } + session.updated_at = SystemTime::now(); + Ok(()) + })?; - // Update active session instance if it exists - if let Some(instance) = self.active_sessions.get_mut(&state.session_id) { + if let Some(instance) = self.active_sessions.get_mut(&session_id) { instance.session = session; } - Ok(()) } @@ -1924,6 +1933,214 @@ mod tests { (manager, dir) } + #[test] + fn checkpoint_preserves_all_session_settings_after_external_changes() { + let (mut manager, dir) = build_manager(false); + let id = manager.create_session(None).unwrap(); + let captured = SessionState::from_messages( + id.clone(), + "run", + vec![Message::new_user("task")], + SessionConfig::default(), + ); + // Another manager owns the settings, independently of the run's manager. + manager.save_session_state(captured.clone()).unwrap(); + let mut settings = SessionManager::new( + FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()), + SessionConfig::default(), + "test-model".into(), + crate::tools::test_registry(), + crate::session::event_stream::EventStream::new(), + ); + settings + .set_session_permission_tier(&id, tools_core::PermissionTier::AllTools) + .unwrap(); + settings + .set_session_sandbox_policy(&id, SandboxPolicy::ReadOnly) + .unwrap(); + settings + .set_session_worktree( + &id, + Some(dir.path().join("worktree")), + Some("new-branch".into()), + ) + .unwrap(); + settings + .set_session_disabled_mcp_servers(&id, vec!["disabled-server".into()]) + .unwrap(); + let mut expected = settings + .persistence + .load_chat_session(&id) + .unwrap() + .unwrap(); + expected.config.init_path = Some(dir.path().join("project")); + expected.config.initial_project = "session-owned-project".into(); + expected.config.tool_syntax = crate::types::ToolSyntax::Caret; + expected.config.use_diff_blocks = true; + expected.model_config = Some(SessionModelConfig::new("new-model".into())); + expected.plan_collapsed = true; + settings.persistence.save_chat_session(&expected).unwrap(); + + manager.save_session_state(captured).unwrap(); + let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + assert_eq!( + serde_json::to_value(&saved.config).unwrap(), + serde_json::to_value(&expected.config).unwrap() + ); + assert_eq!(saved.model_config.as_ref().unwrap().model_name, "new-model"); + assert!(saved.plan_collapsed); + assert_eq!(saved.get_active_messages().len(), 1); + assert_eq!( + serde_json::to_value(&manager.get_session(&id).unwrap().session.config).unwrap(), + serde_json::to_value(&expected.config).unwrap() + ); + } + + #[test] + fn checkpoint_cannot_initialize_project_from_run_config() { + let (mut manager, _dir) = build_manager(false); + let id = manager.create_session(None).unwrap(); + let config = SessionConfig { + initial_project: "run-only".into(), + ..Default::default() + }; + manager + .save_session_state(SessionState::from_messages( + id.clone(), + "run", + vec![Message::new_user("task")], + config, + )) + .unwrap(); + let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + assert!(saved.config.initial_project.is_empty()); + } + + #[test] + fn checkpoint_settings_wait_for_entry_lock_and_read_latest_conversation() { + use std::sync::mpsc; + use std::time::Duration; + let (mut manager, _dir) = build_manager(false); + let id = manager.create_session(None).unwrap(); + let lock_path = manager + .persistence + .sessions_dir() + .unwrap() + .join(format!("{id}.entry.lock")); + let lock = file_utils::lock_exclusive(&lock_path).unwrap(); + let mut latest = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + let path = manager + .persistence + .sessions_dir() + .unwrap() + .join(format!("{id}.json")); + let (started_tx, started_rx) = mpsc::channel(); + let (done_tx, done_rx) = mpsc::channel(); + let worker_id = id.clone(); + let worker = std::thread::spawn(move || { + started_tx.send(()).unwrap(); + manager + .set_session_permission_tier(&worker_id, tools_core::PermissionTier::AllTools) + .unwrap(); + done_tx.send(()).unwrap(); + }); + started_rx.recv().unwrap(); + let blocked = done_rx.recv_timeout(Duration::from_millis(200)).is_err(); + // The holder of the entry lock writes a newer conversation before releasing it. + latest.add_message(Message::new_user("arrived during settings update")); + file_utils::atomic_write_json(&path, &latest).unwrap(); + drop(lock); + worker.join().unwrap(); + assert!( + blocked, + "settings must hold the entry lock across load and save" + ); + let saved: ChatSession = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert_eq!(saved.get_active_messages().len(), 1); + assert_eq!( + saved.config.permission_tier, + tools_core::PermissionTier::AllTools + ); + } + + #[test] + fn checkpoint_project_initialization_is_session_owned_and_stable() { + let (mut manager, _dir) = build_manager(false); + let config = SessionConfig { + init_path: Some(PathBuf::from("./project")), + ..Default::default() + }; + let id = manager + .create_session_with_config(None, Some(config.clone()), None) + .unwrap(); + let projects = crate::mocks::MockProjectManager::new(); + manager.initialize_session_project(&id, &projects).unwrap(); + let initialized = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + assert!(!initialized.config.initial_project.is_empty()); + manager + .set_session_worktree( + &id, + Some(PathBuf::from("./worktree")), + Some("branch".into()), + ) + .unwrap(); + manager.initialize_session_project(&id, &projects).unwrap(); + manager + .save_session_state(SessionState::from_messages( + id.clone(), + "run", + vec![Message::new_user("task")], + config, + )) + .unwrap(); + let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + assert_eq!( + saved.config.initial_project, + initialized.config.initial_project + ); + assert_eq!( + saved.config.worktree_path, + Some(PathBuf::from("./worktree")) + ); + } + + #[test] + fn checkpoint_keeps_branches_and_execution_records_not_loaded_by_run() { + let (mut manager, _dir) = build_manager(false); + let id = manager.create_session(None).unwrap(); + let captured = SessionState::from_messages( + id.clone(), + "run", + vec![Message::new_user("task")], + SessionConfig::default(), + ); + manager.save_session_state(captured.clone()).unwrap(); + let mut session = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + let branch = session.add_message(Message::new_assistant("another branch")); + session.message_nodes.get_mut(&branch).unwrap().extension = + Some(serde_json::json!({"snapshot": true})); + session.tool_executions.push( + agent_core::types::ToolExecution::create_parse_error( + "unavailable-tool".into(), + "retained evidence".into(), + ) + .serialize() + .unwrap(), + ); + manager.persistence.save_chat_session(&session).unwrap(); + manager.save_session_state(captured).unwrap(); + let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); + assert_eq!( + serde_json::to_value(&saved.message_nodes[&branch]).unwrap(), + serde_json::to_value(&session.message_nodes[&branch]).unwrap() + ); + assert_eq!(saved.next_node_id, session.next_node_id); + assert_eq!( + serde_json::to_value(&saved.tool_executions).unwrap(), + serde_json::to_value(&session.tool_executions).unwrap() + ); + } + fn provider_config( provider: &str, config: serde_json::Value, From ca6b02e6a8e62a94d446ef78f964fce86bb9c15f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 9 Sep 2026 08:32:26 +0200 Subject: [PATCH 02/15] refactor(session): own run preparation and cancellation outside command lanes --- crates/agent_core/src/runtime.rs | 66 +- .../code_assistant_core/src/agent/runner.rs | 4 + .../src/session/instance.rs | 27 +- .../src/session/manager.rs | 287 +++++-- .../src/session/permissions.rs | 78 +- .../src/session/service.rs | 722 +++++++++++++----- .../src/session/service/recovery_tests.rs | 316 ++++++++ .../src/session/sleep_inhibitor.rs | 19 +- .../code_assistant_core/src/session/turn.rs | 15 +- crates/tools_core/Cargo.toml | 1 + crates/tools_core/src/cancellation.rs | 75 ++ crates/tools_core/src/lib.rs | 2 + crates/tools_core/src/permissions.rs | 18 +- 13 files changed, 1322 insertions(+), 308 deletions(-) create mode 100644 crates/code_assistant_core/src/session/service/recovery_tests.rs create mode 100644 crates/tools_core/src/cancellation.rs diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index 27d0b5e2..a64e2ea1 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -90,6 +90,7 @@ pub struct AgentRuntime { permission_handler: Option>, permissions: ToolPermissions, + cancellation: tools_core::RunCancellation, conversation: Conversation, /// Run-local LLM projection. Never included in a checkpoint. @@ -160,6 +161,7 @@ impl AgentRuntime { services_provider, permission_handler, permissions, + cancellation: tools_core::RunCancellation::default(), conversation: Conversation::default(), prompt_projection: PromptProjection::default(), tool_executions: Vec::new(), @@ -171,6 +173,11 @@ impl AgentRuntime { } } + pub fn set_cancellation(&mut self, cancellation: tools_core::RunCancellation) { + self.permissions.set_cancellation(cancellation.clone()); + self.cancellation = cancellation; + } + /// Replace the dialect (e.g. after the embedding application reloaded a /// session with a different tool syntax). pub fn set_dialect(&mut self, dialect: Arc) { @@ -332,9 +339,17 @@ impl AgentRuntime { /// Run a single iteration of the agent loop without waiting for user input /// This is used in the new on-demand agent architecture pub async fn run_single_iteration(&mut self) -> Result<()> { + match self.run_until_complete().await { + Err(error) if error.is::() => Ok(()), + result => result, + } + } + + async fn run_until_complete(&mut self) -> Result<()> { let mut streaming_retry_count: u32 = 0; loop { + self.cancellation.check()?; // Check for pending user message and add it to history at start of each iteration if let Some(pending_blocks) = self.get_and_clear_pending_message() { let text_summary = text_summary_from_blocks(&pending_blocks); @@ -375,6 +390,7 @@ impl AgentRuntime { // `continue` restarts the loop, which re-renders the messages and // retries get_next_assistant_message. (StreamingStopped was already // sent by get_next_assistant_message in its error path.) + Err(e) if e.is::() => return Err(e), Err(e) => match self.hooks.recovery.classify(&e, streaming_retry_count) { RecoveryAction::ReduceContext => { self.recover_from_oversized_prompt().await?; @@ -388,13 +404,18 @@ impl AgentRuntime { streaming_retry_count = attempt; self.prepare_streaming_retry(&e, attempt, max_attempts, delay) .await; - tokio::time::sleep(delay).await; + tokio::select! { + biased; + _ = self.cancellation.cancelled() => return Err(tools_core::Cancelled.into()), + _ = tokio::time::sleep(delay) => {} + } continue; } RecoveryAction::Fail => return Err(e), }, }; + self.cancellation.check()?; // 2. Add original LLM response to message history using the pre-allocated node_id if !llm_response.content.is_empty() { self.append_message_with_node_id( @@ -549,6 +570,9 @@ impl AgentRuntime { // Process results in original order for (idx, tool_request) in tool_requests.iter().enumerate() { + if self.cancellation.is_cancelled() { + break; + } let result_block = if parallel_indices.len() > 1 && parallel_indices.contains(&idx) { // This request ran in parallel - get result from parallel execution @@ -615,6 +639,7 @@ impl AgentRuntime { let command_executor = self.command_executor.clone(); let permission_handler = self.permission_handler.clone(); let permissions = self.permissions.clone(); + let cancellation = self.cancellation.clone(); let services_provider = self.services_provider.clone(); let scope_tag = self.tool_capability.clone(); let excluded_capabilities = self.excluded_tool_capabilities.clone(); @@ -630,6 +655,7 @@ impl AgentRuntime { command_executor, permission_handler, permissions, + cancellation, services_provider, scope_tag, excluded_capabilities, @@ -681,6 +707,7 @@ impl AgentRuntime { command_executor: Arc, permission_handler: Option>, permissions: ToolPermissions, + cancellation: tools_core::RunCancellation, services_provider: Arc, scope_tag: String, excluded_capabilities: Vec, @@ -729,6 +756,7 @@ impl AgentRuntime { .await { Err(e) => Err(e), + Ok(()) if cancellation.is_cancelled() => Err(tools_core::Cancelled.into()), Ok(()) => { let mut services = services_provider.detached(&tool_request.id); let mut context = ToolContext { @@ -970,11 +998,14 @@ impl AgentRuntime { ))); let ui_for_callback = self.ui.clone(); + let cancellation = self.cancellation.clone(); let streaming_callback: StreamingCallback = Box::new(move |chunk: &StreamingChunk| { + cancellation.check()?; // Check if streaming should continue if !ui_for_callback.should_streaming_continue() { debug!("Streaming should stop - user requested cancellation"); - return Err(anyhow::anyhow!("Streaming cancelled by user")); + cancellation.cancel(); + return Err(tools_core::Cancelled.into()); } let mut processor_guard = processor @@ -986,15 +1017,16 @@ impl AgentRuntime { }); // Send message to LLM provider - let response = match self - .llm_provider - .send_message(request, Some(&streaming_callback)) - .await - { + let response_result = tokio::select! { + biased; + _ = self.cancellation.cancelled() => Err(tools_core::Cancelled.into()), + result = self.llm_provider.send_message(request, Some(&streaming_callback)) => result, + }; + let response = match response_result { Ok(response) => response, Err(e) => { // Check for streaming cancelled error - if e.to_string().contains("Streaming cancelled by user") { + if self.cancellation.is_cancelled() || e.is::() { debug!("Streaming cancelled by user in LLM request {}", request_id); // End LLM request with cancelled=true let _ = self @@ -1004,15 +1036,7 @@ impl AgentRuntime { error: None, }) .await; - // Return empty response - return Ok(( - llm::LLMResponse { - content: Vec::new(), - usage: llm::Usage::zero(), - rate_limit_info: None, - }, - request_id, - )); + return Err(tools_core::Cancelled.into()); } // For other errors, still end the request but not cancelled @@ -1097,7 +1121,11 @@ impl AgentRuntime { session_id: self.session_id.clone().unwrap_or_default(), }; - let response = self.llm_provider.send_message(request, None).await?; + let response = tokio::select! { + biased; + _ = self.cancellation.cancelled() => return Err(tools_core::Cancelled.into()), + result = self.llm_provider.send_message(request, None) => result?, + }; debug!( "Compaction response usage — Input: {}, Output: {}, Cache Read: {}", @@ -1649,6 +1677,7 @@ impl AgentRuntime { } async fn execute_tool(&mut self, tool_request: &ToolRequest) -> Result { + self.cancellation.check()?; debug!( "Executing tool request: {} (id: {})", tool_request.name, tool_request.id @@ -1734,6 +1763,7 @@ impl AgentRuntime { return Err(e); } + self.cancellation.check()?; // Create a tool context. The services provider builds the application // extension for this invocation (state such as the plan may move in // for the duration) and takes it back afterwards. diff --git a/crates/code_assistant_core/src/agent/runner.rs b/crates/code_assistant_core/src/agent/runner.rs index 174e37b6..ec6a1c24 100644 --- a/crates/code_assistant_core/src/agent/runner.rs +++ b/crates/code_assistant_core/src/agent/runner.rs @@ -147,6 +147,10 @@ impl Agent { } } + pub fn set_cancellation(&mut self, cancellation: tools_core::RunCancellation) { + self.runtime.set_cancellation(cancellation); + } + /// The application state riding on the loop. fn app_state(&self) -> &AgentAppState { AgentAppState::of_ref(self.runtime.extensions()) diff --git a/crates/code_assistant_core/src/session/instance.rs b/crates/code_assistant_core/src/session/instance.rs index b9ba8d0e..02e2e1a1 100644 --- a/crates/code_assistant_core/src/session/instance.rs +++ b/crates/code_assistant_core/src/session/instance.rs @@ -146,6 +146,10 @@ pub struct SessionInstance { // We only track the task handle, not the agent itself /// Task handle for the running agent (None if not running) pub task_handle: Option>>, + /// Owned preparation (LLM construction, MCP trust and registry). Never detached. + pub(crate) setup_task: Option>, + pub(crate) sleep_guard: Option, + pub(crate) cancellation: tools_core::RunCancellation, /// In-flight DisplayFragments of the currently streaming response. /// Written by the [`SessionEventPublisher`]; included in snapshots so a @@ -226,6 +230,12 @@ pub struct SessionInstance { pub tool_registry: Arc, } +impl Drop for SessionInstance { + fn drop(&mut self) { + self.terminate_agent(); + } +} + impl SessionInstance { /// Create a new session instance pub fn new(session: ChatSession, tool_registry: Arc) -> Self { @@ -240,6 +250,9 @@ impl SessionInstance { Self { session, task_handle: None, + setup_task: None, + sleep_guard: None, + cancellation: tools_core::RunCancellation::default(), fragment_buffer: Arc::new(Mutex::new(VecDeque::new())), tool_status_buffer: Arc::new(Mutex::new(HashMap::new())), in_flight_node_id: Arc::new(Mutex::new(None)), @@ -272,6 +285,7 @@ impl SessionInstance { /// Pending permission requests resolve as denied so the agent does not /// stay blocked waiting for an answer. pub fn request_stop(&self) { + self.cancellation.cancel(); self.stop_requested .store(true, std::sync::atomic::Ordering::Relaxed); self.pending_permission_requests.deny_all(); @@ -280,9 +294,10 @@ impl SessionInstance { /// Reset per-run state when a new agent starts: clears a previous stop /// request, the live tool-status map of the prior run, and any stale /// permission requests. - pub fn begin_agent_run(&self) { - self.stop_requested - .store(false, std::sync::atomic::Ordering::Relaxed); + pub fn begin_agent_run(&mut self) { + self.cancellation = tools_core::RunCancellation::default(); + self.activity = SessionActivity::default(); + self.stop_requested = Arc::new(std::sync::atomic::AtomicBool::new(false)); if let Ok(mut buf) = self.tool_status_buffer.lock() { buf.clear(); } @@ -321,12 +336,18 @@ impl SessionInstance { /// Terminate the running agent and release the cross-process agent lock. pub fn terminate_agent(&mut self) { + self.request_stop(); + if let Some(handle) = self.setup_task.take() { + handle.abort(); + } if let Some(handle) = self.task_handle.take() { handle.abort(); self.clear_fragment_buffer(); } // Release the cross-process agent lock self.agent_lock = None; + self.sleep_guard = None; + self.set_activity_state(SessionActivityState::Idle); } /// Add a message with optional branching support. diff --git a/crates/code_assistant_core/src/session/manager.rs b/crates/code_assistant_core/src/session/manager.rs index 4191f22f..a1ca5b67 100644 --- a/crates/code_assistant_core/src/session/manager.rs +++ b/crates/code_assistant_core/src/session/manager.rs @@ -34,6 +34,13 @@ pub struct RegistryRequest { pub include_local_mcp: bool, } +/// Immutable configuration captured when a run reserves its session. +/// Settings changed during preparation belong to a subsequent run. +pub(crate) struct RunConfig { + pub session: SessionConfig, + pub model: Option, +} + /// Provides the tool registry for the next agent run. Consulted at the /// start of every run, so embedders can rebuild the registry from their /// current configuration (e.g. reconnect MCP servers after a settings @@ -244,6 +251,75 @@ impl SessionManager { &self.tool_registry } + /// Clone the registry loader under the lock; poll it only after releasing it. + pub fn registry_loader(&self) -> ToolRegistryProvider { + self.tool_registry_provider.clone().unwrap_or_else(|| { + let registry = self.tool_registry.clone(); + Arc::new(move |_| { + let registry = registry.clone(); + Box::pin(async move { registry }) + }) + }) + } + + pub(crate) fn reserve_agent_run( + &mut self, + session_id: &str, + cancellation: tools_core::RunCancellation, + ) -> Result<()> { + self.ensure_session_loaded(session_id)?; + let instance = self.active_sessions.get_mut(session_id).unwrap(); + anyhow::ensure!( + instance.get_activity_state().is_terminal(), + "Session is already running" + ); + let lock = + file_utils::try_acquire_agent_lock(&self.persistence.sessions_dir()?, session_id)? + .ok_or_else(|| anyhow::anyhow!("Session is running in another instance"))?; + instance.begin_agent_run(); + instance.cancellation = cancellation; + instance.agent_lock = Some(lock); + instance.set_activity_state(crate::session::instance::SessionActivityState::AgentRunning); + instance.sleep_guard = Some(self.sleep_inhibitor.agent_guard()); + self.events.publish_ui( + session_id, + UiEvent::UpdateSessionActivityState { + session_id: session_id.to_string(), + activity_state: crate::session::instance::SessionActivityState::AgentRunning, + }, + ); + Ok(()) + } + + pub(crate) fn finish_failed_setup( + &mut self, + session_id: &str, + cancellation: &tools_core::RunCancellation, + error: String, + ) { + let Some(instance) = self.active_sessions.get_mut(session_id) else { + return; + }; + if !instance.cancellation.same_run(cancellation) || instance.agent_lock.is_none() { + return; + } + instance.agent_lock = None; + instance.sleep_guard = None; + let state = if cancellation.is_cancelled() { + crate::session::instance::SessionActivityState::Idle + } else { + crate::session::instance::SessionActivityState::Errored { message: error } + }; + instance.set_activity_state(state.clone()); + self.events.publish_ui( + session_id, + UiEvent::UpdateSessionActivityState { + session_id: session_id.to_string(), + activity_state: state, + }, + ); + } + /// Pull the registry for `req` from the provider (no-op without one) and /// assign it to `session_id`'s instance — and only that one. Sessions in /// other projects keep the registry of their own last run; MCP tool @@ -854,26 +930,74 @@ impl SessionManager { turn_recorder: Option>, registry_request: RegistryRequest, ) -> Result<()> { + let cancellation = turn_recorder + .as_ref() + .map(|r| r.cancellation.clone()) + .unwrap_or_default(); + self.reserve_agent_run(session_id, cancellation.clone())?; + let instance = self.active_sessions.get(session_id).unwrap(); + let run_config = RunConfig { + session: instance.session.config.clone(), + model: instance.session.model_config.clone(), + }; + self.refresh_tool_registry(session_id, registry_request) + .await; + let registry = self.tool_registry.clone(); + let result = self + .start_reserved_agent_for_session( + session_id, + llm_provider, + project_manager, + command_executor, + permission_handler, + tool_scope_override, + turn_recorder, + registry, + cancellation.clone(), + run_config, + ) + .await; + if let Err(error) = &result { + self.finish_failed_setup(session_id, &cancellation, format!("{error:#}")); + } + result + } + + /// Commit a prepared run. No MCP/LLM network work may occur under this lock. + #[allow(clippy::too_many_arguments)] + pub(crate) async fn start_reserved_agent_for_session( + &mut self, + session_id: &str, + llm_provider: Box, + project_manager: Box, + command_executor: Box, + permission_handler: Option>, + tool_scope_override: Option, + turn_recorder: Option>, + registry: Arc, + cancellation: tools_core::RunCancellation, + run_config: RunConfig, + ) -> Result<()> { + cancellation.check()?; + let instance = self + .active_sessions + .get(session_id) + .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; + anyhow::ensure!( + instance.cancellation.same_run(&cancellation), + "Run superseded" + ); + self.tool_registry = registry.clone(); + self.active_sessions + .get_mut(session_id) + .unwrap() + .tool_registry = registry; // A new run is the point where configuration changes take effect: pull // the registry the caller resolved for this run (scoped to the // session's project, and to whether its local `.mcp.json` was trusted) // before anything below binds to it. Trust is resolved by the caller — // not here — because prompting must happen without the manager lock // that wraps this call, or the response could never be delivered. - self.refresh_tool_registry(session_id, registry_request) - .await; - - // Acquire exclusive cross-process agent lock. - // This prevents another code-assistant instance from running an agent - // for the same session concurrently. - let sessions_dir = self.persistence.sessions_dir()?; - let agent_lock = file_utils::try_acquire_agent_lock(&sessions_dir, session_id)? - .ok_or_else(|| { - anyhow::anyhow!( - "Cannot start agent for session {session_id}: \ - another code-assistant instance is already running an agent for this session" - ) - })?; // Project grouping is session-owned initialization, not a side effect // of a later agent checkpoint carrying its stale run configuration. @@ -900,10 +1024,8 @@ impl SessionManager { // Clone all needed data to avoid borrowing conflicts let name = session_instance.session.name.clone(); - let session_config = session_instance.session.config.clone(); - // A new agent run supersedes any prior stop request and the - // previous run's live tool statuses. - session_instance.begin_agent_run(); + let session_config = run_config.session; + // The reservation already installed this run's fresh cancellation token. let publisher = session_instance.create_publisher(self.events.clone(), turn_recorder.clone()); @@ -949,7 +1071,7 @@ impl SessionManager { active_skills: session_instance.session.active_skills.clone(), config: session_config.clone(), next_request_id: Some(session_instance.session.next_request_id), - model_config: session_instance.session.model_config.clone(), + model_config: run_config.model, }; // Set activity state @@ -1099,22 +1221,7 @@ impl SessionManager { // Set the shared pending message reference agent.set_pending_message_ref(pending_message_ref); - // Load the session state into the agent - agent.load_from_session_state(session_state).await?; - - // Apply the per-run scope after the load, which derives the scope - // from the session config and would otherwise win. - if let Some(scope) = tool_scope_override { - agent.set_tool_scope(scope); - agent.invalidate_system_message_cache(); - } - - // Announce the restored plan to the UI - let _ = publisher - .send_event(UiEvent::UpdatePlan { - plan: agent.plan().clone(), - }) - .await; + agent.set_cancellation(cancellation.clone()); // Spawn the agent task. // @@ -1123,11 +1230,22 @@ impl SessionManager { // automatically on completion, error, panic, or task abort. let session_id_clone = session_id.to_string(); let events_clone = self.events.clone(); - let sleep_inhibitor = self.sleep_inhibitor.clone(); - sleep_inhibitor.agent_started(); - + let sleep_guard = self + .active_sessions + .get_mut(session_id) + .unwrap() + .sleep_guard + .take(); + let agent_lock = self + .active_sessions + .get_mut(session_id) + .unwrap() + .agent_lock + .take() + .ok_or_else(|| anyhow::anyhow!("Run reservation lost"))?; let task_handle = tokio::spawn(async move { let _agent_lock = agent_lock; // moved in — released on drop + debug!("Starting agent for session {}", session_id_clone); // Use catch_unwind to ensure cleanup runs even if the agent panics. @@ -1135,7 +1253,20 @@ impl SessionManager { // AgentRunning state because the cleanup code below is never reached. let result = { use futures::FutureExt; - let iteration_future = std::panic::AssertUnwindSafe(agent.run_single_iteration()); + let iteration_future = std::panic::AssertUnwindSafe(async { + cancellation.check()?; + agent.load_from_session_state(session_state).await?; + if let Some(scope) = tool_scope_override { + agent.set_tool_scope(scope); + agent.invalidate_system_message_cache(); + } + let _ = publisher + .send_event(UiEvent::UpdatePlan { + plan: agent.plan().clone(), + }) + .await; + agent.run_single_iteration().await + }); match iteration_future.catch_unwind().await { Ok(result) => result, Err(panic_payload) => { @@ -1155,6 +1286,26 @@ impl SessionManager { } }; + let result = if cancellation.is_cancelled() { + Ok(()) + } else { + result + }; + // Read usage while this run still owns the session. After publishing + // Idle a successor may already append messages and save its own usage. + let final_usage = if turn_recorder.is_some() { + let mut manager = manager_for_outcome.lock().await; + manager + .ensure_session_loaded(&session_id_clone) + .ok() + .and_then(|_| manager.get_session(&session_id_clone)) + .map(|instance| instance.calculate_total_usage()) + } else { + None + }; + // Release before publishing idle: a dispatch observing idle can acquire it. + drop(_agent_lock); + drop(sleep_guard); // Log the completion with detailed error information if failed match &result { Ok(()) => { @@ -1211,21 +1362,9 @@ impl SessionManager { // publisher is synchronous), so the record is complete here; the // token delta comes from the state the run persisted. if let Some(recorder) = turn_recorder { - let final_usage = { - let mut manager = manager_for_outcome.lock().await; - manager - .ensure_session_loaded(&session_id_clone) - .ok() - .and_then(|_| manager.get_session(&session_id_clone)) - .map(|instance| instance.calculate_total_usage()) - }; recorder.finish(result.as_ref().err().map(|e| format!("{e:#}")), final_usage); } - // Signal that this agent is no longer running so the system sleep - // inhibition can be released once all agents have finished. - sleep_inhibitor.agent_stopped(); - result }); @@ -1246,15 +1385,7 @@ impl SessionManager { pub fn delete_session(&mut self, session_id: &str) -> Result<()> { // Remove from active sessions if let Some(mut session_instance) = self.active_sessions.remove(session_id) { - let agent_is_running = !session_instance.get_activity_state().is_terminal(); session_instance.terminate_agent(); - // When aborting a task, the cleanup code inside the task (including - // agent_stopped) won't run, so we signal completion here instead. - // We check the activity state rather than task_handle.is_some() - // because the handle persists even after the task has completed. - if agent_is_running { - self.sleep_inhibitor.agent_stopped(); - } } // Clear active session if it was the deleted one @@ -1286,18 +1417,10 @@ impl SessionManager { /// Terminate a running agent for a session (e.g. on user cancel). /// - /// This is the proper way to abort an agent task from outside `SessionManager`, - /// because it also updates the sleep-inhibition reference count. Calling - /// `session.terminate_agent()` directly would leak a count. + /// Run-owned guards release the process/wake locks on every exit path. pub fn terminate_session_agent(&mut self, session_id: &str) { if let Some(session) = self.active_sessions.get_mut(session_id) { - // Check the activity state rather than task_handle.is_some() because - // the handle persists even after the task has completed naturally. - let agent_is_running = !session.get_activity_state().is_terminal(); session.terminate_agent(); - if agent_is_running { - self.sleep_inhibitor.agent_stopped(); - } } } @@ -1311,6 +1434,31 @@ impl SessionManager { self.active_sessions.get_mut(session_id) } + /// Heal a removed model without overwriting a selection made after the + /// run was reserved, including selections written by another process. + pub(crate) fn persist_model_fallback( + &mut self, + session_id: &str, + expected_model: &str, + fallback: &SessionModelConfig, + ) -> Result<()> { + let session = self.persistence.update_entry(session_id, |session| { + if session + .model_config + .as_ref() + .map(|model| model.model_name.as_str()) + == Some(expected_model) + { + session.model_config = Some(fallback.clone()); + } + Ok(()) + })?; + if let Some(instance) = self.active_sessions.get_mut(session_id) { + instance.session.model_config = session.model_config; + } + Ok(()) + } + /// Get the model config for a session, if any pub fn get_session_model_config(&self, session_id: &str) -> Result> { if let Some(instance) = self.active_sessions.get(session_id) { @@ -1572,7 +1720,8 @@ impl SessionManager { self.events.clone(), instance.pending_permission_requests.clone(), self.permission_timeout, - ), + ) + .with_cancellation(instance.cancellation.clone()), )) } diff --git a/crates/code_assistant_core/src/session/permissions.rs b/crates/code_assistant_core/src/session/permissions.rs index 62c47119..715786b7 100644 --- a/crates/code_assistant_core/src/session/permissions.rs +++ b/crates/code_assistant_core/src/session/permissions.rs @@ -171,6 +171,7 @@ pub struct SessionPermissionMediator { /// on a prompt nobody is there to answer. `None` waits indefinitely (the /// right default for an interactive frontend with a human present). timeout: Option, + cancellation: tools_core::RunCancellation, } impl SessionPermissionMediator { @@ -185,9 +186,15 @@ impl SessionPermissionMediator { events, pending, timeout, + cancellation: tools_core::RunCancellation::default(), } } + pub fn with_cancellation(mut self, cancellation: tools_core::RunCancellation) -> Self { + self.cancellation = cancellation; + self + } + fn next_request_id() -> String { static COUNTER: AtomicU64 = AtomicU64::new(1); format!("perm-{}", COUNTER.fetch_add(1, Ordering::Relaxed)) @@ -256,36 +263,63 @@ impl PermissionMediator for SessionPermissionMediator { ) -> Result { let data = Self::request_data(&request); let request_id = data.request_id.clone(); - let rx = self.pending.insert(data.clone()); - - self.events.publish_ui( - &self.session_id, - UiEvent::RequestToolPermission { request: data }, - ); + let rx = self.cancellation.if_active(|| { + let rx = self.pending.insert(data.clone()); + self.events.publish_ui( + &self.session_id, + UiEvent::RequestToolPermission { request: data }, + ); + rx + })?; + // Also clean up when an enclosing cancellation select drops this future. + let _guard = RequestGuard { + mediator: self, + request_id: request_id.clone(), + }; // A dropped responder (stop request, new agent run) counts as denial. // With a timeout, an unanswered prompt also fails closed: drop the // pending entry (so a late answer is a no-op) and deny, freeing the // lane's turn instead of blocking it forever. - let decision = match self.timeout { - Some(dur) => match tokio::time::timeout(dur, rx).await { - Ok(result) => result.unwrap_or(PermissionDecision::Denied), - Err(_elapsed) => { - self.pending - .resolve(&request_id, PermissionDecision::Denied); - PermissionDecision::Denied - } - }, - None => rx.await.unwrap_or(PermissionDecision::Denied), + let wait = async { + match self.timeout { + Some(dur) => match tokio::time::timeout(dur, rx).await { + Ok(result) => result.unwrap_or(PermissionDecision::Denied), + Err(_elapsed) => { + self.pending + .resolve(&request_id, PermissionDecision::Denied); + PermissionDecision::Denied + } + }, + None => rx.await.unwrap_or(PermissionDecision::Denied), + } }; + let decision = tokio::select! { + biased; + _ = self.cancellation.cancelled() => return Err(tools_core::Cancelled.into()), + decision = wait => decision, + }; + self.cancellation.check()?; + Ok(decision) + } +} - // Tell every view the request is settled so open prompts dismiss. - self.events.publish_ui( - &self.session_id, - UiEvent::ToolPermissionRequestResolved { request_id }, - ); +struct RequestGuard<'a> { + mediator: &'a SessionPermissionMediator, + request_id: String, +} - Ok(decision) +impl Drop for RequestGuard<'_> { + fn drop(&mut self) { + self.mediator + .pending + .resolve(&self.request_id, PermissionDecision::Denied); + self.mediator.events.publish_ui( + &self.mediator.session_id, + UiEvent::ToolPermissionRequestResolved { + request_id: self.request_id.clone(), + }, + ); } } diff --git a/crates/code_assistant_core/src/session/service.rs b/crates/code_assistant_core/src/session/service.rs index f5885838..b4b5c47a 100644 --- a/crates/code_assistant_core/src/session/service.rs +++ b/crates/code_assistant_core/src/session/service.rs @@ -6,12 +6,10 @@ //! caller gets *its* answer or *its* error — no correlation over shared //! channels. //! -//! Internally the service is an actor: methods enqueue a closure onto a -//! command channel and await a oneshot reply. A single worker future (see -//! [`SessionService::new`]) executes commands strictly in order on the -//! backend's tokio runtime, preserving the serialization of session -//! mutations and keeping the caller's executor (e.g. GPUI) decoupled from -//! tokio. Core→UI notifications keep flowing through [`UiEvent`] and are +//! Bounded dispatch serializes mutations per session on the backend tokio +//! runtime. Slow queries use independent tasks; a separate bounded control +//! mailbox handles stop and permission replies. The caller executor (e.g. +//! GPUI) stays decoupled from tokio. Core→UI notifications keep flowing through [`UiEvent`] and are //! not part of this API. use crate::config::{DefaultProjectManager, ProjectManager, save_project}; @@ -227,10 +225,21 @@ impl ServiceCtx { type BoxedCommandFuture = Pin + Send>>; type Command = Box BoxedCommandFuture + Send>; +struct Dispatch { + lane: String, + command: Command, + _permit: tokio::sync::OwnedSemaphorePermit, +} + +// At most 64 submitted commands, INCLUDING active and lane-queued work. +const COMMAND_CAPACITY: usize = 64; + /// Cloneable handle to the session command worker. See module docs. #[derive(Clone)] pub struct SessionService { - tx: async_channel::Sender, + tx: async_channel::Sender, + control_tx: async_channel::Sender, + admission: Arc, events: EventStream, } @@ -247,20 +256,77 @@ impl SessionService { runtime: Arc, events: EventStream, ) -> (Self, impl Future) { - let (tx, rx) = async_channel::unbounded::(); + let (tx, rx) = async_channel::bounded::(COMMAND_CAPACITY); + let (control_tx, control_rx) = async_channel::bounded::(16); + let admission = Arc::new(tokio::sync::Semaphore::new(COMMAND_CAPACITY)); let ctx = ServiceCtx { manager, runtime, events: events.clone(), }; let worker = async move { - debug!("Session service worker started"); - while let Ok(command) = rx.recv().await { - command(ctx.clone()).await; - } - debug!("Session service worker stopped"); + let control_ctx = ctx.clone(); + let control = async move { + while let Ok(command) = control_rx.recv().await { + command(control_ctx.clone()).await; + } + }; + let dispatch = async move { + use futures::FutureExt; + use std::collections::{HashMap, VecDeque}; + let mut lanes: HashMap> = HashMap::new(); + let mut tasks = tokio::task::JoinSet::new(); + let mut closed = false; + loop { + tokio::select! { + received = rx.recv(), if !closed => match received { + Ok(dispatch) => { + let lane = dispatch.lane.clone(); + if let Some(queue) = lanes.get_mut(&lane) { + queue.push_back(dispatch); + } else { + lanes.insert(lane.clone(), VecDeque::new()); + let ctx = ctx.clone(); + tasks.spawn(async move { + let _permit = dispatch._permit; + let _ = std::panic::AssertUnwindSafe((dispatch.command)(ctx)).catch_unwind().await; + lane + }); + } + } + Err(_) => closed = true, + }, + completed = tasks.join_next(), if !tasks.is_empty() => { + if let Some(Ok(lane)) = completed { + if let Some(dispatch) = lanes.get_mut(&lane).and_then(|queue| queue.pop_front()) { + let ctx = ctx.clone(); + tasks.spawn(async move { + let _permit = dispatch._permit; + let _ = std::panic::AssertUnwindSafe((dispatch.command)(ctx)).catch_unwind().await; + lane + }); + } else { + lanes.remove(&lane); + } + } + } + } + if closed && tasks.is_empty() { + break; + } + } + }; + tokio::join!(dispatch, control); }; - (Self { tx, events }, worker) + ( + Self { + tx, + control_tx, + admission, + events, + }, + worker, + ) } /// Subscribe to the core→UI broadcast stream. @@ -271,7 +337,7 @@ impl SessionService { /// Request that the running agent of a session stops at the next /// opportunity (streaming checkpoint). No-op if no agent is running. pub async fn request_stop(&self, session_id: String) -> Result<()> { - self.call(move |ctx| async move { + self.call_control(move |ctx| async move { let manager = ctx.manager.lock().await; let session = manager .get_session(&session_id) @@ -289,7 +355,7 @@ impl SessionService { /// command. Returns `Ok(())` regardless of whether a match was found /// (the process may have already finished). pub async fn interrupt_terminal(&self, session_id: String, tool_id: String) -> Result<()> { - self.call(move |ctx| async move { + self.call_control(move |ctx| async move { let manager = ctx.manager.lock().await; let session = manager .get_session(&session_id) @@ -305,15 +371,77 @@ impl SessionService { .await } - /// Enqueue a command and await its typed reply. + /// Enqueue a mutation on the global (non-session) lane. async fn call(&self, f: F) -> Result where F: FnOnce(ServiceCtx) -> Fut + Send + 'static, Fut: Future> + Send + 'static, T: Send + 'static, { + self.call_lane("global".into(), f).await + } + + async fn call_session(&self, session_id: String, f: F) -> Result + where + F: FnOnce(ServiceCtx) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + T: Send + 'static, + { + self.call_lane(format!("session:{session_id}"), f).await + } + + /// Read-only slow IO gets an independent, supervised task on the backend. + /// Even synchronous libgit/filesystem work cannot stall a single-thread runtime. + async fn call_io(&self, f: F) -> Result + where + F: FnOnce(ServiceCtx) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + T: Send + 'static, + { + static NEXT_IO: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); + let id = NEXT_IO.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.call_lane(format!("io:{id}"), move |ctx| async move { + let runtime = tokio::runtime::Handle::current(); + tokio::task::spawn_blocking(move || runtime.block_on(f(ctx))).await? + }) + .await + } + + async fn call_lane(&self, lane: String, f: F) -> Result + where + F: FnOnce(ServiceCtx) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + T: Send + 'static, + { + let permit = self.admission.clone().acquire_owned().await?; let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); self.tx + .send(Dispatch { + lane, + _permit: permit, + command: Box::new(move |ctx| { + Box::pin(async move { + let _ = reply_tx.send(f(ctx).await); + }) + }), + }) + .await + .map_err(|_| anyhow!("session service is not running"))?; + reply_rx + .await + .map_err(|_| anyhow!("session service dropped the request"))? + } + + // Separate bounded mailbox: control never waits behind lane or IO capacity. + // These closures still read/mutate the authoritative manager under its lock. + async fn call_control(&self, f: F) -> Result + where + F: FnOnce(ServiceCtx) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + T: Send + 'static, + { + let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); + self.control_tx .send(Box::new(move |ctx| { Box::pin(async move { let _ = reply_tx.send(f(ctx).await); @@ -404,7 +532,7 @@ impl SessionService { /// system-initiated turn — such as an isolated supervised-delegation child — /// is cancelled from the outside. pub async fn terminate_agent(&self, session_id: String) -> Result<()> { - self.call(move |ctx| async move { + self.call_control(move |ctx| async move { let mut manager = ctx.manager.lock().await; manager.terminate_session_agent(&session_id); Ok(()) @@ -420,7 +548,7 @@ impl SessionService { session_id: String, edit_until_node_id: Option, ) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let snapshot = { let mut manager = ctx.manager.lock().await; manager @@ -433,7 +561,7 @@ impl SessionService { } pub async fn delete_session(&self, session_id: String) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let mut manager = ctx.manager.lock().await; manager.delete_session(&session_id) }) @@ -453,7 +581,7 @@ impl SessionService { /// as [`UiEvent`]s; falls back to a full reload if that fails. pub async fn refresh_session(&self, session_id: String) -> Result<()> { let refreshed = self - .call({ + .call_session(session_id.clone(), { let session_id = session_id.clone(); move |ctx| async move { let ui_events = { @@ -478,7 +606,7 @@ impl SessionService { /// Clear the Errored state on a session (user dismissed the error banner). pub async fn clear_session_error(&self, session_id: String) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { { let mut manager = ctx.manager.lock().await; if let Some(session) = manager.get_session_mut(&session_id) { @@ -506,7 +634,7 @@ impl SessionService { /// Clear the conversation context (messages) for a session. The session /// itself is kept alive; only the message history is wiped. pub async fn clear_context(&self, session_id: String) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { { let mut manager = ctx.manager.lock().await; if let Some(session) = manager.get_session_mut(&session_id) { @@ -551,7 +679,7 @@ impl SessionService { attachments: Vec, branch_parent_id: Option, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { send_user_message_impl( &ctx, &session_id, @@ -577,7 +705,7 @@ impl SessionService { attachments: Vec, tool_scope: crate::tools::core::ToolScope, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { send_user_message_impl( &ctx, &session_id, @@ -604,7 +732,7 @@ impl SessionService { message: String, attachments: Vec, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { send_or_queue_user_message_impl(&ctx, &session_id, &message, &attachments).await }) .await @@ -621,7 +749,7 @@ impl SessionService { message: String, attachments: Vec, ) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { try_send_user_message_if_idle_impl(&ctx, &session_id, &message, &attachments).await }) .await @@ -644,9 +772,8 @@ impl SessionService { session_id: String, request: crate::session::TurnRequest, ) -> Result { - let service = self.clone(); - self.call(move |ctx| async move { - start_turn_if_idle_impl(&ctx, service, session_id, request).await + self.call_session(session_id.clone(), move |ctx| async move { + start_turn_if_idle_impl(&ctx, session_id, request).await }) .await } @@ -663,7 +790,7 @@ impl SessionService { /// this rather than a mirror, or it will act on a session the atomic send /// then refuses. Loads the session on demand, like the send paths. pub async fn is_session_busy(&self, session_id: String) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let mut manager = ctx.manager.lock().await; manager.ensure_session_loaded(&session_id)?; let instance = manager @@ -682,7 +809,7 @@ impl SessionService { message: String, attachments: Vec, ) -> Result> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let content_blocks = content_blocks_from(&message, &attachments); let mut manager = ctx.manager.lock().await; manager.queue_structured_user_message(&session_id, content_blocks)?; @@ -694,7 +821,7 @@ impl SessionService { /// Take the pending message out of the queue for editing. Returns its /// text, or `None` if nothing was queued. pub async fn take_pending_message(&self, session_id: String) -> Result> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let mut manager = ctx.manager.lock().await; manager.request_pending_message_for_edit(&session_id) }) @@ -704,8 +831,10 @@ impl SessionService { /// Resume a session that ended in a state where the agent should run /// against the existing message history (no new user message is added). pub async fn resume_session(&self, session_id: String) -> Result<()> { - self.call(move |ctx| async move { resume_session_impl(&ctx, &session_id).await }) - .await + self.call_session(session_id.clone(), move |ctx| async move { + resume_session_impl(&ctx, &session_id).await + }) + .await } /// Deliver a fired wakeup (see [`crate::session::wakeup`]): inject @@ -714,14 +843,16 @@ impl SessionService { /// when an agent is currently running. A session that no longer exists /// swallows the wakeup silently. pub async fn inject_wakeup(&self, session_id: String, message: String) -> Result<()> { - self.call(move |ctx| async move { inject_wakeup_impl(&ctx, &session_id, &message).await }) - .await + self.call_session(session_id.clone(), move |ctx| async move { + inject_wakeup_impl(&ctx, &session_id, &message).await + }) + .await } /// Cancel a running sub-agent by its tool id. Returns `true` if a /// sub-agent was actually cancelled, `false` if it had already finished. pub async fn cancel_sub_agent(&self, session_id: String, tool_id: String) -> Result { - self.call(move |ctx| async move { + self.call_control(move |ctx| async move { let manager = ctx.manager.lock().await; manager.cancel_sub_agent(&session_id, &tool_id) }) @@ -735,7 +866,7 @@ impl SessionService { /// List the skills available to a session (across project / user / /// system scopes), for the input-area skill picker. pub async fn list_skills(&self, session_id: String) -> Result> { - self.call(move |ctx| async move { + self.call_io(move |ctx| async move { let project_name = { let manager = ctx.manager.lock().await; manager @@ -769,7 +900,7 @@ impl SessionService { scope: String, name: String, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let config = SkillsConfig::load(); let pm = (ctx.runtime.project_manager_factory)(); let payload = load_skill_payload(pm.as_ref(), &scope, &name, &config) @@ -804,7 +935,7 @@ impl SessionService { session_id: String, model_name: String, ) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let config_system = ConfigurationSystem::load().context("Failed to load model configuration")?; if config_system.get_model(&model_name).is_none() { @@ -852,7 +983,7 @@ impl SessionService { session_id: String, policy: SandboxPolicy, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { { let mut manager = ctx.manager.lock().await; manager.set_session_sandbox_policy(&session_id, policy.clone())?; @@ -871,7 +1002,7 @@ impl SessionService { session_id: String, tier: tools_core::PermissionTier, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { { let mut manager = ctx.manager.lock().await; manager.set_session_permission_tier(&session_id, tier)?; @@ -892,7 +1023,7 @@ impl SessionService { session_id: String, disabled: Vec, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let servers = { let mut manager = ctx.manager.lock().await; manager.set_session_disabled_mcp_servers(&session_id, disabled)? @@ -912,7 +1043,7 @@ impl SessionService { request_id: String, decision: tools_core::PermissionDecision, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_control(move |ctx| async move { let manager = ctx.manager.lock().await; manager.resolve_permission_request(&session_id, &request_id, decision)?; Ok(()) @@ -931,7 +1062,7 @@ impl SessionService { session_id: String, node_id: NodeId, ) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let manager = ctx.manager.lock().await; let session_instance = manager .get_session(&session_id) @@ -1003,7 +1134,7 @@ impl SessionService { session_id: String, new_node_id: NodeId, ) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let mut manager = ctx.manager.lock().await; let session_instance = manager .get_session_mut(&session_id) @@ -1034,7 +1165,7 @@ impl SessionService { /// Abort a message edit and return the full transcript of the active /// path. pub async fn cancel_message_edit(&self, session_id: String) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let manager = ctx.manager.lock().await; let session_instance = manager .get_session(&session_id) @@ -1049,7 +1180,7 @@ impl SessionService { // ======================================================================== pub async fn list_branches_and_worktrees(&self, session_id: String) -> Result { - self.call(move |ctx| async move { + self.call_io(move |ctx| async move { let project_root = { let manager = ctx.manager.lock().await; session_project_root(&manager, &session_id)? @@ -1092,7 +1223,7 @@ impl SessionService { /// Resolves the session's on-disk directory with `effective_project_path` /// (worktree-aware). An empty result means "not a git project". pub async fn list_review_repos(&self, session_id: String) -> Result> { - self.call(move |ctx| async move { + self.call_io(move |ctx| async move { let project_root = { let manager = ctx.manager.lock().await; session_effective_path(&manager, &session_id)? @@ -1113,7 +1244,7 @@ impl SessionService { mode: ReviewMode, base_override: Option, ) -> Result { - self.call(move |_ctx| async move { + self.call_io(move |_ctx| async move { let repo = git::GitRepository::open(&repo_root).context("Failed to open git repository")?; let current_branch = repo.current_branch(); @@ -1168,7 +1299,7 @@ impl SessionService { base: Option, file: git::ChangedFile, ) -> Result { - self.call(move |_ctx| async move { + self.call_io(move |_ctx| async move { let repo = git::GitRepository::open(&repo_root).context("Failed to open git repository")?; @@ -1196,7 +1327,7 @@ impl SessionService { worktree_path: Option, branch: Option, ) -> Result<()> { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let mut manager = ctx.manager.lock().await; manager.set_session_worktree(&session_id, worktree_path, branch) }) @@ -1211,7 +1342,7 @@ impl SessionService { branch_name: String, base_branch: Option, ) -> Result { - self.call(move |ctx| async move { + self.call_session(session_id.clone(), move |ctx| async move { let project_root = { let manager = ctx.manager.lock().await; session_project_root(&manager, &session_id)? @@ -1560,16 +1691,43 @@ async fn send_user_message_impl( // Headless dispatch (channel adapters, schedulers) reaches sessions // no frontend has opened since the restart — load on demand. manager.ensure_session_loaded(session_id)?; - let node_id = manager - .add_user_message(session_id, content_blocks, branch_parent_id) - .context("Failed to add user message")?; - // If we created a branch, get branch info updates for all siblings. - let updates = if branch_parent_id.is_some() { - manager.get_sibling_branch_infos(session_id, node_id) - } else { - Vec::new() - }; - (node_id, updates) + anyhow::ensure!( + manager + .get_session(session_id) + .unwrap() + .get_activity_state() + .is_terminal(), + "Session is already running" + ); + let cancellation = turn_recorder + .as_ref() + .map(|recorder| recorder.cancellation.clone()) + .unwrap_or_default(); + // Claim the cross-process writer before touching the conversation. + manager.reserve_agent_run(session_id, cancellation.clone())?; + let prepared: Result<_> = (|| { + let node_id = manager + .add_user_message(session_id, content_blocks, branch_parent_id) + .context("Failed to add user message")?; + let updates = if branch_parent_id.is_some() { + manager.get_sibling_branch_infos(session_id, node_id) + } else { + Vec::new() + }; + schedule_agent_impl( + ctx, + &mut manager, + session_id, + tool_scope_override, + turn_recorder, + cancellation.clone(), + )?; + Ok((node_id, updates)) + })(); + if let Err(error) = &prepared { + manager.finish_failed_setup(session_id, &cancellation, format!("{error:#}")); + } + prepared? }; // Now display the user message with the correct node_id. @@ -1594,7 +1752,7 @@ async fn send_user_message_impl( ); } - start_agent_impl(ctx, session_id, tool_scope_override, turn_recorder).await + Ok(()) } /// Shared by [`SessionService::send_or_queue_user_message`] and the wakeup @@ -1658,7 +1816,6 @@ async fn try_send_user_message_if_idle_impl( async fn start_turn_if_idle_impl( ctx: &ServiceCtx, - service: SessionService, session_id: String, request: crate::session::TurnRequest, ) -> Result { @@ -1694,7 +1851,7 @@ async fn start_turn_if_idle_impl( Ok(TurnDispatch::Started(TurnHandle::new( session_id, parts.turn_id, - service, + parts.cancellation, parts.outcome, ))) } @@ -1765,139 +1922,168 @@ async fn start_agent_impl( tool_scope_override: Option, turn_recorder: Option>, ) -> Result<()> { - let (session_config, default_model_name) = { - let manager = ctx.manager.lock().await; - ( - manager.get_session_model_config(session_id).unwrap_or(None), - manager.default_model_name().to_string(), - ) - }; - let Some(mut session_config) = session_config else { - bail!( - "Session has no model configuration. Please ensure all sessions are created with a model." - ); - }; - - // Validation is fail-soft: without a loadable configuration the client - // construction below reports the error. - if let Ok(config_system) = ConfigurationSystem::load() { - session_config = runnable_model_config(session_config, &default_model_name, |model| { - config_system.get_model(model).is_some() - }); + let mut manager = ctx.manager.lock().await; + let cancellation = turn_recorder + .as_ref() + .map(|recorder| recorder.cancellation.clone()) + .unwrap_or_default(); + manager.reserve_agent_run(session_id, cancellation.clone())?; + let result = schedule_agent_impl( + ctx, + &mut manager, + session_id, + tool_scope_override, + turn_recorder, + cancellation.clone(), + ); + if let Err(error) = &result { + manager.finish_failed_setup(session_id, &cancellation, format!("{error:#}")); } + result +} - let llm_client = match &ctx.runtime.llm_client_factory { - Some(factory) => factory(&session_config.model_name) - .context("Failed to create LLM client from injected factory")?, - None => create_llm_client_from_model( - &session_config.model_name, - ctx.runtime.playback_path.clone(), - ctx.runtime.fast_playback, - ctx.runtime.record_path.clone(), - ) - .await - .context("Failed to create LLM client")?, - }; - - let project_manager = (ctx.runtime.project_manager_factory)(); - let command_executor = (ctx.runtime.command_executor_factory)(session_id); - - // Resolve local-`.mcp.json` trust, then start the run. The trust prompt is - // answered via `respond_permission` — another command on this service's - // single-threaded worker — so when a prompt is actually needed we must not - // block the worker awaiting it (that would starve the response and hang the - // prompt forever). In that case the whole resolve+start runs on a detached - // task; the common no-prompt case stays inline and synchronous. - let (project_dir, disabled_mcp_servers, permission_handler) = { - let manager = ctx.manager.lock().await; - let (project_dir, disabled) = manager - .get_session(session_id) - .map(|instance| { - ( - instance.session.config.effective_project_path().cloned(), - instance.session.config.disabled_mcp_servers.clone(), - ) - }) - .unwrap_or((None, Vec::new())); - ( - project_dir, - disabled, - manager.permission_mediator(session_id).ok(), - ) - }; - let needs_prompt = project_dir - .as_deref() - .is_some_and(crate::tools::mcp_trust::needs_prompt); - +fn schedule_agent_impl( + ctx: &ServiceCtx, + manager: &mut SessionManager, + session_id: &str, + tool_scope_override: Option, + turn_recorder: Option>, + cancellation: tools_core::RunCancellation, +) -> Result<()> { + // Caller owns the reservation and handles synchronous setup failures. + let loader = manager.registry_loader(); + let session_config = manager.get_session_model_config(session_id)?; + let default_model_name = manager.default_model_name().to_string(); + let instance = manager.get_session(session_id).unwrap(); + let run_session_config = instance.session.config.clone(); + let project_dir = run_session_config.effective_project_path().cloned(); + let disabled = run_session_config.disabled_mcp_servers.clone(); + let permission_handler = manager.permission_mediator(session_id)?; + // The manager owns setup_task; the task must not own the manager back. + let owner = Arc::downgrade(&ctx.manager); + let runtime = ctx.runtime.clone(); let events = ctx.events.clone(); - let ctx = ctx.clone(); let session_id = session_id.to_string(); - let err_session_id = session_id.clone(); - let run = async move { - let include_local_mcp = SessionManager::resolve_local_mcp_trust( - project_dir.as_deref(), - permission_handler.as_deref(), - ) - .await; - // Trust for this project is now resolved; refresh the input-bar MCP - // list so a project's `.mcp.json` servers appear once trusted (the - // session snapshot was computed before the trust prompt). - ctx.notify_session( - &session_id, - UiEvent::UpdateMcpServers { - servers: crate::tools::mcp::session_mcp_servers( - project_dir.as_deref(), - &disabled_mcp_servers, - ), - }, - ); - let registry_request = crate::session::manager::RegistryRequest { - project_dir, - include_local_mcp, - }; - let mut manager = ctx.manager.lock().await; - manager - .set_session_model_config(&session_id, Some(session_config)) - .context("Failed to persist model config")?; - manager - .start_agent_for_session( - &session_id, - llm_client, - project_manager, - command_executor, - permission_handler, - tool_scope_override, - turn_recorder, - registry_request, + let task_session = session_id.clone(); + let task = tokio::spawn(async move { + use futures::FutureExt; + let prepare = async { + let mut session_config = + session_config.ok_or_else(|| anyhow!("Session has no model configuration"))?; + let original_model = session_config.model_name.clone(); + if let Ok(config) = ConfigurationSystem::load() { + session_config = + runnable_model_config(session_config, &default_model_name, |model| { + config.get_model(model).is_some() + }); + } + let llm_client = match &runtime.llm_client_factory { + Some(factory) => { + let factory = factory.clone(); + let model = session_config.model_name.clone(); + // A synchronous injected constructor must not stall the backend. + // Dropping this wait cannot kill blocking user code, but its result + // has no authority to install a run after cancellation. + tokio::task::spawn_blocking(move || factory(&model)).await?? + } + None => { + create_llm_client_from_model( + &session_config.model_name, + runtime.playback_path.clone(), + runtime.fast_playback, + runtime.record_path.clone(), + ) + .await? + } + }; + cancellation.check()?; + let include_local_mcp = SessionManager::resolve_local_mcp_trust( + project_dir.as_deref(), + Some(permission_handler.as_ref()), ) - .await - .context("Failed to start agent")?; - debug!("Agent started for session {}", session_id); - Ok(()) - }; - - if needs_prompt { - // Detach so the worker is free to process the prompt's response. - tokio::spawn(async move { - if let Err(error) = run.await { - let message = format!("Failed to start agent: {error:#}"); - error!("{message}"); - events.publish_ui(&err_session_id, UiEvent::DisplayError { message }); + .await; + cancellation.check()?; + let registry = loader(crate::session::manager::RegistryRequest { + project_dir: project_dir.clone(), + include_local_mcp, + }) + .await; + cancellation.check()?; + let project_manager = (runtime.project_manager_factory)(); + let command_executor = (runtime.command_executor_factory)(&task_session); + let owner = owner + .upgrade() + .ok_or_else(|| anyhow!("Session service shut down"))?; + let mut manager = owner.lock().await; + // Stop/delete may have won while preparation was in flight. + cancellation.check()?; + // Persist a fallback only while the selection still matches the + // one captured at reservation. A newer selection belongs to the next run. + if session_config.model_name != original_model { + manager.persist_model_fallback(&task_session, &original_model, &session_config)?; } - }); - Ok(()) - } else { - run.await - } + events.publish_ui( + &task_session, + UiEvent::UpdateMcpServers { + servers: crate::tools::mcp::session_mcp_servers( + project_dir.as_deref(), + &disabled, + ), + }, + ); + manager + .start_reserved_agent_for_session( + &task_session, + llm_client, + project_manager, + command_executor, + Some(permission_handler.clone()), + tool_scope_override, + turn_recorder.clone(), + registry, + cancellation.clone(), + crate::session::manager::RunConfig { + session: run_session_config, + model: Some(session_config), + }, + ) + .await + }; + let result = tokio::select! { + biased; + _ = cancellation.cancelled() => Err(tools_core::Cancelled.into()), + result = std::panic::AssertUnwindSafe(prepare).catch_unwind() => { + result.unwrap_or_else(|_| Err(anyhow!("Run preparation panicked"))) + } + }; + if let Err(error) = result { + let message = format!("Failed to start agent: {error:#}"); + if let Some(owner) = owner.upgrade() { + owner.lock().await.finish_failed_setup( + &task_session, + &cancellation, + message.clone(), + ); + } + if let Some(recorder) = turn_recorder { + recorder.finish(Some(message), None); + } + } + }); + manager.get_session_mut(&session_id).unwrap().setup_task = Some(task); + Ok(()) } +#[cfg(test)] +mod recovery_tests; + #[cfg(test)] mod tests { use super::*; use crate::persistence::FileSessionPersistence; use crate::session::SessionConfig; - fn test_service_with_manager( + pub(super) fn test_service_with_manager( root: &std::path::Path, ) -> (SessionService, Arc>) { let events = EventStream::new(); @@ -1961,7 +2147,7 @@ mod tests { /// Service whose agent runs use the injected LLM factory instead of the /// configured providers. - fn test_service_with_llm( + pub(super) fn test_service_with_llm( root: &std::path::Path, factory: LlmClientFactory, ) -> (SessionService, Arc>) { @@ -1989,6 +2175,168 @@ mod tests { (service, manager) } + #[tokio::test(flavor = "multi_thread")] + async fn checkpoint2_setup_is_reserved_and_control_stays_responsive() { + use crate::session::{TurnDispatch, TurnRequest, TurnStatus}; + use std::time::Duration; + let tmp = tempfile::tempdir().unwrap(); + let (service, manager) = test_service_with_llm( + tmp.path(), + Arc::new(|_| { + Ok(Box::new(StreamingScriptedProvider { + text: "done".into(), + })) + }), + ); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + manager.lock().await.set_tool_registry_provider({ + let entered = entered.clone(); + let release = release.clone(); + Arc::new(move |_| { + let entered = entered.clone(); + let release = release.clone(); + Box::pin(async move { + entered.notify_one(); + release.notified().await; + crate::tools::test_registry() + }) + }) + }); + let id = service.create_session(None, None).await.unwrap(); + let other = service.create_session(None, None).await.unwrap(); + let first = tokio::spawn({ + let service = service.clone(); + let id = id.clone(); + async move { + service + .start_turn_if_idle(id, TurnRequest::text("first")) + .await + } + }); + entered.notified().await; + // Every check is bounded; release the fixture even on RED so no task hangs. + let second = tokio::time::timeout( + Duration::from_millis(500), + service.start_turn_if_idle(id.clone(), TurnRequest::text("second")), + ) + .await; + let unrelated = tokio::time::timeout( + Duration::from_millis(500), + service.queue_user_message(other, "unrelated".into(), vec![]), + ) + .await; + let stopped = + tokio::time::timeout(Duration::from_millis(500), service.request_stop(id.clone())) + .await; + release.notify_waiters(); + assert!( + matches!(second, Ok(Ok(TurnDispatch::Busy))), + "pending setup must reserve the run, without blocking dispatch" + ); + assert!(unrelated.is_ok(), "setup blocked another session"); + assert!(stopped.is_ok(), "setup blocked stop"); + let TurnDispatch::Started(handle) = first.await.unwrap().unwrap() else { + panic!("first busy") + }; + let outcome = tokio::time::timeout(Duration::from_secs(2), handle.wait()) + .await + .unwrap() + .unwrap(); + assert_eq!(outcome.status, TurnStatus::Cancelled); + assert!(!service.is_session_busy(id).await.unwrap()); + } + + struct WaitingProvider; + #[async_trait::async_trait] + impl llm::LLMProvider for WaitingProvider { + async fn send_message( + &mut self, + _: llm::LLMRequest, + _: Option<&llm::StreamingCallback>, + ) -> Result { + std::future::pending().await + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn checkpoint2_stop_interrupts_a_provider_without_chunks() { + use crate::session::{TurnDispatch, TurnRequest, TurnStatus}; + let tmp = tempfile::tempdir().unwrap(); + let (service, _) = + test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(WaitingProvider)))); + let id = service.create_session(None, None).await.unwrap(); + let TurnDispatch::Started(handle) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("wait")) + .await + .unwrap() + else { + panic!("busy") + }; + service.request_stop(id.clone()).await.unwrap(); + let outcome = + tokio::time::timeout(std::time::Duration::from_millis(500), handle.wait()).await; + service.terminate_agent(id).await.unwrap(); + assert_eq!( + outcome + .expect("stop must wake a pending provider") + .unwrap() + .status, + TurnStatus::Cancelled + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn checkpoint2_old_turn_handle_does_not_stop_a_new_run() { + use crate::session::{TurnDispatch, TurnRequest}; + let tmp = tempfile::tempdir().unwrap(); + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let (service, manager) = test_service_with_llm( + tmp.path(), + Arc::new(move |_| { + if calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + Ok(Box::new(StreamingScriptedProvider { + text: "first done".into(), + })) + } else { + Ok(Box::new(WaitingProvider)) + } + }), + ); + let id = service.create_session(None, None).await.unwrap(); + let TurnDispatch::Started(old) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("first")) + .await + .unwrap() + else { + panic!("busy") + }; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while service.is_session_busy(id.clone()).await.unwrap() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let TurnDispatch::Started(_new) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("second")) + .await + .unwrap() + else { + panic!("busy") + }; + old.cancel().await.unwrap(); + let stopped = manager + .lock() + .await + .get_session(&id) + .unwrap() + .stop_requested + .load(std::sync::atomic::Ordering::Relaxed); + service.terminate_agent(id).await.unwrap(); + assert!(!stopped, "an old turn handle cancelled its successor"); + } + #[tokio::test(flavor = "multi_thread")] async fn start_turn_if_idle_resolves_the_exact_outcome() { let tmp = tempfile::tempdir().unwrap(); diff --git a/crates/code_assistant_core/src/session/service/recovery_tests.rs b/crates/code_assistant_core/src/session/service/recovery_tests.rs new file mode 100644 index 00000000..2d36d431 --- /dev/null +++ b/crates/code_assistant_core/src/session/service/recovery_tests.rs @@ -0,0 +1,316 @@ +use super::tests::{test_service_with_llm, test_service_with_manager}; +use super::*; +use crate::session::{TurnDispatch, TurnRequest}; +use std::time::Duration; + +struct Done; +#[async_trait::async_trait] +impl llm::LLMProvider for Done { + async fn send_message( + &mut self, + _: llm::LLMRequest, + _: Option<&llm::StreamingCallback>, + ) -> Result { + Ok(llm::LLMResponse { + content: vec![llm::ContentBlock::new_text("done")], + usage: llm::Usage::zero(), + rate_limit_info: None, + }) + } +} + +fn blocked_registry( + entered: Arc, + release: Arc, +) -> crate::session::manager::ToolRegistryProvider { + Arc::new(move |_| { + let entered = entered.clone(); + let release = release.clone(); + Box::pin(async move { + entered.notify_one(); + release.notified().await; + crate::tools::test_registry() + }) + }) +} + +#[tokio::test(flavor = "multi_thread")] +async fn recovery_pending_setup_does_not_keep_its_owner_alive() { + let tmp = tempfile::tempdir().unwrap(); + let (service, manager) = test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(Done)))); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + manager + .lock() + .await + .set_tool_registry_provider(blocked_registry(entered.clone(), release)); + let id = service.create_session(None, None).await.unwrap(); + let TurnDispatch::Started(handle) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("task")) + .await + .unwrap() + else { + panic!("fresh session busy") + }; + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + let inhibitor = manager.lock().await.sleep_inhibitor(); + assert_eq!(inhibitor.running_count(), 1); + let weak = Arc::downgrade(&manager); + drop(manager); + drop(service); + let released = tokio::time::timeout(Duration::from_millis(250), async { + while weak.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await; + // Cleanup even on RED: do not leave a pending task or file lock behind. + if let Some(manager) = weak.upgrade() { + manager.lock().await.terminate_session_agent(&id); + } + assert!( + released.is_ok(), + "setup retained its owner in a reference cycle" + ); + tokio::time::timeout(Duration::from_secs(2), handle.wait()) + .await + .unwrap() + .unwrap(); + assert_eq!( + inhibitor.running_count(), + 0, + "shutdown leaked the run's wake lock" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn recovery_setup_preserves_a_newer_model_selection() { + let tmp = tempfile::tempdir().unwrap(); + let (service, manager) = test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(Done)))); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + manager + .lock() + .await + .set_tool_registry_provider(blocked_registry(entered.clone(), release.clone())); + let id = service.create_session(None, None).await.unwrap(); + let TurnDispatch::Started(handle) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("task")) + .await + .unwrap() + else { + panic!("fresh session busy") + }; + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + // Simulate a selection arriving during preparation without depending on + // the developer's models.json or any real provider configuration. + let next_model = SessionModelConfig::new("selected-during-setup".into()); + { + let mut manager = manager.lock().await; + manager.get_session_mut(&id).unwrap().session.model_config = Some(next_model.clone()); + let mut store = + crate::persistence::FileSessionPersistence::new_with_root_dir(tmp.path().to_path_buf()); + store + .update_entry(&id, |session| { + session.model_config = Some(next_model); + Ok(()) + }) + .unwrap(); + } + release.notify_one(); + tokio::time::timeout(Duration::from_secs(2), handle.wait()) + .await + .unwrap() + .unwrap(); + let snapshot = service.load_session(id, None).await.unwrap(); + assert_eq!(snapshot.current_model, "selected-during-setup"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn recovery_external_run_rejection_does_not_append_a_message() { + let tmp = tempfile::tempdir().unwrap(); + let (service, _) = test_service_with_manager(tmp.path()); + let id = service.create_session(None, None).await.unwrap(); + let lock = crate::utils::file_utils::try_acquire_agent_lock(&tmp.path().join("sessions"), &id) + .unwrap() + .unwrap(); + let sent = service + .send_user_message(id.clone(), "must not land".into(), vec![], None) + .await; + drop(lock); + assert!(sent.is_err()); + let snapshot = service.load_session(id, None).await.unwrap(); + assert!( + snapshot.messages.is_empty(), + "a refused turn modified the conversation" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn recovery_stop_during_permission_does_not_open_the_next_prompt() { + struct TwoCalls; + #[async_trait::async_trait] + impl llm::LLMProvider for TwoCalls { + async fn send_message( + &mut self, + _: llm::LLMRequest, + _: Option<&llm::StreamingCallback>, + ) -> Result { + Ok(llm::LLMResponse { + content: ["first", "second"] + .into_iter() + .map(|id| { + llm::ContentBlock::new_tool_use( + id, + "read_files", + serde_json::json!({"project":"test", "paths":["a.rs"]}), + ) + }) + .collect(), + usage: llm::Usage::zero(), + rate_limit_info: None, + }) + } + } + let tmp = tempfile::tempdir().unwrap(); + let (service, manager) = + test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(TwoCalls)))); + let id = service.create_session(None, None).await.unwrap(); + service + .change_permission_tier(id.clone(), tools_core::PermissionTier::AllTools) + .await + .unwrap(); + let mut events = service.subscribe(); + let TurnDispatch::Started(handle) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("read")) + .await + .unwrap() + else { + panic!("busy") + }; + let first = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let crate::session::EventPayload::Ui(UiEvent::RequestToolPermission { request }) = + events.recv().await.unwrap().payload + { + break request; + } + } + }) + .await + .unwrap(); + assert_eq!(first.tool_id.as_deref(), Some("first")); + service.request_stop(id.clone()).await.unwrap(); + let outcome = tokio::time::timeout(Duration::from_secs(2), handle.wait()) + .await + .unwrap() + .unwrap(); + assert_eq!(outcome.status, crate::session::TurnStatus::Cancelled); + assert!( + manager + .lock() + .await + .get_session(&id) + .unwrap() + .pending_permission_requests + .snapshot() + .is_empty() + ); + // An event after completion fences every earlier prompt publication. + service.clear_session_error(id).await.unwrap(); + loop { + match events.recv().await.unwrap().payload { + crate::session::EventPayload::Ui(UiEvent::RequestToolPermission { .. }) => { + panic!("a second prompt opened after stop") + } + crate::session::EventPayload::Ui(UiEvent::UpdateSessionActivityState { + activity_state, + .. + }) if activity_state.is_terminal() => break, + _ => {} + } + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn recovery_stop_wakes_an_already_waiting_silent_provider() { + struct Silent(Arc); + #[async_trait::async_trait] + impl llm::LLMProvider for Silent { + async fn send_message( + &mut self, + _: llm::LLMRequest, + _: Option<&llm::StreamingCallback>, + ) -> Result { + self.0.notify_one(); + std::future::pending().await + } + } + let tmp = tempfile::tempdir().unwrap(); + let entered = Arc::new(tokio::sync::Notify::new()); + let (service, _) = test_service_with_llm(tmp.path(), { + let entered = entered.clone(); + Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) + }); + let id = service.create_session(None, None).await.unwrap(); + let TurnDispatch::Started(handle) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("wait")) + .await + .unwrap() + else { + panic!("busy") + }; + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + handle.cancel().await.unwrap(); + let outcome = tokio::time::timeout(Duration::from_secs(2), handle.wait()) + .await + .unwrap() + .unwrap(); + assert_eq!(outcome.status, crate::session::TurnStatus::Cancelled); + assert!(!service.is_session_busy(id).await.unwrap()); +} + +#[tokio::test] +async fn recovery_slow_io_does_not_block_session_control() { + let tmp = tempfile::tempdir().unwrap(); + let (service, manager) = test_service_with_manager(tmp.path()); + let id = service.create_session(None, None).await.unwrap(); + let entered = Arc::new(tokio::sync::Notify::new()); + let (release, released) = std::sync::mpsc::channel(); + let task = tokio::spawn({ + let service = service.clone(); + let entered = entered.clone(); + async move { + service + .call_io(move |_| async move { + entered.notify_one(); + released.recv_timeout(Duration::from_secs(3)).unwrap(); + Ok(()) + }) + .await + } + }); + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + let stopped = + tokio::time::timeout(Duration::from_millis(250), service.request_stop(id.clone())).await; + release.send(()).unwrap(); + task.await.unwrap().unwrap(); + stopped.expect("slow IO blocked stop").unwrap(); + assert!( + manager + .lock() + .await + .get_session(&id) + .unwrap() + .cancellation + .is_cancelled() + ); +} diff --git a/crates/code_assistant_core/src/session/sleep_inhibitor.rs b/crates/code_assistant_core/src/session/sleep_inhibitor.rs index 4086f915..c407e871 100644 --- a/crates/code_assistant_core/src/session/sleep_inhibitor.rs +++ b/crates/code_assistant_core/src/session/sleep_inhibitor.rs @@ -1,5 +1,5 @@ -use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use tracing::{debug, warn}; /// Prevents the system from going to idle sleep while any agent is running. @@ -17,7 +17,24 @@ pub struct SleepInhibitor { wake_lock: Mutex>, } +/// One run's wake-lock ownership. Released on every exit, including task abort. +pub(crate) struct AgentSleepGuard(Arc); + +impl Drop for AgentSleepGuard { + fn drop(&mut self) { + self.0.agent_stopped(); + } +} + impl SleepInhibitor { + pub(crate) fn agent_guard(self: &Arc) -> AgentSleepGuard { + self.agent_started(); + AgentSleepGuard(self.clone()) + } + #[cfg(test)] + pub(crate) fn running_count(&self) -> usize { + self.running_count.load(Ordering::SeqCst) + } /// Called when an agent starts running. Acquires the system wake lock /// if this is the first active agent. pub fn agent_started(&self) { diff --git a/crates/code_assistant_core/src/session/turn.rs b/crates/code_assistant_core/src/session/turn.rs index a6399816..ff2202cf 100644 --- a/crates/code_assistant_core/src/session/turn.rs +++ b/crates/code_assistant_core/src/session/turn.rs @@ -147,7 +147,7 @@ pub struct TurnOutcome { pub struct TurnHandle { session_id: String, turn_id: u64, - service: crate::session::SessionService, + cancellation: tools_core::RunCancellation, outcome: oneshot::Receiver, } @@ -155,13 +155,13 @@ impl TurnHandle { pub(crate) fn new( session_id: String, turn_id: u64, - service: crate::session::SessionService, + cancellation: tools_core::RunCancellation, outcome: oneshot::Receiver, ) -> Self { Self { session_id, turn_id, - service, + cancellation, outcome, } } @@ -186,7 +186,8 @@ impl TurnHandle { /// Ask the running agent to stop at its next checkpoint. The outcome /// still resolves (normally as [`TurnStatus::Cancelled`]). pub async fn cancel(&self) -> Result<()> { - self.service.request_stop(self.session_id.clone()).await + self.cancellation.cancel(); + Ok(()) } } @@ -197,6 +198,7 @@ impl TurnHandle { /// /// [`SessionService::start_turn_if_idle`]: crate::session::SessionService::start_turn_if_idle pub struct TurnRecorder { + pub(crate) cancellation: tools_core::RunCancellation, turn_id: u64, started: Instant, inner: Mutex, @@ -223,7 +225,9 @@ impl TurnRecorder { pub(crate) fn arm(baseline_usage: llm::Usage) -> (std::sync::Arc, TurnParts) { let (tx, rx) = oneshot::channel(); let turn_id = NEXT_TURN_ID.fetch_add(1, Ordering::Relaxed); + let cancellation = tools_core::RunCancellation::default(); let recorder = std::sync::Arc::new(Self { + cancellation: cancellation.clone(), turn_id, started: Instant::now(), inner: Mutex::new(RecorderInner { @@ -242,6 +246,7 @@ impl TurnRecorder { ( recorder, TurnParts { + cancellation, turn_id, outcome: rx, }, @@ -348,6 +353,7 @@ impl TurnRecorder { return; }; let status = match error { + _ if self.cancellation.is_cancelled() => TurnStatus::Cancelled, Some(error) => TurnStatus::Failed { error }, None if inner.cancelled => TurnStatus::Cancelled, None => TurnStatus::Completed, @@ -405,6 +411,7 @@ impl Drop for TurnRecorder { /// The handle-side parts produced by [`TurnRecorder::arm`]. pub(crate) struct TurnParts { + pub(crate) cancellation: tools_core::RunCancellation, pub(crate) turn_id: u64, pub(crate) outcome: oneshot::Receiver, } diff --git a/crates/tools_core/Cargo.toml b/crates/tools_core/Cargo.toml index ad1e3a05..7cc1e9fb 100644 --- a/crates/tools_core/Cargo.toml +++ b/crates/tools_core/Cargo.toml @@ -13,6 +13,7 @@ regex = "1.12" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" thiserror = "1.0" +tokio = { version = "1", features = ["sync", "macros"] } [dev-dependencies] tokio = { version = "1", features = ["macros", "rt"] } diff --git a/crates/tools_core/src/cancellation.rs b/crates/tools_core/src/cancellation.rs new file mode 100644 index 00000000..3923850a --- /dev/null +++ b/crates/tools_core/src/cancellation.rs @@ -0,0 +1,75 @@ +//! One-shot run cancellation, independent of frontend events and error text. +use anyhow::Result; +use std::sync::Arc; +use tokio::sync::watch; + +#[derive(Debug, thiserror::Error)] +#[error("run cancelled")] +pub struct Cancelled; + +/// Clones address the same run. A subsequent run must use a fresh token. +#[derive(Clone, Debug)] +pub struct RunCancellation(Arc>); + +impl Default for RunCancellation { + fn default() -> Self { + Self(Arc::new(watch::channel(false).0)) + } +} + +impl RunCancellation { + pub fn cancel(&self) { + self.0.send_replace(true); + } + + pub fn is_cancelled(&self) -> bool { + *self.0.borrow() + } + + pub fn same_run(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } + + pub fn check(&self) -> Result<()> { + if self.is_cancelled() { + Err(Cancelled.into()) + } else { + Ok(()) + } + } + + pub async fn cancelled(&self) { + let mut receiver = self.0.subscribe(); + let _ = receiver.wait_for(|cancelled| *cancelled).await; + } + + /// Linearize a short synchronous publication with cancellation. Never + /// await or cancel this token from inside `f`. + pub fn if_active(&self, f: impl FnOnce() -> T) -> Result { + let cancelled = self.0.borrow(); + if *cancelled { + return Err(Cancelled.into()); + } + Ok(f()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn cancellation_is_sticky_and_run_local() { + let first = RunCancellation::default(); + let second = RunCancellation::default(); + first.cancel(); + first.cancelled().await; + assert!( + first + .if_active(|| panic!("publication after stop")) + .is_err() + ); + assert!(!second.is_cancelled()); + assert!(!first.same_run(&second)); + } +} diff --git a/crates/tools_core/src/lib.rs b/crates/tools_core/src/lib.rs index 6e76072c..985f09b5 100644 --- a/crates/tools_core/src/lib.rs +++ b/crates/tools_core/src/lib.rs @@ -9,6 +9,8 @@ //! [`ToolSpec`]; the crate prescribes no scoping vocabulary beyond the //! generic tags in [`spec::capabilities`]. +pub mod cancellation; +pub use cancellation::{Cancelled, RunCancellation}; pub mod coerce; pub mod dyn_tool; pub mod image; diff --git a/crates/tools_core/src/permissions.rs b/crates/tools_core/src/permissions.rs index 6ab55b11..85a5f4a4 100644 --- a/crates/tools_core/src/permissions.rs +++ b/crates/tools_core/src/permissions.rs @@ -52,6 +52,7 @@ impl PermissionTier { pub struct ToolPermissions { pub tier: PermissionTier, granted_tools: Arc>>, + cancellation: crate::RunCancellation, } impl ToolPermissions { @@ -59,9 +60,14 @@ impl ToolPermissions { Self { tier, granted_tools: Arc::default(), + cancellation: crate::RunCancellation::default(), } } + pub fn set_cancellation(&mut self, cancellation: crate::RunCancellation) { + self.cancellation = cancellation; + } + pub fn is_granted(&self, tool_name: &str) -> bool { self.granted_tools.lock().unwrap().contains(tool_name) } @@ -85,6 +91,7 @@ impl ToolPermissions { tool_id: Option<&str>, params: &serde_json::Value, ) -> Result<()> { + self.cancellation.check()?; if !self.tier.requires_permission(spec) { return Ok(()); } @@ -98,13 +105,16 @@ impl ToolPermissions { spec.name ); }; - let decision = handler - .request_permission(PermissionRequest { + let decision = tokio::select! { + biased; + _ = self.cancellation.cancelled() => return Err(crate::Cancelled.into()), + decision = handler.request_permission(PermissionRequest { tool_id, tool_name: &spec.name, reason: PermissionRequestReason::ToolInvocation { params }, - }) - .await?; + }) => decision?, + }; + self.cancellation.check()?; match decision { PermissionDecision::GrantedOnce => Ok(()), PermissionDecision::GrantedSession | PermissionDecision::GrantedPersistent => { From 9ff102f89dc954704044cfbafc04913cd9af1bc4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 9 Sep 2026 10:04:54 +0200 Subject: [PATCH 03/15] refactor(agent): unify tool dispatch and journal individual outcomes --- crates/agent_core/src/execution.rs | 69 +++ crates/agent_core/src/hooks.rs | 13 +- crates/agent_core/src/lib.rs | 1 + crates/agent_core/src/runtime.rs | 558 +----------------- crates/agent_core/src/runtime/tests.rs | 2 + .../src/runtime/tests/dispatch_tests.rs | 358 +++++++++++ .../agent_core/src/runtime/tool_execution.rs | 410 +++++++++++++ crates/agent_core/src/types.rs | 38 +- crates/tools_core/src/dyn_tool.rs | 10 + 9 files changed, 889 insertions(+), 570 deletions(-) create mode 100644 crates/agent_core/src/execution.rs create mode 100644 crates/agent_core/src/runtime/tests/dispatch_tests.rs create mode 100644 crates/agent_core/src/runtime/tool_execution.rs diff --git a/crates/agent_core/src/execution.rs b/crates/agent_core/src/execution.rs new file mode 100644 index 00000000..eea76958 --- /dev/null +++ b/crates/agent_core/src/execution.rs @@ -0,0 +1,69 @@ +//! Self-describing journal entries for calls without a concrete tool result. +//! Successful/functional-error tool outputs retain their existing codecs. + +use serde::{Deserialize, Serialize}; +use tools_core::{Render, ResourcesTracker, ToolResult}; + +pub(crate) const RUNTIME_OUTPUT_CODEC: &str = "__agent_runtime_outcome_v1"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExecutionState { + NotStarted, + /// Persisted BEFORE invoking a tool. After interruption we cannot tell + /// whether its effects happened, including the save/invoke crash window. + Started, + Succeeded, + Failed, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RuntimeToolOutput { + pub state: ExecutionState, + pub message: String, +} + +impl RuntimeToolOutput { + pub fn not_started(reason: impl AsRef) -> Self { + Self { + state: ExecutionState::NotStarted, + message: format!("Tool execution has not started. {}", reason.as_ref()), + } + } + + pub fn started() -> Self { + Self { + state: ExecutionState::Started, + message: "Tool execution may have started, but its outcome is unknown. Verify the state before retrying any side effects.".into(), + } + } + + pub fn failed(message: impl Into) -> Self { + Self { + state: ExecutionState::Failed, + message: message.into(), + } + } +} + +impl Render for RuntimeToolOutput { + fn status(&self) -> String { + match self.state { + ExecutionState::NotStarted => "Not started", + ExecutionState::Started => "Outcome unknown", + ExecutionState::Succeeded => "Success", + ExecutionState::Failed => "Error", + } + .into() + } + + fn render(&self, _: &mut ResourcesTracker) -> String { + self.message.clone() + } +} + +impl ToolResult for RuntimeToolOutput { + fn is_success(&self) -> bool { + self.state == ExecutionState::Succeeded + } +} diff --git a/crates/agent_core/src/hooks.rs b/crates/agent_core/src/hooks.rs index c73343e0..d8819f23 100644 --- a/crates/agent_core/src/hooks.rs +++ b/crates/agent_core/src/hooks.rs @@ -36,13 +36,16 @@ pub struct LoopCtx<'a> { /// Intercepts tool requests that the application handles itself instead of /// dispatching them to the registry, and observes successful executions. pub trait ToolInterceptor: Send + Sync { - /// Returns `Some(result)` when the request was handled here. Intercepted - /// tools do not appear in the UI. + /// Returns `Some(result)` when the request was handled here. Scope and + /// permission checks and the start checkpoint always precede this hook, + /// including for parallel batches. Intercepted tools do not appear in the UI. + /// Hooks may append a ToolExecution; otherwise a generic outcome is recorded. fn try_intercept(&self, _request: &ToolRequest, _ctx: &mut LoopCtx) -> Option> { None } - /// Invoked after any tool executed successfully (standard path included). + /// Invoked on the state owner after any successful tool (including + /// intercepted and parallel calls), with its final, possibly rewritten input. fn after_tool_success(&self, _request: &ToolRequest, _ctx: &mut LoopCtx) {} } @@ -74,7 +77,9 @@ pub trait IterationHook: Send + Sync { /// Decides which tool requests of a turn may execute concurrently. pub trait ToolDispatchPolicy: Send + Sync { - /// Indices of the requests that may execute concurrently with each other. + /// Indices of requests that support detached services and may overlap. + /// Only adjacent selected requests overlap: unselected calls are ordering + /// barriers. Authorization and completion hooks still run on the state owner. fn parallel_indices(&self, requests: &[ToolRequest]) -> Vec; } diff --git a/crates/agent_core/src/lib.rs b/crates/agent_core/src/lib.rs index dadd2202..682adbae 100644 --- a/crates/agent_core/src/lib.rs +++ b/crates/agent_core/src/lib.rs @@ -11,6 +11,7 @@ //! dyn-Any approach) — no generics infect the embedding application. pub mod dialect; +pub mod execution; pub mod hooks; pub mod native; pub mod persistence; diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index a64e2ea1..e358c050 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -3,6 +3,7 @@ #[cfg(test)] mod tests; +mod tool_execution; use crate::dialect::ToolDialect; use crate::hooks::{ContextSnapshot, HookRegistry, LoopCtx, RecoveryAction, ToolServicesProvider}; @@ -300,7 +301,11 @@ impl AgentRuntime { active_path: self.conversation.path().clone(), next_node_id: self.conversation.next_id(), messages: self.conversation.history().to_vec(), - tool_executions: self.tool_executions.clone(), + tool_executions: self + .tool_executions + .iter() + .map(ToolExecution::try_clone) + .collect::>()?, next_request_id: self.next_request_id, }; self.state_persistence @@ -545,295 +550,6 @@ impl AgentRuntime { } } - /// Executes a list of tool requests and appends tool results to message history. - /// Requests selected by the dispatch policy are executed concurrently. - async fn manage_tool_execution(&mut self, tool_requests: &[ToolRequest]) -> Result { - let parallel_indices = self.hooks.dispatch.parallel_indices(tool_requests); - - // Execute the policy-selected tools concurrently if we have multiple - let parallel_results = if parallel_indices.len() > 1 { - debug!("Running {} tools in parallel", parallel_indices.len()); - self.execute_tools_in_parallel( - parallel_indices - .iter() - .map(|i| &tool_requests[*i]) - .collect(), - ) - .await - } else { - Vec::new() - }; - - // Build content blocks in original order - let mut content_blocks: Vec> = vec![None; tool_requests.len()]; - let mut parallel_result_iter = parallel_results.into_iter(); - - // Process results in original order - for (idx, tool_request) in tool_requests.iter().enumerate() { - if self.cancellation.is_cancelled() { - break; - } - let result_block = if parallel_indices.len() > 1 && parallel_indices.contains(&idx) { - // This request ran in parallel - get result from parallel execution - - parallel_result_iter.next().unwrap_or_else(|| { - let start_time = Some(SystemTime::now()); - ContentBlock::ToolResult { - tool_use_id: tool_request.id.clone(), - content: ToolResultContent::text("Internal error: missing parallel result"), - is_error: Some(true), - start_time, - end_time: Some(SystemTime::now()), - } - }) - } else { - // Sequential execution - let start_time = Some(SystemTime::now()); - match self.execute_tool(tool_request).await { - Ok(success) => ContentBlock::ToolResult { - tool_use_id: tool_request.id.clone(), - content: ToolResultContent::text(""), - is_error: if success { None } else { Some(true) }, - start_time, - end_time: Some(SystemTime::now()), - }, - Err(e) => { - let error_text = Self::format_error_for_user(&e); - ContentBlock::ToolResult { - tool_use_id: tool_request.id.clone(), - content: ToolResultContent::text(error_text), - is_error: Some(true), - start_time, - end_time: Some(SystemTime::now()), - } - } - } - }; - content_blocks[idx] = Some(result_block); - } - - // Flatten and add message - let final_blocks: Vec<_> = content_blocks.into_iter().flatten().collect(); - if !final_blocks.is_empty() { - let result_message = Message::new_user_content(final_blocks); - self.append_message(result_message)?; - } - Ok(LoopFlow::Continue) - } - - /// Execute multiple tool requests in parallel. - /// Returns ContentBlocks in the same order as input. - async fn execute_tools_in_parallel( - &mut self, - tool_requests: Vec<&ToolRequest>, - ) -> Vec { - use futures::future::join_all; - - // Create futures for each tool request - let futures: Vec<_> = tool_requests - .iter() - .map(|tool_request| { - let request = (*tool_request).clone(); - let ui = self.ui.clone(); - let registry = self.registry.clone(); - let command_executor = self.command_executor.clone(); - let permission_handler = self.permission_handler.clone(); - let permissions = self.permissions.clone(); - let cancellation = self.cancellation.clone(); - let services_provider = self.services_provider.clone(); - let scope_tag = self.tool_capability.clone(); - let excluded_capabilities = self.excluded_tool_capabilities.clone(); - let session_id = self.session_id.clone(); - - async move { - let start_time = Some(SystemTime::now()); - - let (is_success, tool_execution) = Self::execute_tool_request_detached( - request, - ui, - registry, - command_executor, - permission_handler, - permissions, - cancellation, - services_provider, - scope_tag, - excluded_capabilities, - session_id, - ) - .await; - - let end_time = Some(SystemTime::now()); - - let content_block = ContentBlock::ToolResult { - tool_use_id: tool_execution.tool_request.id.clone(), - content: ToolResultContent::text(""), - is_error: if is_success { None } else { Some(true) }, - start_time, - end_time, - }; - (content_block, tool_execution) - } - }) - .collect(); - - // Execute all in parallel - let results = join_all(futures).await; - - // Collect results and tool executions - let mut content_blocks = Vec::new(); - for (content_block, tool_execution) in results { - debug!( - "Parallel tool {} ({}) completed", - tool_execution.tool_request.name, tool_execution.tool_request.id - ); - self.tool_executions.push(tool_execution); - content_blocks.push(content_block); - } - - content_blocks - } - - /// Execute a single tool request without exclusive access to the agent - /// state, for the parallel branch decided by the dispatch policy. - /// Compared to the sequential path, interceptors do not run, the plan is - /// unavailable, and input modifications are not propagated back into the - /// message history. - #[allow(clippy::too_many_arguments)] - async fn execute_tool_request_detached( - tool_request: ToolRequest, - ui: Arc, - registry: Arc, - command_executor: Arc, - permission_handler: Option>, - permissions: ToolPermissions, - cancellation: tools_core::RunCancellation, - services_provider: Arc, - scope_tag: String, - excluded_capabilities: Vec, - session_id: Option, - ) -> (bool, ToolExecution) { - let is_hidden = registry.is_tool_hidden(&tool_request.name, &scope_tag); - - // Update UI to show running status (skip for hidden tools) - if !is_hidden { - let _ = ui - .send_event(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status: crate::ui::ToolStatus::Running, - message: None, - output: None, - duration_seconds: None, - images: vec![], - }) - .await; - } - - let execution_start = std::time::Instant::now(); - - let invoke_result = match registry.get(&tool_request.name) { - None => Err(ToolError::UnknownTool(tool_request.name.clone()).into()), - Some(_) - if !registry.tool_has_capability(&tool_request.name, &scope_tag) - || excluded_capabilities - .iter() - .any(|cap| registry.tool_has_capability(&tool_request.name, cap)) => - { - Err(anyhow::anyhow!( - "Tool '{}' is not available in the current scope", - tool_request.name - )) - } - Some(tool) => { - // Tier-based permission gate, mirroring the sequential path. - match permissions - .check( - permission_handler.as_deref(), - &tool.spec(), - Some(&tool_request.id), - &tool_request.input, - ) - .await - { - Err(e) => Err(e), - Ok(()) if cancellation.is_cancelled() => Err(tools_core::Cancelled.into()), - Ok(()) => { - let mut services = services_provider.detached(&tool_request.id); - let mut context = ToolContext { - command_executor: command_executor.as_ref(), - tool_id: Some(tool_request.id.clone()), - session_id, - permission_handler: permission_handler.as_deref(), - extensions: Some(services.as_mut()), - }; - let mut input = tool_request.input.clone(); - tool.invoke(&mut context, &mut input).await - } - } - } - }; - - let execution_duration = Some(execution_start.elapsed().as_secs_f64()); - - match invoke_result { - Ok(result) => { - let success = result.is_success(); - let status = if success { - crate::ui::ToolStatus::Success - } else { - crate::ui::ToolStatus::Error - }; - - let status_msg = result.as_render().status(); - let mut resources_tracker = ResourcesTracker::new(); - let ui_output = result.as_render().render_for_ui(&mut resources_tracker); - let images = result.render_images(); - - if !is_hidden { - let _ = ui - .send_event(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status, - message: Some(status_msg), - output: Some(ui_output), - duration_seconds: execution_duration, - images, - }) - .await; - } - - ( - success, - ToolExecution { - tool_request, - result, - }, - ) - } - Err(e) => { - let error_text = Self::format_error_for_user(&e); - - if !is_hidden { - let _ = ui - .send_event(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status: crate::ui::ToolStatus::Error, - message: Some(error_text.clone()), - output: Some(error_text.clone()), - duration_seconds: execution_duration, - images: vec![], - }) - .await; - } - - ( - false, - ToolExecution::create_parse_error(tool_request.id, error_text), - ) - } - } - } - /// Get the appropriate system prompt based on tool mode fn get_system_prompt(&mut self) -> String { let cache_key = self @@ -1649,268 +1365,6 @@ impl AgentRuntime { }); } - /// A tool may rewrite its own input while executing (e.g. format-on-save). - /// Propagates the updated input to the UI and rewrites the originating tool - /// call in the message history so that follow-up requests see the final input. - async fn propagate_modified_tool_input( - &mut self, - original_request: &ToolRequest, - final_request: &ToolRequest, - is_hidden: bool, - ) -> Result<()> { - if !is_hidden { - self.notify_tool_parameter_updates( - &original_request.input, - &final_request.input, - &original_request.id, - ) - .await?; - } - - if let Err(e) = self.update_message_history_with_formatted_tool(final_request) { - warn!( - "Failed to update message history after input modification: {}", - e - ); - } - Ok(()) - } - - async fn execute_tool(&mut self, tool_request: &ToolRequest) -> Result { - self.cancellation.check()?; - debug!( - "Executing tool request: {} (id: {})", - tool_request.name, tool_request.id - ); - - if let Some(result) = self.intercept_tool(tool_request) { - return result; - } - - // Check if this is a hidden tool - let is_hidden = self - .registry - .as_ref() - .is_tool_hidden(&tool_request.name, self.tool_capability.as_str()); - - // Update status to Running before execution (skip for hidden tools) - if !is_hidden { - self.send_ui(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status: crate::ui::ToolStatus::Running, - message: None, - output: None, - duration_seconds: None, - images: vec![], - }) - .await?; - } - - // Get the tool - could fail with UnknownTool - let tool = match self.registry.as_ref().get(&tool_request.name) { - Some(tool) => tool, - None => return Err(ToolError::UnknownTool(tool_request.name.clone()).into()), - }; - - // Verify the tool is allowed in the current scope. - // The scope filtering on the tool list offered to the LLM is not sufficient on its own, - // because models may hallucinate tool calls they know from training even when the tool - // is not in the provided tool list (e.g. a sub-agent calling write_file). - if !self - .registry - .as_ref() - .tool_has_capability(&tool_request.name, self.tool_capability.as_str()) - || self.excluded_tool_capabilities.iter().any(|cap| { - self.registry - .as_ref() - .tool_has_capability(&tool_request.name, cap) - }) - { - return Err(anyhow::anyhow!( - "Tool '{}' is not available in the current scope", - tool_request.name - )); - } - - // Tier-based permission gate: ask the user before dispatching when - // the active tier requires it for this tool. - if let Err(e) = self - .permissions - .check( - self.permission_handler.as_deref(), - &tool.spec(), - Some(&tool_request.id), - &tool_request.input, - ) - .await - { - let error_text = Self::format_error_for_user(&e); - if !is_hidden { - self.send_ui(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status: crate::ui::ToolStatus::Error, - message: Some(error_text.clone()), - output: Some(error_text.clone()), - duration_seconds: None, - images: vec![], - }) - .await?; - } - self.tool_executions.push(ToolExecution::create_parse_error( - tool_request.id.clone(), - error_text, - )); - return Err(e); - } - - self.cancellation.check()?; - // Create a tool context. The services provider builds the application - // extension for this invocation (state such as the plan may move in - // for the duration) and takes it back afterwards. - let mut services = self - .services_provider - .begin(self.extensions.as_mut(), &tool_request.id); - let mut context = ToolContext { - command_executor: self.command_executor.as_ref(), - tool_id: Some(tool_request.id.clone()), - session_id: self.session_id.clone(), - permission_handler: self.permission_handler.as_deref(), - extensions: Some(services.as_mut()), - }; - - // Execute the tool - could fail with ParseError or other errors - let mut input = tool_request.input.clone(); - let execution_start = std::time::Instant::now(); - - let invoke_result = tool.invoke(&mut context, &mut input).await; - drop(context); - self.services_provider - .end(self.extensions.as_mut(), services); - - match invoke_result { - Ok(result) => { - let execution_duration = Some(execution_start.elapsed().as_secs_f64()); - - // Tool executed successfully (but may have failed functionally) - let success = result.is_success(); - - // Check if input parameters were modified during execution - let input_modified = input != tool_request.input; - - // Determine UI status based on result - let status = if success { - crate::ui::ToolStatus::Success - } else { - crate::ui::ToolStatus::Error - }; - - // Generate status string from result - let short_output = result.as_render().status(); - - // Generate output for UI display (may differ from LLM output for some tools) - let mut resources_tracker = ResourcesTracker::new(); - let ui_output = result.as_render().render_for_ui(&mut resources_tracker); - - // Collect image data from tools that produce visual output - let images = result.render_images(); - - // Update tool status with result (skip for hidden tools) - if !is_hidden { - self.send_ui(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status, - message: Some(short_output), - output: Some(ui_output), - duration_seconds: execution_duration, - images, - }) - .await?; - } - - // Create the tool request with potentially updated input - let final_tool_request = if input_modified { - debug!("Tool input was modified during execution"); - ToolRequest { - id: tool_request.id.clone(), - name: tool_request.name.clone(), - input: input.clone(), - start_offset: tool_request.start_offset, - end_offset: tool_request.end_offset, - } - } else { - tool_request.clone() - }; - - // Create and store the ToolExecution record - let tool_execution = ToolExecution { - tool_request: final_tool_request.clone(), - result, - }; - - // Store the execution record - self.tool_executions.push(tool_execution); - - if success { - self.after_tool_success(tool_request); - } - - if input_modified { - self.propagate_modified_tool_input( - tool_request, - &final_tool_request, - is_hidden, - ) - .await?; - } - - Ok(success) - } - - Err(e) => { - let execution_duration = Some(execution_start.elapsed().as_secs_f64()); - - // Tool execution failed (parameter error, etc.) - let error_text = Self::format_error_for_user(&e); - - // Update UI status to error (skip for hidden tools) - if !is_hidden { - self.send_ui(AgentUiEvent::UpdateToolStatus { - tool_id: tool_request.id.clone(), - status: crate::ui::ToolStatus::Error, - message: Some(error_text.clone()), - output: Some(error_text.clone()), - duration_seconds: execution_duration, - images: vec![], - }) - .await?; - } - - // Create a ToolExecution record for the error - let tool_execution = if let Some(tool_error) = e.downcast_ref::() { - match tool_error { - ToolError::ParseError(_) => { - // For parse errors, create a parse error execution - ToolExecution::create_parse_error(tool_request.id.clone(), error_text) - } - ToolError::UnknownTool(_) => { - // This shouldn't happen since we check above, but handle it - ToolExecution::create_parse_error(tool_request.id.clone(), error_text) - } - } - } else { - // For other error types, also create a parse error record - ToolExecution::create_parse_error(tool_request.id.clone(), error_text) - }; - - // Store the execution record - self.tool_executions.push(tool_execution); - - // Return the error to be handled by manage_tool_execution - Err(e) - } - } - } - async fn notify_tool_parameter_updates( &self, original: &serde_json::Value, diff --git a/crates/agent_core/src/runtime/tests.rs b/crates/agent_core/src/runtime/tests.rs index 8bd5e5b8..62b70a4f 100644 --- a/crates/agent_core/src/runtime/tests.rs +++ b/crates/agent_core/src/runtime/tests.rs @@ -1,3 +1,5 @@ +mod dispatch_tests; + use super::*; use crate::hooks::*; use serde_json::json; diff --git a/crates/agent_core/src/runtime/tests/dispatch_tests.rs b/crates/agent_core/src/runtime/tests/dispatch_tests.rs new file mode 100644 index 00000000..24a88676 --- /dev/null +++ b/crates/agent_core/src/runtime/tests/dispatch_tests.rs @@ -0,0 +1,358 @@ +use super::*; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; +use tools_core::{Render, Tool, ToolResult, ToolSpec}; + +#[derive(serde::Serialize, serde::Deserialize)] +struct Output(String); +impl ToolResult for Output { + fn is_success(&self) -> bool { + true + } +} +impl Render for Output { + fn status(&self) -> String { + "done".into() + } + fn render(&self, _: &mut ResourcesTracker) -> String { + self.0.clone() + } +} + +struct Probe { + calls: Arc>>, + entered: Arc, + release: Arc, +} +#[async_trait::async_trait] +impl Tool for Probe { + type Input = serde_json::Value; + type Output = Output; + fn spec(&self) -> ToolSpec { + ToolSpec { + name: "probe".into(), + description: "test".into(), + parameters_schema: json!({"type":"object"}), + annotations: None, + capabilities: ToolSpec::capabilities(&["test"]), + multiline_params: &[], + hidden: false, + title_template: None, + } + } + async fn execute<'a>( + &self, + _: &mut ToolContext<'a>, + input: &mut Self::Input, + ) -> Result { + let id = input["id"].as_str().unwrap().to_string(); + self.calls.lock().unwrap().push(id.clone()); + if input["wait"] == true { + self.entered.notify_one(); + self.release.notified().await; + } + if input["rewrite"] == true { + input["formatted"] = true.into(); + } + Ok(Output(format!("result for {id}"))) + } +} +struct Fixture { + agent: AgentRuntime, + saved: Capture, + calls: Arc>>, + entered: Arc, + release: Arc, +} +fn fixture(requests: &[ToolRequest]) -> Fixture { + let (mut agent, saved) = runtime(); + let calls = Arc::new(Mutex::new(vec![])); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let mut registry = ToolRegistry::new(); + registry.register(Box::new(Probe { + calls: calls.clone(), + entered: entered.clone(), + release: release.clone(), + })); + agent.registry = Arc::new(registry); + agent.tool_capability = "test".into(); + agent + .append_message(Message::new_assistant_content( + requests + .iter() + .map(|r| ContentBlock::new_tool_use(&r.id, &r.name, r.input.clone())) + .collect(), + )) + .unwrap(); + Fixture { + agent, + saved, + calls, + entered, + release, + } +} +fn request(id: &str, wait: bool) -> ToolRequest { + ToolRequest { + id: id.into(), + name: "probe".into(), + input: json!({"id":id, "wait":wait}), + start_offset: None, + end_offset: None, + } +} +struct Parallel(Vec); +impl ToolDispatchPolicy for Parallel { + fn parallel_indices(&self, _: &[ToolRequest]) -> Vec { + self.0.clone() + } +} +struct Observer { + attempts: Arc, + successes: Arc>>, + intercept: bool, +} +impl ToolInterceptor for Observer { + fn try_intercept(&self, _: &ToolRequest, _: &mut LoopCtx) -> Option> { + self.attempts.fetch_add(1, Ordering::SeqCst); + self.intercept.then_some(Ok(true)) + } + fn after_tool_success(&self, request: &ToolRequest, _: &mut LoopCtx) { + self.successes.lock().unwrap().push(request.clone()); + } +} + +#[tokio::test] +async fn dispatch_interceptors_cannot_bypass_scope_or_permission_checks() { + for restricted_scope in [true, false] { + let requests = vec![request("one", false)]; + let mut f = fixture(&requests); + let attempts = Arc::new(AtomicUsize::new(0)); + f.agent.hooks.interceptors.push(Box::new(Observer { + attempts: attempts.clone(), + successes: Arc::default(), + intercept: true, + })); + if restricted_scope { + f.agent.tool_capability = "other".into(); + } else { + f.agent.permissions = + tools_core::ToolPermissions::new(tools_core::PermissionTier::AllTools); + } + f.agent.manage_tool_execution(&requests).await.unwrap(); + assert_eq!( + attempts.load(Ordering::SeqCst), + 0, + "mandatory checks must precede interception" + ); + assert!(f.calls.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn dispatch_parallel_hooks_and_formatted_inputs_match_sequential_contract() { + let mut requests = vec![request("one", false), request("two", false)]; + for request in &mut requests { + request.input["rewrite"] = true.into(); + } + let mut f = fixture(&requests); + f.agent.hooks.dispatch = Box::new(Parallel(vec![0, 1])); + let attempts = Arc::new(AtomicUsize::new(0)); + let successes = Arc::new(Mutex::new(vec![])); + f.agent.hooks.interceptors.push(Box::new(Observer { + attempts: attempts.clone(), + successes: successes.clone(), + intercept: false, + })); + f.agent.manage_tool_execution(&requests).await.unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + let successes = successes.lock().unwrap(); + assert_eq!(successes.len(), 2); + assert!(successes.iter().all(|r| r.input["formatted"] == true)); + let saved = f.saved.0.lock().unwrap(); + let snapshot = saved.as_ref().unwrap(); + assert!( + snapshot + .tool_executions + .iter() + .all(|e| e.tool_request.input["formatted"] == true) + ); + assert!( + matches!(&snapshot.messages[0].content, MessageContent::Structured(blocks) + if blocks.iter().all(|b| matches!(b, ContentBlock::ToolUse { input, .. } if input["formatted"] == true))) + ); +} + +async fn completion_is_checkpointed_while_sibling_waits(parallel: bool) { + let requests = vec![request("one", false), request("two", true)]; + let mut f = fixture(&requests); + if parallel { + f.agent.hooks.dispatch = Box::new(Parallel(vec![0, 1])); + } + let task = tokio::spawn(async move { f.agent.manage_tool_execution(&requests).await }); + tokio::time::timeout(Duration::from_secs(2), f.entered.notified()) + .await + .unwrap(); + let checkpointed = tokio::time::timeout(Duration::from_millis(250), async { + loop { + if f.saved + .0 + .lock() + .unwrap() + .as_ref() + .unwrap() + .tool_executions + .iter() + .any(|e| e.tool_request.id == "one" && e.result.is_success()) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .is_ok(); + f.release.notify_one(); + task.await.unwrap().unwrap(); + assert!( + checkpointed, + "a completed tool must be durable before the batch finishes" + ); +} +#[tokio::test] +async fn dispatch_sequential_completions_are_saved_individually() { + completion_is_checkpointed_while_sibling_waits(false).await; +} +#[tokio::test] +async fn dispatch_parallel_completions_are_saved_individually() { + completion_is_checkpointed_while_sibling_waits(true).await; +} + +#[tokio::test] +async fn dispatch_journal_distinguishes_unstarted_from_uncertain_after_reload() { + let requests = vec![request("one", true), request("two", false)]; + let f = fixture(&requests); + let mut agent = f.agent; + let task = tokio::spawn(async move { agent.manage_tool_execution(&requests).await }); + tokio::time::timeout(Duration::from_secs(2), f.entered.notified()) + .await + .unwrap(); + let journal = f + .saved + .0 + .lock() + .unwrap() + .as_ref() + .unwrap() + .tool_executions + .clone(); + f.release.notify_one(); + task.await.unwrap().unwrap(); + assert_eq!( + journal.len(), + 2, + "journal records must precede tool invocation" + ); + // Runtime records must be self-describing, even when the tool disappears. + let registry = ToolRegistry::new(); + let restored: Vec<_> = journal + .iter() + .map(|entry| entry.serialize().unwrap().deserialize(®istry).unwrap()) + .collect(); + assert_eq!(restored[0].tool_request.name, "probe"); + assert!( + restored[0] + .result + .as_render() + .render(&mut ResourcesTracker::new()) + .contains("unknown") + ); + assert!( + restored[1] + .result + .as_render() + .render(&mut ResourcesTracker::new()) + .contains("not started") + ); +} + +struct FailingUi; +#[async_trait::async_trait] +impl AgentUi for FailingUi { + async fn send_event(&self, event: AgentUiEvent) -> Result<(), UIError> { + if matches!( + event, + AgentUiEvent::UpdateToolStatus { + status: crate::ui::ToolStatus::Success, + .. + } + ) { + return Err(UIError::IOError(std::io::Error::other("UI disconnected"))); + } + Ok(()) + } + fn display_fragment(&self, _: &DisplayFragment) -> Result<(), UIError> { + Ok(()) + } + fn should_streaming_continue(&self) -> bool { + true + } + fn notify_rate_limit(&self, _: u64) {} + fn clear_rate_limit(&self) {} +} +#[tokio::test] +async fn dispatch_ui_failure_does_not_erase_successful_tool_evidence() { + let requests = vec![request("one", false)]; + let mut f = fixture(&requests); + f.agent.ui = Arc::new(FailingUi); + let _ = f.agent.manage_tool_execution(&requests).await; + let snapshot = f.saved.0.lock().unwrap(); + assert!( + snapshot + .as_ref() + .unwrap() + .tool_executions + .iter() + .any(|e| e.tool_request.id == "one" && e.result.is_success()) + ); +} + +#[tokio::test] +async fn dispatch_parallel_groups_do_not_cross_sequential_barriers() { + let requests = vec![ + request("barrier", false), + request("one", false), + request("two", false), + ]; + let mut f = fixture(&requests); + f.agent.hooks.dispatch = Box::new(Parallel(vec![1, 2])); + f.agent.manage_tool_execution(&requests).await.unwrap(); + assert_eq!(f.calls.lock().unwrap()[0], "barrier"); +} + +struct FailCompletionSave; +impl SnapshotPersistence for FailCompletionSave { + fn save(&mut self, snapshot: AgentSnapshot, _: &(dyn Any + Send)) -> Result<()> { + anyhow::ensure!( + !snapshot + .tool_executions + .iter() + .any(|e| e.result.is_success()), + "disk full" + ); + Ok(()) + } +} +#[tokio::test] +async fn dispatch_checkpoint_failure_prevents_further_side_effects() { + let requests = vec![request("one", false), request("two", false)]; + let mut f = fixture(&requests); + f.agent.state_persistence = Box::new(FailCompletionSave); + assert!(f.agent.manage_tool_execution(&requests).await.is_err()); + assert_eq!( + &*f.calls.lock().unwrap(), + &["one"], + "persistence failure is not an ordinary tool error" + ); +} diff --git a/crates/agent_core/src/runtime/tool_execution.rs b/crates/agent_core/src/runtime/tool_execution.rs new file mode 100644 index 00000000..c7110b57 --- /dev/null +++ b/crates/agent_core/src/runtime/tool_execution.rs @@ -0,0 +1,410 @@ +//! One authorization/invocation/commit pipeline. Scheduling only decides +//! which adjacent invocations may overlap; it does not change their hooks, +//! permission checks, input correction or persistence semantics. +use super::*; +use crate::execution::{ExecutionState, RuntimeToolOutput}; +use futures::{FutureExt, StreamExt, stream::FuturesUnordered}; +use std::time::Instant; +use tools_core::AnyOutput; + +struct Invocation { + request: ToolRequest, + registry: Arc, + command_executor: Arc, + permission_handler: Option>, + cancellation: tools_core::RunCancellation, + session_id: Option, + services: Box, + detached: bool, +} + +struct Completion { + original: ToolRequest, + execution: ToolExecution, + services: Option>, + started_at: Option, + duration: Option, + intercepted: bool, +} + +impl Completion { + fn runtime(request: &ToolRequest, output: RuntimeToolOutput) -> Self { + Self { + original: request.clone(), + execution: ToolExecution { + tool_request: request.clone(), + result: Box::new(output), + }, + services: None, + started_at: None, + duration: None, + intercepted: false, + } + } +} + +impl Invocation { + async fn run(mut self) -> Completion { + let original = self.request.clone(); + let started_at = SystemTime::now(); + let start = Instant::now(); + let result = if self.cancellation.is_cancelled() { + Ok( + Box::new(RuntimeToolOutput::not_started("The run was cancelled.")) + as Box, + ) + } else { + let mut context = ToolContext { + command_executor: self.command_executor.as_ref(), + tool_id: Some(self.request.id.clone()), + session_id: self.session_id, + permission_handler: self.permission_handler.as_deref(), + extensions: Some(self.services.as_mut()), + }; + // The registry is immutable for this run and was checked in prepare. + let tool = self + .registry + .get(&self.request.name) + .expect("authorized tool"); + match std::panic::AssertUnwindSafe(tool.invoke(&mut context, &mut self.request.input)) + .catch_unwind() + .await + { + Ok(result) => result, + Err(_) => Err(anyhow::anyhow!( + "Tool panicked; partial side effects may have occurred. Verify the state before retrying." + )), + } + }; + Completion { + original, + execution: ToolExecution { + tool_request: self.request, + result: result.unwrap_or_else(|error| { + Box::new(RuntimeToolOutput::failed( + AgentRuntime::format_error_for_user(&error), + )) + }), + }, + services: if self.detached { + None + } else { + Some(self.services) + }, + started_at: Some(started_at), + duration: Some(start.elapsed().as_secs_f64()), + intercepted: false, + } + } +} + +impl AgentRuntime { + pub(super) async fn manage_tool_execution( + &mut self, + requests: &[ToolRequest], + ) -> Result { + let mut seen = std::collections::HashSet::new(); + for request in requests { + anyhow::ensure!( + seen.insert(&request.id), + "Duplicate tool call id: {}", + request.id + ); + anyhow::ensure!( + !self + .tool_executions + .iter() + .any(|e| e.tool_request.id == request.id), + "Tool call id has already been recorded: {}", + request.id + ); + } + // Write intent for the whole batch before any tool can have effects. + // A restart can distinguish an unstarted sibling from an uncertain call. + for request in requests { + self.store_execution(ToolExecution { + tool_request: request.clone(), + result: Box::new(RuntimeToolOutput::not_started("No invocation was made.")), + }); + } + self.save_state()?; + + let mut parallel = vec![false; requests.len()]; + for index in self.hooks.dispatch.parallel_indices(requests) { + anyhow::ensure!( + index < requests.len(), + "Dispatch policy returned invalid index {index}" + ); + parallel[index] = true; + } + let mut blocks: Vec> = vec![None; requests.len()]; + let mut index = 0; + while index < requests.len() && !self.cancellation.is_cancelled() { + let mut end = index + 1; + if parallel[index] { + while end < requests.len() && parallel[end] { + end += 1; + } + } + let detached = end - index > 1; + let mut pending = FuturesUnordered::new(); + for i in index..end { + if self.cancellation.is_cancelled() { + break; + } + match self.prepare_invocation(&requests[i], detached).await? { + Ok(invocation) => pending.push(async move { (i, invocation.run().await) }), + Err(completed) => blocks[i] = Some(self.commit_completion(completed).await?), + } + } + // Consume completion order, not request order. A slow sibling must + // never keep a completed side effect out of the durable checkpoint. + // We do not drop active tools on ordinary stop: tools with effects + // finish or implement their own cooperative cancellation. + while let Some((i, completed)) = pending.next().await { + blocks[i] = Some(self.commit_completion(completed).await?); + } + index = end; + } + + for (index, request) in requests.iter().enumerate() { + if blocks[index].is_none() { + // This request never got past preparation. Settle the UI too; + // the model may already have streamed a pending tool card. + blocks[index] = Some( + self.commit_completion(Completion::runtime( + request, + RuntimeToolOutput::not_started("The run was cancelled before invocation."), + )) + .await?, + ); + } + } + if !blocks.is_empty() { + self.append_message(Message::new_user_content( + blocks.into_iter().flatten().collect(), + ))?; + } + Ok(LoopFlow::Continue) + } + + async fn authorize_invocation(&self, request: &ToolRequest) -> Result<()> { + self.cancellation.check()?; + let tool = self + .registry + .get(&request.name) + .ok_or_else(|| ToolError::UnknownTool(request.name.clone()))?; + anyhow::ensure!( + self.registry + .tool_has_capability(&request.name, &self.tool_capability) + && !self + .excluded_tool_capabilities + .iter() + .any(|cap| self.registry.tool_has_capability(&request.name, cap)), + "Tool '{}' is not available in the current scope", + request.name + ); + self.permissions + .check( + self.permission_handler.as_deref(), + &tool.spec(), + Some(&request.id), + &request.input, + ) + .await?; + self.cancellation.check() + } + + /// Outer errors are infrastructure failures and abort dispatch. A rejected + /// or intercepted tool is a normal completion, handled by the same commit path. + async fn prepare_invocation( + &mut self, + request: &ToolRequest, + detached: bool, + ) -> Result> { + if let Err(error) = self.authorize_invocation(request).await { + return Ok(Err(Completion::runtime( + request, + RuntimeToolOutput::not_started(Self::format_error_for_user(&error)), + ))); + } + self.store_execution(ToolExecution { + tool_request: request.clone(), + result: Box::new(RuntimeToolOutput::started()), + }); + self.save_state()?; + + // Interceptors execute on the state owner, even for a parallel group, + // and only after scope/permission checks and the start checkpoint. + if let Some(result) = self.intercept_tool(request) { + let output: Box = match result { + Ok(success) => { + let entry = self + .tool_executions + .iter() + .rev() + .find(|entry| entry.tool_request.id == request.id) + .expect("journaled invocation"); + let still_started = entry + .result + .as_any() + .and_then(|out| out.downcast_ref::()) + .is_some_and(|out| out.state == ExecutionState::Started); + if still_started { + Box::new(RuntimeToolOutput { + state: if success { + ExecutionState::Succeeded + } else { + ExecutionState::Failed + }, + message: "Handled by the application's tool interceptor.".into(), + }) + } else { + entry.result.try_clone()? + } + } + Err(error) => Box::new(RuntimeToolOutput::failed(Self::format_error_for_user( + &error, + ))), + }; + return Ok(Err(Completion { + original: request.clone(), + execution: ToolExecution { + tool_request: request.clone(), + result: output, + }, + services: None, + started_at: Some(SystemTime::now()), + duration: None, + intercepted: true, + })); + } + if !self + .registry + .is_tool_hidden(&request.name, &self.tool_capability) + { + // UI is a projection, not the owner of execution or evidence. + let _ = self + .send_ui(AgentUiEvent::UpdateToolStatus { + tool_id: request.id.clone(), + status: crate::ui::ToolStatus::Running, + message: None, + output: None, + duration_seconds: None, + images: vec![], + }) + .await; + } + let services = if detached { + self.services_provider.detached(&request.id) + } else { + self.services_provider + .begin(self.extensions.as_mut(), &request.id) + }; + Ok(Ok(Invocation { + request: request.clone(), + registry: self.registry.clone(), + command_executor: self.command_executor.clone(), + permission_handler: self.permission_handler.clone(), + cancellation: self.cancellation.clone(), + session_id: self.session_id.clone(), + services, + detached, + })) + } + + /// Replace the journal slot, retaining deterministic request order. Legacy + /// interceptors may append their own record; collapse that duplicate by id. + fn store_execution(&mut self, execution: ToolExecution) { + let id = execution.tool_request.id.clone(); + if let Some(index) = self + .tool_executions + .iter() + .position(|entry| entry.tool_request.id == id) + { + self.tool_executions[index] = execution; + let mut first = true; + self.tool_executions.retain(|entry| { + if entry.tool_request.id != id { + return true; + } + std::mem::replace(&mut first, false) + }); + } else { + self.tool_executions.push(execution); + } + } + + async fn commit_completion(&mut self, mut completed: Completion) -> Result { + if let Some(services) = completed.services.take() { + self.services_provider + .end(self.extensions.as_mut(), services); + } + let request = completed.execution.tool_request.clone(); + let success = completed.execution.result.is_success(); + let changed = request.input != completed.original.input; + let hidden = completed.intercepted + || self + .registry + .is_tool_hidden(&request.name, &self.tool_capability); + self.store_execution(completed.execution); + // Commit evidence before hooks/rendering/UI can fail. The active call's + // Started record remains on disk if this commit itself fails. + self.save_state()?; + if changed { + self.update_message_history_with_formatted_tool(&request)?; + } + if success { + self.after_tool_success(&request); + self.save_state()?; + } + let execution = self + .tool_executions + .iter() + .find(|entry| entry.tool_request.id == request.id) + .expect("committed outcome"); + let content = execution + .result + .as_any() + .and_then(|out| out.downcast_ref::()) + .map(|out| out.message.clone()) + .unwrap_or_default(); + if !hidden { + let _ = self + .send_ui(AgentUiEvent::UpdateToolStatus { + tool_id: request.id.clone(), + status: if success { + crate::ui::ToolStatus::Success + } else { + crate::ui::ToolStatus::Error + }, + message: Some(execution.result.as_render().status()), + output: Some( + execution + .result + .as_render() + .render_for_ui(&mut ResourcesTracker::new()), + ), + duration_seconds: completed.duration, + images: execution.result.render_images(), + }) + .await; + if changed { + let _ = self + .notify_tool_parameter_updates( + &completed.original.input, + &request.input, + &request.id, + ) + .await; + } + } + Ok(ContentBlock::ToolResult { + tool_use_id: request.id, + content: ToolResultContent::text(content), + is_error: if success { None } else { Some(true) }, + start_time: completed.started_at, + end_time: Some(SystemTime::now()), + }) + } +} diff --git a/crates/agent_core/src/types.rs b/crates/agent_core/src/types.rs index 91167cf7..45bfc3aa 100644 --- a/crates/agent_core/src/types.rs +++ b/crates/agent_core/src/types.rs @@ -2,13 +2,13 @@ //! records, and the placeholder outputs the loop itself produces (parse //! errors, prompt-too-long replacements). +use crate::execution::{RUNTIME_OUTPUT_CODEC, RuntimeToolOutput}; use anyhow::Result; use serde::{Deserialize, Serialize}; use serde_json::Value; use tools_core::{ AnnotatedToolDefinition, AnyOutput, Render, ResourcesTracker, ToolRegistry, ToolResult, }; -use tracing::debug; /// Convert a basic ToolDefinition (without annotations) for LLM providers pub fn to_tool_definition(tool: &AnnotatedToolDefinition) -> llm::ToolDefinition { @@ -196,23 +196,23 @@ impl ToolExecution { /// Serialize the tool execution to a storable format pub fn serialize(&self) -> Result { - // Try to serialize the result, but fallback to a simple representation if it fails - let result_json = match self.result.to_json() { - Ok(json) => json, - Err(e) => { - debug!("Failed to serialize tool result, using fallback: {}", e); - serde_json::json!({ - "error": "Failed to serialize result", - "success": self.result.is_success(), - "details": format!("{}", e) - }) - } + // An undecodable fallback is not a committed outcome. Fail the save so + // the runtime stops dispatch instead of losing evidence silently. + let result_json = self.result.to_json()?; + let tool_name = if self + .result + .as_any() + .is_some_and(|output| output.is::()) + { + RUNTIME_OUTPUT_CODEC + } else { + &self.tool_request.name }; Ok(SerializedToolExecution { tool_request: self.tool_request.clone(), result_json, - tool_name: self.tool_request.name.clone(), + tool_name: tool_name.to_string(), }) } } @@ -234,11 +234,21 @@ impl SerializedToolExecution { /// (a reconfigured MCP server, a removed integration). Callers check /// this to skip the record instead of failing the whole session load. pub fn tool_available(&self, registry: &ToolRegistry) -> bool { - self.tool_name == "parse_error" || registry.get(&self.tool_name).is_some() + self.tool_name == "parse_error" + || self.tool_name == RUNTIME_OUTPUT_CODEC + || registry.get(&self.tool_name).is_some() } /// Deserialize back to a ToolExecution pub fn deserialize(&self, registry: &ToolRegistry) -> Result { + if self.tool_name == RUNTIME_OUTPUT_CODEC { + return Ok(ToolExecution { + tool_request: self.tool_request.clone(), + result: Box::new(serde_json::from_value::( + self.result_json.clone(), + )?), + }); + } // Special handling for parse errors if self.tool_name == "parse_error" { let parse_error: ParseError = serde_json::from_value(self.result_json.clone())?; diff --git a/crates/tools_core/src/dyn_tool.rs b/crates/tools_core/src/dyn_tool.rs index 0b2b0556..c8bafa1d 100644 --- a/crates/tools_core/src/dyn_tool.rs +++ b/crates/tools_core/src/dyn_tool.rs @@ -9,6 +9,12 @@ use serde_json::Value; /// Type-erased tool output that can be rendered and determined for success pub trait AnyOutput: Send + Sync { + /// Optional concrete-type access for runtime-owned, self-describing outcomes. + /// Custom implementations can keep the default; ordinary tool outputs + /// continue to be decoded by their tool's registry entry. + fn as_any(&self) -> Option<&dyn std::any::Any> { + None + } /// Get a reference to the output as a Render trait object fn as_render(&self) -> &dyn Render; @@ -33,6 +39,10 @@ impl AnyOutput for T where T: Render + ToolResult + Serialize + DeserializeOwned + Send + Sync + 'static, { + fn as_any(&self) -> Option<&dyn std::any::Any> { + Some(self) + } + fn as_render(&self) -> &dyn Render { self } From aecda2364a62d265c08df79d270dba38537f8b74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 9 Sep 2026 10:30:11 +0200 Subject: [PATCH 04/15] fix(streaming): avoid draining cleared buffers for empty tool inputs --- .../agent_core/src/native/json_processor.rs | 6 +++--- .../src/ui/streaming/json_processor_tests.rs | 21 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/crates/agent_core/src/native/json_processor.rs b/crates/agent_core/src/native/json_processor.rs index 8260308f..a2c9ed1c 100644 --- a/crates/agent_core/src/native/json_processor.rs +++ b/crates/agent_core/src/native/json_processor.rs @@ -420,9 +420,9 @@ impl JsonStreamProcessor { self.emit_fragment(DisplayFragment::ToolEnd { id: tool_id })?; self.state.json_parsing_state = JsonParsingState::ExpectOpenBrace; // Reset for next potential JSON object self.state.current_key = None; - self.state.buffer.clear(); // Object done, clear buffer of this object. This might be too aggressive if there's trailing content. - // Let's refine: only clear if this was the *only* content, or handle trailing chars. - // For now, `drain` handles consumed chars. + // The loop below consumes this closing brace. Clearing + // here would both discard trailing input and make that + // drain panic for an empty object ({}). } else if char_to_process == ',' { // This is for cases like {"a":"b",} -> expecting a key next. // If we see `,,,` this will just loop. Assuming valid JSON structure mostly. diff --git a/crates/code_assistant_core/src/ui/streaming/json_processor_tests.rs b/crates/code_assistant_core/src/ui/streaming/json_processor_tests.rs index e59d0a07..e67746e4 100644 --- a/crates/code_assistant_core/src/ui/streaming/json_processor_tests.rs +++ b/crates/code_assistant_core/src/ui/streaming/json_processor_tests.rs @@ -55,6 +55,27 @@ fn process_json_chunks(chunks: &[String], tool_name: &str, tool_id: &str) -> Vec mod tests { use super::*; + #[test] + fn empty_tool_input_closes_once_without_panicking() { + for size in 1..=4 { + let fragments = process_json_chunks(&chunk_str("{} \n", size), "probe", "empty"); + assert_eq!( + fragments + .iter() + .filter(|fragment| matches!(fragment, + DisplayFragment::ToolEnd { id } if id == "empty" + )) + .count(), + 1 + ); + assert!( + !fragments + .iter() + .any(|fragment| matches!(fragment, DisplayFragment::ToolParameter { .. })) + ); + } + } + #[test] fn test_basic_json_param_parsing() { let json = r#"{"path": "src/main.rs"}"#; From 35916ad39b7a3e1cb96c42aa0f8297e4bf6e8863 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 9 Sep 2026 10:44:01 +0200 Subject: [PATCH 05/15] fix(agent): finalize sub-agent failures and propagate run cancellation --- .../src/agent/checkpoint_tests.rs | 74 ++++ .../src/agent/sub_agent.rs | 271 ++++++------- .../src/agent/sub_agent/run.rs | 123 ++++++ .../src/agent/sub_agent/tests.rs | 364 ++++++++++++++++++ .../src/session/manager.rs | 8 +- .../src/tools/impls/spawn_agent.rs | 23 +- 6 files changed, 723 insertions(+), 140 deletions(-) create mode 100644 crates/code_assistant_core/src/agent/sub_agent/run.rs create mode 100644 crates/code_assistant_core/src/agent/sub_agent/tests.rs diff --git a/crates/code_assistant_core/src/agent/checkpoint_tests.rs b/crates/code_assistant_core/src/agent/checkpoint_tests.rs index 32bef056..181780f3 100644 --- a/crates/code_assistant_core/src/agent/checkpoint_tests.rs +++ b/crates/code_assistant_core/src/agent/checkpoint_tests.rs @@ -161,6 +161,80 @@ async fn formatted_roundtrip(syntax: ToolSyntax) -> Result<()> { Ok(()) } +#[test] +fn journal_outcomes_survive_disk_reload_without_the_original_tools() -> Result<()> { + use agent_core::execution::{ExecutionState, RuntimeToolOutput}; + let dir = tempdir()?; + let mut session = ChatSession::new_empty( + "journal".into(), + "test".into(), + SessionConfig::default(), + None, + ); + let outputs = [ + RuntimeToolOutput { + state: ExecutionState::Succeeded, + message: "completed before interruption".into(), + }, + RuntimeToolOutput::started(), + RuntimeToolOutput::not_started("No invocation was made."), + ]; + for (i, output) in outputs.into_iter().enumerate() { + let id = format!("call-{i}"); + let request = agent_core::ToolRequest { + id: id.clone(), + name: "unavailable-tool".into(), + input: serde_json::json!({"path":"evidence.txt"}), + start_offset: None, + end_offset: None, + }; + session.add_message(Message::new_assistant_content(vec![ + ContentBlock::new_tool_use(&id, &request.name, request.input.clone()), + ])); + session.tool_executions.push( + agent_core::ToolExecution { + tool_request: request, + result: Box::new(output), + } + .serialize()?, + ); + } + let mut store = FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()); + store.save_chat_session(&session)?; + let loaded = store.load_chat_session("journal")?.unwrap(); + let registry = tools_core::ToolRegistry::new(); + let restored: Vec<_> = loaded + .tool_executions + .iter() + .map(|entry| { + assert!(crate::tools::mcp::execution_renderable(entry, ®istry)); + crate::tools::mcp::deserialize_tool_execution(entry, ®istry) + }) + .collect::>()?; + assert_eq!(restored.len(), 3); + for entry in &restored { + assert_eq!(entry.tool_request.name, "unavailable-tool"); + assert_eq!(entry.tool_request.input["path"], "evidence.txt"); + } + assert!(restored[0].result.is_success()); + let mut tracker = tools_core::ResourcesTracker::new(); + assert!( + restored[1] + .result + .as_render() + .render(&mut tracker) + .contains("unknown") + ); + assert!( + restored[2] + .result + .as_render() + .render(&mut tracker) + .contains("not started") + ); + Ok(()) +} + #[tokio::test] async fn checkpoint_native_formatted_roundtrip() -> Result<()> { formatted_roundtrip(ToolSyntax::Native).await diff --git a/crates/code_assistant_core/src/agent/sub_agent.rs b/crates/code_assistant_core/src/agent/sub_agent.rs index ebd10a46..e96ffde4 100644 --- a/crates/code_assistant_core/src/agent/sub_agent.rs +++ b/crates/code_assistant_core/src/agent/sub_agent.rs @@ -1,3 +1,7 @@ +mod run; +#[cfg(test)] +mod tests; + use crate::agent::persistence::AgentStatePersistence; use crate::agent::{Agent, AgentComponents}; use crate::config::DefaultProjectManager; @@ -16,21 +20,71 @@ use tools_core::permissions::{PermissionMediator, ToolPermissions}; /// Cancellation registry keyed by the parent `spawn_agent` tool id. #[derive(Default)] pub struct SubAgentCancellationRegistry { - flags: Mutex>>, + flags: Mutex>, +} + +#[derive(Clone)] +struct ChildCancellation { + flag: Arc, + token: tools_core::RunCancellation, +} + +impl ChildCancellation { + fn cancel(&self) { + self.flag.store(true, Ordering::SeqCst); + self.token.cancel(); + } +} + +struct ChildRegistration<'a> { + registry: &'a SubAgentCancellationRegistry, + tool_id: String, + cancellation: ChildCancellation, +} + +impl Drop for ChildRegistration<'_> { + fn drop(&mut self) { + let mut entries = self.registry.flags.lock().unwrap(); + // A delayed old task must not unregister a replacement with the same id. + if entries + .get(&self.tool_id) + .is_some_and(|entry| entry.token.same_run(&self.cancellation.token)) + { + entries.remove(&self.tool_id); + } + } } impl SubAgentCancellationRegistry { + fn insert(&self, tool_id: String) -> ChildCancellation { + let child = ChildCancellation { + flag: Arc::new(AtomicBool::new(false)), + token: tools_core::RunCancellation::default(), + }; + if let Some(previous) = self.flags.lock().unwrap().insert(tool_id, child.clone()) { + previous.cancel(); + } + child + } + + /// Compatibility flag for callers that observe cancellation synchronously. + /// Use `cancel` to also wake asynchronous waiters. pub fn register(&self, tool_id: String) -> Arc { - let flag = Arc::new(AtomicBool::new(false)); - let mut flags = self.flags.lock().unwrap(); - flags.insert(tool_id, flag.clone()); - flag + self.insert(tool_id).flag + } + + fn register_run(&self, tool_id: &str) -> ChildRegistration<'_> { + ChildRegistration { + registry: self, + tool_id: tool_id.to_string(), + cancellation: self.insert(tool_id.to_string()), + } } pub fn cancel(&self, tool_id: &str) -> bool { let flags = self.flags.lock().unwrap(); - if let Some(flag) = flags.get(tool_id) { - flag.store(true, Ordering::SeqCst); + if let Some(child) = flags.get(tool_id) { + child.cancel(); true } else { false @@ -114,6 +168,21 @@ pub struct SubAgentResult { pub ui_output: String, } +/// A failed child still has useful structured evidence. The spawn tool persists +/// this output instead of replacing it with an unstructured error string. +#[derive(Debug)] +pub struct SubAgentFailure { + pub message: String, + pub ui_output: String, +} + +impl std::fmt::Display for SubAgentFailure { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.message) + } +} +impl std::error::Error for SubAgentFailure {} + /// Execution mode for a sub-agent, selected by the `spawn_agent` tool input. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SubAgentMode { @@ -156,6 +225,9 @@ pub struct DefaultSubAgentRunner { /// Hook factory sub-agents run with (shared with the parent agent); /// `None` uses code-assistant's default hooks. hooks_factory: Option, + llm_client_factory: Option, + project_manager_factory: crate::session::service::ProjectManagerFactory, + parent_cancellation: tools_core::RunCancellation, } impl DefaultSubAgentRunner { @@ -185,9 +257,33 @@ impl DefaultSubAgentRunner { tool_registry, session_source, hooks_factory, + llm_client_factory: None, + project_manager_factory: Arc::new(|| Box::new(DefaultProjectManager::new())), + parent_cancellation: tools_core::RunCancellation::default(), } } + pub fn with_llm_client_factory( + mut self, + factory: crate::session::service::LlmClientFactory, + ) -> Self { + self.llm_client_factory = Some(factory); + self + } + + pub fn with_project_manager_factory( + mut self, + factory: crate::session::service::ProjectManagerFactory, + ) -> Self { + self.project_manager_factory = factory; + self + } + + pub fn with_parent_cancellation(mut self, cancellation: tools_core::RunCancellation) -> Self { + self.parent_cancellation = cancellation; + self + } + fn build_sub_agent_ui( &self, parent_ui: Arc, @@ -209,12 +305,21 @@ impl DefaultSubAgentRunner { permission_handler: Option>, ) -> Result { // Create a fresh LLM provider (avoid requiring Clone). - let llm_provider = - llm::factory::create_llm_client_from_model(&self.model_name, None, false, None).await?; + let llm_provider = match &self.llm_client_factory { + Some(factory) => { + let factory = factory.clone(); + let model = self.model_name.clone(); + tokio::task::spawn_blocking(move || factory(&model)).await?? + } + None => { + llm::factory::create_llm_client_from_model(&self.model_name, None, false, None) + .await? + } + }; // Create a fresh project manager, copying init_path if set. let project_manager: Arc = - Arc::new(DefaultProjectManager::new()); + Arc::from((self.project_manager_factory)()); if let Some(path) = self.session_config.effective_project_path().cloned() { let _ = project_manager.add_temporary_project(path); } @@ -280,124 +385,6 @@ fn tool_scope_for_subagent() -> ToolScope { ToolScope::SubAgentReadOnly } -#[async_trait::async_trait] -impl SubAgentRunner for DefaultSubAgentRunner { - async fn run( - &self, - parent_tool_id: &str, - instructions: String, - mode: SubAgentMode, - require_file_references: bool, - ) -> Result { - // Sub-agents inherit the parent session's edit-tool layout: with the - // diff-format edit tool, the sub-agent uses `replace_in_file` instead - // of `edit`. - let tool_scope = match mode { - SubAgentMode::ReadOnly => ToolScope::SubAgentReadOnly, - SubAgentMode::Default if self.session_config.use_diff_blocks => { - ToolScope::SubAgentDefaultWithDiffBlocks - } - SubAgentMode::Default => ToolScope::SubAgentDefault, - }; - - let cancelled = self - .cancellation_registry - .register(parent_tool_id.to_string()); - let sub_ui = self.build_sub_agent_ui( - self.ui.clone(), - parent_tool_id.to_string(), - cancelled.clone(), - ); - - // Keep a clone of the adapter so we can set the final response - let sub_ui_adapter = sub_ui.clone(); - - let mut agent = self - .build_agent( - parent_tool_id, - sub_ui as Arc, - self.permission_handler.clone(), - ) - .await?; - agent.set_tool_scope(tool_scope); - - // Start with a single user message containing the full instructions. - agent.append_message(Message::new_user(instructions))?; - - // Run 1+ iterations if we need to enforce file references. - let mut last_answer = String::new(); - let mut was_cancelled = false; - - for attempt in 0..=2 { - // Check for cancellation before starting iteration - if cancelled.load(Ordering::SeqCst) { - was_cancelled = true; - break; - } - - agent.run_single_iteration().await?; - - // Update usage after each iteration so the UI ring indicator - // reflects current token consumption while the sub-agent is still running. - let usage = compute_sub_agent_usage(agent.message_history(), &self.model_name); - sub_ui_adapter.set_usage(usage); - sub_ui_adapter.send_output_update().await; - - // Check for cancellation after iteration completes - // (cancellation may have occurred during streaming/tool execution) - if cancelled.load(Ordering::SeqCst) { - was_cancelled = true; - break; - } - - last_answer = extract_last_assistant_text(agent.message_history()).unwrap_or_default(); - - if !require_file_references { - break; - } - - if has_file_references_with_line_ranges(&last_answer) { - break; - } - - if attempt >= 2 { - // Best-effort: return with warning. - last_answer = format!( - "{last_answer}\n\n(Warning: requested file references with line ranges, but the sub-agent did not include them.)" - ); - break; - } - - // Ask the same sub-agent to revise. - agent.append_message(Message::new_user( - "Please revise your last answer to include exact file references with line ranges (e.g. `path/to/file.rs:10-20`).".to_string(), - ))?; - } - - self.cancellation_registry.unregister(parent_tool_id); - - // Handle cancellation: return error - if was_cancelled { - sub_ui_adapter.set_cancelled(); - return Err(anyhow::anyhow!("Cancelled by user")); - } - - // Collect token usage from the sub-agent's message history - let usage = compute_sub_agent_usage(agent.message_history(), &self.model_name); - sub_ui_adapter.set_usage(usage); - - // Set the final response in the adapter and get the complete JSON output - // This preserves the tools list along with the final response for rendering - sub_ui_adapter.set_response(last_answer.clone()); - let final_json = sub_ui_adapter.get_final_output(); - - Ok(SubAgentResult { - answer: last_answer, - ui_output: final_json, - }) - } -} - /// Compute aggregated token usage from a sub-agent's message history. fn compute_sub_agent_usage(messages: &[Message], model_name: &str) -> SubAgentUsage { let mut total = SubAgentUsage::default(); @@ -599,9 +586,21 @@ impl SubAgentUiAdapter { } async fn send_output_update(&self) { - let (json, tool_count, activity) = { + let (json, tool_count, activity, status, message) = { let output = self.output.lock().unwrap(); - (output.to_json(), output.tools.len(), output.activity) + let (status, message) = match output.activity { + Some(SubAgentActivity::Completed) => (ToolStatus::Success, "Sub-agent completed"), + Some(SubAgentActivity::Failed) => (ToolStatus::Error, "Sub-agent failed"), + Some(SubAgentActivity::Cancelled) => (ToolStatus::Error, "Sub-agent cancelled"), + _ => (ToolStatus::Running, "Sub-agent running"), + }; + ( + output.to_json(), + output.tools.len(), + output.activity, + status, + message, + ) }; tracing::debug!( @@ -615,8 +614,8 @@ impl SubAgentUiAdapter { .parent .send_event(UiEvent::UpdateToolStatus { tool_id: self.parent_tool_id.clone(), - status: ToolStatus::Running, - message: Some("Sub-agent running".to_string()), + status, + message: Some(message.to_string()), output: Some(json), styled_output: None, duration_seconds: None, @@ -724,6 +723,12 @@ impl SubAgentUiAdapter { let mut output = self.output.lock().unwrap(); output.error = Some(error); output.activity = Some(SubAgentActivity::Failed); + for tool in &mut output.tools { + if tool.status == SubAgentToolStatus::Running { + tool.status = SubAgentToolStatus::Error; + tool.message = Some("Sub-agent ended without a recorded outcome for this tool; effects are unknown.".into()); + } + } } fn set_activity(&self, activity: SubAgentActivity) { @@ -735,6 +740,8 @@ impl SubAgentUiAdapter { let mut output = self.output.lock().unwrap(); output.response = Some(response); output.activity = Some(SubAgentActivity::Completed); + output.error = None; + output.cancelled = None; } fn set_usage(&self, usage: SubAgentUsage) { diff --git a/crates/code_assistant_core/src/agent/sub_agent/run.rs b/crates/code_assistant_core/src/agent/sub_agent/run.rs new file mode 100644 index 00000000..56c628bb --- /dev/null +++ b/crates/code_assistant_core/src/agent/sub_agent/run.rs @@ -0,0 +1,123 @@ +use super::*; +use futures::FutureExt; + +#[async_trait::async_trait] +impl SubAgentRunner for DefaultSubAgentRunner { + async fn run( + &self, + parent_tool_id: &str, + instructions: String, + mode: SubAgentMode, + require_file_references: bool, + ) -> Result { + let registration = self.cancellation_registry.register_run(parent_tool_id); + let child = registration.cancellation.clone(); + let sub_ui = self.build_sub_agent_ui( + self.ui.clone(), + parent_tool_id.to_string(), + child.flag.clone(), + ); + let work = async { + self.parent_cancellation.check()?; + child.token.check()?; + let mut agent = tokio::select! { + biased; + _ = self.parent_cancellation.cancelled() => return Err(tools_core::Cancelled.into()), + _ = child.token.cancelled() => return Err(tools_core::Cancelled.into()), + agent = self.build_agent(parent_tool_id, sub_ui.clone(), self.permission_handler.clone()) => agent?, + }; + agent.set_cancellation(child.token.clone()); + let scope = match mode { + SubAgentMode::ReadOnly => ToolScope::SubAgentReadOnly, + SubAgentMode::Default if self.session_config.use_diff_blocks => { + ToolScope::SubAgentDefaultWithDiffBlocks + } + SubAgentMode::Default => ToolScope::SubAgentDefault, + }; + agent.set_tool_scope(scope); + agent.append_message(Message::new_user(instructions))?; + let mut answer = String::new(); + for attempt in 0..=2 { + child.token.check()?; + let iteration = agent.run_single_iteration().await; + // Earlier requests and tools can have completed before a later + // request fails. Preserve their usage as well as their tool list. + sub_ui.set_usage(compute_sub_agent_usage( + agent.message_history(), + &self.model_name, + )); + iteration?; + child.token.check()?; + answer = extract_last_assistant_text(agent.message_history()).unwrap_or_default(); + if !require_file_references || has_file_references_with_line_ranges(&answer) { + break; + } + if attempt == 2 { + answer.push_str("\n\n(Warning: requested file references with line ranges, but the sub-agent did not include them.)"); + break; + } + sub_ui.send_output_update().await; + agent.append_message(Message::new_user( + "Please revise your last answer to include exact file references with line ranges (e.g. `path/to/file.rs:10-20`).", + ))?; + } + Ok::<_, anyhow::Error>(answer) + }; + + // Propagate a parent stop without dropping an already executing child + // tool. The child's runtime stops new calls and wakes provider/permission + // waits; existing side effects finish or cooperate with cancellation. + let propagate_stop = async { + self.parent_cancellation.cancelled().await; + child.cancel(); + std::future::pending::<()>().await; + }; + let result = tokio::select! { + biased; + _ = propagate_stop => unreachable!("stop propagation never completes"), + result = std::panic::AssertUnwindSafe(work).catch_unwind() => { + result.unwrap_or_else(|_| Err(anyhow::anyhow!("Sub-agent panicked; partial work may have occurred"))) + } + }; + // Drop handles registration cleanup on every exit, including abort/panic. + // Explicitly unregister before publishing the terminal child status. + drop(registration); + match result { + Ok(answer) + if !child.token.is_cancelled() && !self.parent_cancellation.is_cancelled() => + { + sub_ui.set_response(answer.clone()); + sub_ui.send_output_update().await; + Ok(SubAgentResult { + answer, + ui_output: sub_ui.get_final_output(), + }) + } + result => { + let cancelled = + child.token.is_cancelled() || self.parent_cancellation.is_cancelled(); + let message = if cancelled { + "Sub-agent cancelled. Partial work may have occurred; verify side effects before restarting.".to_string() + } else { + format!( + "{:#}. Partial work may have occurred; verify side effects before restarting.", + result.unwrap_err() + ) + }; + if cancelled { + sub_ui.set_cancelled(); + } + sub_ui.set_error(message.clone()); + if cancelled { + sub_ui.set_activity(SubAgentActivity::Cancelled); + } + sub_ui.send_output_update().await; + Err(SubAgentFailure { + message, + ui_output: sub_ui.get_final_output(), + } + .into()) + } + } + } +} diff --git a/crates/code_assistant_core/src/agent/sub_agent/tests.rs b/crates/code_assistant_core/src/agent/sub_agent/tests.rs new file mode 100644 index 00000000..d39e14bd --- /dev/null +++ b/crates/code_assistant_core/src/agent/sub_agent/tests.rs @@ -0,0 +1,364 @@ +use super::*; +use crate::mocks::{MockProjectManager, MockUI}; +use crate::session::service::LlmClientFactory; +use std::time::Duration; + +fn runner(factory: LlmClientFactory) -> (Arc, Arc) { + let ui = Arc::new(MockUI::default()); + let runner = DefaultSubAgentRunner::new( + "test-sub-agent".into(), + SessionConfig::default(), + Arc::default(), + Arc::default(), + ui.clone(), + None, + ToolPermissions::default(), + crate::tools::test_registry(), + None, + None, + ) + .with_llm_client_factory(factory) + .with_project_manager_factory(Arc::new(|| Box::new(MockProjectManager::new()))); + (Arc::new(runner), ui) +} +struct Fails; +#[async_trait::async_trait] +impl llm::LLMProvider for Fails { + async fn send_message( + &mut self, + _: llm::LLMRequest, + _: Option<&llm::StreamingCallback>, + ) -> Result { + anyhow::bail!("fatal provider failure") + } +} +struct Silent(Arc); +#[async_trait::async_trait] +impl llm::LLMProvider for Silent { + async fn send_message( + &mut self, + _: llm::LLMRequest, + _: Option<&llm::StreamingCallback>, + ) -> Result { + self.0.notify_one(); + std::future::pending().await + } +} + +#[tokio::test] +async fn sub_agent_failure_unregisters_and_publishes_terminal_output() { + let (runner, ui) = runner(Arc::new(|_| Ok(Box::new(Fails)))); + let result = runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await; + assert!(result.is_err()); + assert!( + !runner.cancellation_registry.cancel("child"), + "failed child remained registered as busy" + ); + let events = ui.events(); + let last = events.iter().rev().find_map(|event| match event { + UiEvent::UpdateToolStatus { status, output, .. } => Some((status, output)), + _ => None, + }); + let Some((ToolStatus::Error, Some(json))) = last else { + panic!("failed child did not publish terminal structured output") + }; + let output = SubAgentOutput::from_json(json).unwrap(); + assert_eq!(output.activity, Some(SubAgentActivity::Failed)); + assert!(output.error.unwrap().contains("fatal provider failure")); +} + +#[tokio::test] +async fn sub_agent_construction_failure_unregisters() { + let (runner, _) = runner(Arc::new(|_| anyhow::bail!("constructor failed"))); + assert!( + runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + .is_err() + ); + assert!(!runner.cancellation_registry.cancel("child")); +} + +#[tokio::test] +async fn sub_agent_drop_unregisters_even_without_a_return_value() { + let entered = Arc::new(tokio::sync::Notify::new()); + let (runner, _) = runner({ + let entered = entered.clone(); + Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) + }); + let task = tokio::spawn({ + let runner = runner.clone(); + async move { + runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + } + }); + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + task.abort(); + let _ = task.await; + assert!( + !runner.cancellation_registry.cancel("child"), + "aborted child leaked its registration" + ); +} + +#[tokio::test] +async fn sub_agent_cancel_wakes_a_provider_without_chunks() { + let entered = Arc::new(tokio::sync::Notify::new()); + let (runner, _) = runner({ + let entered = entered.clone(); + Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) + }); + let mut task = tokio::spawn({ + let runner = runner.clone(); + async move { + runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + } + }); + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + assert!(runner.cancellation_registry.cancel("child")); + let ended = tokio::time::timeout(Duration::from_millis(250), &mut task).await; + if ended.is_err() { + task.abort(); + let _ = task.await; + } + assert!( + matches!(ended, Ok(Ok(Err(_)))), + "child cancellation waited for another provider chunk" + ); + assert!(!runner.cancellation_registry.cancel("child")); +} + +#[tokio::test] +async fn sub_agent_parent_cancel_wakes_a_provider_without_chunks() { + let entered = Arc::new(tokio::sync::Notify::new()); + let (runner, _) = runner({ + let entered = entered.clone(); + Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) + }); + let parent = runner.parent_cancellation.clone(); + let mut task = tokio::spawn({ + let runner = runner.clone(); + async move { + runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + } + }); + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + parent.cancel(); + let ended = tokio::time::timeout(Duration::from_millis(250), &mut task).await; + if ended.is_err() { + task.abort(); + let _ = task.await; + } + assert!( + matches!(ended, Ok(Ok(Err(_)))), + "parent cancellation did not reach the child runtime" + ); + assert!(!runner.cancellation_registry.cancel("child")); +} + +#[tokio::test] +async fn sub_agent_tool_retains_structured_failure_output() { + use crate::tools::impls::spawn_agent::{SpawnAgentInput, SpawnAgentTool}; + use tools_core::Tool; + let (runner, ui) = runner(Arc::new(|_| Ok(Box::new(Fails)))); + let mut services = crate::tools::ToolServices { + project_manager: Arc::new(MockProjectManager::new()), + plan: None, + ui: Some(ui), + sub_agent_runner: Some(runner), + wakeups: None, + pty_sessions: None, + terminal_interrupts: None, + browser_sessions: None, + session_source: None, + }; + let executor = command_executor::DefaultCommandExecutor; + let mut context = tools_core::ToolContext { + command_executor: &executor, + tool_id: Some("child".into()), + session_id: None, + permission_handler: None, + extensions: Some(&mut services), + }; + let mut input = SpawnAgentInput { + instructions: "inspect".into(), + require_file_references: false, + mode: "read_only".into(), + }; + let output = SpawnAgentTool + .execute(&mut context, &mut input) + .await + .unwrap(); + assert!(output.error.is_some()); + let json = output + .ui_output + .expect("persist the structured failed child, not just its error text"); + assert_eq!( + SubAgentOutput::from_json(&json).unwrap().activity, + Some(SubAgentActivity::Failed) + ); +} + +#[tokio::test] +async fn sub_agent_keeps_completed_child_evidence_when_later_request_fails() { + use tools_core::{Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolSpec}; + #[derive(serde::Serialize, serde::Deserialize)] + struct Written; + impl ToolResult for Written { + fn is_success(&self) -> bool { + true + } + } + impl Render for Written { + fn status(&self) -> String { + "written".into() + } + fn render(&self, _: &mut ResourcesTracker) -> String { + "changed test resource".into() + } + } + struct Probe; + #[async_trait::async_trait] + impl Tool for Probe { + type Input = serde_json::Value; + type Output = Written; + fn spec(&self) -> ToolSpec { + ToolSpec { + name: "probe".into(), + description: "test".into(), + parameters_schema: serde_json::json!({"type":"object"}), + annotations: None, + capabilities: ToolSpec::capabilities(&[ToolScope::SubAgentReadOnly.tag()]), + multiline_params: &[], + hidden: false, + title_template: None, + } + } + async fn execute<'a>( + &self, + _: &mut ToolContext<'a>, + _: &mut Self::Input, + ) -> Result { + Ok(Written) + } + } + struct FailAfterTool(bool); + #[async_trait::async_trait] + impl llm::LLMProvider for FailAfterTool { + async fn send_message( + &mut self, + _: llm::LLMRequest, + callback: Option<&llm::StreamingCallback>, + ) -> Result { + if std::mem::replace(&mut self.0, true) { + anyhow::bail!("fatal provider failure after tool") + } + if let Some(callback) = callback { + callback(&llm::StreamingChunk::InputJson { + content: "{}".into(), + tool_name: Some("probe".into()), + tool_id: Some("inner".into()), + })?; + callback(&llm::StreamingChunk::StreamingComplete)?; + } + Ok(llm::LLMResponse { + content: vec![llm::ContentBlock::new_tool_use( + "inner", + "probe", + serde_json::json!({}), + )], + usage: llm::Usage { + input_tokens: 10, + output_tokens: 5, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + rate_limit_info: None, + }) + } + } + let (mut runner, _) = runner(Arc::new(|_| Ok(Box::new(FailAfterTool(false))))); + let mut registry = tools_core::ToolRegistry::new(); + registry.register(Box::new(Probe)); + Arc::get_mut(&mut runner).unwrap().tool_registry = Arc::new(registry); + let error = runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + .unwrap_err(); + let failure = error.downcast_ref::().unwrap(); + let output = SubAgentOutput::from_json(&failure.ui_output).unwrap(); + assert_eq!(output.activity, Some(SubAgentActivity::Failed)); + assert_eq!(output.tools.len(), 1); + assert_eq!(output.tools[0].name, "probe"); + assert_eq!(output.tools[0].status, SubAgentToolStatus::Success); + assert_eq!( + output + .usage + .expect("usage before failure must survive") + .output_tokens, + 5 + ); +} + +#[tokio::test] +async fn sub_agent_factory_panic_is_a_terminal_failure() { + let (runner, ui) = runner(Arc::new(|_| panic!("factory panicked"))); + let error = runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + .unwrap_err(); + assert!(error.downcast_ref::().is_some()); + assert!(!runner.cancellation_registry.cancel("child")); + assert!(ui.events().iter().any(|event| matches!( + event, + UiEvent::UpdateToolStatus { + status: ToolStatus::Error, + .. + } + ))); +} + +#[test] +fn sub_agent_success_clears_a_recovered_stream_error() { + let ui = SubAgentUiAdapter::new( + Arc::new(MockUI::default()), + "child".into(), + Arc::new(AtomicBool::new(false)), + crate::tools::test_registry(), + ); + ui.set_error("temporary stream error".into()); + ui.set_response("recovered".into()); + let output = SubAgentOutput::from_json(&ui.get_final_output()).unwrap(); + assert_eq!(output.activity, Some(SubAgentActivity::Completed)); + assert!( + output.error.is_none(), + "successful retry retained an error banner" + ); +} + +#[test] +fn sub_agent_old_registration_cannot_remove_its_replacement() { + let registry = SubAgentCancellationRegistry::default(); + let old = registry.register_run("child"); + let new = registry.register_run("child"); + assert!(old.cancellation.token.is_cancelled()); + drop(old); + assert!(registry.cancel("child")); + assert!(new.cancellation.token.is_cancelled()); + drop(new); + assert!(!registry.cancel("child")); +} diff --git a/crates/code_assistant_core/src/session/manager.rs b/crates/code_assistant_core/src/session/manager.rs index a1ca5b67..762ad336 100644 --- a/crates/code_assistant_core/src/session/manager.rs +++ b/crates/code_assistant_core/src/session/manager.rs @@ -1170,8 +1170,8 @@ impl SessionManager { permissions }; - let sub_agent_runner: Arc = - Arc::new(DefaultSubAgentRunner::new( + let sub_agent_runner: Arc = Arc::new( + DefaultSubAgentRunner::new( model_name_for_subagent, session_config.clone(), sandbox_context_clone, @@ -1182,7 +1182,9 @@ impl SessionManager { self.tool_registry.clone(), Some(Arc::new(self.persistence.clone())), self.hooks_factory.clone(), - )); + ) + .with_parent_cancellation(cancellation.clone()), + ); let components = AgentComponents { llm_provider, diff --git a/crates/code_assistant_core/src/tools/impls/spawn_agent.rs b/crates/code_assistant_core/src/tools/impls/spawn_agent.rs index 5aed07c6..45091923 100644 --- a/crates/code_assistant_core/src/tools/impls/spawn_agent.rs +++ b/crates/code_assistant_core/src/tools/impls/spawn_agent.rs @@ -181,11 +181,24 @@ impl Tool for SpawnAgentTool { error: None, ui_output: Some(sub_result.ui_output), }), - Err(e) => Ok(SpawnAgentOutput { - answer: String::new(), - error: Some(e.to_string()), - ui_output: None, - }), + Err(e) => { + let ui_output = match e.downcast_ref::() { + Some(failure) => failure.ui_output.clone(), + None => { + // Alternate runners may only return an error. Still + // replace the live card with a terminal structured result. + let mut output = crate::agent::sub_agent::SubAgentOutput::new(); + output.activity = Some(crate::agent::sub_agent::SubAgentActivity::Failed); + output.error = Some(e.to_string()); + output.to_json() + } + }; + Ok(SpawnAgentOutput { + answer: String::new(), + error: Some(e.to_string()), + ui_output: Some(ui_output), + }) + } } } } From bf62b86a8a7b6273b3614fbb5c0e2c292783e96d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 9 Sep 2026 15:06:34 +0200 Subject: [PATCH 06/15] test: synchronize session output instead of racing shell startup --- .../src/tools/impls/execute_command.rs | 35 +++++++++++++++---- crates/pty_session/src/session.rs | 30 ++++++++++++++-- 2 files changed, 55 insertions(+), 10 deletions(-) diff --git a/crates/code_assistant_core/src/tools/impls/execute_command.rs b/crates/code_assistant_core/src/tools/impls/execute_command.rs index 42960247..22d760b7 100644 --- a/crates/code_assistant_core/src/tools/impls/execute_command.rs +++ b/crates/code_assistant_core/src/tools/impls/execute_command.rs @@ -879,16 +879,19 @@ mod tests { let mut fixture = session_mode_fixture(dir.path()); let mut context = fixture.context(); - let mut input = session_mode_input("echo started; sleep 30", 500); - let result = ExecuteCommandTool.execute(&mut context, &mut input).await?; + // Do not assume the login shell prints within the first yield window. + // Gate output on stdin so a silent first response is exercised on every + // run, not just when a loaded CI runner starts the shell slowly. + let mut input = session_mode_input("read -r release; echo started; read -r finish", 500); + let result = tokio::time::timeout( + std::time::Duration::from_secs(10), + ExecuteCommandTool.execute(&mut context, &mut input), + ) + .await + .expect("session mode must yield without waiting for command output")?; assert!(result.running, "process should still be running"); assert!(result.success, "a running session is not a failure"); - assert!( - result.output.contains("started"), - "output: {}", - result.output - ); let session_id = result .pty_session_id .expect("session id for running process"); @@ -897,7 +900,25 @@ mod tests { drop(context); let manager = fixture.pty_sessions().unwrap(); let session = manager.get(session_id).expect("session should be tracked"); + session.write(b"go\n")?; + let mut output = result.output; + let observed = tokio::time::timeout(std::time::Duration::from_secs(10), async { + while !output.contains("started") { + let chunk = session + .collect_output(std::time::Duration::from_millis(100)) + .await; + output.push_str(&chunk.output); + if matches!(chunk.status, pty_session::PtySessionStatus::Exited(_)) { + break; + } + } + }) + .await; session.terminate(); + assert!( + observed.is_ok() && output.contains("started"), + "output should arrive after releasing the command: {output:?}" + ); Ok(()) } diff --git a/crates/pty_session/src/session.rs b/crates/pty_session/src/session.rs index ba66ad4f..bd845071 100644 --- a/crates/pty_session/src/session.rs +++ b/crates/pty_session/src/session.rs @@ -598,12 +598,36 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn long_running_command_yields_while_running() { - let session = PtySession::spawn(shell("echo started; sleep 30", true)).unwrap(); - let out = session.collect_output(Duration::from_millis(500)).await; - assert!(out.output.contains("started"), "output: {}", out.output); + // A yield window is not a shell-startup deadline. Gate the first output + // on stdin so the initial collect must work even with no output yet. + let session = + PtySession::spawn(shell("read -r release; echo started; read -r finish", true)) + .unwrap(); + let out = tokio::time::timeout( + Duration::from_secs(10), + session.collect_output(Duration::from_millis(500)), + ) + .await + .expect("collect must yield without waiting for command output"); assert_eq!(out.status, PtySessionStatus::Running); + session.write(b"go\n").unwrap(); + let mut output = out.output; + let observed = tokio::time::timeout(Duration::from_secs(10), async { + while !output.contains("started") { + let chunk = session.collect_output(Duration::from_millis(100)).await; + output.push_str(&chunk.output); + if matches!(chunk.status, PtySessionStatus::Exited(_)) { + break; + } + } + }) + .await; session.terminate(); + assert!( + observed.is_ok() && output.contains("started"), + "output should arrive after releasing the command: {output:?}" + ); let out = session.collect_output(Duration::from_secs(10)).await; assert!(matches!(out.status, PtySessionStatus::Exited(_))); } From 5f9abd9ca40f27cc3665fd381de03c0090834783 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:24:08 +0200 Subject: [PATCH 07/15] refactor(agent): interceptors return their output instead of writing the journal --- crates/agent_core/src/execution.rs | 9 ++- crates/agent_core/src/hooks.rs | 19 ++++--- crates/agent_core/src/runtime.rs | 11 ++-- crates/agent_core/src/runtime/tests.rs | 12 +++- .../src/runtime/tests/dispatch_tests.rs | 9 ++- .../agent_core/src/runtime/tool_execution.rs | 56 ++++--------------- .../src/agent/checkpoint_tests.rs | 12 +++- .../src/plugins/name_session.rs | 43 +++++++------- .../src/plugins/skill_snapshot.rs | 5 -- 9 files changed, 79 insertions(+), 97 deletions(-) diff --git a/crates/agent_core/src/execution.rs b/crates/agent_core/src/execution.rs index eea76958..23cef9e8 100644 --- a/crates/agent_core/src/execution.rs +++ b/crates/agent_core/src/execution.rs @@ -9,14 +9,18 @@ pub(crate) const RUNTIME_OUTPUT_CODEC: &str = "__agent_runtime_outcome_v1"; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum ExecutionState { + /// The call was rejected or cancelled before anything ran. NotStarted, /// Persisted BEFORE invoking a tool. After interruption we cannot tell /// whether its effects happened, including the save/invoke crash window. Started, - Succeeded, + /// The invocation itself failed (as opposed to a tool reporting an error + /// through its own output type). Failed, } +/// The loop's own outcome record for a call. Never a success: successful +/// tools journal their real output. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RuntimeToolOutput { pub state: ExecutionState, @@ -51,7 +55,6 @@ impl Render for RuntimeToolOutput { match self.state { ExecutionState::NotStarted => "Not started", ExecutionState::Started => "Outcome unknown", - ExecutionState::Succeeded => "Success", ExecutionState::Failed => "Error", } .into() @@ -64,6 +67,6 @@ impl Render for RuntimeToolOutput { impl ToolResult for RuntimeToolOutput { fn is_success(&self) -> bool { - self.state == ExecutionState::Succeeded + false } } diff --git a/crates/agent_core/src/hooks.rs b/crates/agent_core/src/hooks.rs index d8819f23..0c96a864 100644 --- a/crates/agent_core/src/hooks.rs +++ b/crates/agent_core/src/hooks.rs @@ -9,17 +9,16 @@ use crate::dialect::ToolDialect; use crate::tree::{ConversationPath, MessageNode, NodeId}; -use crate::types::{ToolExecution, ToolRequest}; +use crate::types::ToolRequest; use anyhow::Result; use llm::Message; use std::any::Any; use std::collections::BTreeMap; use std::time::Duration; -use tools_core::ToolRegistry; +use tools_core::{AnyOutput, ToolRegistry}; /// View of the agent state that hooks may read and act on. pub struct LoopCtx<'a> { - pub tool_executions: &'a mut Vec, pub message_nodes: &'a mut BTreeMap, pub active_path: &'a ConversationPath, /// The session this agent runs, `None` while no session is assigned yet. @@ -36,11 +35,15 @@ pub struct LoopCtx<'a> { /// Intercepts tool requests that the application handles itself instead of /// dispatching them to the registry, and observes successful executions. pub trait ToolInterceptor: Send + Sync { - /// Returns `Some(result)` when the request was handled here. Scope and - /// permission checks and the start checkpoint always precede this hook, - /// including for parallel batches. Intercepted tools do not appear in the UI. - /// Hooks may append a ToolExecution; otherwise a generic outcome is recorded. - fn try_intercept(&self, _request: &ToolRequest, _ctx: &mut LoopCtx) -> Option> { + /// Handles the request in the application instead of the registry and + /// returns the output the loop journals for it. Scope and permission + /// checks always precede this hook, for parallel batches as well. + /// Intercepted tools do not appear in the UI. + fn try_intercept( + &self, + _request: &ToolRequest, + _ctx: &mut LoopCtx, + ) -> Option>> { None } diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index e358c050..edc52a2b 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -618,7 +618,6 @@ impl AgentRuntime { self.conversation .with_nodes_mut(|message_nodes, active_path| { let ctx = LoopCtx { - tool_executions: &mut self.tool_executions, message_nodes, active_path, session_id: self.session_id.as_deref(), @@ -1324,14 +1323,15 @@ impl AgentRuntime { messages } - /// Executes a tool and catches all errors, returning them as Results /// Gives the registered interceptors a chance to handle the request - /// before the standard dispatch. Returns `Some(result)` when one did. - fn intercept_tool(&mut self, tool_request: &ToolRequest) -> Option> { + /// instead of the registry. Returns the output one of them produced. + fn intercept_tool( + &mut self, + tool_request: &ToolRequest, + ) -> Option>> { self.conversation .with_nodes_mut(|message_nodes, active_path| { let mut ctx = LoopCtx { - tool_executions: &mut self.tool_executions, message_nodes, active_path, session_id: self.session_id.as_deref(), @@ -1352,7 +1352,6 @@ impl AgentRuntime { self.conversation .with_nodes_mut(|message_nodes, active_path| { let mut ctx = LoopCtx { - tool_executions: &mut self.tool_executions, message_nodes, active_path, session_id: self.session_id.as_deref(), diff --git a/crates/agent_core/src/runtime/tests.rs b/crates/agent_core/src/runtime/tests.rs index 62b70a4f..0454967c 100644 --- a/crates/agent_core/src/runtime/tests.rs +++ b/crates/agent_core/src/runtime/tests.rs @@ -137,9 +137,15 @@ fn checkpoint_legacy_history_is_imported_only_without_a_tree() { fn checkpoint_hook_message_corrections_rebuild_cache_even_on_early_return() { struct Correction; impl ToolInterceptor for Correction { - fn try_intercept(&self, _: &ToolRequest, ctx: &mut LoopCtx) -> Option> { + fn try_intercept( + &self, + _: &ToolRequest, + ctx: &mut LoopCtx, + ) -> Option>> { ctx.message_nodes.get_mut(&1).unwrap().message = Message::new_user("corrected by hook"); - Some(Ok(true)) + Some(Ok(Box::new(crate::types::ParseError::new( + "handled".into(), + )))) } } let (mut agent, saved) = runtime(); @@ -149,7 +155,7 @@ fn checkpoint_hook_message_corrections_rebuild_cache_even_on_early_return() { agent .intercept_tool(&ToolRequest::from(&call("a"))) .unwrap() - .unwrap() + .is_ok() ); agent.save_state().unwrap(); let snapshot = saved.0.lock().unwrap().take().unwrap(); diff --git a/crates/agent_core/src/runtime/tests/dispatch_tests.rs b/crates/agent_core/src/runtime/tests/dispatch_tests.rs index 24a88676..f45b6b63 100644 --- a/crates/agent_core/src/runtime/tests/dispatch_tests.rs +++ b/crates/agent_core/src/runtime/tests/dispatch_tests.rs @@ -114,9 +114,14 @@ struct Observer { intercept: bool, } impl ToolInterceptor for Observer { - fn try_intercept(&self, _: &ToolRequest, _: &mut LoopCtx) -> Option> { + fn try_intercept( + &self, + _: &ToolRequest, + _: &mut LoopCtx, + ) -> Option>> { self.attempts.fetch_add(1, Ordering::SeqCst); - self.intercept.then_some(Ok(true)) + self.intercept + .then(|| Ok(Box::new(Output("intercepted".into())) as Box)) } fn after_tool_success(&self, request: &ToolRequest, _: &mut LoopCtx) { self.successes.lock().unwrap().push(request.clone()); diff --git a/crates/agent_core/src/runtime/tool_execution.rs b/crates/agent_core/src/runtime/tool_execution.rs index c7110b57..3d8de30e 100644 --- a/crates/agent_core/src/runtime/tool_execution.rs +++ b/crates/agent_core/src/runtime/tool_execution.rs @@ -2,7 +2,7 @@ //! which adjacent invocations may overlap; it does not change their hooks, //! permission checks, input correction or persistence semantics. use super::*; -use crate::execution::{ExecutionState, RuntimeToolOutput}; +use crate::execution::RuntimeToolOutput; use futures::{FutureExt, StreamExt, stream::FuturesUnordered}; use std::time::Instant; use tools_core::AnyOutput; @@ -237,36 +237,11 @@ impl AgentRuntime { // Interceptors execute on the state owner, even for a parallel group, // and only after scope/permission checks and the start checkpoint. if let Some(result) = self.intercept_tool(request) { - let output: Box = match result { - Ok(success) => { - let entry = self - .tool_executions - .iter() - .rev() - .find(|entry| entry.tool_request.id == request.id) - .expect("journaled invocation"); - let still_started = entry - .result - .as_any() - .and_then(|out| out.downcast_ref::()) - .is_some_and(|out| out.state == ExecutionState::Started); - if still_started { - Box::new(RuntimeToolOutput { - state: if success { - ExecutionState::Succeeded - } else { - ExecutionState::Failed - }, - message: "Handled by the application's tool interceptor.".into(), - }) - } else { - entry.result.try_clone()? - } - } - Err(error) => Box::new(RuntimeToolOutput::failed(Self::format_error_for_user( + let output = result.unwrap_or_else(|error| { + Box::new(RuntimeToolOutput::failed(Self::format_error_for_user( &error, - ))), - }; + ))) + }); return Ok(Err(Completion { original: request.clone(), execution: ToolExecution { @@ -313,25 +288,16 @@ impl AgentRuntime { })) } - /// Replace the journal slot, retaining deterministic request order. Legacy - /// interceptors may append their own record; collapse that duplicate by id. + /// Replace the journal slot of this call, keeping request order. fn store_execution(&mut self, execution: ToolExecution) { - let id = execution.tool_request.id.clone(); - if let Some(index) = self + let id = &execution.tool_request.id; + match self .tool_executions .iter() - .position(|entry| entry.tool_request.id == id) + .position(|entry| &entry.tool_request.id == id) { - self.tool_executions[index] = execution; - let mut first = true; - self.tool_executions.retain(|entry| { - if entry.tool_request.id != id { - return true; - } - std::mem::replace(&mut first, false) - }); - } else { - self.tool_executions.push(execution); + Some(index) => self.tool_executions[index] = execution, + None => self.tool_executions.push(execution), } } diff --git a/crates/code_assistant_core/src/agent/checkpoint_tests.rs b/crates/code_assistant_core/src/agent/checkpoint_tests.rs index 181780f3..66506155 100644 --- a/crates/code_assistant_core/src/agent/checkpoint_tests.rs +++ b/crates/code_assistant_core/src/agent/checkpoint_tests.rs @@ -173,8 +173,8 @@ fn journal_outcomes_survive_disk_reload_without_the_original_tools() -> Result<( ); let outputs = [ RuntimeToolOutput { - state: ExecutionState::Succeeded, - message: "completed before interruption".into(), + state: ExecutionState::Failed, + message: "failed before interruption".into(), }, RuntimeToolOutput::started(), RuntimeToolOutput::not_started("No invocation was made."), @@ -216,8 +216,14 @@ fn journal_outcomes_survive_disk_reload_without_the_original_tools() -> Result<( assert_eq!(entry.tool_request.name, "unavailable-tool"); assert_eq!(entry.tool_request.input["path"], "evidence.txt"); } - assert!(restored[0].result.is_success()); let mut tracker = tools_core::ResourcesTracker::new(); + assert!( + restored[0] + .result + .as_render() + .render(&mut tracker) + .contains("failed before interruption") + ); assert!( restored[1] .result diff --git a/crates/code_assistant_core/src/plugins/name_session.rs b/crates/code_assistant_core/src/plugins/name_session.rs index 12eaaac5..40924364 100644 --- a/crates/code_assistant_core/src/plugins/name_session.rs +++ b/crates/code_assistant_core/src/plugins/name_session.rs @@ -3,10 +3,11 @@ use crate::plugins::AgentAppState; use crate::tools::ToolRequest; +use crate::tools::impls::name_session::NameSessionOutput; use agent_core::hooks::{IterationHook, LoopCtx, ToolInterceptor}; -use agent_core::types::ToolExecution; use anyhow::Result; use llm::{ContentBlock, Message, MessageContent, MessageRole}; +use tools_core::AnyOutput; use tracing::{trace, warn}; /// Handles the `name_session` tool at the agent level: the title is session @@ -14,7 +15,11 @@ use tracing::{trace, warn}; pub struct NameSessionInterceptor; impl ToolInterceptor for NameSessionInterceptor { - fn try_intercept(&self, request: &ToolRequest, ctx: &mut LoopCtx) -> Option> { + fn try_intercept( + &self, + request: &ToolRequest, + ctx: &mut LoopCtx, + ) -> Option>> { if request.name != "name_session" { return None; } @@ -22,27 +27,21 @@ impl ToolInterceptor for NameSessionInterceptor { } } -fn apply_session_name(request: &ToolRequest, ctx: &mut LoopCtx) -> Result { - if let Some(title) = request.input["title"].as_str() { - let title = title.trim(); - if !title.is_empty() { - trace!("Obtained session title from LLM: {}", title); - AgentAppState::of(ctx.extensions).session_name = title.to_string(); +fn apply_session_name(request: &ToolRequest, ctx: &mut LoopCtx) -> Result> { + let title = request.input["title"] + .as_str() + .map(str::trim) + .filter(|title| !title.is_empty()); + let Some(title) = title else { + warn!("name_session was called without a usable title"); + return Err(anyhow::anyhow!("Invalid session title provided")); + }; - ctx.tool_executions.push(ToolExecution { - tool_request: request.clone(), - result: Box::new(crate::tools::impls::name_session::NameSessionOutput { - title: title.to_string(), - }), - }); - return Ok(true); - } else { - warn!("Title for name_session is empty after trimming"); - } - } else { - warn!("No 'title' field found in name_session input or it's not a string"); - } - Err(anyhow::anyhow!("Invalid session title provided")) + trace!("Obtained session title from LLM: {}", title); + AgentAppState::of(ctx.extensions).session_name = title.to_string(); + Ok(Box::new(NameSessionOutput { + title: title.to_string(), + })) } /// Appends a system reminder to the last actual user message while the diff --git a/crates/code_assistant_core/src/plugins/skill_snapshot.rs b/crates/code_assistant_core/src/plugins/skill_snapshot.rs index 6d04a81d..b7dc31d4 100644 --- a/crates/code_assistant_core/src/plugins/skill_snapshot.rs +++ b/crates/code_assistant_core/src/plugins/skill_snapshot.rs @@ -84,12 +84,10 @@ mod tests { }, ); let active_path = vec![1]; - let mut tool_executions = Vec::new(); let mut state = AgentAppState::new(SessionConfig::default()); { let mut ctx = LoopCtx { - tool_executions: &mut tool_executions, message_nodes: &mut message_nodes, active_path: &active_path, session_id: None, @@ -108,7 +106,6 @@ mod tests { // Re-activating the same skill does not duplicate it. { let mut ctx = LoopCtx { - tool_executions: &mut tool_executions, message_nodes: &mut message_nodes, active_path: &active_path, session_id: None, @@ -125,7 +122,6 @@ mod tests { let registry = crate::tools::test_registry(); let mut message_nodes = BTreeMap::new(); let active_path: Vec = Vec::new(); - let mut tool_executions = Vec::new(); let mut state = AgentAppState::new(SessionConfig::default()); let request = ToolRequest { @@ -137,7 +133,6 @@ mod tests { }; { let mut ctx = LoopCtx { - tool_executions: &mut tool_executions, message_nodes: &mut message_nodes, active_path: &active_path, session_id: None, From acce7c931736ada1c9492f13e7b30c215e430770 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:37:47 +0200 Subject: [PATCH 08/15] refactor(agent): checkpoint only the changes since the previous one The runtime hands its persistence the nodes and journal entries that changed since the last checkpoint instead of a full copy of the session. Conversation tracks edited nodes, ToolJournal tracks recorded entries, and the session store merges the delta by id. This removes the per-checkpoint JSON round-trip of every tool result, the dead linear-history copy, and the metadata-notifying persistence wrapper: the manager publishes the session metadata itself after the merge. --- crates/agent_core/src/execution.rs | 67 ++++- crates/agent_core/src/hooks.rs | 8 +- crates/agent_core/src/lib.rs | 6 +- crates/agent_core/src/persistence.rs | 47 +-- crates/agent_core/src/runtime.rs | 271 ++++++++---------- crates/agent_core/src/runtime/tests.rs | 231 ++++++++++----- .../src/runtime/tests/dispatch_tests.rs | 43 +-- .../agent_core/src/runtime/tool_execution.rs | 38 +-- crates/agent_core/src/tree.rs | 111 ++++--- .../src/agent/checkpoint_tests.rs | 85 +++--- .../src/agent/persistence.rs | 150 +++------- .../code_assistant_core/src/agent/runner.rs | 7 +- .../src/agent/sub_agent.rs | 10 +- .../src/agent/sub_agent/run.rs | 4 +- crates/code_assistant_core/src/agent/tests.rs | 44 +-- crates/code_assistant_core/src/persistence.rs | 29 ++ .../code_assistant_core/src/plugins/plan.rs | 19 +- .../src/plugins/skill_snapshot.rs | 32 +-- .../src/session/instance.rs | 35 +-- .../src/session/manager.rs | 203 ++++--------- crates/code_assistant_core/src/session/mod.rs | 68 ++--- 21 files changed, 732 insertions(+), 776 deletions(-) diff --git a/crates/agent_core/src/execution.rs b/crates/agent_core/src/execution.rs index 23cef9e8..63f8aec9 100644 --- a/crates/agent_core/src/execution.rs +++ b/crates/agent_core/src/execution.rs @@ -1,7 +1,11 @@ -//! Self-describing journal entries for calls without a concrete tool result. -//! Successful/functional-error tool outputs retain their existing codecs. +//! The tool journal: every call's outcome record in request order, and the +//! loop's own self-describing entries for calls without a concrete tool +//! result. Successful and functionally failed tool outputs keep the codec +//! of their tool. +use crate::types::ToolExecution; use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; use tools_core::{Render, ResourcesTracker, ToolResult}; pub(crate) const RUNTIME_OUTPUT_CODEC: &str = "__agent_runtime_outcome_v1"; @@ -70,3 +74,62 @@ impl ToolResult for RuntimeToolOutput { false } } + +/// The run's record of tool calls in request order, plus which entries +/// changed since the last checkpoint. +#[derive(Default)] +pub struct ToolJournal { + entries: Vec, + changed: BTreeSet, +} + +impl ToolJournal { + /// Restore the journal from persisted entries; nothing counts as changed. + pub fn restore(entries: Vec) -> Self { + Self { + entries, + changed: BTreeSet::new(), + } + } + + pub fn entries(&self) -> &[ToolExecution] { + &self.entries + } + + pub fn find(&self, id: &str) -> Option<&ToolExecution> { + self.entries + .iter() + .find(|entry| entry.tool_request.id == id) + } + + pub fn contains(&self, id: &str) -> bool { + self.find(id).is_some() + } + + /// Record the entry for a call, replacing an earlier entry with the same + /// id in place so request order is preserved. + pub fn record(&mut self, execution: ToolExecution) { + let id = execution.tool_request.id.clone(); + match self + .entries + .iter() + .position(|entry| entry.tool_request.id == id) + { + Some(index) => self.entries[index] = execution, + None => self.entries.push(execution), + } + self.changed.insert(id); + } + + /// Entries recorded or updated since the last checkpoint, in journal order. + pub fn changed(&self) -> impl Iterator + '_ { + self.entries + .iter() + .filter(|entry| self.changed.contains(&entry.tool_request.id)) + } + + /// Forget the change marks after a successful checkpoint. + pub fn mark_checkpointed(&mut self) { + self.changed.clear(); + } +} diff --git a/crates/agent_core/src/hooks.rs b/crates/agent_core/src/hooks.rs index 0c96a864..a3c99a36 100644 --- a/crates/agent_core/src/hooks.rs +++ b/crates/agent_core/src/hooks.rs @@ -8,19 +8,19 @@ //! same dyn-Any approach `ToolContext` uses. use crate::dialect::ToolDialect; -use crate::tree::{ConversationPath, MessageNode, NodeId}; +use crate::tree::Conversation; use crate::types::ToolRequest; use anyhow::Result; use llm::Message; use std::any::Any; -use std::collections::BTreeMap; use std::time::Duration; use tools_core::{AnyOutput, ToolRegistry}; /// View of the agent state that hooks may read and act on. pub struct LoopCtx<'a> { - pub message_nodes: &'a mut BTreeMap, - pub active_path: &'a ConversationPath, + /// The conversation tree. Edits through it are part of the next + /// checkpoint. + pub conversation: &'a mut Conversation, /// The session this agent runs, `None` while no session is assigned yet. /// Lets shared hook state (built once per process) be keyed per session — /// same role `PromptCtx::session_id` plays for system-prompt providers. diff --git a/crates/agent_core/src/lib.rs b/crates/agent_core/src/lib.rs index 682adbae..4461a304 100644 --- a/crates/agent_core/src/lib.rs +++ b/crates/agent_core/src/lib.rs @@ -3,7 +3,7 @@ //! Applications embed [`runtime::AgentRuntime`] and bring their own tools //! (via a `tools_core::ToolRegistry`), their own behavior plugins (the hook //! traits in [`hooks`]), their own UI adapter ([`ui::AgentUi`]), their own -//! persistence ([`persistence::SnapshotPersistence`]), and — optionally — +//! persistence ([`persistence::CheckpointPersistence`]), and — optionally — //! their own tool invocation format ([`dialect::ToolDialect`]; the built-in //! default is native tool calling, [`native::NativeDialect`]). //! @@ -21,9 +21,9 @@ pub mod types; pub mod ui; pub use dialect::ToolDialect; -pub use persistence::{AgentSnapshot, SnapshotPersistence}; +pub use persistence::{AgentCheckpoint, CheckpointPersistence}; pub use runtime::{AgentRuntime, AgentRuntimeComponents}; -pub use tree::{ConversationPath, MessageNode, NodeId}; +pub use tree::{Conversation, ConversationPath, MessageNode, NodeId}; pub use types::{ ParseError, PromptTooLongError, SerializedToolExecution, ToolExecution, ToolRequest, text_summary_from_blocks, to_tool_definition, to_tool_definitions, diff --git a/crates/agent_core/src/persistence.rs b/crates/agent_core/src/persistence.rs index e98bb3d0..01e11b44 100644 --- a/crates/agent_core/src/persistence.rs +++ b/crates/agent_core/src/persistence.rs @@ -1,31 +1,36 @@ -//! Core-shaped persistence: the loop saves what it owns — the conversation -//! tree, its derived linear history, the tool executions, and the id counters. -//! Prompt-only repairs and context-recovery projections are never checkpointed. -//! Application-level fields travel separately through the extension state -//! and are assembled into the application's storage format by its adapter. +//! Core-shaped persistence: after every change the loop hands its +//! persistence the delta since the previous checkpoint — the nodes and +//! journal entries that changed, plus the small always-current fields. +//! Prompt-only repairs and context-recovery projections are never part of +//! it. Application-level fields travel separately through the extension +//! state and are assembled into the application's storage format by its +//! adapter. -use crate::tree::{ConversationPath, MessageNode, NodeId}; +use crate::tree::{MessageNode, NodeId}; use crate::types::ToolExecution; use anyhow::Result; use std::any::Any; -use std::collections::BTreeMap; -/// What the agent loop itself knows about and persists. -pub struct AgentSnapshot { - pub session_id: Option, - pub message_nodes: BTreeMap, - pub active_path: ConversationPath, +/// What changed since the previous checkpoint of this run. +pub struct AgentCheckpoint<'a> { + pub session_id: &'a str, + /// Nodes appended or edited since the previous checkpoint. + pub changed_nodes: Vec<&'a MessageNode>, + pub active_path: &'a [NodeId], pub next_node_id: NodeId, - /// Canonical linear history derived from `active_path`, retained for API - /// compatibility. Never the rendered/repaired LLM prompt. A supplied tree - /// takes precedence over this field on restore. - pub messages: Vec, - pub tool_executions: Vec, + /// Journal entries recorded or updated since the previous checkpoint. + pub changed_executions: Vec<&'a ToolExecution>, pub next_request_id: u64, } -/// Persistence used by the agent loop: it saves the loop's snapshot, with -/// the application fields supplied by the extension state. -pub trait SnapshotPersistence: Send + Sync { - fn save(&mut self, snapshot: AgentSnapshot, extensions: &(dyn Any + Send)) -> Result<()>; +/// Persistence used by the agent loop. +pub trait CheckpointPersistence: Send + Sync { + /// Merge the checkpoint into the stored session. A call is atomic: on + /// `Err` nothing of it is stored, and the loop keeps the changes marked + /// for its next attempt. + fn commit( + &mut self, + checkpoint: &AgentCheckpoint<'_>, + extensions: &(dyn Any + Send), + ) -> Result<()>; } diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index edc52a2b..47d359f5 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -6,8 +6,9 @@ mod tests; mod tool_execution; use crate::dialect::ToolDialect; +use crate::execution::ToolJournal; use crate::hooks::{ContextSnapshot, HookRegistry, LoopCtx, RecoveryAction, ToolServicesProvider}; -use crate::persistence::{AgentSnapshot, SnapshotPersistence}; +use crate::persistence::{AgentCheckpoint, CheckpointPersistence}; use crate::tree::{Conversation, ConversationPath, MessageNode, NodeId}; use crate::types::{ToolExecution, ToolRequest, text_summary_from_blocks, to_tool_definitions}; use crate::ui::{AgentActivity, AgentUi, AgentUiEvent, DisplayFragment, HiddenTools, UIError}; @@ -49,7 +50,7 @@ pub struct AgentRuntimeComponents { pub permissions: ToolPermissions, /// Builds the application services handed to each tool invocation. pub services_provider: Arc, - pub state_persistence: Box, + pub state_persistence: Box, pub hooks: HookRegistry, /// Application-specific loop state, exposed to the hooks type-erased. /// `Sync` because the loop holds `&self` across awaits. @@ -85,7 +86,7 @@ pub struct AgentRuntime { stream_hidden_tools: HiddenTools, command_executor: Arc, ui: Arc, - state_persistence: Box, + state_persistence: Box, /// Builds the application services handed to each tool invocation. services_provider: Arc, @@ -97,8 +98,8 @@ pub struct AgentRuntime { /// Run-local LLM projection. Never included in a checkpoint. prompt_projection: PromptProjection, - // Store the history of tool executions - tool_executions: Vec, + /// Every tool call of the conversation with its outcome. + journal: ToolJournal, // Cached system prompts keyed by model hint cached_system_prompts: HashMap, // Optional model identifier used for prompt selection @@ -165,7 +166,7 @@ impl AgentRuntime { cancellation: tools_core::RunCancellation::default(), conversation: Conversation::default(), prompt_projection: PromptProjection::default(), - tool_executions: Vec::new(), + journal: ToolJournal::default(), cached_system_prompts: HashMap::new(), next_request_id: 1, // Start from 1 session_id: None, @@ -209,8 +210,9 @@ impl AgentRuntime { self.extensions.as_mut() } - /// Restore the conversation (tree, active path, id counter, linearized - /// history) from persisted state. + /// Restore the conversation (tree, active path, id counter) from + /// persisted state. `messages` is the legacy linear history, imported + /// only when the session has no tree yet. pub fn restore_conversation( &mut self, message_nodes: BTreeMap, @@ -225,7 +227,7 @@ impl AgentRuntime { /// Restore the tool execution records from persisted state. pub fn set_tool_executions(&mut self, tool_executions: Vec) { - self.tool_executions = tool_executions; + self.journal = ToolJournal::restore(tool_executions); } /// Restore the request id counter from persisted state. @@ -258,8 +260,8 @@ impl AgentRuntime { } } - /// Get a reference to the message history - pub fn message_history(&self) -> &[Message] { + /// The messages on the active path, in order. + pub fn message_history(&self) -> Vec { self.conversation.history() } @@ -287,29 +289,30 @@ impl AgentRuntime { self.ui.send_event(event).await } - /// Save the current state (message history and tool executions) - fn save_state(&mut self) -> Result<()> { - trace!( - "saving {} messages to persistence (tree nodes: {})", - self.conversation.history().len(), - self.conversation.nodes().len() - ); - - let snapshot = AgentSnapshot { - session_id: self.session_id.clone(), - message_nodes: self.conversation.nodes().clone(), - active_path: self.conversation.path().clone(), + /// Persist everything that changed since the previous checkpoint. An + /// agent without a session keeps nothing. + fn checkpoint(&mut self) -> Result<()> { + let Some(session_id) = self.session_id.as_deref() else { + return Ok(()); + }; + let checkpoint = AgentCheckpoint { + session_id, + changed_nodes: self.conversation.changed_nodes().collect(), + active_path: self.conversation.path(), next_node_id: self.conversation.next_id(), - messages: self.conversation.history().to_vec(), - tool_executions: self - .tool_executions - .iter() - .map(ToolExecution::try_clone) - .collect::>()?, + changed_executions: self.journal.changed().collect(), next_request_id: self.next_request_id, }; + trace!( + "checkpoint: {} changed node(s), {} changed execution(s)", + checkpoint.changed_nodes.len(), + checkpoint.changed_executions.len() + ); self.state_persistence - .save(snapshot, self.extensions.as_ref()) + .commit(&checkpoint, self.extensions.as_ref())?; + self.conversation.mark_checkpointed(); + self.journal.mark_checkpointed(); + Ok(()) } /// Pre-allocate the next node_id without creating a node. @@ -329,7 +332,7 @@ impl AgentRuntime { observer.on_message(self.session_id.as_deref(), &message); } - self.save_state()?; + self.checkpoint()?; Ok(()) } @@ -458,7 +461,7 @@ impl AgentRuntime { let flow = self.manage_tool_execution(&tool_requests).await?; // Save state after tool executions - self.save_state()?; + self.checkpoint()?; match flow { LoopFlow::Continue => { /* Continue to the next iteration */ } @@ -476,20 +479,18 @@ impl AgentRuntime { } } - /// Compatibility entry point. Restores no longer delete incomplete tool - /// calls: the tree/cache retain evidence, and prompt rendering supplies - /// missing outcomes without guessing that a user cancelled the operation. - pub fn normalize_loaded_message_history(&mut self) {} - + /// Replace the last assistant message with the parser's corrected + /// version (e.g. text after the tool call truncated). fn correct_last_assistant_response(&mut self, response: &llm::LLMResponse) -> Result<()> { - if let Some(id) = self.conversation.path().last().copied() { - self.conversation.edit_message(id, |message| { - if message.role == MessageRole::Assistant { - message.content = MessageContent::Structured(response.content.clone()); - message.usage = Some(response.usage.clone()); - } - }); - self.save_state()?; + let Some(id) = self.conversation.path().last().copied() else { + return Ok(()); + }; + if let Some(node) = self.conversation.node_mut(id) + && node.message.role == MessageRole::Assistant + { + node.message.content = MessageContent::Structured(response.content.clone()); + node.message.usage = Some(response.usage.clone()); + self.checkpoint()?; } Ok(()) } @@ -529,10 +530,10 @@ impl AgentRuntime { // Generate normal tool ID for consistency with UI expectations let tool_id = format!("tool-{request_counter}-1"); - // Create and store a ToolExecution for the parse error - let tool_execution = - ToolExecution::create_parse_error(tool_id.clone(), error_text.clone()); - self.tool_executions.push(tool_execution); + self.journal.record(ToolExecution::create_parse_error( + tool_id.clone(), + error_text.clone(), + )); Message::new_user_content(vec![ContentBlock::ToolResult { tool_use_id: tool_id, @@ -615,21 +616,17 @@ impl AgentRuntime { /// Runs the iteration hooks over the rendered messages right before they /// are sent to the LLM (e.g. to inject system reminders). pub fn shape_request_messages(&mut self, mut messages: Vec) -> Vec { - self.conversation - .with_nodes_mut(|message_nodes, active_path| { - let ctx = LoopCtx { - message_nodes, - active_path, - session_id: self.session_id.as_deref(), - registry: self.registry.as_ref(), - extensions: self.extensions.as_mut(), - }; - for hook in &self.hooks.iteration_hooks { - if let Err(e) = hook.shape_request(&mut messages, &ctx) { - warn!("Iteration hook failed to shape the request: {}", e); - } - } - }); + let ctx = LoopCtx { + conversation: &mut self.conversation, + session_id: self.session_id.as_deref(), + registry: self.registry.as_ref(), + extensions: self.extensions.as_mut(), + }; + for hook in &self.hooks.iteration_hooks { + if let Err(e) = hook.shape_request(&mut messages, &ctx) { + warn!("Iteration hook failed to shape the request: {}", e); + } + } messages } @@ -881,13 +878,15 @@ impl AgentRuntime { } } - fn active_messages(&self) -> &[Message] { - let history = self.conversation.history(); - let start = history + /// The active-path messages from the last compaction summary onwards. + fn active_messages(&self) -> Vec<&Message> { + let mut messages: Vec<&Message> = self.conversation.active_messages().collect(); + let start = messages .iter() .rposition(|message| message.is_compaction_summary) .unwrap_or(0); - &history[start..] + messages.drain(..start); + messages } fn prompt_messages(&self) -> Vec { @@ -1058,7 +1057,7 @@ impl AgentRuntime { // Render each current-turn tool output to measure its size let mut sizes: Vec<(usize, usize)> = Vec::new(); // (index, byte_size) let mut tracker = ResourcesTracker::new(); - for (i, exec) in self.tool_executions.iter().enumerate() { + for (i, exec) in self.journal.entries().iter().enumerate() { if !current_turn_ids.contains(&exec.tool_request.id) || self .prompt_projection @@ -1083,8 +1082,8 @@ impl AgentRuntime { if byte_size < MIN_REPLACE_THRESHOLD { break; } - let tool_name = self.tool_executions[idx].tool_request.name.clone(); - let tool_id = self.tool_executions[idx].tool_request.id.clone(); + let request = &self.journal.entries()[idx].tool_request; + let (tool_name, tool_id) = (request.name.clone(), request.id.clone()); warn!( "Replacing tool result for '{}' ({}KB) with prompt-too-long error", tool_name, @@ -1230,7 +1229,7 @@ impl AgentRuntime { ids }) .collect(); - for execution in self.tool_executions.iter().rev() { + for execution in self.journal.entries().iter().rev() { let id = &execution.tool_request.id; if !visible_ids.contains(id) || outputs.contains_key(id) { continue; @@ -1329,39 +1328,29 @@ impl AgentRuntime { &mut self, tool_request: &ToolRequest, ) -> Option>> { - self.conversation - .with_nodes_mut(|message_nodes, active_path| { - let mut ctx = LoopCtx { - message_nodes, - active_path, - session_id: self.session_id.as_deref(), - registry: self.registry.as_ref(), - extensions: self.extensions.as_mut(), - }; - for interceptor in &self.hooks.interceptors { - if let Some(result) = interceptor.try_intercept(tool_request, &mut ctx) { - return Some(result); - } - } - None - }) + let mut ctx = LoopCtx { + conversation: &mut self.conversation, + session_id: self.session_id.as_deref(), + registry: self.registry.as_ref(), + extensions: self.extensions.as_mut(), + }; + self.hooks + .interceptors + .iter() + .find_map(|interceptor| interceptor.try_intercept(tool_request, &mut ctx)) } /// Notifies the registered interceptors that a tool executed successfully. fn after_tool_success(&mut self, tool_request: &ToolRequest) { - self.conversation - .with_nodes_mut(|message_nodes, active_path| { - let mut ctx = LoopCtx { - message_nodes, - active_path, - session_id: self.session_id.as_deref(), - registry: self.registry.as_ref(), - extensions: self.extensions.as_mut(), - }; - for interceptor in &self.hooks.interceptors { - interceptor.after_tool_success(tool_request, &mut ctx); - } - }); + let mut ctx = LoopCtx { + conversation: &mut self.conversation, + session_id: self.session_id.as_deref(), + registry: self.registry.as_ref(), + extensions: self.extensions.as_mut(), + }; + for interceptor in &self.hooks.interceptors { + interceptor.after_tool_success(tool_request, &mut ctx); + } } async fn notify_tool_parameter_updates( @@ -1413,63 +1402,51 @@ impl AgentRuntime { ) -> Result<()> { let dialect = self.dialect.clone(); let registry = self.registry.clone(); - let Some(id) = self - .conversation - .path() - .iter() - .rev() - .find(|id| { - self.conversation - .nodes() - .get(id) - .is_some_and(|node| node.message.role == MessageRole::Assistant) - }) - .copied() - else { + let Some(node) = self.conversation.last_assistant_node_mut() else { return Ok(()); }; - let mut updated = false; - self.conversation.edit_message(id, |message| { - let request_id = message.request_id.unwrap_or(0); - match &mut message.content { - MessageContent::Structured(blocks) => { - for block in blocks.iter_mut() { - if let ContentBlock::ToolUse { - id, name, input, .. - } = block - && id == &updated_request.id - && name == &updated_request.name - { - *input = updated_request.input.clone(); - updated = true; - return; - } + let message = &mut node.message; + let request_id = message.request_id.unwrap_or(0); + let updated = match &mut message.content { + MessageContent::Structured(blocks) => { + let native_call = blocks.iter_mut().find_map(|block| match block { + ContentBlock::ToolUse { + id, name, input, .. + } if id == &updated_request.id && name == &updated_request.name => Some(input), + _ => None, + }); + match native_call { + Some(input) => { + *input = updated_request.input.clone(); + true } - if !dialect.uses_native_tools() { - updated = Self::update_tool_call_in_text_blocks( - blocks, - updated_request, - request_id, - dialect.as_ref(), - registry.as_ref(), - ); - } - } - MessageContent::Text(text) => { - if let Ok(replacement) = Self::update_tool_call_in_text_static( - text, + None if !dialect.uses_native_tools() => Self::update_tool_call_in_text_blocks( + blocks, updated_request, + request_id, dialect.as_ref(), registry.as_ref(), - ) { + ), + None => false, + } + } + MessageContent::Text(text) => { + match Self::update_tool_call_in_text_static( + text, + updated_request, + dialect.as_ref(), + registry.as_ref(), + ) { + Ok(replacement) => { *text = replacement; - updated = true; + true } + Err(_) => false, } } - }); + }; if updated { - self.save_state()?; + self.checkpoint()?; } else { warn!("Could not find tool call {} to update", updated_request.id); } diff --git a/crates/agent_core/src/runtime/tests.rs b/crates/agent_core/src/runtime/tests.rs index 0454967c..a6c44909 100644 --- a/crates/agent_core/src/runtime/tests.rs +++ b/crates/agent_core/src/runtime/tests.rs @@ -4,7 +4,9 @@ use super::*; use crate::hooks::*; use serde_json::json; +/// Inert implementation of every collaborator the runtime needs. struct Stub; + #[async_trait::async_trait] impl LLMProvider for Stub { async fn send_message( @@ -15,66 +17,124 @@ impl LLMProvider for Stub { anyhow::bail!("unexpected LLM call") } } + #[async_trait::async_trait] impl AgentUi for Stub { async fn send_event(&self, _: AgentUiEvent) -> Result<(), UIError> { Ok(()) } + fn display_fragment(&self, _: &DisplayFragment) -> Result<(), UIError> { Ok(()) } + fn should_streaming_continue(&self) -> bool { true } + fn notify_rate_limit(&self, _: u64) {} + fn clear_rate_limit(&self) {} } + impl ToolServicesProvider for Stub { fn begin(&self, _: &mut (dyn Any + Send), _: &str) -> Box { Box::new(()) } + fn end(&self, _: &mut (dyn Any + Send), _: Box) {} + fn detached(&self, _: &str) -> Box { Box::new(()) } } + impl ToolDispatchPolicy for Stub { fn parallel_indices(&self, _: &[ToolRequest]) -> Vec { vec![] } } + impl CompactionPolicy for Stub { fn context_limit(&self, _: &(dyn Any + Send)) -> Result> { Ok(None) } + fn should_compact(&self, _: &ContextSnapshot) -> bool { false } + fn compaction_prompt(&self) -> &str { "summarize" } } + impl RecoveryPolicy for Stub { fn classify(&self, _: &anyhow::Error, _: u32) -> RecoveryAction { RecoveryAction::Fail } } + impl SystemPromptProvider for Stub { fn build(&self, _: &PromptCtx) -> String { String::new() } } + +/// What a session store would hold after merging every checkpoint. #[derive(Clone, Default)] -struct Capture(Arc>>); -impl SnapshotPersistence for Capture { - fn save(&mut self, snapshot: AgentSnapshot, _: &(dyn Any + Send)) -> Result<()> { - *self.0.lock().unwrap() = Some(snapshot); +struct Saved { + nodes: BTreeMap, + active_path: ConversationPath, + next_node_id: NodeId, + executions: Vec, + /// Sizes of the most recent checkpoint: changed nodes, changed executions. + last_delta: (usize, usize), + commits: usize, +} + +/// Merges checkpoints the way a session store does. +#[derive(Clone, Default)] +struct Capture(Arc>); + +impl Capture { + fn saved(&self) -> Saved { + self.0.lock().unwrap().clone() + } +} + +impl CheckpointPersistence for Capture { + fn commit(&mut self, checkpoint: &AgentCheckpoint<'_>, _: &(dyn Any + Send)) -> Result<()> { + let mut saved = self.0.lock().unwrap(); + for node in &checkpoint.changed_nodes { + saved.nodes.insert(node.id, (*node).clone()); + } + saved.active_path = checkpoint.active_path.to_vec(); + saved.next_node_id = checkpoint.next_node_id; + for execution in &checkpoint.changed_executions { + let execution = execution.try_clone()?; + let id = &execution.tool_request.id; + match saved + .executions + .iter() + .position(|entry| &entry.tool_request.id == id) + { + Some(index) => saved.executions[index] = execution, + None => saved.executions.push(execution), + } + } + saved.last_delta = ( + checkpoint.changed_nodes.len(), + checkpoint.changed_executions.len(), + ); + saved.commits += 1; Ok(()) } } + fn runtime() -> (AgentRuntime, Capture) { let capture = Capture::default(); - let runtime = AgentRuntime::new(AgentRuntimeComponents { + let mut runtime = AgentRuntime::new(AgentRuntimeComponents { llm_provider: Box::new(Stub), dialect: Arc::new(crate::native::NativeDialect), ui: Arc::new(Stub), @@ -98,11 +158,24 @@ fn runtime() -> (AgentRuntime, Capture) { }, extensions: Box::new(()), }); + runtime.set_session_id(Some("test-session".into())); (runtime, capture) } + +/// A runtime restored from what the store holds, through the serialized +/// tree rather than an in-memory alias of its messages. +fn reload(saved: Saved) -> AgentRuntime { + let (mut restored, _) = runtime(); + let nodes = serde_json::from_value(serde_json::to_value(saved.nodes).unwrap()).unwrap(); + restored.restore_conversation(nodes, saved.active_path, saved.next_node_id, Vec::new()); + restored.set_tool_executions(saved.executions); + restored +} + fn call(id: &str) -> ContentBlock { ContentBlock::new_tool_use(id, "write_file", json!({"content": "unformatted"})) } + fn result(id: &str) -> ContentBlock { ContentBlock::ToolResult { tool_use_id: id.into(), @@ -112,6 +185,61 @@ fn result(id: &str) -> ContentBlock { end_time: None, } } + +fn text(message: &Message) -> &str { + match &message.content { + MessageContent::Text(text) => text, + MessageContent::Structured(_) => panic!("expected a text message"), + } +} + +#[test] +fn checkpoint_carries_only_the_changes_since_the_previous_one() { + let (mut agent, saved) = runtime(); + agent.append_message(Message::new_user("one")).unwrap(); + agent.append_message(Message::new_assistant("two")).unwrap(); + assert_eq!(saved.saved().last_delta, (1, 0)); + + agent + .journal + .record(ToolExecution::create_parse_error("a".into(), "x".into())); + agent.checkpoint().unwrap(); + assert_eq!(saved.saved().last_delta, (0, 1)); + + agent.checkpoint().unwrap(); + let state = saved.saved(); + assert_eq!(state.last_delta, (0, 0)); + assert_eq!(state.nodes.len(), 2); + assert_eq!(state.executions.len(), 1); +} + +#[test] +fn checkpoint_failure_keeps_the_changes_for_the_next_attempt() { + struct FailOnce(Capture, bool); + + impl CheckpointPersistence for FailOnce { + fn commit( + &mut self, + checkpoint: &AgentCheckpoint<'_>, + extensions: &(dyn Any + Send), + ) -> Result<()> { + if std::mem::replace(&mut self.1, false) { + anyhow::bail!("disk full"); + } + self.0.commit(checkpoint, extensions) + } + } + + let (mut agent, saved) = runtime(); + agent.state_persistence = Box::new(FailOnce(saved.clone(), true)); + assert!(agent.append_message(Message::new_user("one")).is_err()); + agent.append_message(Message::new_assistant("two")).unwrap(); + let state = saved.saved(); + assert_eq!(state.commits, 1); + assert_eq!(state.last_delta, (2, 0)); + assert_eq!(state.nodes.len(), 2); +} + #[test] fn checkpoint_legacy_history_is_imported_only_without_a_tree() { let (mut agent, saved) = runtime(); @@ -122,32 +250,35 @@ fn checkpoint_legacy_history_is_imported_only_without_a_tree() { vec![Message::new_user("legacy")], ); agent.append_message(Message::new_assistant("new")).unwrap(); - let mut snapshot = saved.0.lock().unwrap().take().unwrap(); - assert_eq!(snapshot.message_nodes.len(), 2); - assert_eq!(snapshot.message_nodes[&2].parent_id, Some(1)); + let mut state = saved.saved(); + assert_eq!(state.nodes.len(), 2); + assert_eq!(state.nodes[&2].parent_id, Some(1)); + // A nonempty tree with an intentionally empty active path is authoritative // too: neither reactivate a branch nor import stale linear messages. - snapshot.active_path.clear(); - let restored = reload(snapshot); + state.active_path.clear(); + let restored = reload(state); assert!(restored.message_history().is_empty()); assert_eq!(restored.conversation.nodes().len(), 2); } #[test] -fn checkpoint_hook_message_corrections_rebuild_cache_even_on_early_return() { +fn checkpoint_persists_hook_edits_to_the_tree() { struct Correction; + impl ToolInterceptor for Correction { fn try_intercept( &self, _: &ToolRequest, ctx: &mut LoopCtx, ) -> Option>> { - ctx.message_nodes.get_mut(&1).unwrap().message = Message::new_user("corrected by hook"); + ctx.conversation.node_mut(1).unwrap().message = Message::new_user("corrected by hook"); Some(Ok(Box::new(crate::types::ParseError::new( "handled".into(), )))) } } + let (mut agent, saved) = runtime(); agent.append_message(Message::new_user("before")).unwrap(); agent.hooks.interceptors.push(Box::new(Correction)); @@ -157,45 +288,11 @@ fn checkpoint_hook_message_corrections_rebuild_cache_even_on_early_return() { .unwrap() .is_ok() ); - agent.save_state().unwrap(); - let snapshot = saved.0.lock().unwrap().take().unwrap(); - assert_eq!( - serde_json::to_value(&snapshot.message_nodes[&1].message).unwrap(), - serde_json::to_value(&snapshot.messages[0]).unwrap() - ); - assert!( - matches!(&snapshot.messages[0].content, MessageContent::Text(text) if text == "corrected by hook") - ); -} - -fn reload(snapshot: AgentSnapshot) -> AgentRuntime { - let (mut restored, _) = runtime(); - // Exercise the serialized tree, not an in-memory alias of its messages. - let nodes = - serde_json::from_value(serde_json::to_value(snapshot.message_nodes).unwrap()).unwrap(); - restored.restore_conversation( - nodes, - snapshot.active_path, - snapshot.next_node_id, - snapshot.messages, - ); - restored.set_tool_executions(snapshot.tool_executions); - restored.normalize_loaded_message_history(); - restored -} - -#[test] -fn checkpoint_tree_wins_over_stale_linear_history() { - let (mut agent, saved) = runtime(); - agent - .append_message(Message::new_user("canonical")) - .unwrap(); - let mut snapshot = saved.0.lock().unwrap().take().unwrap(); - snapshot.messages = vec![Message::new_user("stale")]; - let restored = reload(snapshot); - assert!( - matches!(&restored.message_history()[0].content, MessageContent::Text(text) if text == "canonical") - ); + agent.checkpoint().unwrap(); + let state = saved.saved(); + assert_eq!(state.last_delta, (1, 0)); + assert_eq!(text(&state.nodes[&1].message), "corrected by hook"); + assert_eq!(text(&agent.message_history()[0]), "corrected by hook"); } #[test] @@ -212,12 +309,8 @@ fn checkpoint_formatted_input_survives_roundtrip() { agent .append_message(Message::new_user_content(vec![result("a")])) .unwrap(); - let snapshot = saved.0.lock().unwrap().take().unwrap(); - assert_eq!( - serde_json::to_value(&snapshot.message_nodes[&1].message.content).unwrap(), - serde_json::to_value(&snapshot.messages[0].content).unwrap() - ); - let restored = reload(snapshot); + + let restored = reload(saved.saved()); assert!(matches!(&restored.message_history()[0].content, MessageContent::Structured(blocks) if matches!(&blocks[0], ContentBlock::ToolUse { input, .. } if input == &request.input))); } @@ -232,7 +325,8 @@ fn checkpoint_dangling_calls_survive_reload_with_unknown_prompt_outcome() { agent .append_message(Message::new_user_content(vec![result("a")])) .unwrap(); - let mut restored = reload(saved.0.lock().unwrap().take().unwrap()); + + let restored = reload(saved.saved()); let before = serde_json::to_value(restored.message_history()).unwrap(); let prompt = restored.render_tool_results_in_messages(); let MessageContent::Structured(blocks) = &prompt[1].content else { @@ -245,7 +339,6 @@ fn checkpoint_dangling_calls_survive_reload_with_unknown_prompt_outcome() { ); assert!(blocks.iter().any(|block| matches!(block, ContentBlock::ToolResult { tool_use_id, content, .. } if tool_use_id == "b" && content.contains("unknown") && !content.contains("cancelled by user")))); - restored.normalize_loaded_message_history(); assert_eq!( before, serde_json::to_value(restored.message_history()).unwrap() @@ -259,7 +352,8 @@ fn checkpoint_dangling_tail_is_not_deleted() { agent .append_message(Message::new_assistant_content(vec![call("a")])) .unwrap(); - let restored = reload(saved.0.lock().unwrap().take().unwrap()); + + let restored = reload(saved.saved()); assert_eq!(restored.message_history().len(), 2); assert_eq!(restored.render_tool_results_in_messages().len(), 3); } @@ -274,12 +368,14 @@ fn checkpoint_recovery_keeps_canonical_messages_and_tool_evidence() { .append_message(Message::new_user_content(vec![result("a")])) .unwrap(); // A serializable large output suffices to exercise size-based recovery. - agent.set_tool_executions(vec![ToolExecution::create_parse_error( + agent.journal.record(ToolExecution::create_parse_error( "a".into(), "x".repeat(60 * 1024), - )]); + )); + agent.checkpoint().unwrap(); let before = serde_json::to_value(agent.message_history()).unwrap(); - let evidence = agent.tool_executions[0].serialize().unwrap(); + let evidence = agent.journal.entries()[0].serialize().unwrap(); + assert_eq!(agent.replace_large_tool_results().len(), 1); let projected = agent.render_tool_results_in_messages(); assert!(serde_json::to_string(&projected).unwrap().len() < 10 * 1024); @@ -293,17 +389,18 @@ fn checkpoint_recovery_keeps_canonical_messages_and_tool_evidence() { ); assert_eq!( serde_json::to_value(&evidence).unwrap(), - serde_json::to_value(agent.tool_executions[0].serialize().unwrap()).unwrap() + serde_json::to_value(agent.journal.entries()[0].serialize().unwrap()).unwrap() ); + agent.drop_last_tool_exchange(); assert!(agent.render_tool_results_in_messages().is_empty()); - agent.save_state().unwrap(); - let restored = reload(saved.0.lock().unwrap().take().unwrap()); + agent.checkpoint().unwrap(); + let restored = reload(saved.saved()); assert_eq!( before, serde_json::to_value(restored.message_history()).unwrap() ); - assert_eq!(restored.tool_executions.len(), 1); + assert_eq!(restored.journal.entries().len(), 1); assert!( serde_json::to_string(&restored.render_tool_results_in_messages()) .unwrap() diff --git a/crates/agent_core/src/runtime/tests/dispatch_tests.rs b/crates/agent_core/src/runtime/tests/dispatch_tests.rs index f45b6b63..16e3144c 100644 --- a/crates/agent_core/src/runtime/tests/dispatch_tests.rs +++ b/crates/agent_core/src/runtime/tests/dispatch_tests.rs @@ -175,16 +175,15 @@ async fn dispatch_parallel_hooks_and_formatted_inputs_match_sequential_contract( let successes = successes.lock().unwrap(); assert_eq!(successes.len(), 2); assert!(successes.iter().all(|r| r.input["formatted"] == true)); - let saved = f.saved.0.lock().unwrap(); - let snapshot = saved.as_ref().unwrap(); + let saved = f.saved.saved(); assert!( - snapshot - .tool_executions + saved + .executions .iter() .all(|e| e.tool_request.input["formatted"] == true) ); assert!( - matches!(&snapshot.messages[0].content, MessageContent::Structured(blocks) + matches!(&saved.nodes[&1].message.content, MessageContent::Structured(blocks) if blocks.iter().all(|b| matches!(b, ContentBlock::ToolUse { input, .. } if input["formatted"] == true))) ); } @@ -202,12 +201,8 @@ async fn completion_is_checkpointed_while_sibling_waits(parallel: bool) { let checkpointed = tokio::time::timeout(Duration::from_millis(250), async { loop { if f.saved - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .tool_executions + .saved() + .executions .iter() .any(|e| e.tool_request.id == "one" && e.result.is_success()) { @@ -243,15 +238,7 @@ async fn dispatch_journal_distinguishes_unstarted_from_uncertain_after_reload() tokio::time::timeout(Duration::from_secs(2), f.entered.notified()) .await .unwrap(); - let journal = f - .saved - .0 - .lock() - .unwrap() - .as_ref() - .unwrap() - .tool_executions - .clone(); + let journal = f.saved.saved().executions; f.release.notify_one(); task.await.unwrap().unwrap(); assert_eq!( @@ -312,12 +299,10 @@ async fn dispatch_ui_failure_does_not_erase_successful_tool_evidence() { let mut f = fixture(&requests); f.agent.ui = Arc::new(FailingUi); let _ = f.agent.manage_tool_execution(&requests).await; - let snapshot = f.saved.0.lock().unwrap(); assert!( - snapshot - .as_ref() - .unwrap() - .tool_executions + f.saved + .saved() + .executions .iter() .any(|e| e.tool_request.id == "one" && e.result.is_success()) ); @@ -337,11 +322,11 @@ async fn dispatch_parallel_groups_do_not_cross_sequential_barriers() { } struct FailCompletionSave; -impl SnapshotPersistence for FailCompletionSave { - fn save(&mut self, snapshot: AgentSnapshot, _: &(dyn Any + Send)) -> Result<()> { +impl CheckpointPersistence for FailCompletionSave { + fn commit(&mut self, checkpoint: &AgentCheckpoint<'_>, _: &(dyn Any + Send)) -> Result<()> { anyhow::ensure!( - !snapshot - .tool_executions + !checkpoint + .changed_executions .iter() .any(|e| e.result.is_success()), "disk full" diff --git a/crates/agent_core/src/runtime/tool_execution.rs b/crates/agent_core/src/runtime/tool_execution.rs index 3d8de30e..bb67e5d1 100644 --- a/crates/agent_core/src/runtime/tool_execution.rs +++ b/crates/agent_core/src/runtime/tool_execution.rs @@ -111,10 +111,7 @@ impl AgentRuntime { request.id ); anyhow::ensure!( - !self - .tool_executions - .iter() - .any(|e| e.tool_request.id == request.id), + !self.journal.contains(&request.id), "Tool call id has already been recorded: {}", request.id ); @@ -122,12 +119,12 @@ impl AgentRuntime { // Write intent for the whole batch before any tool can have effects. // A restart can distinguish an unstarted sibling from an uncertain call. for request in requests { - self.store_execution(ToolExecution { + self.journal.record(ToolExecution { tool_request: request.clone(), result: Box::new(RuntimeToolOutput::not_started("No invocation was made.")), }); } - self.save_state()?; + self.checkpoint()?; let mut parallel = vec![false; requests.len()]; for index in self.hooks.dispatch.parallel_indices(requests) { @@ -228,11 +225,11 @@ impl AgentRuntime { RuntimeToolOutput::not_started(Self::format_error_for_user(&error)), ))); } - self.store_execution(ToolExecution { + self.journal.record(ToolExecution { tool_request: request.clone(), result: Box::new(RuntimeToolOutput::started()), }); - self.save_state()?; + self.checkpoint()?; // Interceptors execute on the state owner, even for a parallel group, // and only after scope/permission checks and the start checkpoint. @@ -288,19 +285,6 @@ impl AgentRuntime { })) } - /// Replace the journal slot of this call, keeping request order. - fn store_execution(&mut self, execution: ToolExecution) { - let id = &execution.tool_request.id; - match self - .tool_executions - .iter() - .position(|entry| &entry.tool_request.id == id) - { - Some(index) => self.tool_executions[index] = execution, - None => self.tool_executions.push(execution), - } - } - async fn commit_completion(&mut self, mut completed: Completion) -> Result { if let Some(services) = completed.services.take() { self.services_provider @@ -313,22 +297,18 @@ impl AgentRuntime { || self .registry .is_tool_hidden(&request.name, &self.tool_capability); - self.store_execution(completed.execution); + self.journal.record(completed.execution); // Commit evidence before hooks/rendering/UI can fail. The active call's // Started record remains on disk if this commit itself fails. - self.save_state()?; + self.checkpoint()?; if changed { self.update_message_history_with_formatted_tool(&request)?; } if success { self.after_tool_success(&request); - self.save_state()?; + self.checkpoint()?; } - let execution = self - .tool_executions - .iter() - .find(|entry| entry.tool_request.id == request.id) - .expect("committed outcome"); + let execution = self.journal.find(&request.id).expect("committed outcome"); let content = execution .result .as_any() diff --git a/crates/agent_core/src/tree.rs b/crates/agent_core/src/tree.rs index be97da9d..e9e3c8ab 100644 --- a/crates/agent_core/src/tree.rs +++ b/crates/agent_core/src/tree.rs @@ -3,6 +3,7 @@ use llm::Message; use serde::{Deserialize, Serialize}; +use std::collections::{BTreeMap, BTreeSet}; use std::time::SystemTime; /// Unique identifier for a message node within a session @@ -36,25 +37,31 @@ pub struct MessageNode { )] pub extension: Option, } -/// Runtime-owned conversation. Only the tree is writable; the linear history -/// is a derived cache, never a prompt recovery workspace or restore authority. -/// Kept crate-private so the persisted/public tree representation stays stable. -pub(crate) struct Conversation { - nodes: std::collections::BTreeMap, + +/// The conversation a running agent owns: the tree, the active path, the id +/// counter, and which nodes changed since the last checkpoint. +/// +/// Only the tree is written to. The linear history is derived from the +/// active path whenever it is needed, so it can never disagree with the tree. +pub struct Conversation { + nodes: BTreeMap, path: ConversationPath, next_id: NodeId, - history: Vec, + changed: BTreeSet, } impl Default for Conversation { fn default() -> Self { - Self::restore(Default::default(), Vec::new(), 1, Vec::new()) + Self::restore(BTreeMap::new(), Vec::new(), 1, Vec::new()) } } impl Conversation { - pub(crate) fn restore( - nodes: std::collections::BTreeMap, + /// Rebuild the conversation from persisted state. A session without a + /// tree imports its legacy linear messages as a single branch; those + /// imported nodes count as changed so the next checkpoint persists them. + pub fn restore( + nodes: BTreeMap, path: ConversationPath, next_id: NodeId, legacy_messages: Vec, @@ -63,7 +70,7 @@ impl Conversation { next_id: next_id.max(nodes.keys().next_back().copied().unwrap_or(0) + 1), nodes, path, - history: Vec::new(), + changed: BTreeSet::new(), }; if conversation.nodes.is_empty() { conversation.path.clear(); @@ -72,33 +79,68 @@ impl Conversation { conversation.append(message, id); } } - conversation.rebuild_history(); conversation } - pub(crate) fn nodes(&self) -> &std::collections::BTreeMap { + pub fn nodes(&self) -> &BTreeMap { &self.nodes } - pub(crate) fn path(&self) -> &ConversationPath { + pub fn path(&self) -> &[NodeId] { &self.path } - pub(crate) fn next_id(&self) -> NodeId { + pub fn next_id(&self) -> NodeId { self.next_id } - pub(crate) fn history(&self) -> &[Message] { - &self.history + pub fn node(&self, id: NodeId) -> Option<&MessageNode> { + self.nodes.get(&id) + } + + /// Mutable access to a node. The node counts as changed for the next + /// checkpoint; callers are responsible for keeping parent links intact. + pub fn node_mut(&mut self, id: NodeId) -> Option<&mut MessageNode> { + let node = self.nodes.get_mut(&id)?; + self.changed.insert(id); + Some(node) + } + + /// The most recent assistant message on the active path. + pub fn last_assistant_id(&self) -> Option { + self.path.iter().rev().copied().find(|id| { + self.nodes + .get(id) + .is_some_and(|node| node.message.role == llm::MessageRole::Assistant) + }) + } + + pub fn last_assistant_node_mut(&mut self) -> Option<&mut MessageNode> { + let id = self.last_assistant_id()?; + self.node_mut(id) + } + + /// The messages on the active path, in order. + pub fn active_messages(&self) -> impl Iterator + '_ { + self.path + .iter() + .filter_map(|id| self.nodes.get(id)) + .map(|node| &node.message) } - pub(crate) fn reserve_id(&mut self) -> NodeId { + /// The linear history as an owned list. + pub fn history(&self) -> Vec { + self.active_messages().cloned().collect() + } + + pub fn reserve_id(&mut self) -> NodeId { let id = self.next_id; self.next_id += 1; id } - pub(crate) fn append(&mut self, message: Message, id: NodeId) { + /// Append a message to the active path as a child of its last node. + pub fn append(&mut self, message: Message, id: NodeId) { assert!( !self.nodes.contains_key(&id), "message node id already exists" @@ -115,35 +157,16 @@ impl Conversation { }, ); self.path.push(id); - self.rebuild_history(); + self.changed.insert(id); } - /// Persistent correction of one active-path message (content, usage, etc.). - /// Node identity, parent links, extensions and inactive branches survive. - pub(crate) fn edit_message(&mut self, id: NodeId, edit: impl FnOnce(&mut Message)) { - if let Some(node) = self.nodes.get_mut(&id) { - edit(&mut node.message); - self.rebuild_history(); - } + /// Nodes appended or edited since the last checkpoint, in id order. + pub fn changed_nodes(&self) -> impl Iterator + '_ { + self.changed.iter().filter_map(|id| self.nodes.get(id)) } - /// Compatibility boundary for existing hooks that take mutable tree nodes. - /// Re-derive history once after the hook batch, including early results. - pub(crate) fn with_nodes_mut( - &mut self, - edit: impl FnOnce(&mut std::collections::BTreeMap, &ConversationPath) -> T, - ) -> T { - let result = edit(&mut self.nodes, &self.path); - self.rebuild_history(); - result - } - - fn rebuild_history(&mut self) { - self.history = self - .path - .iter() - .filter_map(|id| self.nodes.get(id)) - .map(|node| node.message.clone()) - .collect(); + /// Forget the change marks after a successful checkpoint. + pub fn mark_checkpointed(&mut self) { + self.changed.clear(); } } diff --git a/crates/code_assistant_core/src/agent/checkpoint_tests.rs b/crates/code_assistant_core/src/agent/checkpoint_tests.rs index 66506155..8c6ee333 100644 --- a/crates/code_assistant_core/src/agent/checkpoint_tests.rs +++ b/crates/code_assistant_core/src/agent/checkpoint_tests.rs @@ -1,29 +1,45 @@ use super::*; use crate::agent::persistence::AgentStatePersistence; use crate::persistence::{ChatSession, FileSessionPersistence, MessageNode}; +use crate::session::SessionCheckpoint; use std::sync::Mutex; -#[derive(Clone, Default)] -struct Capture(Arc>>); +/// Applies checkpoints to an in-memory session exactly like the store does. +#[derive(Clone)] +struct Capture(Arc>); + +impl Capture { + fn new(session: ChatSession) -> Self { + Self(Arc::new(Mutex::new(session))) + } + + fn session(&self) -> ChatSession { + self.0.lock().unwrap().clone() + } +} + impl AgentStatePersistence for Capture { - fn save_agent_state(&mut self, state: SessionState) -> Result<()> { - *self.0.lock().unwrap() = Some(state); + fn commit_checkpoint(&mut self, checkpoint: SessionCheckpoint<'_>) -> Result<()> { + self.0.lock().unwrap().apply_checkpoint(&checkpoint); Ok(()) } } /// Deterministic format-on-save without depending on an installed formatter. struct FormattingTool; + #[async_trait::async_trait] impl tools_core::Tool for FormattingTool { type Input = serde_json::Value; type Output = agent_core::types::ParseError; + fn spec(&self) -> tools_core::ToolSpec { crate::tools::test_registry() .get("write_file") .unwrap() .spec() } + async fn execute<'a>( &self, _: &mut tools_core::ToolContext<'a>, @@ -71,29 +87,12 @@ async fn formatted_roundtrip(syntax: ToolSyntax) -> Result<()> { rate_limit_info: None, }), ]); - let captured = Capture::default(); - let components = AgentComponents { - llm_provider: Box::new(mock_llm), - project_manager: Arc::new(MockProjectManager::new()), - command_executor: Arc::new(create_command_executor_mock()), - ui: Arc::new(MockUI::default()), - state_persistence: Box::new(captured.clone()), - permission_handler: None, - permissions: Default::default(), - tool_registry: registry.clone(), - sub_agent_runner: None, - wakeups: None, - pty_sessions: None, - browser_sessions: None, - terminal_interrupts: None, - session_source: None, - hooks_factory: None, - }; let config = SessionConfig { tool_syntax: syntax, ..Default::default() }; - let mut agent = Agent::new(components, config.clone()); + // The stored session: one user message plus an inactive branch the run + // never touches. let mut initial = SessionState::from_messages( "checkpoint", "test", @@ -111,20 +110,36 @@ async fn formatted_roundtrip(syntax: ToolSyntax) -> Result<()> { }, ); initial.next_node_id = 100; + let mut stored = + ChatSession::new_empty("checkpoint".into(), "test".into(), config.clone(), None); + stored.message_nodes = initial.message_nodes.clone(); + stored.active_path = initial.active_path.clone(); + stored.next_node_id = initial.next_node_id; + let captured = Capture::new(stored); + + let components = AgentComponents { + llm_provider: Box::new(mock_llm), + project_manager: Arc::new(MockProjectManager::new()), + command_executor: Arc::new(create_command_executor_mock()), + ui: Arc::new(MockUI::default()), + state_persistence: Box::new(captured.clone()), + permission_handler: None, + permissions: Default::default(), + tool_registry: registry.clone(), + sub_agent_runner: None, + wakeups: None, + pty_sessions: None, + browser_sessions: None, + terminal_interrupts: None, + session_source: None, + hooks_factory: None, + }; + let mut agent = Agent::new(components, config.clone()); agent.load_from_session_state(initial).await?; agent.run_single_iteration().await?; - let state = captured.0.lock().unwrap().take().unwrap(); - let mut session = ChatSession::new_empty("checkpoint".into(), "test".into(), config, None); - session.message_nodes = state.message_nodes; - session.active_path = state.active_path; - session.next_node_id = state.next_node_id; - session.tool_executions = state - .tool_executions - .iter() - .map(|e| e.serialize()) - .collect::>()?; + let mut persistence = FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()); - persistence.save_chat_session(&session)?; + persistence.save_chat_session(&captured.session())?; let loaded = persistence.load_chat_session("checkpoint")?.unwrap(); assert_eq!( loaded.message_nodes[&99].extension, @@ -134,7 +149,7 @@ async fn formatted_roundtrip(syntax: ToolSyntax) -> Result<()> { let canonical = loaded.get_active_messages_cloned(); assert_eq!( serde_json::to_value(&canonical)?, - serde_json::to_value(&state.messages)? + serde_json::to_value(agent.message_history())? ); let tool_message = &canonical[1]; let MessageContent::Structured(blocks) = &tool_message.content else { diff --git a/crates/code_assistant_core/src/agent/persistence.rs b/crates/code_assistant_core/src/agent/persistence.rs index 4aab4dd1..bf693b14 100644 --- a/crates/code_assistant_core/src/agent/persistence.rs +++ b/crates/code_assistant_core/src/agent/persistence.rs @@ -1,31 +1,24 @@ +//! How a running agent's checkpoints reach the session store. + +use crate::session::{SessionCheckpoint, SessionManager}; use anyhow::Result; use std::sync::Arc; use tokio::sync::Mutex; -#[cfg(test)] -use crate::types::PlanState; -#[cfg(test)] -use agent_core::types::ToolExecution; -#[cfg(test)] -use llm::Message; - -use crate::session::{SessionManager, SessionState}; - -// The snapshot shape and trait the loop persists through live in the agent +// The checkpoint shape and trait the loop persists through live in the agent // core. -pub use agent_core::{AgentSnapshot, SnapshotPersistence}; +pub use agent_core::{AgentCheckpoint, CheckpointPersistence}; -/// Trait for persisting agent state -/// This abstracts away the storage mechanism from the Agent implementation +/// Application-side checkpoint sink: receives the loop's delta together with +/// the application fields that ride along. pub trait AgentStatePersistence: Send + Sync { - /// Save the current agent state - fn save_agent_state(&mut self, state: SessionState) -> Result<()>; + /// Merge one checkpoint into the stored session. Atomic per call. + fn commit_checkpoint(&mut self, checkpoint: SessionCheckpoint<'_>) -> Result<()>; } -/// Assembles code-assistant's [`SessionState`] from the loop snapshot plus -/// [`crate::plugins::AgentAppState`], and forwards it to an -/// [`AgentStatePersistence`] backend. Snapshots without a session id are not -/// persisted (matching the loop's previous behavior for anonymous agents). +/// Assembles code-assistant's [`SessionCheckpoint`] from the loop checkpoint +/// plus [`crate::plugins::AgentAppState`], and forwards it to an +/// [`AgentStatePersistence`] backend. pub struct SessionStateAdapter { inner: Box, } @@ -36,104 +29,45 @@ impl SessionStateAdapter { } } -impl SnapshotPersistence for SessionStateAdapter { - fn save( +impl CheckpointPersistence for SessionStateAdapter { + fn commit( &mut self, - snapshot: AgentSnapshot, + checkpoint: &AgentCheckpoint<'_>, extensions: &(dyn std::any::Any + Send), ) -> Result<()> { - let Some(session_id) = snapshot.session_id else { - return Ok(()); - }; let state = crate::plugins::AgentAppState::of_ref(extensions); - - self.inner.save_agent_state(SessionState { - session_id, - name: state.session_name.clone(), - message_nodes: snapshot.message_nodes, - active_path: snapshot.active_path, - next_node_id: snapshot.next_node_id, - messages: snapshot.messages, - - tool_executions: snapshot.tool_executions, - plan: state.plan.clone(), - active_skills: state.active_skills.clone(), - // Compatibility restore inputs only. The session manager owns the - // persistent settings and never applies a run's config on save. - config: state.session_config.clone(), - next_request_id: Some(snapshot.next_request_id), - model_config: state.model_config.clone(), + let changed_executions = checkpoint + .changed_executions + .iter() + .map(|execution| execution.serialize()) + .collect::>()?; + + self.inner.commit_checkpoint(SessionCheckpoint { + session_id: checkpoint.session_id, + name: &state.session_name, + changed_nodes: &checkpoint.changed_nodes, + active_path: checkpoint.active_path, + next_node_id: checkpoint.next_node_id, + changed_executions, + plan: &state.plan, + active_skills: &state.active_skills, + next_request_id: checkpoint.next_request_id, }) } } -/// Mock implementation for testing -#[cfg(test)] -#[derive(Default)] -pub struct MockStatePersistence { - pub save_count: usize, - pub last_saved_messages: Option>, - pub last_saved_tool_executions: Option>, - pub last_saved_plan: Option, -} - -#[cfg(test)] -impl MockStatePersistence { - pub fn new() -> Self { - Self::default() - } -} - -#[cfg(test)] -impl AgentStatePersistence for MockStatePersistence { - fn save_agent_state(&mut self, state: SessionState) -> Result<()> { - self.save_count += 1; - self.last_saved_messages = Some(state.messages); - self.last_saved_tool_executions = Some(state.tool_executions); - self.last_saved_plan = Some(state.plan); - Ok(()) - } -} - -/// Decorates a persistence backend with the session-metadata UI update that -/// accompanies every save while an agent runs. Keeps the agent loop free of -/// the `ChatMetadata` concern: the metadata is derived from the saved state. -pub struct MetadataNotifyingPersistence { - inner: Box, - ui: Arc, -} - -impl MetadataNotifyingPersistence { - pub fn new( - inner: Box, - ui: Arc, - ) -> Self { - Self { inner, ui } - } -} - -impl AgentStatePersistence for MetadataNotifyingPersistence { - fn save_agent_state(&mut self, state: SessionState) -> Result<()> { - let metadata = state.build_metadata(); - self.inner.save_agent_state(state)?; - - // Send updated session metadata to UI (fire-and-forget) - let _ = tokio::runtime::Handle::try_current().map(|handle| { - let ui = self.ui.clone(); - handle.spawn(async move { - let _ = ui - .send_event(crate::ui::UiEvent::UpdateSessionMetadata { metadata }) - .await; - }); - }); +/// Discards checkpoints. For agents whose conversation is not a session of +/// its own (sub-agents) and for tests that do not look at persistence. +pub struct NoOpStatePersistence; +impl AgentStatePersistence for NoOpStatePersistence { + fn commit_checkpoint(&mut self, _: SessionCheckpoint<'_>) -> Result<()> { Ok(()) } } -/// Session-specific wrapper that implements AgentStatePersistence -/// This allows agents to save state to a specific session without the SessionManager -/// needing to track a single "current" session (which would break concurrent agents) +/// Commits checkpoints through the session manager, which owns the session +/// entry on disk and the active instance. pub struct SessionStatePersistence { session_manager: Arc>, } @@ -145,12 +79,12 @@ impl SessionStatePersistence { } impl AgentStatePersistence for SessionStatePersistence { - fn save_agent_state(&mut self, state: SessionState) -> Result<()> { - // Use blocking_lock to avoid async context issues - // This is safe because we're in a background task context + fn commit_checkpoint(&mut self, checkpoint: SessionCheckpoint<'_>) -> Result<()> { + // The agent loop is synchronous at this point but runs on the tokio + // runtime; block in place instead of holding an async lock across it. let mut session_manager = tokio::task::block_in_place(|| { tokio::runtime::Handle::current().block_on(self.session_manager.lock()) }); - session_manager.save_session_state(state) + session_manager.commit_checkpoint(checkpoint) } } diff --git a/crates/code_assistant_core/src/agent/runner.rs b/crates/code_assistant_core/src/agent/runner.rs index ec6a1c24..284b6754 100644 --- a/crates/code_assistant_core/src/agent/runner.rs +++ b/crates/code_assistant_core/src/agent/runner.rs @@ -200,8 +200,8 @@ impl Agent { self.app_state_mut().session_name = session_name; } - /// Get a reference to the message history - pub fn message_history(&self) -> &[Message] { + /// The messages on the active path, in order. + pub fn message_history(&self) -> Vec { self.runtime.message_history() } @@ -279,7 +279,6 @@ impl Agent { self.set_tool_scope(ToolScope::Agent); self.invalidate_system_message_cache(); } - self.runtime.normalize_loaded_message_history(); { let state = self.app_state_mut(); state.session_name = session_state.name; @@ -439,7 +438,7 @@ impl Agent { } #[cfg(test)] - pub fn message_history_for_tests(&self) -> &[Message] { + pub fn message_history_for_tests(&self) -> Vec { self.runtime.message_history() } diff --git a/crates/code_assistant_core/src/agent/sub_agent.rs b/crates/code_assistant_core/src/agent/sub_agent.rs index e96ffde4..a2b1a213 100644 --- a/crates/code_assistant_core/src/agent/sub_agent.rs +++ b/crates/code_assistant_core/src/agent/sub_agent.rs @@ -2,7 +2,7 @@ mod run; #[cfg(test)] mod tests; -use crate::agent::persistence::AgentStatePersistence; +use crate::agent::persistence::NoOpStatePersistence; use crate::agent::{Agent, AgentComponents}; use crate::config::DefaultProjectManager; use crate::persistence::SessionModelConfig; @@ -97,14 +97,6 @@ impl SubAgentCancellationRegistry { } } -/// Minimal in-memory persistence used for sub-agents. -struct NoOpStatePersistence; - -impl AgentStatePersistence for NoOpStatePersistence { - fn save_agent_state(&mut self, _state: crate::session::SessionState) -> Result<()> { - Ok(()) - } -} /// Aggregated token usage for a sub-agent run. #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] pub struct SubAgentUsage { diff --git a/crates/code_assistant_core/src/agent/sub_agent/run.rs b/crates/code_assistant_core/src/agent/sub_agent/run.rs index 56c628bb..0f80c060 100644 --- a/crates/code_assistant_core/src/agent/sub_agent/run.rs +++ b/crates/code_assistant_core/src/agent/sub_agent/run.rs @@ -43,12 +43,12 @@ impl SubAgentRunner for DefaultSubAgentRunner { // Earlier requests and tools can have completed before a later // request fails. Preserve their usage as well as their tool list. sub_ui.set_usage(compute_sub_agent_usage( - agent.message_history(), + &agent.message_history(), &self.model_name, )); iteration?; child.token.check()?; - answer = extract_last_assistant_text(agent.message_history()).unwrap_or_default(); + answer = extract_last_assistant_text(&agent.message_history()).unwrap_or_default(); if !require_file_references || has_file_references_with_line_ranges(&answer) { break; } diff --git a/crates/code_assistant_core/src/agent/tests.rs b/crates/code_assistant_core/src/agent/tests.rs index 5fc7c32b..20671545 100644 --- a/crates/code_assistant_core/src/agent/tests.rs +++ b/crates/code_assistant_core/src/agent/tests.rs @@ -2,7 +2,7 @@ mod checkpoint_tests; use super::*; -use crate::agent::persistence::MockStatePersistence; +use crate::agent::persistence::NoOpStatePersistence; use crate::mocks::MockLLMProvider; use crate::mocks::{ MockProjectManager, MockUI, create_command_executor_mock, create_test_response, @@ -52,7 +52,7 @@ async fn test_unknown_tool_error_handling() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -183,7 +183,7 @@ async fn test_invalid_xml_tool_error_handling() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -318,7 +318,7 @@ async fn test_parse_error_handling() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -466,7 +466,7 @@ async fn test_write_file_outside_root_error_masks_paths() -> Result<()> { project_manager: Arc::new(project_manager), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -575,7 +575,7 @@ async fn test_context_compaction_inserts_summary() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: ui.clone(), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -737,7 +737,7 @@ async fn test_compaction_reminds_about_active_skills() -> Result<()> { project_manager: Arc::new(project_manager), command_executor: Arc::new(create_command_executor_mock()), ui: ui.clone(), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -831,7 +831,7 @@ async fn test_compaction_prompt_not_persisted_in_history() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: ui.clone(), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -967,7 +967,7 @@ async fn test_context_compaction_uses_only_messages_after_previous_summary() -> project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: ui.clone(), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -1173,7 +1173,7 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { let project_manager = Arc::new(MockProjectManager::default()); let command_executor = Arc::new(create_command_executor_mock()); let ui = Arc::new(MockUI::default()); - let state_persistence = Box::new(MockStatePersistence::new()); + let state_persistence = Box::new(NoOpStatePersistence); let components = AgentComponents { llm_provider, @@ -1554,7 +1554,7 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -1617,7 +1617,7 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -1692,7 +1692,7 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -1753,7 +1753,7 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -1812,7 +1812,7 @@ async fn test_render_tool_results_generates_unknown_results_for_missing_executio project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -1923,7 +1923,7 @@ async fn test_render_tool_results_preserves_existing_tool_results() -> Result<() project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -2014,7 +2014,7 @@ async fn test_render_tool_results_handles_multiple_unknown_tools() -> Result<()> project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -2176,7 +2176,7 @@ async fn test_prompt_too_long_replaces_large_tool_results() -> Result<()> { project_manager: Arc::new(project_manager), command_executor: Arc::new(create_command_executor_mock()), ui: ui.clone(), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -2312,7 +2312,7 @@ async fn test_prompt_too_long_fallback_drops_exchange_and_compacts() -> Result<( project_manager: Arc::new(mock_project_manager), command_executor: Arc::new(create_command_executor_mock()), ui: ui.clone(), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: None, permissions: Default::default(), tool_registry: crate::tools::test_registry(), @@ -2475,7 +2475,7 @@ async fn test_write_tier_denied_tool_reports_error_to_llm() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: Some(mediator.clone()), permissions: tools_core::ToolPermissions::new(tools_core::PermissionTier::WriteTools), tool_registry: crate::tools::test_registry(), @@ -2541,7 +2541,7 @@ async fn test_write_tier_does_not_ask_for_read_only_tools() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: Some(mediator.clone()), permissions: tools_core::ToolPermissions::new(tools_core::PermissionTier::WriteTools), tool_registry: crate::tools::test_registry(), @@ -2596,7 +2596,7 @@ async fn test_granted_session_asks_only_once_per_tool() -> Result<()> { project_manager: Arc::new(MockProjectManager::new()), command_executor: Arc::new(create_command_executor_mock()), ui: Arc::new(MockUI::default()), - state_persistence: Box::new(MockStatePersistence::new()), + state_persistence: Box::new(NoOpStatePersistence), permission_handler: Some(mediator.clone()), permissions: tools_core::ToolPermissions::new(tools_core::PermissionTier::AllTools), tool_registry: crate::tools::test_registry(), diff --git a/crates/code_assistant_core/src/persistence.rs b/crates/code_assistant_core/src/persistence.rs index 8f8767bb..e7076292 100644 --- a/crates/code_assistant_core/src/persistence.rs +++ b/crates/code_assistant_core/src/persistence.rs @@ -583,6 +583,35 @@ impl ChatSession { self.message_nodes.len() } + /// Merge a running agent's checkpoint. Nodes and journal entries are + /// replaced by id, so branches and records the run never touched + /// survive; counters only ever grow. + pub fn apply_checkpoint(&mut self, checkpoint: &crate::session::SessionCheckpoint<'_>) { + self.name = checkpoint.name.to_string(); + for node in checkpoint.changed_nodes { + self.message_nodes.insert(node.id, (*node).clone()); + } + self.active_path = checkpoint.active_path.to_vec(); + self.next_node_id = self.next_node_id.max(checkpoint.next_node_id); + // The tree is authoritative once a checkpoint has been applied. + self.messages.clear(); + for execution in &checkpoint.changed_executions { + let id = &execution.tool_request.id; + match self + .tool_executions + .iter() + .position(|entry| &entry.tool_request.id == id) + { + Some(index) => self.tool_executions[index] = execution.clone(), + None => self.tool_executions.push(execution.clone()), + } + } + self.plan = checkpoint.plan.clone(); + self.active_skills = checkpoint.active_skills.to_vec(); + self.next_request_id = self.next_request_id.max(checkpoint.next_request_id); + self.updated_at = SystemTime::now(); + } + /// Returns true if the session looks like it failed mid-flight and could /// usefully be "resumed" by re-running the agent against the existing /// message history. diff --git a/crates/code_assistant_core/src/plugins/plan.rs b/crates/code_assistant_core/src/plugins/plan.rs index 85cf4abe..d9952329 100644 --- a/crates/code_assistant_core/src/plugins/plan.rs +++ b/crates/code_assistant_core/src/plugins/plan.rs @@ -18,21 +18,12 @@ impl ToolInterceptor for PlanSnapshotHook { } let plan = AgentAppState::of_ref(&*ctx.extensions).plan.clone(); - - // Find the last assistant message in the active path - for &node_id in ctx.active_path.iter().rev() { - if let Some(node) = ctx.message_nodes.get(&node_id) - && node.message.role == llm::MessageRole::Assistant - { - // Found it - set the snapshot - if let Some(node_mut) = ctx.message_nodes.get_mut(&node_id) { - node_mut.set_plan_snapshot(plan); - trace!("Saved plan snapshot to assistant message node {}", node_id); - } - return; + match ctx.conversation.last_assistant_node_mut() { + Some(node) => { + node.set_plan_snapshot(plan); + trace!("Saved plan snapshot to assistant message node {}", node.id); } + None => trace!("No assistant message found to save plan snapshot"), } - // No assistant message found - this shouldn't happen in normal flow - trace!("No assistant message found to save plan snapshot"); } } diff --git a/crates/code_assistant_core/src/plugins/skill_snapshot.rs b/crates/code_assistant_core/src/plugins/skill_snapshot.rs index b7dc31d4..14cdb836 100644 --- a/crates/code_assistant_core/src/plugins/skill_snapshot.rs +++ b/crates/code_assistant_core/src/plugins/skill_snapshot.rs @@ -33,18 +33,13 @@ impl ToolInterceptor for SkillSnapshotHook { let active_skills = state.active_skills.clone(); // Snapshot onto the last assistant node for branch reconstruction. - for &node_id in ctx.active_path.iter().rev() { - if let Some(node) = ctx.message_nodes.get(&node_id) - && node.message.role == llm::MessageRole::Assistant - { - if let Some(node_mut) = ctx.message_nodes.get_mut(&node_id) { - node_mut.set_active_skills_snapshot(active_skills); - trace!("Saved active-skills snapshot to assistant node {}", node_id); - } - return; + match ctx.conversation.last_assistant_node_mut() { + Some(node) => { + node.set_active_skills_snapshot(active_skills); + trace!("Saved active-skills snapshot to assistant node {}", node.id); } + None => trace!("No assistant message found to save active-skills snapshot"), } - trace!("No assistant message found to save active-skills snapshot"); } } @@ -53,6 +48,7 @@ mod tests { use super::*; use crate::persistence::MessageNode; use crate::session::SessionConfig; + use agent_core::Conversation; use agent_core::hooks::LoopCtx; use llm::Message; use serde_json::json; @@ -83,13 +79,12 @@ mod tests { extension: None, }, ); - let active_path = vec![1]; + let mut conversation = Conversation::restore(message_nodes, vec![1], 2, Vec::new()); let mut state = AgentAppState::new(SessionConfig::default()); { let mut ctx = LoopCtx { - message_nodes: &mut message_nodes, - active_path: &active_path, + conversation: &mut conversation, session_id: None, registry: registry.as_ref(), extensions: &mut state, @@ -99,15 +94,14 @@ mod tests { assert_eq!(state.active_skills, vec!["alpha".to_string()]); assert_eq!( - message_nodes.get(&1).unwrap().active_skills_snapshot(), + conversation.node(1).unwrap().active_skills_snapshot(), Some(vec!["alpha".to_string()]) ); // Re-activating the same skill does not duplicate it. { let mut ctx = LoopCtx { - message_nodes: &mut message_nodes, - active_path: &active_path, + conversation: &mut conversation, session_id: None, registry: registry.as_ref(), extensions: &mut state, @@ -120,8 +114,7 @@ mod tests { #[test] fn ignores_other_tools() { let registry = crate::tools::test_registry(); - let mut message_nodes = BTreeMap::new(); - let active_path: Vec = Vec::new(); + let mut conversation = Conversation::default(); let mut state = AgentAppState::new(SessionConfig::default()); let request = ToolRequest { @@ -133,8 +126,7 @@ mod tests { }; { let mut ctx = LoopCtx { - message_nodes: &mut message_nodes, - active_path: &active_path, + conversation: &mut conversation, session_id: None, registry: registry.as_ref(), extensions: &mut state, diff --git a/crates/code_assistant_core/src/session/instance.rs b/crates/code_assistant_core/src/session/instance.rs index 02e2e1a1..cdffd0b4 100644 --- a/crates/code_assistant_core/src/session/instance.rs +++ b/crates/code_assistant_core/src/session/instance.rs @@ -425,6 +425,24 @@ impl SessionInstance { llm::Usage::zero() } + /// The session-list entry describing the current state of this session. + pub fn metadata(&self) -> ChatMetadata { + ChatMetadata { + id: self.session.id.clone(), + name: self.session.name.clone(), + created_at: self.session.created_at, + updated_at: self.session.updated_at, + message_count: self.session.get_active_messages().len(), + total_usage: self.calculate_total_usage(), + last_usage: self.get_last_usage(), + tokens_limit: None, // Will be updated by persistence layer if available + tool_syntax: self.session.config.tool_syntax, + initial_project: self.session.config.initial_project.clone(), + plan_collapsed: self.session.plan_collapsed, + is_resumable: self.session.is_resumable(), + } + } + /// Reload session data from persistence /// This ensures SessionInstance has the latest state even if agents have made changes pub fn reload_from_persistence( @@ -510,22 +528,7 @@ impl SessionInstance { }); } - let metadata = ChatMetadata { - id: self.session.id.clone(), - name: self.session.name.clone(), - created_at: self.session.created_at, - updated_at: self.session.updated_at, - - message_count: self.session.get_active_messages().len(), - total_usage: self.calculate_total_usage(), - last_usage: self.get_last_usage(), - - tokens_limit: None, // Will be updated by persistence layer if available - tool_syntax: self.session.config.tool_syntax, - initial_project: self.session.config.initial_project.clone(), - plan_collapsed: self.session.plan_collapsed, - is_resumable: self.session.is_resumable(), - }; + let metadata = self.metadata(); let pending_message = self.pending_message.lock().ok().and_then(|pending| { pending diff --git a/crates/code_assistant_core/src/session/manager.rs b/crates/code_assistant_core/src/session/manager.rs index 762ad336..3184498e 100644 --- a/crates/code_assistant_core/src/session/manager.rs +++ b/crates/code_assistant_core/src/session/manager.rs @@ -3,7 +3,6 @@ use llm::{ContentBlock, Message}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; -use std::time::SystemTime; use tokio::sync::Mutex; use crate::agent::{Agent, AgentComponents, DefaultSubAgentRunner, SubAgentCancellationRegistry}; @@ -13,7 +12,7 @@ use crate::persistence::{ }; use crate::session::instance::SessionInstance; use crate::session::sleep_inhibitor::SleepInhibitor; -use crate::session::{SessionConfig, SessionState}; +use crate::session::{SessionCheckpoint, SessionConfig}; use crate::ui::ui_events::UiEvent; use crate::utils::file_utils; use command_executor::{CommandExecutor, SandboxedCommandExecutor}; @@ -1088,9 +1087,6 @@ impl SessionManager { ) }; - // Now save the session state with the user message (outside the borrow scope) - self.save_session_state(session_state.clone())?; - // Broadcast the initial state change self.events.publish_ui( session_id, @@ -1116,13 +1112,6 @@ impl SessionManager { let state_storage = Box::new(crate::agent::persistence::SessionStatePersistence::new( session_manager_ref, )); - // Saves announce the refreshed session metadata to the UI - let state_storage = Box::new( - crate::agent::persistence::MetadataNotifyingPersistence::new( - state_storage, - publisher.clone(), - ), - ); let sandbox_context_clone = sandbox_context.clone(); @@ -1854,46 +1843,24 @@ impl SessionManager { .save_chat_session(&session_instance.session) } - /// Save only run-owned conversation state. `state.config` and - /// `state.model_config` are restore/run inputs, never checkpoint writes. - pub fn save_session_state(&mut self, state: SessionState) -> Result<()> { - let session_id = state.session_id.clone(); - let executions = state - .tool_executions - .into_iter() - .map(|te| te.serialize()) - .collect::>>()?; - let session = self.persistence.update_entry(&session_id, |session| { - session.name = state.name; - // Retain branches not carried by this run. Active-path corrections - // replace nodes by id; concurrent conversation writers still require - // the existing single-agent/branch guards, not just this disk lock. - session.message_nodes.extend(state.message_nodes); - session.active_path = state.active_path; - session.next_node_id = session.next_node_id.max(state.next_node_id); - session.messages.clear(); - for execution in executions { - if let Some(existing) = session - .tool_executions - .iter_mut() - .find(|existing| existing.tool_request.id == execution.tool_request.id) - { - *existing = execution; - } else { - session.tool_executions.push(execution); - } - } - session.plan = state.plan; - session.active_skills = state.active_skills; - if let Some(next_id) = state.next_request_id { - session.next_request_id = session.next_request_id.max(next_id); - } - session.updated_at = SystemTime::now(); - Ok(()) - })?; + /// Merge a running agent's checkpoint into the stored session. Only the + /// conversation delta and run-owned fields are written; settings changed + /// meanwhile, by this or another process, stay untouched. + pub fn commit_checkpoint(&mut self, checkpoint: SessionCheckpoint<'_>) -> Result<()> { + let session = self + .persistence + .update_entry(checkpoint.session_id, |session| { + session.apply_checkpoint(&checkpoint); + Ok(()) + })?; - if let Some(instance) = self.active_sessions.get_mut(&session_id) { + if let Some(instance) = self.active_sessions.get_mut(checkpoint.session_id) { instance.session = session; + let metadata = instance.metadata(); + self.events.publish_ui( + checkpoint.session_id, + UiEvent::UpdateSessionMetadata { metadata }, + ); } Ok(()) } @@ -2059,7 +2026,9 @@ impl SessionManager { #[cfg(test)] mod tests { use super::*; + use crate::persistence::{MessageNode, NodeId}; use std::collections::HashMap; + use std::time::SystemTime; use tempfile::TempDir; fn temp_persistence() -> (FileSessionPersistence, TempDir) { @@ -2068,6 +2037,37 @@ mod tests { (persistence, dir) } + /// Commit `messages` as a linear conversation, the way a run's first + /// checkpoints would. + fn commit_messages(manager: &mut SessionManager, session_id: &str, messages: Vec) { + let nodes: Vec = messages + .into_iter() + .enumerate() + .map(|(index, message)| MessageNode { + id: index as NodeId + 1, + message, + parent_id: (index > 0).then_some(index as NodeId), + created_at: SystemTime::now(), + extension: None, + }) + .collect(); + let changed_nodes: Vec<&MessageNode> = nodes.iter().collect(); + let active_path: Vec = nodes.iter().map(|node| node.id).collect(); + manager + .commit_checkpoint(SessionCheckpoint { + session_id, + name: "run", + changed_nodes: &changed_nodes, + active_path: &active_path, + next_node_id: nodes.len() as NodeId + 1, + changed_executions: Vec::new(), + plan: &crate::types::PlanState::default(), + active_skills: &[], + next_request_id: 1, + }) + .unwrap(); + } + fn build_manager(force_diff: bool) -> (SessionManager, TempDir) { let (persistence, dir) = temp_persistence(); let template = SessionConfig { @@ -2088,14 +2088,8 @@ mod tests { fn checkpoint_preserves_all_session_settings_after_external_changes() { let (mut manager, dir) = build_manager(false); let id = manager.create_session(None).unwrap(); - let captured = SessionState::from_messages( - id.clone(), - "run", - vec![Message::new_user("task")], - SessionConfig::default(), - ); + commit_messages(&mut manager, &id, vec![Message::new_user("task")]); // Another manager owns the settings, independently of the run's manager. - manager.save_session_state(captured.clone()).unwrap(); let mut settings = SessionManager::new( FileSessionPersistence::new_with_root_dir(dir.path().to_path_buf()), SessionConfig::default(), @@ -2132,7 +2126,7 @@ mod tests { expected.plan_collapsed = true; settings.persistence.save_chat_session(&expected).unwrap(); - manager.save_session_state(captured).unwrap(); + commit_messages(&mut manager, &id, vec![Message::new_user("task")]); let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); assert_eq!( serde_json::to_value(&saved.config).unwrap(), @@ -2147,26 +2141,6 @@ mod tests { ); } - #[test] - fn checkpoint_cannot_initialize_project_from_run_config() { - let (mut manager, _dir) = build_manager(false); - let id = manager.create_session(None).unwrap(); - let config = SessionConfig { - initial_project: "run-only".into(), - ..Default::default() - }; - manager - .save_session_state(SessionState::from_messages( - id.clone(), - "run", - vec![Message::new_user("task")], - config, - )) - .unwrap(); - let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); - assert!(saved.config.initial_project.is_empty()); - } - #[test] fn checkpoint_settings_wait_for_entry_lock_and_read_latest_conversation() { use std::sync::mpsc; @@ -2236,14 +2210,7 @@ mod tests { ) .unwrap(); manager.initialize_session_project(&id, &projects).unwrap(); - manager - .save_session_state(SessionState::from_messages( - id.clone(), - "run", - vec![Message::new_user("task")], - config, - )) - .unwrap(); + commit_messages(&mut manager, &id, vec![Message::new_user("task")]); let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); assert_eq!( saved.config.initial_project, @@ -2259,13 +2226,7 @@ mod tests { fn checkpoint_keeps_branches_and_execution_records_not_loaded_by_run() { let (mut manager, _dir) = build_manager(false); let id = manager.create_session(None).unwrap(); - let captured = SessionState::from_messages( - id.clone(), - "run", - vec![Message::new_user("task")], - SessionConfig::default(), - ); - manager.save_session_state(captured.clone()).unwrap(); + commit_messages(&mut manager, &id, vec![Message::new_user("task")]); let mut session = manager.persistence.load_chat_session(&id).unwrap().unwrap(); let branch = session.add_message(Message::new_assistant("another branch")); session.message_nodes.get_mut(&branch).unwrap().extension = @@ -2279,7 +2240,7 @@ mod tests { .unwrap(), ); manager.persistence.save_chat_session(&session).unwrap(); - manager.save_session_state(captured).unwrap(); + commit_messages(&mut manager, &id, vec![Message::new_user("task")]); let saved = manager.persistence.load_chat_session(&id).unwrap().unwrap(); assert_eq!( serde_json::to_value(&saved.message_nodes[&branch]).unwrap(), @@ -2433,60 +2394,6 @@ mod tests { assert!(check.allowed); } - #[test] - fn save_session_state_preserves_model_switch_for_next_iteration() { - let (mut manager, _dir) = build_manager(false); - let session_id = manager - .create_session_with_config( - Some("test".to_string()), - None, - Some(SessionModelConfig::new("old-model".to_string())), - ) - .expect("create session"); - - let mut persisted = manager - .persistence - .load_chat_session(&session_id) - .expect("load session") - .expect("session exists"); - persisted.model_config = Some(SessionModelConfig::new("new-model".to_string())); - persisted.config.use_diff_blocks = true; - manager - .persistence - .save_chat_session(&persisted) - .expect("save switched session"); - - let captured_config = SessionConfig { - use_diff_blocks: false, - ..Default::default() - }; - let mut captured_state = SessionState::from_messages( - session_id.clone(), - "test".to_string(), - vec![Message::new_user("hello")], - captured_config, - ); - captured_state.model_config = Some(SessionModelConfig::new("old-model".to_string())); - - manager - .save_session_state(captured_state) - .expect("save session state"); - - let saved = manager - .persistence - .load_chat_session(&session_id) - .expect("load saved session") - .expect("session exists"); - assert_eq!( - saved - .model_config - .as_ref() - .map(|config| config.model_name.as_str()), - Some("new-model") - ); - assert!(saved.config.use_diff_blocks); - } - /// CLI override (`--use-diff-format`) takes precedence regardless of /// the model's `edit_format` preference. #[test] diff --git a/crates/code_assistant_core/src/session/mod.rs b/crates/code_assistant_core/src/session/mod.rs index d642acae..40cef09d 100644 --- a/crates/code_assistant_core/src/session/mod.rs +++ b/crates/code_assistant_core/src/session/mod.rs @@ -1,5 +1,6 @@ use crate::persistence::{ConversationPath, MessageNode, NodeId, SessionModelConfig}; use crate::types::{PlanState, ToolSyntax}; +use agent_core::types::SerializedToolExecution; use agent_core::types::ToolExecution; use llm::Message; use sandbox::SandboxPolicy; @@ -217,58 +218,21 @@ pub struct SessionState { pub model_config: Option, } -impl SessionState { - /// Build the chat metadata that describes this state, as shown in the - /// session list. `created_at`/`updated_at` and the token limit are - /// placeholders the persistence layer overrides. - pub fn build_metadata(&self) -> crate::persistence::ChatMetadata { - use std::time::SystemTime; - - // Calculate total usage and find last usage across all messages - let mut total_usage = llm::Usage::zero(); - let mut last_usage = llm::Usage::zero(); - - for message in &self.messages { - if let Some(usage) = &message.usage { - total_usage.input_tokens += usage.input_tokens; - total_usage.output_tokens += usage.output_tokens; - total_usage.cache_creation_input_tokens += usage.cache_creation_input_tokens; - total_usage.cache_read_input_tokens += usage.cache_read_input_tokens; - - // For assistant messages, update last usage (most recent wins) - if matches!(message.role, llm::MessageRole::Assistant) { - last_usage = usage.clone(); - } - } - } - - // Compute resumability from the current in-memory history. - // While the agent is running this is largely cosmetic — the UI - // only acts on it once the session is idle — but we still want - // it to reflect the truth as soon as a save runs after the - // agent finishes. - let messages_ref: Vec<&llm::Message> = self.messages.iter().collect(); - let is_resumable = crate::persistence::is_resumable_from_messages(messages_ref.as_slice()); - - crate::persistence::ChatMetadata { - id: self.session_id.clone(), - name: self.name.clone(), // Empty string if not named yet - created_at: SystemTime::now(), // Will be overridden by persistence - updated_at: SystemTime::now(), - message_count: self.messages.len(), - total_usage, - last_usage, - tokens_limit: None, // Will be updated by the persistence layer - tool_syntax: self.config.tool_syntax, - initial_project: if self.config.initial_project.is_empty() { - "unknown".to_string() - } else { - self.config.initial_project.clone() - }, - plan_collapsed: false, // Agent doesn't track UI state - is_resumable, - } - } +/// What a running agent commits after each change: the conversation delta +/// since its previous checkpoint plus the run-owned fields. Session settings +/// are never part of it; the session manager owns those. +pub struct SessionCheckpoint<'a> { + pub session_id: &'a str, + pub name: &'a str, + /// Nodes appended or edited since the previous checkpoint. + pub changed_nodes: &'a [&'a MessageNode], + pub active_path: &'a [NodeId], + pub next_node_id: NodeId, + /// Journal entries recorded or updated since the previous checkpoint. + pub changed_executions: Vec, + pub plan: &'a PlanState, + pub active_skills: &'a [String], + pub next_request_id: u64, } #[cfg(test)] From f15de689d2f60a48f80080a467c9722f449bc0ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:41:21 +0200 Subject: [PATCH 09/15] refactor(agent): journal started calls only for tools with effects Read-only tools no longer write a started record before running; a missing record means the call did not happen and repeating it is harmless. The batch-wide not-started records are gone for the same reason. A completion is committed with one checkpoint that also carries what the hooks derived from it. --- crates/agent_core/src/runtime.rs | 18 ++-- crates/agent_core/src/runtime/tests.rs | 8 +- .../src/runtime/tests/dispatch_tests.rs | 102 +++++++++++------- .../agent_core/src/runtime/tool_execution.rs | 38 ++++--- crates/code_assistant_core/src/agent/tests.rs | 14 +-- 5 files changed, 100 insertions(+), 80 deletions(-) diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index 47d359f5..387aa72e 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -1268,7 +1268,7 @@ impl AgentRuntime { }) }).map(|id| ContentBlock::ToolResult { content: outputs.get(&id).map(|(content, _)| content.clone()).unwrap_or_else(|| { - ToolResultContent::text("Tool result is missing; execution outcome is unknown. Verify the state before retrying any side effects.") + ToolResultContent::text("No result was recorded for this tool call, so it did not run. Call it again if its result is still needed.") }), is_error: Some(outputs.get(&id).map(|(_, error)| *error).unwrap_or(true)), tool_use_id: id, @@ -1395,15 +1395,14 @@ impl AgentRuntime { Ok(()) } - /// Persist formatted inputs through the conversation mutation boundary. - fn update_message_history_with_formatted_tool( - &mut self, - updated_request: &ToolRequest, - ) -> Result<()> { + /// Rewrite the originating tool call with the input the tool settled on + /// (e.g. format-on-save), so follow-up requests see the final input. The + /// caller checkpoints. + fn update_message_history_with_formatted_tool(&mut self, updated_request: &ToolRequest) { let dialect = self.dialect.clone(); let registry = self.registry.clone(); let Some(node) = self.conversation.last_assistant_node_mut() else { - return Ok(()); + return; }; let message = &mut node.message; let request_id = message.request_id.unwrap_or(0); @@ -1445,12 +1444,9 @@ impl AgentRuntime { } } }; - if updated { - self.checkpoint()?; - } else { + if !updated { warn!("Could not find tool call {} to update", updated_request.id); } - Ok(()) } fn update_tool_call_in_text_blocks( diff --git a/crates/agent_core/src/runtime/tests.rs b/crates/agent_core/src/runtime/tests.rs index a6c44909..2aba3b67 100644 --- a/crates/agent_core/src/runtime/tests.rs +++ b/crates/agent_core/src/runtime/tests.rs @@ -303,9 +303,7 @@ fn checkpoint_formatted_input_survives_roundtrip() { .unwrap(); let mut request = ToolRequest::from(&call("a")); request.input = json!({"content": "formatted"}); - agent - .update_message_history_with_formatted_tool(&request) - .unwrap(); + agent.update_message_history_with_formatted_tool(&request); agent .append_message(Message::new_user_content(vec![result("a")])) .unwrap(); @@ -316,7 +314,7 @@ fn checkpoint_formatted_input_survives_roundtrip() { } #[test] -fn checkpoint_dangling_calls_survive_reload_with_unknown_prompt_outcome() { +fn checkpoint_dangling_calls_without_a_record_render_as_not_run() { let (mut agent, saved) = runtime(); agent .append_message(Message::new_assistant_content(vec![call("a"), call("b")])) @@ -338,7 +336,7 @@ fn checkpoint_dangling_calls_survive_reload_with_unknown_prompt_outcome() { "partial tool results must be repaired by id" ); assert!(blocks.iter().any(|block| matches!(block, ContentBlock::ToolResult { tool_use_id, content, .. } - if tool_use_id == "b" && content.contains("unknown") && !content.contains("cancelled by user")))); + if tool_use_id == "b" && content.contains("did not run") && !content.contains("cancelled by user")))); assert_eq!( before, serde_json::to_value(restored.message_history()).unwrap() diff --git a/crates/agent_core/src/runtime/tests/dispatch_tests.rs b/crates/agent_core/src/runtime/tests/dispatch_tests.rs index 16e3144c..0c439c6f 100644 --- a/crates/agent_core/src/runtime/tests/dispatch_tests.rs +++ b/crates/agent_core/src/runtime/tests/dispatch_tests.rs @@ -23,6 +23,7 @@ struct Probe { calls: Arc>>, entered: Arc, release: Arc, + capabilities: Vec>, } #[async_trait::async_trait] impl Tool for Probe { @@ -34,7 +35,7 @@ impl Tool for Probe { description: "test".into(), parameters_schema: json!({"type":"object"}), annotations: None, - capabilities: ToolSpec::capabilities(&["test"]), + capabilities: self.capabilities.clone(), multiline_params: &[], hidden: false, title_template: None, @@ -64,19 +65,29 @@ struct Fixture { entered: Arc, release: Arc, } -fn fixture(requests: &[ToolRequest]) -> Fixture { - let (mut agent, saved) = runtime(); - let calls = Arc::new(Mutex::new(vec![])); - let entered = Arc::new(tokio::sync::Notify::new()); - let release = Arc::new(tokio::sync::Notify::new()); +fn probe_registry(f: &Fixture, capabilities: &[&str]) -> ToolRegistry { let mut registry = ToolRegistry::new(); registry.register(Box::new(Probe { - calls: calls.clone(), - entered: entered.clone(), - release: release.clone(), + calls: f.calls.clone(), + entered: f.entered.clone(), + release: f.release.clone(), + capabilities: capabilities.iter().map(|c| c.to_string().into()).collect(), })); - agent.registry = Arc::new(registry); - agent.tool_capability = "test".into(); + registry +} + +fn fixture(requests: &[ToolRequest]) -> Fixture { + let (agent, saved) = runtime(); + let mut f = Fixture { + agent, + saved, + calls: Arc::new(Mutex::new(vec![])), + entered: Arc::new(tokio::sync::Notify::new()), + release: Arc::new(tokio::sync::Notify::new()), + }; + f.agent.registry = Arc::new(probe_registry(&f, &["test"])); + f.agent.tool_capability = "test".into(); + let agent = &mut f.agent; agent .append_message(Message::new_assistant_content( requests @@ -85,13 +96,7 @@ fn fixture(requests: &[ToolRequest]) -> Fixture { .collect(), )) .unwrap(); - Fixture { - agent, - saved, - calls, - entered, - release, - } + f } fn request(id: &str, wait: bool) -> ToolRequest { ToolRequest { @@ -230,7 +235,7 @@ async fn dispatch_parallel_completions_are_saved_individually() { } #[tokio::test] -async fn dispatch_journal_distinguishes_unstarted_from_uncertain_after_reload() { +async fn dispatch_journals_an_effectful_call_before_it_runs() { let requests = vec![request("one", true), request("two", false)]; let f = fixture(&requests); let mut agent = f.agent; @@ -241,32 +246,55 @@ async fn dispatch_journal_distinguishes_unstarted_from_uncertain_after_reload() let journal = f.saved.saved().executions; f.release.notify_one(); task.await.unwrap().unwrap(); + assert_eq!( - journal.len(), - 2, - "journal records must precede tool invocation" + journal + .iter() + .map(|e| e.tool_request.id.as_str()) + .collect::>(), + ["one"], + "the running call is journaled, its unstarted sibling is not" ); - // Runtime records must be self-describing, even when the tool disappears. - let registry = ToolRegistry::new(); - let restored: Vec<_> = journal - .iter() - .map(|entry| entry.serialize().unwrap().deserialize(®istry).unwrap()) - .collect(); - assert_eq!(restored[0].tool_request.name, "probe"); + // The record must be self-describing, even when the tool disappears. + let restored = journal[0] + .serialize() + .unwrap() + .deserialize(&ToolRegistry::new()) + .unwrap(); + assert_eq!(restored.tool_request.name, "probe"); assert!( - restored[0] + restored .result .as_render() .render(&mut ResourcesTracker::new()) .contains("unknown") ); - assert!( - restored[1] - .result - .as_render() - .render(&mut ResourcesTracker::new()) - .contains("not started") - ); +} + +#[tokio::test] +async fn dispatch_read_only_calls_checkpoint_once() { + let requests = vec![request("one", false)]; + let mut effectful = fixture(&requests); + effectful + .agent + .manage_tool_execution(&requests) + .await + .unwrap(); + // Assistant message, started record, outcome, result message. + assert_eq!(effectful.saved.saved().commits, 4); + + let mut read_only = fixture(&requests); + read_only.agent.registry = Arc::new(probe_registry( + &read_only, + &["test", tools_core::spec::capabilities::READ_ONLY], + )); + read_only + .agent + .manage_tool_execution(&requests) + .await + .unwrap(); + // No started record: the outcome is the first thing the journal sees. + assert_eq!(read_only.saved.saved().commits, 3); } struct FailingUi; diff --git a/crates/agent_core/src/runtime/tool_execution.rs b/crates/agent_core/src/runtime/tool_execution.rs index bb67e5d1..44ea1726 100644 --- a/crates/agent_core/src/runtime/tool_execution.rs +++ b/crates/agent_core/src/runtime/tool_execution.rs @@ -116,16 +116,6 @@ impl AgentRuntime { request.id ); } - // Write intent for the whole batch before any tool can have effects. - // A restart can distinguish an unstarted sibling from an uncertain call. - for request in requests { - self.journal.record(ToolExecution { - tool_request: request.clone(), - result: Box::new(RuntimeToolOutput::not_started("No invocation was made.")), - }); - } - self.checkpoint()?; - let mut parallel = vec![false; requests.len()]; for index in self.hooks.dispatch.parallel_indices(requests) { anyhow::ensure!( @@ -225,11 +215,20 @@ impl AgentRuntime { RuntimeToolOutput::not_started(Self::format_error_for_user(&error)), ))); } - self.journal.record(ToolExecution { - tool_request: request.clone(), - result: Box::new(RuntimeToolOutput::started()), - }); - self.checkpoint()?; + // A tool with effects is journaled as started before it runs, so a + // crash mid-call leaves an "outcome unknown" record instead of + // nothing. Read-only tools skip this: a missing record means the + // call did not happen, and repeating it is harmless either way. + if !self + .registry + .tool_has_capability(&request.name, tools_core::spec::capabilities::READ_ONLY) + { + self.journal.record(ToolExecution { + tool_request: request.clone(), + result: Box::new(RuntimeToolOutput::started()), + }); + self.checkpoint()?; + } // Interceptors execute on the state owner, even for a parallel group, // and only after scope/permission checks and the start checkpoint. @@ -298,16 +297,15 @@ impl AgentRuntime { .registry .is_tool_hidden(&request.name, &self.tool_capability); self.journal.record(completed.execution); - // Commit evidence before hooks/rendering/UI can fail. The active call's - // Started record remains on disk if this commit itself fails. - self.checkpoint()?; if changed { - self.update_message_history_with_formatted_tool(&request)?; + self.update_message_history_with_formatted_tool(&request); } if success { self.after_tool_success(&request); - self.checkpoint()?; } + // One checkpoint commits the outcome together with everything the + // hooks derived from it; the UI is only told afterwards. + self.checkpoint()?; let execution = self.journal.find(&request.id).expect("committed outcome"); let content = execution .result diff --git a/crates/code_assistant_core/src/agent/tests.rs b/crates/code_assistant_core/src/agent/tests.rs index 20671545..2712097d 100644 --- a/crates/code_assistant_core/src/agent/tests.rs +++ b/crates/code_assistant_core/src/agent/tests.rs @@ -1603,7 +1603,7 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { assert!(matches!(history[1].role, MessageRole::Assistant)); let prompt = agent.render_tool_results_in_messages(); assert_eq!(prompt.len(), 3); - assert!(serde_json::to_string(&prompt[2])?.contains("unknown")); + assert!(serde_json::to_string(&prompt[2])?.contains("did not run")); assert!(matches!(history[0].role, MessageRole::User)); Ok(()) @@ -1739,7 +1739,7 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { assert!(matches!(history[1].role, MessageRole::Assistant)); let prompt = agent.render_tool_results_in_messages(); assert_eq!(prompt.len(), 3); - assert!(serde_json::to_string(&prompt[2])?.contains("unknown")); + assert!(serde_json::to_string(&prompt[2])?.contains("did not run")); assert!(matches!(history[0].role, MessageRole::User)); Ok(()) @@ -1801,7 +1801,7 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() } #[tokio::test] -async fn test_render_tool_results_generates_unknown_results_for_missing_executions() -> Result<()> { +async fn test_render_tool_results_generates_not_run_results_for_missing_executions() -> Result<()> { // This test verifies that when an assistant message contains ToolUse blocks // but no corresponding ToolResult, the prompt supplies an unknown outcome // rather than claiming cancellation or silently repeating side effects. @@ -1891,7 +1891,7 @@ async fn test_render_tool_results_generates_unknown_results_for_missing_executio } = &blocks[0] { assert_eq!(tool_use_id, "tool-1-1"); - assert!(content.contains("unknown")); + assert!(content.contains("did not run")); assert!(is_error.unwrap_or(false)); } else { panic!("Expected ToolResult block"); @@ -2005,7 +2005,7 @@ async fn test_render_tool_results_preserves_existing_tool_results() -> Result<() } #[tokio::test] -async fn test_render_tool_results_handles_multiple_unknown_tools() -> Result<()> { +async fn test_render_tool_results_handles_multiple_missing_tools() -> Result<()> { // This test verifies that multiple unknown tool calls are all handled correctly. let mock_llm = MockLLMProvider::new(vec![]); @@ -2093,7 +2093,7 @@ async fn test_render_tool_results_handles_multiple_unknown_tools() -> Result<()> } = &blocks[0] { assert_eq!(tool_use_id, "tool-1-1"); - assert!(content.contains("unknown")); + assert!(content.contains("did not run")); assert!(is_error.unwrap_or(false)); } else { panic!("Expected ToolResult block for first unknown tool"); @@ -2108,7 +2108,7 @@ async fn test_render_tool_results_handles_multiple_unknown_tools() -> Result<()> } = &blocks[1] { assert_eq!(tool_use_id, "tool-1-2"); - assert!(content.contains("unknown")); + assert!(content.contains("did not run")); assert!(is_error.unwrap_or(false)); } else { panic!("Expected ToolResult block for second unknown tool"); From 7dd1ca3146ab346f5bcb505885fc17dd419ae6a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:46:09 +0200 Subject: [PATCH 10/15] refactor(tools_core): one cancellation authority per run, linked child tokens The runtime races its own token against the permission check; the mediator and ToolPermissions no longer carry a copy of it, and the mediator's guard settles a dropped prompt. Sub-agents get a child token that is cancelled with the parent run, which replaces the propagation task and the double select. --- crates/agent_core/src/runtime.rs | 1 - .../agent_core/src/runtime/tool_execution.rs | 23 +++-- .../src/agent/sub_agent.rs | 93 +++++++---------- .../src/agent/sub_agent/run.rs | 59 +++++------ .../src/agent/sub_agent/tests.rs | 11 ++- .../src/session/manager.rs | 3 +- .../src/session/permissions.rs | 56 ++++------- .../src/tests/sub_agent_tests.rs | 27 ----- crates/tools_core/Cargo.toml | 2 +- crates/tools_core/src/cancellation.rs | 99 +++++++++++++++---- crates/tools_core/src/permissions.rs | 18 +--- 11 files changed, 187 insertions(+), 205 deletions(-) diff --git a/crates/agent_core/src/runtime.rs b/crates/agent_core/src/runtime.rs index 387aa72e..088f77a8 100644 --- a/crates/agent_core/src/runtime.rs +++ b/crates/agent_core/src/runtime.rs @@ -176,7 +176,6 @@ impl AgentRuntime { } pub fn set_cancellation(&mut self, cancellation: tools_core::RunCancellation) { - self.permissions.set_cancellation(cancellation.clone()); self.cancellation = cancellation; } diff --git a/crates/agent_core/src/runtime/tool_execution.rs b/crates/agent_core/src/runtime/tool_execution.rs index 44ea1726..b24e25da 100644 --- a/crates/agent_core/src/runtime/tool_execution.rs +++ b/crates/agent_core/src/runtime/tool_execution.rs @@ -191,15 +191,20 @@ impl AgentRuntime { "Tool '{}' is not available in the current scope", request.name ); - self.permissions - .check( - self.permission_handler.as_deref(), - &tool.spec(), - Some(&request.id), - &request.input, - ) - .await?; - self.cancellation.check() + // A stop must not wait for the user to answer a permission prompt; + // dropping the mediator's future settles the prompt as denied. + let spec = tool.spec(); + let check = self.permissions.check( + self.permission_handler.as_deref(), + &spec, + Some(&request.id), + &request.input, + ); + tokio::select! { + biased; + _ = self.cancellation.cancelled() => Err(tools_core::Cancelled.into()), + result = check => result, + } } /// Outer errors are infrastructure failures and abort dispatch. A rejected diff --git a/crates/code_assistant_core/src/agent/sub_agent.rs b/crates/code_assistant_core/src/agent/sub_agent.rs index a2b1a213..da78acab 100644 --- a/crates/code_assistant_core/src/agent/sub_agent.rs +++ b/crates/code_assistant_core/src/agent/sub_agent.rs @@ -14,87 +14,70 @@ use command_executor::{CommandExecutor, DefaultCommandExecutor, SandboxedCommand use llm::Message; use sandbox::{SandboxContext, SandboxPolicy}; use std::collections::HashMap; -use std::sync::{Arc, Mutex, atomic::AtomicBool, atomic::Ordering}; +use std::sync::{Arc, Mutex}; use tools_core::permissions::{PermissionMediator, ToolPermissions}; -/// Cancellation registry keyed by the parent `spawn_agent` tool id. +/// Cancellation tokens of the running sub-agents, keyed by the parent +/// `spawn_agent` tool id, so a child can be cancelled from the UI. #[derive(Default)] pub struct SubAgentCancellationRegistry { - flags: Mutex>, -} - -#[derive(Clone)] -struct ChildCancellation { - flag: Arc, - token: tools_core::RunCancellation, -} - -impl ChildCancellation { - fn cancel(&self) { - self.flag.store(true, Ordering::SeqCst); - self.token.cancel(); - } + children: Mutex>, } +/// A child's entry in the registry; removed when the run ends. struct ChildRegistration<'a> { registry: &'a SubAgentCancellationRegistry, tool_id: String, - cancellation: ChildCancellation, + token: tools_core::RunCancellation, } impl Drop for ChildRegistration<'_> { fn drop(&mut self) { - let mut entries = self.registry.flags.lock().unwrap(); + let mut children = self.registry.children.lock().unwrap(); // A delayed old task must not unregister a replacement with the same id. - if entries + if children .get(&self.tool_id) - .is_some_and(|entry| entry.token.same_run(&self.cancellation.token)) + .is_some_and(|token| token.same_run(&self.token)) { - entries.remove(&self.tool_id); + children.remove(&self.tool_id); } } } impl SubAgentCancellationRegistry { - fn insert(&self, tool_id: String) -> ChildCancellation { - let child = ChildCancellation { - flag: Arc::new(AtomicBool::new(false)), - token: tools_core::RunCancellation::default(), - }; - if let Some(previous) = self.flags.lock().unwrap().insert(tool_id, child.clone()) { + /// Register a child of `parent`. An earlier child with the same id is + /// cancelled and replaced. + fn register_run( + &self, + tool_id: &str, + parent: &tools_core::RunCancellation, + ) -> ChildRegistration<'_> { + let token = parent.child(); + if let Some(previous) = self + .children + .lock() + .unwrap() + .insert(tool_id.to_string(), token.clone()) + { previous.cancel(); } - child - } - - /// Compatibility flag for callers that observe cancellation synchronously. - /// Use `cancel` to also wake asynchronous waiters. - pub fn register(&self, tool_id: String) -> Arc { - self.insert(tool_id).flag - } - - fn register_run(&self, tool_id: &str) -> ChildRegistration<'_> { ChildRegistration { registry: self, tool_id: tool_id.to_string(), - cancellation: self.insert(tool_id.to_string()), + token, } } + /// Cancel the child running for `tool_id`. Returns `false` if none is. pub fn cancel(&self, tool_id: &str) -> bool { - let flags = self.flags.lock().unwrap(); - if let Some(child) = flags.get(tool_id) { - child.cancel(); - true - } else { - false + match self.children.lock().unwrap().get(tool_id) { + Some(token) => { + token.cancel(); + true + } + None => false, } } - - pub fn unregister(&self, tool_id: &str) { - let mut flags = self.flags.lock().unwrap(); - flags.remove(tool_id); - } } /// Aggregated token usage for a sub-agent run. @@ -280,12 +263,12 @@ impl DefaultSubAgentRunner { &self, parent_ui: Arc, parent_tool_id: String, - cancelled: Arc, + cancellation: tools_core::RunCancellation, ) -> Arc { Arc::new(SubAgentUiAdapter::new( parent_ui, parent_tool_id, - cancelled, + cancellation, self.tool_registry.clone(), )) } @@ -552,7 +535,7 @@ impl Default for SubAgentOutput { struct SubAgentUiAdapter { parent: Arc, parent_tool_id: String, - cancelled: Arc, + cancellation: tools_core::RunCancellation, output: Mutex, /// Map from tool_id to index in output.tools for fast lookup tool_id_to_index: Mutex>, @@ -564,13 +547,13 @@ impl SubAgentUiAdapter { fn new( parent: Arc, parent_tool_id: String, - cancelled: Arc, + cancellation: tools_core::RunCancellation, tool_registry: Arc, ) -> Self { Self { parent, parent_tool_id, - cancelled, + cancellation, output: Mutex::new(SubAgentOutput::new()), tool_id_to_index: Mutex::new(std::collections::HashMap::new()), tool_registry, @@ -846,7 +829,7 @@ impl UserInterface for SubAgentUiAdapter { } fn should_streaming_continue(&self) -> bool { - !self.cancelled.load(Ordering::SeqCst) && self.parent.should_streaming_continue() + !self.cancellation.is_cancelled() && self.parent.should_streaming_continue() } fn notify_rate_limit(&self, _seconds_remaining: u64) {} diff --git a/crates/code_assistant_core/src/agent/sub_agent/run.rs b/crates/code_assistant_core/src/agent/sub_agent/run.rs index 0f80c060..cd04c3b9 100644 --- a/crates/code_assistant_core/src/agent/sub_agent/run.rs +++ b/crates/code_assistant_core/src/agent/sub_agent/run.rs @@ -10,23 +10,27 @@ impl SubAgentRunner for DefaultSubAgentRunner { mode: SubAgentMode, require_file_references: bool, ) -> Result { - let registration = self.cancellation_registry.register_run(parent_tool_id); - let child = registration.cancellation.clone(); + // The child's token is cancelled with the parent run and on its own + // by the UI; its runtime stops new calls and wakes provider and + // permission waits. A tool already executing finishes or cooperates. + let registration = self + .cancellation_registry + .register_run(parent_tool_id, &self.parent_cancellation); + let cancellation = registration.token.clone(); let sub_ui = self.build_sub_agent_ui( self.ui.clone(), parent_tool_id.to_string(), - child.flag.clone(), + cancellation.clone(), ); + let work = async { - self.parent_cancellation.check()?; - child.token.check()?; + cancellation.check()?; let mut agent = tokio::select! { biased; - _ = self.parent_cancellation.cancelled() => return Err(tools_core::Cancelled.into()), - _ = child.token.cancelled() => return Err(tools_core::Cancelled.into()), + _ = cancellation.cancelled() => return Err(tools_core::Cancelled.into()), agent = self.build_agent(parent_tool_id, sub_ui.clone(), self.permission_handler.clone()) => agent?, }; - agent.set_cancellation(child.token.clone()); + agent.set_cancellation(cancellation.clone()); let scope = match mode { SubAgentMode::ReadOnly => ToolScope::SubAgentReadOnly, SubAgentMode::Default if self.session_config.use_diff_blocks => { @@ -36,9 +40,10 @@ impl SubAgentRunner for DefaultSubAgentRunner { }; agent.set_tool_scope(scope); agent.append_message(Message::new_user(instructions))?; + let mut answer = String::new(); for attempt in 0..=2 { - child.token.check()?; + cancellation.check()?; let iteration = agent.run_single_iteration().await; // Earlier requests and tools can have completed before a later // request fails. Preserve their usage as well as their tool list. @@ -47,7 +52,7 @@ impl SubAgentRunner for DefaultSubAgentRunner { &self.model_name, )); iteration?; - child.token.check()?; + cancellation.check()?; answer = extract_last_assistant_text(&agent.message_history()).unwrap_or_default(); if !require_file_references || has_file_references_with_line_ranges(&answer) { break; @@ -64,28 +69,19 @@ impl SubAgentRunner for DefaultSubAgentRunner { Ok::<_, anyhow::Error>(answer) }; - // Propagate a parent stop without dropping an already executing child - // tool. The child's runtime stops new calls and wakes provider/permission - // waits; existing side effects finish or cooperate with cancellation. - let propagate_stop = async { - self.parent_cancellation.cancelled().await; - child.cancel(); - std::future::pending::<()>().await; - }; - let result = tokio::select! { - biased; - _ = propagate_stop => unreachable!("stop propagation never completes"), - result = std::panic::AssertUnwindSafe(work).catch_unwind() => { - result.unwrap_or_else(|_| Err(anyhow::anyhow!("Sub-agent panicked; partial work may have occurred"))) - } - }; - // Drop handles registration cleanup on every exit, including abort/panic. - // Explicitly unregister before publishing the terminal child status. + let result = std::panic::AssertUnwindSafe(work) + .catch_unwind() + .await + .unwrap_or_else(|_| { + Err(anyhow::anyhow!( + "Sub-agent panicked; partial work may have occurred" + )) + }); + // Unregister before publishing the terminal child status. drop(registration); + match result { - Ok(answer) - if !child.token.is_cancelled() && !self.parent_cancellation.is_cancelled() => - { + Ok(answer) if !cancellation.is_cancelled() => { sub_ui.set_response(answer.clone()); sub_ui.send_output_update().await; Ok(SubAgentResult { @@ -94,8 +90,7 @@ impl SubAgentRunner for DefaultSubAgentRunner { }) } result => { - let cancelled = - child.token.is_cancelled() || self.parent_cancellation.is_cancelled(); + let cancelled = cancellation.is_cancelled(); let message = if cancelled { "Sub-agent cancelled. Partial work may have occurred; verify side effects before restarting.".to_string() } else { diff --git a/crates/code_assistant_core/src/agent/sub_agent/tests.rs b/crates/code_assistant_core/src/agent/sub_agent/tests.rs index d39e14bd..2d316ff8 100644 --- a/crates/code_assistant_core/src/agent/sub_agent/tests.rs +++ b/crates/code_assistant_core/src/agent/sub_agent/tests.rs @@ -337,7 +337,7 @@ fn sub_agent_success_clears_a_recovered_stream_error() { let ui = SubAgentUiAdapter::new( Arc::new(MockUI::default()), "child".into(), - Arc::new(AtomicBool::new(false)), + tools_core::RunCancellation::default(), crate::tools::test_registry(), ); ui.set_error("temporary stream error".into()); @@ -353,12 +353,13 @@ fn sub_agent_success_clears_a_recovered_stream_error() { #[test] fn sub_agent_old_registration_cannot_remove_its_replacement() { let registry = SubAgentCancellationRegistry::default(); - let old = registry.register_run("child"); - let new = registry.register_run("child"); - assert!(old.cancellation.token.is_cancelled()); + let parent = tools_core::RunCancellation::default(); + let old = registry.register_run("child", &parent); + let new = registry.register_run("child", &parent); + assert!(old.token.is_cancelled()); drop(old); assert!(registry.cancel("child")); - assert!(new.cancellation.token.is_cancelled()); + assert!(new.token.is_cancelled()); drop(new); assert!(!registry.cancel("child")); } diff --git a/crates/code_assistant_core/src/session/manager.rs b/crates/code_assistant_core/src/session/manager.rs index 3184498e..300d10f0 100644 --- a/crates/code_assistant_core/src/session/manager.rs +++ b/crates/code_assistant_core/src/session/manager.rs @@ -1711,8 +1711,7 @@ impl SessionManager { self.events.clone(), instance.pending_permission_requests.clone(), self.permission_timeout, - ) - .with_cancellation(instance.cancellation.clone()), + ), )) } diff --git a/crates/code_assistant_core/src/session/permissions.rs b/crates/code_assistant_core/src/session/permissions.rs index 715786b7..ee76f4f5 100644 --- a/crates/code_assistant_core/src/session/permissions.rs +++ b/crates/code_assistant_core/src/session/permissions.rs @@ -171,7 +171,6 @@ pub struct SessionPermissionMediator { /// on a prompt nobody is there to answer. `None` waits indefinitely (the /// right default for an interactive frontend with a human present). timeout: Option, - cancellation: tools_core::RunCancellation, } impl SessionPermissionMediator { @@ -186,15 +185,9 @@ impl SessionPermissionMediator { events, pending, timeout, - cancellation: tools_core::RunCancellation::default(), } } - pub fn with_cancellation(mut self, cancellation: tools_core::RunCancellation) -> Self { - self.cancellation = cancellation; - self - } - fn next_request_id() -> String { static COUNTER: AtomicU64 = AtomicU64::new(1); format!("perm-{}", COUNTER.fetch_add(1, Ordering::Relaxed)) @@ -263,47 +256,34 @@ impl PermissionMediator for SessionPermissionMediator { ) -> Result { let data = Self::request_data(&request); let request_id = data.request_id.clone(); - let rx = self.cancellation.if_active(|| { - let rx = self.pending.insert(data.clone()); - self.events.publish_ui( - &self.session_id, - UiEvent::RequestToolPermission { request: data }, - ); - rx - })?; - // Also clean up when an enclosing cancellation select drops this future. - let _guard = RequestGuard { + let rx = self.pending.insert(data.clone()); + self.events.publish_ui( + &self.session_id, + UiEvent::RequestToolPermission { request: data }, + ); + // Settles the request on every exit, including a caller that drops + // this future because its run was cancelled. + let _settled = RequestGuard { mediator: self, - request_id: request_id.clone(), + request_id, }; // A dropped responder (stop request, new agent run) counts as denial. - // With a timeout, an unanswered prompt also fails closed: drop the - // pending entry (so a late answer is a no-op) and deny, freeing the + // With a timeout, an unanswered prompt also fails closed, freeing the // lane's turn instead of blocking it forever. - let wait = async { - match self.timeout { - Some(dur) => match tokio::time::timeout(dur, rx).await { - Ok(result) => result.unwrap_or(PermissionDecision::Denied), - Err(_elapsed) => { - self.pending - .resolve(&request_id, PermissionDecision::Denied); - PermissionDecision::Denied - } - }, - None => rx.await.unwrap_or(PermissionDecision::Denied), - } - }; - let decision = tokio::select! { - biased; - _ = self.cancellation.cancelled() => return Err(tools_core::Cancelled.into()), - decision = wait => decision, + let decision = match self.timeout { + Some(dur) => match tokio::time::timeout(dur, rx).await { + Ok(result) => result.unwrap_or(PermissionDecision::Denied), + Err(_elapsed) => PermissionDecision::Denied, + }, + None => rx.await.unwrap_or(PermissionDecision::Denied), }; - self.cancellation.check()?; Ok(decision) } } +/// Removes the pending entry (so a late answer is a no-op) and tells every +/// view the request is settled, whichever way the wait ended. struct RequestGuard<'a> { mediator: &'a SessionPermissionMediator, request_id: String, diff --git a/crates/code_assistant_core/src/tests/sub_agent_tests.rs b/crates/code_assistant_core/src/tests/sub_agent_tests.rs index f59f91f3..ab37bb4a 100644 --- a/crates/code_assistant_core/src/tests/sub_agent_tests.rs +++ b/crates/code_assistant_core/src/tests/sub_agent_tests.rs @@ -1,6 +1,5 @@ //! Tests for the sub-agent feature (spawn_agent tool). -use crate::agent::SubAgentCancellationRegistry; use crate::agent::SubAgentMode; use crate::agent::sub_agent::{SubAgentResult, SubAgentRunner}; use crate::tools::core::ToolScope; @@ -172,32 +171,6 @@ fn test_spawn_agent_input_parsing() { assert_eq!(input.mode, "read_only"); // default } -#[test] -fn test_cancellation_registry() { - let registry = SubAgentCancellationRegistry::default(); - - // Register a new tool - let flag1 = registry.register("tool-1".to_string()); - assert!(!flag1.load(Ordering::SeqCst)); - - let flag2 = registry.register("tool-2".to_string()); - assert!(!flag2.load(Ordering::SeqCst)); - - // Cancel tool-1 - assert!(registry.cancel("tool-1")); - assert!(flag1.load(Ordering::SeqCst)); - assert!(!flag2.load(Ordering::SeqCst)); - - // Cancel non-existent tool returns false - assert!(!registry.cancel("tool-3")); - - // Unregister tool-1 - registry.unregister("tool-1"); - - // Cancel after unregister returns false - assert!(!registry.cancel("tool-1")); -} - #[tokio::test] async fn test_mock_sub_agent_runner() { let runner = MockSubAgentRunner::new(10, "Test response"); diff --git a/crates/tools_core/Cargo.toml b/crates/tools_core/Cargo.toml index 7cc1e9fb..1184467d 100644 --- a/crates/tools_core/Cargo.toml +++ b/crates/tools_core/Cargo.toml @@ -13,7 +13,7 @@ regex = "1.12" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" thiserror = "1.0" -tokio = { version = "1", features = ["sync", "macros"] } +tokio = { version = "1", features = ["sync"] } [dev-dependencies] tokio = { version = "1", features = ["macros", "rt"] } diff --git a/crates/tools_core/src/cancellation.rs b/crates/tools_core/src/cancellation.rs index 3923850a..a36fe974 100644 --- a/crates/tools_core/src/cancellation.rs +++ b/crates/tools_core/src/cancellation.rs @@ -1,33 +1,59 @@ //! One-shot run cancellation, independent of frontend events and error text. + use anyhow::Result; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; +use std::task::Poll; use tokio::sync::watch; #[derive(Debug, thiserror::Error)] #[error("run cancelled")] pub struct Cancelled; -/// Clones address the same run. A subsequent run must use a fresh token. +/// Cancellation of one run. Clones address the same run; a subsequent run +/// must use a fresh token. A child token (see [`RunCancellation::child`]) +/// is cancelled together with its parent but can also be cancelled alone. #[derive(Clone, Debug)] -pub struct RunCancellation(Arc>); +pub struct RunCancellation { + own: Arc>, + /// The tokens of the enclosing runs. + ancestors: Vec>>, +} impl Default for RunCancellation { fn default() -> Self { - Self(Arc::new(watch::channel(false).0)) + Self { + own: Arc::new(watch::channel(false).0), + ancestors: Vec::new(), + } } } impl RunCancellation { + /// A token for work nested inside this run: cancelling the parent + /// cancels the child, cancelling the child leaves the parent running. + pub fn child(&self) -> Self { + let mut ancestors = self.ancestors.clone(); + ancestors.push(self.own.clone()); + Self { + own: Arc::new(watch::channel(false).0), + ancestors, + } + } + pub fn cancel(&self) { - self.0.send_replace(true); + self.own.send_replace(true); } pub fn is_cancelled(&self) -> bool { - *self.0.borrow() + self.senders().any(|sender| *sender.borrow()) } + /// Whether both tokens belong to the same run (not merely to related + /// runs). pub fn same_run(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.0, &other.0) + Arc::ptr_eq(&self.own, &other.own) } pub fn check(&self) -> Result<()> { @@ -38,19 +64,32 @@ impl RunCancellation { } } + /// Resolves once this run or any enclosing run is cancelled. pub async fn cancelled(&self) { - let mut receiver = self.0.subscribe(); - let _ = receiver.wait_for(|cancelled| *cancelled).await; + let mut waits: Vec + Send + '_>>> = self + .senders() + .map(|sender| { + Box::pin(async move { + let mut receiver = sender.subscribe(); + let _ = receiver.wait_for(|cancelled| *cancelled).await; + }) as Pin + Send + '_>> + }) + .collect(); + std::future::poll_fn(|cx| { + if waits + .iter_mut() + .any(|wait| wait.as_mut().poll(cx).is_ready()) + { + Poll::Ready(()) + } else { + Poll::Pending + } + }) + .await } - /// Linearize a short synchronous publication with cancellation. Never - /// await or cancel this token from inside `f`. - pub fn if_active(&self, f: impl FnOnce() -> T) -> Result { - let cancelled = self.0.borrow(); - if *cancelled { - return Err(Cancelled.into()); - } - Ok(f()) + fn senders(&self) -> impl Iterator>> + '_ { + std::iter::once(&self.own).chain(self.ancestors.iter()) } } @@ -64,12 +103,30 @@ mod tests { let second = RunCancellation::default(); first.cancel(); first.cancelled().await; - assert!( - first - .if_active(|| panic!("publication after stop")) - .is_err() - ); + assert!(first.check().is_err()); assert!(!second.is_cancelled()); assert!(!first.same_run(&second)); } + + #[tokio::test] + async fn a_child_is_cancelled_with_its_parent_but_not_the_other_way_round() { + let parent = RunCancellation::default(); + let child = parent.child(); + let grandchild = child.child(); + assert!(!child.same_run(&parent)); + + child.cancel(); + assert!(child.is_cancelled()); + assert!(grandchild.is_cancelled()); + assert!(!parent.is_cancelled()); + + let sibling = parent.child(); + let waiting = tokio::spawn({ + let sibling = sibling.clone(); + async move { sibling.cancelled().await } + }); + parent.cancel(); + waiting.await.unwrap(); + assert!(sibling.is_cancelled()); + } } diff --git a/crates/tools_core/src/permissions.rs b/crates/tools_core/src/permissions.rs index 85a5f4a4..6ab55b11 100644 --- a/crates/tools_core/src/permissions.rs +++ b/crates/tools_core/src/permissions.rs @@ -52,7 +52,6 @@ impl PermissionTier { pub struct ToolPermissions { pub tier: PermissionTier, granted_tools: Arc>>, - cancellation: crate::RunCancellation, } impl ToolPermissions { @@ -60,14 +59,9 @@ impl ToolPermissions { Self { tier, granted_tools: Arc::default(), - cancellation: crate::RunCancellation::default(), } } - pub fn set_cancellation(&mut self, cancellation: crate::RunCancellation) { - self.cancellation = cancellation; - } - pub fn is_granted(&self, tool_name: &str) -> bool { self.granted_tools.lock().unwrap().contains(tool_name) } @@ -91,7 +85,6 @@ impl ToolPermissions { tool_id: Option<&str>, params: &serde_json::Value, ) -> Result<()> { - self.cancellation.check()?; if !self.tier.requires_permission(spec) { return Ok(()); } @@ -105,16 +98,13 @@ impl ToolPermissions { spec.name ); }; - let decision = tokio::select! { - biased; - _ = self.cancellation.cancelled() => return Err(crate::Cancelled.into()), - decision = handler.request_permission(PermissionRequest { + let decision = handler + .request_permission(PermissionRequest { tool_id, tool_name: &spec.name, reason: PermissionRequestReason::ToolInvocation { params }, - }) => decision?, - }; - self.cancellation.check()?; + }) + .await?; match decision { PermissionDecision::GrantedOnce => Ok(()), PermissionDecision::GrantedSession | PermissionDecision::GrantedPersistent => { From d4d28ea42c27c6b10ca203ea4ce01f1cfb875f39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:51:09 +0200 Subject: [PATCH 11/15] refactor(session): run commands as lane tasks, commit runs without interruption The dispatcher spawns one task per lane that runs its commands in arrival order and plain tasks for read-only queries; the select loop with its queue map and join set is gone, and so is the block_on-inside- spawn_blocking detour. Run preparation is still abandoned on stop, but the commit under the manager lock now runs to completion. --- .../src/session/service.rs | 263 ++++++++++-------- .../src/session/service/recovery_tests.rs | 7 +- 2 files changed, 145 insertions(+), 125 deletions(-) diff --git a/crates/code_assistant_core/src/session/service.rs b/crates/code_assistant_core/src/session/service.rs index b4b5c47a..b889af11 100644 --- a/crates/code_assistant_core/src/session/service.rs +++ b/crates/code_assistant_core/src/session/service.rs @@ -6,11 +6,13 @@ //! caller gets *its* answer or *its* error — no correlation over shared //! channels. //! -//! Bounded dispatch serializes mutations per session on the backend tokio -//! runtime. Slow queries use independent tasks; a separate bounded control -//! mailbox handles stop and permission replies. The caller executor (e.g. -//! GPUI) stays decoupled from tokio. Core→UI notifications keep flowing through [`UiEvent`] and are -//! not part of this API. +//! Internally each method hands a command to a worker on the backend tokio +//! runtime, which keeps the caller's executor (e.g. GPUI) decoupled from +//! tokio. Commands that mutate a session run in that session's lane, one +//! after the other; read-only queries run as independent tasks; and a +//! separate control mailbox takes stop requests and permission replies so +//! they never wait behind queued work. Core→UI notifications keep flowing +//! through [`UiEvent`] and are not part of this API. use crate::config::{DefaultProjectManager, ProjectManager, save_project}; use crate::persistence::{ChatMetadata, DraftAttachment, NodeId, SessionModelConfig}; @@ -225,13 +227,18 @@ impl ServiceCtx { type BoxedCommandFuture = Pin + Send>>; type Command = Box BoxedCommandFuture + Send>; +/// A command on its way to the worker, holding its admission permit until +/// it has run. struct Dispatch { - lane: String, + /// Commands in the same lane run one after the other, in arrival order. + /// `None` runs the command on its own. + lane: Option, command: Command, - _permit: tokio::sync::OwnedSemaphorePermit, + permit: tokio::sync::OwnedSemaphorePermit, } -// At most 64 submitted commands, INCLUDING active and lane-queued work. +/// At most this many commands are admitted at once, queued or running. +/// Callers beyond that wait before their command is even accepted. const COMMAND_CAPACITY: usize = 64; /// Cloneable handle to the session command worker. See module docs. @@ -265,54 +272,28 @@ impl SessionService { events: events.clone(), }; let worker = async move { - let control_ctx = ctx.clone(); - let control = async move { - while let Ok(command) = control_rx.recv().await { - command(control_ctx.clone()).await; + let control = { + let ctx = ctx.clone(); + async move { + while let Ok(command) = control_rx.recv().await { + command(ctx.clone()).await; + } } }; let dispatch = async move { - use futures::FutureExt; - use std::collections::{HashMap, VecDeque}; - let mut lanes: HashMap> = HashMap::new(); - let mut tasks = tokio::task::JoinSet::new(); - let mut closed = false; - loop { - tokio::select! { - received = rx.recv(), if !closed => match received { - Ok(dispatch) => { - let lane = dispatch.lane.clone(); - if let Some(queue) = lanes.get_mut(&lane) { - queue.push_back(dispatch); - } else { - lanes.insert(lane.clone(), VecDeque::new()); - let ctx = ctx.clone(); - tasks.spawn(async move { - let _permit = dispatch._permit; - let _ = std::panic::AssertUnwindSafe((dispatch.command)(ctx)).catch_unwind().await; - lane - }); - } - } - Err(_) => closed = true, - }, - completed = tasks.join_next(), if !tasks.is_empty() => { - if let Some(Ok(lane)) = completed { - if let Some(dispatch) = lanes.get_mut(&lane).and_then(|queue| queue.pop_front()) { - let ctx = ctx.clone(); - tasks.spawn(async move { - let _permit = dispatch._permit; - let _ = std::panic::AssertUnwindSafe((dispatch.command)(ctx)).catch_unwind().await; - lane - }); - } else { - lanes.remove(&lane); - } - } + let mut lanes = std::collections::HashMap::new(); + while let Ok(dispatch) = rx.recv().await { + match dispatch.lane { + None => { + let ctx = ctx.clone(); + tokio::spawn(run_command(ctx, dispatch.command, dispatch.permit)); + } + Some(name) => { + let lane = lanes.entry(name).or_insert_with(|| spawn_lane(ctx.clone())); + // Only fails once the lane task is gone, which + // takes the reply channel with it. + let _ = lane.send((dispatch.command, dispatch.permit)); } - } - if closed && tasks.is_empty() { - break; } } }; @@ -371,43 +352,39 @@ impl SessionService { .await } - /// Enqueue a mutation on the global (non-session) lane. + /// Run a mutation that is not tied to one session, after earlier ones. async fn call(&self, f: F) -> Result where F: FnOnce(ServiceCtx) -> Fut + Send + 'static, Fut: Future> + Send + 'static, T: Send + 'static, { - self.call_lane("global".into(), f).await + self.dispatch(Some("global".into()), f).await } + /// Run a mutation of one session, after that session's earlier ones. async fn call_session(&self, session_id: String, f: F) -> Result where F: FnOnce(ServiceCtx) -> Fut + Send + 'static, Fut: Future> + Send + 'static, T: Send + 'static, { - self.call_lane(format!("session:{session_id}"), f).await + self.dispatch(Some(format!("session:{session_id}")), f) + .await } - /// Read-only slow IO gets an independent, supervised task on the backend. - /// Even synchronous libgit/filesystem work cannot stall a single-thread runtime. + /// Run a read-only query on its own, so slow git or filesystem work + /// never holds up session commands. async fn call_io(&self, f: F) -> Result where F: FnOnce(ServiceCtx) -> Fut + Send + 'static, Fut: Future> + Send + 'static, T: Send + 'static, { - static NEXT_IO: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); - let id = NEXT_IO.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - self.call_lane(format!("io:{id}"), move |ctx| async move { - let runtime = tokio::runtime::Handle::current(); - tokio::task::spawn_blocking(move || runtime.block_on(f(ctx))).await? - }) - .await + self.dispatch(None, f).await } - async fn call_lane(&self, lane: String, f: F) -> Result + async fn dispatch(&self, lane: Option, f: F) -> Result where F: FnOnce(ServiceCtx) -> Fut + Send + 'static, Fut: Future> + Send + 'static, @@ -418,7 +395,7 @@ impl SessionService { self.tx .send(Dispatch { lane, - _permit: permit, + permit, command: Box::new(move |ctx| { Box::pin(async move { let _ = reply_tx.send(f(ctx).await); @@ -432,8 +409,9 @@ impl SessionService { .map_err(|_| anyhow!("session service dropped the request"))? } - // Separate bounded mailbox: control never waits behind lane or IO capacity. - // These closures still read/mutate the authoritative manager under its lock. + /// Control commands (stop, permission replies) bypass admission and the + /// lanes: they must reach the manager even while it is busy. They still + /// take the manager lock like everything else. async fn call_control(&self, f: F) -> Result where F: FnOnce(ServiceCtx) -> Fut + Send + 'static, @@ -1228,7 +1206,9 @@ impl SessionService { let manager = ctx.manager.lock().await; session_effective_path(&manager, &session_id)? }; - Ok(discover_review_repos(&project_root)) + tokio::task::spawn_blocking(move || discover_review_repos(&project_root)) + .await + .context("Review repo scan was aborted") }) .await } @@ -1691,14 +1671,6 @@ async fn send_user_message_impl( // Headless dispatch (channel adapters, schedulers) reaches sessions // no frontend has opened since the restart — load on demand. manager.ensure_session_loaded(session_id)?; - anyhow::ensure!( - manager - .get_session(session_id) - .unwrap() - .get_activity_state() - .is_terminal(), - "Session is already running" - ); let cancellation = turn_recorder .as_ref() .map(|recorder| recorder.cancellation.clone()) @@ -1967,6 +1939,10 @@ fn schedule_agent_impl( let task_session = session_id.clone(); let task = tokio::spawn(async move { use futures::FutureExt; + + // Everything that may take long or wait for the user: the LLM + // client, the MCP trust prompt, the tool registry. Abandoned the + // moment the run is cancelled. let prepare = async { let mut session_config = session_config.ok_or_else(|| anyhow!("Session has no model configuration"))?; @@ -1982,8 +1958,6 @@ fn schedule_agent_impl( let factory = factory.clone(); let model = session_config.model_name.clone(); // A synchronous injected constructor must not stall the backend. - // Dropping this wait cannot kill blocking user code, but its result - // has no authority to install a run after cancellation. tokio::task::spawn_blocking(move || factory(&model)).await?? } None => { @@ -1996,66 +1970,88 @@ fn schedule_agent_impl( .await? } }; - cancellation.check()?; let include_local_mcp = SessionManager::resolve_local_mcp_trust( project_dir.as_deref(), Some(permission_handler.as_ref()), ) .await; - cancellation.check()?; let registry = loader(crate::session::manager::RegistryRequest { project_dir: project_dir.clone(), include_local_mcp, }) .await; - cancellation.check()?; - let project_manager = (runtime.project_manager_factory)(); - let command_executor = (runtime.command_executor_factory)(&task_session); - let owner = owner - .upgrade() - .ok_or_else(|| anyhow!("Session service shut down"))?; - let mut manager = owner.lock().await; - // Stop/delete may have won while preparation was in flight. - cancellation.check()?; - // Persist a fallback only while the selection still matches the - // one captured at reservation. A newer selection belongs to the next run. - if session_config.model_name != original_model { - manager.persist_model_fallback(&task_session, &original_model, &session_config)?; - } - events.publish_ui( - &task_session, - UiEvent::UpdateMcpServers { - servers: crate::tools::mcp::session_mcp_servers( - project_dir.as_deref(), - &disabled, - ), - }, - ); - manager - .start_reserved_agent_for_session( - &task_session, - llm_client, - project_manager, - command_executor, - Some(permission_handler.clone()), - tool_scope_override, - turn_recorder.clone(), - registry, - cancellation.clone(), - crate::session::manager::RunConfig { - session: run_session_config, - model: Some(session_config), - }, - ) - .await + Ok::<_, anyhow::Error>((llm_client, registry, session_config, original_model)) }; - let result = tokio::select! { + let prepared = tokio::select! { biased; _ = cancellation.cancelled() => Err(tools_core::Cancelled.into()), result = std::panic::AssertUnwindSafe(prepare).catch_unwind() => { result.unwrap_or_else(|_| Err(anyhow!("Run preparation panicked"))) } }; + + // Committing the run mutates the manager and therefore runs to + // completion once started; it checks for cancellation itself. + let commit = |(llm_client, registry, session_config, original_model): ( + Box, + Arc, + SessionModelConfig, + String, + )| { + let owner = owner.clone(); + let cancellation = cancellation.clone(); + let turn_recorder = turn_recorder.clone(); + let task_session = task_session.clone(); + let events = events.clone(); + async move { + let owner = owner + .upgrade() + .ok_or_else(|| anyhow!("Session service shut down"))?; + let mut manager = owner.lock().await; + // Stop/delete may have won while preparation was in flight. + cancellation.check()?; + // Persist a fallback only while the selection still matches the + // one captured at reservation. A newer selection belongs to the next run. + if session_config.model_name != original_model { + manager.persist_model_fallback( + &task_session, + &original_model, + &session_config, + )?; + } + events.publish_ui( + &task_session, + UiEvent::UpdateMcpServers { + servers: crate::tools::mcp::session_mcp_servers( + project_dir.as_deref(), + &disabled, + ), + }, + ); + manager + .start_reserved_agent_for_session( + &task_session, + llm_client, + (runtime.project_manager_factory)(), + (runtime.command_executor_factory)(&task_session), + Some(permission_handler.clone()), + tool_scope_override, + turn_recorder, + registry, + cancellation, + crate::session::manager::RunConfig { + session: run_session_config, + model: Some(session_config), + }, + ) + .await + } + }; + let result = match prepared { + Ok(prepared) => commit(prepared).await, + Err(error) => Err(error), + }; + if let Err(error) = result { let message = format!("Failed to start agent: {error:#}"); if let Some(owner) = owner.upgrade() { @@ -2074,6 +2070,29 @@ fn schedule_agent_impl( Ok(()) } +type LaneSender = tokio::sync::mpsc::UnboundedSender<(Command, tokio::sync::OwnedSemaphorePermit)>; + +/// A lane: one task that runs its commands strictly in arrival order. Lives +/// until the worker shuts down; a lane per session id is cheap. +fn spawn_lane(ctx: ServiceCtx) -> LaneSender { + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + tokio::spawn(async move { + while let Some((command, permit)) = rx.recv().await { + run_command(ctx.clone(), command, permit).await; + } + }); + tx +} + +/// Run one command, containing a panic so it fails only its own caller. +async fn run_command(ctx: ServiceCtx, command: Command, permit: tokio::sync::OwnedSemaphorePermit) { + use futures::FutureExt; + let _ = std::panic::AssertUnwindSafe(command(ctx)) + .catch_unwind() + .await; + drop(permit); +} + #[cfg(test)] mod recovery_tests; diff --git a/crates/code_assistant_core/src/session/service/recovery_tests.rs b/crates/code_assistant_core/src/session/service/recovery_tests.rs index 2d36d431..bca12ced 100644 --- a/crates/code_assistant_core/src/session/service/recovery_tests.rs +++ b/crates/code_assistant_core/src/session/service/recovery_tests.rs @@ -282,15 +282,16 @@ async fn recovery_slow_io_does_not_block_session_control() { let (service, manager) = test_service_with_manager(tmp.path()); let id = service.create_session(None, None).await.unwrap(); let entered = Arc::new(tokio::sync::Notify::new()); - let (release, released) = std::sync::mpsc::channel(); + let release = Arc::new(tokio::sync::Notify::new()); let task = tokio::spawn({ let service = service.clone(); let entered = entered.clone(); + let release = release.clone(); async move { service .call_io(move |_| async move { entered.notify_one(); - released.recv_timeout(Duration::from_secs(3)).unwrap(); + release.notified().await; Ok(()) }) .await @@ -301,7 +302,7 @@ async fn recovery_slow_io_does_not_block_session_control() { .unwrap(); let stopped = tokio::time::timeout(Duration::from_millis(250), service.request_stop(id.clone())).await; - release.send(()).unwrap(); + release.notify_one(); task.await.unwrap().unwrap(); stopped.expect("slow IO blocked stop").unwrap(); assert!( From 4daac70ae0b71f442f518e7bc85a67f3683968ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:54:05 +0200 Subject: [PATCH 12/15] feat(session): /clear continues in a fresh session with the same settings clear_context wiped the in-memory conversation without persisting it, and the next message reloaded the session from disk, so the history came back. History must never be lost anyway: /clear now creates a new session that inherits project, worktree, model, sandbox policy and permission tier, and the terminal switches to it. --- .../src/session/service.rs | 85 ++++++++++++++----- crates/ui_terminal/src/app.rs | 10 ++- 2 files changed, 72 insertions(+), 23 deletions(-) diff --git a/crates/code_assistant_core/src/session/service.rs b/crates/code_assistant_core/src/session/service.rs index b889af11..509caa66 100644 --- a/crates/code_assistant_core/src/session/service.rs +++ b/crates/code_assistant_core/src/session/service.rs @@ -609,23 +609,23 @@ impl SessionService { .await } - /// Clear the conversation context (messages) for a session. The session - /// itself is kept alive; only the message history is wiped. - pub async fn clear_context(&self, session_id: String) -> Result<()> { - self.call_session(session_id.clone(), move |ctx| async move { - { - let mut manager = ctx.manager.lock().await; - if let Some(session) = manager.get_session_mut(&session_id) { - let chat = &mut session.session; - chat.message_nodes.clear(); - chat.active_path.clear(); - chat.next_node_id = 1; - chat.messages.clear(); - chat.plan = Default::default(); - } - } - ctx.notify_session(&session_id, UiEvent::ClearMessages); - Ok(()) + /// Start over in a new session that inherits this session's settings: + /// project, worktree, model, sandbox policy and permission tier. The + /// current session keeps its history untouched. Returns the new id. + pub async fn start_fresh_session(&self, session_id: String) -> Result { + self.call(move |ctx| async move { + let mut manager = ctx.manager.lock().await; + manager.ensure_session_loaded(&session_id)?; + let (config, model) = { + let instance = manager + .get_session(&session_id) + .ok_or_else(|| anyhow!("Session {session_id} not found"))?; + ( + instance.session.config.clone(), + instance.session.model_config.clone(), + ) + }; + manager.create_session_with_config(None, Some(config), model) }) .await } @@ -2484,6 +2484,48 @@ mod tests { assert!(service.list_sessions().await.unwrap().is_empty()); } + #[tokio::test(flavor = "multi_thread")] + async fn start_fresh_session_keeps_the_history_and_copies_the_settings() { + use crate::session::{TurnDispatch, TurnRequest}; + let tmp = tempfile::tempdir().unwrap(); + let (service, _) = test_service_with_llm( + tmp.path(), + Arc::new(|_| { + Ok(Box::new(StreamingScriptedProvider { + text: "done".into(), + })) + }), + ); + let id = service.create_session(None, None).await.unwrap(); + service + .change_permission_tier(id.clone(), tools_core::PermissionTier::AllTools) + .await + .unwrap(); + service + .change_sandbox_policy(id.clone(), SandboxPolicy::ReadOnly) + .await + .unwrap(); + let TurnDispatch::Started(handle) = service + .start_turn_if_idle(id.clone(), TurnRequest::text("task")) + .await + .unwrap() + else { + panic!("busy") + }; + handle.wait().await.unwrap(); + + let fresh = service.start_fresh_session(id.clone()).await.unwrap(); + assert_ne!(fresh, id); + + let old = service.load_session(id, None).await.unwrap(); + assert_eq!(old.messages.len(), 2, "the old session keeps its history"); + let new = service.load_session(fresh, None).await.unwrap(); + assert!(new.messages.is_empty()); + assert_eq!(new.permission_tier, tools_core::PermissionTier::AllTools); + assert_eq!(new.sandbox_policy, SandboxPolicy::ReadOnly); + assert_eq!(new.current_model, old.current_model); + } + #[tokio::test] async fn load_session_returns_snapshot() { let tmp = tempfile::tempdir().unwrap(); @@ -2762,15 +2804,16 @@ mod tests { let id = service.create_session(None, None).await.unwrap(); let mut subscription = service.subscribe(); - service.clear_context(id.clone()).await.unwrap(); + service.clear_session_error(id.clone()).await.unwrap(); - // The ClearMessages notification arrives session-tagged on the - // broadcast stream. + // The notification arrives session-tagged on the broadcast stream. loop { let event = subscription.recv().await.unwrap(); if matches!( event.payload, - crate::session::event_stream::EventPayload::Ui(UiEvent::ClearMessages) + crate::session::event_stream::EventPayload::Ui( + UiEvent::UpdateSessionActivityState { .. } + ) ) { assert_eq!(event.session_id.as_deref(), Some(id.as_str())); break; diff --git a/crates/ui_terminal/src/app.rs b/crates/ui_terminal/src/app.rs index 21215583..9041401c 100644 --- a/crates/ui_terminal/src/app.rs +++ b/crates/ui_terminal/src/app.rs @@ -220,14 +220,20 @@ impl Actions { }); } + /// `/clear`: continue in a fresh session with the same settings. The + /// current session and its history stay available in the session list. fn clear_context(&self, session_id: String) { let this = self.clone(); tokio::spawn(async move { if this.refuse_if_view_only().await { return; } - if let Err(e) = this.service.clear_context(session_id).await { - this.display_error(format!("Failed to clear context: {e:#}")); + match this.service.start_fresh_session(session_id).await { + Ok(fresh) => { + this.switch_session(fresh); + this.refresh_chat_list(); + } + Err(e) => this.display_error(format!("Failed to start a fresh session: {e:#}")), } }); } From 35214738d4d7230aef93d61347b7d85590724acb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 10 Sep 2026 23:59:21 +0200 Subject: [PATCH 13/15] test: share the LLM mocks and tidy the new test modules MockLLMProvider can replay its responses through the streaming callback and hand out a client factory; PendingLLMProvider replaces the three never-answering stubs. The scripted, waiting, failing and two-call providers the new tests had defined for themselves are gone, and the test modules read like the rest of the code base. --- .../src/runtime/tests/dispatch_tests.rs | 22 ++ .../src/agent/checkpoint_tests.rs | 2 + .../src/agent/sub_agent/tests.rs | 288 ++++++++---------- crates/code_assistant_core/src/mocks.rs | 77 ++++- .../src/session/service.rs | 91 ++---- .../src/session/service/recovery_tests.rs | 119 +++----- 6 files changed, 291 insertions(+), 308 deletions(-) diff --git a/crates/agent_core/src/runtime/tests/dispatch_tests.rs b/crates/agent_core/src/runtime/tests/dispatch_tests.rs index 0c439c6f..cb9057fb 100644 --- a/crates/agent_core/src/runtime/tests/dispatch_tests.rs +++ b/crates/agent_core/src/runtime/tests/dispatch_tests.rs @@ -5,15 +5,18 @@ use tools_core::{Render, Tool, ToolResult, ToolSpec}; #[derive(serde::Serialize, serde::Deserialize)] struct Output(String); + impl ToolResult for Output { fn is_success(&self) -> bool { true } } + impl Render for Output { fn status(&self) -> String { "done".into() } + fn render(&self, _: &mut ResourcesTracker) -> String { self.0.clone() } @@ -25,6 +28,7 @@ struct Probe { release: Arc, capabilities: Vec>, } + #[async_trait::async_trait] impl Tool for Probe { type Input = serde_json::Value; @@ -41,6 +45,7 @@ impl Tool for Probe { title_template: None, } } + async fn execute<'a>( &self, _: &mut ToolContext<'a>, @@ -58,6 +63,7 @@ impl Tool for Probe { Ok(Output(format!("result for {id}"))) } } + struct Fixture { agent: AgentRuntime, saved: Capture, @@ -65,6 +71,7 @@ struct Fixture { entered: Arc, release: Arc, } + fn probe_registry(f: &Fixture, capabilities: &[&str]) -> ToolRegistry { let mut registry = ToolRegistry::new(); registry.register(Box::new(Probe { @@ -98,6 +105,7 @@ fn fixture(requests: &[ToolRequest]) -> Fixture { .unwrap(); f } + fn request(id: &str, wait: bool) -> ToolRequest { ToolRequest { id: id.into(), @@ -107,17 +115,21 @@ fn request(id: &str, wait: bool) -> ToolRequest { end_offset: None, } } + struct Parallel(Vec); + impl ToolDispatchPolicy for Parallel { fn parallel_indices(&self, _: &[ToolRequest]) -> Vec { self.0.clone() } } + struct Observer { attempts: Arc, successes: Arc>>, intercept: bool, } + impl ToolInterceptor for Observer { fn try_intercept( &self, @@ -128,6 +140,7 @@ impl ToolInterceptor for Observer { self.intercept .then(|| Ok(Box::new(Output("intercepted".into())) as Box)) } + fn after_tool_success(&self, request: &ToolRequest, _: &mut LoopCtx) { self.successes.lock().unwrap().push(request.clone()); } @@ -225,10 +238,12 @@ async fn completion_is_checkpointed_while_sibling_waits(parallel: bool) { "a completed tool must be durable before the batch finishes" ); } + #[tokio::test] async fn dispatch_sequential_completions_are_saved_individually() { completion_is_checkpointed_while_sibling_waits(false).await; } + #[tokio::test] async fn dispatch_parallel_completions_are_saved_individually() { completion_is_checkpointed_while_sibling_waits(true).await; @@ -298,6 +313,7 @@ async fn dispatch_read_only_calls_checkpoint_once() { } struct FailingUi; + #[async_trait::async_trait] impl AgentUi for FailingUi { async fn send_event(&self, event: AgentUiEvent) -> Result<(), UIError> { @@ -312,15 +328,19 @@ impl AgentUi for FailingUi { } Ok(()) } + fn display_fragment(&self, _: &DisplayFragment) -> Result<(), UIError> { Ok(()) } + fn should_streaming_continue(&self) -> bool { true } + fn notify_rate_limit(&self, _: u64) {} fn clear_rate_limit(&self) {} } + #[tokio::test] async fn dispatch_ui_failure_does_not_erase_successful_tool_evidence() { let requests = vec![request("one", false)]; @@ -350,6 +370,7 @@ async fn dispatch_parallel_groups_do_not_cross_sequential_barriers() { } struct FailCompletionSave; + impl CheckpointPersistence for FailCompletionSave { fn commit(&mut self, checkpoint: &AgentCheckpoint<'_>, _: &(dyn Any + Send)) -> Result<()> { anyhow::ensure!( @@ -362,6 +383,7 @@ impl CheckpointPersistence for FailCompletionSave { Ok(()) } } + #[tokio::test] async fn dispatch_checkpoint_failure_prevents_further_side_effects() { let requests = vec![request("one", false), request("two", false)]; diff --git a/crates/code_assistant_core/src/agent/checkpoint_tests.rs b/crates/code_assistant_core/src/agent/checkpoint_tests.rs index 8c6ee333..fa8b008f 100644 --- a/crates/code_assistant_core/src/agent/checkpoint_tests.rs +++ b/crates/code_assistant_core/src/agent/checkpoint_tests.rs @@ -260,10 +260,12 @@ fn journal_outcomes_survive_disk_reload_without_the_original_tools() -> Result<( async fn checkpoint_native_formatted_roundtrip() -> Result<()> { formatted_roundtrip(ToolSyntax::Native).await } + #[tokio::test] async fn checkpoint_xml_formatted_roundtrip() -> Result<()> { formatted_roundtrip(ToolSyntax::Xml).await } + #[tokio::test] async fn checkpoint_caret_formatted_roundtrip() -> Result<()> { formatted_roundtrip(ToolSyntax::Caret).await diff --git a/crates/code_assistant_core/src/agent/sub_agent/tests.rs b/crates/code_assistant_core/src/agent/sub_agent/tests.rs index 2d316ff8..6e3c65f3 100644 --- a/crates/code_assistant_core/src/agent/sub_agent/tests.rs +++ b/crates/code_assistant_core/src/agent/sub_agent/tests.rs @@ -1,5 +1,5 @@ use super::*; -use crate::mocks::{MockProjectManager, MockUI}; +use crate::mocks::{MockLLMProvider, MockProjectManager, MockUI, PendingLLMProvider}; use crate::session::service::LlmClientFactory; use std::time::Duration; @@ -21,33 +21,52 @@ fn runner(factory: LlmClientFactory) -> (Arc, Arc .with_project_manager_factory(Arc::new(|| Box::new(MockProjectManager::new()))); (Arc::new(runner), ui) } -struct Fails; -#[async_trait::async_trait] -impl llm::LLMProvider for Fails { - async fn send_message( - &mut self, - _: llm::LLMRequest, - _: Option<&llm::StreamingCallback>, - ) -> Result { - anyhow::bail!("fatal provider failure") - } + +fn failing_provider() -> LlmClientFactory { + Arc::new(|_| { + Ok(Box::new(MockLLMProvider::new(vec![Err(anyhow::anyhow!( + "fatal provider failure" + ))]))) + }) } -struct Silent(Arc); -#[async_trait::async_trait] -impl llm::LLMProvider for Silent { - async fn send_message( - &mut self, - _: llm::LLMRequest, - _: Option<&llm::StreamingCallback>, - ) -> Result { - self.0.notify_one(); - std::future::pending().await + +/// Start the child on its own task and wait until it is parked inside the +/// provider. +async fn parked_child( + runner: &Arc, + entered: &tokio::sync::Notify, +) -> tokio::task::JoinHandle> { + let task = tokio::spawn({ + let runner = runner.clone(); + async move { + runner + .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) + .await + } + }); + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + task +} + +/// The child must end with an error within a short while, without a +/// provider chunk ever arriving. +async fn assert_ends_cancelled(mut task: tokio::task::JoinHandle>) { + let ended = tokio::time::timeout(Duration::from_millis(250), &mut task).await; + if ended.is_err() { + task.abort(); + let _ = task.await; } + assert!( + matches!(ended, Ok(Ok(Err(_)))), + "cancellation waited for another provider chunk" + ); } #[tokio::test] async fn sub_agent_failure_unregisters_and_publishes_terminal_output() { - let (runner, ui) = runner(Arc::new(|_| Ok(Box::new(Fails)))); + let (runner, ui) = runner(failing_provider()); let result = runner .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) .await; @@ -56,6 +75,7 @@ async fn sub_agent_failure_unregisters_and_publishes_terminal_output() { !runner.cancellation_registry.cancel("child"), "failed child remained registered as busy" ); + let events = ui.events(); let last = events.iter().rev().find_map(|event| match event { UiEvent::UpdateToolStatus { status, output, .. } => Some((status, output)), @@ -83,22 +103,10 @@ async fn sub_agent_construction_failure_unregisters() { #[tokio::test] async fn sub_agent_drop_unregisters_even_without_a_return_value() { - let entered = Arc::new(tokio::sync::Notify::new()); - let (runner, _) = runner({ - let entered = entered.clone(); - Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) - }); - let task = tokio::spawn({ - let runner = runner.clone(); - async move { - runner - .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) - .await - } - }); - tokio::time::timeout(Duration::from_secs(2), entered.notified()) - .await - .unwrap(); + let provider = PendingLLMProvider::default(); + let entered = provider.entered.clone(); + let (runner, _) = runner(provider.into_factory()); + let task = parked_child(&runner, &entered).await; task.abort(); let _ = task.await; assert!( @@ -109,64 +117,24 @@ async fn sub_agent_drop_unregisters_even_without_a_return_value() { #[tokio::test] async fn sub_agent_cancel_wakes_a_provider_without_chunks() { - let entered = Arc::new(tokio::sync::Notify::new()); - let (runner, _) = runner({ - let entered = entered.clone(); - Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) - }); - let mut task = tokio::spawn({ - let runner = runner.clone(); - async move { - runner - .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) - .await - } - }); - tokio::time::timeout(Duration::from_secs(2), entered.notified()) - .await - .unwrap(); + let provider = PendingLLMProvider::default(); + let entered = provider.entered.clone(); + let (runner, _) = runner(provider.into_factory()); + let task = parked_child(&runner, &entered).await; assert!(runner.cancellation_registry.cancel("child")); - let ended = tokio::time::timeout(Duration::from_millis(250), &mut task).await; - if ended.is_err() { - task.abort(); - let _ = task.await; - } - assert!( - matches!(ended, Ok(Ok(Err(_)))), - "child cancellation waited for another provider chunk" - ); + assert_ends_cancelled(task).await; assert!(!runner.cancellation_registry.cancel("child")); } #[tokio::test] async fn sub_agent_parent_cancel_wakes_a_provider_without_chunks() { - let entered = Arc::new(tokio::sync::Notify::new()); - let (runner, _) = runner({ - let entered = entered.clone(); - Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) - }); + let provider = PendingLLMProvider::default(); + let entered = provider.entered.clone(); + let (runner, _) = runner(provider.into_factory()); let parent = runner.parent_cancellation.clone(); - let mut task = tokio::spawn({ - let runner = runner.clone(); - async move { - runner - .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) - .await - } - }); - tokio::time::timeout(Duration::from_secs(2), entered.notified()) - .await - .unwrap(); + let task = parked_child(&runner, &entered).await; parent.cancel(); - let ended = tokio::time::timeout(Duration::from_millis(250), &mut task).await; - if ended.is_err() { - task.abort(); - let _ = task.await; - } - assert!( - matches!(ended, Ok(Ok(Err(_)))), - "parent cancellation did not reach the child runtime" - ); + assert_ends_cancelled(task).await; assert!(!runner.cancellation_registry.cancel("child")); } @@ -174,7 +142,8 @@ async fn sub_agent_parent_cancel_wakes_a_provider_without_chunks() { async fn sub_agent_tool_retains_structured_failure_output() { use crate::tools::impls::spawn_agent::{SpawnAgentInput, SpawnAgentTool}; use tools_core::Tool; - let (runner, ui) = runner(Arc::new(|_| Ok(Box::new(Fails)))); + + let (runner, ui) = runner(failing_provider()); let mut services = crate::tools::ToolServices { project_manager: Arc::new(MockProjectManager::new()), plan: None, @@ -213,88 +182,81 @@ async fn sub_agent_tool_retains_structured_failure_output() { ); } -#[tokio::test] -async fn sub_agent_keeps_completed_child_evidence_when_later_request_fails() { - use tools_core::{Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolSpec}; - #[derive(serde::Serialize, serde::Deserialize)] - struct Written; - impl ToolResult for Written { - fn is_success(&self) -> bool { - true - } +/// A read-only probe tool whose success the sub-agent's UI must remember. +struct Probe; + +#[derive(serde::Serialize, serde::Deserialize)] +struct Written; + +impl tools_core::ToolResult for Written { + fn is_success(&self) -> bool { + true } - impl Render for Written { - fn status(&self) -> String { - "written".into() - } - fn render(&self, _: &mut ResourcesTracker) -> String { - "changed test resource".into() - } +} + +impl tools_core::Render for Written { + fn status(&self) -> String { + "written".into() } - struct Probe; - #[async_trait::async_trait] - impl Tool for Probe { - type Input = serde_json::Value; - type Output = Written; - fn spec(&self) -> ToolSpec { - ToolSpec { - name: "probe".into(), - description: "test".into(), - parameters_schema: serde_json::json!({"type":"object"}), - annotations: None, - capabilities: ToolSpec::capabilities(&[ToolScope::SubAgentReadOnly.tag()]), - multiline_params: &[], - hidden: false, - title_template: None, - } - } - async fn execute<'a>( - &self, - _: &mut ToolContext<'a>, - _: &mut Self::Input, - ) -> Result { - Ok(Written) - } + + fn render(&self, _: &mut tools_core::ResourcesTracker) -> String { + "changed test resource".into() } - struct FailAfterTool(bool); - #[async_trait::async_trait] - impl llm::LLMProvider for FailAfterTool { - async fn send_message( - &mut self, - _: llm::LLMRequest, - callback: Option<&llm::StreamingCallback>, - ) -> Result { - if std::mem::replace(&mut self.0, true) { - anyhow::bail!("fatal provider failure after tool") - } - if let Some(callback) = callback { - callback(&llm::StreamingChunk::InputJson { - content: "{}".into(), - tool_name: Some("probe".into()), - tool_id: Some("inner".into()), - })?; - callback(&llm::StreamingChunk::StreamingComplete)?; - } - Ok(llm::LLMResponse { - content: vec![llm::ContentBlock::new_tool_use( - "inner", - "probe", - serde_json::json!({}), - )], - usage: llm::Usage { - input_tokens: 10, - output_tokens: 5, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - }, - rate_limit_info: None, - }) +} + +#[async_trait::async_trait] +impl tools_core::Tool for Probe { + type Input = serde_json::Value; + type Output = Written; + + fn spec(&self) -> tools_core::ToolSpec { + tools_core::ToolSpec { + name: "probe".into(), + description: "test".into(), + parameters_schema: serde_json::json!({"type":"object"}), + annotations: None, + capabilities: tools_core::ToolSpec::capabilities(&[ToolScope::SubAgentReadOnly.tag()]), + multiline_params: &[], + hidden: false, + title_template: None, } } - let (mut runner, _) = runner(Arc::new(|_| Ok(Box::new(FailAfterTool(false))))); + + async fn execute<'a>( + &self, + _: &mut tools_core::ToolContext<'a>, + _: &mut Self::Input, + ) -> Result { + Ok(Written) + } +} + +#[tokio::test] +async fn sub_agent_keeps_completed_child_evidence_when_later_request_fails() { + // Served in reverse: first a tool call, then a fatal failure. + let provider = MockLLMProvider::new(vec![ + Err(anyhow::anyhow!("fatal provider failure after tool")), + Ok(llm::LLMResponse { + content: vec![llm::ContentBlock::new_tool_use( + "inner", + "probe", + serde_json::json!({}), + )], + usage: llm::Usage { + input_tokens: 10, + output_tokens: 5, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + rate_limit_info: None, + }), + ]) + .streaming(); + let (mut runner, _) = runner(provider.into_factory()); let mut registry = tools_core::ToolRegistry::new(); registry.register(Box::new(Probe)); Arc::get_mut(&mut runner).unwrap().tool_registry = Arc::new(registry); + let error = runner .run("child", "inspect".into(), SubAgentMode::ReadOnly, false) .await diff --git a/crates/code_assistant_core/src/mocks.rs b/crates/code_assistant_core/src/mocks.rs index 8761c3be..f181a78a 100644 --- a/crates/code_assistant_core/src/mocks.rs +++ b/crates/code_assistant_core/src/mocks.rs @@ -13,7 +13,9 @@ use fs_explorer::{ reconstruct_formatted_replacements, }, }; -use llm::{LLMProvider, LLMRequest, StreamingCallback as LlmStreamingCallback, types::*}; +use llm::{ + LLMProvider, LLMRequest, StreamingCallback as LlmStreamingCallback, StreamingChunk, types::*, +}; use regex::RegexBuilder; use std::collections::HashMap; use std::path::{Path, PathBuf}; @@ -21,11 +23,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use tools_core::permissions::PermissionMediator; -// New MockLLMProvider that works with the trait-based tool system +/// Scripted LLM provider: serves prepared responses and records requests. #[derive(Default, Clone)] pub struct MockLLMProvider { requests: Arc>>, responses: Arc>>>, + streaming: bool, } impl MockLLMProvider { @@ -36,9 +39,41 @@ impl MockLLMProvider { Self { requests: Arc::new(Mutex::new(Vec::new())), responses: Arc::new(Mutex::new(responses)), + streaming: false, } } + /// Replay each response's content through the streaming callback before + /// returning it, the way a real provider does. Needed by tests that + /// watch display fragments or tool cards; the default stays silent. + pub fn streaming(mut self) -> Self { + self.streaming = true; + self + } + + /// A client factory serving this provider to every run. Clones share the + /// response stack, so the responses are consumed across runs. + pub fn into_factory(self) -> crate::session::service::LlmClientFactory { + Arc::new(move |_| Ok(Box::new(self.clone()))) + } + + fn replay(response: &LLMResponse, callback: &LlmStreamingCallback) -> Result<()> { + for block in &response.content { + match block { + ContentBlock::Text { text, .. } => callback(&StreamingChunk::Text(text.clone()))?, + ContentBlock::ToolUse { + id, name, input, .. + } => callback(&StreamingChunk::InputJson { + content: input.to_string(), + tool_name: Some(name.clone()), + tool_id: Some(id.clone()), + })?, + _ => {} + } + } + callback(&StreamingChunk::StreamingComplete) + } + // Get access to the stored requests pub fn get_requests(&self) -> Vec { self.requests.lock().unwrap().clone() @@ -71,14 +106,46 @@ impl LLMProvider for MockLLMProvider { async fn send_message( &mut self, request: LLMRequest, - _streaming_callback: Option<&LlmStreamingCallback>, + streaming_callback: Option<&LlmStreamingCallback>, ) -> Result { self.requests.lock().unwrap().push(request); - self.responses + let response = self + .responses .lock() .unwrap() .pop() - .unwrap_or(Err(anyhow::anyhow!("No more mock responses"))) + .unwrap_or(Err(anyhow::anyhow!("No more mock responses")))?; + if let Some(callback) = streaming_callback.filter(|_| self.streaming) { + Self::replay(&response, callback)?; + } + Ok(response) + } +} + +/// A provider that never answers. Tests that need a run parked inside the +/// provider wait on `entered` to know it got there. +#[derive(Default, Clone)] +pub struct PendingLLMProvider { + pub entered: Arc, +} + +impl PendingLLMProvider { + /// A client factory serving clones of this provider, which share + /// `entered`. + pub fn into_factory(self) -> crate::session::service::LlmClientFactory { + Arc::new(move |_| Ok(Box::new(self.clone()))) + } +} + +#[async_trait] +impl LLMProvider for PendingLLMProvider { + async fn send_message( + &mut self, + _request: LLMRequest, + _streaming_callback: Option<&LlmStreamingCallback>, + ) -> Result { + self.entered.notify_one(); + std::future::pending().await } } diff --git a/crates/code_assistant_core/src/session/service.rs b/crates/code_assistant_core/src/session/service.rs index 509caa66..f7968afe 100644 --- a/crates/code_assistant_core/src/session/service.rs +++ b/crates/code_assistant_core/src/session/service.rs @@ -2099,6 +2099,7 @@ mod recovery_tests; #[cfg(test)] mod tests { use super::*; + use crate::mocks::{MockLLMProvider, PendingLLMProvider, create_test_response_text}; use crate::persistence::FileSessionPersistence; use crate::session::SessionConfig; @@ -2133,26 +2134,12 @@ mod tests { test_service_with_manager(root).0 } - /// A provider that streams its scripted text through the callback (like - /// a real provider) and returns it as the response — enough to drive a - /// complete agent turn without any network. - struct StreamingScriptedProvider { - text: String, - } - - #[async_trait::async_trait] - impl llm::LLMProvider for StreamingScriptedProvider { - async fn send_message( - &mut self, - _request: llm::LLMRequest, - streaming_callback: Option<&llm::StreamingCallback>, - ) -> Result { - if let Some(callback) = streaming_callback { - callback(&llm::StreamingChunk::Text(self.text.clone()))?; - callback(&llm::StreamingChunk::StreamingComplete)?; - } - Ok(llm::LLMResponse { - content: vec![llm::ContentBlock::new_text(&self.text)], + /// One scripted turn per run: the text is streamed like a real provider + /// would, with a small token usage so outcomes can be checked. + pub(super) fn scripted_turn(text: &'static str) -> LlmClientFactory { + Arc::new(move |_| { + let response = llm::LLMResponse { + content: vec![llm::ContentBlock::new_text(text)], usage: llm::Usage { input_tokens: 10, output_tokens: 5, @@ -2160,8 +2147,11 @@ mod tests { cache_read_input_tokens: 0, }, rate_limit_info: None, - }) - } + }; + Ok(Box::new( + MockLLMProvider::new(vec![Ok(response)]).streaming(), + )) + }) } /// Service whose agent runs use the injected LLM factory instead of the @@ -2195,18 +2185,11 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn checkpoint2_setup_is_reserved_and_control_stays_responsive() { + async fn setup_is_reserved_and_control_stays_responsive() { use crate::session::{TurnDispatch, TurnRequest, TurnStatus}; use std::time::Duration; let tmp = tempfile::tempdir().unwrap(); - let (service, manager) = test_service_with_llm( - tmp.path(), - Arc::new(|_| { - Ok(Box::new(StreamingScriptedProvider { - text: "done".into(), - })) - }), - ); + let (service, manager) = test_service_with_llm(tmp.path(), scripted_turn("done")); let entered = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); manager.lock().await.set_tool_registry_provider({ @@ -2266,24 +2249,12 @@ mod tests { assert!(!service.is_session_busy(id).await.unwrap()); } - struct WaitingProvider; - #[async_trait::async_trait] - impl llm::LLMProvider for WaitingProvider { - async fn send_message( - &mut self, - _: llm::LLMRequest, - _: Option<&llm::StreamingCallback>, - ) -> Result { - std::future::pending().await - } - } - #[tokio::test(flavor = "multi_thread")] - async fn checkpoint2_stop_interrupts_a_provider_without_chunks() { + async fn stop_interrupts_a_provider_without_chunks() { use crate::session::{TurnDispatch, TurnRequest, TurnStatus}; let tmp = tempfile::tempdir().unwrap(); let (service, _) = - test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(WaitingProvider)))); + test_service_with_llm(tmp.path(), PendingLLMProvider::default().into_factory()); let id = service.create_session(None, None).await.unwrap(); let TurnDispatch::Started(handle) = service .start_turn_if_idle(id.clone(), TurnRequest::text("wait")) @@ -2306,7 +2277,7 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn checkpoint2_old_turn_handle_does_not_stop_a_new_run() { + async fn an_old_turn_handle_does_not_stop_a_new_run() { use crate::session::{TurnDispatch, TurnRequest}; let tmp = tempfile::tempdir().unwrap(); let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); @@ -2314,11 +2285,12 @@ mod tests { tmp.path(), Arc::new(move |_| { if calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { - Ok(Box::new(StreamingScriptedProvider { - text: "first done".into(), - })) + Ok(Box::new( + MockLLMProvider::new(vec![Ok(create_test_response_text("first done"))]) + .streaming(), + )) } else { - Ok(Box::new(WaitingProvider)) + Ok(Box::new(PendingLLMProvider::default())) } }), ); @@ -2359,14 +2331,8 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn start_turn_if_idle_resolves_the_exact_outcome() { let tmp = tempfile::tempdir().unwrap(); - let (service, _) = test_service_with_llm( - tmp.path(), - Arc::new(|_model| { - Ok(Box::new(StreamingScriptedProvider { - text: "Considered it carefully; done.".to_string(), - })) - }), - ); + let (service, _) = + test_service_with_llm(tmp.path(), scripted_turn("Considered it carefully; done.")); let id = service.create_session(None, None).await.unwrap(); let dispatch = service @@ -2488,14 +2454,7 @@ mod tests { async fn start_fresh_session_keeps_the_history_and_copies_the_settings() { use crate::session::{TurnDispatch, TurnRequest}; let tmp = tempfile::tempdir().unwrap(); - let (service, _) = test_service_with_llm( - tmp.path(), - Arc::new(|_| { - Ok(Box::new(StreamingScriptedProvider { - text: "done".into(), - })) - }), - ); + let (service, _) = test_service_with_llm(tmp.path(), scripted_turn("done")); let id = service.create_session(None, None).await.unwrap(); service .change_permission_tier(id.clone(), tools_core::PermissionTier::AllTools) diff --git a/crates/code_assistant_core/src/session/service/recovery_tests.rs b/crates/code_assistant_core/src/session/service/recovery_tests.rs index bca12ced..4aa4dddb 100644 --- a/crates/code_assistant_core/src/session/service/recovery_tests.rs +++ b/crates/code_assistant_core/src/session/service/recovery_tests.rs @@ -1,24 +1,13 @@ -use super::tests::{test_service_with_llm, test_service_with_manager}; +//! Run setup and cancellation under adverse conditions: slow preparation, +//! settings changed meanwhile, external locks, stops at awkward moments. + +use super::tests::{scripted_turn, test_service_with_llm, test_service_with_manager}; use super::*; +use crate::mocks::{MockLLMProvider, PendingLLMProvider}; use crate::session::{TurnDispatch, TurnRequest}; use std::time::Duration; -struct Done; -#[async_trait::async_trait] -impl llm::LLMProvider for Done { - async fn send_message( - &mut self, - _: llm::LLMRequest, - _: Option<&llm::StreamingCallback>, - ) -> Result { - Ok(llm::LLMResponse { - content: vec![llm::ContentBlock::new_text("done")], - usage: llm::Usage::zero(), - rate_limit_info: None, - }) - } -} - +/// A registry provider that parks run preparation until released. fn blocked_registry( entered: Arc, release: Arc, @@ -35,9 +24,9 @@ fn blocked_registry( } #[tokio::test(flavor = "multi_thread")] -async fn recovery_pending_setup_does_not_keep_its_owner_alive() { +async fn pending_setup_does_not_keep_its_owner_alive() { let tmp = tempfile::tempdir().unwrap(); - let (service, manager) = test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(Done)))); + let (service, manager) = test_service_with_llm(tmp.path(), scripted_turn("done")); let entered = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); manager @@ -57,6 +46,7 @@ async fn recovery_pending_setup_does_not_keep_its_owner_alive() { .unwrap(); let inhibitor = manager.lock().await.sleep_inhibitor(); assert_eq!(inhibitor.running_count(), 1); + let weak = Arc::downgrade(&manager); drop(manager); drop(service); @@ -66,7 +56,8 @@ async fn recovery_pending_setup_does_not_keep_its_owner_alive() { } }) .await; - // Cleanup even on RED: do not leave a pending task or file lock behind. + // Clean up even when the assertion below fails: no pending task or file + // lock may outlive the test. if let Some(manager) = weak.upgrade() { manager.lock().await.terminate_session_agent(&id); } @@ -86,9 +77,9 @@ async fn recovery_pending_setup_does_not_keep_its_owner_alive() { } #[tokio::test(flavor = "multi_thread")] -async fn recovery_setup_preserves_a_newer_model_selection() { +async fn setup_preserves_a_newer_model_selection() { let tmp = tempfile::tempdir().unwrap(); - let (service, manager) = test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(Done)))); + let (service, manager) = test_service_with_llm(tmp.path(), scripted_turn("done")); let entered = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); manager @@ -106,8 +97,9 @@ async fn recovery_setup_preserves_a_newer_model_selection() { tokio::time::timeout(Duration::from_secs(2), entered.notified()) .await .unwrap(); - // Simulate a selection arriving during preparation without depending on - // the developer's models.json or any real provider configuration. + + // A selection arriving during preparation, written the way another + // process would write it. No models.json involved. let next_model = SessionModelConfig::new("selected-during-setup".into()); { let mut manager = manager.lock().await; @@ -131,7 +123,7 @@ async fn recovery_setup_preserves_a_newer_model_selection() { } #[tokio::test(flavor = "multi_thread")] -async fn recovery_external_run_rejection_does_not_append_a_message() { +async fn a_run_refused_by_an_external_lock_does_not_append_a_message() { let tmp = tempfile::tempdir().unwrap(); let (service, _) = test_service_with_manager(tmp.path()); let id = service.create_session(None, None).await.unwrap(); @@ -151,34 +143,26 @@ async fn recovery_external_run_rejection_does_not_append_a_message() { } #[tokio::test(flavor = "multi_thread")] -async fn recovery_stop_during_permission_does_not_open_the_next_prompt() { - struct TwoCalls; - #[async_trait::async_trait] - impl llm::LLMProvider for TwoCalls { - async fn send_message( - &mut self, - _: llm::LLMRequest, - _: Option<&llm::StreamingCallback>, - ) -> Result { - Ok(llm::LLMResponse { - content: ["first", "second"] - .into_iter() - .map(|id| { - llm::ContentBlock::new_tool_use( - id, - "read_files", - serde_json::json!({"project":"test", "paths":["a.rs"]}), - ) - }) - .collect(), - usage: llm::Usage::zero(), - rate_limit_info: None, +async fn a_stop_during_a_permission_prompt_does_not_open_the_next_one() { + let two_calls = llm::LLMResponse { + content: ["first", "second"] + .into_iter() + .map(|id| { + llm::ContentBlock::new_tool_use( + id, + "read_files", + serde_json::json!({"project":"test", "paths":["a.rs"]}), + ) }) - } - } + .collect(), + usage: llm::Usage::zero(), + rate_limit_info: None, + }; let tmp = tempfile::tempdir().unwrap(); - let (service, manager) = - test_service_with_llm(tmp.path(), Arc::new(|_| Ok(Box::new(TwoCalls)))); + let (service, manager) = test_service_with_llm( + tmp.path(), + MockLLMProvider::new(vec![Ok(two_calls)]).into_factory(), + ); let id = service.create_session(None, None).await.unwrap(); service .change_permission_tier(id.clone(), tools_core::PermissionTier::AllTools) @@ -204,6 +188,7 @@ async fn recovery_stop_during_permission_does_not_open_the_next_prompt() { .await .unwrap(); assert_eq!(first.tool_id.as_deref(), Some("first")); + service.request_stop(id.clone()).await.unwrap(); let outcome = tokio::time::timeout(Duration::from_secs(2), handle.wait()) .await @@ -220,7 +205,7 @@ async fn recovery_stop_during_permission_does_not_open_the_next_prompt() { .snapshot() .is_empty() ); - // An event after completion fences every earlier prompt publication. + // An event published after completion fences every earlier prompt. service.clear_session_error(id).await.unwrap(); loop { match events.recv().await.unwrap().payload { @@ -237,25 +222,11 @@ async fn recovery_stop_during_permission_does_not_open_the_next_prompt() { } #[tokio::test(flavor = "multi_thread")] -async fn recovery_stop_wakes_an_already_waiting_silent_provider() { - struct Silent(Arc); - #[async_trait::async_trait] - impl llm::LLMProvider for Silent { - async fn send_message( - &mut self, - _: llm::LLMRequest, - _: Option<&llm::StreamingCallback>, - ) -> Result { - self.0.notify_one(); - std::future::pending().await - } - } +async fn a_stop_wakes_a_provider_that_never_sends_a_chunk() { let tmp = tempfile::tempdir().unwrap(); - let entered = Arc::new(tokio::sync::Notify::new()); - let (service, _) = test_service_with_llm(tmp.path(), { - let entered = entered.clone(); - Arc::new(move |_| Ok(Box::new(Silent(entered.clone())))) - }); + let provider = PendingLLMProvider::default(); + let entered = provider.entered.clone(); + let (service, _) = test_service_with_llm(tmp.path(), provider.into_factory()); let id = service.create_session(None, None).await.unwrap(); let TurnDispatch::Started(handle) = service .start_turn_if_idle(id.clone(), TurnRequest::text("wait")) @@ -277,13 +248,13 @@ async fn recovery_stop_wakes_an_already_waiting_silent_provider() { } #[tokio::test] -async fn recovery_slow_io_does_not_block_session_control() { +async fn a_slow_query_does_not_block_session_control() { let tmp = tempfile::tempdir().unwrap(); let (service, manager) = test_service_with_manager(tmp.path()); let id = service.create_session(None, None).await.unwrap(); let entered = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); - let task = tokio::spawn({ + let query = tokio::spawn({ let service = service.clone(); let entered = entered.clone(); let release = release.clone(); @@ -303,8 +274,8 @@ async fn recovery_slow_io_does_not_block_session_control() { let stopped = tokio::time::timeout(Duration::from_millis(250), service.request_stop(id.clone())).await; release.notify_one(); - task.await.unwrap().unwrap(); - stopped.expect("slow IO blocked stop").unwrap(); + query.await.unwrap().unwrap(); + stopped.expect("a slow query blocked stop").unwrap(); assert!( manager .lock() From 488e3c1d22da5ba7e43156c1a15bdf102f5b8572 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Fri, 11 Sep 2026 00:31:54 +0200 Subject: [PATCH 14/15] test: name the dangling-call load tests after what a restore does now --- crates/code_assistant_core/src/agent/tests.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/code_assistant_core/src/agent/tests.rs b/crates/code_assistant_core/src/agent/tests.rs index 2712097d..d9da41f3 100644 --- a/crates/code_assistant_core/src/agent/tests.rs +++ b/crates/code_assistant_core/src/agent/tests.rs @@ -1547,7 +1547,7 @@ fn test_update_tool_call_in_text_fallback_mode() -> Result<()> { } #[tokio::test] -async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { +async fn test_load_keeps_native_dangling_tool_request_and_repairs_the_prompt() -> Result<()> { let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { llm_provider: Box::new(mock_llm), @@ -1610,7 +1610,7 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { } #[tokio::test] -async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() -> Result<()> { +async fn test_load_keeps_native_dangling_tool_request_before_a_followup_user_message() -> Result<()> { let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { llm_provider: Box::new(mock_llm), @@ -1685,7 +1685,7 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() } #[tokio::test] -async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { +async fn test_load_keeps_xml_dangling_tool_request_and_repairs_the_prompt() -> Result<()> { let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { llm_provider: Box::new(mock_llm), From 27e338a9e64c4f10048d4f1770f95bfdb012dcc5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Fri, 11 Sep 2026 00:34:12 +0200 Subject: [PATCH 15/15] refactor(dialect): drop the invocation sniffing a restore no longer needs message_contains_invocation existed only for the history normalization that deleted dangling tool calls on load; restores keep them now. --- crates/agent_core/src/dialect.rs | 6 +--- crates/agent_core/src/native/mod.rs | 12 +------- crates/code_assistant_core/src/agent/tests.rs | 3 +- .../src/tool_dialects/caret/mod.rs | 28 ++----------------- .../src/tool_dialects/mod.rs | 19 ------------- .../src/tool_dialects/xml/mod.rs | 26 ++--------------- 6 files changed, 8 insertions(+), 86 deletions(-) diff --git a/crates/agent_core/src/dialect.rs b/crates/agent_core/src/dialect.rs index 58b26a30..9ea1705e 100644 --- a/crates/agent_core/src/dialect.rs +++ b/crates/agent_core/src/dialect.rs @@ -11,7 +11,7 @@ use crate::types::ToolRequest; use crate::ui::{AgentUi, HiddenTools, StreamProcessorTrait}; use anyhow::Result; -use llm::{LLMResponse, Message}; +use llm::LLMResponse; use std::sync::Arc; use tools_core::ToolRegistry; @@ -65,8 +65,4 @@ pub trait ToolDialect: Send + Sync { registry: &ToolRegistry, capability: &str, ) -> Option; - - /// Whether an already stored message contains a tool invocation in this - /// dialect (used to normalize the history when loading a session). - fn message_contains_invocation(&self, message: &Message, registry: &ToolRegistry) -> bool; } diff --git a/crates/agent_core/src/native/mod.rs b/crates/agent_core/src/native/mod.rs index 281cf6ef..aed9909e 100644 --- a/crates/agent_core/src/native/mod.rs +++ b/crates/agent_core/src/native/mod.rs @@ -11,7 +11,7 @@ use crate::dialect::ToolDialect; use crate::types::ToolRequest; use crate::ui::{AgentUi, HiddenTools, StreamProcessorTrait}; use anyhow::Result; -use llm::{ContentBlock, LLMResponse, Message, MessageContent}; +use llm::{ContentBlock, LLMResponse}; use std::sync::Arc; use tools_core::ToolRegistry; @@ -82,14 +82,4 @@ impl ToolDialect for NativeDialect { // Native mode uses API-provided tool definitions, no custom documentation needed None } - - fn message_contains_invocation(&self, message: &Message, _registry: &ToolRegistry) -> bool { - if let MessageContent::Structured(blocks) = &message.content { - blocks - .iter() - .any(|block| matches!(block, ContentBlock::ToolUse { .. })) - } else { - false - } - } } diff --git a/crates/code_assistant_core/src/agent/tests.rs b/crates/code_assistant_core/src/agent/tests.rs index d9da41f3..b1dee6b2 100644 --- a/crates/code_assistant_core/src/agent/tests.rs +++ b/crates/code_assistant_core/src/agent/tests.rs @@ -1610,7 +1610,8 @@ async fn test_load_keeps_native_dangling_tool_request_and_repairs_the_prompt() - } #[tokio::test] -async fn test_load_keeps_native_dangling_tool_request_before_a_followup_user_message() -> Result<()> { +async fn test_load_keeps_native_dangling_tool_request_before_a_followup_user_message() -> Result<()> +{ let mock_llm = MockLLMProvider::new(vec![]); let components = AgentComponents { llm_provider: Box::new(mock_llm), diff --git a/crates/code_assistant_core/src/tool_dialects/caret/mod.rs b/crates/code_assistant_core/src/tool_dialects/caret/mod.rs index fef926ca..93bd00ae 100644 --- a/crates/code_assistant_core/src/tool_dialects/caret/mod.rs +++ b/crates/code_assistant_core/src/tool_dialects/caret/mod.rs @@ -13,16 +13,15 @@ mod tests; pub use parser::parse_caret_tool_invocations; pub use stream::CaretStreamProcessor; -use crate::tool_dialects::{example_placeholder, is_multiline_param, message_text_segments}; +use crate::tool_dialects::{example_placeholder, is_multiline_param}; use crate::tools::ToolRequest; use crate::tools::core::ToolRegistry; use crate::tools::tool_use_filter::SmartToolFilter; use agent_core::dialect::ToolDialect; use agent_core::ui::{AgentUi, HiddenTools, StreamProcessorTrait}; use anyhow::Result; -use llm::{ContentBlock, LLMResponse, Message}; +use llm::{ContentBlock, LLMResponse}; use std::sync::Arc; -use tracing::debug; /// Parse Caret tool requests from LLM response and return both requests and truncated response after first tool fn parse_and_truncate_caret_response( @@ -131,29 +130,6 @@ impl ToolDialect for CaretDialect { fn render_format_section_for_prompt(&self) -> Option { Some(self.generate_caret_syntax_documentation()) } - - fn message_contains_invocation(&self, message: &Message, registry: &ToolRegistry) -> bool { - let request_id = message.request_id.unwrap_or(0); - for text in message_text_segments(message) { - if !text.contains("^^^") { - continue; - } - match parse_caret_tool_invocations(text, request_id, 0, None, registry) { - Ok((requests, _)) => { - if !requests.is_empty() { - return true; - } - } - Err(error) => { - debug!( - "Failed to parse Caret tool invocation while inspecting message: {error}" - ); - return true; - } - } - } - false - } } impl CaretDialect { diff --git a/crates/code_assistant_core/src/tool_dialects/mod.rs b/crates/code_assistant_core/src/tool_dialects/mod.rs index 5ff41d7f..5fc0773b 100644 --- a/crates/code_assistant_core/src/tool_dialects/mod.rs +++ b/crates/code_assistant_core/src/tool_dialects/mod.rs @@ -17,7 +17,6 @@ use crate::tools::core::ToolRegistry; use crate::types::ToolSyntax; use agent_core::ToolDialect; use anyhow::{Result, anyhow}; -use llm::{ContentBlock, Message, MessageContent}; use serde_json::{Value, json}; use std::collections::HashMap; use std::sync::Arc; @@ -30,24 +29,6 @@ pub fn dialect_for(syntax: ToolSyntax) -> Arc { ToolSyntax::Caret => Arc::new(CaretDialect), } } - -/// The text segments of a message, for invocation sniffing. -pub(crate) fn message_text_segments(message: &Message) -> Vec<&str> { - match &message.content { - MessageContent::Text(text) => vec![text.as_str()], - MessageContent::Structured(blocks) => blocks - .iter() - .filter_map(|block| { - if let ContentBlock::Text { text, .. } = block { - Some(text.as_str()) - } else { - None - } - }) - .collect(), - } -} - /// Whether the named parameter of the given tool typically spans multiple /// lines (block syntax in the text dialects). pub(crate) fn is_multiline_param( diff --git a/crates/code_assistant_core/src/tool_dialects/xml/mod.rs b/crates/code_assistant_core/src/tool_dialects/xml/mod.rs index 280e13f5..bc8b77ea 100644 --- a/crates/code_assistant_core/src/tool_dialects/xml/mod.rs +++ b/crates/code_assistant_core/src/tool_dialects/xml/mod.rs @@ -12,16 +12,15 @@ mod tests; pub use parser::parse_xml_tool_invocations; pub use stream::XmlStreamProcessor; -use crate::tool_dialects::{example_placeholder, is_multiline_param, message_text_segments}; +use crate::tool_dialects::{example_placeholder, is_multiline_param}; use crate::tools::ToolRequest; use crate::tools::core::ToolRegistry; use crate::tools::tool_use_filter::SmartToolFilter; use agent_core::dialect::ToolDialect; use agent_core::ui::{AgentUi, HiddenTools, StreamProcessorTrait}; use anyhow::Result; -use llm::{ContentBlock, LLMResponse, Message}; +use llm::{ContentBlock, LLMResponse}; use std::sync::Arc; -use tracing::debug; /// Parse XML tool requests from LLM response and return both requests and truncated response after first tool fn parse_and_truncate_xml_response( @@ -130,27 +129,6 @@ impl ToolDialect for XmlDialect { fn render_format_section_for_prompt(&self) -> Option { Some(self.generate_xml_syntax_documentation()) } - - fn message_contains_invocation(&self, message: &Message, registry: &ToolRegistry) -> bool { - let request_id = message.request_id.unwrap_or(0); - for text in message_text_segments(message) { - if !text.contains(" { - if !requests.is_empty() { - return true; - } - } - Err(error) => { - debug!("Failed to parse XML tool invocation while inspecting message: {error}"); - return true; - } - } - } - false - } } impl XmlDialect {