agent_base/tool/
subagent.rs

1use async_trait::async_trait;
2use serde_json::{Value, json};
3use tokio::sync::Mutex;
4
5use super::{Tool, ToolContext, ToolControlFlow, ToolOutput};
6use crate::engine::AgentRuntime;
7use crate::types::{AgentError, AgentResult, RuntimeEvent, SessionId, UserEvent};
8
9/// Sub-Agent session policy
10#[derive(Clone, Debug)]
11pub enum SubAgentSessionPolicy {
12    /// Create a new session per call (default)
13    Ephemeral,
14    /// Reuse the same session; sub-agent accumulates history
15    Persistent,
16}
17
18pub struct SubAgentTool {
19    name: &'static str,
20    description: &'static str,
21    sub_runtime: Mutex<AgentRuntime>,
22    sub_session_id: Mutex<Option<SessionId>>,
23    session_policy: SubAgentSessionPolicy,
24}
25
26impl SubAgentTool {
27    pub fn new(name: &'static str, description: &'static str, sub_runtime: AgentRuntime) -> Self {
28        Self {
29            name,
30            description,
31            sub_runtime: Mutex::new(sub_runtime),
32            sub_session_id: Mutex::new(None),
33            session_policy: SubAgentSessionPolicy::Ephemeral,
34        }
35    }
36
37    pub fn with_persistent(
38        name: &'static str,
39        description: &'static str,
40        sub_runtime: AgentRuntime,
41    ) -> Self {
42        Self {
43            name,
44            description,
45            sub_runtime: Mutex::new(sub_runtime),
46            sub_session_id: Mutex::new(None),
47            session_policy: SubAgentSessionPolicy::Persistent,
48        }
49    }
50}
51
52#[async_trait]
53impl Tool for SubAgentTool {
54    fn name(&self) -> &'static str {
55        self.name
56    }
57
58    fn definition(&self) -> Value {
59        json!({
60            "type": "function",
61            "function": {
62                "name": self.name,
63                "description": self.description,
64                "parameters": {
65                    "type": "object",
66                    "properties": {
67                        "task": {
68                            "type": "string",
69                            "description": "Task description to delegate to the sub-agent"
70                        }
71                    },
72                    "required": ["task"]
73                }
74            }
75        })
76    }
77
78    async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<ToolOutput> {
79        let task = args.get("task").and_then(Value::as_str).ok_or_else(|| {
80            AgentError::ToolArgsInvalid {
81                name: self.name.to_string(),
82                raw: args.to_string(),
83            }
84        })?;
85
86        if task.is_empty() {
87            return Ok(ToolOutput {
88                summary: "Task description is empty, cannot execute".to_string(),
89                raw: None,
90                control_flow: ToolControlFlow::Break,
91                truncation: None,
92            });
93        }
94
95        let user_event_tx = ctx.user_event_tx.clone();
96
97        let sub_session_id = match self.session_policy {
98            SubAgentSessionPolicy::Ephemeral => {
99                let runtime = self.sub_runtime.lock().await;
100                let new_id = runtime.create_session().await;
101                let mut sid_guard = self.sub_session_id.lock().await;
102                *sid_guard = Some(new_id.clone());
103                tracing::debug!(
104                    subagent = self.name,
105                    sub_session = new_id.id,
106                    "sub-agent ephemeral session created"
107                );
108                new_id
109            }
110            SubAgentSessionPolicy::Persistent => {
111                let mut sid_guard = self.sub_session_id.lock().await;
112                if let Some(id) = sid_guard.clone() {
113                    tracing::debug!(
114                        subagent = self.name,
115                        sub_session = id.id,
116                        "sub-agent reusing persistent session"
117                    );
118                    id
119                } else {
120                    let runtime = self.sub_runtime.lock().await;
121                    let new_id = runtime.create_session().await;
122                    *sid_guard = Some(new_id.clone());
123                    tracing::debug!(
124                        subagent = self.name,
125                        sub_session = new_id.id,
126                        "sub-agent persistent session created"
127                    );
128                    new_id
129                }
130            }
131        };
132
133        // Add the user task as a message to the sub-agent session
134        {
135            let runtime = self.sub_runtime.lock().await;
136            runtime
137                .add_user_message(&sub_session_id, task)
138                .await
139                .map_err(|e| AgentError::ToolExecution {
140                    name: self.name.to_string(),
141                    source: Box::new(e),
142                })?;
143        }
144
145        let mut runtime_events = Vec::new();
146        let _outcome = {
147            let runtime = self.sub_runtime.lock().await;
148            runtime
149                .run(sub_session_id, |event| {
150                    runtime_events.push(event.clone());
151                    Ok(())
152                })
153                .await
154                .map_err(|e| AgentError::ToolExecution {
155                    name: self.name.to_string(),
156                    source: Box::new(e),
157                })?
158        };
159
160        let mut final_text = String::new();
161        for event in &runtime_events {
162            match event {
163                RuntimeEvent::TextDelta { text, .. } => {
164                    final_text.push_str(&text);
165                }
166                _ => {}
167            }
168            // Forward each sub-agent event to the parent via UserEvent::SubAgentEvent
169            let _ = user_event_tx.send(UserEvent::SubAgentEvent {
170                subagent: self.name.to_string(),
171                event: Box::new(event.clone()),
172            });
173        }
174
175        let summary = if final_text.is_empty() {
176            format!("Sub-agent [{}] finished", self.name)
177        } else {
178            final_text
179        };
180
181        tracing::debug!(
182            subagent = self.name,
183            text_len = summary.len(),
184            event_count = runtime_events.len(),
185            "sub-agent completed"
186        );
187
188        Ok(ToolOutput {
189            summary,
190            raw: None,
191            control_flow: ToolControlFlow::Continue,
192            truncation: None,
193        })
194    }
195}
agent_base/tool/subagent.rs

agent_base/tool/
subagent.rs
