gemini_adk_rs/tool/
dispatcher.rs

1//! Tool dispatcher — routes function calls to the right tool implementation.
2
3use std::collections::{BTreeMap, HashMap};
4use std::sync::Arc;
5use std::time::Duration;
6
7use tokio_util::sync::CancellationToken;
8
9use gemini_genai_rs::prelude::{FunctionCall, FunctionDeclaration, FunctionResponse, Tool};
10
11use crate::error::ToolError;
12
13use super::{ActiveStreamingTool, DEFAULT_TOOL_TIMEOUT, ToolClass, ToolFunction, ToolKind};
14
15/// Routes function calls to the right tool implementation.
16pub struct ToolDispatcher {
17    /// Ordered by name, so declarations are identical from run to run.
18    tools: BTreeMap<String, ToolKind>,
19    active: Arc<tokio::sync::Mutex<HashMap<String, ActiveStreamingTool>>>,
20    default_timeout: Duration,
21    /// Tool declarations, computed on first access and cleared whenever a tool
22    /// is registered.
23    cached_declarations: std::sync::OnceLock<Vec<Tool>>,
24    /// Optional provider consulted before running confirmation-gated tools.
25    confirmation_provider: Option<Arc<dyn crate::confirmation::ConfirmationProvider>>,
26}
27
28impl ToolDispatcher {
29    /// Create a new empty tool dispatcher with the default 30-second timeout.
30    ///
31    /// # Examples
32    ///
33    /// ```rust,ignore
34    /// use gemini_adk_rs::tool::{ToolDispatcher, SimpleTool};
35    /// use serde_json::json;
36    ///
37    /// let mut dispatcher = ToolDispatcher::new();
38    /// dispatcher.register(SimpleTool::new(
39    ///     "echo", "Echo input", None,
40    ///     |args| async move { Ok(args) },
41    /// ));
42    /// ```
43    pub fn new() -> Self {
44        Self {
45            tools: BTreeMap::new(),
46            active: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
47            default_timeout: DEFAULT_TOOL_TIMEOUT,
48            cached_declarations: std::sync::OnceLock::new(),
49            confirmation_provider: None,
50        }
51    }
52
53    /// Set the default timeout for tool calls.
54    pub fn with_timeout(mut self, timeout: Duration) -> Self {
55        self.default_timeout = timeout;
56        self
57    }
58
59    /// Attach a confirmation provider (builder form).
60    ///
61    /// Once set, any tool reporting
62    /// [`requires_confirmation`](crate::tool::ToolFunction::requires_confirmation)
63    /// — e.g. one built with `T::confirm(..)` — is checked against the provider
64    /// before it executes; a denied decision returns a `ToolError` instead of
65    /// running the tool. With no provider configured, confirmation-gated tools
66    /// run normally (enforcement is opt-in).
67    pub fn with_confirmation_provider(
68        mut self,
69        provider: Arc<dyn crate::confirmation::ConfirmationProvider>,
70    ) -> Self {
71        self.confirmation_provider = Some(provider);
72        self
73    }
74
75    /// Attach a confirmation provider in place. See
76    /// [`with_confirmation_provider`](Self::with_confirmation_provider).
77    pub fn set_confirmation_provider(
78        &mut self,
79        provider: Arc<dyn crate::confirmation::ConfirmationProvider>,
80    ) {
81        self.confirmation_provider = Some(provider);
82    }
83
84    /// Whether a confirmation provider is configured.
85    pub fn has_confirmation_provider(&self) -> bool {
86        self.confirmation_provider.is_some()
87    }
88
89    /// Consult the confirmation provider for a gated tool. Returns `Ok(())`
90    /// when the tool is not gated, no provider is set, or the call is approved;
91    /// returns a `ToolError` when the provider denies it.
92    async fn ensure_confirmed(
93        &self,
94        func: &Arc<dyn ToolFunction>,
95        args: &serde_json::Value,
96    ) -> Result<(), ToolError> {
97        if !func.requires_confirmation() {
98            return Ok(());
99        }
100        let Some(provider) = &self.confirmation_provider else {
101            return Ok(());
102        };
103        let request = crate::confirmation::ConfirmationRequest {
104            tool_name: func.name().to_string(),
105            args: args.clone(),
106            message: func.confirmation_message().map(str::to_string),
107        };
108        let decision = provider.confirm(request).await;
109        if decision.confirmed {
110            Ok(())
111        } else {
112            Err(ToolError::Declined(
113                decision.hint.unwrap_or_else(|| "no reason given".into()),
114            ))
115        }
116    }
117
118    /// Returns the configured default timeout.
119    pub fn default_timeout(&self) -> Duration {
120        self.default_timeout
121    }
122
123    /// Register a tool that implements [`ToolFunction`].
124    pub fn register(&mut self, tool: impl ToolFunction) {
125        self.insert(ToolKind::Function(Arc::new(tool)));
126    }
127
128    /// Register a regular function tool (pre-wrapped in Arc).
129    pub fn register_function(&mut self, tool: Arc<dyn ToolFunction>) {
130        self.insert(ToolKind::Function(tool));
131    }
132
133    /// Register a streaming tool.
134    pub fn register_streaming(&mut self, tool: Arc<dyn super::StreamingTool>) {
135        self.insert(ToolKind::Streaming(tool));
136    }
137
138    /// Register an input-streaming tool.
139    pub fn register_input_streaming(&mut self, tool: Arc<dyn super::InputStreamingTool>) {
140        self.insert(ToolKind::InputStream(tool));
141    }
142
143    /// Register `tool` under its name, replacing a tool of the same name.
144    fn insert(&mut self, tool: ToolKind) {
145        let name = match &tool {
146            ToolKind::Function(f) => f.name(),
147            ToolKind::Streaming(s) => s.name(),
148            ToolKind::InputStream(i) => i.name(),
149        }
150        .to_string();
151        self.tools.insert(name, tool);
152        self.cached_declarations.take();
153    }
154
155    /// Add every tool of `other` whose name this dispatcher does not already
156    /// have. This dispatcher's own tools, timeout and confirmation provider
157    /// win.
158    pub fn merge(&mut self, other: ToolDispatcher) {
159        for (name, tool) in other.tools {
160            self.tools.entry(name).or_insert(tool);
161        }
162        self.cached_declarations.take();
163    }
164
165    /// The function tools that ask for confirmation before they run
166    /// (`T::confirm(..)`), which need a confirmation provider to be gated.
167    pub fn gated_tools(&self) -> impl Iterator<Item = &str> {
168        self.tools.iter().filter_map(|(name, tool)| match tool {
169            ToolKind::Function(f) if f.requires_confirmation() => Some(name.as_str()),
170            _ => None,
171        })
172    }
173
174    /// The names of the registered tools, in declaration order.
175    pub fn names(&self) -> impl Iterator<Item = &str> {
176        self.tools.keys().map(String::as_str)
177    }
178
179    /// Get a tool by name (for introspection/streaming tool spawning).
180    pub fn get_tool(&self, name: &str) -> Option<&ToolKind> {
181        self.tools.get(name)
182    }
183
184    /// Classify a tool by name.
185    pub fn classify(&self, name: &str) -> Option<ToolClass> {
186        self.tools.get(name).map(|t| match t {
187            ToolKind::Function(_) => ToolClass::Regular,
188            ToolKind::Streaming(_) => ToolClass::Streaming,
189            ToolKind::InputStream(_) => ToolClass::InputStream,
190        })
191    }
192
193    /// Call a regular function tool by name, using the default timeout.
194    ///
195    /// The tool gets a detached [`ToolContext`](super::ToolContext) (fresh
196    /// state, no call id). Inside a session use
197    /// [`call_function_in`](Self::call_function_in).
198    pub async fn call_function(
199        &self,
200        name: &str,
201        args: serde_json::Value,
202    ) -> Result<serde_json::Value, ToolError> {
203        self.call_function_with_timeout(name, args, self.default_timeout)
204            .await
205    }
206
207    /// Call a regular function tool by name within a session: the tool
208    /// receives `ctx` (see [`ToolContext`](super::ToolContext)). The default
209    /// timeout applies, and cancelling `ctx.cancel` drops the call with
210    /// [`ToolError::Cancelled`].
211    pub async fn call_function_in(
212        &self,
213        name: &str,
214        args: serde_json::Value,
215        ctx: super::ToolContext,
216    ) -> Result<serde_json::Value, ToolError> {
217        let func = self.function(name)?;
218        self.ensure_confirmed(&func, &args).await?;
219        let timeout = self.default_timeout;
220        let cancel = ctx.cancel.clone();
221        tokio::select! {
222            biased;
223            () = cancel.cancelled() => Err(ToolError::Cancelled),
224            result = tokio::time::timeout(timeout, func.call_with_context(args, ctx)) => {
225                result.unwrap_or(Err(ToolError::Timeout(timeout)))
226            }
227        }
228    }
229
230    /// The regular function tool registered as `name`.
231    fn function(&self, name: &str) -> Result<Arc<dyn super::ToolFunction>, ToolError> {
232        match self.tools.get(name) {
233            Some(ToolKind::Function(f)) => Ok(f.clone()),
234            Some(_) => Err(ToolError::Other(format!(
235                "{name} is not a regular function tool"
236            ))),
237            None => Err(ToolError::NotFound(name.to_string())),
238        }
239    }
240
241    /// Call a regular function tool by name with an explicit timeout.
242    ///
243    /// If the tool does not complete within the given duration, its future is
244    /// dropped (cancelling it) and `ToolError::Timeout` is returned.
245    pub async fn call_function_with_timeout(
246        &self,
247        name: &str,
248        args: serde_json::Value,
249        timeout: Duration,
250    ) -> Result<serde_json::Value, ToolError> {
251        let func = match self.tools.get(name) {
252            Some(ToolKind::Function(f)) => f.clone(),
253            Some(_) => {
254                return Err(ToolError::Other(format!(
255                    "{name} is not a regular function tool"
256                )));
257            }
258            None => return Err(ToolError::NotFound(name.to_string())),
259        };
260
261        self.ensure_confirmed(&func, &args).await?;
262
263        match tokio::time::timeout(
264            timeout,
265            func.call_with_context(args, super::ToolContext::detached()),
266        )
267        .await
268        {
269            Ok(result) => result,
270            Err(_elapsed) => Err(ToolError::Timeout(timeout)),
271        }
272    }
273
274    /// Call a regular function tool by name, racing against a cancellation token.
275    ///
276    /// If the token is cancelled before the tool completes, its future is
277    /// dropped and `ToolError::Cancelled` is returned.
278    pub async fn call_function_with_cancel(
279        &self,
280        name: &str,
281        args: serde_json::Value,
282        cancel: CancellationToken,
283    ) -> Result<serde_json::Value, ToolError> {
284        let func = match self.tools.get(name) {
285            Some(ToolKind::Function(f)) => f.clone(),
286            Some(_) => {
287                return Err(ToolError::Other(format!(
288                    "{name} is not a regular function tool"
289                )));
290            }
291            None => return Err(ToolError::NotFound(name.to_string())),
292        };
293
294        self.ensure_confirmed(&func, &args).await?;
295
296        tokio::select! {
297            result = func.call_with_context(
298                args,
299                super::ToolContext::detached().with_cancel(cancel.clone()),
300            ) => result,
301            _ = cancel.cancelled() => Err(ToolError::Cancelled),
302        }
303    }
304
305    /// Build a FunctionResponse from a FunctionCall result.
306    pub fn build_response(
307        call: &FunctionCall,
308        result: Result<serde_json::Value, ToolError>,
309    ) -> FunctionResponse {
310        match result {
311            Ok(value) => FunctionResponse {
312                name: call.name.clone(),
313                response: value,
314                id: call.id.clone(),
315                scheduling: None,
316            },
317            Err(e) => FunctionResponse {
318                name: call.name.clone(),
319                response: serde_json::json!({"error": e.to_string()}),
320                id: call.id.clone(),
321                scheduling: None,
322            },
323        }
324    }
325
326    /// Cancel a streaming tool by name.
327    pub async fn cancel_streaming(&self, name: &str) {
328        let mut active = self.active.lock().await;
329        if let Some(tool) = active.remove(name) {
330            tool.cancel.cancel();
331            tool.task.abort();
332        }
333    }
334
335    /// Store an active streaming tool (for cancellation tracking).
336    pub(crate) async fn store_active(&self, id: String, tool: ActiveStreamingTool) {
337        self.active.lock().await.insert(id, tool);
338    }
339
340    /// Cancel streaming tools by IDs.
341    pub async fn cancel_by_ids(&self, ids: &[String]) {
342        let mut active = self.active.lock().await;
343        for id in ids {
344            if let Some(tool) = active.remove(id.as_str()) {
345                tool.cancel.cancel();
346                tool.task.abort();
347            }
348        }
349    }
350
351    /// Generate Tool declarations for the setup message.
352    ///
353    /// Declarations are ordered by tool name and cached until the next
354    /// `register*()` or [`merge`](Self::merge).
355    pub fn to_tool_declarations(&self) -> Vec<Tool> {
356        self.cached_declarations
357            .get_or_init(|| {
358                let declarations: Vec<FunctionDeclaration> = self
359                    .tools
360                    .values()
361                    .map(|t| {
362                        let (name, desc, params) = match t {
363                            ToolKind::Function(f) => (f.name(), f.description(), f.parameters()),
364                            ToolKind::Streaming(s) => (s.name(), s.description(), s.parameters()),
365                            ToolKind::InputStream(i) => (i.name(), i.description(), i.parameters()),
366                        };
367                        FunctionDeclaration {
368                            name: name.to_string(),
369                            description: desc.to_string(),
370                            parameters: params,
371                            behavior: None,
372                        }
373                    })
374                    .collect();
375
376                if declarations.is_empty() {
377                    vec![]
378                } else {
379                    vec![Tool::functions(declarations)]
380                }
381            })
382            .clone()
383    }
384
385    /// Number of registered tools.
386    pub fn len(&self) -> usize {
387        self.tools.len()
388    }
389
390    /// Whether no tools are registered.
391    pub fn is_empty(&self) -> bool {
392        self.tools.is_empty()
393    }
394}
395
396impl Default for ToolDispatcher {
397    fn default() -> Self {
398        Self::new()
399    }
400}
401
402impl gemini_genai_rs::prelude::ToolProvider for ToolDispatcher {
403    fn declarations(&self) -> Vec<gemini_genai_rs::prelude::Tool> {
404        self.to_tool_declarations()
405    }
406}
407
408#[cfg(test)]
409mod confirmation_tests {
410    use super::*;
411    use crate::confirmation::StaticConfirmation;
412    use crate::tool::{PolicyTool, SimpleTool, policy::ToolPolicy};
413    use serde_json::json;
414    use std::sync::atomic::{AtomicUsize, Ordering};
415
416    /// A counting tool wrapped in a confirm policy.
417    fn confirm_tool(runs: Arc<AtomicUsize>) -> Arc<dyn ToolFunction> {
418        let inner: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new(
419            "danger",
420            "does something sensitive",
421            None,
422            move |_| {
423                let runs = runs.clone();
424                async move {
425                    runs.fetch_add(1, Ordering::SeqCst);
426                    Ok(json!({ "ok": true }))
427                }
428            },
429        ));
430        Arc::new(PolicyTool::new(
431            inner,
432            ToolPolicy::new().with_confirm(Some("delete production data?".into())),
433        ))
434    }
435
436    #[tokio::test]
437    async fn denied_confirmation_blocks_execution() {
438        let runs = Arc::new(AtomicUsize::new(0));
439        let mut d = ToolDispatcher::new();
440        d.register_function(confirm_tool(runs.clone()));
441        d.set_confirmation_provider(StaticConfirmation::deny_all("blocked by policy"));
442
443        let result = d.call_function("danger", json!({})).await;
444        assert!(
445            matches!(&result, Err(ToolError::Declined(reason)) if reason == "blocked by policy"),
446            "the model must learn why: {result:?}"
447        );
448        assert_eq!(
449            runs.load(Ordering::SeqCst),
450            0,
451            "tool must not run when denied"
452        );
453    }
454
455    #[tokio::test]
456    async fn approved_confirmation_runs() {
457        let runs = Arc::new(AtomicUsize::new(0));
458        let mut d = ToolDispatcher::new();
459        d.register_function(confirm_tool(runs.clone()));
460        d.set_confirmation_provider(StaticConfirmation::allow_all());
461
462        let out = d.call_function("danger", json!({})).await.unwrap();
463        assert_eq!(out["ok"], true);
464        assert_eq!(runs.load(Ordering::SeqCst), 1);
465    }
466
467    #[tokio::test]
468    async fn no_provider_runs_optin() {
469        // Enforcement is opt-in: a confirm-gated tool runs when no provider is set.
470        let runs = Arc::new(AtomicUsize::new(0));
471        let mut d = ToolDispatcher::new();
472        d.register_function(confirm_tool(runs.clone()));
473
474        let out = d.call_function("danger", json!({})).await.unwrap();
475        assert_eq!(out["ok"], true);
476        assert_eq!(runs.load(Ordering::SeqCst), 1);
477    }
478
479    #[tokio::test]
480    async fn provider_sees_request_and_ignores_non_gated_tools() {
481        // A non-confirm tool is never sent to the (deny-all) provider.
482        let mut d = ToolDispatcher::new();
483        d.register(SimpleTool::new(
484            "plain",
485            "no confirmation",
486            None,
487            |_| async move { Ok(json!({ "ran": true })) },
488        ));
489        d.set_confirmation_provider(StaticConfirmation::deny_all("should not be consulted"));
490
491        let out = d.call_function("plain", json!({})).await.unwrap();
492        assert_eq!(out["ran"], true);
493    }
494
495    #[tokio::test]
496    async fn nested_policy_wrapper_does_not_bypass_confirmation() {
497        // T::cached(T::confirm(tool)): an outer cache PolicyTool (confirm=false)
498        // wraps an inner confirm PolicyTool. The gate must still fire.
499        let runs = Arc::new(AtomicUsize::new(0));
500        let inner_confirm = confirm_tool(runs.clone()); // Arc<PolicyTool{confirm}>
501        let outer_cached: Arc<dyn ToolFunction> = Arc::new(PolicyTool::new(
502            inner_confirm,
503            ToolPolicy::new().with_cache(),
504        ));
505        assert!(
506            outer_cached.requires_confirmation(),
507            "must propagate through nesting"
508        );
509
510        let mut d = ToolDispatcher::new();
511        d.register_function(outer_cached);
512        d.set_confirmation_provider(StaticConfirmation::deny_all("blocked"));
513
514        let result = d.call_function("danger", json!({})).await;
515        assert!(matches!(result, Err(ToolError::Declined(_))));
516        assert_eq!(
517            runs.load(Ordering::SeqCst),
518            0,
519            "nested confirm must not run when denied"
520        );
521    }
522
523    #[tokio::test]
524    async fn closure_provider_can_gate_by_name() {
525        let runs = Arc::new(AtomicUsize::new(0));
526        let mut d = ToolDispatcher::new();
527        d.register_function(confirm_tool(runs.clone()));
528        d.set_confirmation_provider(Arc::new(
529            |req: crate::confirmation::ConfirmationRequest| async move {
530                if req.tool_name == "danger" {
531                    crate::confirmation::ToolConfirmation::denied("name-gated")
532                } else {
533                    crate::confirmation::ToolConfirmation::confirmed()
534                }
535            },
536        ));
537
538        assert!(d.call_function("danger", json!({})).await.is_err());
539        assert_eq!(runs.load(Ordering::SeqCst), 0);
540    }
541}
542
543#[cfg(test)]
544mod declaration_tests {
545    use super::*;
546    use crate::tool::SimpleTool;
547
548    fn named(name: &'static str) -> Arc<dyn ToolFunction> {
549        Arc::new(SimpleTool::new(name, name, None, |_| async {
550            Ok(serde_json::json!({}))
551        }))
552    }
553
554    fn declared(d: &ToolDispatcher) -> Vec<String> {
555        d.to_tool_declarations()
556            .iter()
557            .flat_map(|t| t.function_declarations.iter().flatten())
558            .map(|f| f.name.clone())
559            .collect()
560    }
561
562    /// Registering after the declarations were read must change them; the
563    /// cache used to be computed once and never cleared.
564    #[test]
565    fn registering_a_tool_refreshes_the_declarations() {
566        let mut d = ToolDispatcher::new();
567        d.register_function(named("a"));
568        assert_eq!(declared(&d), ["a"]);
569        d.register_function(named("b"));
570        assert_eq!(declared(&d), ["a", "b"]);
571    }
572
573    /// Declarations are ordered by name, whatever the registration order, so a
574    /// setup message is byte-identical from run to run.
575    #[test]
576    fn declarations_are_deterministic() {
577        let mut d = ToolDispatcher::new();
578        for name in ["zeta", "alpha", "mid"] {
579            d.register_function(named(name));
580        }
581        assert_eq!(declared(&d), ["alpha", "mid", "zeta"]);
582        assert_eq!(d.names().collect::<Vec<_>>(), ["alpha", "mid", "zeta"]);
583    }
584
585    #[test]
586    fn merge_keeps_both_and_this_dispatcher_wins_a_clash() {
587        let mut mine = ToolDispatcher::new().with_timeout(Duration::from_secs(3));
588        mine.register_function(Arc::new(SimpleTool::new(
589            "shared",
590            "mine",
591            None,
592            |_| async { Ok(serde_json::json!({})) },
593        )));
594        let mut theirs = ToolDispatcher::new();
595        theirs.register_function(named("shared"));
596        theirs.register_function(named("extra"));
597        mine.merge(theirs);
598        assert_eq!(declared(&mine), ["extra", "shared"]);
599        match mine.get_tool("shared") {
600            Some(ToolKind::Function(f)) => assert_eq!(f.description(), "mine"),
601            _ => panic!("shared must stay a function tool"),
602        }
603        assert_eq!(mine.default_timeout(), Duration::from_secs(3));
604    }
605}
gemini_adk_rs/tool/dispatcher.rs

gemini_adk_rs/tool/
dispatcher.rs
