agent_engine/tools/subagent/

oneshot.rs

use super::super::{resolve_agent_prompt, Tool, ToolContext, NEXT_SUBAGENT_ID};
pub use crate::runtime::subagent::SubagentResult;
use crate::{AgentEvent, LlmEvent, Result, RuntimeError, SessionEvent};
use serde_json::{json, Value};
use std::sync::atomic::Ordering;
use std::time::Duration;

pub struct SubagentTool;

#[async_trait::async_trait]
impl Tool for SubagentTool {
    fn origin(&self) -> crate::tools::ToolOrigin {
        crate::tools::ToolOrigin::Builtin
    }

    fn name(&self) -> &str {
        "subagent"
    }

    fn description(&self) -> &str {
        "Dispatch a one-shot subagent with a specific system prompt to perform a task. The subagent gets its own tool suite (bash, read, write, edit, grep, find, ls) and runs autonomously until done. Use this when you need the result before continuing. Blocks until done. For parallel work, use subagent_start instead. Provide either an agent name (resolves from ~/.synaps-cli/agents/<name>.md) or a system_prompt string directly."
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "agent": {
                    "type": "string",
                    "description": "Agent name — resolves to ~/.synaps-cli/agents/<name>.md. Mutually exclusive with system_prompt."
                },
                "system_prompt": {
                    "type": "string",
                    "description": "Inline system prompt for the subagent. Use when you don't have a named agent file."
                },
                "task": {
                    "type": "string",
                    "description": "The task/prompt to send to the subagent."
                },
                "model": {"type": "string", "description": "Omit to inherit the session foreground qualified identity. Explicit values must be one of subagent_models' listed exact choices."},
                "role": {"type": "string", "enum": ["planner", "implementer", "tester", "reviewer", "researcher", "debugger"]},
                "write_policy": {
                    "oneOf": [
                        {"type": "object", "properties": {"mode": {"const": "read_only"}}, "required": ["mode"]},
                        {"type": "object", "properties": {"mode": {"const": "isolated_worktree"}}, "required": ["mode"]},
                        {"type": "object", "properties": {"mode": {"const": "non_overlapping_paths"}, "scopes": {"type": "array", "items": {"type": "string"}}}, "required": ["mode", "scopes"]}
                    ]
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds (default: 300). Increase for long-running tasks."
                }
            },
            "required": ["task"]
        })
    }

    async fn execute(&self, params: Value, ctx: ToolContext) -> Result<String> {
        let task = params["task"]
            .as_str()
            .ok_or_else(|| RuntimeError::Tool("Missing 'task' parameter".to_string()))?
            .to_string();

        // Treat blank / whitespace / control-char strings as absent — see
        // subagent_start.rs for full rationale. Models occasionally pass `agent: ""`
        // (or "\u{0}", or " ") alongside a real `system_prompt`; without this filter
        // we'd try to resolve an empty agent name and fail, instead of falling
        // through to the inline prompt.
        let is_blank = |s: &String| s.chars().all(|c| c.is_whitespace() || c.is_control());
        let agent_name = params["agent"]
            .as_str()
            .map(|s| s.to_string())
            .filter(|s| !is_blank(s));
        let inline_prompt = params["system_prompt"]
            .as_str()
            .map(|s| s.to_string())
            .filter(|s| !is_blank(s));
        let requested_model = params["model"].as_str();
        let system_prompt = match (&agent_name, &inline_prompt) {
            (Some(name), _) => resolve_agent_prompt(name).map_err(RuntimeError::Tool)?,
            (None, Some(prompt)) => prompt.clone(),
            (None, None) => {
                return Err(RuntimeError::Tool(
                    "Must provide either 'agent' (name) or 'system_prompt' (inline). Got neither."
                        .to_string(),
                ));
            }
        };
        let subagent_id = NEXT_SUBAGENT_ID.fetch_add(1, Ordering::Relaxed);
        let orchestration_id = format!("sa_{}", subagent_id);
        let decision = ctx
            .capabilities
            .orchestration
            .as_ref()
            .ok_or_else(|| RuntimeError::Tool("delegation policy unavailable".into()))?
            .resolve_and_authorize(&orchestration_id, requested_model)
            .map_err(|error| RuntimeError::Tool(error.to_string()))?;
        let model = decision.model.as_str().to_owned();
        let codex_parent_plan = ctx.capabilities.codex_parent_plan.clone();
        let memory_backend = ctx.capabilities.memory_backend.clone();
        let timeout_secs = params["timeout"]
            .as_u64()
            .unwrap_or(ctx.limits.subagent_timeout);

        let label = agent_name.as_deref().unwrap_or("inline").to_string();
        let task_preview: String = task.chars().take(80).collect();

        tracing::info!(
            "Dispatching subagent '{}' (id={}) with model {}",
            label,
            subagent_id,
            model
        );

        if let Some(ref tx) = ctx.channels.tx_events {
            let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentStart {
                subagent_id,
                agent_name: label.clone(),
                task_preview: task_preview.clone(),
            }));
        }

        let start_time = std::time::Instant::now();

        let (result_tx, result_rx) =
            tokio::sync::oneshot::channel::<std::result::Result<SubagentResult, String>>();
        let label_inner = label.clone();
        let model_inner = model.clone();
        let tx_events_inner = ctx.channels.tx_events.clone();

        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
        let orchestration = ctx.capabilities.orchestration.as_ref().unwrap();
        if let Err(error) = orchestration.mark_starting(&orchestration_id) {
            orchestration.rollback(&orchestration_id);
            return Err(RuntimeError::Tool(error));
        }

        let thread_handle = match std::thread::Builder::new().spawn(move || {
            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                let rt = match tokio::runtime::Builder::new_current_thread()
                    .enable_all()
                    .build()
                {
                    Ok(rt) => rt,
                    Err(_) => {
                        let _ = result_tx.send(Err("runtime initialization failed".into()));
                        return;
                    }
                };

                let result = rt.block_on(async move {
                    use futures::StreamExt;

                    // Host-built worker (shared client/creds/token cache, cached
                    // registry) or the legacy fresh runtime — see `spawn_runtime`.
                    let mut runtime = match super::spawn_runtime().await {
                        Ok(r) => r,
                        Err(_) => return Err("subagent runtime initialization failed".into()),
                    };

                    // Apply subagent spawn policy: worker role, 5m cache TTL, worker
                    // turn budget. Subagents are short-lived one-shots — paying the 1h
                    // write premium (~2× input price) on them is unrecoverable waste
                    // (~$0.23 per 10-spawn fan-out). (#110)
                    super::apply_subagent_runtime_policy(&mut runtime, &crate::config::load_config(), memory_backend.as_ref());
                    runtime.set_system_prompt(super::compose_system_prompt(system_prompt, runtime.memory_backend_is_axel()));
                    runtime.set_model(model);
                    super::apply_codex_worker_reasoning(&mut runtime, codex_parent_plan.as_ref());

                    let cancel = crate::CancellationToken::new();
                    let cancel_inner = cancel.clone();

                    tokio::spawn(async move {
                        let _ = shutdown_rx.await;
                        cancel_inner.cancel();
                    });

                    let mut stream = runtime.run_stream(task, cancel).await;

                let mut final_text = String::new();
                let mut tool_count = 0u32;
                let mut tool_log: Vec<String> = Vec::new();
                let mut response_baseline: (usize, usize, u32) = (0, 0, 0);
                let mut total_input_tokens = 0u64;
                let mut total_output_tokens = 0u64;
                let mut total_cache_read = 0u64;
                let mut total_cache_creation = 0u64;
                // TTL split: None only if no turn ever reported one; otherwise summed.
                let mut total_cache_5m: Option<u64> = None;
                let mut total_cache_1h: Option<u64> = None;

                let timeout_fut = tokio::time::sleep(Duration::from_secs(timeout_secs));
                tokio::pin!(timeout_fut);

                loop {
                    tokio::select! {
                        event = stream.next() => {
                            let Some(event) = event else { break };
                            match event {
                                crate::StreamEvent::Llm(LlmEvent::Thinking(_)) => {
                                    if let Some(ref tx) = tx_events_inner {
                                        let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentUpdate {
                                            subagent_id,
                                            agent_name: label_inner.clone(),
                                            status: "💭 thinking...".to_string(),
                                        }));
                                    }
                                }
                                crate::StreamEvent::Llm(LlmEvent::ResponseStart) => {
                                    response_baseline = (final_text.len(), tool_log.len(), tool_count);
                                }
                                crate::StreamEvent::Llm(LlmEvent::ResponseReset) => {
                                    final_text.truncate(response_baseline.0);
                                    tool_log.truncate(response_baseline.1);
                                    tool_count = response_baseline.2;
                                }
                                crate::StreamEvent::Llm(LlmEvent::Text(text)) => {
                                    final_text.push_str(&text);
                                }
                                crate::StreamEvent::Llm(LlmEvent::ToolUseStart { tool_name: name, .. }) => {
                                    tool_count += 1;
                                    if let Some(ref tx) = tx_events_inner {
                                        let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentUpdate {
                                            subagent_id,
                                            agent_name: label_inner.clone(),
                                            status: format!("⚙ {} (tool #{})", name, tool_count),
                                        }));
                                    }
                                }
                                crate::StreamEvent::Llm(LlmEvent::ToolUse { tool_name, input, .. }) => {
                                    let input_str = input.to_string();
                                    let input_preview: String = input_str.chars().take(200).collect();
                                    tool_log.push(format!("[tool_use]: {} — {}", tool_name, input_preview));
                                    // Build a rich status from the tool input
                                    let detail = match tool_name.as_str() {
                                        "bash" => {
                                            let cmd = input["command"].as_str().unwrap_or("");
                                            let preview: String = cmd.chars().take(60).collect();
                                            format!("$ {}", preview)
                                        }
                                        "read" => {
                                            let path = input["path"].as_str().unwrap_or("?");
                                            let short = path.rsplit('/').next().unwrap_or(path);
                                            format!("reading {}", short)
                                        }
                                        "write" => {
                                            let path = input["path"].as_str().unwrap_or("?");
                                            let short = path.rsplit('/').next().unwrap_or(path);
                                            format!("writing {}", short)
                                        }
                                        "edit" => {
                                            let path = input["path"].as_str().unwrap_or("?");
                                            let short = path.rsplit('/').next().unwrap_or(path);
                                            format!("editing {}", short)
                                        }
                                        "grep" => {
                                            let pat = input["pattern"].as_str().unwrap_or("?");
                                            let preview: String = pat.chars().take(30).collect();
                                            format!("grep /{}/", preview)
                                        }
                                        "find" => {
                                            let pat = input["pattern"].as_str().unwrap_or("?");
                                            format!("find {}", pat)
                                        }
                                        "ls" => {
                                            let path = input["path"].as_str().unwrap_or(".");
                                            let short = path.rsplit('/').next().unwrap_or(path);
                                            format!("ls {}", short)
                                        }
                                        "subagent" => {
                                            let name = input["agent"].as_str()
                                                .or_else(|| input["system_prompt"].as_str().map(|s| if s.len() > 20 { "inline" } else { s }))
                                                .unwrap_or("?");
                                            format!("spawning {}", name)
                                        }
                                        other => {
                                            // MCP or unknown tools — show tool name + first param
                                            let short_name = if other.starts_with("ext__") {
                                                other.splitn(3, "__").last().unwrap_or(other)
                                            } else {
                                                other
                                            };
                                            short_name.to_string()
                                        }
                                    };
                                    if let Some(ref tx) = tx_events_inner {
                                        let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentUpdate {
                                            subagent_id,
                                            agent_name: label_inner.clone(),
                                            status: detail,
                                        }));
                                    }
                                }
                                crate::StreamEvent::Llm(LlmEvent::ToolResult { result, .. }) => {
                                    let preview: String = result.chars().take(300).collect();
                                    tool_log.push(format!("[tool_result]: {}", preview));
                                }
                                crate::StreamEvent::Session(SessionEvent::Usage {
                                    input_tokens, output_tokens,
                                    cache_read_input_tokens, cache_creation_input_tokens,
                                    cache_creation_5m, cache_creation_1h,
                                    model: _,
                                }) => {
                                    total_input_tokens += input_tokens;
                                    total_output_tokens += output_tokens;
                                    total_cache_read += cache_read_input_tokens;
                                    total_cache_creation += cache_creation_input_tokens;
                                    crate::core::rpc_dispatch::merge_split(&mut total_cache_5m, cache_creation_5m);
                                    crate::core::rpc_dispatch::merge_split(&mut total_cache_1h, cache_creation_1h);
                                }
                                crate::StreamEvent::Session(SessionEvent::Error(e)) => {
                                    // Typed outcome from the engine — keep the
                                    // terminal category + correlation ID in the
                                    // subagent's failure state (spec §5.2).
                                    return Err(format!("{} [{}]", e.message, e.category_label()));
                                }
                                crate::StreamEvent::Session(SessionEvent::Done) => break,
                                _ => {}
                            }
                        }
                        _ = &mut timeout_fut => {
                            // Return partial work instead of just an error
                            let mut partial = format!("[TIMED OUT after {}s — partial results below]\n\n", timeout_secs);
                            if !tool_log.is_empty() {
                                partial.push_str(&tool_log.join("\n"));
                                partial.push('\n');
                            }
                            if !final_text.is_empty() {
                                partial.push_str("\n[partial response]:\n");
                                partial.push_str(&final_text);
                            }
                            return Ok(SubagentResult {
                                text: partial,
                                model: model_inner,
                                input_tokens: total_input_tokens,
                                output_tokens: total_output_tokens,
                                cache_read: total_cache_read,
                                cache_creation: total_cache_creation,
                                cache_creation_5m: total_cache_5m,
                                cache_creation_1h: total_cache_1h,
                                tool_count,
                                timed_out: true,
                            });
                        }
                    }
                }

                Ok(SubagentResult {
                    text: final_text,
                    model: model_inner,
                    input_tokens: total_input_tokens,
                    output_tokens: total_output_tokens,
                    cache_read: total_cache_read,
                    cache_creation: total_cache_creation,
                    cache_creation_5m: total_cache_5m,
                    cache_creation_1h: total_cache_1h,
                    tool_count,
                    timed_out: false,
                })
            });

                let _ = result_tx.send(result);
            }));

            if let Err(panic_info) = result {
                let msg = if let Some(s) = panic_info.downcast_ref::<&str>() {
                    s.to_string()
                } else if let Some(s) = panic_info.downcast_ref::<String>() {
                    s.clone()
                } else {
                    "unknown panic".to_string()
                };
                tracing::error!("Subagent thread panicked: {}", msg);
                // result_tx is consumed inside the closure, so we can't send here —
                // the oneshot receiver will see a RecvError, handled below.
            }
        }) {
            Ok(handle) => handle,
            Err(_) => {
                orchestration.rollback(&orchestration_id);
                return Err(RuntimeError::Tool("worker thread initialization failed".into()));
            }
        };
        if let Err(error) = orchestration.mark_running(&orchestration_id) {
            drop(shutdown_tx);
            let _ = thread_handle.join();
            orchestration.rollback(&orchestration_id);
            return Err(RuntimeError::Tool(error));
        }

        let result = result_rx.await;
        let elapsed = start_time.elapsed().as_secs_f64();

        drop(shutdown_tx);
        let _ = thread_handle.join();

        // One-shot workers have no separate collect call: every outcome must pass
        // through terminal, collected, and reconciled before this tool exits.
        if let Some(policy) = &ctx.capabilities.orchestration {
            let terminal = match &result {
                Ok(Ok(result)) if result.timed_out => {
                    agent_core::orchestration::WorkerTerminal::TimedOut
                }
                Ok(Ok(_)) => agent_core::orchestration::WorkerTerminal::Completed,
                Ok(Err(_)) | Err(_) => agent_core::orchestration::WorkerTerminal::Failed,
            };
            policy
                .finish_one_shot(&orchestration_id, terminal)
                .map_err(RuntimeError::Tool)?;
        }

        let log_dir = crate::config::base_dir().join("logs").join("subagents");
        let _ = tokio::fs::create_dir_all(&log_dir).await;
        let timestamp = chrono::Local::now().format("%Y%m%d-%H%M%S");

        match result {
            Ok(Ok(sa_result)) => {
                let preview: String = sa_result.text.chars().take(120).collect();

                if let Some(ref tx) = ctx.channels.tx_events {
                    let _ = tx.send(crate::StreamEvent::Session(SessionEvent::Usage {
                        input_tokens: sa_result.input_tokens,
                        output_tokens: sa_result.output_tokens,
                        cache_read_input_tokens: sa_result.cache_read,
                        cache_creation_input_tokens: sa_result.cache_creation,
                        // TTL split aggregated across the subagent's turns —
                        // forwarded so the parent bills 1h writes at 2.0x, not 1.25x.
                        cache_creation_5m: sa_result.cache_creation_5m,
                        cache_creation_1h: sa_result.cache_creation_1h,
                        model: Some(sa_result.model),
                    }));
                    let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentDone {
                        subagent_id,
                        agent_name: label.clone(),
                        result_preview: preview,
                        duration_secs: elapsed,
                    }));
                }

                let log_content = format!(
                    "# Subagent: {}\nDate: {}\nModel: {}\nTask: {}\nDuration: {:.1}s\nTokens: {}in/{}out ({}cr/{}cw)\nTools used: {}\n\n## Result\n\n{}\n",
                    label, timestamp, params["model"].as_str().unwrap_or("sonnet"),
                    task_preview, elapsed,
                    sa_result.input_tokens, sa_result.output_tokens,
                    sa_result.cache_read, sa_result.cache_creation,
                    sa_result.tool_count, sa_result.text,
                );
                let log_path = log_dir.join(format!("{}-{}.md", timestamp, label));
                let _ = tokio::fs::write(&log_path, &log_content).await;

                Ok(format!("[subagent:{}] {}", label, sa_result.text))
            }
            Ok(Err(e)) => {
                if let Some(ref tx) = ctx.channels.tx_events {
                    let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentDone {
                        subagent_id,
                        agent_name: label.clone(),
                        result_preview: format!("ERROR: {}", e),
                        duration_secs: elapsed,
                    }));
                }
                let log_path = log_dir.join(format!("{}-{}-error.md", timestamp, label));
                let _ = tokio::fs::write(
                    &log_path,
                    format!(
                        "# Subagent ERROR: {}\nTask: {}\nError: {}\n",
                        label, task_preview, e
                    ),
                )
                .await;
                Ok(format!("[subagent:{} ERROR] {}", label, e))
            }
            Err(_) => {
                if let Some(ref tx) = ctx.channels.tx_events {
                    let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentDone {
                        subagent_id,
                        agent_name: label.clone(),
                        result_preview: "Task panicked or dropped".to_string(),
                        duration_secs: elapsed,
                    }));
                }
                Ok(format!(
                    "[subagent:{} ERROR] Subagent task panicked or was dropped",
                    label
                ))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tools::test_helpers::create_tool_context;
    use crate::tools::Tool;
    use serde_json::json;

    #[test]
    fn test_subagent_tool_schema() {
        let tool = SubagentTool;
        assert_eq!(tool.name(), "subagent");
        assert!(!tool.description().is_empty());

        let params = tool.parameters();
        assert_eq!(params["type"], "object");
        assert!(params["properties"].is_object());
        assert!(params["required"].is_array());
    }

    #[tokio::test]
    async fn test_subagent_blank_agent_uses_system_prompt() {
        let tool = SubagentTool;
        let ctx = create_tool_context();
        let params = json!({
            "agent": "",
            "system_prompt": "You are a concise test subagent. Reply with only: ok",
            "task": "Say ok",
            "model": "anthropic/claude-sonnet-4-6",
            "timeout": 1
        });

        let result = tool.execute(params, ctx).await;
        assert!(
            result.is_ok(),
            "blank agent should not be resolved as ~/.synaps-cli/agents/.md: {result:?}"
        );
    }
}