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 covers confirmation and execution. Cancelling `ctx.cancel`
210    /// drops either wait with [`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        self.call_function_with_context_timeout(name, args, ctx, self.default_timeout)
218            .await
219    }
220
221    /// The regular function tool registered as `name`.
222    fn function(&self, name: &str) -> Result<Arc<dyn super::ToolFunction>, ToolError> {
223        match self.tools.get(name) {
224            Some(ToolKind::Function(f)) => Ok(f.clone()),
225            Some(_) => Err(ToolError::Other(format!(
226                "{name} is not a regular function tool"
227            ))),
228            None => Err(ToolError::NotFound(name.to_string())),
229        }
230    }
231
232    /// Call a regular function tool by name with an explicit timeout.
233    ///
234    /// The duration includes confirmation and execution. When it expires,
235    /// the pending future is dropped and `ToolError::Timeout` is returned.
236    pub async fn call_function_with_timeout(
237        &self,
238        name: &str,
239        args: serde_json::Value,
240        timeout: Duration,
241    ) -> Result<serde_json::Value, ToolError> {
242        self.call_function_with_context_timeout(name, args, super::ToolContext::detached(), timeout)
243            .await
244    }
245
246    /// Call a regular function tool by name, racing against a cancellation token.
247    ///
248    /// The default timeout covers confirmation and execution. If the token
249    /// is cancelled during either wait, its future is dropped and
250    /// `ToolError::Cancelled` is returned.
251    pub async fn call_function_with_cancel(
252        &self,
253        name: &str,
254        args: serde_json::Value,
255        cancel: CancellationToken,
256    ) -> Result<serde_json::Value, ToolError> {
257        self.call_function_in(
258            name,
259            args,
260            super::ToolContext::detached().with_cancel(cancel),
261        )
262        .await
263    }
264
265    async fn call_function_with_context_timeout(
266        &self,
267        name: &str,
268        args: serde_json::Value,
269        ctx: super::ToolContext,
270        timeout: Duration,
271    ) -> Result<serde_json::Value, ToolError> {
272        let func = self.function(name)?;
273        let cancel = ctx.cancel.clone();
274        tokio::select! {
275            biased;
276            () = cancel.cancelled() => Err(ToolError::Cancelled),
277            result = tokio::time::timeout(timeout, async {
278                self.ensure_confirmed(&func, &args).await?;
279                if cancel.is_cancelled() {
280                    return Err(ToolError::Cancelled);
281                }
282                func.call_with_context(args, ctx).await
283            }) => result.unwrap_or(Err(ToolError::Timeout(timeout))),
284        }
285    }
286
287    /// Build a FunctionResponse from a FunctionCall result.
288    pub fn build_response(
289        call: &FunctionCall,
290        result: Result<serde_json::Value, ToolError>,
291    ) -> FunctionResponse {
292        match result {
293            Ok(value) => FunctionResponse {
294                name: call.name.clone(),
295                response: value,
296                id: call.id.clone(),
297                scheduling: None,
298            },
299            Err(e) => FunctionResponse {
300                name: call.name.clone(),
301                response: serde_json::json!({"error": e.to_string()}),
302                id: call.id.clone(),
303                scheduling: None,
304            },
305        }
306    }
307
308    /// Cancel a streaming tool by name.
309    pub async fn cancel_streaming(&self, name: &str) {
310        let mut active = self.active.lock().await;
311        if let Some(tool) = active.remove(name) {
312            tool.cancel.cancel();
313            tool.task.abort();
314        }
315    }
316
317    /// Store an active streaming tool (for cancellation tracking).
318    pub(crate) async fn store_active(&self, id: String, tool: ActiveStreamingTool) {
319        self.active.lock().await.insert(id, tool);
320    }
321
322    /// Cancel streaming tools by IDs.
323    pub async fn cancel_by_ids(&self, ids: &[String]) {
324        let mut active = self.active.lock().await;
325        for id in ids {
326            if let Some(tool) = active.remove(id.as_str()) {
327                tool.cancel.cancel();
328                tool.task.abort();
329            }
330        }
331    }
332
333    /// Generate Tool declarations for the setup message.
334    ///
335    /// Declarations are ordered by tool name and cached until the next
336    /// `register*()` or [`merge`](Self::merge).
337    pub fn to_tool_declarations(&self) -> Vec<Tool> {
338        self.cached_declarations
339            .get_or_init(|| {
340                let declarations: Vec<FunctionDeclaration> = self
341                    .tools
342                    .values()
343                    .map(|t| {
344                        let (name, desc, params) = match t {
345                            ToolKind::Function(f) => (f.name(), f.description(), f.parameters()),
346                            ToolKind::Streaming(s) => (s.name(), s.description(), s.parameters()),
347                            ToolKind::InputStream(i) => (i.name(), i.description(), i.parameters()),
348                        };
349                        FunctionDeclaration {
350                            name: name.to_string(),
351                            description: desc.to_string(),
352                            parameters: params,
353                            behavior: None,
354                        }
355                    })
356                    .collect();
357
358                if declarations.is_empty() {
359                    vec![]
360                } else {
361                    vec![Tool::functions(declarations)]
362                }
363            })
364            .clone()
365    }
366
367    /// Number of registered tools.
368    pub fn len(&self) -> usize {
369        self.tools.len()
370    }
371
372    /// Whether no tools are registered.
373    pub fn is_empty(&self) -> bool {
374        self.tools.is_empty()
375    }
376}
377
378impl Default for ToolDispatcher {
379    fn default() -> Self {
380        Self::new()
381    }
382}
383
384impl gemini_genai_rs::prelude::ToolProvider for ToolDispatcher {
385    fn declarations(&self) -> Vec<gemini_genai_rs::prelude::Tool> {
386        self.to_tool_declarations()
387    }
388}
389
390#[cfg(test)]
391mod confirmation_tests {
392    use super::*;
393    use crate::confirmation::StaticConfirmation;
394    use crate::tool::{PolicyTool, SimpleTool, policy::ToolPolicy};
395    use serde_json::json;
396    use std::sync::atomic::{AtomicUsize, Ordering};
397
398    /// A counting tool wrapped in a confirm policy.
399    fn confirm_tool(runs: Arc<AtomicUsize>) -> Arc<dyn ToolFunction> {
400        let inner: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new(
401            "danger",
402            "does something sensitive",
403            None,
404            move |_| {
405                let runs = runs.clone();
406                async move {
407                    runs.fetch_add(1, Ordering::SeqCst);
408                    Ok(json!({ "ok": true }))
409                }
410            },
411        ));
412        Arc::new(PolicyTool::new(
413            inner,
414            ToolPolicy::new().with_confirm(Some("delete production data?".into())),
415        ))
416    }
417
418    #[tokio::test]
419    async fn denied_confirmation_blocks_execution() {
420        let runs = Arc::new(AtomicUsize::new(0));
421        let mut d = ToolDispatcher::new();
422        d.register_function(confirm_tool(runs.clone()));
423        d.set_confirmation_provider(StaticConfirmation::deny_all("blocked by policy"));
424
425        let result = d.call_function("danger", json!({})).await;
426        assert!(
427            matches!(&result, Err(ToolError::Declined(reason)) if reason == "blocked by policy"),
428            "the model must learn why: {result:?}"
429        );
430        assert_eq!(
431            runs.load(Ordering::SeqCst),
432            0,
433            "tool must not run when denied"
434        );
435    }
436
437    #[tokio::test]
438    async fn approved_confirmation_runs() {
439        let runs = Arc::new(AtomicUsize::new(0));
440        let mut d = ToolDispatcher::new();
441        d.register_function(confirm_tool(runs.clone()));
442        d.set_confirmation_provider(StaticConfirmation::allow_all());
443
444        let out = d.call_function("danger", json!({})).await.unwrap();
445        assert_eq!(out["ok"], true);
446        assert_eq!(runs.load(Ordering::SeqCst), 1);
447    }
448
449    fn awaiting_confirmation(timeout: Duration) -> (ToolDispatcher, Arc<AtomicUsize>) {
450        let runs = Arc::new(AtomicUsize::new(0));
451        let mut dispatcher = ToolDispatcher::new().with_timeout(timeout);
452        dispatcher.register_function(confirm_tool(runs.clone()));
453        dispatcher.set_confirmation_provider(Arc::new(
454            |_: crate::confirmation::ConfirmationRequest| {
455                std::future::pending::<crate::confirmation::ToolConfirmation>()
456            },
457        ));
458        (dispatcher, runs)
459    }
460
461    #[tokio::test]
462    async fn session_cancellation_stops_pending_confirmation() {
463        let (dispatcher, runs) = awaiting_confirmation(Duration::from_secs(30));
464        let cancel = CancellationToken::new();
465        let ctx = crate::tool::ToolContext::detached().with_cancel(cancel.clone());
466        let mut call =
467            tokio_test::task::spawn(dispatcher.call_function_in("danger", json!({}), ctx));
468        tokio_test::assert_pending!(call.poll());
469
470        cancel.cancel();
471
472        assert!(matches!(
473            tokio_test::assert_ready!(call.poll()),
474            Err(ToolError::Cancelled)
475        ));
476        assert_eq!(runs.load(Ordering::SeqCst), 0);
477    }
478
479    #[tokio::test]
480    async fn token_cancellation_stops_pending_confirmation() {
481        let (dispatcher, runs) = awaiting_confirmation(Duration::from_secs(30));
482        let cancel = CancellationToken::new();
483        let mut call = tokio_test::task::spawn(dispatcher.call_function_with_cancel(
484            "danger",
485            json!({}),
486            cancel.clone(),
487        ));
488        tokio_test::assert_pending!(call.poll());
489
490        cancel.cancel();
491
492        assert!(matches!(
493            tokio_test::assert_ready!(call.poll()),
494            Err(ToolError::Cancelled)
495        ));
496        assert_eq!(runs.load(Ordering::SeqCst), 0);
497    }
498
499    async fn assert_confirmation_timeout(
500        call: impl std::future::Future<Output = Result<serde_json::Value, ToolError>>,
501        timeout: Duration,
502    ) {
503        let mut call = tokio_test::task::spawn(call);
504        tokio_test::assert_pending!(call.poll());
505        tokio::time::advance(timeout).await;
506        assert!(matches!(
507            tokio_test::assert_ready!(call.poll()),
508            Err(ToolError::Timeout(elapsed)) if elapsed == timeout
509        ));
510    }
511
512    #[tokio::test(start_paused = true)]
513    async fn session_timeout_includes_pending_confirmation() {
514        let timeout = Duration::from_secs(5);
515        let (dispatcher, runs) = awaiting_confirmation(timeout);
516        assert_confirmation_timeout(
517            dispatcher.call_function_in("danger", json!({}), crate::tool::ToolContext::detached()),
518            timeout,
519        )
520        .await;
521        assert_eq!(runs.load(Ordering::SeqCst), 0);
522    }
523
524    #[tokio::test(start_paused = true)]
525    async fn cancellable_call_timeout_includes_pending_confirmation() {
526        let timeout = Duration::from_secs(5);
527        let (dispatcher, runs) = awaiting_confirmation(timeout);
528        assert_confirmation_timeout(
529            dispatcher.call_function_with_cancel("danger", json!({}), CancellationToken::new()),
530            timeout,
531        )
532        .await;
533        assert_eq!(runs.load(Ordering::SeqCst), 0);
534    }
535
536    #[tokio::test(start_paused = true)]
537    async fn explicit_timeout_includes_pending_confirmation() {
538        let timeout = Duration::from_secs(5);
539        let (dispatcher, runs) = awaiting_confirmation(Duration::from_secs(30));
540        assert_confirmation_timeout(
541            dispatcher.call_function_with_timeout("danger", json!({}), timeout),
542            timeout,
543        )
544        .await;
545        assert_eq!(runs.load(Ordering::SeqCst), 0);
546    }
547
548    #[tokio::test]
549    async fn cancellation_during_approval_prevents_execution() {
550        let runs = Arc::new(AtomicUsize::new(0));
551        let cancel = CancellationToken::new();
552        let approval_cancel = cancel.clone();
553        let mut dispatcher = ToolDispatcher::new();
554        dispatcher.register_function(confirm_tool(runs.clone()));
555        dispatcher.set_confirmation_provider(Arc::new(
556            move |_: crate::confirmation::ConfirmationRequest| {
557                let cancel = approval_cancel.clone();
558                async move {
559                    cancel.cancel();
560                    crate::confirmation::ToolConfirmation::confirmed()
561                }
562            },
563        ));
564
565        let result = dispatcher
566            .call_function_in(
567                "danger",
568                json!({}),
569                crate::tool::ToolContext::detached().with_cancel(cancel),
570            )
571            .await;
572        assert!(matches!(result, Err(ToolError::Cancelled)));
573        assert_eq!(runs.load(Ordering::SeqCst), 0);
574    }
575
576    #[tokio::test]
577    async fn pending_confirmation_can_approve_execution() {
578        let runs = Arc::new(AtomicUsize::new(0));
579        let approve = Arc::new(tokio::sync::Notify::new());
580        let approval = approve.clone();
581        let mut dispatcher = ToolDispatcher::new();
582        dispatcher.register_function(confirm_tool(runs.clone()));
583        dispatcher.set_confirmation_provider(Arc::new(
584            move |_: crate::confirmation::ConfirmationRequest| {
585                let approve = approval.clone();
586                async move {
587                    approve.notified().await;
588                    crate::confirmation::ToolConfirmation::confirmed()
589                }
590            },
591        ));
592        let mut call = tokio_test::task::spawn(dispatcher.call_function_with_cancel(
593            "danger",
594            json!({}),
595            CancellationToken::new(),
596        ));
597        tokio_test::assert_pending!(call.poll());
598        assert_eq!(runs.load(Ordering::SeqCst), 0);
599
600        approve.notify_one();
601
602        let result = tokio_test::assert_ready!(call.poll()).unwrap();
603        assert_eq!(result["ok"], true);
604        assert_eq!(runs.load(Ordering::SeqCst), 1);
605    }
606
607    #[tokio::test(start_paused = true)]
608    async fn approval_and_execution_share_one_timeout_budget() {
609        let mut dispatcher = ToolDispatcher::new();
610        dispatcher.register(PolicyTool::new(
611            Arc::new(SimpleTool::new(
612                "danger",
613                "slow operation",
614                None,
615                |_| async {
616                    tokio::time::sleep(Duration::from_secs(2)).await;
617                    Ok(json!({ "ok": true }))
618                },
619            )),
620            ToolPolicy::new().with_confirm(None),
621        ));
622        dispatcher.set_confirmation_provider(Arc::new(
623            |_: crate::confirmation::ConfirmationRequest| async {
624                tokio::time::sleep(Duration::from_secs(4)).await;
625                crate::confirmation::ToolConfirmation::confirmed()
626            },
627        ));
628
629        let timeout = Duration::from_secs(5);
630        let started = tokio::time::Instant::now();
631        let result = dispatcher
632            .call_function_with_timeout("danger", json!({}), timeout)
633            .await;
634        assert!(matches!(result, Err(ToolError::Timeout(elapsed)) if elapsed == timeout));
635        assert_eq!(started.elapsed(), timeout);
636    }
637
638    #[tokio::test]
639    async fn no_provider_runs_optin() {
640        // Enforcement is opt-in: a confirm-gated tool runs when no provider is set.
641        let runs = Arc::new(AtomicUsize::new(0));
642        let mut d = ToolDispatcher::new();
643        d.register_function(confirm_tool(runs.clone()));
644
645        let out = d.call_function("danger", json!({})).await.unwrap();
646        assert_eq!(out["ok"], true);
647        assert_eq!(runs.load(Ordering::SeqCst), 1);
648    }
649
650    #[tokio::test]
651    async fn provider_sees_request_and_ignores_non_gated_tools() {
652        // A non-confirm tool is never sent to the (deny-all) provider.
653        let mut d = ToolDispatcher::new();
654        d.register(SimpleTool::new(
655            "plain",
656            "no confirmation",
657            None,
658            |_| async move { Ok(json!({ "ran": true })) },
659        ));
660        d.set_confirmation_provider(StaticConfirmation::deny_all("should not be consulted"));
661
662        let out = d.call_function("plain", json!({})).await.unwrap();
663        assert_eq!(out["ran"], true);
664    }
665
666    #[tokio::test]
667    async fn nested_policy_wrapper_does_not_bypass_confirmation() {
668        // T::cached(T::confirm(tool)): an outer cache PolicyTool (confirm=false)
669        // wraps an inner confirm PolicyTool. The gate must still fire.
670        let runs = Arc::new(AtomicUsize::new(0));
671        let inner_confirm = confirm_tool(runs.clone()); // Arc<PolicyTool{confirm}>
672        let outer_cached: Arc<dyn ToolFunction> = Arc::new(PolicyTool::new(
673            inner_confirm,
674            ToolPolicy::new().with_cache(),
675        ));
676        assert!(
677            outer_cached.requires_confirmation(),
678            "must propagate through nesting"
679        );
680
681        let mut d = ToolDispatcher::new();
682        d.register_function(outer_cached);
683        d.set_confirmation_provider(StaticConfirmation::deny_all("blocked"));
684
685        let result = d.call_function("danger", json!({})).await;
686        assert!(matches!(result, Err(ToolError::Declined(_))));
687        assert_eq!(
688            runs.load(Ordering::SeqCst),
689            0,
690            "nested confirm must not run when denied"
691        );
692    }
693
694    #[tokio::test]
695    async fn closure_provider_can_gate_by_name() {
696        let runs = Arc::new(AtomicUsize::new(0));
697        let mut d = ToolDispatcher::new();
698        d.register_function(confirm_tool(runs.clone()));
699        d.set_confirmation_provider(Arc::new(
700            |req: crate::confirmation::ConfirmationRequest| async move {
701                if req.tool_name == "danger" {
702                    crate::confirmation::ToolConfirmation::denied("name-gated")
703                } else {
704                    crate::confirmation::ToolConfirmation::confirmed()
705                }
706            },
707        ));
708
709        assert!(d.call_function("danger", json!({})).await.is_err());
710        assert_eq!(runs.load(Ordering::SeqCst), 0);
711    }
712}
713
714#[cfg(test)]
715mod declaration_tests {
716    use super::*;
717    use crate::tool::SimpleTool;
718
719    fn named(name: &'static str) -> Arc<dyn ToolFunction> {
720        Arc::new(SimpleTool::new(name, name, None, |_| async {
721            Ok(serde_json::json!({}))
722        }))
723    }
724
725    fn declared(d: &ToolDispatcher) -> Vec<String> {
726        d.to_tool_declarations()
727            .iter()
728            .flat_map(|t| t.function_declarations.iter().flatten())
729            .map(|f| f.name.clone())
730            .collect()
731    }
732
733    /// Registering after the declarations were read must change them; the
734    /// cache used to be computed once and never cleared.
735    #[test]
736    fn registering_a_tool_refreshes_the_declarations() {
737        let mut d = ToolDispatcher::new();
738        d.register_function(named("a"));
739        assert_eq!(declared(&d), ["a"]);
740        d.register_function(named("b"));
741        assert_eq!(declared(&d), ["a", "b"]);
742    }
743
744    /// Declarations are ordered by name, whatever the registration order, so a
745    /// setup message is byte-identical from run to run.
746    #[test]
747    fn declarations_are_deterministic() {
748        let mut d = ToolDispatcher::new();
749        for name in ["zeta", "alpha", "mid"] {
750            d.register_function(named(name));
751        }
752        assert_eq!(declared(&d), ["alpha", "mid", "zeta"]);
753        assert_eq!(d.names().collect::<Vec<_>>(), ["alpha", "mid", "zeta"]);
754    }
755
756    #[test]
757    fn merge_keeps_both_and_this_dispatcher_wins_a_clash() {
758        let mut mine = ToolDispatcher::new().with_timeout(Duration::from_secs(3));
759        mine.register_function(Arc::new(SimpleTool::new(
760            "shared",
761            "mine",
762            None,
763            |_| async { Ok(serde_json::json!({})) },
764        )));
765        let mut theirs = ToolDispatcher::new();
766        theirs.register_function(named("shared"));
767        theirs.register_function(named("extra"));
768        mine.merge(theirs);
769        assert_eq!(declared(&mine), ["extra", "shared"]);
770        match mine.get_tool("shared") {
771            Some(ToolKind::Function(f)) => assert_eq!(f.description(), "mine"),
772            _ => panic!("shared must stay a function tool"),
773        }
774        assert_eq!(mine.default_timeout(), Duration::from_secs(3));
775    }
776}
gemini_adk_rs/tool/dispatcher.rs

gemini_adk_rs/tool/
dispatcher.rs
