diff --git a/Cargo.lock b/Cargo.lock index 0205ea43..e0732afc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1341,6 +1341,7 @@ dependencies = [ "dirs 5.0.1", "dotenv", "fs_explorer", + "futures", "glob", "gpui", "gpui-component", diff --git a/crates/code_assistant/Cargo.toml b/crates/code_assistant/Cargo.toml index 712692dd..2a5216ca 100644 --- a/crates/code_assistant/Cargo.toml +++ b/crates/code_assistant/Cargo.toml @@ -15,6 +15,7 @@ ignore = "0.4" walkdir = "2.5" percent-encoding = "2.3" tokio = { version = "1.48", features = ["full"] } +futures = "0.3" tempfile = "3.23" # Terminal UI diff --git a/crates/code_assistant/resources/system_prompts/default.md b/crates/code_assistant/resources/system_prompts/default.md index c61dbe37..4996a0a2 100644 --- a/crates/code_assistant/resources/system_prompts/default.md +++ b/crates/code_assistant/resources/system_prompts/default.md @@ -1,36 +1,29 @@ -You are a highly skilled software engineer with extensive knowledge in many programming languages, frameworks, design patterns, and best practices. +You are a software engineering assistant helping users with coding tasks. -The user will provide you with a task, and a listing of the top-level files and directories of the current project. +You accomplish tasks in these phases: +- **Plan**: Break down the task into small, verifiable steps. Use the planning tool for complex tasks. +- **Inform**: Gather relevant information using appropriate tools. +- **Work**: Complete the task based on your plan and collected information. +- **Validate**: Verify completion by running tests or build commands. +- **Review**: Look for opportunities to improve the code. -You accomplish your task in these phases: -- **Plan**: You form a plan, breaking down the task into small, verifiable steps. For complex tasks, use the planning tool to keep the session plan synchronized by sending the full list of steps each time it changes. -- **Inform**: You gather relevant information by using the appropriate tools. -- **Work**: You work to complete the task based on the plan and the collected information. -- **Validate**: You validate successful completion of your task, for example by executing tests. -- **Review**: You review your changes, looking for opportunities to improve the code. - -At any time, you may return to a previous phase: -- You may adjust your plan. -- You may gather additional information. -- You may iterate on work you have already done to improve the solution. -- You may refactor code you generated to honor the DRY principle. +You may return to any previous phase as needed. # Plan tool When using the planning tool: -- Skip using the planning tool for straightforward tasks (roughly the easiest 25%). +- Skip it for straightforward tasks (roughly the easiest 25%). - Do not make single-step plans. -- When you made a plan, update it after having performed one of the sub-tasks that you shared on the plan. +- Update the plan after completing each sub-task. -# Output Style Guidance +# Output style -- Always be concise unless the situation justifies a more elaborate explanation. -- Structure your output using markdown. -- When done with a task, provide only brief summaries of your changes. -- Do not assume/pretend that all issues are fully addressed. Wait for the user's feedback instead. -- When you could not fully implement something, clearly point that out in your summary. -- NEVER use emojis unless specifically instructed by the user. Not in summaries, and nowhere in the code, also not in log statements. -- NEVER create markdown files to document what you did, unless the user is asking you to create such files. +- Be concise unless the situation requires elaboration. +- Use markdown for structure. +- Provide brief summaries when done; don't claim issues are resolved without verification. +- Clearly state when something could not be fully implemented. +- Never use emojis. +- Never create documentation files unless explicitly requested. ==== @@ -38,95 +31,29 @@ When using the planning tool: {{tools}} -# Tool Use Guidelines - -1. Assess what information you still need to proceed with the task. -2. Choose the most appropriate tool based on the task and the tool descriptions provided. Assess if you need additional information to proceed, and which of the available tools would be most effective for gathering this information. For example using the list_files tool is more effective than running a command like `ls` in the terminal. It's critical that you think about each available tool and use the one that best fits the current step in the task. -3. If multiple actions are needed, use one tool at a time per message to accomplish the task iteratively, with each tool use being informed by the result of the previous tool use. Do not assume the outcome of any tool use. Each step must be informed by the previous step's result. -4. Formulate your tool use using the format specified for each tool. -5. After each tool use, the system will respond with the result of that tool use. This result will provide you with the necessary information to continue your task or make further decisions. - -==== - -WORKFLOW TIPS - -1. Before editing, assess the scope of your changes and decide which tool to use. -2. For targeted edits, use the replace_in_file or edit tool. -3. For major overhauls or initial file creation, rely on write_file. -4. After making edits to code, consider what consequences this may have to other parts of the code, especially in files you have not yet seen. If appropriate, use the search tool to find files that might be affected by your changes. - -By thoughtfully selecting between write_file and edit/replace_in_file, and using the appropriate replacement blocks, you can make your file editing process smoother, safer, and more efficient. - -# Interface Change Considerations - -When modifying code structures, it's essential to understand and address all their usages: - -1. **Identify All References**: After changing any interface, structure, class definition, or feature flag: - - Use `search_files` with targeted regex patterns to find all usages of the changed component - - Look for imports, function calls, inheritances, or any other references to the modified code - - Don't assume you've seen all usage locations without performing a thorough search +# Tool use -2. **Verify Your Changes**: Always validate that your modifications work as expected: - - Run build commands appropriate for the project (e.g., `cargo check`, `npm run build`) - - Execute relevant tests to catch regressions (`cargo test`, `npm test`) - - Address any compiler errors or test failures that result from your changes +- Prefer specialized tools over shell commands (e.g., `list_files` over `ls`, `search_files` over `grep`). +- Use one tool at a time; let each result inform the next action. +- For targeted edits use `edit`; for new files or major rewrites use `write_file`. +- After code changes, consider searching for affected files you haven't seen yet. -3. **Track Modified Files**: Keep an overview of what you've changed: - - Use `execute_command` with git commands like `git status` to see which files have been modified - - Use `execute_command` with `git diff` to review specific changes within files - - This helps ensure all necessary updates are made consistently - -Remember that refactoring is not complete until all dependent code has been updated to work with your changes. - -# Code Review and Improvement - -After implementing working functionality, take time to review and improve the code that relates to your change, not unrelated imperfections. - -1. **Functionality Review**: Verify your implementation fully meets requirements: - - Double-check all acceptance criteria have been met - - Test edge cases and error conditions - - Verify all components interact correctly - -2. **Code Quality Improvements**: - - Look for repeated code that could be refactored into reusable functions - - Improve variable and function names for clarity - - Add or improve comments for complex logic - - Check for proper error handling - - Ensure consistent style and formatting - -3. **Performance Considerations**: - - Identify any inefficient operations or algorithms - - Consider resource usage (memory, CPU, network, disk) - - Look for unnecessary operations that could be optimized - -4. **Security and Robustness**: - - Check for input validation and sanitization - - Validate assumptions about data and environment - - Look for potential security issues - -Remember that the first working solution is rarely the best solution. Take time to refine your code once the core functionality is working. - -==== +# Git safety -WEB RESEARCH +- Never revert changes you didn't make unless explicitly requested. +- If you notice unexpected changes in files you're working on, stop and ask the user. +- Avoid destructive commands like `git reset --hard` or `git checkout --` without user approval. -When conducting web research, follow these steps: +# Long-running processes -1. Initial Search - - Start with web_search using specific, targeted queries - - Review search results to identify promising pages, taking into account the credibility and relevance of each source +When running dev servers, watchers, or long-running tests, always background them: -2. Deep Dive - - Use web_fetch to load full content of relevant pages - - Look for links to additional relevant resources within fetched pages - - Use web_fetch again to follow those links if needed - - Combine information from multiple sources +```bash +command > /path/to/log 2>&1 & +``` -Example scenarios when to use web research: -- Fetching the latest API or library documentation -- Reading source code on GitHub or other version control platforms -- Compiling accurate information from multiple sources +Never run blocking commands in the foreground. ==== -ALWAYS respond with your thoughts about what to do next first, then call the appropriate tool according to your reasoning. +When referencing files, use inline code with optional line numbers: `src/app.ts:42` diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index d4d2fb7a..78f2a400 100644 --- a/crates/code_assistant/src/acp/ui.rs +++ b/crates/code_assistant/src/acp/ui.rs @@ -9,7 +9,6 @@ use tokio::sync::{mpsc, oneshot}; use serde_json::{Map as JsonMap, Value as JsonValue}; use crate::acp::types::{fragment_to_content_block, map_tool_kind, map_tool_status}; -use crate::tools::core::registry::ToolRegistry; use crate::ui::{DisplayFragment, UIError, UiEvent, UserInterface}; /// UserInterface implementation that sends session/update notifications via ACP @@ -104,84 +103,18 @@ impl ToolCallState { } fn update_title_from_template(&mut self, tool_name: &str) { - let registry = ToolRegistry::global(); - if let Some(tool) = registry.get(tool_name) { - let spec = tool.spec(); - if let Some(template) = spec.title_template { - if let Some(new_title) = self.generate_title_from_template(template) { - self.title = Some(new_title); - } - } - } - } - - fn generate_title_from_template(&self, template: &str) -> Option { - let mut result = template.to_string(); - let mut has_substitution = false; - - // Find all {parameter_name} patterns and replace them - let re = regex::Regex::new(r"\{([^}]+)\}").ok()?; - - result = re - .replace_all(&result, |caps: ®ex::Captures| { - let param_name = &caps[1]; - if let Some(param_value) = self.parameters.get(param_name) { - let formatted_value = self.format_parameter_for_title(¶m_value.value); - if !formatted_value.trim().is_empty() { - has_substitution = true; - formatted_value - } else { - caps[0].to_string() // Keep placeholder if value is empty - } - } else { - caps[0].to_string() // Keep placeholder if parameter not found - } - }) - .to_string(); + // Convert parameters to HashMap for shared title function + let params: std::collections::HashMap = self + .parameters + .iter() + .map(|(k, v)| (k.clone(), v.value.clone())) + .collect(); - // Only return the new title if we actually made substitutions - if has_substitution { - Some(result) - } else { - None + if let Some(new_title) = crate::tools::core::generate_tool_title(tool_name, ¶ms) { + self.title = Some(new_title); } } - fn format_parameter_for_title(&self, value: &str) -> String { - const MAX_TITLE_LENGTH: usize = 50; - - let trimmed = value.trim(); - if trimmed.is_empty() { - return String::new(); - } - - // Try to parse as JSON and extract meaningful parts - if let Ok(json_val) = serde_json::from_str::(trimmed) { - match json_val { - serde_json::Value::Array(arr) if !arr.is_empty() => { - let first = arr[0].as_str().unwrap_or("...").to_string(); - if arr.len() > 1 { - format!("{} and {} more", first, arr.len() - 1) - } else { - first - } - } - serde_json::Value::String(s) => s, - _ => trimmed.to_string(), - } - } else { - trimmed.to_string() - } - .chars() - .take(MAX_TITLE_LENGTH) - .collect::() - + if trimmed.len() > MAX_TITLE_LENGTH { - "..." - } else { - "" - } - } - fn update_status( &mut self, status: acp::ToolCallStatus, @@ -302,7 +235,16 @@ impl ToolCallState { // For all other tools, put the full output as the primary content if let Some(output) = self.output_text() { if !output.is_empty() { - content.push(text_content(output)); + // For spawn_agent, try to render as markdown + if self.tool_name.as_deref() == Some("spawn_agent") { + if let Some(markdown) = render_sub_agent_output_as_markdown(&output) { + content.push(text_content(markdown)); + } else { + content.push(text_content(output)); + } + } else { + content.push(text_content(output)); + } } } } @@ -410,6 +352,57 @@ fn text_content(text: String) -> acp::ToolCallContent { } } +/// Render SubAgentOutput JSON as markdown for ACP display. +/// Returns None if the JSON is not valid SubAgentOutput. +fn render_sub_agent_output_as_markdown(json_str: &str) -> Option { + use crate::agent::sub_agent::{SubAgentOutput, SubAgentToolStatus}; + + let output: SubAgentOutput = serde_json::from_str(json_str).ok()?; + + let mut lines = Vec::new(); + // Render tool calls as a bullet list + if !output.tools.is_empty() { + for tool in &output.tools { + // Use title if available, otherwise tool name + let display_text = tool + .title + .as_ref() + .filter(|t| !t.is_empty()) + .cloned() + .or_else(|| tool.message.as_ref().filter(|m| !m.is_empty()).cloned()) + .unwrap_or_else(|| tool.name.replace('_', " ")); + + let suffix = match tool.status { + SubAgentToolStatus::Error => " (failed)", + _ => "", + }; + + lines.push(format!("- {display_text}{suffix}")); + } + } + + // Render error if present + if let Some(error) = &output.error { + lines.push(format!("**Error:** {error}")); + } + + // Render final response + if let Some(response) = &output.response { + if !response.is_empty() { + if !lines.is_empty() { + lines.push(String::new()); // Blank line before response + } + lines.push(response.clone()); + } + } + + if lines.is_empty() { + None + } else { + Some(lines.join("\n")) + } +} + fn resolve_path(path: &str, base_path: Option<&Path>) -> PathBuf { let candidate = PathBuf::from(path); if candidate.is_absolute() { @@ -715,7 +708,8 @@ impl UserInterface for ACPUserUI { | UiEvent::UpdatePendingMessage { .. } | UiEvent::ClearError | UiEvent::UpdateCurrentModel { .. } - | UiEvent::UpdateSandboxPolicy { .. } => { + | UiEvent::UpdateSandboxPolicy { .. } + | UiEvent::CancelSubAgent { .. } => { // These are UI management events, not relevant for ACP } UiEvent::DisplayError { message } => { diff --git a/crates/code_assistant/src/agent/mod.rs b/crates/code_assistant/src/agent/mod.rs index d0a2e493..0a297879 100644 --- a/crates/code_assistant/src/agent/mod.rs +++ b/crates/code_assistant/src/agent/mod.rs @@ -3,9 +3,11 @@ mod tests; pub mod persistence; pub mod runner; +pub mod sub_agent; pub mod types; pub use crate::types::ToolSyntax; // pub use persistence::FileStatePersistence; pub use runner::{Agent, AgentComponents}; +pub use sub_agent::{DefaultSubAgentRunner, SubAgentCancellationRegistry, SubAgentRunner}; pub use types::ToolExecution; diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index ee41e68f..3a84c9c1 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -5,7 +5,7 @@ use crate::permissions::PermissionMediator; use crate::persistence::{ChatMetadata, SessionModelConfig}; use crate::session::instance::SessionActivityState; use crate::session::SessionConfig; -use crate::tools::core::{ResourcesTracker, ToolContext, ToolRegistry, ToolScope}; +use crate::tools::core::{Render, ResourcesTracker, ToolContext, ToolRegistry, ToolScope}; use crate::tools::{generate_system_message, ParserRegistry, ToolRequest}; use crate::types::*; use crate::ui::{DisplayFragment, UiEvent, UserInterface}; @@ -30,6 +30,9 @@ pub struct AgentComponents { pub ui: Arc, pub state_persistence: Box, pub permission_handler: Option>, + + /// Optional sub-agent runner used by the `spawn_agent` tool. + pub sub_agent_runner: Option>, } use super::ToolSyntax; @@ -53,7 +56,9 @@ pub struct Agent { command_executor: Box, ui: Arc, state_persistence: Box, + permission_handler: Option>, + sub_agent_runner: Option>, // Store all messages exchanged message_history: Vec, // Store the history of tool executions @@ -111,6 +116,7 @@ impl Agent { ui, state_persistence, permission_handler, + sub_agent_runner, } = components; let mut this = Self { @@ -123,6 +129,7 @@ impl Agent { command_executor, state_persistence, permission_handler, + sub_agent_runner, message_history: Vec::new(), tool_executions: Vec::new(), cached_system_prompts: HashMap::new(), @@ -177,6 +184,27 @@ impl Agent { } } + /// Set the tool scope for this agent + pub fn set_tool_scope(&mut self, scope: ToolScope) { + self.tool_scope = scope; + } + + /// Set the session model configuration + pub fn set_session_model_config(&mut self, config: SessionModelConfig) { + self.session_model_config = Some(config); + } + + /// Set the session identity (id and name) for this agent + pub fn set_session_identity(&mut self, session_id: String, session_name: String) { + self.session_id = Some(session_id); + self.session_name = session_name; + } + + /// Get a reference to the message history + pub fn message_history(&self) -> &[Message] { + &self.message_history + } + /// Disable naming reminders (used for tests) #[cfg(test)] pub fn disable_naming_reminders(&mut self) { @@ -629,46 +657,282 @@ impl Agent { /// Executes a list of tool requests. /// Handles the "complete_task" action and appends tool results to message history. + /// Multiple `spawn_agent` read-only calls are executed concurrently for efficiency. async fn manage_tool_execution(&mut self, tool_requests: &[ToolRequest]) -> Result { - let mut content_blocks = Vec::new(); + // Check for complete_task first + if tool_requests.iter().any(|r| r.name == "complete_task") { + debug!("Task completed"); + return Ok(LoopFlow::Break); + } - for tool_request in tool_requests { - if tool_request.name == "complete_task" { - debug!("Task completed"); - return Ok(LoopFlow::Break); - } + // Partition into spawn_agent (read_only) and other tools, preserving indices - let start_time = Some(SystemTime::now()); - let result_block = match self.execute_tool(tool_request).await { - Ok(success) => ContentBlock::ToolResult { - tool_use_id: tool_request.id.clone(), - content: String::new(), // Will be filled dynamically in prepare_messages - 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: error_text, - is_error: Some(true), - start_time, - end_time: Some(SystemTime::now()), + let (parallel_indices, _sequential_indices): (Vec<_>, Vec<_>) = tool_requests + .iter() + .enumerate() + .partition(|(_, req)| self.can_run_in_parallel(req)); + + // Execute parallel spawn_agent tools concurrently if we have multiple + let parallel_results = if parallel_indices.len() > 1 { + debug!( + "Running {} spawn_agent 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() { + let result_block = + if parallel_indices.len() > 1 && parallel_indices.iter().any(|(i, _)| *i == idx) { + // This was a parallel spawn_agent - 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: "Internal error: missing parallel result".to_string(), + 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: String::new(), + 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: error_text, + is_error: Some(true), + start_time, + end_time: Some(SystemTime::now()), + } + } } - } - }; - content_blocks.push(result_block); + }; + content_blocks[idx] = Some(result_block); } - // Only add message if there were actual tool executions (not just complete_task) - if !content_blocks.is_empty() { - let result_message = Message::new_user_content(content_blocks); + // 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) } + /// Check if a tool request can be run in parallel with others. + /// Currently only spawn_agent with read_only mode is parallelizable. + fn can_run_in_parallel(&self, tool_request: &ToolRequest) -> bool { + if tool_request.name != "spawn_agent" { + return false; + } + // Check if mode is read_only (default if not specified) + let mode = tool_request.input["mode"].as_str().unwrap_or("read_only"); + mode == "read_only" + } + + /// Execute multiple spawn_agent tools 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; + + // Prepare context components that can be shared across parallel executions + let sub_agent_runner = self.sub_agent_runner.clone(); + let ui = self.ui.clone(); + + // Create futures for each spawn_agent tool + let futures: Vec<_> = tool_requests + .iter() + .map(|tool_request| { + let tool_id = tool_request.id.clone(); + let input = tool_request.input.clone(); + let runner = sub_agent_runner.clone(); + let ui_clone = ui.clone(); + + async move { + let start_time = Some(SystemTime::now()); + + // Execute spawn_agent directly without going through execute_tool + let result = Self::execute_spawn_agent_parallel( + runner.as_deref(), + &tool_id, + input, + ui_clone.as_ref(), + ) + .await; + + let (is_success, tool_execution) = result; + let end_time = Some(SystemTime::now()); + + ( + tool_id.clone(), + ContentBlock::ToolResult { + tool_use_id: tool_id, + content: String::new(), + is_error: if is_success { None } else { Some(true) }, + start_time, + end_time, + }, + 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 (tool_id, content_block, tool_execution) in results { + debug!("Parallel spawn_agent {} completed", tool_id); + self.tool_executions.push(tool_execution); + content_blocks.push(content_block); + } + + content_blocks + } + + /// Execute a spawn_agent tool in parallel without requiring &mut self. + async fn execute_spawn_agent_parallel( + sub_agent_runner: Option<&dyn crate::agent::SubAgentRunner>, + tool_id: &str, + input: serde_json::Value, + ui: &dyn UserInterface, + ) -> (bool, ToolExecution) { + use crate::tools::core::Tool; + + use crate::tools::core::ToolResult; + use crate::tools::impls::spawn_agent::{SpawnAgentInput, SpawnAgentTool}; + + // Update UI to show running status + let _ = ui + .send_event(UiEvent::UpdateToolStatus { + tool_id: tool_id.to_string(), + status: crate::ui::ToolStatus::Running, + message: None, + output: None, + }) + .await; + + // Parse input + let parsed_input: Result = serde_json::from_value(input.clone()); + let mut parsed_input = match parsed_input { + Ok(input) => input, + Err(e) => { + let error_msg = format!("Failed to parse spawn_agent input: {e}"); + let _ = ui + .send_event(UiEvent::UpdateToolStatus { + tool_id: tool_id.to_string(), + status: crate::ui::ToolStatus::Error, + message: Some(error_msg.clone()), + output: Some(error_msg.clone()), + }) + .await; + return ( + false, + ToolExecution::create_parse_error(tool_id.to_string(), error_msg), + ); + } + }; + + // Create a minimal context for spawn_agent (it only needs sub_agent_runner and tool_id) + // We use a dummy project manager and command executor since spawn_agent doesn't use them + let dummy_project_manager = crate::config::DefaultProjectManager::new(); + let dummy_command_executor = command_executor::DefaultCommandExecutor; + + let mut context = ToolContext { + project_manager: &dummy_project_manager, + command_executor: &dummy_command_executor, + plan: None, // spawn_agent doesn't use plan + ui: Some(ui), + tool_id: Some(tool_id.to_string()), + permission_handler: None, // Will be handled by sub-agent runner + sub_agent_runner, + }; + + let tool = SpawnAgentTool; + match tool.execute(&mut context, &mut parsed_input).await { + Ok(output) => { + let success = output.is_success(); + let status = if success { + crate::ui::ToolStatus::Success + } else { + crate::ui::ToolStatus::Error + }; + + let status_msg = output.status(); + let mut resources_tracker = ResourcesTracker::new(); + // Use render_for_ui() which returns JSON for spawn_agent (for custom renderer) + let ui_output = output.render_for_ui(&mut resources_tracker); + + let _ = ui + .send_event(UiEvent::UpdateToolStatus { + tool_id: tool_id.to_string(), + status, + message: Some(status_msg), + output: Some(ui_output), + }) + .await; + + let tool_execution = ToolExecution { + tool_request: ToolRequest { + id: tool_id.to_string(), + name: "spawn_agent".to_string(), + input, + start_offset: None, + end_offset: None, + }, + result: Box::new(output), + }; + + (success, tool_execution) + } + Err(e) => { + let error_msg = format!("spawn_agent failed: {e}"); + let _ = ui + .send_event(UiEvent::UpdateToolStatus { + tool_id: tool_id.to_string(), + status: crate::ui::ToolStatus::Error, + message: Some(error_msg.clone()), + output: Some(error_msg.clone()), + }) + .await; + + ( + false, + ToolExecution::create_parse_error(tool_id.to_string(), error_msg), + ) + } + } + } + /// Start a new agent task #[cfg(test)] pub async fn start_with_task(&mut self, task: String) -> Result<()> { @@ -1415,7 +1679,9 @@ impl Agent { plan: Some(&mut self.plan), ui: Some(self.ui.as_ref()), tool_id: Some(tool_request.id.clone()), + permission_handler: self.permission_handler.as_deref(), + sub_agent_runner: self.sub_agent_runner.as_deref(), }; // Execute the tool - could fail with ParseError or other errors @@ -1438,9 +1704,9 @@ impl Agent { // Generate status string from result let short_output = result.as_render().status(); - // Generate isolated output from result + // Generate output for UI display (may differ from LLM output for some tools) let mut resources_tracker = ResourcesTracker::new(); - let output = result.as_render().render(&mut resources_tracker); + let ui_output = result.as_render().render_for_ui(&mut resources_tracker); // Update tool status with result (skip for hidden tools) if !is_hidden { @@ -1449,7 +1715,7 @@ impl Agent { tool_id: tool_request.id.clone(), status, message: Some(short_output), - output: Some(output), + output: Some(ui_output), }) .await?; } diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs new file mode 100644 index 00000000..4a0a1ed1 --- /dev/null +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -0,0 +1,688 @@ +use crate::agent::persistence::AgentStatePersistence; +use crate::agent::{Agent, AgentComponents}; +use crate::config::DefaultProjectManager; +use crate::permissions::PermissionMediator; +use crate::persistence::SessionModelConfig; +use crate::session::SessionConfig; +use crate::tools::core::ToolScope; +use crate::ui::{ToolStatus, UiEvent, UserInterface}; +use anyhow::Result; +use command_executor::{CommandExecutor, DefaultCommandExecutor, SandboxedCommandExecutor}; +use llm::Message; +use sandbox::{SandboxContext, SandboxPolicy}; +use std::collections::HashMap; +use std::sync::{atomic::AtomicBool, atomic::Ordering, Arc, Mutex}; + +/// Cancellation registry keyed by the parent `spawn_agent` tool id. +#[derive(Default)] +pub struct SubAgentCancellationRegistry { + flags: Mutex>>, +} + +impl SubAgentCancellationRegistry { + 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 + } + + 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); + true + } else { + false + } + } + + pub fn unregister(&self, tool_id: &str) { + let mut flags = self.flags.lock().unwrap(); + flags.remove(tool_id); + } +} + +/// 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(()) + } +} + +/// Result from a sub-agent run, containing both the answer and UI output +#[derive(Debug, Clone)] +pub struct SubAgentResult { + /// The plain text answer for LLM context + pub answer: String, + /// The JSON output for UI display (tools list + response) + pub ui_output: String, +} + +/// Runs sub-agents with isolated history and streams a compact progress view into the parent tool UI. +#[async_trait::async_trait] +pub trait SubAgentRunner: Send + Sync { + async fn run( + &self, + parent_tool_id: &str, + instructions: String, + tool_scope: ToolScope, + require_file_references: bool, + ) -> Result; +} + +pub struct DefaultSubAgentRunner { + model_name: String, + session_config: SessionConfig, + sandbox_policy: SandboxPolicy, + sandbox_context: Arc, + cancellation_registry: Arc, + /// The parent UI to stream progress updates to. + ui: Arc, + /// Optional permission handler for sub-agent tool invocations. + permission_handler: Option>, +} + +impl DefaultSubAgentRunner { + pub fn new( + model_name: String, + session_config: SessionConfig, + sandbox_context: Arc, + cancellation_registry: Arc, + ui: Arc, + permission_handler: Option>, + ) -> Self { + let sandbox_policy = session_config.sandbox_policy.clone(); + Self { + model_name, + session_config, + sandbox_policy, + sandbox_context, + cancellation_registry, + ui, + permission_handler, + } + } + + fn build_sub_agent_ui( + &self, + parent_ui: Arc, + parent_tool_id: String, + cancelled: Arc, + ) -> Arc { + Arc::new(SubAgentUiAdapter::new(parent_ui, parent_tool_id, cancelled)) + } + + async fn build_agent( + &self, + parent_tool_id: &str, + ui: Arc, + 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?; + + // Create a fresh project manager, copying init_path if set. + let mut project_manager: Box = + Box::new(DefaultProjectManager::new()); + if let Some(path) = self.session_config.init_path.clone() { + let _ = project_manager.add_temporary_project(path); + } + + let command_executor: Box = { + let base: Box = Box::new(DefaultCommandExecutor); + if self.sandbox_policy.requires_restrictions() { + Box::new(SandboxedCommandExecutor::new( + base, + self.sandbox_policy.clone(), + Some(self.sandbox_context.clone()), + Some(format!("sub-agent:{parent_tool_id}")), + )) + } else { + base + } + }; + + let components = AgentComponents { + llm_provider, + project_manager, + command_executor, + ui, + state_persistence: Box::new(NoOpStatePersistence), + permission_handler, + sub_agent_runner: None, + }; + + let mut agent = Agent::new(components, self.session_config.clone()); + + // Configure for sub-agent use. + agent.set_tool_scope(tool_scope_for_subagent()); + + // Ensure it uses the same model name for prompt selection. + agent.set_session_model_config(SessionModelConfig::new(self.model_name.clone())); + + // Provide a stable session id so UI components that key off it don't break. + agent.set_session_identity(format!("sub-agent:{parent_tool_id}"), String::new()); + + // Initialize project trees, etc. + agent.init_project_context()?; + + Ok(agent) + } +} + +fn tool_scope_for_subagent() -> ToolScope { + // default; actual scope is set by caller via Agent::set_tool_scope + ToolScope::SubAgentReadOnly +} + +#[async_trait::async_trait] +impl SubAgentRunner for DefaultSubAgentRunner { + async fn run( + &self, + parent_tool_id: &str, + instructions: String, + tool_scope: ToolScope, + require_file_references: bool, + ) -> Result { + 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?; + + // 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")); + } + + // 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, + }) + } +} + +fn extract_last_assistant_text(messages: &[Message]) -> Option { + messages + .iter() + .rev() + .find(|m| matches!(m.role, llm::MessageRole::Assistant)) + .map(extract_text_from_message) +} + +/// Extract just the text content from a message, ignoring tool calls and other blocks +fn extract_text_from_message(message: &Message) -> String { + match &message.content { + llm::MessageContent::Text(text) => text.clone(), + llm::MessageContent::Structured(blocks) => { + let mut text_parts = Vec::new(); + for block in blocks { + match block { + llm::ContentBlock::Text { text, .. } => { + text_parts.push(text.as_str()); + } + _ => { + // Skip tool uses, thinking, tool results, images, etc. + } + } + } + text_parts.join("\n\n") + } + } +} + +fn has_file_references_with_line_ranges(text: &str) -> bool { + // Very lightweight heuristic: + // - backticked `path:10-20` OR raw path:10-20 + // - accept common extensions. + // Note: Rust regex doesn't support backreferences, so we use alternation instead. + let pattern = r"(?m)(`[\w./-]+\.(rs|ts|tsx|js|jsx|py|go|java|kt|swift|c|cc|cpp|h|hpp|md|toml|json|yaml|yml):(\d+)(-\d+)?`|[\w./-]+\.(rs|ts|tsx|js|jsx|py|go|java|kt|swift|c|cc|cpp|h|hpp|md|toml|json|yaml|yml):(\d+)(-\d+)?)"; + regex::Regex::new(pattern) + .map(|r| r.is_match(text)) + .unwrap_or(false) +} + +/// Structured representation of a sub-agent tool call for UI display and persistence. +/// This is serialized to JSON as the spawn_agent tool output. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct SubAgentToolCall { + pub name: String, + pub status: SubAgentToolStatus, + /// Human-readable title generated from tool's title_template and parameters + #[serde(skip_serializing_if = "Option::is_none")] + pub title: Option, + /// Status message (e.g., "Successfully loaded 2 file(s)") + #[serde(skip_serializing_if = "Option::is_none")] + pub message: Option, + /// Parameters collected during streaming (used to generate title) + #[serde(default, skip_serializing_if = "std::collections::HashMap::is_empty")] + pub parameters: std::collections::HashMap, +} + +/// Status of a sub-agent tool call +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SubAgentToolStatus { + Running, + Success, + Error, +} + +impl From for SubAgentToolStatus { + fn from(status: ToolStatus) -> Self { + match status { + ToolStatus::Pending | ToolStatus::Running => SubAgentToolStatus::Running, + ToolStatus::Success => SubAgentToolStatus::Success, + ToolStatus::Error => SubAgentToolStatus::Error, + } + } +} + +/// Current activity state of the sub-agent +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SubAgentActivity { + /// Waiting for LLM to start streaming + WaitingForLlm, + /// LLM is streaming response + Streaming, + /// Executing tools + ExecutingTools, + /// Completed successfully + Completed, + /// Was cancelled + Cancelled, + /// Encountered an error + Failed, +} + +/// Structured output for spawn_agent tool, serialized as JSON +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct SubAgentOutput { + pub tools: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub activity: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cancelled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + /// Final response from the sub-agent (set when completed) + #[serde(skip_serializing_if = "Option::is_none")] + pub response: Option, +} + +impl SubAgentOutput { + pub fn new() -> Self { + Self { + tools: Vec::new(), + activity: Some(SubAgentActivity::WaitingForLlm), + cancelled: None, + error: None, + response: None, + } + } + + pub fn to_json(&self) -> String { + serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string()) + } + + pub fn from_json(json: &str) -> Option { + serde_json::from_str(json).ok() + } +} + +impl Default for SubAgentOutput { + fn default() -> Self { + Self::new() + } +} + +/// A minimal UI adapter that captures sub-agent activity as structured data and streams it +/// into the parent `spawn_agent` tool block. +struct SubAgentUiAdapter { + parent: Arc, + parent_tool_id: String, + cancelled: Arc, + output: Mutex, + /// Map from tool_id to index in output.tools for fast lookup + tool_id_to_index: Mutex>, +} + +impl SubAgentUiAdapter { + fn new( + parent: Arc, + parent_tool_id: String, + cancelled: Arc, + ) -> Self { + Self { + parent, + parent_tool_id, + cancelled, + output: Mutex::new(SubAgentOutput::new()), + tool_id_to_index: Mutex::new(std::collections::HashMap::new()), + } + } + + async fn send_output_update(&self) { + let (json, tool_count, activity) = { + let output = self.output.lock().unwrap(); + (output.to_json(), output.tools.len(), output.activity) + }; + + tracing::debug!( + "SubAgentUiAdapter: Sending output update - {} tools, activity={:?}, json_len={}", + tool_count, + activity, + json.len() + ); + + let _ = self + .parent + .send_event(UiEvent::UpdateToolStatus { + tool_id: self.parent_tool_id.clone(), + status: ToolStatus::Running, + message: Some("Sub-agent running".to_string()), + output: Some(json), + }) + .await; + } + + fn add_tool_start(&self, name: &str, id: &str) { + let mut output = self.output.lock().unwrap(); + let mut id_map = self.tool_id_to_index.lock().unwrap(); + + // Check if tool already exists (avoid duplicates) + if id_map.contains_key(id) { + tracing::debug!( + "SubAgentUiAdapter: Tool already exists, skipping add: {} ({})", + name, + id + ); + return; + } + + // Add new tool as running + let index = output.tools.len(); + output.tools.push(SubAgentToolCall { + name: name.to_string(), + status: SubAgentToolStatus::Running, + title: None, + message: None, + parameters: std::collections::HashMap::new(), + }); + id_map.insert(id.to_string(), index); + tracing::debug!( + "SubAgentUiAdapter: Added tool {} ({}) at index {}, total tools: {}", + name, + id, + index, + output.tools.len() + ); + } + + fn add_tool_parameter(&self, tool_id: &str, name: &str, value: &str) { + let mut output = self.output.lock().unwrap(); + let id_map = self.tool_id_to_index.lock().unwrap(); + + if let Some(&index) = id_map.get(tool_id) { + if let Some(tool) = output.tools.get_mut(index) { + // Append to existing parameter value (streaming may send chunks) + let entry = tool.parameters.entry(name.to_string()).or_default(); + entry.push_str(value); + + // Update title from template using collected parameters + if let Some(new_title) = + crate::tools::core::generate_tool_title(&tool.name, &tool.parameters) + { + tool.title = Some(new_title); + } + } + } + } + + fn update_tool_status(&self, tool_id: &str, status: ToolStatus, message: Option) { + let mut output = self.output.lock().unwrap(); + let mut id_map = self.tool_id_to_index.lock().unwrap(); + + // Find tool by id and update its status + if let Some(&index) = id_map.get(tool_id) { + if let Some(tool) = output.tools.get_mut(index) { + tracing::debug!( + "SubAgentUiAdapter: Updating tool {} status to {:?}", + tool.name, + status + ); + tool.status = status.into(); + tool.message = message; + } + } else { + // Tool not found - this can happen if UpdateToolStatus arrives before ToolName + // In this case, we should add the tool + tracing::warn!( + "SubAgentUiAdapter: UpdateToolStatus for unknown tool_id={}, status={:?}. Adding placeholder.", + tool_id, + status + ); + let index = output.tools.len(); + output.tools.push(SubAgentToolCall { + name: format!("tool_{}", tool_id.chars().take(8).collect::()), + status: status.into(), + title: None, + message, + parameters: std::collections::HashMap::new(), + }); + id_map.insert(tool_id.to_string(), index); + } + } + + fn set_cancelled(&self) { + let mut output = self.output.lock().unwrap(); + output.cancelled = Some(true); + } + + fn set_error(&self, error: String) { + let mut output = self.output.lock().unwrap(); + output.error = Some(error); + output.activity = Some(SubAgentActivity::Failed); + } + + fn set_activity(&self, activity: SubAgentActivity) { + let mut output = self.output.lock().unwrap(); + output.activity = Some(activity); + } + + fn set_response(&self, response: String) { + let mut output = self.output.lock().unwrap(); + output.response = Some(response); + output.activity = Some(SubAgentActivity::Completed); + } + + /// Get the final JSON output including response + fn get_final_output(&self) -> String { + let output = self.output.lock().unwrap(); + output.to_json() + } +} + +#[async_trait::async_trait] +impl UserInterface for SubAgentUiAdapter { + async fn send_event(&self, event: UiEvent) -> Result<(), crate::ui::UIError> { + match &event { + UiEvent::UpdateToolStatus { + tool_id, + status, + message, + .. + } => { + tracing::debug!( + "SubAgentUiAdapter: UpdateToolStatus event - tool_id={}, status={:?}", + tool_id, + status + ); + self.update_tool_status(tool_id, *status, message.clone()); + // If a tool is running, we're executing tools + if *status == ToolStatus::Running { + self.set_activity(SubAgentActivity::ExecutingTools); + } + self.send_output_update().await; + } + + UiEvent::StreamingStarted(_) => { + tracing::debug!("SubAgentUiAdapter: StreamingStarted"); + self.set_activity(SubAgentActivity::Streaming); + self.send_output_update().await; + } + UiEvent::StreamingStopped { + cancelled, error, .. + } => { + tracing::debug!( + "SubAgentUiAdapter: StreamingStopped - cancelled={}, error={:?}", + cancelled, + error + ); + if *cancelled { + self.set_cancelled(); + self.set_activity(SubAgentActivity::Cancelled); + } else if let Some(err) = error { + self.set_error(err.clone()); + // activity already set to Failed in set_error + } else { + // Streaming stopped normally - will likely execute tools or complete + // The activity will be updated by tool execution or completion + } + self.send_output_update().await; + } + _ => { + // Ignore other events; they belong to the sub-agent's isolated transcript. + } + } + + Ok(()) + } + + fn display_fragment( + &self, + fragment: &crate::ui::DisplayFragment, + ) -> Result<(), crate::ui::UIError> { + use crate::ui::DisplayFragment; + + match fragment { + DisplayFragment::ToolName { name, id } => { + // A sub-agent tool is starting - capture it in our internal state. + // This is called during LLM streaming when the tool name is parsed. + // + // Note: We don't notify the parent UI here because display_fragment() is + // synchronous. The parent UI will be notified when runner.rs sends + // UiEvent::UpdateToolStatus with Running status just before tool execution + // starts. At that point, send_event() calls send_output_update() which + // forwards our accumulated state (including this tool) to the parent. + tracing::debug!( + "SubAgentUiAdapter: ToolName fragment received: {} ({})", + name, + id + ); + self.add_tool_start(name, id); + } + DisplayFragment::ToolParameter { + name, + value, + tool_id, + } => { + // Capture parameters to generate tool titles + self.add_tool_parameter(tool_id, name, value); + } + _ => { + // Ignore other fragments (text, etc.) + // They belong to the sub-agent's isolated transcript + } + } + + Ok(()) + } + + fn should_streaming_continue(&self) -> bool { + !self.cancelled.load(Ordering::SeqCst) && self.parent.should_streaming_continue() + } + + fn notify_rate_limit(&self, _seconds_remaining: u64) {} + + fn clear_rate_limit(&self) {} + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index 254023e0..f04b167b 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -395,6 +395,7 @@ async fn test_unknown_tool_error_handling() -> Result<()> { ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -516,6 +517,7 @@ async fn test_invalid_xml_tool_error_handling() -> Result<()> { ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -640,6 +642,7 @@ async fn test_parse_error_handling() -> Result<()> { ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -778,6 +781,7 @@ async fn test_write_file_outside_root_error_masks_paths() -> Result<()> { ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -878,6 +882,7 @@ async fn test_context_compaction_inserts_summary() -> Result<()> { ui: ui.clone(), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -995,6 +1000,7 @@ async fn test_compaction_prompt_not_persisted_in_history() -> Result<()> { ui: ui.clone(), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -1121,6 +1127,7 @@ async fn test_context_compaction_uses_only_messages_after_previous_summary() -> ui: ui.clone(), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -1596,6 +1603,7 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { ui, state_persistence, permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -1944,6 +1952,7 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -1998,6 +2007,7 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -2066,6 +2076,7 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { @@ -2118,6 +2129,7 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() ui: Arc::new(MockUI::default()), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, }; let session_config = SessionConfig { diff --git a/crates/code_assistant/src/mcp/handler.rs b/crates/code_assistant/src/mcp/handler.rs index 8a1ad144..78cab502 100644 --- a/crates/code_assistant/src/mcp/handler.rs +++ b/crates/code_assistant/src/mcp/handler.rs @@ -269,6 +269,7 @@ impl MessageHandler { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Invoke the tool diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 2bbed095..175aa54e 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -5,6 +5,7 @@ use std::sync::{Arc, Mutex}; use tokio::task::JoinHandle; // Agent instances are created on-demand, no need to import +use crate::agent::SubAgentCancellationRegistry; use crate::persistence::{ChatMetadata, ChatSession}; use crate::ui::gpui::elements::MessageRole; use crate::ui::streaming::create_stream_processor; @@ -53,6 +54,9 @@ pub struct SessionInstance { /// Tracks sandbox-approved roots for this session pub sandbox_context: Arc, + + /// Cancellation registry for sub-agents running in agent tasks + pub sub_agent_cancellation_registry: Arc, } impl SessionInstance { @@ -71,9 +75,16 @@ impl SessionInstance { activity_state: Arc::new(Mutex::new(SessionActivityState::Idle)), pending_message: Arc::new(Mutex::new(None)), sandbox_context, + sub_agent_cancellation_registry: Arc::new(SubAgentCancellationRegistry::default()), } } + /// Cancel a running sub-agent by its tool ID + /// Returns true if a sub-agent was found and cancelled, false otherwise + pub fn cancel_sub_agent(&self, tool_id: &str) -> bool { + self.sub_agent_cancellation_registry.cancel(tool_id) + } + /// Get the current activity state pub fn get_activity_state(&self) -> SessionActivityState { self.activity_state.lock().unwrap().clone() @@ -398,7 +409,11 @@ impl SessionInstance { }; let short_output = execution.result.as_render().status(); - let output = execution.result.as_render().render(&mut resources_tracker); + // Use render_for_ui() to get the UI-specific output (e.g., JSON for spawn_agent) + let output = execution + .result + .as_render() + .render_for_ui(&mut resources_tracker); tool_results.push(crate::ui::ui_events::ToolResultData { tool_id: execution.tool_request.id, diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index 84a085f9..d079313c 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use std::time::SystemTime; use tokio::sync::Mutex; -use crate::agent::{Agent, AgentComponents}; +use crate::agent::{Agent, AgentComponents, DefaultSubAgentRunner, SubAgentCancellationRegistry}; use crate::config::ProjectManager; use crate::permissions::PermissionMediator; use crate::persistence::{ @@ -315,9 +315,32 @@ impl SessionManager { let sandboxed_project_manager = Box::new(crate::config::SandboxAwareProjectManager::new( project_manager, - sandbox_context_clone, + sandbox_context_clone.clone(), )); + // Get the cancellation registry from the session instance + let sub_agent_cancellation_registry = self + .active_sessions + .get(session_id) + .map(|s| s.sub_agent_cancellation_registry.clone()) + .unwrap_or_else(|| Arc::new(SubAgentCancellationRegistry::default())); + + // Create sub-agent runner + let model_name_for_subagent = session_state + .model_config + .as_ref() + .map(|c| c.model_name.clone()) + .unwrap_or_else(|| self.default_model_name.clone()); + let sub_agent_runner: Arc = + Arc::new(DefaultSubAgentRunner::new( + model_name_for_subagent, + session_config.clone(), + sandbox_context_clone, + sub_agent_cancellation_registry.clone(), + proxy_ui.clone(), + permission_handler.clone(), + )); + let components = AgentComponents { llm_provider, project_manager: sandboxed_project_manager, @@ -325,6 +348,7 @@ impl SessionManager { ui: proxy_ui, state_persistence: state_storage, permission_handler, + sub_agent_runner: Some(sub_agent_runner), }; let mut agent = Agent::new(components, session_config.clone()); @@ -437,6 +461,17 @@ impl SessionManager { Ok(()) } + /// Cancel a running sub-agent by its tool ID + /// Returns Ok(true) if the sub-agent was found and cancelled, + /// Ok(false) if the sub-agent was not found (may have already completed) + pub fn cancel_sub_agent(&self, session_id: &str, tool_id: &str) -> Result { + if let Some(session_instance) = self.active_sessions.get(session_id) { + Ok(session_instance.cancel_sub_agent(tool_id)) + } else { + Err(anyhow::anyhow!("Session not found: {}", session_id)) + } + } + /// Get a session instance by ID pub fn get_session(&self, session_id: &str) -> Option<&SessionInstance> { self.active_sessions.get(session_id) diff --git a/crates/code_assistant/src/tests/format_on_save_tests.rs b/crates/code_assistant/src/tests/format_on_save_tests.rs index d3952d21..9d568592 100644 --- a/crates/code_assistant/src/tests/format_on_save_tests.rs +++ b/crates/code_assistant/src/tests/format_on_save_tests.rs @@ -198,6 +198,7 @@ async fn test_edit_tool_parameter_update_after_formatting() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Test editing: search is formatted (matches file), replacement is unformatted @@ -271,6 +272,7 @@ async fn test_write_file_with_format_on_save() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Test writing a Rust file @@ -343,6 +345,7 @@ async fn test_replace_in_file_with_format_on_save() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Diff has two SEARCH/REPLACE blocks; replacements are unformatted (missing spaces around '=') @@ -417,6 +420,7 @@ async fn test_no_format_when_pattern_doesnt_match() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Test editing a .txt file (should not be formatted) @@ -493,6 +497,7 @@ async fn test_format_on_save_multiple_patterns() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Test editing JS file @@ -589,6 +594,7 @@ async fn test_format_on_save_glob_patterns() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; let tool = EditTool; @@ -678,6 +684,7 @@ async fn test_format_on_save_with_conflicting_matches() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, }; // Test that the tool handles potential conflicts gracefully diff --git a/crates/code_assistant/src/tests/mocks.rs b/crates/code_assistant/src/tests/mocks.rs index b081109f..0eaa8c08 100644 --- a/crates/code_assistant/src/tests/mocks.rs +++ b/crates/code_assistant/src/tests/mocks.rs @@ -231,6 +231,7 @@ pub fn create_test_tool_context<'a>( ui, tool_id, permission_handler: None, + sub_agent_runner: None, } } @@ -1128,6 +1129,7 @@ impl ToolTestFixture { ui: self.ui.as_ref().map(|ui| ui as &dyn UserInterface), tool_id: self.tool_id.clone(), permission_handler: self.permission_handler.as_deref(), + sub_agent_runner: None, } } diff --git a/crates/code_assistant/src/tests/mod.rs b/crates/code_assistant/src/tests/mod.rs index 1e173302..0f02f429 100644 --- a/crates/code_assistant/src/tests/mod.rs +++ b/crates/code_assistant/src/tests/mod.rs @@ -3,4 +3,5 @@ pub mod gitignore_tests; pub mod integration_tests; pub mod mocks; pub mod sandbox_tests; +pub mod sub_agent_tests; pub mod utils; diff --git a/crates/code_assistant/src/tests/sub_agent_tests.rs b/crates/code_assistant/src/tests/sub_agent_tests.rs new file mode 100644 index 00000000..89629917 --- /dev/null +++ b/crates/code_assistant/src/tests/sub_agent_tests.rs @@ -0,0 +1,374 @@ +//! Tests for the sub-agent feature (spawn_agent tool). + +use crate::agent::sub_agent::{SubAgentResult, SubAgentRunner}; +use crate::agent::SubAgentCancellationRegistry; +use crate::tools::core::ToolScope; +use crate::tools::impls::spawn_agent::{SpawnAgentInput, SpawnAgentOutput}; +use anyhow::Result; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +/// A mock sub-agent runner for testing. +struct MockSubAgentRunner { + call_count: AtomicUsize, + concurrent_count: AtomicUsize, + max_concurrent: AtomicUsize, + delay_ms: u64, + response: String, +} + +impl MockSubAgentRunner { + fn new(delay_ms: u64, response: &str) -> Self { + Self { + call_count: AtomicUsize::new(0), + concurrent_count: AtomicUsize::new(0), + max_concurrent: AtomicUsize::new(0), + delay_ms, + response: response.to_string(), + } + } + + fn call_count(&self) -> usize { + self.call_count.load(Ordering::SeqCst) + } + + fn max_concurrent(&self) -> usize { + self.max_concurrent.load(Ordering::SeqCst) + } +} + +#[async_trait::async_trait] +impl SubAgentRunner for MockSubAgentRunner { + async fn run( + &self, + _parent_tool_id: &str, + _instructions: String, + _tool_scope: ToolScope, + _require_file_references: bool, + ) -> Result { + // Track concurrent executions + let current = self.concurrent_count.fetch_add(1, Ordering::SeqCst) + 1; + + // Update max concurrent count + let mut max = self.max_concurrent.load(Ordering::SeqCst); + while current > max { + match self.max_concurrent.compare_exchange_weak( + max, + current, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => break, + Err(m) => max = m, + } + } + + self.call_count.fetch_add(1, Ordering::SeqCst); + + // Simulate work + if self.delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(self.delay_ms)).await; + } + + // Track completion + self.concurrent_count.fetch_sub(1, Ordering::SeqCst); + + Ok(SubAgentResult { + answer: self.response.clone(), + ui_output: format!(r#"{{"tools":[],"response":"{}"}}"#, self.response), + }) + } +} + +#[test] +fn test_spawn_agent_output_render() { + use crate::tools::core::{Render, ResourcesTracker}; + + // Test successful output + let output = SpawnAgentOutput { + answer: "The answer is 42.".to_string(), + error: None, + ui_output: None, + }; + + let mut tracker = ResourcesTracker::new(); + let rendered = output.render(&mut tracker); + assert_eq!(rendered, "The answer is 42."); + assert_eq!(output.status(), "Sub-agent completed"); + + // Test cancelled output (cancellation is now just an error) + let output = SpawnAgentOutput { + answer: String::new(), + error: Some("Cancelled by user".to_string()), + ui_output: None, + }; + + let rendered = output.render(&mut tracker); + assert_eq!(rendered, "Sub-agent error: Cancelled by user"); + assert!(output.status().contains("Cancelled by user")); + + // Test error output + let output = SpawnAgentOutput { + answer: String::new(), + error: Some("Connection failed".to_string()), + ui_output: None, + }; + + let rendered = output.render(&mut tracker); + assert_eq!(rendered, "Sub-agent error: Connection failed"); + assert!(output.status().contains("Connection failed")); + + // Test render_for_ui with ui_output set + let output_with_ui = SpawnAgentOutput { + answer: "Plain text answer".to_string(), + error: None, + ui_output: Some(r#"{"tools":[],"response":"Markdown response"}"#.to_string()), + }; + + let mut tracker = ResourcesTracker::new(); + // render() returns plain text for LLM + assert_eq!(output_with_ui.render(&mut tracker), "Plain text answer"); + // render_for_ui() returns JSON for UI + assert_eq!( + output_with_ui.render_for_ui(&mut tracker), + r#"{"tools":[],"response":"Markdown response"}"# + ); + + // Test render_for_ui falls back to render() when ui_output is None + let output_without_ui = SpawnAgentOutput { + answer: "Just the answer".to_string(), + error: None, + ui_output: None, + }; + assert_eq!( + output_without_ui.render_for_ui(&mut tracker), + "Just the answer" + ); +} + +#[test] +fn test_spawn_agent_input_parsing() { + // Test with all parameters + let json = serde_json::json!({ + "instructions": "Find all TODO comments", + "require_file_references": true, + "mode": "read_only" + }); + + let input: SpawnAgentInput = serde_json::from_value(json).unwrap(); + assert_eq!(input.instructions, "Find all TODO comments"); + assert!(input.require_file_references); + assert_eq!(input.mode, "read_only"); + + // Test with defaults + let json = serde_json::json!({ + "instructions": "Search for patterns" + }); + + let input: SpawnAgentInput = serde_json::from_value(json).unwrap(); + assert_eq!(input.instructions, "Search for patterns"); + assert!(!input.require_file_references); + 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"); + + let result = runner + .run( + "tool-1", + "test".to_string(), + ToolScope::SubAgentReadOnly, + false, + ) + .await; + + assert!(result.is_ok()); + assert_eq!(result.unwrap().answer, "Test response"); + assert_eq!(runner.call_count(), 1); +} + +#[tokio::test] +async fn test_parallel_sub_agent_execution() { + use futures::future::join_all; + + let runner = Arc::new(MockSubAgentRunner::new(50, "Parallel result")); + + // Run multiple sub-agents in parallel + let futures: Vec<_> = (0..4) + .map(|i| { + let runner = runner.clone(); + async move { + runner + .run( + &format!("tool-{i}"), + format!("Task {i}"), + ToolScope::SubAgentReadOnly, + false, + ) + .await + } + }) + .collect(); + + let results = join_all(futures).await; + + // All should succeed + assert_eq!(results.len(), 4); + for result in results { + assert!(result.is_ok()); + assert_eq!(result.unwrap().answer, "Parallel result"); + } + + // Verify they ran in parallel (max concurrent should be > 1) + assert_eq!(runner.call_count(), 4); + // With a 50ms delay and parallel execution, we should see concurrency + assert!( + runner.max_concurrent() > 1, + "Expected parallel execution, but max concurrent was {}", + runner.max_concurrent() + ); +} + +#[test] +fn test_tool_scope_for_sub_agent() { + use crate::tools::core::ToolRegistry; + + // spawn_agent should only be available in main agent scopes + let registry = ToolRegistry::global(); + + // Check spawn_agent is available in normal agent scope + let tool = registry.get("spawn_agent"); + assert!(tool.is_some(), "spawn_agent should be registered"); + + // Helper to get tool names for a scope + let get_tools_for_scope = |scope: ToolScope| -> Vec { + registry + .get_tool_definitions_for_scope(scope) + .iter() + .map(|t| t.name.clone()) + .collect() + }; + + // Check spawn_agent is NOT available in sub-agent scope + let available = get_tools_for_scope(ToolScope::SubAgentReadOnly); + assert!( + !available.contains(&"spawn_agent".to_string()), + "spawn_agent should not be available in SubAgentReadOnly scope" + ); + + let available = get_tools_for_scope(ToolScope::SubAgentDefault); + assert!( + !available.contains(&"spawn_agent".to_string()), + "spawn_agent should not be available in SubAgentDefault scope" + ); + + // Check read-only tools are available in SubAgentReadOnly scope + let available = get_tools_for_scope(ToolScope::SubAgentReadOnly); + assert!(available.contains(&"search_files".to_string())); + assert!(available.contains(&"read_files".to_string())); + assert!(available.contains(&"list_files".to_string())); + assert!(available.contains(&"glob_files".to_string())); + assert!(available.contains(&"web_fetch".to_string())); + assert!(available.contains(&"web_search".to_string())); + assert!(available.contains(&"perplexity_ask".to_string())); + + // Check write tools are NOT available in SubAgentReadOnly scope + assert!(!available.contains(&"write_file".to_string())); + assert!(!available.contains(&"edit".to_string())); + assert!(!available.contains(&"delete_files".to_string())); + assert!(!available.contains(&"execute_command".to_string())); +} + +#[test] +fn test_can_run_in_parallel_logic() { + use crate::tools::ToolRequest; + + // This tests the logic of can_run_in_parallel without needing a full agent + + // spawn_agent with read_only mode should be parallelizable + let request = ToolRequest { + id: "test-1".to_string(), + name: "spawn_agent".to_string(), + input: serde_json::json!({ + "instructions": "test", + "mode": "read_only" + }), + start_offset: None, + end_offset: None, + }; + + let mode = request.input["mode"].as_str().unwrap_or("read_only"); + assert_eq!(mode, "read_only"); + assert!(request.name == "spawn_agent" && mode == "read_only"); + + // spawn_agent with default mode (no mode specified) should default to read_only + let request = ToolRequest { + id: "test-2".to_string(), + name: "spawn_agent".to_string(), + input: serde_json::json!({ + "instructions": "test" + }), + start_offset: None, + end_offset: None, + }; + + let mode = request.input["mode"].as_str().unwrap_or("read_only"); + assert_eq!(mode, "read_only"); + + // spawn_agent with non-read_only mode should NOT be parallelizable + let request = ToolRequest { + id: "test-3".to_string(), + name: "spawn_agent".to_string(), + input: serde_json::json!({ + "instructions": "test", + "mode": "default" + }), + start_offset: None, + end_offset: None, + }; + + let mode = request.input["mode"].as_str().unwrap_or("read_only"); + assert_eq!(mode, "default"); + assert!(!(request.name == "spawn_agent" && mode == "read_only")); + + // Other tools should NOT be parallelizable + let request = ToolRequest { + id: "test-4".to_string(), + name: "read_files".to_string(), + input: serde_json::json!({ + "paths": ["test.txt"] + }), + start_offset: None, + end_offset: None, + }; + + assert!(request.name != "spawn_agent"); +} diff --git a/crates/code_assistant/src/tools/core/mod.rs b/crates/code_assistant/src/tools/core/mod.rs index 707c8191..085de914 100644 --- a/crates/code_assistant/src/tools/core/mod.rs +++ b/crates/code_assistant/src/tools/core/mod.rs @@ -4,6 +4,7 @@ pub mod registry; pub mod render; pub mod result; pub mod spec; +pub mod title; pub mod tool; // Re-export all core components for easier imports @@ -12,4 +13,5 @@ pub use registry::ToolRegistry; pub use render::{Render, ResourcesTracker}; pub use result::ToolResult; pub use spec::{ToolScope, ToolSpec}; +pub use title::generate_tool_title; pub use tool::{Tool, ToolContext}; diff --git a/crates/code_assistant/src/tools/core/registry.rs b/crates/code_assistant/src/tools/core/registry.rs index fb2a5d43..ec5d47b4 100644 --- a/crates/code_assistant/src/tools/core/registry.rs +++ b/crates/code_assistant/src/tools/core/registry.rs @@ -73,7 +73,8 @@ impl ToolRegistry { use crate::tools::impls::{ DeleteFilesTool, EditTool, ExecuteCommandTool, GlobFilesTool, ListFilesTool, ListProjectsTool, NameSessionTool, PerplexityAskTool, ReadFilesTool, ReplaceInFileTool, - SearchFilesTool, UpdatePlanTool, WebFetchTool, WebSearchTool, WriteFileTool, + SearchFilesTool, SpawnAgentTool, UpdatePlanTool, WebFetchTool, WebSearchTool, + WriteFileTool, }; // Register all tools - the ToolScope system will filter which ones are available @@ -88,6 +89,7 @@ impl ToolRegistry { self.register(Box::new(ReadFilesTool)); self.register(Box::new(ReplaceInFileTool)); self.register(Box::new(SearchFilesTool)); + self.register(Box::new(SpawnAgentTool)); self.register(Box::new(UpdatePlanTool)); self.register(Box::new(WebFetchTool)); self.register(Box::new(WebSearchTool)); diff --git a/crates/code_assistant/src/tools/core/render.rs b/crates/code_assistant/src/tools/core/render.rs index 91732741..30955242 100644 --- a/crates/code_assistant/src/tools/core/render.rs +++ b/crates/code_assistant/src/tools/core/render.rs @@ -5,9 +5,17 @@ pub trait Render: Send + Sync + 'static { /// Generate a short status message for display in action history fn status(&self) -> String; - /// Format the detailed output, with awareness of other tool results + /// Format the detailed output for LLM context (tool result in conversation) /// The resources_tracker helps detect and handle redundant output fn render(&self, resources_tracker: &mut ResourcesTracker) -> String; + + /// Format the output for UI display in tool blocks. + /// By default, returns the same as render(). + /// Override this for tools that need different UI representation (e.g., spawn_agent + /// returns JSON for custom rendering while render() returns plain text for LLM). + fn render_for_ui(&self, resources_tracker: &mut ResourcesTracker) -> String { + self.render(resources_tracker) + } } /// Tracks resources that have been included in tool outputs to prevent redundant display diff --git a/crates/code_assistant/src/tools/core/spec.rs b/crates/code_assistant/src/tools/core/spec.rs index f84ab44f..55368c3a 100644 --- a/crates/code_assistant/src/tools/core/spec.rs +++ b/crates/code_assistant/src/tools/core/spec.rs @@ -7,6 +7,10 @@ pub enum ToolScope { Agent, /// Tool can be used in the agent when configured for diff blocks format AgentWithDiffBlocks, + /// Tool scope for sub-agents running in a restricted, read-only mode + SubAgentReadOnly, + /// Tool scope for sub-agents running with broader permissions (reserved for future use) + SubAgentDefault, } /// Specification for a tool, including metadata diff --git a/crates/code_assistant/src/tools/core/title.rs b/crates/code_assistant/src/tools/core/title.rs new file mode 100644 index 00000000..bfbae898 --- /dev/null +++ b/crates/code_assistant/src/tools/core/title.rs @@ -0,0 +1,147 @@ +//! Tool title generation utilities +//! +//! This module provides functions for generating human-readable tool titles +//! from tool specs and their parameters. Used by ACP mode and sub-agent output rendering. + +use crate::tools::core::ToolRegistry; +use std::collections::HashMap; + +/// Maximum length for title parameter values before truncation +const MAX_TITLE_LENGTH: usize = 50; + +/// Generate a tool title from a tool's title_template and its parameters. +/// +/// Returns the generated title if template substitution was successful, +/// otherwise returns None (caller should fall back to tool name or other default). +pub fn generate_tool_title( + tool_name: &str, + parameters: &HashMap, +) -> Option { + let registry = ToolRegistry::global(); + let tool = registry.get(tool_name)?; + let spec = tool.spec(); + let template = spec.title_template?; + + generate_title_from_template(template, parameters) +} + +/// Generate a title by substituting {parameter_name} placeholders in the template. +/// +/// Returns Some(title) if at least one substitution was made, None otherwise. +pub fn generate_title_from_template( + template: &str, + parameters: &HashMap, +) -> Option { + let mut result = template.to_string(); + let mut has_substitution = false; + + // Find all {parameter_name} patterns and replace them + let re = regex::Regex::new(r"\{([^}]+)\}").ok()?; + + result = re + .replace_all(&result, |caps: ®ex::Captures| { + let param_name = &caps[1]; + if let Some(param_value) = parameters.get(param_name) { + let formatted_value = format_parameter_for_title(param_value); + if !formatted_value.trim().is_empty() { + has_substitution = true; + formatted_value + } else { + caps[0].to_string() // Keep placeholder if value is empty + } + } else { + caps[0].to_string() // Keep placeholder if parameter not found + } + }) + .to_string(); + + // Only return the new title if we actually made substitutions + if has_substitution { + Some(result) + } else { + None + } +} + +/// Format a parameter value for display in a title. +/// +/// Handles JSON arrays (showing first item + count), truncates long values, etc. +pub fn format_parameter_for_title(value: &str) -> String { + let trimmed = value.trim(); + if trimmed.is_empty() { + return String::new(); + } + + // Try to parse as JSON and extract meaningful parts + let formatted = if let Ok(json_val) = serde_json::from_str::(trimmed) { + match json_val { + serde_json::Value::Array(arr) if !arr.is_empty() => { + let first = arr[0].as_str().unwrap_or("...").to_string(); + if arr.len() > 1 { + format!("{} and {} more", first, arr.len() - 1) + } else { + first + } + } + serde_json::Value::String(s) => s, + _ => trimmed.to_string(), + } + } else { + trimmed.to_string() + }; + + // Truncate if needed + if formatted.len() > MAX_TITLE_LENGTH { + format!( + "{}...", + formatted.chars().take(MAX_TITLE_LENGTH).collect::() + ) + } else { + formatted + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_format_parameter_simple() { + assert_eq!(format_parameter_for_title("hello"), "hello"); + assert_eq!(format_parameter_for_title(" hello "), "hello"); + assert_eq!(format_parameter_for_title(""), ""); + } + + #[test] + fn test_format_parameter_json_array() { + assert_eq!(format_parameter_for_title(r#"["file1.txt"]"#), "file1.txt"); + assert_eq!( + format_parameter_for_title(r#"["file1.txt", "file2.txt", "file3.txt"]"#), + "file1.txt and 2 more" + ); + } + + #[test] + fn test_format_parameter_truncation() { + let long_value = "a".repeat(100); + let result = format_parameter_for_title(&long_value); + assert!(result.len() <= MAX_TITLE_LENGTH + 3); // +3 for "..." + assert!(result.ends_with("...")); + } + + #[test] + fn test_generate_title_from_template() { + let mut params = HashMap::new(); + params.insert("path".to_string(), "src/main.rs".to_string()); + + let result = generate_title_from_template("Editing {path}", ¶ms); + assert_eq!(result, Some("Editing src/main.rs".to_string())); + } + + #[test] + fn test_generate_title_no_substitution() { + let params = HashMap::new(); + let result = generate_title_from_template("Editing {path}", ¶ms); + assert_eq!(result, None); // No substitution made + } +} diff --git a/crates/code_assistant/src/tools/core/tool.rs b/crates/code_assistant/src/tools/core/tool.rs index 8236958e..4b00994d 100644 --- a/crates/code_assistant/src/tools/core/tool.rs +++ b/crates/code_assistant/src/tools/core/tool.rs @@ -21,6 +21,9 @@ pub struct ToolContext<'a> { pub tool_id: Option, /// Optional permission handler for potentially sensitive operations pub permission_handler: Option<&'a dyn PermissionMediator>, + + /// Optional sub-agent runner used by the `spawn_agent` tool. + pub sub_agent_runner: Option<&'a dyn crate::agent::SubAgentRunner>, } /// Core trait for tools, defining the execution interface diff --git a/crates/code_assistant/src/tools/impls/glob_files.rs b/crates/code_assistant/src/tools/impls/glob_files.rs index 7cedc7ce..f0759c75 100644 --- a/crates/code_assistant/src/tools/impls/glob_files.rs +++ b/crates/code_assistant/src/tools/impls/glob_files.rs @@ -89,10 +89,13 @@ impl Tool for GlobFilesTool { annotations: Some(json!({ "readOnlyHint": true })), + supported_scopes: &[ ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], hidden: false, title_template: Some("Finding files matching '{pattern}'"), diff --git a/crates/code_assistant/src/tools/impls/list_files.rs b/crates/code_assistant/src/tools/impls/list_files.rs index 3178e359..13a9aae5 100644 --- a/crates/code_assistant/src/tools/impls/list_files.rs +++ b/crates/code_assistant/src/tools/impls/list_files.rs @@ -113,6 +113,8 @@ impl Tool for ListFilesTool { ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], hidden: false, title_template: Some("Listing files in {paths}"), diff --git a/crates/code_assistant/src/tools/impls/list_projects.rs b/crates/code_assistant/src/tools/impls/list_projects.rs index a86e696d..114b0768 100644 --- a/crates/code_assistant/src/tools/impls/list_projects.rs +++ b/crates/code_assistant/src/tools/impls/list_projects.rs @@ -72,6 +72,8 @@ impl Tool for ListProjectsTool { annotations: Some(json!({ "readOnlyHint": true })), + // This tool is only needed in MCP mode where we don't control the system message. + // The regular code-assistant will insert known projects into the system message. supported_scopes: &[ToolScope::McpServer], hidden: false, title_template: None, // Uses default tool name diff --git a/crates/code_assistant/src/tools/impls/mod.rs b/crates/code_assistant/src/tools/impls/mod.rs index de990491..591ea609 100644 --- a/crates/code_assistant/src/tools/impls/mod.rs +++ b/crates/code_assistant/src/tools/impls/mod.rs @@ -10,6 +10,7 @@ pub mod perplexity_ask; pub mod read_files; pub mod replace_in_file; pub mod search_files; +pub mod spawn_agent; pub mod update_plan; pub mod web_fetch; pub mod web_search; @@ -27,6 +28,7 @@ pub use perplexity_ask::PerplexityAskTool; pub use read_files::ReadFilesTool; pub use replace_in_file::ReplaceInFileTool; pub use search_files::SearchFilesTool; +pub use spawn_agent::SpawnAgentTool; pub use update_plan::UpdatePlanTool; pub use web_fetch::WebFetchTool; pub use web_search::WebSearchTool; diff --git a/crates/code_assistant/src/tools/impls/perplexity_ask.rs b/crates/code_assistant/src/tools/impls/perplexity_ask.rs index 94f68810..68f6dd33 100644 --- a/crates/code_assistant/src/tools/impls/perplexity_ask.rs +++ b/crates/code_assistant/src/tools/impls/perplexity_ask.rs @@ -109,11 +109,15 @@ impl Tool for PerplexityAskTool { "idempotentHint": false, "openWorldHint": true })), + supported_scopes: &[ ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], + // Note: can be disabled in read-only sub-agents if needed later. hidden: false, title_template: None, // Uses default tool name } diff --git a/crates/code_assistant/src/tools/impls/read_files.rs b/crates/code_assistant/src/tools/impls/read_files.rs index 046c8c6f..4ff47a0d 100644 --- a/crates/code_assistant/src/tools/impls/read_files.rs +++ b/crates/code_assistant/src/tools/impls/read_files.rs @@ -138,6 +138,8 @@ impl Tool for ReadFilesTool { ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], hidden: false, title_template: Some("Reading {paths}"), diff --git a/crates/code_assistant/src/tools/impls/search_files.rs b/crates/code_assistant/src/tools/impls/search_files.rs index 3aaa1873..3eddaa14 100644 --- a/crates/code_assistant/src/tools/impls/search_files.rs +++ b/crates/code_assistant/src/tools/impls/search_files.rs @@ -334,10 +334,13 @@ impl Tool for SearchFilesTool { annotations: Some(json!({ "readOnlyHint": true })), + supported_scopes: &[ ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], hidden: false, title_template: Some("Searching for '{regex}'"), diff --git a/crates/code_assistant/src/tools/impls/spawn_agent.rs b/crates/code_assistant/src/tools/impls/spawn_agent.rs new file mode 100644 index 00000000..6fbd7edc --- /dev/null +++ b/crates/code_assistant/src/tools/impls/spawn_agent.rs @@ -0,0 +1,184 @@ +use crate::tools::core::{ + Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolScope, ToolSpec, +}; +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +/// Input parameters for the spawn_agent tool. +#[derive(Deserialize, Serialize, Clone, PartialEq)] +pub struct SpawnAgentInput { + /// The instructions to give to the sub-agent. + pub instructions: String, + /// If true, instruct the sub-agent to include file references with line ranges. + #[serde(default, skip_serializing_if = "is_false")] + pub require_file_references: bool, + /// The mode for the sub-agent: "read_only" or "default". + #[serde(default = "default_mode", skip_serializing_if = "is_default_mode")] + pub mode: String, +} + +fn is_false(b: &bool) -> bool { + !b +} + +fn is_default_mode(mode: &str) -> bool { + mode == "read_only" +} + +fn default_mode() -> String { + "read_only".to_string() +} + +/// Output from the spawn_agent tool. +#[derive(Serialize, Deserialize)] +pub struct SpawnAgentOutput { + /// The final answer from the sub-agent (plain text for LLM context). + pub answer: String, + /// Error message if the sub-agent failed or was cancelled. + pub error: Option, + /// JSON output for UI display (includes tools list + response for custom renderer). + /// This is separate from `answer` because the UI needs structured data while + /// the LLM needs plain text. Persisted so sessions can restore the custom UI. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ui_output: Option, +} + +impl Render for SpawnAgentOutput { + fn status(&self) -> String { + if let Some(e) = &self.error { + format!("Sub-agent error: {e}") + } else { + "Sub-agent completed".to_string() + } + } + + /// Returns plain text for LLM context + fn render(&self, _tracker: &mut ResourcesTracker) -> String { + if let Some(e) = &self.error { + return format!("Sub-agent error: {e}"); + } + self.answer.clone() + } + + /// Returns JSON for UI display (custom renderer) + fn render_for_ui(&self, tracker: &mut ResourcesTracker) -> String { + // Use structured JSON output if available, otherwise fall back to plain text + self.ui_output + .clone() + .unwrap_or_else(|| self.render(tracker)) + } +} + +impl ToolResult for SpawnAgentOutput { + fn is_success(&self) -> bool { + self.error.is_none() + } +} + +/// The spawn_agent tool launches a sub-agent with isolated context. +pub struct SpawnAgentTool; + +#[async_trait::async_trait] +impl Tool for SpawnAgentTool { + type Input = SpawnAgentInput; + type Output = SpawnAgentOutput; + + fn spec(&self) -> ToolSpec { + let description = concat!( + "Spawns a sub-agent to execute a task with isolated context/history. ", + "The sub-agent runs independently and only the final answer is returned. ", + "Use this for exploratory or repetitive work that shouldn't pollute the main conversation history.\n\n", + "The sub-agent has access to read-only tools by default (file reading, searching, web access). ", + "Progress is streamed as the sub-agent works." + ); + + ToolSpec { + name: "spawn_agent", + description, + parameters_schema: json!({ + "type": "object", + "properties": { + "instructions": { + "type": "string", + "description": "The instructions/task to give to the sub-agent" + }, + "require_file_references": { + "type": "boolean", + "default": false, + "description": "If true, the sub-agent will be instructed to include exact file references with line ranges (e.g. `path/to/file.rs:10-20`)" + }, + "mode": { + "type": "string", + "enum": ["read_only", "default"], + "default": "read_only", + "description": "The mode for the sub-agent. 'read_only' restricts to read-only tools. 'default' allows broader tools (reserved for future use)." + } + }, + "required": ["instructions"] + }), + annotations: None, + // Exclude sub-agent scopes to prevent nesting + supported_scopes: &[ToolScope::Agent, ToolScope::AgentWithDiffBlocks], + hidden: false, + title_template: Some("Running sub-agent"), + } + } + + async fn execute<'a>( + &self, + context: &mut ToolContext<'a>, + input: &mut Self::Input, + ) -> Result { + // Get the sub-agent runner from context + let sub_agent_runner = context.sub_agent_runner.ok_or_else(|| { + anyhow!("Sub-agent runner not available. This tool requires the agent to be configured with sub-agent support.") + })?; + // Get tool_id for progress streaming + let tool_id = context + .tool_id + .clone() + .ok_or_else(|| anyhow!("Tool ID not available"))?; + + // Determine tool scope based on mode + let tool_scope = match input.mode.as_str() { + "read_only" => ToolScope::SubAgentReadOnly, + "default" => ToolScope::SubAgentDefault, + _ => ToolScope::SubAgentReadOnly, // Default to read-only for safety + }; + + // Build final instructions + let mut final_instructions = input.instructions.clone(); + if input.require_file_references { + final_instructions.push_str( + "\n\n---\n\ + IMPORTANT: When referencing code or files in your answer, always include exact file paths with line ranges.\n\ + Use the format: `path/to/file.rs:10-20` for ranges or `path/to/file.rs:15` for single lines.\n\ + This is required for your response to be considered complete.", + ); + } + + // Run the sub-agent + let result = sub_agent_runner + .run( + &tool_id, + final_instructions, + tool_scope, + input.require_file_references, + ) + .await; + + match result { + Ok(sub_result) => Ok(SpawnAgentOutput { + answer: sub_result.answer, + error: None, + ui_output: Some(sub_result.ui_output), + }), + Err(e) => Ok(SpawnAgentOutput { + answer: String::new(), + error: Some(e.to_string()), + ui_output: None, + }), + } + } +} diff --git a/crates/code_assistant/src/tools/impls/web_fetch.rs b/crates/code_assistant/src/tools/impls/web_fetch.rs index 8b6c5d12..2d862469 100644 --- a/crates/code_assistant/src/tools/impls/web_fetch.rs +++ b/crates/code_assistant/src/tools/impls/web_fetch.rs @@ -91,11 +91,15 @@ impl Tool for WebFetchTool { "idempotentHint": true, "openWorldHint": true })), + supported_scopes: &[ ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], + // Note: can be disabled in read-only sub-agents if needed later. hidden: false, title_template: Some("Fetching {url}"), } diff --git a/crates/code_assistant/src/tools/impls/web_search.rs b/crates/code_assistant/src/tools/impls/web_search.rs index 35943fd5..ddcc248d 100644 --- a/crates/code_assistant/src/tools/impls/web_search.rs +++ b/crates/code_assistant/src/tools/impls/web_search.rs @@ -108,11 +108,15 @@ impl Tool for WebSearchTool { "idempotentHint": true, "openWorldHint": true })), + supported_scopes: &[ ToolScope::McpServer, ToolScope::Agent, ToolScope::AgentWithDiffBlocks, + ToolScope::SubAgentReadOnly, + ToolScope::SubAgentDefault, ], + // Note: can be disabled in read-only sub-agents if needed later. hidden: false, title_template: Some("Searching web for '{query}'"), } diff --git a/crates/code_assistant/src/ui/backend.rs b/crates/code_assistant/src/ui/backend.rs index d51cfb86..a3511c5b 100644 --- a/crates/code_assistant/src/ui/backend.rs +++ b/crates/code_assistant/src/ui/backend.rs @@ -52,6 +52,12 @@ pub enum BackendEvent { session_id: String, policy: SandboxPolicy, }, + + // Sub-agent management + CancelSubAgent { + session_id: String, + tool_id: String, + }, } // Response from backend to UI @@ -83,10 +89,15 @@ pub enum BackendResponse { session_id: String, model_name: String, }, + SandboxPolicyChanged { session_id: String, policy: SandboxPolicy, }, + SubAgentCancelled { + session_id: String, + tool_id: String, + }, } #[derive(Debug, Clone)] @@ -164,6 +175,11 @@ pub async fn handle_backend_events( BackendEvent::ChangeSandboxPolicy { session_id, policy } => Some( handle_change_sandbox_policy(&multi_session_manager, &session_id, policy).await, ), + + BackendEvent::CancelSubAgent { + session_id, + tool_id, + } => Some(handle_cancel_sub_agent(&multi_session_manager, &session_id, &tool_id).await), }; // Send response back to UI only if there is one @@ -545,3 +561,52 @@ async fn handle_change_sandbox_policy( }, } } + +async fn handle_cancel_sub_agent( + multi_session_manager: &Arc>, + session_id: &str, + tool_id: &str, +) -> BackendResponse { + debug!( + "Cancelling sub-agent {} for session {}", + tool_id, session_id + ); + + let result = { + let manager = multi_session_manager.lock().await; + manager.cancel_sub_agent(session_id, tool_id) + }; + + match result { + Ok(true) => { + info!( + "Successfully cancelled sub-agent {} for session {}", + tool_id, session_id + ); + BackendResponse::SubAgentCancelled { + session_id: session_id.to_string(), + tool_id: tool_id.to_string(), + } + } + Ok(false) => { + debug!( + "Sub-agent {} not found or already completed for session {}", + tool_id, session_id + ); + // Not really an error - the sub-agent may have already completed + BackendResponse::SubAgentCancelled { + session_id: session_id.to_string(), + tool_id: tool_id.to_string(), + } + } + Err(e) => { + error!( + "Failed to cancel sub-agent {} for session {}: {}", + tool_id, session_id, e + ); + BackendResponse::Error { + message: format!("Failed to cancel sub-agent: {e}"), + } + } + } +} diff --git a/crates/code_assistant/src/ui/gpui/elements.rs b/crates/code_assistant/src/ui/gpui/elements.rs index a8c367b5..7cea2bd1 100644 --- a/crates/code_assistant/src/ui/gpui/elements.rs +++ b/crates/code_assistant/src/ui/gpui/elements.rs @@ -242,18 +242,15 @@ impl MessageContainer { if let Some(tool) = view.block.as_tool_mut() { if tool.id == tool_id { let was_generating = view.is_generating; - let had_streaming_output = - tool.output.as_ref().map(|o| !o.is_empty()).unwrap_or(false); tool.status = status; tool.status_message = message.clone(); - // Set output if provided and we don't already have streaming output + // Update output if provided + // Note: UpdateToolStatus always replaces output (used by spawn_agent for JSON updates) + // AppendToolOutput is used for streaming append behavior if let Some(ref new_output) = output { - if !had_streaming_output { - tool.output = Some(new_output.clone()); - } - // If we had streaming output, keep it (don't overwrite) + tool.output = Some(new_output.clone()); } // Update generating state based on tool completion @@ -1355,25 +1352,55 @@ impl Render for BlockView { block.state == ToolBlockState::Expanded }; + if should_show_output { - let output_color = - if block.status == crate::ui::ToolStatus::Error { - cx.theme().danger - } else { - cx.theme().foreground - }; - - expandable_elements.push( - div() - .p_2() - .mt_1() - .w_full() - .text_color(output_color) - .text_size(px(13.)) - .whitespace_normal() - .child(output_content.clone()) - .into_any(), - ); + // Try custom tool output renderer first + // Clone theme to avoid borrow conflict with cx + let theme = cx.theme().clone(); + let custom_output = crate::ui::gpui::tool_output_renderers::ToolOutputRendererRegistry::global() + .and_then(|registry| { + registry.render_output( + &block.name, + &block.id, + output_content, + &block.status, + &theme, + window, + cx, + ) + }); + + if let Some(custom_element) = custom_output { + // Use custom renderer output + expandable_elements.push( + div() + .px_2() + .mt_1() + .w_full() + .child(custom_element) + .into_any(), + ); + } else { + // Default output rendering + let output_color = + if block.status == crate::ui::ToolStatus::Error { + cx.theme().danger + } else { + cx.theme().foreground + }; + + expandable_elements.push( + div() + .p_2() + .mt_1() + .w_full() + .text_color(output_color) + .text_size(px(13.)) + .whitespace_normal() + .child(output_content.clone()) + .into_any(), + ); + } } } } diff --git a/crates/code_assistant/src/ui/gpui/file_icons.rs b/crates/code_assistant/src/ui/gpui/file_icons.rs index 09fcc219..261f0896 100644 --- a/crates/code_assistant/src/ui/gpui/file_icons.rs +++ b/crates/code_assistant/src/ui/gpui/file_icons.rs @@ -52,6 +52,7 @@ pub const TOOL_OPEN_PROJECT: &str = "expanded_folder"; // folder_open.svg pub const TOOL_USER_INPUT: &str = "person"; // person.svg pub const TOOL_COMPLETE_TASK: &str = "check_circle"; // check_circle.svg pub const TOOL_UPDATE_PLAN: &str = "file_generic"; // file_generic.svg +pub const TOOL_SPAWN_AGENT: &str = "rerun"; // rerun.svg - for spawning sub-agents pub const TOOL_GENERIC: &str = "file_code"; // file_code.svg const FILE_TYPES_ASSET: &str = "icons/file_icons/file_types.json"; @@ -137,7 +138,9 @@ impl FileIcons { TOOL_DELETE_FILES => Some("icons/trash.svg"), TOOL_USER_INPUT => Some("icons/person.svg"), TOOL_COMPLETE_TASK => Some("icons/check_circle.svg"), + TOOL_UPDATE_PLAN => Some("icons/file_generic.svg"), + TOOL_SPAWN_AGENT => Some("icons/rerun.svg"), TOOL_GENERIC => Some("icons/file_code.svg"), // For file_types.json types we missed _ => None, @@ -173,7 +176,9 @@ impl FileIcons { TOOL_OPEN_PROJECT => Some(SharedString::from("đź“‚")), TOOL_USER_INPUT => Some(SharedString::from("👤")), TOOL_COMPLETE_TASK => Some(SharedString::from("âś…")), + TOOL_UPDATE_PLAN => Some(SharedString::from("📝")), + TOOL_SPAWN_AGENT => Some(SharedString::from("🔄")), TOOL_GENERIC => Some(SharedString::from("đź”§")), _ => Some(SharedString::from("đź“„")), // Default fallback } @@ -195,8 +200,10 @@ impl FileIcons { "delete_files" => TOOL_DELETE_FILES, "open_project" => TOOL_OPEN_PROJECT, "user_input" => TOOL_USER_INPUT, + "complete_task" => TOOL_COMPLETE_TASK, "update_plan" => TOOL_UPDATE_PLAN, + "spawn_agent" => TOOL_SPAWN_AGENT, _ => TOOL_GENERIC, }; diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index 3d527ac0..85add068 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -16,7 +16,9 @@ mod plan_banner; mod root; pub mod sandbox_selector; pub mod simple_renderers; +pub mod spawn_agent_renderer; pub mod theme; +pub mod tool_output_renderers; use crate::persistence::{ChatMetadata, DraftStorage}; use crate::types::PlanState; @@ -27,6 +29,8 @@ use crate::ui::gpui::{ elements::MessageRole, parameter_renderers::{DefaultParameterRenderer, ParameterRendererRegistry}, simple_renderers::SimpleParameterRenderer, + spawn_agent_renderer::SpawnAgentInstructionsRenderer, + tool_output_renderers::{SpawnAgentOutputRenderer, ToolOutputRendererRegistry}, }; use crate::ui::{async_trait, DisplayFragment, UIError, UiEvent, UserInterface}; use assets::Assets; @@ -235,12 +239,20 @@ impl Gpui { false, // These are not full-width ))); + // Register spawn_agent instructions renderer (full-width markdown) + registry.register_renderer(Box::new(SpawnAgentInstructionsRenderer)); + // Wrap the registry in Arc for sharing let parameter_renderers = Arc::new(registry); // Set the global registry ParameterRendererRegistry::set_global(parameter_renderers.clone()); + // Initialize tool output renderers registry + let mut tool_output_registry = ToolOutputRendererRegistry::new(); + tool_output_registry.register_renderer(Box::new(SpawnAgentOutputRenderer)); + ToolOutputRendererRegistry::set_global(Arc::new(tool_output_registry)); + // Create a channel to send and receive UiEvents let (tx, rx) = async_channel::unbounded::(); let event_sender = Arc::new(Mutex::new(tx)); @@ -1001,6 +1013,20 @@ impl Gpui { path.display() ); } + UiEvent::CancelSubAgent { tool_id } => { + debug!("UI: CancelSubAgent event for tool_id: {}", tool_id); + // Forward to backend with current session ID + if let Some(session_id) = self.current_session_id.lock().unwrap().clone() { + if let Some(sender) = self.backend_event_sender.lock().unwrap().as_ref() { + let _ = sender.try_send(BackendEvent::CancelSubAgent { + session_id, + tool_id, + }); + } + } else { + warn!("UI: CancelSubAgent requested but no active session"); + } + } } } @@ -1323,6 +1349,7 @@ impl Gpui { ); } } + BackendResponse::SandboxPolicyChanged { session_id, policy } => { let current_session_id = self.current_session_id.lock().unwrap().clone(); if current_session_id.as_deref() == Some(session_id.as_str()) { @@ -1338,6 +1365,17 @@ impl Gpui { ); } } + BackendResponse::SubAgentCancelled { + session_id, + tool_id, + } => { + debug!( + "Received BackendResponse::SubAgentCancelled for tool {} in session {}", + tool_id, session_id + ); + // The sub-agent will update its own UI state via the normal tool output mechanism + // No additional UI update needed here + } } } } diff --git a/crates/code_assistant/src/ui/gpui/spawn_agent_renderer.rs b/crates/code_assistant/src/ui/gpui/spawn_agent_renderer.rs new file mode 100644 index 00000000..4ce2bc61 --- /dev/null +++ b/crates/code_assistant/src/ui/gpui/spawn_agent_renderer.rs @@ -0,0 +1,53 @@ +//! Custom parameter renderer for spawn_agent tool instructions + +use crate::ui::gpui::parameter_renderers::ParameterRenderer; +use gpui::{div, px, FontWeight, IntoElement, ParentElement, Styled}; + +/// Renderer for spawn_agent instructions parameter +/// Renders as full-width with a modest heading +pub struct SpawnAgentInstructionsRenderer; + +impl ParameterRenderer for SpawnAgentInstructionsRenderer { + fn supported_parameters(&self) -> Vec<(String, String)> { + vec![("spawn_agent".to_string(), "instructions".to_string())] + } + + fn render( + &self, + _tool_name: &str, + _param_name: &str, + param_value: &str, + theme: &gpui_component::theme::Theme, + ) -> gpui::AnyElement { + div() + .w_full() + .rounded_md() + .px_2() + .py_1() + .text_size(px(13.)) + .bg(crate::ui::gpui::theme::colors::tool_parameter_bg(theme)) + .child( + div() + .flex() + .flex_col() + .gap_1() + .child( + div() + .font_weight(FontWeight(500.0)) + .text_size(px(12.)) + .text_color(crate::ui::gpui::theme::colors::tool_parameter_label(theme)) + .child("Instructions:"), + ) + .child( + div() + .text_color(crate::ui::gpui::theme::colors::tool_parameter_value(theme)) + .child(param_value.to_string()), + ), + ) + .into_any_element() + } + + fn is_full_width(&self, _tool_name: &str, _param_name: &str) -> bool { + true + } +} diff --git a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs new file mode 100644 index 00000000..270029f2 --- /dev/null +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -0,0 +1,509 @@ +use crate::agent::sub_agent::{SubAgentActivity, SubAgentOutput, SubAgentToolStatus}; +use crate::ui::ToolStatus; +use gpui::{ + bounce, div, ease_in_out, percentage, px, svg, Animation, AnimationExt, ClickEvent, Context, + Element, InteractiveElement, ParentElement, SharedString, StatefulInteractiveElement, Styled, + Transformation, Window, +}; +use gpui_component::text::TextView; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; +use tracing::warn; + +/// A unique key for tool name +pub type ToolKey = String; + +/// Trait for tool output renderers that can provide custom rendering for tool output +pub trait ToolOutputRenderer: Send + Sync { + /// List of tool names this renderer supports + fn supported_tools(&self) -> Vec; + + /// Render the tool output as a UI element + /// Returns None if the default rendering should be used + fn render( + &self, + tool_id: &str, + output: &str, + status: &ToolStatus, + theme: &gpui_component::theme::Theme, + window: &mut Window, + cx: &mut Context, + ) -> Option; +} + +/// Registry for tool output renderers +pub struct ToolOutputRendererRegistry { + renderers: HashMap>>, +} + +// Global registry singleton using OnceLock (thread-safe) +static GLOBAL_REGISTRY: OnceLock>>> = OnceLock::new(); + +impl ToolOutputRendererRegistry { + /// Set the global registry + pub fn set_global(registry: Arc) { + let global_mutex = GLOBAL_REGISTRY.get_or_init(|| Mutex::new(None)); + if let Ok(mut guard) = global_mutex.lock() { + *guard = Some(registry); + } else { + warn!("Failed to acquire lock for setting global tool output registry"); + } + } + + /// Get a reference to the global registry + pub fn global() -> Option> { + if let Some(global_mutex) = GLOBAL_REGISTRY.get() { + if let Ok(guard) = global_mutex.lock() { + return guard.clone(); + } + } + None + } + + /// Create a new empty registry + pub fn new() -> Self { + Self { + renderers: HashMap::new(), + } + } + + /// Register a new renderer for its supported tools + pub fn register_renderer(&mut self, renderer: Box) { + let renderer_arc = Arc::new(renderer); + for tool_name in renderer_arc.supported_tools() { + if self.renderers.contains_key(&tool_name) { + warn!( + "Overriding existing output renderer for tool: {}", + tool_name + ); + } + self.renderers.insert(tool_name, renderer_arc.clone()); + } + } + + /// Check if a custom renderer exists for a tool + #[allow(dead_code)] + pub fn has_renderer(&self, tool_name: &str) -> bool { + self.renderers.contains_key(tool_name) + } + + /// Render tool output using the appropriate renderer + /// Returns None if no custom renderer is registered (use default rendering) + #[allow(clippy::too_many_arguments)] + pub fn render_output( + &self, + tool_name: &str, + tool_id: &str, + output: &str, + status: &ToolStatus, + theme: &gpui_component::theme::Theme, + window: &mut Window, + cx: &mut Context, + ) -> Option { + self.renderers + .get(tool_name) + .and_then(|renderer| renderer.render(tool_id, output, status, theme, window, cx)) + } +} + +impl Default for ToolOutputRendererRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Renderer for spawn_agent tool output +/// Displays sub-agent tool calls in a compact, Zed-like style with markdown response +pub struct SpawnAgentOutputRenderer; + +impl SpawnAgentOutputRenderer { + /// Get the icon for a tool name (matching file_icons.rs logic) + fn get_tool_icon(tool_name: &str) -> Option { + use super::file_icons; + file_icons::get().get_tool_icon(tool_name) + } + + /// Render a single compact tool line showing tool title (one-liner) + fn render_tool_line( + tool: &crate::agent::sub_agent::SubAgentToolCall, + theme: &gpui_component::theme::Theme, + ) -> gpui::AnyElement { + use super::file_icons; + + let icon = Self::get_tool_icon(&tool.name); + + // Status-based colors + let (icon_color, text_color) = match tool.status { + SubAgentToolStatus::Running => (theme.info, theme.muted_foreground), + SubAgentToolStatus::Success => (theme.success, theme.muted_foreground), + SubAgentToolStatus::Error => (theme.danger, theme.danger), + }; + + // Use title (generated from template) if available, + // otherwise fall back to status message, then tool name + let display_text = tool + .title + .as_ref() + .filter(|t| !t.is_empty()) + .cloned() + .or_else(|| tool.message.as_ref().filter(|m| !m.is_empty()).cloned()) + .unwrap_or_else(|| tool.name.replace('_', " ")); + + div() + .flex() + .flex_row() + .items_center() + .gap_2() + .py(px(2.)) + .children(vec![ + // Icon + file_icons::render_icon_container(&icon, 14.0, icon_color, "đź”§").into_any(), + // Title text (one-liner from template or fallback) + div() + .text_size(px(13.)) + .text_color(text_color) + .child(display_text) + .into_any(), + ]) + .into_any() + } + + /// Render activity line with a cancel button for running sub-agents + fn render_activity_with_cancel( + activity: SubAgentActivity, + tool_id: &str, + theme: &gpui_component::theme::Theme, + cx: &mut Context, + ) -> Option { + let (text, color, show_spinner, is_cancellable) = match activity { + SubAgentActivity::WaitingForLlm => ("Waiting...", theme.muted_foreground, true, true), + SubAgentActivity::Streaming => ("Responding...", theme.info, true, true), + SubAgentActivity::ExecutingTools => ("Executing tools...", theme.info, true, true), + SubAgentActivity::Completed => return None, // Don't show for completed + SubAgentActivity::Cancelled => ("Cancelled", theme.warning, false, false), + SubAgentActivity::Failed => return None, // Error shown separately + }; + + let mut children: Vec = Vec::new(); + + // Add spinning arrow if active + if show_spinner { + children.push( + svg() + .size(px(14.)) + .path(SharedString::from("icons/arrow_circle.svg")) + .text_color(color) + .with_animation( + "sub_agent_spinner", + Animation::new(Duration::from_secs(2)) + .repeat() + .with_easing(bounce(ease_in_out)), + |svg, delta| { + svg.with_transformation(Transformation::rotate(percentage(delta))) + }, + ) + .into_any(), + ); + } + + // Add status text + children.push( + div() + .text_size(px(13.)) + .text_color(color) + .child(text) + .into_any(), + ); + + // Add cancel button for active sub-agents + if is_cancellable { + let tool_id_owned = tool_id.to_string(); + let cancel_color = theme.muted_foreground; + let cancel_hover_color = theme.danger; + + children.push( + div() + .id(SharedString::from(format!("cancel-{}", tool_id))) + .ml_2() + .px_2() + .py(px(1.)) + .rounded(px(4.)) + .text_size(px(11.)) + .text_color(cancel_color) + .cursor_pointer() + .hover(|s| { + s.bg(cancel_hover_color.opacity(0.15)) + .text_color(cancel_hover_color) + }) + .on_click(cx.listener(move |_view, _event: &ClickEvent, _window, cx| { + // Send cancel event through the UI event system + if let Some(sender) = cx.try_global::() { + let _ = sender.0.try_send(crate::ui::UiEvent::CancelSubAgent { + tool_id: tool_id_owned.clone(), + }); + } + })) + .child("Cancel") + .into_any(), + ); + } + + Some( + div() + .flex() + .flex_row() + .items_center() + .gap_2() + .py(px(2.)) + .children(children) + .into_any(), + ) + } + + /// Render error/cancelled status if present + fn render_status_line( + output: &SubAgentOutput, + theme: &gpui_component::theme::Theme, + ) -> Option { + if output.cancelled == Some(true) { + return Some( + div() + .flex() + .flex_row() + .items_center() + .gap_2() + .py(px(2.)) + .child( + div() + .text_size(px(13.)) + .text_color(theme.warning) + .child("Sub-agent cancelled"), + ) + .into_any(), + ); + } + + if let Some(error) = &output.error { + return Some( + div() + .flex() + .flex_row() + .items_center() + .gap_2() + .py(px(2.)) + .child( + div() + .text_size(px(13.)) + .text_color(theme.danger) + .child(format!("Error: {error}")), + ) + .into_any(), + ); + } + + None + } + + /// Render the final response as markdown + fn render_response( + response: &str, + theme: &gpui_component::theme::Theme, + window: &mut Window, + cx: &mut Context, + ) -> gpui::AnyElement { + div() + .mt_2() + .pt_2() + .border_t_1() + .border_color(theme.border) + .text_color(theme.foreground) + .child(TextView::markdown( + "sub-agent-response", + response.to_string(), + window, + cx, + )) + .into_any() + } +} + +impl ToolOutputRenderer for SpawnAgentOutputRenderer { + fn supported_tools(&self) -> Vec { + vec!["spawn_agent".to_string()] + } + + fn render( + &self, + tool_id: &str, + output: &str, + _status: &ToolStatus, + theme: &gpui_component::theme::Theme, + window: &mut Window, + cx: &mut Context, + ) -> Option { + if output.is_empty() { + return None; + } + + // Parse JSON output + let parsed = match SubAgentOutput::from_json(output) { + Some(p) => p, + None => { + // If not valid JSON, return None to use default text rendering + // This handles backwards compatibility with any old markdown format + return None; + } + }; + + // Always render if we have activity state, response, or tools + let has_activity = parsed.activity.is_some(); + let has_response = parsed.response.is_some(); + if parsed.tools.is_empty() + && parsed.cancelled.is_none() + && parsed.error.is_none() + && !has_activity + && !has_response + { + return None; + } + + let mut elements: Vec = Vec::new(); + + // Render compact list of tool calls first (history of what was done) + for tool in &parsed.tools { + elements.push(Self::render_tool_line(tool, theme)); + } + + // Add activity line with cancel button (shows current state with spinner) + // This is where new output will appear + if let Some(activity) = parsed.activity { + if let Some(activity_element) = + Self::render_activity_with_cancel(activity, tool_id, theme, cx) + { + elements.push(activity_element); + } + } + + // Add error/cancelled status line if present + if let Some(status_line) = Self::render_status_line(&parsed, theme) { + elements.push(status_line); + } + + // Add the final response rendered as markdown (after completion) + if let Some(ref response) = parsed.response { + if !response.is_empty() { + elements.push(Self::render_response(response, theme, window, cx)); + } + } + + if elements.is_empty() { + return None; + } + + Some( + div() + .flex() + .flex_col() + .gap_0() + .mt_1() + .children(elements) + .into_any(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_json_output() { + let json = r#"{"tools":[{"name":"read_files","status":"success"},{"name":"search_files","status":"running","message":"Searching..."}]}"#; + + let parsed = SubAgentOutput::from_json(json).unwrap(); + assert_eq!(parsed.tools.len(), 2); + assert_eq!(parsed.tools[0].name, "read_files"); + assert_eq!(parsed.tools[0].status, SubAgentToolStatus::Success); + assert_eq!(parsed.tools[1].name, "search_files"); + assert_eq!(parsed.tools[1].status, SubAgentToolStatus::Running); + assert_eq!(parsed.tools[1].message.as_deref(), Some("Searching...")); + } + + #[test] + fn test_parse_json_with_cancelled() { + let json = r#"{"tools":[{"name":"read_files","status":"success"}],"cancelled":true}"#; + + let parsed = SubAgentOutput::from_json(json).unwrap(); + assert_eq!(parsed.tools.len(), 1); + assert_eq!(parsed.cancelled, Some(true)); + } + + #[test] + fn test_parse_json_with_error() { + let json = r#"{"tools":[],"error":"Connection failed"}"#; + + let parsed = SubAgentOutput::from_json(json).unwrap(); + assert_eq!(parsed.error.as_deref(), Some("Connection failed")); + } + + #[test] + fn test_roundtrip() { + let mut output = SubAgentOutput::new(); + output + .tools + .push(crate::agent::sub_agent::SubAgentToolCall { + name: "read_files".to_string(), + status: SubAgentToolStatus::Success, + title: None, + message: None, + parameters: std::collections::HashMap::new(), + }); + output + .tools + .push(crate::agent::sub_agent::SubAgentToolCall { + name: "search_files".to_string(), + status: SubAgentToolStatus::Running, + title: Some("Searching for 'pattern'".to_string()), + message: Some("Searching for pattern".to_string()), + parameters: std::collections::HashMap::new(), + }); + + let json = output.to_json(); + let parsed = SubAgentOutput::from_json(&json).unwrap(); + + assert_eq!(parsed.tools.len(), 2); + assert_eq!(parsed.tools[0].name, "read_files"); + assert_eq!( + parsed.tools[1].message.as_deref(), + Some("Searching for pattern") + ); + assert_eq!( + parsed.tools[1].title.as_deref(), + Some("Searching for 'pattern'") + ); + } + + #[test] + fn test_invalid_json_returns_none() { + let invalid = "### Sub-agent activity\n- Calling tool read_files"; + assert!(SubAgentOutput::from_json(invalid).is_none()); + } + + #[test] + fn test_parse_json_with_response() { + let json = r#"{"tools":[{"name":"read_files","status":"success","message":"Successfully loaded 2 file(s)"}],"activity":"completed","response":"Here is the answer based on the files I read."}"#; + + let parsed = SubAgentOutput::from_json(json).unwrap(); + assert_eq!(parsed.tools.len(), 1); + assert_eq!( + parsed.tools[0].message.as_deref(), + Some("Successfully loaded 2 file(s)") + ); + assert_eq!(parsed.activity, Some(SubAgentActivity::Completed)); + assert_eq!( + parsed.response.as_deref(), + Some("Here is the answer based on the files I read.") + ); + } +} diff --git a/crates/code_assistant/src/ui/terminal/app.rs b/crates/code_assistant/src/ui/terminal/app.rs index f1c75eba..20482c38 100644 --- a/crates/code_assistant/src/ui/terminal/app.rs +++ b/crates/code_assistant/src/ui/terminal/app.rs @@ -430,6 +430,7 @@ impl TerminalTuiApp { "Switched to model: {model_name}", ))); } + BackendResponse::SandboxPolicyChanged { session_id: _, policy, @@ -441,6 +442,13 @@ impl TerminalTuiApp { policy ))); } + BackendResponse::SubAgentCancelled { + session_id: _, + tool_id: _, + } => { + // Sub-agent cancellation handled; the sub-agent will + // update its tool output via the normal mechanism + } } } }); diff --git a/crates/code_assistant/src/ui/terminal/message.rs b/crates/code_assistant/src/ui/terminal/message.rs index 3351efd2..d1d8ca05 100644 --- a/crates/code_assistant/src/ui/terminal/message.rs +++ b/crates/code_assistant/src/ui/terminal/message.rs @@ -161,9 +161,11 @@ impl MessageBlock { height += 1; } - // Output - if block.output.is_some() && block.status == ToolStatus::Success { - height += 1; + // Output (used by spawn_agent for streaming sub-agent activity) + if let Some(ref output) = block.output { + if !output.is_empty() { + height += output.lines().count() as u16; + } } height diff --git a/crates/code_assistant/src/ui/terminal/tool_widget.rs b/crates/code_assistant/src/ui/terminal/tool_widget.rs index 3b913a34..b5127ad5 100644 --- a/crates/code_assistant/src/ui/terminal/tool_widget.rs +++ b/crates/code_assistant/src/ui/terminal/tool_widget.rs @@ -228,6 +228,111 @@ impl<'a> Widget for ToolWidget<'a> { display_text, Style::default().fg(Color::LightRed), ); + current_y += 1; + } + } + + // Tool output (used by spawn_agent for streaming sub-agent activity) + if let Some(ref output) = self.tool_block.output { + if current_y < area.y + area.height && !output.is_empty() { + // Try to parse as JSON (new structured format) + if let Some(sub_agent_output) = + crate::agent::sub_agent::SubAgentOutput::from_json(output) + { + // Render structured sub-agent tools + for tool in &sub_agent_output.tools { + if current_y >= area.y + area.height { + break; + } + + let (status_symbol, color) = match tool.status { + crate::agent::sub_agent::SubAgentToolStatus::Running => { + ("â—Ź", Color::Blue) + } + crate::agent::sub_agent::SubAgentToolStatus::Success => { + ("â—Ź", Color::Green) + } + crate::agent::sub_agent::SubAgentToolStatus::Error => ("â—Ź", Color::Red), + }; + + // Use title (generated from template) if available, + // otherwise fall back to status message, then tool name + let display_text = tool + .title + .as_ref() + .filter(|t| !t.is_empty()) + .cloned() + .or_else(|| tool.message.as_ref().filter(|m| !m.is_empty()).cloned()) + .unwrap_or_else(|| tool.name.replace('_', " ")); + + let full_text = format!("{status_symbol} {display_text}"); + + let truncated = if full_text.len() > (area.width.saturating_sub(4)) as usize + { + format!( + "{}...", + &full_text[..(area.width.saturating_sub(7)) as usize] + ) + } else { + full_text + }; + + buf.set_string( + area.x + 2, + current_y, + &truncated, + Style::default().fg(color), + ); + current_y += 1; + } + + // Render cancelled/error status + if sub_agent_output.cancelled == Some(true) && current_y < area.y + area.height + { + buf.set_string( + area.x + 2, + current_y, + "Sub-agent cancelled", + Style::default().fg(Color::Yellow), + ); + current_y += 1; + } + + if let Some(error) = &sub_agent_output.error { + if current_y < area.y + area.height { + let error_text = format!("Error: {error}"); + buf.set_string( + area.x + 2, + current_y, + &error_text, + Style::default().fg(Color::Red), + ); + // current_y is not incremented after this as it's the last possible item + } + } + } else { + // Fallback: render as plain text (for backwards compatibility) + for line in output.lines() { + if current_y >= area.y + area.height { + break; + } + + let color = Color::Gray; + let truncated = if line.len() > (area.width.saturating_sub(4)) as usize { + format!("{}...", &line[..(area.width.saturating_sub(7)) as usize]) + } else { + line.to_string() + }; + + buf.set_string( + area.x + 2, + current_y, + &truncated, + Style::default().fg(color), + ); + current_y += 1; + } + } } } } diff --git a/crates/code_assistant/src/ui/ui_events.rs b/crates/code_assistant/src/ui/ui_events.rs index b63bd0ae..6c716973 100644 --- a/crates/code_assistant/src/ui/ui_events.rs +++ b/crates/code_assistant/src/ui/ui_events.rs @@ -118,6 +118,8 @@ pub enum UiEvent { UpdateCurrentModel { model_name: String }, /// Update the current sandbox selection in the UI UpdateSandboxPolicy { policy: SandboxPolicy }, + /// Cancel a running sub-agent by its tool id + CancelSubAgent { tool_id: String }, // === Resource Events (for tool operations) === /// A file was loaded/read by a tool diff --git a/docs/sub-agents-feature.md b/docs/sub-agents-feature.md new file mode 100644 index 00000000..8a778cde --- /dev/null +++ b/docs/sub-agents-feature.md @@ -0,0 +1,283 @@ +# Sub-agents feature (`spawn_agent` tool) + +## Implementation Status + +### Completed +- [x] `SubAgentRunner` trait and `DefaultSubAgentRunner` implementation (`sub_agent.rs`) +- [x] `SubAgentCancellationRegistry` for managing cancellation tokens +- [x] `SubAgentUiAdapter` for streaming progress to parent tool block +- [x] `ToolScope::SubAgentReadOnly` and `ToolScope::SubAgentDefault` variants +- [x] Tool scope updates for read-only tools (search_files, read_files, glob_files, list_files, web_fetch, web_search, perplexity_ask) +- [x] `spawn_agent` tool implementation (`spawn_agent.rs`) +- [x] Tool registration in `mod.rs` and `registry.rs` +- [x] Session manager wiring of `DefaultSubAgentRunner` with UI and permission handler +- [x] Added required `Agent` methods: `set_tool_scope`, `set_session_model_config`, `set_session_identity`, `message_history` +- [x] File reference enforcement with retry logic (up to 2 retries) + +### Completed (phase 2) +- [x] **Terminal UI rendering**: Update terminal tool block to show streaming sub-agent activity + - Added output rendering in `ToolWidget` with color coding for activity lines + - Updated height calculation to account for multi-line output +- [x] **Parallel execution**: Update `manage_tool_execution()` to run multiple `spawn_agent` calls concurrently + - Multiple `spawn_agent` read-only tools now run in parallel using `futures::join_all` + - Results are collected in deterministic order matching original tool request ordering + - Only `read_only` mode spawn_agents are parallelized for safety + +### Completed (phase 3) +- [x] **Integration tests**: Added tests in `tests/sub_agent_tests.rs`: + - `test_spawn_agent_output_render` - output rendering for success/cancel/error + - `test_spawn_agent_input_parsing` - input parsing with defaults + - `test_cancellation_registry` - cancellation registration and triggering + - `test_mock_sub_agent_runner` - basic mock runner execution + - `test_parallel_sub_agent_execution` - verifies concurrent execution + - `test_tool_scope_for_sub_agent` - verifies tool availability per scope + - `test_can_run_in_parallel_logic` - parallel execution eligibility logic + + + +### Completed (phase 4) +- [x] **Structured output format**: Sub-agent activity is now captured as structured JSON + - Added `SubAgentOutput`, `SubAgentToolCall`, `SubAgentToolStatus` types in `sub_agent.rs` + - JSON format is persisted with the session and can be restored/re-rendered + - Format: `{"tools": [{"name": "read_files", "status": "success", "message": null}], "cancelled": null, "error": null}` +- [x] **SubAgentUiAdapter refactored**: Emits structured JSON instead of markdown + - Tracks current tool state and updates status properly + - Handles cancellation and error states +- [x] **GPUI rendering**: Custom tool output renderer for sub-agent progress display + - Added `ToolOutputRendererRegistry` pattern (similar to `ParameterRendererRegistry`) + - Implemented `SpawnAgentOutputRenderer` that parses JSON sub-agent output + - Renders sub-tool calls in compact Zed-like style with icons and status colors + - Falls back to plain text rendering for backwards compatibility + - Located in `crates/code_assistant/src/ui/gpui/tool_output_renderers.rs` +- [x] **Terminal UI rendering**: Updated to parse JSON format + - Renders sub-tool calls with status symbols and colors + - Falls back to plain text for backwards compatibility +- [x] **ACP mode support**: Sub-agent activity streams through existing tool output mechanisms + - Added `spawn_agent` icon mapping in `file_icons.rs` (uses `rerun.svg`) + - JSON output flows as `ToolCallUpdate` content for display in Zed's ACP panel + + +### Completed (phase 5) +- [x] **UI integration for cancellation**: Cancel button per running `spawn_agent` block + - Added `CancelSubAgent` event to `UiEvent` enum (`ui_events.rs`) + - Added `CancelSubAgent` backend event and handler (`backend.rs`) + - Added `sub_agent_cancellation_registry` field to `SessionInstance` for tracking active registries + - Added `cancel_sub_agent` method to `SessionManager` to trigger cancellation + - Added cancel button in `SpawnAgentOutputRenderer` that appears during active sub-agent execution + - Button triggers `CancelSubAgent` event which flows through backend to the registry + +### Completed (phase 6) +- [x] **ACP spawn_agent markdown rendering**: Convert SubAgentOutput JSON to markdown for Zed display + - Added `render_sub_agent_output_as_markdown()` helper function in `acp/ui.rs` + - Renders tool calls with status symbols (âźł/âś“/âś—), activity state, errors, and final response + - Falls back to raw output if JSON parsing fails + - Located in `build_content()` method, triggered when tool_name is "spawn_agent" + +### Pending +- [ ] **Permission attribution**: Show permission requests as originating from sub-agent context (inline or popover) + +### Notes +- The `sub_agent_cancellation_registry` field in `ToolContext` is reserved for future tool implementations that may need to cancel sub-agents programmatically. + +--- + +## Goal + +Add a new tool, `spawn_agent`, that launches a sub-agent to execute a task with **isolated context/history**, returning only the **final output** back to the main agent as the tool result. + +The primary motivation is **context window management**: the sub-agent can perform repetitive/exploratory work without polluting the main agent’s conversation history. Additionally, sub-agent progress should be visible in the UI **inside the `spawn_agent` tool block**, and multiple sub-agents should be able to run **concurrently**. + +## UX requirements (what the user sees) + +### During execution +- The main chat history shows a single tool call: `spawn_agent`. +- While the sub-agent runs, the tool block output is continuously updated with a reduced view of sub-agent activity, such as: + - progress lines (e.g. “Searching…”, “Reading file…”) and/or + - a compact list of the sub-agent’s tool calls. + +This output is **UI-only** and should not be appended to the main agent’s message history. However, when persisting the `span_agent` tool execution, it allows for re-creating the custom tool output UI. + +### Completion +- When the sub-agent completes, the `spawn_agent` tool block gets its final output. +- The tool result handed back into the main agent context is **only** the sub-agent’s final answer (or a cancellation message). + +### Parallel sub-agents +- If the main agent requests multiple `spawn_agent` tool calls in a single response, they should run **truly concurrently**. +- Recommended safety default: only allow concurrent sub-agents in **read-only** mode. + +### Cancellation +- Users can cancel a specific running sub-agent. +- The main agent sees a deterministic result: e.g. `"Sub-agent cancelled by user."` + +### Permissions +- Sub-agents may use permission-gated tools. +- Permission requests should be attributed to the sub-agent (and ideally shown inline in the tool block, or as a popover anchored to it). + +## Tool API + +### Tool name +- `spawn_agent` + +### Parameters +Keep parameters intentionally small: + +- `instructions: string` (required) +- `require_file_references: bool` (optional, default `false`) + - If true, the implementation appends a static instruction suffix telling the sub-agent to include exact file references with line ranges. +- `mode: "read_only" | "default"` (optional, default TBD) + - `read_only` maps to a restricted tool scope. + - `default` may allow broader tools (future/optional). + +### Tool result returned to main agent +- For the main agent, the tool result content is the **final answer string**. +- If cancelled: return a cancellation string. +- Optional future extension: structured output (locations list), but not required for v1. + +## File references (current code points likely to change) + +The following files were identified during planning as the key implementation points: + +### Tool registration & implementation +- `crates/code_assistant/src/tools/core/registry.rs` + - Add registration for the new `spawn_agent` tool in `register_default_tools()`. +- `crates/code_assistant/src/tools/impls/mod.rs` + - Add `pub mod spawn_agent;` and re-export `SpawnAgentTool`. +- New file: `crates/code_assistant/src/tools/impls/spawn_agent.rs` + - Implement tool schema, invoke sub-agent runner, stream progress updates, and produce final tool output. + +### Tool context / runtime plumbing +- `crates/code_assistant/src/tools/core/tool.rs` + - Extend `ToolContext` to include a sub-agent spawning capability (e.g. `sub_agent_runner` or `agent_factory`). + +### Agent loop changes (parallel tool execution) +- `crates/code_assistant/src/agent/runner.rs` + - `manage_tool_execution()` currently executes tool requests sequentially. + - Update it to run multiple `spawn_agent` tool calls concurrently (likely using `tokio::spawn` + join), while preserving deterministic ordering of tool-result blocks. + - Add cancellation wiring hooks (tool-level cancellation signal). + +### Session/UI integration (where agent is created) +- `crates/code_assistant/src/session/manager.rs` + - This is where the main `AgentComponents` are built and the agent task is spawned. + - Likely place to wire in factories/services needed by `spawn_agent` (e.g. LLM provider factory) and/or cancellation routing. + +### Tool scopes +- `crates/code_assistant/src/tools/core/spec.rs` + - Add new `ToolScope` variants: + - `SubAgentReadOnly` + - `SubAgentDefault` +- Tool implementations under `crates/code_assistant/src/tools/impls/*.rs` + - Update each tool’s `ToolSpec.supported_scopes` to include/exclude the new sub-agent scopes. + - Ensure edit/write/delete tools do **not** include `SubAgentReadOnly`. + - The new `spawn_agent` tool includes neither scope, to prevent nesting agents. + +### Permissions +- `crates/code_assistant/src/permissions/*` (exact files depend on current mediator implementation) + - Ensure permission requests can be attributed to a sub-agent execution (ideally include parent tool id or sub-agent id in metadata). + +### UI rendering for streaming sub-agent activity +- GPUI tool widget rendering: + - `crates/code_assistant/src/ui/gpui/*` (tool widget / tool output renderer) + - `crates/code_assistant/src/ui/terminal/tool_widget.rs` (terminal UI tool block rendering) + +(Exact UI files depend on how tool outputs are currently rendered and updated; the runtime emits `UiEvent::UpdateToolStatus` in `crates/code_assistant/src/agent/runner.rs`.) + +## Implementation approach + +### 1) Add a sub-agent runner abstraction + +Create an internal service (name flexible): +- `SubAgentRunner` (trait) or `AgentSpawner` + +Responsibilities: +- Construct an in-process nested `Agent` configured for sub-agent operation. +- Ensure sub-agent state is isolated (no persistence into parent session). +- Provide a UI adapter that streams sub-agent activity into the parent `spawn_agent` tool block. +- Support cancellation via a per-sub-agent cancellation token. + +### 2) LLM provider creation + +Avoid requiring `LLMProvider: Clone`. + +Implement a small factory interface used by the sub-agent runner: +- `LlmProviderFactory::create(model_name) -> Box` + +Use existing model configuration and construction: +- `crates/llm/src/factory.rs` provides `create_llm_client_from_model(...)`. + +### 3) Tool scopes for sub-agent read-only mode + +Add a new tool scope: +- `ToolScope::SubAgentReadOnly` + +Update tools: +- Read/search/list/glob: include `SubAgentReadOnly`. +- Write/edit/delete/replace: exclude `SubAgentReadOnly`. + +The sub-agent runner selects `tool_scope` based on `mode`. + +### 4) `spawn_agent` tool implementation + +In `crates/code_assistant/src/tools/impls/spawn_agent.rs`: + +- Parse input parameters. +- Build final sub-agent instructions: + - Always include `instructions`. + - If `require_file_references`, append a fixed instruction block requesting references with line ranges. +- Ask runner to spawn sub-agent with: + - chosen scope (`SubAgentReadOnly` for read-only mode), + - cancellation token, + - UI adapter target = the current tool id. +- Stream sub-agent activity into tool output. +- Return the final answer as the tool result content. + +### 5) File reference enforcement (without rerunning) + +If `require_file_references=true`: + +- After sub-agent produces a candidate final answer, validate the presence of file references with line ranges. +- If missing, ask the *same* sub-agent to revise by appending a corrective user message. +- Bound number of retries (e.g. 2). +- If still missing, return a failure (or return best-effort answer with a warning — final behavior TBD). + +This avoids rerunning from scratch and avoids involving the main agent in “please try again” loops. + +### 6) Parallel execution of multiple `spawn_agent` tool calls + +Update `Agent::manage_tool_execution()` (`crates/code_assistant/src/agent/runner.rs`): + +- When tool requests include multiple `spawn_agent` calls, execute them concurrently. +- Keep deterministic ordering for the tool result blocks appended to the message history (match original tool request ordering). +- Recommended v1 safety: only run concurrently when `mode=read_only`. + +### 7) Cancellation + +- Each `spawn_agent` execution registers a cancellation handle keyed by the tool call id. +- UI exposes cancel control per running `spawn_agent` block. +- On cancellation: + - cancel the sub-agent LLM request/tool loop, + - mark tool status as cancelled, + - return tool result text: `"Sub-agent cancelled by user."` + +### 8) Permissions + +- Sub-agent tool invocations may require permission. +- Ensure permission requests can be shown as originating from the sub-agent context. +- UX: inline within the tool block or popover. + +## Testing plan + +### Unit tests +- `spawn_agent` input validation. +- Rendering of tool result (only final output is returned to main agent). + +### Integration tests +- Isolation: verify main agent message history contains only the `spawn_agent` tool result, not sub-agent transcript. +- Parallelism: multiple `spawn_agent` calls execute concurrently (use artificial delays/mocks). +- Cancellation: cancelling one sub-agent yields deterministic cancelled output and doesn’t cancel others. +- Permission routing: sub-agent permission prompts surface and block correctly. +- File reference enforcement: missing refs triggers in-sub-agent revision (bounded). + +## Rollout notes + +- Implement the UI streaming inside tool block first; ACP can initially show a markdown list. +- Keep the tool API minimal; add structured file reference extraction/enforcement later if needed.