diff --git a/crates/agent/src/cli.rs b/crates/agent/src/cli.rs index 87763d5c7..c0ab22e53 100644 --- a/crates/agent/src/cli.rs +++ b/crates/agent/src/cli.rs @@ -1,6 +1,7 @@ use crate::{ AgentEvent, AnthropicProfile, GeminiProfile, LocalExecutionEnvironment, OpenAiProfile, ProviderProfile, Session, SessionConfig, ToolApprovalFn, Turn, + subagent::{SessionFactory, SubAgentManager}, }; use clap::{Parser, ValueEnum}; use llm::client::Client; @@ -65,6 +66,8 @@ fn tool_category(name: &str) -> &'static str { match name { "read_file" | "read_many_files" | "grep" | "glob" | "list_dir" => "read", "write_file" | "edit_file" | "apply_patch" => "write", + // subagent tools inherit parent permissions, always allowed + "spawn_agent" | "send_input" | "wait" | "close_agent" => "subagent", // shell and unknown tools require highest permission _ => "shell", } @@ -74,6 +77,7 @@ fn is_auto_approved(level: PermissionLevel, category: &str) -> bool { matches!( (level, category), (_, "read") + | (_, "subagent") | (PermissionLevel::ReadWrite | PermissionLevel::Full, "write") | (PermissionLevel::Full, "shell") ) @@ -128,12 +132,12 @@ fn build_tool_approval( }) } -fn build_profile(provider: &str, model: &str) -> Arc { +fn build_profile(provider: &str, model: &str) -> Box { match provider { - "openai" => Arc::new(OpenAiProfile::new(model)), - "gemini" => Arc::new(GeminiProfile::new(model)), + "openai" => Box::new(OpenAiProfile::new(model)), + "gemini" => Box::new(GeminiProfile::new(model)), // anthropic and unknown providers - _ => Arc::new(AnthropicProfile::new(model)), + _ => Box::new(AnthropicProfile::new(model)), } } @@ -321,12 +325,12 @@ pub async fn run() -> anyhow::Result<()> { "{}Using model: {model}{}", styles.dim, styles.reset, ); - let profile = build_profile(&cli.provider, model); + let mut profile = build_profile(&cli.provider, model); // Build execution environment let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); let cwd_str = cwd.to_string_lossy().to_string(); - let env = Arc::new(LocalExecutionEnvironment::new(cwd)); + let env: Arc = Arc::new(LocalExecutionEnvironment::new(cwd)); // Build tool approval callback let is_interactive = std::io::stdin().is_terminal() && !cli.auto_approve; @@ -338,6 +342,34 @@ pub async fn run() -> anyhow::Result<()> { ..SessionConfig::default() }; + // Register subagent tools + let manager = Arc::new(tokio::sync::Mutex::new( + SubAgentManager::new(config.max_subagent_depth), + )); + let factory_client = client.clone(); + let factory_provider = cli.provider.clone(); + let factory_model = model.to_string(); + let factory_env = Arc::clone(&env); + let factory_approval = config.tool_approval.clone(); + let factory: SessionFactory = Arc::new(move || { + let child_profile: Arc = match factory_provider.as_str() { + "openai" => Arc::new(OpenAiProfile::new(&factory_model)), + "gemini" => Arc::new(GeminiProfile::new(&factory_model)), + _ => Arc::new(AnthropicProfile::new(&factory_model)), + }; + Session::new( + factory_client.clone(), + child_profile, + Arc::clone(&factory_env), + SessionConfig { + tool_approval: factory_approval.clone(), + ..SessionConfig::default() + }, + ) + }); + profile.register_subagent_tools(manager, factory, 0); + let profile: Arc = Arc::from(profile); + let mut session = Session::new(client, profile, env, config); // SIGINT handler @@ -433,6 +465,14 @@ mod tests { assert_eq!(tool_category("shell"), "shell"); } + #[test] + fn tool_category_subagent_tools() { + assert_eq!(tool_category("spawn_agent"), "subagent"); + assert_eq!(tool_category("send_input"), "subagent"); + assert_eq!(tool_category("wait"), "subagent"); + assert_eq!(tool_category("close_agent"), "subagent"); + } + #[test] fn tool_category_unknown_defaults_to_shell() { assert_eq!(tool_category("some_random_tool"), "shell"); @@ -443,6 +483,7 @@ mod tests { #[test] fn is_auto_approved_read_only() { assert!(is_auto_approved(PermissionLevel::ReadOnly, "read")); + assert!(is_auto_approved(PermissionLevel::ReadOnly, "subagent")); assert!(!is_auto_approved(PermissionLevel::ReadOnly, "write")); assert!(!is_auto_approved(PermissionLevel::ReadOnly, "shell")); } @@ -450,6 +491,7 @@ mod tests { #[test] fn is_auto_approved_read_write() { assert!(is_auto_approved(PermissionLevel::ReadWrite, "read")); + assert!(is_auto_approved(PermissionLevel::ReadWrite, "subagent")); assert!(is_auto_approved(PermissionLevel::ReadWrite, "write")); assert!(!is_auto_approved(PermissionLevel::ReadWrite, "shell")); } @@ -457,6 +499,7 @@ mod tests { #[test] fn is_auto_approved_full() { assert!(is_auto_approved(PermissionLevel::Full, "read")); + assert!(is_auto_approved(PermissionLevel::Full, "subagent")); assert!(is_auto_approved(PermissionLevel::Full, "write")); assert!(is_auto_approved(PermissionLevel::Full, "shell")); } @@ -534,4 +577,22 @@ mod tests { let profile = build_profile("gemini", "model"); assert_eq!(profile.id(), "gemini"); } + + // subagent tool registration tests + + #[test] + fn build_profile_can_register_subagent_tools() { + let mut profile = build_profile("anthropic", "model"); + let manager = Arc::new(tokio::sync::Mutex::new(SubAgentManager::new(1))); + let factory: SessionFactory = Arc::new(|| { + panic!("factory should not be called in this test"); + }); + profile.register_subagent_tools(manager, factory, 0); + + let names = profile.tool_registry().names(); + assert!(names.contains(&"spawn_agent".to_string())); + assert!(names.contains(&"send_input".to_string())); + assert!(names.contains(&"wait".to_string())); + assert!(names.contains(&"close_agent".to_string())); + } } diff --git a/crates/agent/src/subagent.rs b/crates/agent/src/subagent.rs index 3d64bab64..b9197ee11 100644 --- a/crates/agent/src/subagent.rs +++ b/crates/agent/src/subagent.rs @@ -18,23 +18,11 @@ pub struct SubAgentResult { } pub struct SubAgent { - #[allow(dead_code)] - id: String, - #[allow(dead_code)] - depth: usize, task: Option>>, followup_queue: Arc>>, cancel_token: CancellationToken, } -#[cfg(test)] -impl SubAgent { - #[must_use] - pub fn depth(&self) -> usize { - self.depth - } -} - pub struct SubAgentManager { agents: HashMap, max_depth: usize, @@ -83,8 +71,6 @@ impl SubAgentManager { self.agents.insert( agent_id.clone(), SubAgent { - id: agent_id.clone(), - depth, task: Some(task), followup_queue, cancel_token, @@ -339,7 +325,6 @@ mod tests { let agent_id = result.unwrap(); assert!(!agent_id.is_empty()); assert!(manager.get(&agent_id).is_some()); - assert_eq!(manager.get(&agent_id).unwrap().depth(), 0); } #[tokio::test] diff --git a/crates/attractor/src/cli/backend.rs b/crates/attractor/src/cli/backend.rs index 533d99507..c89fa2d52 100644 --- a/crates/attractor/src/cli/backend.rs +++ b/crates/attractor/src/cli/backend.rs @@ -6,6 +6,7 @@ use async_trait::async_trait; use agent::{ AgentEvent, AnthropicProfile, ExecutionEnvironment, GeminiProfile, OpenAiProfile, ProviderProfile, Session, SessionConfig, Turn, + subagent::{SessionFactory, SubAgentManager}, }; use llm::client::Client; use terminal::Styles; @@ -54,22 +55,51 @@ impl AgentBackend { .await .map_err(|e| AttractorError::Handler(format!("Failed to create LLM client: {e}")))?; - let profile = self.build_profile(); + let mut profile = self.build_profile(); let config = SessionConfig { reasoning_effort: Some(node.reasoning_effort().to_string()), ..SessionConfig::default() }; + let manager = Arc::new(tokio::sync::Mutex::new( + SubAgentManager::new(config.max_subagent_depth), + )); + + // Build factory that creates child sessions WITHOUT subagent tools + let factory_client = client.clone(); + let factory_provider = self.provider.clone(); + let factory_model = self.model.clone(); + let factory_env = Arc::clone(execution_env); + let factory: SessionFactory = Arc::new(move || { + let child_profile = { + let provider = factory_provider.as_deref().unwrap_or("anthropic"); + match provider { + "openai" => Arc::new(OpenAiProfile::new(&factory_model)) as Arc, + "gemini" => Arc::new(GeminiProfile::new(&factory_model)) as Arc, + _ => Arc::new(AnthropicProfile::new(&factory_model)) as Arc, + } + }; + Session::new( + factory_client.clone(), + child_profile, + Arc::clone(&factory_env), + SessionConfig::default(), + ) + }); + + profile.register_subagent_tools(manager, factory, 0); + let profile: Arc = Arc::from(profile); + Ok(Session::new(client, profile, Arc::clone(execution_env), config)) } - fn build_profile(&self) -> Arc { + fn build_profile(&self) -> Box { let provider = self.provider.as_deref().unwrap_or("anthropic"); match provider { - "openai" => Arc::new(OpenAiProfile::new(&self.model)), - "gemini" => Arc::new(GeminiProfile::new(&self.model)), - _ => Arc::new(AnthropicProfile::new(&self.model)), + "openai" => Box::new(OpenAiProfile::new(&self.model)), + "gemini" => Box::new(GeminiProfile::new(&self.model)), + _ => Box::new(AnthropicProfile::new(&self.model)), } } } @@ -416,6 +446,7 @@ fn format_tool_args(args: &serde_json::Value) -> String { #[cfg(test)] mod tests { use super::*; + use agent::subagent::SessionFactory; #[test] fn agent_backend_stores_config() { @@ -442,4 +473,27 @@ mod tests { ); assert!(backend.sessions.lock().unwrap().is_empty()); } + + #[test] + fn build_profile_can_register_subagent_tools() { + let styles = Box::leak(Box::new(Styles::new(false))); + let backend = AgentBackend::new( + "claude-opus-4-6".to_string(), + None, + 0, + styles, + ); + let mut profile = backend.build_profile(); + let manager = Arc::new(tokio::sync::Mutex::new(SubAgentManager::new(1))); + let factory: SessionFactory = Arc::new(|| { + panic!("factory should not be called in this test"); + }); + profile.register_subagent_tools(manager, factory, 0); + + let names = profile.tool_registry().names(); + assert!(names.contains(&"spawn_agent".to_string())); + assert!(names.contains(&"send_input".to_string())); + assert!(names.contains(&"wait".to_string())); + assert!(names.contains(&"close_agent".to_string())); + } }