gemini_adk_rs/
agent_tool.rs

1//! AgentTool — wraps an Agent as a ToolFunction for "agent as a tool" dispatch.
2//!
3//! When the live model calls this tool, the wrapped agent runs in an isolated
4//! context (no live WebSocket). The agent's text output is collected and returned
5//! as the tool result. State changes propagate back to the parent context.
6//!
7//! This bridges live<->non-live: the wrapped agent can use regular Gemini API,
8//! external services, or pure computation — it doesn't need a WebSocket.
9
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use serde_json::json;
14use tokio::sync::broadcast;
15
16use gemini_genai_rs::session::SessionEvent;
17
18use crate::agent::Agent;
19use crate::agent_session::{AgentSession, NoOpSessionWriter};
20use crate::context::{AgentEvent, InvocationContext};
21use crate::error::ToolError;
22use crate::tool::ToolFunction;
23
24/// Wraps an Agent as a ToolFunction for "agent as a tool" dispatch.
25///
26/// When the live model calls this tool, the wrapped agent runs in an isolated
27/// context (no live WebSocket). The agent's text output is collected and returned
28/// as the tool result.
29pub struct AgentTool {
30    agent: Arc<dyn Agent>,
31    description: String,
32    parameters: Option<serde_json::Value>,
33}
34
35impl AgentTool {
36    /// Create a new AgentTool wrapping the given agent.
37    pub fn new(agent: impl Agent + 'static) -> Self {
38        let description = format!("Delegate to the {} agent", agent.name());
39        Self {
40            agent: Arc::new(agent),
41            description,
42            parameters: Some(json!({
43                "type": "object",
44                "properties": {
45                    "request": {
46                        "type": "string",
47                        "description": "The request to send to the agent"
48                    }
49                },
50                "required": ["request"]
51            })),
52        }
53    }
54
55    /// Create from an already-Arc'd agent.
56    pub fn from_arc(agent: Arc<dyn Agent>) -> Self {
57        let description = format!("Delegate to the {} agent", agent.name());
58        Self {
59            agent,
60            description,
61            parameters: Some(json!({
62                "type": "object",
63                "properties": {
64                    "request": {
65                        "type": "string",
66                        "description": "The request to send to the agent"
67                    }
68                },
69                "required": ["request"]
70            })),
71        }
72    }
73
74    /// Override the tool description.
75    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
76        self.description = desc.into();
77        self
78    }
79
80    /// Override the tool parameters schema.
81    pub fn with_parameters(mut self, params: serde_json::Value) -> Self {
82        self.parameters = Some(params);
83        self
84    }
85}
86
87#[async_trait]
88impl ToolFunction for AgentTool {
89    fn name(&self) -> &str {
90        self.agent.name()
91    }
92
93    fn description(&self) -> &str {
94        &self.description
95    }
96
97    fn parameters(&self) -> Option<serde_json::Value> {
98        self.parameters.clone()
99    }
100
101    async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
102        let start = std::time::Instant::now();
103        let agent_name = self.agent.name().to_string();
104
105        // Telemetry
106        crate::telemetry::logging::log_agent_tool_dispatch("parent", &agent_name);
107
108        // 1. Create isolated context with NoOpSessionWriter
109        let (event_tx, _) = broadcast::channel::<SessionEvent>(64);
110        let noop_writer: Arc<dyn gemini_genai_rs::session::SessionWriter> =
111            Arc::new(NoOpSessionWriter);
112        let isolated_session = AgentSession::from_writer(noop_writer, event_tx);
113
114        // 2. Inject args into state
115        if let Some(request) = args.get("request").and_then(|r| r.as_str()) {
116            let _ = isolated_session.state().set("request_text", request);
117        }
118        let _ = isolated_session.state().set("request", &args);
119
120        // 3. Create isolated InvocationContext
121        let mut ctx = InvocationContext::new(isolated_session);
122
123        // 4. Subscribe to events before running (to collect text output)
124        let mut events = ctx.subscribe();
125
126        // 5. Run the agent
127        let agent = self.agent.clone();
128        let run_result = tokio::spawn(async move { agent.run_live(&mut ctx).await }).await;
129
130        // 6. Collect text output from events
131        let mut output_parts = Vec::new();
132        while let Ok(event) = events.try_recv() {
133            match event {
134                AgentEvent::Session(SessionEvent::TextDelta(text)) => {
135                    output_parts.push(text);
136                }
137                AgentEvent::Session(SessionEvent::TextComplete(text)) => {
138                    if output_parts.is_empty() {
139                        output_parts.push(text);
140                    }
141                    // If we already have deltas, TextComplete is the full assembled text
142                    // Don't double-count — deltas already captured incrementally
143                }
144                _ => {}
145            }
146        }
147
148        let elapsed = start.elapsed();
149        crate::telemetry::metrics::record_agent_tool_dispatch(
150            "parent",
151            &agent_name,
152            elapsed.as_millis() as f64,
153        );
154
155        // 7. Handle result
156        match run_result {
157            Ok(Ok(())) => {
158                let output = if output_parts.is_empty() {
159                    json!({"status": "completed"})
160                } else {
161                    json!({"result": output_parts.join("")})
162                };
163                Ok(output)
164            }
165            Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!(
166                "Agent '{agent_name}' failed: {e}"
167            ))),
168            Err(e) => Err(ToolError::ExecutionFailed(format!(
169                "Agent '{agent_name}' task panicked: {e}"
170            ))),
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use crate::error::AgentError;
179
180    struct EchoAgent {
181        name: String,
182    }
183
184    #[async_trait]
185    impl Agent for EchoAgent {
186        fn name(&self) -> &str {
187            &self.name
188        }
189        async fn run_live(&self, ctx: &mut InvocationContext) -> Result<(), AgentError> {
190            // Read the request from state and echo it back as a text event
191            let request = ctx
192                .state()
193                .get::<String>("request_text")
194                .unwrap_or_else(|| "no request".to_string());
195            ctx.emit(AgentEvent::Session(SessionEvent::TextDelta(format!(
196                "Echo: {request}"
197            ))));
198            ctx.emit(AgentEvent::Session(SessionEvent::TurnComplete));
199            Ok(())
200        }
201    }
202
203    struct FailingAgent;
204
205    #[async_trait]
206    impl Agent for FailingAgent {
207        fn name(&self) -> &str {
208            "failing"
209        }
210        async fn run_live(&self, _ctx: &mut InvocationContext) -> Result<(), AgentError> {
211            Err(AgentError::Other("intentional failure".to_string()))
212        }
213    }
214
215    struct SilentAgent;
216
217    #[async_trait]
218    impl Agent for SilentAgent {
219        fn name(&self) -> &str {
220            "silent"
221        }
222        async fn run_live(&self, _ctx: &mut InvocationContext) -> Result<(), AgentError> {
223            Ok(())
224        }
225    }
226
227    #[tokio::test]
228    async fn agent_tool_runs_agent_in_isolation() {
229        let agent = EchoAgent {
230            name: "echo".to_string(),
231        };
232        let tool = AgentTool::new(agent);
233
234        assert_eq!(tool.name(), "echo");
235        assert!(tool.description().contains("echo"));
236    }
237
238    #[tokio::test]
239    async fn agent_tool_collects_text_output() {
240        let agent = EchoAgent {
241            name: "echo".to_string(),
242        };
243        let tool = AgentTool::new(agent);
244
245        let result = tool.call(json!({"request": "hello world"})).await.unwrap();
246        assert_eq!(result["result"], "Echo: hello world");
247    }
248
249    #[tokio::test]
250    async fn agent_tool_propagates_errors() {
251        let tool = AgentTool::new(FailingAgent);
252        let result = tool.call(json!({"request": "test"})).await;
253        assert!(result.is_err());
254        let err = result.unwrap_err();
255        match err {
256            ToolError::ExecutionFailed(msg) => {
257                assert!(msg.contains("intentional failure"));
258            }
259            other => panic!("expected ExecutionFailed, got: {other:?}"),
260        }
261    }
262
263    #[tokio::test]
264    async fn agent_tool_returns_completed_when_no_output() {
265        let tool = AgentTool::new(SilentAgent);
266        let result = tool.call(json!({"request": "test"})).await.unwrap();
267        assert_eq!(result["status"], "completed");
268    }
269
270    #[tokio::test]
271    async fn agent_tool_state_injection() {
272        // Verify that args are injected into state
273        struct StateCheckAgent;
274
275        #[async_trait]
276        impl Agent for StateCheckAgent {
277            fn name(&self) -> &str {
278                "state_check"
279            }
280            async fn run_live(&self, ctx: &mut InvocationContext) -> Result<(), AgentError> {
281                let request_text = ctx.state().get::<String>("request_text");
282                let request = ctx.state().get::<serde_json::Value>("request");
283
284                assert!(request_text.is_some());
285                assert!(request.is_some());
286                assert_eq!(request_text.unwrap(), "check state");
287
288                ctx.emit(AgentEvent::Session(SessionEvent::TextDelta(
289                    "state ok".to_string(),
290                )));
291                Ok(())
292            }
293        }
294
295        let tool = AgentTool::new(StateCheckAgent);
296        let result = tool.call(json!({"request": "check state"})).await.unwrap();
297        assert_eq!(result["result"], "state ok");
298    }
299
300    #[tokio::test]
301    async fn agent_tool_with_custom_description() {
302        let tool = AgentTool::new(SilentAgent).with_description("Custom description");
303        assert_eq!(tool.description(), "Custom description");
304    }
305
306    #[tokio::test]
307    async fn agent_tool_with_custom_parameters() {
308        let params = json!({
309            "type": "object",
310            "properties": {
311                "query": { "type": "string" }
312            }
313        });
314        let tool = AgentTool::new(SilentAgent).with_parameters(params.clone());
315        assert_eq!(tool.parameters().unwrap(), params);
316    }
317}
gemini_adk_rs/agent_tool.rs

gemini_adk_rs/
agent_tool.rs
