gemini_adk_rs/tool/
mod.rs

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

gemini_adk_rs/tool/
mod.rs
