recursive/tools/

sub_agent.rs

//! Sub-agent tool: spawn a fresh agent loop with a restricted tool subset.
//!
//! The parent agent can delegate focused sub-tasks to a child agent that
//! starts with an empty transcript. This prevents the parent's context
//! window from growing with intermediate exploration steps.
//!
//! Recursive safety: a depth limit (env `RECURSIVE_SUBAGENT_MAX_DEPTH`,
//! default 2) prevents unbounded nesting. Each nested invocation increments
//! a counter; when the limit is reached, the tool returns an error string
//! instead of spawning.
//!
//! # Agent types (`subagent_type`)
//!
//! Inspired by fake-cc's `AgentTool`, this tool accepts an optional
//! `subagent_type` parameter that selects a named agent personality:
//!
//! - `"explore"`: read-only tools only; declared [`ReadOnly`](crate::tools::ToolSideEffect::ReadOnly)
//!   so the dispatch layer can run multiple explore sub-agents **in parallel**.
//! - `"general_purpose"` (default): full tool access; declared `External`.

use async_trait::async_trait;
use serde_json::{json, Value};
use std::sync::Arc;

use crate::agent::{FinishReason, PlanningMode};
use crate::error::{Error, Result};
use crate::kernel::{AgentKernel, TurnContext};
use crate::llm::{LlmProvider, ToolSpec};
use crate::message::Message;
use crate::permissions::PermissionMode;
use crate::tools::PermissionHook;
use crate::tools::{Tool, ToolRegistry};

// ---------------------------------------------------------------------------
// AgentType — named sub-agent personality
// ---------------------------------------------------------------------------

/// Named sub-agent personality, aligned with fake-cc's `subagent_type`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AgentType {
    /// Read-only exploration agent. Restricts tool access to read-only tools.
    /// Declared `ReadOnly` so the parallel-dispatch layer can run multiple
    /// explore sub-agents concurrently.
    Explore,
    /// General-purpose agent with access to the full (parent-provided) tool
    /// registry. Declared `External` (the conservative default).
    GeneralPurpose,
}

impl AgentType {
    /// Parse from the `subagent_type` JSON string.
    /// Returns `None` for unknown values (caller falls back to `GeneralPurpose`).
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            "explore" => Some(Self::Explore),
            "general_purpose" => Some(Self::GeneralPurpose),
            _ => None,
        }
    }

    /// `true` for agent types that only read, never write.
    pub fn is_read_only(self) -> bool {
        matches!(self, Self::Explore)
    }

    /// System prompt fragment to prepend for this agent type.
    pub fn system_prompt_hint(self) -> &'static str {
        match self {
            Self::Explore => {
                "You are an exploration sub-agent. Use only read tools to gather \
                 information. Do NOT write or modify files. Be thorough and concise."
            }
            Self::GeneralPurpose => {
                "You are a focused sub-agent. Complete the given task using the \
                 available tools. Be concise."
            }
        }
    }

    /// Restricted tool names for this agent type, or `None` for "use caller-supplied list".
    pub fn allowed_tool_names(self) -> Option<Vec<String>> {
        match self {
            Self::Explore => Some(vec![
                "read_file".to_string(),
                "list_dir".to_string(),
                "search_files".to_string(),
                "recall".to_string(),
                "web_fetch".to_string(),
                "sub_agent".to_string(),
            ]),
            Self::GeneralPurpose => None,
        }
    }
}

/// The sub-agent tool.
///
/// Constructed with:
/// - `workspace`: path for sandboxed tools
/// - `provider`: the LLM provider (shared Arc from parent)
/// - `all_tools`: the full tool registry from which the sub-agent can draw
/// - `max_depth`: absolute depth limit (env-configured)
/// - `current_depth`: how deep we already are (passed from parent)
/// - `permission_hook`: optional permission hook inherited from the parent agent
pub struct SubAgent {
    workspace: std::path::PathBuf,
    provider: Arc<dyn LlmProvider>,
    all_tools: ToolRegistry,
    max_depth: usize,
    current_depth: usize,
    permission_hook: Option<Arc<dyn PermissionHook>>,
}

impl SubAgent {
    pub fn new(
        workspace: impl Into<std::path::PathBuf>,
        provider: Arc<dyn LlmProvider>,
        all_tools: ToolRegistry,
        max_depth: usize,
        current_depth: usize,
        permission_hook: Option<Arc<dyn PermissionHook>>,
    ) -> Self {
        Self {
            workspace: workspace.into(),
            provider,
            all_tools,
            max_depth,
            current_depth,
            permission_hook,
        }
    }

    /// Build a restricted tool registry containing only the named tools.
    fn build_sub_registry(&self, tool_names: &[String]) -> ToolRegistry {
        let mut reg = self.all_tools.with_same_transport();
        for name in tool_names {
            if let Some(tool) = self.all_tools.get(name) {
                reg = reg.register(tool);
            }
        }
        reg
    }

    /// Default tool set when no `tools` arg is given: read-only tools.
    fn default_tool_names() -> Vec<String> {
        vec![
            "read_file".to_string(),
            "list_dir".to_string(),
            "search_files".to_string(),
            "web_fetch".to_string(),
        ]
    }
}

#[async_trait]
impl Tool for SubAgent {
    fn spec(&self) -> ToolSpec {
        ToolSpec {
            name: "sub_agent".into(),
            description: "Spawn a fresh agent with its own transcript to complete a focused sub-task. Returns the sub-agent's final response.".into(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "prompt": {
                        "type": "string",
                        "description": "The goal / prompt for the sub-agent"
                    },
                    "subagent_type": {
                        "type": "string",
                        "enum": ["explore", "general_purpose"],
                        "description": "Agent personality. 'explore': read-only tools only, can run in parallel with other explore agents. 'general_purpose' (default): full tool access, runs sequentially."
                    },
                    "max_steps": {
                        "type": "integer",
                        "description": "Maximum steps for the sub-agent (default 30, capped at parent's remaining budget)",
                        "default": 30
                    },
                    "tools": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": "Optional list of tool names to make available to the sub-agent. Ignored when subagent_type is 'explore' (which enforces its own read-only tool set). Default: read_file, list_dir, search_files, web_fetch"
                    }
                },
                "required": ["prompt"]
            }),
        }
    }

    /// Returns `true` when `subagent_type` is `"explore"`, making the dispatch
    /// layer treat this call as read-only and eligible for parallel execution.
    fn is_readonly_for_args(&self, arguments: &serde_json::Value) -> bool {
        arguments
            .get("subagent_type")
            .and_then(|v| v.as_str())
            .and_then(AgentType::parse)
            .map(AgentType::is_read_only)
            .unwrap_or(false)
    }

    async fn execute(&self, arguments: Value) -> Result<String> {
        let prompt = arguments["prompt"]
            .as_str()
            .ok_or_else(|| Error::BadToolArgs {
                name: "sub_agent".into(),
                message: "missing required parameter: prompt".to_string(),
            })?;

        // Resolve agent type (defaults to GeneralPurpose if absent or unknown).
        let agent_type = arguments
            .get("subagent_type")
            .and_then(|v| v.as_str())
            .and_then(AgentType::parse)
            .unwrap_or(AgentType::GeneralPurpose);

        let max_steps = arguments["max_steps"].as_i64().unwrap_or(30).clamp(1, 100) as usize;

        // Tool list: explore type enforces its own read-only set;
        // general_purpose respects the caller-supplied list or the default.
        let tool_names: Vec<String> = if let Some(forced) = agent_type.allowed_tool_names() {
            forced
        } else {
            arguments["tools"]
                .as_array()
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| v.as_str().map(String::from))
                        .collect()
                })
                .unwrap_or_else(Self::default_tool_names)
        };

        // Depth limit check
        if self.current_depth >= self.max_depth {
            return Ok(format!(
                "ERROR: sub-agent depth limit reached (max_depth={}). Cannot spawn deeper sub-agent.",
                self.max_depth
            ));
        }

        // Build the sub-agent's tool registry
        let mut sub_registry = self.build_sub_registry(&tool_names);

        // Register a fresh SubAgent with incremented depth so the child
        // can also spawn sub-agents (up to the limit).
        let child_sub = SubAgent::new(
            &self.workspace,
            self.provider.clone(),
            self.all_tools.clone(),
            self.max_depth,
            self.current_depth + 1,
            self.permission_hook.clone(),
        );
        sub_registry = sub_registry.register(Arc::new(child_sub));

        // Build and run the sub-agent using AgentKernel (stateless, single-turn)
        let kernel = AgentKernel::builder()
            .llm(self.provider.clone())
            .tools(sub_registry)
            .max_steps(max_steps)
            .build()
            .map_err(|e| Error::Tool {
                name: "sub_agent".into(),
                message: format!("failed to build sub-agent kernel: {e}"),
            })?;

        let ctx = TurnContext {
            messages: vec![
                Message::system(agent_type.system_prompt_hint().to_string()),
                Message::user(prompt.to_string()),
            ],
            step_events_tx: None,
            plan_confirmed: false,
            plan_buffer: None,
            tool_specs: kernel.tools().specs(),
            streaming: false,
            permission_hook: self.permission_hook.clone(),
            planning_mode: PlanningMode::default(),
            exploring_plan_mode: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
            permission_mode: PermissionMode::Default,
            mailbox: None,
        };

        let outcome = kernel.run(ctx).await.map_err(|e| Error::Tool {
            name: "sub_agent".into(),
            message: format!("sub-agent failed: {e}"),
        })?;

        let finish_label = match &outcome.finish_reason {
            FinishReason::NoMoreToolCalls => "NoMoreToolCalls",
            FinishReason::BudgetExceeded => "BudgetExceeded",
            FinishReason::ProviderStop(r) => r,
            FinishReason::Stuck { .. } => "Stuck",
            FinishReason::TranscriptLimit { .. } => "TranscriptLimit",
            FinishReason::PlanPending => "PlanPending",
            FinishReason::Cancelled => "Cancelled",
            FinishReason::PermissionDenialLimit => "PermissionDenialLimit",
        };

        let final_text = outcome
            .final_text
            .unwrap_or_else(|| "(no final message)".to_string());

        Ok(format!(
            "[sub-agent finished: {finish_label}]\n{final_text}"
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::{Completion, MockProvider, ToolCall};
    use crate::tools::{
        ApplyPatch, ListDir, LocalTransport, ReadFile, SearchFiles, ToolTransport, WriteFile,
    };

    /// Helper: create a MockProvider with the given scripted completions.
    fn mock_provider(script: Vec<Completion>) -> Arc<dyn LlmProvider> {
        Arc::new(MockProvider::new(script))
    }

    /// Helper: build a full tool registry with read-only + write tools.
    fn full_tool_registry(workspace: &std::path::Path) -> ToolRegistry {
        let transport: Arc<dyn ToolTransport> = Arc::new(LocalTransport);
        ToolRegistry::new(transport)
            .register(Arc::new(ReadFile::new(workspace)))
            .register(Arc::new(ListDir::new(workspace)))
            .register(Arc::new(SearchFiles::new(workspace)))
            .register(Arc::new(WriteFile::new(workspace)))
            .register(Arc::new(ApplyPatch::new(workspace)))
    }

    #[tokio::test]
    async fn sub_agent_basic_dispatch() {
        // Sub-agent gets one completion with no tool calls → NoMoreToolCalls
        let provider = mock_provider(vec![Completion {
            content: "The answer is 42.".to_string(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: None,
            reasoning_content: None,
        }]);

        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());

        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        let result = sub
            .execute(json!({"prompt": "What is the meaning of life?"}))
            .await
            .unwrap();

        assert!(result.contains("NoMoreToolCalls"));
        assert!(result.contains("The answer is 42."));
    }

    #[tokio::test]
    async fn sub_agent_depth_limit_enforced() {
        // Create a sub-agent at depth=2 with max_depth=2.
        // It should refuse to spawn deeper.
        let provider = mock_provider(vec![]);
        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());

        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 2, None);

        let result = sub
            .execute(json!({"prompt": "do something"}))
            .await
            .unwrap();

        assert!(result.contains("depth limit reached"));
        assert!(result.contains("max_depth=2"));
    }

    #[tokio::test]
    async fn sub_agent_tool_subset_respected() {
        // Parent passes tools: ["read_file"]; sub-agent must NOT have apply_patch.
        let provider = mock_provider(vec![Completion {
            content: "done".to_string(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: None,
            reasoning_content: None,
        }]);

        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());

        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        // Execute with only read_file allowed
        let _ = sub
            .execute(json!({"prompt": "read something", "tools": ["read_file"]}))
            .await
            .unwrap();

        // We can't easily inspect the sub-agent's registry from here,
        // but we can verify the sub-agent ran successfully with just read_file.
        // The real test is that apply_patch is NOT in the default set.
        let defaults = SubAgent::default_tool_names();
        assert!(!defaults.contains(&"apply_patch".to_string()));
        assert!(!defaults.contains(&"write_file".to_string()));
        assert!(defaults.contains(&"read_file".to_string()));
    }

    #[tokio::test]
    async fn sub_agent_max_steps_capped() {
        // Sub-agent with max_steps=5 should hit BudgetExceeded when
        // MockProvider keeps asking for tool calls.
        let tmp = tempfile::tempdir().unwrap();
        // Create a file so read_file succeeds
        std::fs::write(tmp.path().join("test.txt"), b"hello").unwrap();

        let mut script = Vec::new();
        for _ in 0..10 {
            script.push(Completion {
                content: "".to_string(),
                tool_calls: vec![ToolCall {
                    id: "c1".into(),
                    name: "read_file".into(),
                    arguments: json!({"path": "test.txt"}),
                }],
                finish_reason: Some("tool_calls".into()),
                usage: None,
                reasoning_content: None,
            });
        }

        let provider = mock_provider(script);
        let all_tools = full_tool_registry(tmp.path());

        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        let result = sub
            .execute(json!({"prompt": "loop", "max_steps": 5}))
            .await
            .unwrap();

        assert!(result.contains("BudgetExceeded"));
    }

    #[tokio::test]
    async fn sub_agent_default_tools_are_read_only() {
        let defaults = SubAgent::default_tool_names();
        assert!(defaults.contains(&"read_file".to_string()));
        assert!(defaults.contains(&"list_dir".to_string()));
        assert!(defaults.contains(&"search_files".to_string()));
        assert!(defaults.contains(&"web_fetch".to_string()));
        assert_eq!(defaults.len(), 4);
    }

    #[tokio::test]
    async fn sub_agent_missing_prompt_returns_error() {
        let provider = mock_provider(vec![]);
        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());

        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        let result = sub.execute(json!({})).await;
        assert!(result.is_err());
        let err = result.unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("missing required parameter: prompt"));
    }

    #[tokio::test]
    async fn sub_agent_nested_depth_works() {
        // Test that a SubAgent at depth=0 can spawn a child (depth=1)
        // which can spawn a grandchild (depth=2) — but the grandchild
        // is at the limit (max_depth=2) so it cannot spawn deeper.
        //
        // Scripted completions consumed in order:
        //   1. Child agent's first call → sub_agent tool call
        //   2. Grandchild agent's first call → sub_agent tool call (denied by depth)
        //   3. Grandchild agent's second call → "grandchild done"
        //   4. Child agent's second call → "child done"
        let tmp = tempfile::tempdir().unwrap();
        std::fs::write(tmp.path().join("test.txt"), b"hello").unwrap();

        let provider = mock_provider(vec![
            // 1. Child agent: calls sub_agent
            Completion {
                content: "".to_string(),
                tool_calls: vec![ToolCall {
                    id: "c1".into(),
                    name: "sub_agent".into(),
                    arguments: json!({"prompt": "grandchild task"}),
                }],
                finish_reason: Some("tool_calls".into()),
                usage: None,
                reasoning_content: None,
            },
            // 2. Grandchild agent: tries to spawn deeper (denied)
            Completion {
                content: "".to_string(),
                tool_calls: vec![ToolCall {
                    id: "c2".into(),
                    name: "sub_agent".into(),
                    arguments: json!({"prompt": "great-grandchild task"}),
                }],
                finish_reason: Some("tool_calls".into()),
                usage: None,
                reasoning_content: None,
            },
            // 3. Grandchild agent: finishes after seeing depth error
            Completion {
                content: "grandchild done".to_string(),
                tool_calls: vec![],
                finish_reason: Some("stop".into()),
                usage: None,
                reasoning_content: None,
            },
            // 4. Child agent: finishes
            Completion {
                content: "child done".to_string(),
                tool_calls: vec![],
                finish_reason: Some("stop".into()),
                usage: None,
                reasoning_content: None,
            },
        ]);

        let all_tools = full_tool_registry(tmp.path());

        // Parent at depth=0, max_depth=2
        let parent = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        let result = parent
            .execute(json!({"prompt": "parent task", "tools": ["sub_agent", "read_file"]}))
            .await
            .unwrap();

        // The parent should complete successfully with child's result
        assert!(result.contains("NoMoreToolCalls"), "result: {result}");
        assert!(result.contains("child done"), "result: {result}");
        // The grandchild's result is embedded in the child's transcript
        // (as a tool result), not in the final message. The depth limit
        // was enforced: the grandchild could not spawn deeper.
        // This test verifies the nesting works without panicking.
    }

    // ── AgentType & is_readonly_for_args tests (Goal 175) ─────────────────

    #[test]
    fn explore_agent_type_is_read_only() {
        assert!(AgentType::Explore.is_read_only());
        assert!(!AgentType::GeneralPurpose.is_read_only());
    }

    #[test]
    fn agent_type_from_str_roundtrip() {
        assert_eq!(AgentType::parse("explore"), Some(AgentType::Explore));
        assert_eq!(
            AgentType::parse("general_purpose"),
            Some(AgentType::GeneralPurpose)
        );
        assert_eq!(AgentType::parse("unknown"), None);
        assert_eq!(AgentType::parse(""), None);
    }

    #[test]
    fn explore_agent_has_restricted_tool_list() {
        let names = AgentType::Explore.allowed_tool_names().unwrap();
        // Must contain read tools
        assert!(names.contains(&"read_file".to_string()));
        assert!(names.contains(&"list_dir".to_string()));
        assert!(names.contains(&"search_files".to_string()));
        // Must NOT contain write tools
        assert!(!names.contains(&"write_file".to_string()));
        assert!(!names.contains(&"apply_patch".to_string()));
        assert!(!names.contains(&"run_shell".to_string()));
    }

    #[test]
    fn general_purpose_agent_has_no_forced_tool_list() {
        assert!(AgentType::GeneralPurpose.allowed_tool_names().is_none());
    }

    #[test]
    fn is_readonly_for_args_explore_returns_true() {
        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());
        let provider = mock_provider(vec![]);
        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        assert!(sub.is_readonly_for_args(&json!({"subagent_type": "explore"})));
    }

    #[test]
    fn is_readonly_for_args_general_purpose_returns_false() {
        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());
        let provider = mock_provider(vec![]);
        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        assert!(!sub.is_readonly_for_args(&json!({"subagent_type": "general_purpose"})));
    }

    #[test]
    fn is_readonly_for_args_missing_type_returns_false() {
        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());
        let provider = mock_provider(vec![]);
        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        assert!(!sub.is_readonly_for_args(&json!({"prompt": "hello"})));
    }

    #[test]
    fn is_readonly_for_args_unknown_type_returns_false() {
        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());
        let provider = mock_provider(vec![]);
        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        assert!(!sub.is_readonly_for_args(&json!({"subagent_type": "super_agent"})));
    }

    #[tokio::test]
    async fn explore_agent_dispatch_succeeds() {
        let provider = mock_provider(vec![Completion {
            content: "Exploration complete.".to_string(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: None,
            reasoning_content: None,
        }]);

        let tmp = tempfile::tempdir().unwrap();
        let all_tools = full_tool_registry(tmp.path());
        let sub = SubAgent::new(tmp.path(), provider, all_tools, 2, 0, None);

        let result = sub
            .execute(json!({
                "prompt": "Explore the workspace",
                "subagent_type": "explore"
            }))
            .await
            .unwrap();

        assert!(result.contains("NoMoreToolCalls"));
        assert!(result.contains("Exploration complete."));
    }
}