gemini_adk_rs/tool/
mod.rs

1//! Tool dispatch — regular, streaming, and input-streaming tools.
2
3pub mod dispatcher;
4pub mod policy;
5pub mod simple;
6pub mod typed;
7
8pub use dispatcher::*;
9pub use policy::*;
10pub use simple::*;
11pub use typed::*;
12
13use std::sync::Arc;
14use std::time::Duration;
15
16pub mod media;
17
18use async_trait::async_trait;
19use tokio::sync::{broadcast, mpsc};
20use tokio::task::JoinHandle;
21use tokio_util::sync::CancellationToken;
22
23use crate::agent_session::InputEvent;
24use crate::error::ToolError;
25
26/// A regular tool — called once, returns a result.
27///
28/// # Examples
29///
30/// ```rust,ignore
31/// use async_trait::async_trait;
32/// use gemini_adk_rs::tool::ToolFunction;
33/// use gemini_adk_rs::error::ToolError;
34///
35/// struct MyTool;
36///
37/// #[async_trait]
38/// impl ToolFunction for MyTool {
39///     fn name(&self) -> &str { "my_tool" }
40///     fn description(&self) -> &str { "Does something useful" }
41///     fn parameters(&self) -> Option<serde_json::Value> { None }
42///     async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
43///         Ok(serde_json::json!({"status": "ok"}))
44///     }
45/// }
46/// ```
47#[async_trait]
48pub trait ToolFunction: Send + Sync + 'static {
49    /// The unique name of this tool.
50    fn name(&self) -> &str;
51    /// Human-readable description of what this tool does.
52    fn description(&self) -> &str;
53    /// JSON Schema for the tool's input parameters, or `None` if parameterless.
54    fn parameters(&self) -> Option<serde_json::Value>;
55    /// Execute the tool with the given arguments and return the result.
56    async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError>;
57
58    /// Whether this tool must be confirmed before it runs. Defaults to `false`.
59    ///
60    /// Tools built with `T::confirm(..)` (a [`PolicyTool`] with a confirm
61    /// policy) return `true`, prompting the [`ToolDispatcher`] to consult its
62    /// configured `ConfirmationProvider` before executing.
63    fn requires_confirmation(&self) -> bool {
64        false
65    }
66
67    /// Optional hint shown when confirmation is requested.
68    fn confirmation_message(&self) -> Option<&str> {
69        None
70    }
71}
72
73/// An `Arc<dyn ToolFunction>` (or `Arc<SomeTool>`) is itself a tool, so a
74/// setter that takes `impl ToolFunction` accepts a shared handle as readily as
75/// a fresh value.
76#[async_trait]
77impl<T: ToolFunction + ?Sized> ToolFunction for Arc<T> {
78    fn name(&self) -> &str {
79        (**self).name()
80    }
81    fn description(&self) -> &str {
82        (**self).description()
83    }
84    fn parameters(&self) -> Option<serde_json::Value> {
85        (**self).parameters()
86    }
87    async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
88        (**self).call(args).await
89    }
90    fn requires_confirmation(&self) -> bool {
91        (**self).requires_confirmation()
92    }
93    fn confirmation_message(&self) -> Option<&str> {
94        (**self).confirmation_message()
95    }
96}
97
98/// A streaming tool — runs in background, yields multiple results.
99#[async_trait]
100pub trait StreamingTool: Send + Sync + 'static {
101    /// The unique name of this tool.
102    fn name(&self) -> &str;
103    /// Human-readable description of what this tool does.
104    fn description(&self) -> &str;
105    /// JSON Schema for the tool's input parameters, or `None` if parameterless.
106    fn parameters(&self) -> Option<serde_json::Value>;
107    /// Execute the tool, sending intermediate results via `yield_tx`.
108    async fn run(
109        &self,
110        args: serde_json::Value,
111        yield_tx: mpsc::Sender<serde_json::Value>,
112    ) -> Result<(), ToolError>;
113}
114
115/// An input-streaming tool — receives duplicated live input while running.
116#[async_trait]
117pub trait InputStreamingTool: Send + Sync + 'static {
118    /// The unique name of this tool.
119    fn name(&self) -> &str;
120    /// Human-readable description of what this tool does.
121    fn description(&self) -> &str;
122    /// JSON Schema for the tool's input parameters, or `None` if parameterless.
123    fn parameters(&self) -> Option<serde_json::Value>;
124    /// Execute the tool, receiving live input via `input_rx` and sending results via `yield_tx`.
125    async fn run(
126        &self,
127        args: serde_json::Value,
128        input_rx: broadcast::Receiver<InputEvent>,
129        yield_tx: mpsc::Sender<serde_json::Value>,
130    ) -> Result<(), ToolError>;
131}
132
133/// Classification of a registered tool.
134#[derive(Debug, Clone, Copy, PartialEq, Eq)]
135pub enum ToolClass {
136    /// A one-shot tool that returns a single result.
137    Regular,
138    /// A tool that yields multiple results over time.
139    Streaming,
140    /// A tool that receives live input while producing results.
141    InputStream,
142}
143
144/// Unified tool storage.
145pub enum ToolKind {
146    /// A regular one-shot function tool.
147    Function(Arc<dyn ToolFunction>),
148    /// A streaming tool that yields multiple results.
149    Streaming(Arc<dyn StreamingTool>),
150    /// An input-streaming tool that receives live input.
151    InputStream(Arc<dyn InputStreamingTool>),
152}
153
154/// Handle to a running streaming tool.
155pub struct ActiveStreamingTool {
156    /// The spawned task handle.
157    pub task: JoinHandle<()>,
158    /// Token to cancel this streaming tool.
159    pub cancel: CancellationToken,
160}
161
162/// Default timeout for tool execution (30 seconds).
163pub(crate) const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(30);
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use gemini_genai_rs::prelude::FunctionCall;
169    use serde_json::json;
170
171    struct MockTool;
172
173    #[async_trait]
174    impl ToolFunction for MockTool {
175        fn name(&self) -> &str {
176            "mock_tool"
177        }
178        fn description(&self) -> &str {
179            "A mock tool"
180        }
181        fn parameters(&self) -> Option<serde_json::Value> {
182            None
183        }
184        async fn call(&self, _args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
185            Ok(json!({"result": "ok"}))
186        }
187    }
188
189    #[tokio::test]
190    async fn register_and_call_function_tool() {
191        let mut dispatcher = ToolDispatcher::new();
192        dispatcher.register_function(Arc::new(MockTool));
193        let result = dispatcher
194            .call_function("mock_tool", json!({}))
195            .await
196            .unwrap();
197        assert_eq!(result["result"], "ok");
198    }
199
200    #[tokio::test]
201    async fn call_unknown_tool_returns_error() {
202        let dispatcher = ToolDispatcher::new();
203        let result = dispatcher.call_function("nonexistent", json!({})).await;
204        assert!(result.is_err());
205    }
206
207    #[test]
208    fn to_tool_declarations() {
209        let mut dispatcher = ToolDispatcher::new();
210        dispatcher.register_function(Arc::new(MockTool));
211        let decls = dispatcher.to_tool_declarations();
212        assert_eq!(decls.len(), 1);
213    }
214
215    #[test]
216    fn classify_tool() {
217        let mut dispatcher = ToolDispatcher::new();
218        dispatcher.register_function(Arc::new(MockTool));
219        assert_eq!(dispatcher.classify("mock_tool"), Some(ToolClass::Regular));
220        assert_eq!(dispatcher.classify("nonexistent"), None);
221    }
222
223    #[test]
224    fn empty_dispatcher() {
225        let dispatcher = ToolDispatcher::new();
226        assert!(dispatcher.is_empty());
227        assert_eq!(dispatcher.len(), 0);
228        assert!(dispatcher.to_tool_declarations().is_empty());
229    }
230
231    #[test]
232    fn build_response_success() {
233        let call = FunctionCall {
234            name: "test".to_string(),
235            args: json!({}),
236            id: Some("call-1".to_string()),
237        };
238        let resp = ToolDispatcher::build_response(&call, Ok(json!({"ok": true})));
239        assert_eq!(resp.name, "test");
240        assert_eq!(resp.response["ok"], true);
241    }
242
243    #[test]
244    fn build_response_error() {
245        let call = FunctionCall {
246            name: "test".to_string(),
247            args: json!({}),
248            id: Some("call-1".to_string()),
249        };
250        let resp = ToolDispatcher::build_response(
251            &call,
252            Err(ToolError::ExecutionFailed("boom".to_string())),
253        );
254        assert!(resp.response["error"].as_str().unwrap().contains("boom"));
255    }
256
257    #[test]
258    fn tool_dispatcher_implements_tool_provider() {
259        use gemini_genai_rs::prelude::ToolProvider;
260        let mut dispatcher = ToolDispatcher::new();
261        dispatcher.register_function(Arc::new(MockTool));
262        let decls = dispatcher.declarations();
263        assert_eq!(decls.len(), 1);
264    }
265
266    #[tokio::test]
267    async fn simple_tool_closure() {
268        let tool = SimpleTool::new(
269            "add",
270            "Add two numbers",
271            Some(
272                json!({"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}),
273            ),
274            |args| async move {
275                let a = args["a"].as_f64().unwrap_or(0.0);
276                let b = args["b"].as_f64().unwrap_or(0.0);
277                Ok(json!({"sum": a + b}))
278            },
279        );
280
281        let mut dispatcher = ToolDispatcher::new();
282        dispatcher.register_function(Arc::new(tool));
283        let result = dispatcher
284            .call_function("add", json!({"a": 3, "b": 4}))
285            .await
286            .unwrap();
287        assert_eq!(result["sum"], 7.0);
288    }
289
290    // --- TypedTool tests ---
291
292    #[derive(serde::Deserialize, schemars::JsonSchema)]
293    struct WeatherArgs {
294        /// The city to get weather for
295        city: String,
296        /// Temperature units (celsius or fahrenheit)
297        #[serde(default = "default_units")]
298        units: String,
299    }
300
301    fn default_units() -> String {
302        "celsius".to_string()
303    }
304
305    #[test]
306    fn typed_tool_auto_generates_schema() {
307        let tool = TypedTool::new(
308            "get_weather",
309            "Get current weather for a city",
310            |_args: WeatherArgs| async move { Ok(json!({})) },
311        );
312
313        let params = tool.parameters().expect("should have parameters");
314
315        // The schema should be an object type with "city" and "units" properties
316        let props = &params["properties"];
317        assert!(
318            props.get("city").is_some(),
319            "schema should contain 'city' property"
320        );
321        assert!(
322            props.get("units").is_some(),
323            "schema should contain 'units' property"
324        );
325
326        // "city" should be required (no default), "units" has a default so may not be
327        let required = params["required"]
328            .as_array()
329            .expect("should have required array");
330        let required_names: Vec<&str> = required.iter().filter_map(|v| v.as_str()).collect();
331        assert!(required_names.contains(&"city"), "city should be required");
332    }
333
334    #[tokio::test]
335    async fn typed_tool_deserializes_args() {
336        let tool = TypedTool::new(
337            "get_weather",
338            "Get current weather for a city",
339            |args: WeatherArgs| async move {
340                Ok(json!({
341                    "temp": 22,
342                    "city": args.city,
343                    "units": args.units,
344                }))
345            },
346        );
347
348        let result = tool
349            .call(json!({"city": "London", "units": "fahrenheit"}))
350            .await
351            .unwrap();
352        assert_eq!(result["city"], "London");
353        assert_eq!(result["units"], "fahrenheit");
354        assert_eq!(result["temp"], 22);
355    }
356
357    #[tokio::test]
358    async fn typed_tool_invalid_args_returns_error() {
359        let tool = TypedTool::new(
360            "get_weather",
361            "Get current weather for a city",
362            |_args: WeatherArgs| async move { Ok(json!({})) },
363        );
364
365        // Missing required field "city"
366        let result = tool.call(json!({"units": "celsius"})).await;
367        assert!(result.is_err(), "should fail with missing required field");
368        let err = result.unwrap_err();
369        match &err {
370            ToolError::InvalidArgs(msg) => {
371                assert!(
372                    msg.contains("city"),
373                    "error message should mention the missing field: {msg}"
374                );
375            }
376            other => panic!("expected ToolError::InvalidArgs, got: {other:?}"),
377        }
378
379        // Wrong type for "city" (number instead of string)
380        let result = tool.call(json!({"city": 12345})).await;
381        assert!(result.is_err(), "should fail with wrong type");
382    }
383
384    #[tokio::test]
385    async fn typed_tool_registers_in_dispatcher() {
386        let tool = TypedTool::new(
387            "get_weather",
388            "Get current weather for a city",
389            |args: WeatherArgs| async move { Ok(json!({"city": args.city})) },
390        );
391
392        let mut dispatcher = ToolDispatcher::new();
393        dispatcher.register_function(Arc::new(tool));
394
395        assert_eq!(dispatcher.classify("get_weather"), Some(ToolClass::Regular));
396        assert_eq!(dispatcher.len(), 1);
397
398        let result = dispatcher
399            .call_function("get_weather", json!({"city": "Paris"}))
400            .await
401            .unwrap();
402        assert_eq!(result["city"], "Paris");
403
404        // Verify it appears in tool declarations
405        let decls = dispatcher.to_tool_declarations();
406        assert_eq!(decls.len(), 1);
407    }
408
409    // --- Timeout and cancellation tests ---
410
411    /// A tool that sleeps forever (until cancelled/timed out).
412    struct SlowTool;
413
414    #[async_trait]
415    impl ToolFunction for SlowTool {
416        fn name(&self) -> &str {
417            "slow_tool"
418        }
419        fn description(&self) -> &str {
420            "A tool that never completes"
421        }
422        fn parameters(&self) -> Option<serde_json::Value> {
423            None
424        }
425        async fn call(&self, _args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
426            // Sleep effectively forever
427            tokio::time::sleep(Duration::from_secs(3600)).await;
428            Ok(json!({"result": "should never reach here"}))
429        }
430    }
431
432    #[tokio::test]
433    async fn tool_timeout_returns_error() {
434        let mut dispatcher = ToolDispatcher::new();
435        dispatcher.register_function(Arc::new(SlowTool));
436
437        let timeout = Duration::from_millis(50);
438        let result = dispatcher
439            .call_function_with_timeout("slow_tool", json!({}), timeout)
440            .await;
441
442        match result {
443            Err(ToolError::Timeout(d)) => assert_eq!(d, timeout),
444            other => panic!("expected ToolError::Timeout, got: {other:?}"),
445        }
446    }
447
448    #[tokio::test]
449    async fn tool_completes_before_timeout() {
450        let mut dispatcher = ToolDispatcher::new();
451        dispatcher.register_function(Arc::new(MockTool));
452
453        let result = dispatcher
454            .call_function_with_timeout("mock_tool", json!({}), Duration::from_secs(5))
455            .await
456            .unwrap();
457        assert_eq!(result["result"], "ok");
458    }
459
460    #[tokio::test]
461    async fn tool_cancelled_returns_error() {
462        let mut dispatcher = ToolDispatcher::new();
463        dispatcher.register_function(Arc::new(SlowTool));
464
465        let cancel = CancellationToken::new();
466        let cancel_clone = cancel.clone();
467
468        // Cancel after a short delay
469        tokio::spawn(async move {
470            tokio::time::sleep(Duration::from_millis(50)).await;
471            cancel_clone.cancel();
472        });
473
474        let result = dispatcher
475            .call_function_with_cancel("slow_tool", json!({}), cancel)
476            .await;
477
478        match result {
479            Err(ToolError::Cancelled) => {} // expected
480            other => panic!("expected ToolError::Cancelled, got: {other:?}"),
481        }
482    }
483
484    #[test]
485    fn default_timeout_is_30s() {
486        let dispatcher = ToolDispatcher::new();
487        assert_eq!(dispatcher.default_timeout(), Duration::from_secs(30));
488    }
489
490    #[test]
491    fn with_timeout_overrides_default() {
492        let dispatcher = ToolDispatcher::new().with_timeout(Duration::from_secs(10));
493        assert_eq!(dispatcher.default_timeout(), Duration::from_secs(10));
494    }
495
496    #[tokio::test]
497    async fn call_function_uses_default_timeout() {
498        // Set a very short default timeout so the slow tool times out
499        let mut dispatcher = ToolDispatcher::new().with_timeout(Duration::from_millis(50));
500        dispatcher.register_function(Arc::new(SlowTool));
501
502        let result = dispatcher.call_function("slow_tool", json!({})).await;
503
504        match result {
505            Err(ToolError::Timeout(d)) => assert_eq!(d, Duration::from_millis(50)),
506            other => panic!("expected ToolError::Timeout, got: {other:?}"),
507        }
508    }
509}
gemini_adk_rs/tool/mod.rs

gemini_adk_rs/tool/
mod.rs
