agent_base/tool/
subagent.rs

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

agent_base/tool/
subagent.rs
