From 367a364ced977388e2b34e08da3c1432c76491de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 12:50:51 +0100 Subject: [PATCH 01/18] Implementation plan --- docs/sub-agents-feature.md | 202 +++++++++++++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 docs/sub-agents-feature.md diff --git a/docs/sub-agents-feature.md b/docs/sub-agents-feature.md new file mode 100644 index 00000000..516043c4 --- /dev/null +++ b/docs/sub-agents-feature.md @@ -0,0 +1,202 @@ +# Sub-agents feature (`spawn_agent` tool) + +## 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. From e373b2e566841f68bb1ff7e2f2e5dd7e2db25f38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 15:37:04 +0100 Subject: [PATCH 02/18] Switch default system message to the recently used claude.md --- .../resources/system_prompts/default.md | 139 +++++------------- 1 file changed, 33 insertions(+), 106 deletions(-) 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` From 3cb4d0dd55889c1cefb8705ac373ca086bd31c7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 15:37:23 +0100 Subject: [PATCH 03/18] Implement sub-agent feature and spawn_agent tool --- crates/code_assistant/src/agent/mod.rs | 2 + crates/code_assistant/src/agent/runner.rs | 40 ++ crates/code_assistant/src/agent/sub_agent.rs | 401 ++++++++++++++++++ crates/code_assistant/src/agent/tests.rs | 24 ++ crates/code_assistant/src/mcp/handler.rs | 2 + crates/code_assistant/src/session/manager.rs | 23 +- .../src/tests/format_on_save_tests.rs | 14 + crates/code_assistant/src/tests/mocks.rs | 4 + .../code_assistant/src/tools/core/registry.rs | 4 +- crates/code_assistant/src/tools/core/spec.rs | 4 + crates/code_assistant/src/tools/core/tool.rs | 4 + .../src/tools/impls/glob_files.rs | 3 + .../src/tools/impls/list_files.rs | 2 + .../src/tools/impls/list_projects.rs | 2 + crates/code_assistant/src/tools/impls/mod.rs | 2 + .../src/tools/impls/perplexity_ask.rs | 4 + .../src/tools/impls/read_files.rs | 2 + .../src/tools/impls/search_files.rs | 3 + .../src/tools/impls/spawn_agent.rs | 172 ++++++++ .../src/tools/impls/web_fetch.rs | 4 + .../src/tools/impls/web_search.rs | 4 + 21 files changed, 717 insertions(+), 3 deletions(-) create mode 100644 crates/code_assistant/src/agent/sub_agent.rs create mode 100644 crates/code_assistant/src/tools/impls/spawn_agent.rs 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..762d0006 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -30,6 +30,10 @@ 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>, + /// Shared cancellation registry so the UI/back-end can cancel running sub-agents by tool id. + pub sub_agent_cancellation_registry: Option>, } use super::ToolSyntax; @@ -54,6 +58,8 @@ pub struct Agent { ui: Arc, state_persistence: Box, permission_handler: Option>, + sub_agent_runner: Option>, + sub_agent_cancellation_registry: Option>, // Store all messages exchanged message_history: Vec, // Store the history of tool executions @@ -111,6 +117,8 @@ impl Agent { ui, state_persistence, permission_handler, + sub_agent_runner, + sub_agent_cancellation_registry, } = components; let mut this = Self { @@ -123,6 +131,8 @@ impl Agent { command_executor, state_persistence, permission_handler, + sub_agent_runner, + sub_agent_cancellation_registry, message_history: Vec::new(), tool_executions: Vec::new(), cached_system_prompts: HashMap::new(), @@ -177,6 +187,34 @@ 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; + } + + /// Set an external cancellation flag that can interrupt streaming + pub fn set_external_cancel_flag(&mut self, _flag: Arc) { + // Note: The LLM streaming checks the UI's should_streaming_continue() method. + // For sub-agents, the SubAgentUiAdapter already checks the cancellation flag. + // This method is a placeholder for future direct cancellation integration. + } + + /// 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) { @@ -1416,6 +1454,8 @@ impl Agent { 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(), + sub_agent_cancellation_registry: self.sub_agent_cancellation_registry.as_deref(), }; // Execute the tool - could fail with ParseError or other errors 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..a8fe1760 --- /dev/null +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -0,0 +1,401 @@ +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}; +use std::time::SystemTime; + +/// 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(()) + } +} + +/// Runs sub-agents with isolated history and streams a compact progress view into the parent tool UI. +#[async_trait::async_trait] +/// 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, + cancelled: 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, + sub_agent_cancellation_registry: Some(self.cancellation_registry.clone()), + }; + + 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()?; + + // Ensure sub-agent cancellation can interrupt streaming. + agent.set_external_cancel_flag(cancelled); + + 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(), + ); + + let mut agent = self + .build_agent( + parent_tool_id, + sub_ui, + cancelled.clone(), + 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(); + for attempt in 0..=2 { + if cancelled.load(Ordering::SeqCst) { + last_answer = "Sub-agent cancelled by user.".to_string(); + break; + } + + agent.run_single_iteration().await?; + + 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); + + // Ensure the parent tool block shows completion. + let _ = self + .ui + .send_event(UiEvent::UpdateToolStatus { + tool_id: parent_tool_id.to_string(), + status: ToolStatus::Success, + message: Some("Sub-agent finished".to_string()), + output: None, + }) + .await; + + Ok(last_answer) + } +} + +fn extract_last_assistant_text(messages: &[Message]) -> Option { + messages + .iter() + .rev() + .find(|m| matches!(m.role, llm::MessageRole::Assistant)) + .map(|m| m.to_string()) +} + +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) +} + +/// A minimal UI adapter that turns sub-agent activity into a compact markdown log and streams it +/// into the parent `spawn_agent` tool block. +struct SubAgentUiAdapter { + parent: Arc, + parent_tool_id: String, + cancelled: Arc, + lines: Mutex>, +} + +impl SubAgentUiAdapter { + fn new( + parent: Arc, + parent_tool_id: String, + cancelled: Arc, + ) -> Self { + Self { + parent, + parent_tool_id, + cancelled, + lines: Mutex::new(Vec::new()), + } + } + + async fn push_line(&self, line: String) { + // Build the output string while holding the lock, then drop it before await + let output = { + let mut lines = self.lines.lock().unwrap(); + if lines.len() > 50 { + // Drop oldest lines to keep UI compact. + let drain = (lines.len() - 50) + 1; + lines.drain(0..drain); + } + lines.push(line); + + format!( + "### Sub-agent activity\n\n{}\n", + lines + .iter() + .map(|l| format!("- {l}")) + .collect::>() + .join("\n") + ) + }; // MutexGuard is dropped here + + 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(output), + }) + .await; + } +} + +#[async_trait::async_trait] +impl UserInterface for SubAgentUiAdapter { + async fn send_event(&self, event: UiEvent) -> Result<(), crate::ui::UIError> { + match event { + UiEvent::StartTool { name, .. } => { + self.push_line(format!("Calling tool `{name}`")).await; + } + UiEvent::UpdateToolStatus { + status, message, .. + } => { + if let Some(message) = message { + self.push_line(format!("Tool status: {status:?} — {message}")) + .await; + } else { + self.push_line(format!("Tool status: {status:?}")).await; + } + } + UiEvent::StreamingStopped { + cancelled, error, .. + } => { + if cancelled { + self.push_line("LLM streaming cancelled".to_string()).await; + } + if let Some(err) = error { + self.push_line(format!("LLM error: {err}")).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> { + 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 + } +} + +/// Helper for tool implementations to mark a tool id as cancelled. +pub fn cancel_sub_agent(registry: &SubAgentCancellationRegistry, tool_id: &str) -> bool { + registry.cancel(tool_id) +} + +/// Helper for tool implementations to create an initial visible tool result block. +pub fn sub_agent_tool_result_timestamps() -> (Option, Option) { + (Some(SystemTime::now()), None) +} diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index 254023e0..94e8fcaf 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -395,6 +395,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -516,6 +518,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -640,6 +644,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -778,6 +784,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -878,6 +886,8 @@ async fn test_context_compaction_inserts_summary() -> Result<()> { ui: ui.clone(), state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -995,6 +1005,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1121,6 +1133,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1596,6 +1610,8 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { ui, state_persistence, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1944,6 +1960,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1998,6 +2016,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -2066,6 +2086,8 @@ 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, + sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -2118,6 +2140,8 @@ 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, + sub_agent_cancellation_registry: 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..ccbb338e 100644 --- a/crates/code_assistant/src/mcp/handler.rs +++ b/crates/code_assistant/src/mcp/handler.rs @@ -269,6 +269,8 @@ impl MessageHandler { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; // Invoke the tool diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index 84a085f9..d1d9951e 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,26 @@ impl SessionManager { let sandboxed_project_manager = Box::new(crate::config::SandboxAwareProjectManager::new( project_manager, - sandbox_context_clone, + sandbox_context_clone.clone(), )); + // Create sub-agent runner and cancellation registry + let sub_agent_cancellation_registry = Arc::new(SubAgentCancellationRegistry::default()); + 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 +342,8 @@ impl SessionManager { ui: proxy_ui, state_persistence: state_storage, permission_handler, + sub_agent_runner: Some(sub_agent_runner), + sub_agent_cancellation_registry: Some(sub_agent_cancellation_registry), }; let mut agent = Agent::new(components, session_config.clone()); 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..7c2e49e9 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,8 @@ async fn test_edit_tool_parameter_update_after_formatting() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; // Test editing: search is formatted (matches file), replacement is unformatted @@ -271,6 +273,8 @@ async fn test_write_file_with_format_on_save() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; // Test writing a Rust file @@ -343,6 +347,8 @@ async fn test_replace_in_file_with_format_on_save() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; // Diff has two SEARCH/REPLACE blocks; replacements are unformatted (missing spaces around '=') @@ -417,6 +423,8 @@ async fn test_no_format_when_pattern_doesnt_match() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; // Test editing a .txt file (should not be formatted) @@ -493,6 +501,8 @@ async fn test_format_on_save_multiple_patterns() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; // Test editing JS file @@ -589,6 +599,8 @@ async fn test_format_on_save_glob_patterns() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, }; let tool = EditTool; @@ -678,6 +690,8 @@ async fn test_format_on_save_with_conflicting_matches() -> Result<()> { ui: None, tool_id: None, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: 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..0231a50f 100644 --- a/crates/code_assistant/src/tests/mocks.rs +++ b/crates/code_assistant/src/tests/mocks.rs @@ -231,6 +231,8 @@ pub fn create_test_tool_context<'a>( ui, tool_id, permission_handler: None, + sub_agent_runner: None, + sub_agent_cancellation_registry: None, } } @@ -1128,6 +1130,8 @@ 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, + sub_agent_cancellation_registry: None, } } 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/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/tool.rs b/crates/code_assistant/src/tools/core/tool.rs index 8236958e..101873cc 100644 --- a/crates/code_assistant/src/tools/core/tool.rs +++ b/crates/code_assistant/src/tools/core/tool.rs @@ -21,6 +21,10 @@ 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>, + /// Optional cancellation registry used for sub-agent cancellation. + pub sub_agent_cancellation_registry: Option<&'a crate::agent::SubAgentCancellationRegistry>, } /// 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..e1af17d1 --- /dev/null +++ b/crates/code_assistant/src/tools/impls/spawn_agent.rs @@ -0,0 +1,172 @@ +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)] + pub require_file_references: bool, + /// The mode for the sub-agent: "read_only" or "default". + #[serde(default = "default_mode")] + pub mode: String, +} + +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. + pub answer: String, + /// Whether the sub-agent was cancelled. + pub cancelled: bool, + /// Error message if the sub-agent failed. + pub error: Option, +} + +impl Render for SpawnAgentOutput { + fn status(&self) -> String { + if let Some(e) = &self.error { + format!("Sub-agent failed: {e}") + } else if self.cancelled { + "Sub-agent cancelled by user".to_string() + } else { + "Sub-agent completed".to_string() + } + } + + fn render(&self, _tracker: &mut ResourcesTracker) -> String { + if let Some(e) = &self.error { + return format!("Sub-agent failed: {e}"); + } + if self.cancelled { + return "Sub-agent cancelled by user.".to_string(); + } + self.answer.clone() + } +} + +impl ToolResult for SpawnAgentOutput { + fn is_success(&self) -> bool { + self.error.is_none() && !self.cancelled + } +} + +/// 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(answer) => { + let cancelled = answer == "Sub-agent cancelled by user."; + Ok(SpawnAgentOutput { + answer, + cancelled, + error: None, + }) + } + Err(e) => Ok(SpawnAgentOutput { + answer: String::new(), + cancelled: false, + error: Some(e.to_string()), + }), + } + } +} 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}'"), } From 69ea17d9fa360e31d9218fd0ad89f026fb8f30d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 15:48:28 +0100 Subject: [PATCH 04/18] Update sub-agents-feature.md --- docs/sub-agents-feature.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/docs/sub-agents-feature.md b/docs/sub-agents-feature.md index 516043c4..d10ddee2 100644 --- a/docs/sub-agents-feature.md +++ b/docs/sub-agents-feature.md @@ -1,5 +1,33 @@ # 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`, `set_external_cancel_flag`, `message_history` +- [x] File reference enforcement with retry logic (up to 2 retries) + +### Pending +- [ ] **UI integration for cancellation**: Expose cancel button per running `spawn_agent` block +- [ ] **Parallel execution**: Update `manage_tool_execution()` to run multiple `spawn_agent` calls concurrently +- [ ] **Permission attribution**: Show permission requests as originating from sub-agent context (inline or popover) +- [ ] **Terminal UI rendering**: Update terminal tool block to show streaming sub-agent activity +- [ ] **GPUI rendering**: Update GPUI tool widget for sub-agent progress display +- [ ] **Integration tests**: Add tests for isolation, parallelism, cancellation, and permission routing + +### Notes +- The cancellation infrastructure is in place (`SubAgentCancellationRegistry`, `cancel` method, `cancel_sub_agent` helper) but the UI hooks to trigger cancellation are not yet implemented. +- The `sub_agent_cancellation_registry` field in `ToolContext` is available for future use when implementing tool-level cancellation from UI. + +--- + ## 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. From 452237c68f0907185a1173cca37fa0e500e23974 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 17:05:20 +0100 Subject: [PATCH 05/18] WIP: spawn_agent output rendering --- Cargo.lock | 1 + crates/code_assistant/Cargo.toml | 1 + crates/code_assistant/src/agent/runner.rs | 296 ++++++++-- crates/code_assistant/src/tests/mod.rs | 1 + .../src/tests/sub_agent_tests.rs | 344 ++++++++++++ crates/code_assistant/src/ui/gpui/elements.rs | 60 ++- .../code_assistant/src/ui/gpui/file_icons.rs | 7 + crates/code_assistant/src/ui/gpui/mod.rs | 7 + .../src/ui/gpui/tool_output_renderers.rs | 507 ++++++++++++++++++ .../code_assistant/src/ui/terminal/message.rs | 8 +- .../src/ui/terminal/tool_widget.rs | 50 ++ docs/sub-agents-feature.md | 34 +- 12 files changed, 1261 insertions(+), 55 deletions(-) create mode 100644 crates/code_assistant/src/tests/sub_agent_tests.rs create mode 100644 crates/code_assistant/src/ui/gpui/tool_output_renderers.rs 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/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 762d0006..bd2b838f 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}; @@ -667,46 +667,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, + sub_agent_cancellation_registry: None, + }; + + 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(); + let output_text = output.render(&mut resources_tracker); + + let _ = ui + .send_event(UiEvent::UpdateToolStatus { + tool_id: tool_id.to_string(), + status, + message: Some(status_msg), + output: Some(output_text), + }) + .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<()> { 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..39c43e45 --- /dev/null +++ b/crates/code_assistant/src/tests/sub_agent_tests.rs @@ -0,0 +1,344 @@ +//! Tests for the sub-agent feature (spawn_agent tool). + +use crate::agent::sub_agent::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(self.response.clone()) + } +} + +#[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(), + cancelled: false, + error: 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 + let output = SpawnAgentOutput { + answer: String::new(), + cancelled: true, + error: None, + }; + + let rendered = output.render(&mut tracker); + assert_eq!(rendered, "Sub-agent cancelled by user."); + assert_eq!(output.status(), "Sub-agent cancelled by user"); + + // Test error output + let output = SpawnAgentOutput { + answer: String::new(), + cancelled: false, + error: Some("Connection failed".to_string()), + }; + + let rendered = output.render(&mut tracker); + assert_eq!(rendered, "Sub-agent failed: Connection failed"); + assert!(output.status().contains("Connection failed")); +} + +#[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(), "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(), "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/ui/gpui/elements.rs b/crates/code_assistant/src/ui/gpui/elements.rs index a8c367b5..1269cd7e 100644 --- a/crates/code_assistant/src/ui/gpui/elements.rs +++ b/crates/code_assistant/src/ui/gpui/elements.rs @@ -1356,24 +1356,48 @@ impl Render for BlockView { }; 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 + let custom_output = crate::ui::gpui::tool_output_renderers::ToolOutputRendererRegistry::global() + .and_then(|registry| { + registry.render_output( + &block.name, + output_content, + &block.status, + cx.theme(), + ) + }); + + 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..f00a6b86 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -17,6 +17,7 @@ mod root; pub mod sandbox_selector; pub mod simple_renderers; pub mod theme; +pub mod tool_output_renderers; use crate::persistence::{ChatMetadata, DraftStorage}; use crate::types::PlanState; @@ -27,6 +28,7 @@ use crate::ui::gpui::{ elements::MessageRole, parameter_renderers::{DefaultParameterRenderer, ParameterRendererRegistry}, simple_renderers::SimpleParameterRenderer, + tool_output_renderers::{SpawnAgentOutputRenderer, ToolOutputRendererRegistry}, }; use crate::ui::{async_trait, DisplayFragment, UIError, UiEvent, UserInterface}; use assets::Assets; @@ -241,6 +243,11 @@ impl Gpui { // 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)); 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..88990312 --- /dev/null +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -0,0 +1,507 @@ +use crate::ui::ToolStatus; +use gpui::{div, px, Element, ParentElement, SharedString, Styled}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; +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_name: &str, + output: &str, + status: &ToolStatus, + theme: &gpui_component::theme::Theme, + ) -> 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 + 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) + pub fn render_output( + &self, + tool_name: &str, + output: &str, + status: &ToolStatus, + theme: &gpui_component::theme::Theme, + ) -> Option { + self.renderers + .get(tool_name) + .and_then(|renderer| renderer.render(tool_name, output, status, theme)) + } +} + +impl Default for ToolOutputRendererRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Parsed sub-agent activity line +#[derive(Debug, Clone)] +pub enum SubAgentActivityLine { + /// Tool call: "Calling tool `tool_name`" + ToolCall { tool_name: String }, + /// Tool status: "Tool status: Success" or "Tool status: Running — message" + ToolStatus { + status: SubAgentToolStatus, + message: Option, + }, + /// LLM streaming cancelled + Cancelled, + /// LLM error + Error { message: String }, + /// Unknown/unparsed line + Other { text: String }, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum SubAgentToolStatus { + Pending, + Running, + Success, + Error, +} + +impl SubAgentActivityLine { + /// Parse a single line from sub-agent output + fn parse(line: &str) -> Self { + let line = line.trim(); + + // Remove leading "- " if present + let line = line.strip_prefix("- ").unwrap_or(line); + + // Check for tool call + if let Some(rest) = line.strip_prefix("Calling tool `") { + if let Some(tool_name) = rest.strip_suffix('`') { + return Self::ToolCall { + tool_name: tool_name.to_string(), + }; + } + } + + // Check for tool status + if let Some(rest) = line.strip_prefix("Tool status: ") { + // Format: "Status" or "Status — message" + let (status_str, message) = if let Some(idx) = rest.find(" — ") { + (&rest[..idx], Some(rest[idx + " — ".len()..].to_string())) + } else { + (rest, None) + }; + + let status = match status_str { + "Pending" => SubAgentToolStatus::Pending, + "Running" => SubAgentToolStatus::Running, + "Success" => SubAgentToolStatus::Success, + "Error" => SubAgentToolStatus::Error, + _ => SubAgentToolStatus::Running, // Default to running for unknown + }; + + return Self::ToolStatus { status, message }; + } + + // Check for LLM events + if line == "LLM streaming cancelled" { + return Self::Cancelled; + } + + if let Some(rest) = line.strip_prefix("LLM error: ") { + return Self::Error { + message: rest.to_string(), + }; + } + + // Fallback + Self::Other { + text: line.to_string(), + } + } +} + +/// Parsed sub-agent output with activity lines +#[derive(Debug, Clone)] +pub struct ParsedSubAgentOutput { + pub activities: Vec, +} + +impl ParsedSubAgentOutput { + /// Parse the full sub-agent output + pub fn parse(output: &str) -> Self { + let mut activities = Vec::new(); + + for line in output.lines() { + let line = line.trim(); + + // Skip the header + if line == "### Sub-agent activity" || line.is_empty() { + continue; + } + + activities.push(SubAgentActivityLine::parse(line)); + } + + Self { activities } + } + + /// Get the most recent tool call that is still "active" (last call without a success/error) + pub fn get_active_tool(&self) -> Option<&str> { + let mut last_tool: Option<&str> = None; + + for activity in &self.activities { + match activity { + SubAgentActivityLine::ToolCall { tool_name } => { + last_tool = Some(tool_name.as_str()); + } + SubAgentActivityLine::ToolStatus { status, .. } + if *status == SubAgentToolStatus::Success + || *status == SubAgentToolStatus::Error => + { + last_tool = None; // Tool completed + } + _ => {} + } + } + + last_tool + } + + /// Get a summary of completed tools for compact display + pub fn get_tool_summary(&self) -> Vec { + let mut tools: Vec = Vec::new(); + let mut current_tool: Option = None; + + for activity in &self.activities { + match activity { + SubAgentActivityLine::ToolCall { tool_name } => { + current_tool = Some(tool_name.clone()); + } + SubAgentActivityLine::ToolStatus { status, message } => { + if let Some(tool_name) = current_tool.take() { + let tool_status = match status { + SubAgentToolStatus::Success => CompactToolStatus::Success, + SubAgentToolStatus::Error => CompactToolStatus::Error, + SubAgentToolStatus::Running => CompactToolStatus::Running, + SubAgentToolStatus::Pending => CompactToolStatus::Pending, + }; + tools.push(CompactToolInfo { + name: tool_name, + status: tool_status, + message: message.clone(), + }); + } + } + _ => {} + } + } + + // If there's still a current tool without status, it's running + if let Some(tool_name) = current_tool { + tools.push(CompactToolInfo { + name: tool_name, + status: CompactToolStatus::Running, + message: None, + }); + } + + tools + } +} + +#[derive(Debug, Clone)] +pub struct CompactToolInfo { + pub name: String, + pub status: CompactToolStatus, + pub message: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum CompactToolStatus { + Pending, + Running, + Success, + Error, +} + +/// Renderer for spawn_agent tool output +/// Displays sub-agent tool calls in a compact, Zed-like style +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) + } + + /// Get a display title for a tool (like Zed's ACP mode) + fn get_tool_title(tool_name: &str) -> String { + match tool_name { + "read_files" => "Reading".to_string(), + "search_files" => "Searching".to_string(), + "list_files" => "Listing".to_string(), + "glob_files" => "Finding files".to_string(), + "write_file" => "Writing".to_string(), + "edit" => "Editing".to_string(), + "replace_in_file" => "Replacing".to_string(), + "delete_files" => "Deleting".to_string(), + "execute_command" => "Executing".to_string(), + "web_fetch" => "Fetching".to_string(), + "web_search" => "Searching web".to_string(), + "perplexity_ask" => "Asking Perplexity".to_string(), + _ => tool_name.replace('_', " "), + } + } + + /// Render a single compact tool line + fn render_tool_line( + tool: &CompactToolInfo, + theme: &gpui_component::theme::Theme, + ) -> gpui::AnyElement { + use super::file_icons; + + let icon = Self::get_tool_icon(&tool.name); + let title = Self::get_tool_title(&tool.name); + + // Status-based colors + let (icon_color, text_color) = match tool.status { + CompactToolStatus::Running => (theme.info, theme.muted_foreground), + CompactToolStatus::Success => (theme.success, theme.muted_foreground), + CompactToolStatus::Error => (theme.danger, theme.danger), + CompactToolStatus::Pending => (theme.muted_foreground, theme.muted_foreground), + }; + + // Build the display text + let display_text = if let Some(msg) = &tool.message { + if msg.is_empty() { + title + } else { + format!("{} {}", title, msg) + } + } else { + title + }; + + 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 + div() + .text_size(px(13.)) + .text_color(text_color) + .child(display_text) + .into_any(), + ]) + .into_any() + } +} + +impl ToolOutputRenderer for SpawnAgentOutputRenderer { + fn supported_tools(&self) -> Vec { + vec!["spawn_agent".to_string()] + } + + fn render( + &self, + _tool_name: &str, + output: &str, + _status: &ToolStatus, + theme: &gpui_component::theme::Theme, + ) -> Option { + if output.is_empty() { + return None; + } + + let parsed = ParsedSubAgentOutput::parse(output); + let tools = parsed.get_tool_summary(); + + if tools.is_empty() { + return None; + } + + // Render compact list of tool calls + let tool_elements: Vec = tools + .iter() + .map(|tool| Self::render_tool_line(tool, theme)) + .collect(); + + Some( + div() + .flex() + .flex_col() + .gap_0() + .mt_1() + .children(tool_elements) + .into_any(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_sub_agent_output() { + let output = r#"### Sub-agent activity + +- Calling tool `read_files` +- Tool status: Success +- Calling tool `search_files` +- Tool status: Running — Searching for pattern +"#; + + let parsed = ParsedSubAgentOutput::parse(output); + assert_eq!(parsed.activities.len(), 4); + + // First activity: tool call + match &parsed.activities[0] { + SubAgentActivityLine::ToolCall { tool_name } => { + assert_eq!(tool_name, "read_files"); + } + _ => panic!("Expected ToolCall"), + } + + // Second activity: status + match &parsed.activities[1] { + SubAgentActivityLine::ToolStatus { status, message } => { + assert_eq!(*status, SubAgentToolStatus::Success); + assert!(message.is_none()); + } + _ => panic!("Expected ToolStatus"), + } + + // Fourth activity: status with message + match &parsed.activities[3] { + SubAgentActivityLine::ToolStatus { status, message } => { + assert_eq!(*status, SubAgentToolStatus::Running); + assert_eq!(message.as_deref(), Some("Searching for pattern")); + } + _ => panic!("Expected ToolStatus"), + } + } + + #[test] + fn test_get_tool_summary() { + let output = r#"### Sub-agent activity + +- Calling tool `read_files` +- Tool status: Success +- Calling tool `search_files` +"#; + + let parsed = ParsedSubAgentOutput::parse(output); + let tools = parsed.get_tool_summary(); + + assert_eq!(tools.len(), 2); + assert_eq!(tools[0].name, "read_files"); + assert_eq!(tools[0].status, CompactToolStatus::Success); + assert_eq!(tools[1].name, "search_files"); + assert_eq!(tools[1].status, CompactToolStatus::Running); + } + + #[test] + fn test_get_active_tool() { + let output = r#"### Sub-agent activity + +- Calling tool `read_files` +- Tool status: Success +- Calling tool `search_files` +"#; + + let parsed = ParsedSubAgentOutput::parse(output); + assert_eq!(parsed.get_active_tool(), Some("search_files")); + } + + #[test] + fn test_parse_error_and_cancel() { + let output = r#"### Sub-agent activity + +- Calling tool `read_files` +- LLM streaming cancelled +- LLM error: Connection failed +"#; + + let parsed = ParsedSubAgentOutput::parse(output); + assert_eq!(parsed.activities.len(), 3); + + match &parsed.activities[1] { + SubAgentActivityLine::Cancelled => {} + _ => panic!("Expected Cancelled"), + } + + match &parsed.activities[2] { + SubAgentActivityLine::Error { message } => { + assert_eq!(message, "Connection failed"); + } + _ => panic!("Expected Error"), + } + } +} 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..496e5449 100644 --- a/crates/code_assistant/src/ui/terminal/tool_widget.rs +++ b/crates/code_assistant/src/ui/terminal/tool_widget.rs @@ -228,6 +228,56 @@ 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() { + // Render each line of the output + for line in output.lines() { + if current_y >= area.y + area.height { + break; + } + + // Determine line style based on content + let (prefix, color) = if line.starts_with("### ") { + // Header line (e.g. "### Sub-agent activity") + ("", Color::Cyan) + } else if line.starts_with("- Calling tool") { + // Tool call line + ("", Color::Yellow) + } else if line.starts_with("- Tool status:") { + // Tool status line + ("", Color::Gray) + } else if line.starts_with("- ") { + // Other list items + ("", Color::White) + } else { + // Regular text + ("", Color::Gray) + }; + + let display_line = format!("{prefix}{line}"); + let truncated = if display_line.len() > (area.width.saturating_sub(4)) as usize + { + format!( + "{}...", + &display_line[..(area.width.saturating_sub(7)) as usize] + ) + } else { + display_line + }; + + buf.set_string( + area.x + 2, + current_y, + &truncated, + Style::default().fg(color), + ); + current_y += 1; + } } } } diff --git a/docs/sub-agents-feature.md b/docs/sub-agents-feature.md index d10ddee2..6874d543 100644 --- a/docs/sub-agents-feature.md +++ b/docs/sub-agents-feature.md @@ -14,13 +14,39 @@ - [x] Added required `Agent` methods: `set_tool_scope`, `set_session_model_config`, `set_session_identity`, `set_external_cancel_flag`, `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] **GPUI rendering**: Custom tool output renderer for sub-agent progress display + - Added `ToolOutputRendererRegistry` pattern (similar to `ParameterRendererRegistry`) + - Implemented `SpawnAgentOutputRenderer` that parses sub-agent activity markdown + - Renders sub-tool calls in compact Zed-like style with icons and status colors + - Located in `crates/code_assistant/src/ui/gpui/tool_output_renderers.rs` +- [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`) + - Output flows as `ToolCallUpdate` content for display in Zed's ACP panel + ### Pending - [ ] **UI integration for cancellation**: Expose cancel button per running `spawn_agent` block -- [ ] **Parallel execution**: Update `manage_tool_execution()` to run multiple `spawn_agent` calls concurrently - [ ] **Permission attribution**: Show permission requests as originating from sub-agent context (inline or popover) -- [ ] **Terminal UI rendering**: Update terminal tool block to show streaming sub-agent activity -- [ ] **GPUI rendering**: Update GPUI tool widget for sub-agent progress display -- [ ] **Integration tests**: Add tests for isolation, parallelism, cancellation, and permission routing ### Notes - The cancellation infrastructure is in place (`SubAgentCancellationRegistry`, `cancel` method, `cancel_sub_agent` helper) but the UI hooks to trigger cancellation are not yet implemented. From 32858e7cc3a0d164bb4c9be7e83e76f673ef24ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 18:03:52 +0100 Subject: [PATCH 06/18] Now tools correctly arrive in the spawn_agent tool block --- crates/code_assistant/src/agent/sub_agent.rs | 209 +++++++-- .../src/ui/gpui/tool_output_renderers.rs | 397 ++++++------------ .../src/ui/terminal/tool_widget.rs | 140 ++++-- docs/sub-agents-feature.md | 16 +- 4 files changed, 405 insertions(+), 357 deletions(-) diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index a8fe1760..a0356940 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -279,13 +279,78 @@ fn has_file_references_with_line_ranges(text: &str) -> bool { .unwrap_or(false) } -/// A minimal UI adapter that turns sub-agent activity into a compact markdown log and streams it +/// 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, + #[serde(skip_serializing_if = "Option::is_none")] + pub message: Option, +} + +/// 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, + } + } +} + +/// 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 cancelled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +impl SubAgentOutput { + pub fn new() -> Self { + Self { + tools: Vec::new(), + cancelled: None, + error: 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, - lines: Mutex>, + output: Mutex, + /// Map from tool_id to index in output.tools for fast lookup + tool_id_to_index: Mutex>, } impl SubAgentUiAdapter { @@ -298,30 +363,16 @@ impl SubAgentUiAdapter { parent, parent_tool_id, cancelled, - lines: Mutex::new(Vec::new()), + output: Mutex::new(SubAgentOutput::new()), + tool_id_to_index: Mutex::new(std::collections::HashMap::new()), } } - async fn push_line(&self, line: String) { - // Build the output string while holding the lock, then drop it before await - let output = { - let mut lines = self.lines.lock().unwrap(); - if lines.len() > 50 { - // Drop oldest lines to keep UI compact. - let drain = (lines.len() - 50) + 1; - lines.drain(0..drain); - } - lines.push(line); - - format!( - "### Sub-agent activity\n\n{}\n", - lines - .iter() - .map(|l| format!("- {l}")) - .collect::>() - .join("\n") - ) - }; // MutexGuard is dropped here + async fn send_output_update(&self) { + let json = { + let output = self.output.lock().unwrap(); + output.to_json() + }; let _ = self .parent @@ -329,38 +380,98 @@ impl SubAgentUiAdapter { tool_id: self.parent_tool_id.clone(), status: ToolStatus::Running, message: Some("Sub-agent running".to_string()), - output: Some(output), + 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(); + + // Add new tool as running + let index = output.tools.len(); + output.tools.push(SubAgentToolCall { + name: name.to_string(), + status: SubAgentToolStatus::Running, + message: None, + }); + id_map.insert(id.to_string(), index); + } + + fn update_tool_status(&self, tool_id: &str, status: ToolStatus, message: Option) { + let mut output = self.output.lock().unwrap(); + let 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) { + tool.status = status.into(); + tool.message = message; + } + } else { + tracing::warn!( + "SubAgentUiAdapter: UpdateToolStatus for unknown tool_id={}", + tool_id + ); + } + } + + 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); + } } #[async_trait::async_trait] impl UserInterface for SubAgentUiAdapter { async fn send_event(&self, event: UiEvent) -> Result<(), crate::ui::UIError> { - match event { - UiEvent::StartTool { name, .. } => { - self.push_line(format!("Calling tool `{name}`")).await; + match &event { + UiEvent::StartTool { name, id } => { + // This event is typically not sent directly (GPUI creates it from DisplayFragment) + // but handle it for completeness in case other UIs send it + tracing::debug!( + "SubAgentUiAdapter: StartTool event received: {} ({})", + name, + id + ); + self.add_tool_start(name, id); + self.send_output_update().await; } UiEvent::UpdateToolStatus { - status, message, .. + tool_id, + status, + message, + .. } => { - if let Some(message) = message { - self.push_line(format!("Tool status: {status:?} — {message}")) - .await; - } else { - self.push_line(format!("Tool status: {status:?}")).await; - } + tracing::debug!( + "SubAgentUiAdapter: UpdateToolStatus event - tool_id={}, status={:?}", + tool_id, + status + ); + self.update_tool_status(tool_id, *status, message.clone()); + self.send_output_update().await; } UiEvent::StreamingStopped { cancelled, error, .. } => { - if cancelled { - self.push_line("LLM streaming cancelled".to_string()).await; + tracing::debug!( + "SubAgentUiAdapter: StreamingStopped - cancelled={}, error={:?}", + cancelled, + error + ); + if *cancelled { + self.set_cancelled(); } if let Some(err) = error { - self.push_line(format!("LLM error: {err}")).await; + self.set_error(err.clone()); } + self.send_output_update().await; } _ => { // Ignore other events; they belong to the sub-agent's isolated transcript. @@ -372,8 +483,30 @@ impl UserInterface for SubAgentUiAdapter { fn display_fragment( &self, - _fragment: &crate::ui::DisplayFragment, + 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 + // Note: This is called during LLM streaming when the tool name is parsed + // The UpdateToolStatus event with Running status comes later when execution starts + tracing::debug!( + "SubAgentUiAdapter: ToolName fragment received: {} ({})", + name, + id + ); + self.add_tool_start(name, id); + // We can't send async update here, but the subsequent UpdateToolStatus + // event will trigger send_output_update() + } + _ => { + // Ignore other fragments (text, parameters, etc.) + // They belong to the sub-agent's isolated transcript + } + } + Ok(()) } diff --git a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs index 88990312..ad968327 100644 --- a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -1,3 +1,4 @@ +use crate::agent::sub_agent::{SubAgentOutput, SubAgentToolStatus}; use crate::ui::ToolStatus; use gpui::{div, px, Element, ParentElement, SharedString, Styled}; use std::collections::HashMap; @@ -74,6 +75,7 @@ impl ToolOutputRendererRegistry { } /// 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) } @@ -99,191 +101,6 @@ impl Default for ToolOutputRendererRegistry { } } -/// Parsed sub-agent activity line -#[derive(Debug, Clone)] -pub enum SubAgentActivityLine { - /// Tool call: "Calling tool `tool_name`" - ToolCall { tool_name: String }, - /// Tool status: "Tool status: Success" or "Tool status: Running — message" - ToolStatus { - status: SubAgentToolStatus, - message: Option, - }, - /// LLM streaming cancelled - Cancelled, - /// LLM error - Error { message: String }, - /// Unknown/unparsed line - Other { text: String }, -} - -#[derive(Debug, Clone, PartialEq)] -pub enum SubAgentToolStatus { - Pending, - Running, - Success, - Error, -} - -impl SubAgentActivityLine { - /// Parse a single line from sub-agent output - fn parse(line: &str) -> Self { - let line = line.trim(); - - // Remove leading "- " if present - let line = line.strip_prefix("- ").unwrap_or(line); - - // Check for tool call - if let Some(rest) = line.strip_prefix("Calling tool `") { - if let Some(tool_name) = rest.strip_suffix('`') { - return Self::ToolCall { - tool_name: tool_name.to_string(), - }; - } - } - - // Check for tool status - if let Some(rest) = line.strip_prefix("Tool status: ") { - // Format: "Status" or "Status — message" - let (status_str, message) = if let Some(idx) = rest.find(" — ") { - (&rest[..idx], Some(rest[idx + " — ".len()..].to_string())) - } else { - (rest, None) - }; - - let status = match status_str { - "Pending" => SubAgentToolStatus::Pending, - "Running" => SubAgentToolStatus::Running, - "Success" => SubAgentToolStatus::Success, - "Error" => SubAgentToolStatus::Error, - _ => SubAgentToolStatus::Running, // Default to running for unknown - }; - - return Self::ToolStatus { status, message }; - } - - // Check for LLM events - if line == "LLM streaming cancelled" { - return Self::Cancelled; - } - - if let Some(rest) = line.strip_prefix("LLM error: ") { - return Self::Error { - message: rest.to_string(), - }; - } - - // Fallback - Self::Other { - text: line.to_string(), - } - } -} - -/// Parsed sub-agent output with activity lines -#[derive(Debug, Clone)] -pub struct ParsedSubAgentOutput { - pub activities: Vec, -} - -impl ParsedSubAgentOutput { - /// Parse the full sub-agent output - pub fn parse(output: &str) -> Self { - let mut activities = Vec::new(); - - for line in output.lines() { - let line = line.trim(); - - // Skip the header - if line == "### Sub-agent activity" || line.is_empty() { - continue; - } - - activities.push(SubAgentActivityLine::parse(line)); - } - - Self { activities } - } - - /// Get the most recent tool call that is still "active" (last call without a success/error) - pub fn get_active_tool(&self) -> Option<&str> { - let mut last_tool: Option<&str> = None; - - for activity in &self.activities { - match activity { - SubAgentActivityLine::ToolCall { tool_name } => { - last_tool = Some(tool_name.as_str()); - } - SubAgentActivityLine::ToolStatus { status, .. } - if *status == SubAgentToolStatus::Success - || *status == SubAgentToolStatus::Error => - { - last_tool = None; // Tool completed - } - _ => {} - } - } - - last_tool - } - - /// Get a summary of completed tools for compact display - pub fn get_tool_summary(&self) -> Vec { - let mut tools: Vec = Vec::new(); - let mut current_tool: Option = None; - - for activity in &self.activities { - match activity { - SubAgentActivityLine::ToolCall { tool_name } => { - current_tool = Some(tool_name.clone()); - } - SubAgentActivityLine::ToolStatus { status, message } => { - if let Some(tool_name) = current_tool.take() { - let tool_status = match status { - SubAgentToolStatus::Success => CompactToolStatus::Success, - SubAgentToolStatus::Error => CompactToolStatus::Error, - SubAgentToolStatus::Running => CompactToolStatus::Running, - SubAgentToolStatus::Pending => CompactToolStatus::Pending, - }; - tools.push(CompactToolInfo { - name: tool_name, - status: tool_status, - message: message.clone(), - }); - } - } - _ => {} - } - } - - // If there's still a current tool without status, it's running - if let Some(tool_name) = current_tool { - tools.push(CompactToolInfo { - name: tool_name, - status: CompactToolStatus::Running, - message: None, - }); - } - - tools - } -} - -#[derive(Debug, Clone)] -pub struct CompactToolInfo { - pub name: String, - pub status: CompactToolStatus, - pub message: Option, -} - -#[derive(Debug, Clone, PartialEq)] -pub enum CompactToolStatus { - Pending, - Running, - Success, - Error, -} - /// Renderer for spawn_agent tool output /// Displays sub-agent tool calls in a compact, Zed-like style pub struct SpawnAgentOutputRenderer; @@ -314,9 +131,9 @@ impl SpawnAgentOutputRenderer { } } - /// Render a single compact tool line + /// Render a single compact tool line from structured data fn render_tool_line( - tool: &CompactToolInfo, + tool: &crate::agent::sub_agent::SubAgentToolCall, theme: &gpui_component::theme::Theme, ) -> gpui::AnyElement { use super::file_icons; @@ -326,10 +143,9 @@ impl SpawnAgentOutputRenderer { // Status-based colors let (icon_color, text_color) = match tool.status { - CompactToolStatus::Running => (theme.info, theme.muted_foreground), - CompactToolStatus::Success => (theme.success, theme.muted_foreground), - CompactToolStatus::Error => (theme.danger, theme.danger), - CompactToolStatus::Pending => (theme.muted_foreground, theme.muted_foreground), + SubAgentToolStatus::Running => (theme.info, theme.muted_foreground), + SubAgentToolStatus::Success => (theme.success, theme.muted_foreground), + SubAgentToolStatus::Error => (theme.danger, theme.danger), }; // Build the display text @@ -337,7 +153,7 @@ impl SpawnAgentOutputRenderer { if msg.is_empty() { title } else { - format!("{} {}", title, msg) + format!("{title} — {msg}") } } else { title @@ -361,6 +177,50 @@ impl SpawnAgentOutputRenderer { ]) .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 + } } impl ToolOutputRenderer for SpawnAgentOutputRenderer { @@ -379,26 +239,39 @@ impl ToolOutputRenderer for SpawnAgentOutputRenderer { return None; } - let parsed = ParsedSubAgentOutput::parse(output); - let tools = parsed.get_tool_summary(); + // 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; + } + }; - if tools.is_empty() { + if parsed.tools.is_empty() && parsed.cancelled.is_none() && parsed.error.is_none() { return None; } // Render compact list of tool calls - let tool_elements: Vec = tools + let mut elements: Vec = parsed + .tools .iter() .map(|tool| Self::render_tool_line(tool, theme)) .collect(); + // Add status line if present + if let Some(status_line) = Self::render_status_line(&parsed, theme) { + elements.push(status_line); + } + Some( div() .flex() .flex_col() .gap_0() .mt_1() - .children(tool_elements) + .children(elements) .into_any(), ) } @@ -409,99 +282,67 @@ mod tests { use super::*; #[test] - fn test_parse_sub_agent_output() { - let output = r#"### Sub-agent activity - -- Calling tool `read_files` -- Tool status: Success -- Calling tool `search_files` -- Tool status: Running — Searching for pattern -"#; - - let parsed = ParsedSubAgentOutput::parse(output); - assert_eq!(parsed.activities.len(), 4); - - // First activity: tool call - match &parsed.activities[0] { - SubAgentActivityLine::ToolCall { tool_name } => { - assert_eq!(tool_name, "read_files"); - } - _ => panic!("Expected ToolCall"), - } - - // Second activity: status - match &parsed.activities[1] { - SubAgentActivityLine::ToolStatus { status, message } => { - assert_eq!(*status, SubAgentToolStatus::Success); - assert!(message.is_none()); - } - _ => panic!("Expected ToolStatus"), - } - - // Fourth activity: status with message - match &parsed.activities[3] { - SubAgentActivityLine::ToolStatus { status, message } => { - assert_eq!(*status, SubAgentToolStatus::Running); - assert_eq!(message.as_deref(), Some("Searching for pattern")); - } - _ => panic!("Expected ToolStatus"), - } + 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_get_tool_summary() { - let output = r#"### Sub-agent activity - -- Calling tool `read_files` -- Tool status: Success -- Calling tool `search_files` -"#; - - let parsed = ParsedSubAgentOutput::parse(output); - let tools = parsed.get_tool_summary(); - - assert_eq!(tools.len(), 2); - assert_eq!(tools[0].name, "read_files"); - assert_eq!(tools[0].status, CompactToolStatus::Success); - assert_eq!(tools[1].name, "search_files"); - assert_eq!(tools[1].status, CompactToolStatus::Running); + 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_get_active_tool() { - let output = r#"### Sub-agent activity + fn test_parse_json_with_error() { + let json = r#"{"tools":[],"error":"Connection failed"}"#; -- Calling tool `read_files` -- Tool status: Success -- Calling tool `search_files` -"#; - - let parsed = ParsedSubAgentOutput::parse(output); - assert_eq!(parsed.get_active_tool(), Some("search_files")); + let parsed = SubAgentOutput::from_json(json).unwrap(); + assert_eq!(parsed.error.as_deref(), Some("Connection failed")); } #[test] - fn test_parse_error_and_cancel() { - let output = r#"### Sub-agent activity - -- Calling tool `read_files` -- LLM streaming cancelled -- LLM error: Connection failed -"#; + fn test_roundtrip() { + let mut output = SubAgentOutput::new(); + output + .tools + .push(crate::agent::sub_agent::SubAgentToolCall { + name: "read_files".to_string(), + status: SubAgentToolStatus::Success, + message: None, + }); + output + .tools + .push(crate::agent::sub_agent::SubAgentToolCall { + name: "search_files".to_string(), + status: SubAgentToolStatus::Running, + message: Some("Searching for pattern".to_string()), + }); - let parsed = ParsedSubAgentOutput::parse(output); - assert_eq!(parsed.activities.len(), 3); + let json = output.to_json(); + let parsed = SubAgentOutput::from_json(&json).unwrap(); - match &parsed.activities[1] { - SubAgentActivityLine::Cancelled => {} - _ => panic!("Expected Cancelled"), - } + 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") + ); + } - match &parsed.activities[2] { - SubAgentActivityLine::Error { message } => { - assert_eq!(message, "Connection failed"); - } - _ => panic!("Expected Error"), - } + #[test] + fn test_invalid_json_returns_none() { + let invalid = "### Sub-agent activity\n- Calling tool read_files"; + assert!(SubAgentOutput::from_json(invalid).is_none()); } } diff --git a/crates/code_assistant/src/ui/terminal/tool_widget.rs b/crates/code_assistant/src/ui/terminal/tool_widget.rs index 496e5449..cbf73ac3 100644 --- a/crates/code_assistant/src/ui/terminal/tool_widget.rs +++ b/crates/code_assistant/src/ui/terminal/tool_widget.rs @@ -235,48 +235,110 @@ impl<'a> Widget for ToolWidget<'a> { // 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() { - // Render each line of the output - for line in output.lines() { - if current_y >= area.y + area.height { - break; + // 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), + }; + + let tool_title = match tool.name.as_str() { + "read_files" => "Reading", + "search_files" => "Searching", + "list_files" => "Listing", + "glob_files" => "Finding files", + "write_file" => "Writing", + "edit" => "Editing", + "execute_command" => "Executing", + "web_fetch" => "Fetching", + "web_search" => "Searching web", + _ => &tool.name, + }; + + let display_text = if let Some(msg) = &tool.message { + format!("{status_symbol} {tool_title} — {msg}") + } else { + format!("{status_symbol} {tool_title}") + }; + + let truncated = + if display_text.len() > (area.width.saturating_sub(4)) as usize { + format!( + "{}...", + &display_text[..(area.width.saturating_sub(7)) as usize] + ) + } else { + display_text + }; + + buf.set_string( + area.x + 2, + current_y, + &truncated, + Style::default().fg(color), + ); + current_y += 1; } - // Determine line style based on content - let (prefix, color) = if line.starts_with("### ") { - // Header line (e.g. "### Sub-agent activity") - ("", Color::Cyan) - } else if line.starts_with("- Calling tool") { - // Tool call line - ("", Color::Yellow) - } else if line.starts_with("- Tool status:") { - // Tool status line - ("", Color::Gray) - } else if line.starts_with("- ") { - // Other list items - ("", Color::White) - } else { - // Regular text - ("", Color::Gray) - }; - - let display_line = format!("{prefix}{line}"); - let truncated = if display_line.len() > (area.width.saturating_sub(4)) as usize + // Render cancelled/error status + if sub_agent_output.cancelled == Some(true) && current_y < area.y + area.height { - format!( - "{}...", - &display_line[..(area.width.saturating_sub(7)) as usize] - ) - } else { - display_line - }; - - buf.set_string( - area.x + 2, - current_y, - &truncated, - Style::default().fg(color), - ); - current_y += 1; + 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/docs/sub-agents-feature.md b/docs/sub-agents-feature.md index 6874d543..0d2b981f 100644 --- a/docs/sub-agents-feature.md +++ b/docs/sub-agents-feature.md @@ -34,15 +34,27 @@ - `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 sub-agent activity markdown + - 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`) - - Output flows as `ToolCallUpdate` content for display in Zed's ACP panel + - JSON output flows as `ToolCallUpdate` content for display in Zed's ACP panel ### Pending - [ ] **UI integration for cancellation**: Expose cancel button per running `spawn_agent` block From a3b3c4fe093a2f8668a82b1fe1167c5760b14016 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 19:15:14 +0100 Subject: [PATCH 07/18] Show sub-agent activity --- crates/code_assistant/src/agent/sub_agent.rs | 46 ++++++++- .../src/ui/gpui/tool_output_renderers.rs | 98 +++++++++++++++++-- 2 files changed, 133 insertions(+), 11 deletions(-) diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index a0356940..6cd8456e 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -308,11 +308,31 @@ impl From for SubAgentToolStatus { } } +/// 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, @@ -322,6 +342,7 @@ impl SubAgentOutput { pub fn new() -> Self { Self { tools: Vec::new(), + activity: Some(SubAgentActivity::WaitingForLlm), cancelled: None, error: None, } @@ -425,6 +446,12 @@ impl SubAgentUiAdapter { 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); } } @@ -443,6 +470,7 @@ impl UserInterface for SubAgentUiAdapter { self.add_tool_start(name, id); self.send_output_update().await; } + UiEvent::UpdateToolStatus { tool_id, status, @@ -455,6 +483,16 @@ impl UserInterface for SubAgentUiAdapter { 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 { @@ -467,9 +505,13 @@ impl UserInterface for SubAgentUiAdapter { ); if *cancelled { self.set_cancelled(); - } - if let Some(err) = error { + 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; } diff --git a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs index ad968327..be96e7c9 100644 --- a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -1,8 +1,12 @@ -use crate::agent::sub_agent::{SubAgentOutput, SubAgentToolStatus}; +use crate::agent::sub_agent::{SubAgentActivity, SubAgentOutput, SubAgentToolStatus}; use crate::ui::ToolStatus; -use gpui::{div, px, Element, ParentElement, SharedString, Styled}; +use gpui::{ + bounce, div, ease_in_out, percentage, px, svg, Animation, AnimationExt, Element, ParentElement, + SharedString, Styled, Transformation, +}; use std::collections::HashMap; use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; use tracing::warn; /// A unique key for tool name @@ -178,6 +182,65 @@ impl SpawnAgentOutputRenderer { .into_any() } + /// Render the current activity status line with spinning indicator + fn render_activity_line( + activity: SubAgentActivity, + theme: &gpui_component::theme::Theme, + ) -> Option { + let (text, color, show_spinner) = match activity { + SubAgentActivity::WaitingForLlm => { + ("Waiting for response...", theme.muted_foreground, true) + } + SubAgentActivity::Streaming => ("Thinking...", theme.info, true), + SubAgentActivity::ExecutingTools => ("Executing tools...", theme.info, true), + SubAgentActivity::Completed => return None, // Don't show for completed + SubAgentActivity::Cancelled => ("Cancelled", theme.warning, 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(), + ); + + 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, @@ -249,22 +312,39 @@ impl ToolOutputRenderer for SpawnAgentOutputRenderer { } }; - if parsed.tools.is_empty() && parsed.cancelled.is_none() && parsed.error.is_none() { + // Always render if we have activity state, even if no tools yet + let has_activity = parsed.activity.is_some(); + if parsed.tools.is_empty() + && parsed.cancelled.is_none() + && parsed.error.is_none() + && !has_activity + { return None; } + let mut elements: Vec = Vec::new(); + + // Add activity line first (shows current state with spinner) + if let Some(activity) = parsed.activity { + if let Some(activity_line) = Self::render_activity_line(activity, theme) { + elements.push(activity_line); + } + } + // Render compact list of tool calls - let mut elements: Vec = parsed - .tools - .iter() - .map(|tool| Self::render_tool_line(tool, theme)) - .collect(); + for tool in &parsed.tools { + elements.push(Self::render_tool_line(tool, theme)); + } - // Add status line if present + // Add error/cancelled status line if present if let Some(status_line) = Self::render_status_line(&parsed, theme) { elements.push(status_line); } + if elements.is_empty() { + return None; + } + Some( div() .flex() From e801fc737c7a627a02dc5f9ccb7f802fe4ae89d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 14 Dec 2025 23:16:04 +0100 Subject: [PATCH 08/18] WIP: better spawn_agent output while it's running --- crates/code_assistant/src/agent/sub_agent.rs | 81 +++++++++-- crates/code_assistant/src/ui/gpui/elements.rs | 18 +-- .../src/ui/gpui/tool_output_renderers.rs | 126 +++++++++++------- 3 files changed, 161 insertions(+), 64 deletions(-) diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index 6cd8456e..a8a2d521 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -105,7 +105,7 @@ impl DefaultSubAgentRunner { parent_ui: Arc, parent_tool_id: String, cancelled: Arc, - ) -> Arc { + ) -> Arc { Arc::new(SubAgentUiAdapter::new(parent_ui, parent_tool_id, cancelled)) } @@ -196,10 +196,13 @@ impl SubAgentRunner for DefaultSubAgentRunner { 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, + sub_ui as Arc, cancelled.clone(), self.permission_handler.clone(), ) @@ -245,14 +248,19 @@ impl SubAgentRunner for DefaultSubAgentRunner { self.cancellation_registry.unregister(parent_tool_id); - // Ensure the parent tool block shows completion. + // Set the final response in the adapter and send 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(); + + // Update the parent tool block with the complete output (tools + response) let _ = self .ui .send_event(UiEvent::UpdateToolStatus { tool_id: parent_tool_id.to_string(), status: ToolStatus::Success, message: Some("Sub-agent finished".to_string()), - output: None, + output: Some(final_json), }) .await; @@ -336,6 +344,9 @@ pub struct SubAgentOutput { 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 { @@ -345,6 +356,7 @@ impl SubAgentOutput { activity: Some(SubAgentActivity::WaitingForLlm), cancelled: None, error: None, + response: None, } } @@ -390,11 +402,18 @@ impl SubAgentUiAdapter { } async fn send_output_update(&self) { - let json = { + let (json, tool_count, activity) = { let output = self.output.lock().unwrap(); - output.to_json() + (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 { @@ -410,6 +429,16 @@ impl SubAgentUiAdapter { 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 { @@ -418,23 +447,45 @@ impl SubAgentUiAdapter { message: None, }); id_map.insert(id.to_string(), index); + tracing::debug!( + "SubAgentUiAdapter: Added tool {} ({}) at index {}, total tools: {}", + name, + id, + index, + output.tools.len() + ); } fn update_tool_status(&self, tool_id: &str, status: ToolStatus, message: Option) { let mut output = self.output.lock().unwrap(); - let id_map = self.tool_id_to_index.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={}", - tool_id + "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(), + message, + }); + id_map.insert(tool_id.to_string(), index); } } @@ -453,6 +504,18 @@ impl SubAgentUiAdapter { 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] diff --git a/crates/code_assistant/src/ui/gpui/elements.rs b/crates/code_assistant/src/ui/gpui/elements.rs index 1269cd7e..3b38c509 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,15 +1352,20 @@ impl Render for BlockView { block.state == ToolBlockState::Expanded }; + if should_show_output { // 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, output_content, &block.status, - cx.theme(), + &theme, + window, + cx, ) }); diff --git a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs index be96e7c9..d587deda 100644 --- a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -1,9 +1,10 @@ use crate::agent::sub_agent::{SubAgentActivity, SubAgentOutput, SubAgentToolStatus}; use crate::ui::ToolStatus; use gpui::{ - bounce, div, ease_in_out, percentage, px, svg, Animation, AnimationExt, Element, ParentElement, - SharedString, Styled, Transformation, + bounce, div, ease_in_out, percentage, px, svg, Animation, AnimationExt, Context, Element, + ParentElement, SharedString, Styled, Transformation, Window, }; +use gpui_component::text::TextView; use std::collections::HashMap; use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; @@ -25,6 +26,8 @@ pub trait ToolOutputRenderer: Send + Sync { output: &str, status: &ToolStatus, theme: &gpui_component::theme::Theme, + window: &mut Window, + cx: &mut Context, ) -> Option; } @@ -92,10 +95,12 @@ impl ToolOutputRendererRegistry { 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_name, output, status, theme)) + .and_then(|renderer| renderer.render(tool_name, output, status, theme, window, cx)) } } @@ -106,7 +111,7 @@ impl Default for ToolOutputRendererRegistry { } /// Renderer for spawn_agent tool output -/// Displays sub-agent tool calls in a compact, Zed-like style +/// Displays sub-agent tool calls in a compact, Zed-like style with markdown response pub struct SpawnAgentOutputRenderer; impl SpawnAgentOutputRenderer { @@ -116,26 +121,7 @@ impl SpawnAgentOutputRenderer { file_icons::get().get_tool_icon(tool_name) } - /// Get a display title for a tool (like Zed's ACP mode) - fn get_tool_title(tool_name: &str) -> String { - match tool_name { - "read_files" => "Reading".to_string(), - "search_files" => "Searching".to_string(), - "list_files" => "Listing".to_string(), - "glob_files" => "Finding files".to_string(), - "write_file" => "Writing".to_string(), - "edit" => "Editing".to_string(), - "replace_in_file" => "Replacing".to_string(), - "delete_files" => "Deleting".to_string(), - "execute_command" => "Executing".to_string(), - "web_fetch" => "Fetching".to_string(), - "web_search" => "Searching web".to_string(), - "perplexity_ask" => "Asking Perplexity".to_string(), - _ => tool_name.replace('_', " "), - } - } - - /// Render a single compact tool line from structured data + /// Render a single compact tool line showing tool status message (one-liner) fn render_tool_line( tool: &crate::agent::sub_agent::SubAgentToolCall, theme: &gpui_component::theme::Theme, @@ -143,7 +129,6 @@ impl SpawnAgentOutputRenderer { use super::file_icons; let icon = Self::get_tool_icon(&tool.name); - let title = Self::get_tool_title(&tool.name); // Status-based colors let (icon_color, text_color) = match tool.status { @@ -152,16 +137,14 @@ impl SpawnAgentOutputRenderer { SubAgentToolStatus::Error => (theme.danger, theme.danger), }; - // Build the display text - let display_text = if let Some(msg) = &tool.message { - if msg.is_empty() { - title - } else { - format!("{title} — {msg}") - } - } else { - title - }; + // Use the status message if available (this is the tool's status() output) + // Otherwise fall back to just the tool name + let display_text = tool + .message + .as_ref() + .filter(|m| !m.is_empty()) + .cloned() + .unwrap_or_else(|| tool.name.replace('_', " ")); div() .flex() @@ -172,7 +155,7 @@ impl SpawnAgentOutputRenderer { .children(vec![ // Icon file_icons::render_icon_container(&icon, 14.0, icon_color, "🔧").into_any(), - // Title text + // Status text (one-liner from tool) div() .text_size(px(13.)) .text_color(text_color) @@ -188,10 +171,8 @@ impl SpawnAgentOutputRenderer { theme: &gpui_component::theme::Theme, ) -> Option { let (text, color, show_spinner) = match activity { - SubAgentActivity::WaitingForLlm => { - ("Waiting for response...", theme.muted_foreground, true) - } - SubAgentActivity::Streaming => ("Thinking...", theme.info, true), + SubAgentActivity::WaitingForLlm => ("Waiting...", theme.muted_foreground, true), + SubAgentActivity::Streaming => ("Responding...", theme.info, true), SubAgentActivity::ExecutingTools => ("Executing tools...", theme.info, true), SubAgentActivity::Completed => return None, // Don't show for completed SubAgentActivity::Cancelled => ("Cancelled", theme.warning, false), @@ -284,6 +265,28 @@ impl SpawnAgentOutputRenderer { 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 { @@ -297,6 +300,8 @@ impl ToolOutputRenderer for SpawnAgentOutputRenderer { output: &str, _status: &ToolStatus, theme: &gpui_component::theme::Theme, + window: &mut Window, + cx: &mut Context, ) -> Option { if output.is_empty() { return None; @@ -312,35 +317,45 @@ impl ToolOutputRenderer for SpawnAgentOutputRenderer { } }; - // Always render if we have activity state, even if no tools yet + // 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(); - // Add activity line first (shows current state with spinner) + // 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 below tools (shows current state with spinner) + // This is where new output will appear if let Some(activity) = parsed.activity { if let Some(activity_line) = Self::render_activity_line(activity, theme) { elements.push(activity_line); } } - // Render compact list of tool calls - for tool in &parsed.tools { - elements.push(Self::render_tool_line(tool, theme)); - } - // 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; } @@ -425,4 +440,21 @@ mod tests { 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.") + ); + } } From 458249a27b1268322d22621a23bf906bb6799726 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 15 Dec 2025 11:37:55 +0100 Subject: [PATCH 09/18] Distinguish tool output for UI and LLM --- crates/code_assistant/src/agent/mod.rs | 4 +- crates/code_assistant/src/agent/runner.rs | 11 ++--- crates/code_assistant/src/agent/sub_agent.rs | 34 +++++++------- .../src/tests/sub_agent_tests.rs | 45 ++++++++++++++++--- .../code_assistant/src/tools/core/render.rs | 10 ++++- .../src/tools/impls/spawn_agent.rs | 24 ++++++++-- 6 files changed, 96 insertions(+), 32 deletions(-) diff --git a/crates/code_assistant/src/agent/mod.rs b/crates/code_assistant/src/agent/mod.rs index 0a297879..83f2782f 100644 --- a/crates/code_assistant/src/agent/mod.rs +++ b/crates/code_assistant/src/agent/mod.rs @@ -9,5 +9,7 @@ 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 sub_agent::{ + DefaultSubAgentRunner, SubAgentCancellationRegistry, SubAgentResult, SubAgentRunner, +}; pub use types::ToolExecution; diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index bd2b838f..7dccaba9 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -900,14 +900,15 @@ impl Agent { let status_msg = output.status(); let mut resources_tracker = ResourcesTracker::new(); - let output_text = output.render(&mut resources_tracker); + // 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(output_text), + output: Some(ui_output), }) .await; @@ -1714,9 +1715,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 { @@ -1725,7 +1726,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 index a8a2d521..5f7e7886 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -53,8 +53,15 @@ impl AgentStatePersistence for NoOpStatePersistence { } } -/// Runs sub-agents with isolated history and streams a compact progress view into the parent tool UI. -#[async_trait::async_trait] +/// 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 { @@ -64,7 +71,7 @@ pub trait SubAgentRunner: Send + Sync { instructions: String, tool_scope: ToolScope, require_file_references: bool, - ) -> Result; + ) -> Result; } pub struct DefaultSubAgentRunner { @@ -186,7 +193,7 @@ impl SubAgentRunner for DefaultSubAgentRunner { instructions: String, tool_scope: ToolScope, require_file_references: bool, - ) -> Result { + ) -> Result { let cancelled = self .cancellation_registry .register(parent_tool_id.to_string()); @@ -248,23 +255,18 @@ impl SubAgentRunner for DefaultSubAgentRunner { self.cancellation_registry.unregister(parent_tool_id); - // Set the final response in the adapter and send the complete JSON output + // 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(); - // Update the parent tool block with the complete output (tools + response) - let _ = self - .ui - .send_event(UiEvent::UpdateToolStatus { - tool_id: parent_tool_id.to_string(), - status: ToolStatus::Success, - message: Some("Sub-agent finished".to_string()), - output: Some(final_json), - }) - .await; + // Note: We don't send UpdateToolStatus here - the caller (runner.rs) will do that + // using render_for_ui() which returns our JSON output - Ok(last_answer) + Ok(SubAgentResult { + answer: last_answer, + ui_output: final_json, + }) } } diff --git a/crates/code_assistant/src/tests/sub_agent_tests.rs b/crates/code_assistant/src/tests/sub_agent_tests.rs index 39c43e45..b18b29aa 100644 --- a/crates/code_assistant/src/tests/sub_agent_tests.rs +++ b/crates/code_assistant/src/tests/sub_agent_tests.rs @@ -1,6 +1,6 @@ //! Tests for the sub-agent feature (spawn_agent tool). -use crate::agent::sub_agent::SubAgentRunner; +use crate::agent::sub_agent::{SubAgentResult, SubAgentRunner}; use crate::agent::SubAgentCancellationRegistry; use crate::tools::core::ToolScope; use crate::tools::impls::spawn_agent::{SpawnAgentInput, SpawnAgentOutput}; @@ -45,7 +45,7 @@ impl SubAgentRunner for MockSubAgentRunner { _instructions: String, _tool_scope: ToolScope, _require_file_references: bool, - ) -> Result { + ) -> Result { // Track concurrent executions let current = self.concurrent_count.fetch_add(1, Ordering::SeqCst) + 1; @@ -73,7 +73,10 @@ impl SubAgentRunner for MockSubAgentRunner { // Track completion self.concurrent_count.fetch_sub(1, Ordering::SeqCst); - Ok(self.response.clone()) + Ok(SubAgentResult { + answer: self.response.clone(), + ui_output: format!(r#"{{"tools":[],"response":"{}"}}"#, self.response), + }) } } @@ -86,6 +89,7 @@ fn test_spawn_agent_output_render() { answer: "The answer is 42.".to_string(), cancelled: false, error: None, + ui_output: None, }; let mut tracker = ResourcesTracker::new(); @@ -98,6 +102,7 @@ fn test_spawn_agent_output_render() { answer: String::new(), cancelled: true, error: None, + ui_output: None, }; let rendered = output.render(&mut tracker); @@ -109,11 +114,41 @@ fn test_spawn_agent_output_render() { answer: String::new(), cancelled: false, error: Some("Connection failed".to_string()), + ui_output: None, }; let rendered = output.render(&mut tracker); assert_eq!(rendered, "Sub-agent failed: 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(), + cancelled: false, + 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(), + cancelled: false, + error: None, + ui_output: None, + }; + assert_eq!( + output_without_ui.render_for_ui(&mut tracker), + "Just the answer" + ); } #[test] @@ -181,7 +216,7 @@ async fn test_mock_sub_agent_runner() { .await; assert!(result.is_ok()); - assert_eq!(result.unwrap(), "Test response"); + assert_eq!(result.unwrap().answer, "Test response"); assert_eq!(runner.call_count(), 1); } @@ -214,7 +249,7 @@ async fn test_parallel_sub_agent_execution() { assert_eq!(results.len(), 4); for result in results { assert!(result.is_ok()); - assert_eq!(result.unwrap(), "Parallel result"); + assert_eq!(result.unwrap().answer, "Parallel result"); } // Verify they ran in parallel (max concurrent should be > 1) 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/impls/spawn_agent.rs b/crates/code_assistant/src/tools/impls/spawn_agent.rs index e1af17d1..4bdcf437 100644 --- a/crates/code_assistant/src/tools/impls/spawn_agent.rs +++ b/crates/code_assistant/src/tools/impls/spawn_agent.rs @@ -25,12 +25,17 @@ fn default_mode() -> String { /// Output from the spawn_agent tool. #[derive(Serialize, Deserialize)] pub struct SpawnAgentOutput { - /// The final answer from the sub-agent. + /// The final answer from the sub-agent (plain text for LLM context). pub answer: String, /// Whether the sub-agent was cancelled. pub cancelled: bool, /// Error message if the sub-agent failed. 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. + #[serde(skip)] + pub ui_output: Option, } impl Render for SpawnAgentOutput { @@ -44,6 +49,7 @@ impl Render for SpawnAgentOutput { } } + /// Returns plain text for LLM context fn render(&self, _tracker: &mut ResourcesTracker) -> String { if let Some(e) = &self.error { return format!("Sub-agent failed: {e}"); @@ -53,6 +59,14 @@ impl Render for SpawnAgentOutput { } 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 { @@ -154,18 +168,20 @@ impl Tool for SpawnAgentTool { .await; match result { - Ok(answer) => { - let cancelled = answer == "Sub-agent cancelled by user."; + Ok(sub_result) => { + let cancelled = sub_result.answer == "Sub-agent cancelled by user."; Ok(SpawnAgentOutput { - answer, + answer: sub_result.answer, cancelled, error: None, + ui_output: Some(sub_result.ui_output), }) } Err(e) => Ok(SpawnAgentOutput { answer: String::new(), cancelled: false, error: Some(e.to_string()), + ui_output: None, }), } } From 7eefbbe4b2b470c8322cf78ebcd3867d3b300b71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 15 Dec 2025 13:31:41 +0100 Subject: [PATCH 10/18] Fix inconsistencies and fix persisting structured output --- crates/code_assistant/src/agent/sub_agent.rs | 25 ++++++------------- crates/code_assistant/src/session/instance.rs | 6 ++++- .../src/tools/impls/spawn_agent.rs | 4 +-- 3 files changed, 15 insertions(+), 20 deletions(-) diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index 5f7e7886..b90f48ff 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -524,18 +524,6 @@ impl SubAgentUiAdapter { impl UserInterface for SubAgentUiAdapter { async fn send_event(&self, event: UiEvent) -> Result<(), crate::ui::UIError> { match &event { - UiEvent::StartTool { name, id } => { - // This event is typically not sent directly (GPUI creates it from DisplayFragment) - // but handle it for completeness in case other UIs send it - tracing::debug!( - "SubAgentUiAdapter: StartTool event received: {} ({})", - name, - id - ); - self.add_tool_start(name, id); - self.send_output_update().await; - } - UiEvent::UpdateToolStatus { tool_id, status, @@ -596,17 +584,20 @@ impl UserInterface for SubAgentUiAdapter { match fragment { DisplayFragment::ToolName { name, id } => { - // A sub-agent tool is starting - capture it - // Note: This is called during LLM streaming when the tool name is parsed - // The UpdateToolStatus event with Running status comes later when execution starts + // 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); - // We can't send async update here, but the subsequent UpdateToolStatus - // event will trigger send_output_update() } _ => { // Ignore other fragments (text, parameters, etc.) diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 2bbed095..2b98957b 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -398,7 +398,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/tools/impls/spawn_agent.rs b/crates/code_assistant/src/tools/impls/spawn_agent.rs index 4bdcf437..e83abdd3 100644 --- a/crates/code_assistant/src/tools/impls/spawn_agent.rs +++ b/crates/code_assistant/src/tools/impls/spawn_agent.rs @@ -33,8 +33,8 @@ pub struct SpawnAgentOutput { 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. - #[serde(skip)] + /// 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, } From 462418a2748cd560d89f67f09a45f003b9ee2cf7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 15 Dec 2025 13:34:00 +0100 Subject: [PATCH 11/18] Fix extracting LLM response --- crates/code_assistant/src/agent/sub_agent.rs | 23 +++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index b90f48ff..f9ba21f2 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -275,7 +275,28 @@ fn extract_last_assistant_text(messages: &[Message]) -> Option { .iter() .rev() .find(|m| matches!(m.role, llm::MessageRole::Assistant)) - .map(|m| m.to_string()) + .map(|m| extract_text_from_message(m)) +} + +/// 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 { From 489a02fc2c620b3e59808643cd3e50a127690d55 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 15 Dec 2025 14:09:05 +0100 Subject: [PATCH 12/18] Centralize "title" rendering for tools... ... and use for spawn_agent progress display. --- crates/code_assistant/src/acp/ui.rs | 83 +--------- crates/code_assistant/src/agent/sub_agent.rs | 41 ++++- crates/code_assistant/src/tools/core/mod.rs | 2 + crates/code_assistant/src/tools/core/title.rs | 147 ++++++++++++++++++ .../src/tools/impls/spawn_agent.rs | 12 +- crates/code_assistant/src/ui/gpui/mod.rs | 5 + .../src/ui/gpui/spawn_agent_renderer.rs | 53 +++++++ .../src/ui/gpui/tool_output_renderers.rs | 21 ++- 8 files changed, 280 insertions(+), 84 deletions(-) create mode 100644 crates/code_assistant/src/tools/core/title.rs create mode 100644 crates/code_assistant/src/ui/gpui/spawn_agent_renderer.rs diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index d4d2fb7a..3b1e9aad 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,82 +103,16 @@ 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(); - - // Only return the new title if we actually made substitutions - if has_substitution { - Some(result) - } else { - None - } - } - - fn format_parameter_for_title(&self, value: &str) -> String { - const MAX_TITLE_LENGTH: usize = 50; + // 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(); - let trimmed = value.trim(); - if trimmed.is_empty() { - return String::new(); + if let Some(new_title) = crate::tools::core::generate_tool_title(tool_name, ¶ms) { + self.title = Some(new_title); } - - // 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( diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index f9ba21f2..07d9537a 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -316,8 +316,15 @@ fn has_file_references_with_line_ranges(text: &str) -> bool { 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 @@ -467,7 +474,9 @@ impl SubAgentUiAdapter { 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!( @@ -479,6 +488,26 @@ impl SubAgentUiAdapter { ); } + 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(); @@ -506,7 +535,9 @@ impl SubAgentUiAdapter { 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); } @@ -620,8 +651,16 @@ impl UserInterface for SubAgentUiAdapter { ); 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, parameters, etc.) + // Ignore other fragments (text, etc.) // They belong to the sub-agent's isolated transcript } } 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/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/impls/spawn_agent.rs b/crates/code_assistant/src/tools/impls/spawn_agent.rs index e83abdd3..f38bc2d7 100644 --- a/crates/code_assistant/src/tools/impls/spawn_agent.rs +++ b/crates/code_assistant/src/tools/impls/spawn_agent.rs @@ -11,13 +11,21 @@ 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)] + #[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")] + #[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() } diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index f00a6b86..679e5775 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -16,6 +16,7 @@ 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; @@ -28,6 +29,7 @@ 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}; @@ -237,6 +239,9 @@ 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); 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 index d587deda..c0d13d03 100644 --- a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -121,7 +121,7 @@ impl SpawnAgentOutputRenderer { file_icons::get().get_tool_icon(tool_name) } - /// Render a single compact tool line showing tool status message (one-liner) + /// 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, @@ -137,13 +137,14 @@ impl SpawnAgentOutputRenderer { SubAgentToolStatus::Error => (theme.danger, theme.danger), }; - // Use the status message if available (this is the tool's status() output) - // Otherwise fall back to just the tool name + // Use title (generated from template) if available, + // otherwise fall back to status message, then tool name let display_text = tool - .message + .title .as_ref() - .filter(|m| !m.is_empty()) + .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() @@ -155,7 +156,7 @@ impl SpawnAgentOutputRenderer { .children(vec![ // Icon file_icons::render_icon_container(&icon, 14.0, icon_color, "🔧").into_any(), - // Status text (one-liner from tool) + // Title text (one-liner from template or fallback) div() .text_size(px(13.)) .text_color(text_color) @@ -414,14 +415,18 @@ mod tests { .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(); @@ -433,6 +438,10 @@ mod tests { parsed.tools[1].message.as_deref(), Some("Searching for pattern") ); + assert_eq!( + parsed.tools[1].title.as_deref(), + Some("Searching for 'pattern'") + ); } #[test] From d0bd716aeaaf31e6e55bb59151c576526cf1cb7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 17 Dec 2025 18:41:03 +0100 Subject: [PATCH 13/18] Implement sub-agent canceling. --- crates/code_assistant/src/acp/ui.rs | 3 +- crates/code_assistant/src/agent/mod.rs | 4 +- crates/code_assistant/src/agent/sub_agent.rs | 13 +--- crates/code_assistant/src/session/instance.rs | 11 +++ crates/code_assistant/src/session/manager.rs | 21 +++++- crates/code_assistant/src/tools/core/tool.rs | 5 +- crates/code_assistant/src/ui/backend.rs | 65 +++++++++++++++++ crates/code_assistant/src/ui/gpui/elements.rs | 1 + crates/code_assistant/src/ui/gpui/mod.rs | 26 +++++++ .../src/ui/gpui/tool_output_renderers.rs | 70 +++++++++++++++---- crates/code_assistant/src/ui/terminal/app.rs | 8 +++ crates/code_assistant/src/ui/ui_events.rs | 2 + docs/sub-agents-feature.md | 14 +++- 13 files changed, 206 insertions(+), 37 deletions(-) diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index 3b1e9aad..92dad0e5 100644 --- a/crates/code_assistant/src/acp/ui.rs +++ b/crates/code_assistant/src/acp/ui.rs @@ -648,7 +648,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 83f2782f..0a297879 100644 --- a/crates/code_assistant/src/agent/mod.rs +++ b/crates/code_assistant/src/agent/mod.rs @@ -9,7 +9,5 @@ pub mod types; pub use crate::types::ToolSyntax; // pub use persistence::FileStatePersistence; pub use runner::{Agent, AgentComponents}; -pub use sub_agent::{ - DefaultSubAgentRunner, SubAgentCancellationRegistry, SubAgentResult, SubAgentRunner, -}; +pub use sub_agent::{DefaultSubAgentRunner, SubAgentCancellationRegistry, SubAgentRunner}; pub use types::ToolExecution; diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index 07d9537a..28dabc2b 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -12,7 +12,6 @@ use llm::Message; use sandbox::{SandboxContext, SandboxPolicy}; use std::collections::HashMap; use std::sync::{atomic::AtomicBool, atomic::Ordering, Arc, Mutex}; -use std::time::SystemTime; /// Cancellation registry keyed by the parent `spawn_agent` tool id. #[derive(Default)] @@ -275,7 +274,7 @@ fn extract_last_assistant_text(messages: &[Message]) -> Option { .iter() .rev() .find(|m| matches!(m.role, llm::MessageRole::Assistant)) - .map(|m| extract_text_from_message(m)) + .map(extract_text_from_message) } /// Extract just the text content from a message, ignoring tool calls and other blocks @@ -680,13 +679,3 @@ impl UserInterface for SubAgentUiAdapter { self } } - -/// Helper for tool implementations to mark a tool id as cancelled. -pub fn cancel_sub_agent(registry: &SubAgentCancellationRegistry, tool_id: &str) -> bool { - registry.cancel(tool_id) -} - -/// Helper for tool implementations to create an initial visible tool result block. -pub fn sub_agent_tool_result_timestamps() -> (Option, Option) { - (Some(SystemTime::now()), None) -} diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 2b98957b..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() diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index d1d9951e..3fd5d810 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -318,8 +318,14 @@ impl SessionManager { sandbox_context_clone.clone(), )); - // Create sub-agent runner and cancellation registry - let sub_agent_cancellation_registry = Arc::new(SubAgentCancellationRegistry::default()); + // 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() @@ -456,6 +462,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/tools/core/tool.rs b/crates/code_assistant/src/tools/core/tool.rs index 101873cc..89378eea 100644 --- a/crates/code_assistant/src/tools/core/tool.rs +++ b/crates/code_assistant/src/tools/core/tool.rs @@ -21,9 +21,12 @@ 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>, - /// Optional cancellation registry used for sub-agent cancellation. + /// Optional cancellation registry for sub-agent cancellation. + /// Reserved for future tool implementations that may need to cancel sub-agents. + #[allow(dead_code)] pub sub_agent_cancellation_registry: Option<&'a crate::agent::SubAgentCancellationRegistry>, } 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 3b38c509..7cea2bd1 100644 --- a/crates/code_assistant/src/ui/gpui/elements.rs +++ b/crates/code_assistant/src/ui/gpui/elements.rs @@ -1361,6 +1361,7 @@ impl Render for BlockView { .and_then(|registry| { registry.render_output( &block.name, + &block.id, output_content, &block.status, &theme, diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index 679e5775..85add068 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -1013,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"); + } + } } } @@ -1335,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()) { @@ -1350,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/tool_output_renderers.rs b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs index c0d13d03..270029f2 100644 --- a/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs +++ b/crates/code_assistant/src/ui/gpui/tool_output_renderers.rs @@ -1,8 +1,9 @@ use crate::agent::sub_agent::{SubAgentActivity, SubAgentOutput, SubAgentToolStatus}; use crate::ui::ToolStatus; use gpui::{ - bounce, div, ease_in_out, percentage, px, svg, Animation, AnimationExt, Context, Element, - ParentElement, SharedString, Styled, Transformation, Window, + 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; @@ -22,7 +23,7 @@ pub trait ToolOutputRenderer: Send + Sync { /// Returns None if the default rendering should be used fn render( &self, - tool_name: &str, + tool_id: &str, output: &str, status: &ToolStatus, theme: &gpui_component::theme::Theme, @@ -89,9 +90,11 @@ impl ToolOutputRendererRegistry { /// 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, @@ -100,7 +103,7 @@ impl ToolOutputRendererRegistry { ) -> Option { self.renderers .get(tool_name) - .and_then(|renderer| renderer.render(tool_name, output, status, theme, window, cx)) + .and_then(|renderer| renderer.render(tool_id, output, status, theme, window, cx)) } } @@ -166,17 +169,19 @@ impl SpawnAgentOutputRenderer { .into_any() } - /// Render the current activity status line with spinning indicator - fn render_activity_line( + /// 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) = match activity { - SubAgentActivity::WaitingForLlm => ("Waiting...", theme.muted_foreground, true), - SubAgentActivity::Streaming => ("Responding...", theme.info, true), - SubAgentActivity::ExecutingTools => ("Executing tools...", theme.info, true), + 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), + SubAgentActivity::Cancelled => ("Cancelled", theme.warning, false, false), SubAgentActivity::Failed => return None, // Error shown separately }; @@ -211,6 +216,39 @@ impl SpawnAgentOutputRenderer { .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() @@ -297,7 +335,7 @@ impl ToolOutputRenderer for SpawnAgentOutputRenderer { fn render( &self, - _tool_name: &str, + tool_id: &str, output: &str, _status: &ToolStatus, theme: &gpui_component::theme::Theme, @@ -337,11 +375,13 @@ impl ToolOutputRenderer for SpawnAgentOutputRenderer { elements.push(Self::render_tool_line(tool, theme)); } - // Add activity line below tools (shows current state with spinner) + // 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_line) = Self::render_activity_line(activity, theme) { - elements.push(activity_line); + if let Some(activity_element) = + Self::render_activity_with_cancel(activity, tool_id, theme, cx) + { + elements.push(activity_element); } } 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/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 index 0d2b981f..0c84491d 100644 --- a/docs/sub-agents-feature.md +++ b/docs/sub-agents-feature.md @@ -56,13 +56,21 @@ - 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 + ### Pending -- [ ] **UI integration for cancellation**: Expose cancel button per running `spawn_agent` block - [ ] **Permission attribution**: Show permission requests as originating from sub-agent context (inline or popover) ### Notes -- The cancellation infrastructure is in place (`SubAgentCancellationRegistry`, `cancel` method, `cancel_sub_agent` helper) but the UI hooks to trigger cancellation are not yet implemented. -- The `sub_agent_cancellation_registry` field in `ToolContext` is available for future use when implementing tool-level cancellation from UI. +- The `sub_agent_cancellation_registry` field in `ToolContext` is reserved for future tool implementations that may need to cancel sub-agents programmatically. --- From f92820c45388552f8e7fa658782a94972303ae3a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 17 Dec 2025 19:09:35 +0100 Subject: [PATCH 14/18] Fix cancelling --- crates/code_assistant/src/agent/sub_agent.rs | 21 ++++++++++--- .../src/tests/sub_agent_tests.rs | 15 ++++------ .../src/tools/impls/spawn_agent.rs | 30 ++++++------------- 3 files changed, 31 insertions(+), 35 deletions(-) diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index 28dabc2b..be50697f 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -220,14 +220,24 @@ impl SubAgentRunner for DefaultSubAgentRunner { // 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) { - last_answer = "Sub-agent cancelled by user.".to_string(); + 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 { @@ -254,14 +264,17 @@ impl SubAgentRunner for DefaultSubAgentRunner { 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(); - // Note: We don't send UpdateToolStatus here - the caller (runner.rs) will do that - // using render_for_ui() which returns our JSON output - Ok(SubAgentResult { answer: last_answer, ui_output: final_json, diff --git a/crates/code_assistant/src/tests/sub_agent_tests.rs b/crates/code_assistant/src/tests/sub_agent_tests.rs index b18b29aa..89629917 100644 --- a/crates/code_assistant/src/tests/sub_agent_tests.rs +++ b/crates/code_assistant/src/tests/sub_agent_tests.rs @@ -87,7 +87,6 @@ fn test_spawn_agent_output_render() { // Test successful output let output = SpawnAgentOutput { answer: "The answer is 42.".to_string(), - cancelled: false, error: None, ui_output: None, }; @@ -97,34 +96,31 @@ fn test_spawn_agent_output_render() { assert_eq!(rendered, "The answer is 42."); assert_eq!(output.status(), "Sub-agent completed"); - // Test cancelled output + // Test cancelled output (cancellation is now just an error) let output = SpawnAgentOutput { answer: String::new(), - cancelled: true, - error: None, + error: Some("Cancelled by user".to_string()), ui_output: None, }; let rendered = output.render(&mut tracker); - assert_eq!(rendered, "Sub-agent cancelled by user."); - assert_eq!(output.status(), "Sub-agent cancelled by user"); + 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(), - cancelled: false, error: Some("Connection failed".to_string()), ui_output: None, }; let rendered = output.render(&mut tracker); - assert_eq!(rendered, "Sub-agent failed: Connection failed"); + 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(), - cancelled: false, error: None, ui_output: Some(r#"{"tools":[],"response":"Markdown response"}"#.to_string()), }; @@ -141,7 +137,6 @@ fn test_spawn_agent_output_render() { // Test render_for_ui falls back to render() when ui_output is None let output_without_ui = SpawnAgentOutput { answer: "Just the answer".to_string(), - cancelled: false, error: None, ui_output: None, }; diff --git a/crates/code_assistant/src/tools/impls/spawn_agent.rs b/crates/code_assistant/src/tools/impls/spawn_agent.rs index f38bc2d7..6fbd7edc 100644 --- a/crates/code_assistant/src/tools/impls/spawn_agent.rs +++ b/crates/code_assistant/src/tools/impls/spawn_agent.rs @@ -35,9 +35,7 @@ fn default_mode() -> String { pub struct SpawnAgentOutput { /// The final answer from the sub-agent (plain text for LLM context). pub answer: String, - /// Whether the sub-agent was cancelled. - pub cancelled: bool, - /// Error message if the sub-agent failed. + /// 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 @@ -49,9 +47,7 @@ pub struct SpawnAgentOutput { impl Render for SpawnAgentOutput { fn status(&self) -> String { if let Some(e) = &self.error { - format!("Sub-agent failed: {e}") - } else if self.cancelled { - "Sub-agent cancelled by user".to_string() + format!("Sub-agent error: {e}") } else { "Sub-agent completed".to_string() } @@ -60,10 +56,7 @@ impl Render for SpawnAgentOutput { /// Returns plain text for LLM context fn render(&self, _tracker: &mut ResourcesTracker) -> String { if let Some(e) = &self.error { - return format!("Sub-agent failed: {e}"); - } - if self.cancelled { - return "Sub-agent cancelled by user.".to_string(); + return format!("Sub-agent error: {e}"); } self.answer.clone() } @@ -79,7 +72,7 @@ impl Render for SpawnAgentOutput { impl ToolResult for SpawnAgentOutput { fn is_success(&self) -> bool { - self.error.is_none() && !self.cancelled + self.error.is_none() } } @@ -176,18 +169,13 @@ impl Tool for SpawnAgentTool { .await; match result { - Ok(sub_result) => { - let cancelled = sub_result.answer == "Sub-agent cancelled by user."; - Ok(SpawnAgentOutput { - answer: sub_result.answer, - cancelled, - error: None, - ui_output: Some(sub_result.ui_output), - }) - } + 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(), - cancelled: false, error: Some(e.to_string()), ui_output: None, }), From 2b0b999dda704c55db4f278f89b6d1e8bed962af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 17 Dec 2025 19:20:48 +0100 Subject: [PATCH 15/18] More cleanup in sub-agent cancellation --- crates/code_assistant/src/agent/runner.rs | 7 ------- crates/code_assistant/src/agent/sub_agent.rs | 5 ----- docs/sub-agents-feature.md | 2 +- 3 files changed, 1 insertion(+), 13 deletions(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 7dccaba9..ecb1acfe 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -203,13 +203,6 @@ impl Agent { self.session_name = session_name; } - /// Set an external cancellation flag that can interrupt streaming - pub fn set_external_cancel_flag(&mut self, _flag: Arc) { - // Note: The LLM streaming checks the UI's should_streaming_continue() method. - // For sub-agents, the SubAgentUiAdapter already checks the cancellation flag. - // This method is a placeholder for future direct cancellation integration. - } - /// Get a reference to the message history pub fn message_history(&self) -> &[Message] { &self.message_history diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index be50697f..d8586ada 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -119,7 +119,6 @@ impl DefaultSubAgentRunner { &self, parent_tool_id: &str, ui: Arc, - cancelled: Arc, permission_handler: Option>, ) -> Result { // Create a fresh LLM provider (avoid requiring Clone). @@ -172,9 +171,6 @@ impl DefaultSubAgentRunner { // Initialize project trees, etc. agent.init_project_context()?; - // Ensure sub-agent cancellation can interrupt streaming. - agent.set_external_cancel_flag(cancelled); - Ok(agent) } } @@ -209,7 +205,6 @@ impl SubAgentRunner for DefaultSubAgentRunner { .build_agent( parent_tool_id, sub_ui as Arc, - cancelled.clone(), self.permission_handler.clone(), ) .await?; diff --git a/docs/sub-agents-feature.md b/docs/sub-agents-feature.md index 0c84491d..a5eb6b29 100644 --- a/docs/sub-agents-feature.md +++ b/docs/sub-agents-feature.md @@ -11,7 +11,7 @@ - [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`, `set_external_cancel_flag`, `message_history` +- [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) From b6410cb4c06b99fd06363f40ec1eb9bd3bf50c2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 17 Dec 2025 19:42:15 +0100 Subject: [PATCH 16/18] Remove sub agent cancellation registry from ToolContext --- crates/code_assistant/src/agent/runner.rs | 10 +++------- crates/code_assistant/src/agent/sub_agent.rs | 1 - crates/code_assistant/src/agent/tests.rs | 12 ------------ crates/code_assistant/src/mcp/handler.rs | 1 - crates/code_assistant/src/session/manager.rs | 1 - .../code_assistant/src/tests/format_on_save_tests.rs | 7 ------- crates/code_assistant/src/tests/mocks.rs | 2 -- crates/code_assistant/src/tools/core/tool.rs | 4 ---- 8 files changed, 3 insertions(+), 35 deletions(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index ecb1acfe..3a84c9c1 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -30,10 +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>, - /// Shared cancellation registry so the UI/back-end can cancel running sub-agents by tool id. - pub sub_agent_cancellation_registry: Option>, } use super::ToolSyntax; @@ -57,9 +56,9 @@ pub struct Agent { command_executor: Box, ui: Arc, state_persistence: Box, + permission_handler: Option>, sub_agent_runner: Option>, - sub_agent_cancellation_registry: Option>, // Store all messages exchanged message_history: Vec, // Store the history of tool executions @@ -118,7 +117,6 @@ impl Agent { state_persistence, permission_handler, sub_agent_runner, - sub_agent_cancellation_registry, } = components; let mut this = Self { @@ -132,7 +130,6 @@ impl Agent { state_persistence, permission_handler, sub_agent_runner, - sub_agent_cancellation_registry, message_history: Vec::new(), tool_executions: Vec::new(), cached_system_prompts: HashMap::new(), @@ -878,7 +875,6 @@ impl Agent { tool_id: Some(tool_id.to_string()), permission_handler: None, // Will be handled by sub-agent runner sub_agent_runner, - sub_agent_cancellation_registry: None, }; let tool = SpawnAgentTool; @@ -1683,9 +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(), - sub_agent_cancellation_registry: self.sub_agent_cancellation_registry.as_deref(), }; // Execute the tool - could fail with ParseError or other errors diff --git a/crates/code_assistant/src/agent/sub_agent.rs b/crates/code_assistant/src/agent/sub_agent.rs index d8586ada..4a0a1ed1 100644 --- a/crates/code_assistant/src/agent/sub_agent.rs +++ b/crates/code_assistant/src/agent/sub_agent.rs @@ -154,7 +154,6 @@ impl DefaultSubAgentRunner { state_persistence: Box::new(NoOpStatePersistence), permission_handler, sub_agent_runner: None, - sub_agent_cancellation_registry: Some(self.cancellation_registry.clone()), }; let mut agent = Agent::new(components, self.session_config.clone()); diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index 94e8fcaf..f04b167b 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -396,7 +396,6 @@ async fn test_unknown_tool_error_handling() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -519,7 +518,6 @@ async fn test_invalid_xml_tool_error_handling() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -645,7 +643,6 @@ async fn test_parse_error_handling() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -785,7 +782,6 @@ async fn test_write_file_outside_root_error_masks_paths() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -887,7 +883,6 @@ async fn test_context_compaction_inserts_summary() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1006,7 +1001,6 @@ async fn test_compaction_prompt_not_persisted_in_history() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1134,7 +1128,6 @@ async fn test_context_compaction_uses_only_messages_after_previous_summary() -> state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1611,7 +1604,6 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { state_persistence, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -1961,7 +1953,6 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -2017,7 +2008,6 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -2087,7 +2077,6 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { @@ -2141,7 +2130,6 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() state_persistence: Box::new(MockStatePersistence::new()), permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let session_config = SessionConfig { diff --git a/crates/code_assistant/src/mcp/handler.rs b/crates/code_assistant/src/mcp/handler.rs index ccbb338e..78cab502 100644 --- a/crates/code_assistant/src/mcp/handler.rs +++ b/crates/code_assistant/src/mcp/handler.rs @@ -270,7 +270,6 @@ impl MessageHandler { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; // Invoke the tool diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index 3fd5d810..d079313c 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -349,7 +349,6 @@ impl SessionManager { state_persistence: state_storage, permission_handler, sub_agent_runner: Some(sub_agent_runner), - sub_agent_cancellation_registry: Some(sub_agent_cancellation_registry), }; let mut agent = Agent::new(components, session_config.clone()); 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 7c2e49e9..9d568592 100644 --- a/crates/code_assistant/src/tests/format_on_save_tests.rs +++ b/crates/code_assistant/src/tests/format_on_save_tests.rs @@ -199,7 +199,6 @@ async fn test_edit_tool_parameter_update_after_formatting() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; // Test editing: search is formatted (matches file), replacement is unformatted @@ -274,7 +273,6 @@ async fn test_write_file_with_format_on_save() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; // Test writing a Rust file @@ -348,7 +346,6 @@ async fn test_replace_in_file_with_format_on_save() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; // Diff has two SEARCH/REPLACE blocks; replacements are unformatted (missing spaces around '=') @@ -424,7 +421,6 @@ async fn test_no_format_when_pattern_doesnt_match() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; // Test editing a .txt file (should not be formatted) @@ -502,7 +498,6 @@ async fn test_format_on_save_multiple_patterns() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; // Test editing JS file @@ -600,7 +595,6 @@ async fn test_format_on_save_glob_patterns() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, }; let tool = EditTool; @@ -691,7 +685,6 @@ async fn test_format_on_save_with_conflicting_matches() -> Result<()> { tool_id: None, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: 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 0231a50f..0eaa8c08 100644 --- a/crates/code_assistant/src/tests/mocks.rs +++ b/crates/code_assistant/src/tests/mocks.rs @@ -232,7 +232,6 @@ pub fn create_test_tool_context<'a>( tool_id, permission_handler: None, sub_agent_runner: None, - sub_agent_cancellation_registry: None, } } @@ -1131,7 +1130,6 @@ impl ToolTestFixture { tool_id: self.tool_id.clone(), permission_handler: self.permission_handler.as_deref(), sub_agent_runner: None, - sub_agent_cancellation_registry: None, } } diff --git a/crates/code_assistant/src/tools/core/tool.rs b/crates/code_assistant/src/tools/core/tool.rs index 89378eea..4b00994d 100644 --- a/crates/code_assistant/src/tools/core/tool.rs +++ b/crates/code_assistant/src/tools/core/tool.rs @@ -24,10 +24,6 @@ pub struct ToolContext<'a> { /// Optional sub-agent runner used by the `spawn_agent` tool. pub sub_agent_runner: Option<&'a dyn crate::agent::SubAgentRunner>, - /// Optional cancellation registry for sub-agent cancellation. - /// Reserved for future tool implementations that may need to cancel sub-agents. - #[allow(dead_code)] - pub sub_agent_cancellation_registry: Option<&'a crate::agent::SubAgentCancellationRegistry>, } /// Core trait for tools, defining the execution interface From 93ddb54fe01e82e9d00be5dd6c469abb6b056d88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 17 Dec 2025 19:53:34 +0100 Subject: [PATCH 17/18] Use tool title rendering also in Terminal UI --- .../src/ui/terminal/tool_widget.rs | 45 ++++++++----------- 1 file changed, 19 insertions(+), 26 deletions(-) diff --git a/crates/code_assistant/src/ui/terminal/tool_widget.rs b/crates/code_assistant/src/ui/terminal/tool_widget.rs index cbf73ac3..b5127ad5 100644 --- a/crates/code_assistant/src/ui/terminal/tool_widget.rs +++ b/crates/code_assistant/src/ui/terminal/tool_widget.rs @@ -255,35 +255,28 @@ impl<'a> Widget for ToolWidget<'a> { crate::agent::sub_agent::SubAgentToolStatus::Error => ("●", Color::Red), }; - let tool_title = match tool.name.as_str() { - "read_files" => "Reading", - "search_files" => "Searching", - "list_files" => "Listing", - "glob_files" => "Finding files", - "write_file" => "Writing", - "edit" => "Editing", - "execute_command" => "Executing", - "web_fetch" => "Fetching", - "web_search" => "Searching web", - _ => &tool.name, - }; - - let display_text = if let Some(msg) = &tool.message { - format!("{status_symbol} {tool_title} — {msg}") + // 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 { - format!("{status_symbol} {tool_title}") + full_text }; - let truncated = - if display_text.len() > (area.width.saturating_sub(4)) as usize { - format!( - "{}...", - &display_text[..(area.width.saturating_sub(7)) as usize] - ) - } else { - display_text - }; - buf.set_string( area.x + 2, current_y, From 23f5aaa3af61dadcfdb17ecd93d70868a4a29324 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Wed, 17 Dec 2025 21:04:54 +0100 Subject: [PATCH 18/18] Fix ACP spawn_agent output rendering --- crates/code_assistant/src/acp/ui.rs | 62 ++++++++++++++++++++++++++++- docs/sub-agents-feature.md | 7 ++++ 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index 92dad0e5..78f2a400 100644 --- a/crates/code_assistant/src/acp/ui.rs +++ b/crates/code_assistant/src/acp/ui.rs @@ -235,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)); + } } } } @@ -343,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() { diff --git a/docs/sub-agents-feature.md b/docs/sub-agents-feature.md index a5eb6b29..8a778cde 100644 --- a/docs/sub-agents-feature.md +++ b/docs/sub-agents-feature.md @@ -66,6 +66,13 @@ - 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)