yoagent/
sub_agent.rs

1//! Sub-agent tool — delegates tasks to a child agent loop.
2//!
3//! The `SubAgentTool` implements `AgentTool` and internally runs `agent_loop()`
4//! with its own system prompt, tools, and provider. The parent LLM invokes it
5//! like any other tool, passing a natural-language `task` string.
6//!
7//! # Design
8//!
9//! - **Context isolation**: each invocation starts a fresh conversation
10//! - **Nesting supported**: sub-agents can contain other SubAgentTools for recursive delegation (use `with_max_turns()` to bound depth)
11//! - **Cancellation propagation**: the parent's cancel token is forwarded
12//! - **Event forwarding**: sub-agent events stream to the parent via `on_update`
13//!
14//! # Example
15//!
16//! ```rust,no_run
17//! use yoagent::sub_agent::SubAgentTool;
18//! use yoagent::provider::ModelConfig;
19//!
20//! // Provider selected from the config's protocol; key from ANTHROPIC_API_KEY.
21//! let researcher = SubAgentTool::from_config(
22//!     "researcher",
23//!     ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"),
24//! )
25//! .with_description("Searches codebases and documents")
26//! .with_system_prompt("You are a research assistant.");
27//! ```
28
29use crate::agent_loop::{agent_loop, AgentLoopConfig};
30use crate::context::ExecutionLimits;
31use crate::provider::model::ModelConfig;
32use crate::provider::StreamProvider;
33use crate::shared_state::SharedState;
34use crate::tools::shared_state_tool::SharedStateTool;
35use crate::types::*;
36use std::sync::Arc;
37use tokio::sync::mpsc;
38
39/// Default max turns for sub-agents (prevents runaway execution).
40const DEFAULT_MAX_TURNS: usize = 10;
41
42/// A tool that delegates work to a child agent loop.
43///
44/// When the parent LLM calls this tool, it spawns a fresh `agent_loop()` with
45/// its own system prompt, tools, and provider. The sub-agent runs to completion
46/// and its final text output is returned as the tool result.
47pub struct SubAgentTool {
48    tool_name: String,
49    tool_description: String,
50    system_prompt: String,
51    skills_prompt: String,
52    model: String,
53    api_key: String,
54    provider: Arc<dyn StreamProvider>,
55    tools: Vec<Arc<dyn AgentTool>>,
56    thinking_level: ThinkingLevel,
57    max_tokens: Option<u32>,
58    temperature: Option<f32>,
59    cache_config: CacheConfig,
60    tool_execution: ToolExecutionStrategy,
61    retry_config: crate::retry::RetryConfig,
62    max_turns: usize,
63    shared_state: Option<SharedState>,
64    turn_delay: Option<std::time::Duration>,
65    model_config: Option<ModelConfig>,
66    tool_middleware: Vec<Arc<dyn ToolMiddleware>>,
67}
68
69impl SubAgentTool {
70    /// Create a new sub-agent tool with a name and provider.
71    #[deprecated(
72        since = "0.10.0",
73        note = "use SubAgentTool::from_config(name, config) — provider + env key \
74                resolved automatically — or SubAgentTool::from_provider(name, provider, config) \
75                for a custom provider; will be removed in 1.0"
76    )]
77    pub fn new(name: impl Into<String>, provider: Arc<dyn StreamProvider>) -> Self {
78        Self::build(name, provider)
79    }
80
81    /// Internal constructor shared by `new` and the `from_*` builders (not
82    /// deprecated, so the builders don't trip the deprecation lint).
83    fn build(name: impl Into<String>, provider: Arc<dyn StreamProvider>) -> Self {
84        let name = name.into();
85        Self {
86            tool_description: format!("Delegate a task to the '{}' sub-agent", name),
87            tool_name: name,
88            system_prompt: String::new(),
89            skills_prompt: String::new(),
90            model: String::new(),
91            api_key: String::new(),
92            provider,
93            tools: Vec::new(),
94            thinking_level: ThinkingLevel::Off,
95            max_tokens: None,
96            temperature: None,
97            cache_config: CacheConfig::default(),
98            tool_execution: ToolExecutionStrategy::default(),
99            retry_config: crate::retry::RetryConfig::default(),
100            max_turns: DEFAULT_MAX_TURNS,
101            shared_state: None,
102            turn_delay: None,
103            model_config: None,
104            tool_middleware: Vec::new(),
105        }
106    }
107
108    /// Create a sub-agent from a name and [`ModelConfig`], selecting the
109    /// built-in provider for the config's protocol.
110    ///
111    /// Mirrors [`Agent::from_config`](crate::Agent::from_config): the model
112    /// id, provider, and pricing come from one config, and the API key is
113    /// resolved from the provider-conventional env var unless set explicitly
114    /// with [`with_api_key`](Self::with_api_key).
115    ///
116    /// # Panics
117    ///
118    /// Never panics — the default registry covers every [`ApiProtocol`]
119    /// variant. Use [`from_config_with`](Self::from_config_with) with a custom
120    /// registry when a protocol may be unregistered and you want a `Result`.
121    ///
122    /// [`ApiProtocol`]: crate::provider::ApiProtocol
123    pub fn from_config(name: impl Into<String>, config: ModelConfig) -> Self {
124        Self::from_config_with(&crate::provider::ProviderRegistry::default(), name, config)
125            .expect("default registry covers all built-in protocols")
126    }
127
128    /// Like [`from_config`](Self::from_config) but resolves the provider from a
129    /// caller-supplied registry, returning an error if the config's protocol
130    /// isn't registered. Mirrors
131    /// [`Agent::from_config_with`](crate::Agent::from_config_with).
132    pub fn from_config_with(
133        registry: &crate::provider::ProviderRegistry,
134        name: impl Into<String>,
135        config: ModelConfig,
136    ) -> Result<Self, crate::AgentBuildError> {
137        let provider = registry
138            .resolve(&config.api)
139            .ok_or(crate::AgentBuildError::NoProviderForProtocol(config.api))?;
140        Ok(Self::build(name, provider).configured_for(config))
141    }
142
143    /// Create a sub-agent from a name, explicit provider, and [`ModelConfig`].
144    ///
145    /// The escape hatch for custom providers and test doubles (pair with
146    /// [`ModelConfig::mock`](crate::provider::ModelConfig::mock)). Mirrors
147    /// [`Agent::from_provider`](crate::Agent::from_provider).
148    pub fn from_provider(
149        name: impl Into<String>,
150        provider: Arc<dyn StreamProvider>,
151        config: ModelConfig,
152    ) -> Self {
153        Self::build(name, provider).configured_for(config)
154    }
155
156    /// Set the model id and stash the config on a freshly-constructed
157    /// sub-agent (provider already wired).
158    fn configured_for(mut self, config: ModelConfig) -> Self {
159        self.model = config.id.clone();
160        self.model_config = Some(config);
161        self
162    }
163
164    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
165        self.tool_description = desc.into();
166        self
167    }
168
169    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
170        self.system_prompt = prompt.into();
171        self
172    }
173
174    /// Attach a skill set so the sub-agent sees the skills index.
175    ///
176    /// Mirrors [`Agent::with_skills`](crate::agent::Agent::with_skills): the skills
177    /// index is formatted as XML per the [AgentSkills standard](https://agentskills.io)
178    /// and appended to the sub-agent's system prompt at dispatch time. The sub-agent
179    /// can then read individual SKILL.md files via the `read_file` tool when it
180    /// decides a skill is relevant (make sure the sub-agent has such a tool).
181    pub fn with_skills(mut self, skills: crate::skills::SkillSet) -> Self {
182        self.skills_prompt = skills.format_for_prompt();
183        self
184    }
185
186    #[deprecated(
187        since = "0.10.0",
188        note = "the model id now comes from the ModelConfig passed to \
189                SubAgentTool::from_config / from_provider; will be removed in 1.0"
190    )]
191    pub fn with_model(mut self, model: impl Into<String>) -> Self {
192        self.model = model.into();
193        self
194    }
195
196    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
197        self.api_key = key.into();
198        self
199    }
200
201    pub fn with_tools(mut self, tools: Vec<Arc<dyn AgentTool>>) -> Self {
202        self.tools = tools;
203        self
204    }
205
206    /// Add a tool middleware for the sub-agent's own tool calls. Mirrors
207    /// [`Agent::with_tool_middleware`](crate::Agent::with_tool_middleware).
208    pub fn with_tool_middleware(mut self, middleware: impl ToolMiddleware + 'static) -> Self {
209        self.tool_middleware.push(Arc::new(middleware));
210        self
211    }
212
213    pub fn with_thinking(mut self, level: ThinkingLevel) -> Self {
214        self.thinking_level = level;
215        self
216    }
217
218    pub fn with_max_tokens(mut self, max: u32) -> Self {
219        self.max_tokens = Some(max);
220        self
221    }
222
223    /// Set the sampling temperature for the sub-agent. Note: the newest
224    /// reasoning models (e.g. Claude Fable 5 / Opus 4.7+) reject sampling
225    /// parameters — leave unset for those.
226    pub fn with_temperature(mut self, temperature: f32) -> Self {
227        self.temperature = Some(temperature);
228        self
229    }
230
231    pub fn with_cache_config(mut self, config: CacheConfig) -> Self {
232        self.cache_config = config;
233        self
234    }
235
236    pub fn with_tool_execution(mut self, strategy: ToolExecutionStrategy) -> Self {
237        self.tool_execution = strategy;
238        self
239    }
240
241    pub fn with_retry_config(mut self, config: crate::retry::RetryConfig) -> Self {
242        self.retry_config = config;
243        self
244    }
245
246    pub fn with_max_turns(mut self, max: usize) -> Self {
247        self.max_turns = max;
248        self
249    }
250
251    /// Attach a shared key-value store. Sub-agents get a `shared_state` tool
252    /// to read/write variables. The parent can also read/write programmatically
253    /// via the `SharedState` handle.
254    pub fn with_shared_state(mut self, state: SharedState) -> Self {
255        self.shared_state = Some(state);
256        self
257    }
258
259    /// Add an inter-turn delay to throttle API requests.
260    /// Useful when using OAuth tokens or providers with low rate limits.
261    /// The delay is applied before each turn except the first.
262    pub fn with_turn_delay(mut self, delay: std::time::Duration) -> Self {
263        self.turn_delay = Some(delay);
264        self
265    }
266
267    /// Set the model configuration for multi-provider support.
268    /// Required for non-Anthropic providers (OpenAI-compat, Google, etc.)
269    /// to specify base URL, compat flags, and other provider-specific settings.
270    #[deprecated(
271        since = "0.10.0",
272        note = "pass the ModelConfig to SubAgentTool::from_config(name, config) or \
273                from_provider(name, provider, config) instead; will be removed in 1.0"
274    )]
275    pub fn with_model_config(mut self, config: ModelConfig) -> Self {
276        self.model_config = Some(config);
277        self
278    }
279}
280
281/// Thin adapter: wraps `Arc<dyn AgentTool>` so it can be placed in a
282/// `Vec<Box<dyn AgentTool>>` (required by `AgentContext`).
283struct ArcToolWrapper(Arc<dyn AgentTool>);
284
285#[async_trait::async_trait]
286impl AgentTool for ArcToolWrapper {
287    fn name(&self) -> &str {
288        self.0.name()
289    }
290    fn label(&self) -> &str {
291        self.0.label()
292    }
293    fn description(&self) -> &str {
294        self.0.description()
295    }
296    fn parameters_schema(&self) -> serde_json::Value {
297        self.0.parameters_schema()
298    }
299    async fn execute(
300        &self,
301        params: serde_json::Value,
302        ctx: ToolContext,
303    ) -> Result<ToolResult, ToolError> {
304        self.0.execute(params, ctx).await
305    }
306}
307
308#[async_trait::async_trait]
309impl AgentTool for SubAgentTool {
310    fn name(&self) -> &str {
311        &self.tool_name
312    }
313
314    fn label(&self) -> &str {
315        &self.tool_name
316    }
317
318    fn description(&self) -> &str {
319        &self.tool_description
320    }
321
322    fn parameters_schema(&self) -> serde_json::Value {
323        serde_json::json!({
324            "type": "object",
325            "properties": {
326                "task": {
327                    "type": "string",
328                    "description": "The task to delegate to this sub-agent"
329                }
330            },
331            "required": ["task"]
332        })
333    }
334
335    async fn execute(
336        &self,
337        params: serde_json::Value,
338        ctx: ToolContext,
339    ) -> Result<ToolResult, ToolError> {
340        let cancel = ctx.cancel;
341        let on_update = ctx.on_update;
342        let on_progress = ctx.on_progress;
343        // Extract the task parameter
344        let task = params
345            .get("task")
346            .and_then(|v| v.as_str())
347            .ok_or_else(|| ToolError::InvalidArgs("Missing required 'task' parameter".into()))?
348            .to_string();
349
350        // Build tool list from Arc wrappers
351        let mut tools: Vec<Box<dyn AgentTool>> = self
352            .tools
353            .iter()
354            .map(|t| Box::new(ArcToolWrapper(Arc::clone(t))) as Box<dyn AgentTool>)
355            .collect();
356
357        // Append the skills index (if any) so the sub-agent can discover skills.
358        let mut system_prompt = self.system_prompt.clone();
359        if !self.skills_prompt.is_empty() {
360            if system_prompt.is_empty() {
361                system_prompt = self.skills_prompt.clone();
362            } else {
363                system_prompt = format!("{}\n\n{}", system_prompt, self.skills_prompt);
364            }
365        }
366
367        // Inject SharedStateTool when shared state is configured
368        if let Some(ref state) = self.shared_state {
369            tools.push(Box::new(SharedStateTool::new(state.clone())));
370            let summary = state.summary().await;
371            system_prompt.push_str(&format!(
372                "\n\n## Shared State\nYou have access to a shared variable store via the `shared_state` tool.\nAvailable: {}",
373                summary
374            ));
375        }
376
377        // Fresh context for the sub-agent
378        let mut context = AgentContext {
379            system_prompt,
380            messages: Vec::new(),
381            tools,
382        };
383
384        // Config with Arc'd provider
385        let config = AgentLoopConfig {
386            provider: self.provider.clone(),
387            model: self.model.clone(),
388            api_key: if self.api_key.is_empty() {
389                crate::provider::resolve_api_key_or_warn(
390                    self.model_config
391                        .as_ref()
392                        .map(|m| m.provider.as_str())
393                        .unwrap_or("anthropic"),
394                )
395            } else {
396                self.api_key.clone()
397            },
398            thinking_level: self.thinking_level,
399            max_tokens: self.max_tokens,
400            temperature: self.temperature,
401            model_config: self.model_config.clone(),
402            convert_to_llm: None,
403            transform_context: None,
404            get_steering_messages: None,
405            get_follow_up_messages: None,
406            context_config: None,
407            compaction_strategy: None,
408            execution_limits: Some(ExecutionLimits {
409                max_turns: self.max_turns,
410                // Generous token/duration limits — turn limit is the primary guard
411                max_total_tokens: 1_000_000,
412                max_duration: std::time::Duration::from_secs(300),
413            }),
414            cache_config: self.cache_config.clone(),
415            tool_execution: self.tool_execution.clone(),
416            retry_config: self.retry_config.clone(),
417            before_turn: None,
418            after_turn: None,
419            on_error: None,
420            input_filters: vec![],
421            tool_middleware: self.tool_middleware.clone(),
422            output_schema: None,
423            turn_delay: self.turn_delay,
424        };
425
426        // Channel for sub-agent events
427        let (tx, mut rx) = mpsc::unbounded_channel();
428
429        // Forward sub-agent events to parent via on_update and on_progress callbacks
430        let forward_handle = if on_update.is_some() || on_progress.is_some() {
431            let tool_name = self.tool_name.clone();
432            Some(tokio::spawn(async move {
433                while let Some(event) = rx.recv().await {
434                    // Forward progress messages via on_progress
435                    if let AgentEvent::ProgressMessage { text, .. } = &event {
436                        if let Some(ref cb) = on_progress {
437                            cb(text.clone());
438                        }
439                    }
440
441                    // Convert interesting events to ToolResult updates for the parent
442                    if let Some(ref on_update) = on_update {
443                        let update_text = match &event {
444                            AgentEvent::MessageUpdate {
445                                delta: StreamDelta::Text { delta },
446                                ..
447                            } => Some(delta.clone()),
448                            AgentEvent::ToolExecutionStart { tool_name, .. } => {
449                                Some(format!("[sub-agent calling tool: {}]", tool_name))
450                            }
451                            _ => None,
452                        };
453
454                        if let Some(text) = update_text {
455                            on_update(ToolResult {
456                                content: vec![Content::Text { text }],
457                                details: serde_json::json!({ "sub_agent": tool_name }),
458                            });
459                        }
460                    }
461                }
462            }))
463        } else {
464            None
465        };
466
467        // Run the sub-agent loop
468        let prompt = AgentMessage::Llm(Message::user(task));
469        let new_messages = agent_loop(vec![prompt], &mut context, &config, tx, cancel).await;
470
471        // Wait for event forwarding to complete
472        if let Some(handle) = forward_handle {
473            let _ = handle.await;
474        }
475
476        // Check if the last message was an error
477        if let Some(error_msg) = extract_error(&new_messages) {
478            return Err(ToolError::Failed(format!(
479                "Sub-agent '{}' failed: {}",
480                self.tool_name, error_msg
481            )));
482        }
483
484        // Extract final assistant text from the returned messages
485        let result_text = extract_final_text(&new_messages);
486
487        // Include full sub-agent conversation in details for debugging
488        let details = serde_json::json!({
489            "sub_agent": self.tool_name,
490            "turns": new_messages.len(),
491        });
492
493        Ok(ToolResult {
494            content: vec![Content::Text { text: result_text }],
495            details,
496        })
497    }
498}
499
500/// Check if the last assistant message was an error, return the error message.
501fn extract_error(messages: &[AgentMessage]) -> Option<String> {
502    for msg in messages.iter().rev() {
503        if let AgentMessage::Llm(Message::Assistant {
504            stop_reason,
505            error_message,
506            ..
507        }) = msg
508        {
509            if *stop_reason == StopReason::Error {
510                return Some(
511                    error_message
512                        .clone()
513                        .unwrap_or_else(|| "Unknown error".into()),
514                );
515            }
516        }
517    }
518    None
519}
520
521/// Extract the final assistant text from agent messages.
522/// Collects text from the last assistant message, or returns a fallback.
523fn extract_final_text(messages: &[AgentMessage]) -> String {
524    for msg in messages.iter().rev() {
525        if let AgentMessage::Llm(Message::Assistant { content, .. }) = msg {
526            let texts: Vec<&str> = content
527                .iter()
528                .filter_map(|c| match c {
529                    Content::Text { text } if !text.is_empty() => Some(text.as_str()),
530                    _ => None,
531                })
532                .collect();
533            if !texts.is_empty() {
534                return texts.join("\n");
535            }
536        }
537    }
538    "(sub-agent produced no text output)".to_string()
539}
yoagent/sub_agent.rs

yoagent/
sub_agent.rs
