gemini_adk_fluent_rs/compose/
tools.rs

1//! T — Tool composition.
2//!
3//! Combine tools with `+`: all of them are offered to the model.
4
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::Arc;
8
9use gemini_adk_rs::text::TextAgent;
10use gemini_adk_rs::tool::{PolicyTool, SimpleTool, ToolFunction, ToolPolicy, TypedTool};
11use gemini_genai_rs::prelude::{FunctionDeclaration, Tool};
12
13/// A tool composite — one or more tool entries.
14///
15/// Built from the `T` namespace and combined with `+`. Any single
16/// [`ToolFunction`] (a `SimpleTool`, a `TypedTool`, the value a `#[tool]`
17/// function returns, or an `Arc<dyn ToolFunction>`) converts into a
18/// one-entry composite via `From`, so `.tools(get_weather())` works without
19/// the namespace.
20#[derive(Clone)]
21#[non_exhaustive]
22pub struct ToolComposite {
23    /// The tool entries in this composite.
24    pub entries: Vec<ToolCompositeEntry>,
25}
26
27/// Async transformer applied to a tool result value.
28pub type TransformFn = Arc<
29    dyn Fn(serde_json::Value) -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>>
30        + Send
31        + Sync,
32>;
33
34/// An entry in a tool composite.
35#[derive(Clone)]
36pub enum ToolCompositeEntry {
37    /// A runtime tool function.
38    Function(Arc<dyn ToolFunction>),
39    /// A built-in Gemini tool declaration.
40    BuiltIn(Tool),
41    /// A text agent wrapped as a tool.
42    Agent {
43        /// Tool name exposed to the model.
44        name: String,
45        /// Tool description exposed to the model.
46        description: String,
47        /// The text agent to invoke.
48        agent: Arc<dyn TextAgent>,
49    },
50    /// An MCP (Model Context Protocol) toolset connection.
51    Mcp {
52        /// Connection params (e.g. URL or command string).
53        params: String,
54    },
55    /// A mock tool that returns a fixed response (useful for testing).
56    Mock {
57        /// Tool name.
58        name: String,
59        /// Tool description.
60        description: String,
61        /// Fixed response to return.
62        response: serde_json::Value,
63    },
64    /// A schema-defined tool (placeholder/marker).
65    Schema {
66        /// Tool name.
67        name: String,
68        /// JSON Schema defining the tool's parameters.
69        schema: serde_json::Value,
70    },
71    /// A tool wrapped with a result transformer.
72    Transform {
73        /// The inner tool entry.
74        inner: Box<ToolCompositeEntry>,
75        /// Transformer function applied to the tool result.
76        transformer: TransformFn,
77    },
78}
79
80impl ToolComposite {
81    /// Create a composite containing a single runtime tool function.
82    pub fn from_function(f: Arc<dyn ToolFunction>) -> Self {
83        Self {
84            entries: vec![ToolCompositeEntry::Function(f)],
85        }
86    }
87
88    /// Create a composite containing a single built-in tool declaration.
89    pub fn from_built_in(tool: Tool) -> Self {
90        Self {
91            entries: vec![ToolCompositeEntry::BuiltIn(tool)],
92        }
93    }
94
95    /// Number of tool entries.
96    pub fn len(&self) -> usize {
97        self.entries.len()
98    }
99
100    /// Whether empty.
101    pub fn is_empty(&self) -> bool {
102        self.entries.is_empty()
103    }
104
105    /// Apply a per-tool [`ToolPolicy`] transform to every function entry.
106    ///
107    /// Each [`ToolCompositeEntry::Function`] is wrapped in a [`PolicyTool`]
108    /// carrying the policy. Successive modifiers nest (e.g. `T::cached(T::timeout(..))`
109    /// applies both timeout and cache), since a `PolicyTool` is itself a
110    /// [`ToolFunction`]. Other entry kinds are left untouched.
111    fn map_function_policy(
112        mut self,
113        f: impl Fn(ToolPolicy) -> ToolPolicy + Send + Sync + 'static,
114    ) -> Self {
115        self.entries = self
116            .entries
117            .into_iter()
118            .map(|entry| match entry {
119                ToolCompositeEntry::Function(func) => {
120                    let policy = f(ToolPolicy::new());
121                    ToolCompositeEntry::Function(PolicyTool::wrap(func, policy))
122                }
123                other => other,
124            })
125            .collect();
126        self
127    }
128}
129
130/// A single tool is a one-entry composite, so `.tools(my_tool)` and
131/// `.tool(my_tool)` accept the same values.
132impl<F: ToolFunction + 'static> From<F> for ToolComposite {
133    fn from(f: F) -> Self {
134        Self::from_function(Arc::new(f))
135    }
136}
137
138/// Combine two tool composites with `+`.
139impl std::ops::Add for ToolComposite {
140    type Output = ToolComposite;
141
142    fn add(mut self, rhs: ToolComposite) -> Self::Output {
143        self.entries.extend(rhs.entries);
144        self
145    }
146}
147
148/// The `T` namespace — static factory methods for tool composition.
149pub struct T;
150
151impl T {
152    /// Register a function tool.
153    pub fn function(f: Arc<dyn ToolFunction>) -> ToolComposite {
154        ToolComposite::from_function(f)
155    }
156
157    /// Add Google Search built-in tool.
158    pub fn google_search() -> ToolComposite {
159        ToolComposite::from_built_in(Tool::google_search())
160    }
161
162    /// Add URL context built-in tool.
163    pub fn url_context() -> ToolComposite {
164        ToolComposite::from_built_in(Tool::url_context())
165    }
166
167    /// Add code execution built-in tool.
168    pub fn code_execution() -> ToolComposite {
169        ToolComposite::from_built_in(Tool::code_execution())
170    }
171
172    /// A tool that takes **no parameters**, from a name, description, and
173    /// async closure.
174    ///
175    /// The declaration the model sees has no parameters, so the model calls
176    /// it with `{}`. For a tool with arguments, use [`#[tool]`](macro@crate::prelude::tool)
177    /// on a documented `async fn`, or [`T::typed`](Self::typed) when the tool
178    /// must capture something (a client, a pool) from its environment.
179    pub fn simple<F, Fut>(
180        name: impl Into<String>,
181        description: impl Into<String>,
182        f: F,
183    ) -> ToolComposite
184    where
185        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
186        Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
187    {
188        let tool = SimpleTool::new(name, description, None, f);
189        ToolComposite::from_function(Arc::new(tool))
190    }
191
192    /// A tool from a closure that also receives the call's
193    /// [`ToolContext`](gemini_adk_rs::tool::ToolContext): the session state,
194    /// the call id, and a cancellation token that fires on barge-in.
195    ///
196    /// ```
197    /// use gemini_adk_fluent_rs::prelude::*;
198    ///
199    /// let balance = T::contextual("balance", "The caller's balance", |_args, ctx| async move {
200    ///     let account: String = ctx.state.get("account_id").unwrap_or_default();
201    ///     Ok(serde_json::json!({ "account": account, "cents": 1200 }))
202    /// });
203    /// ```
204    pub fn contextual<F, Fut>(
205        name: impl Into<String>,
206        description: impl Into<String>,
207        f: F,
208    ) -> ToolComposite
209    where
210        F: Fn(serde_json::Value, gemini_adk_rs::tool::ToolContext) -> Fut + Send + Sync + 'static,
211        Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
212    {
213        let tool = gemini_adk_rs::tool::ContextTool::new(name, description, None, f);
214        ToolComposite::from_function(Arc::new(tool))
215    }
216
217    /// A tool whose arguments are the type `A`, from a closure.
218    ///
219    /// `A`'s JSON Schema (through [`wire_schema`](gemini_adk_rs::tool::wire_schema))
220    /// is the declaration, and doc comments on its fields describe the
221    /// parameters. Reach for this over [`#[tool]`](macro@crate::prelude::tool)
222    /// when the tool needs state from its environment:
223    ///
224    /// ```
225    /// use gemini_adk_fluent_rs::prelude::*;
226    /// use std::sync::Arc;
227    ///
228    /// #[derive(serde::Deserialize, schemars::JsonSchema)]
229    /// struct Lookup {
230    ///     /// The customer's account id.
231    ///     account: String,
232    /// }
233    ///
234    /// let accounts = Arc::new(vec![("a-1".to_string(), 120)]);
235    /// let balance = T::typed("balance", "Look up an account balance", move |args: Lookup| {
236    ///     let accounts = accounts.clone();
237    ///     async move {
238    ///         let found = accounts.iter().find(|(id, _)| *id == args.account);
239    ///         Ok(serde_json::json!({ "balance": found.map(|(_, b)| *b) }))
240    ///     }
241    /// });
242    /// AgentBuilder::new("support").tools(balance);
243    /// ```
244    pub fn typed<A, F, Fut>(
245        name: impl Into<String>,
246        description: impl Into<String>,
247        f: F,
248    ) -> ToolComposite
249    where
250        A: serde::de::DeserializeOwned + schemars::JsonSchema + Send + Sync + 'static,
251        F: Fn(A) -> Fut + Send + Sync + 'static,
252        Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
253    {
254        ToolComposite::from_function(Arc::new(TypedTool::new::<F, Fut>(name, description, f)))
255    }
256
257    /// Alias for [`simple`](Self::simple) — matches upstream Python `T.fn()`.
258    #[deprecated(
259        since = "2.1.0",
260        note = "use `T::simple` (no parameters), `T::typed` or `#[tool]`"
261    )]
262    ///
263    /// Named `fn_tool` because `fn` is a reserved keyword in Rust.
264    pub fn fn_tool<F, Fut>(
265        name: impl Into<String>,
266        description: impl Into<String>,
267        f: F,
268    ) -> ToolComposite
269    where
270        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
271        Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
272    {
273        Self::simple(name, description, f)
274    }
275
276    /// Require user confirmation before each function tool in the composite runs.
277    ///
278    /// The confirmation flag is recorded on the tool's [`ToolPolicy`] and surfaced
279    /// to the runtime via [`PolicyTool::requires_confirmation`] — it is never
280    /// silently dropped. The `message` becomes the confirmation hint. Built-in and
281    /// placeholder entries are left unchanged.
282    pub fn confirm(tool: impl Into<ToolComposite>, message: &str) -> ToolComposite {
283        let tool = tool.into();
284        let msg = if message.is_empty() {
285            None
286        } else {
287            Some(message.to_string())
288        };
289        tool.map_function_policy(move |p| p.with_confirm(msg.clone()))
290    }
291
292    /// Bound each function tool in the composite by a timeout.
293    ///
294    /// At dispatch the tool's future is raced against the duration; on elapse the
295    /// call returns [`ToolError::Timeout`](gemini_adk_rs::ToolError::Timeout).
296    /// Built-in and placeholder entries are left unchanged.
297    pub fn timeout(tool: impl Into<ToolComposite>, duration: std::time::Duration) -> ToolComposite {
298        tool.into()
299            .map_function_policy(move |p| p.with_timeout(duration))
300    }
301
302    /// Memoize each function tool's successful results.
303    ///
304    /// Results are cached by `(tool name, canonical-JSON args)`; repeat calls with
305    /// identical arguments return the cached value without re-invoking the tool.
306    /// Errors are not cached. Built-in/placeholder entries are left unchanged.
307    pub fn cached(tool: impl Into<ToolComposite>) -> ToolComposite {
308        tool.into()
309            .map_function_policy(gemini_adk_rs::tool::ToolPolicy::with_cache)
310    }
311
312    /// Combine multiple tool functions into a single composite.
313    pub fn toolset(tools: Vec<Arc<dyn ToolFunction>>) -> ToolComposite {
314        ToolComposite {
315            entries: tools
316                .into_iter()
317                .map(ToolCompositeEntry::Function)
318                .collect(),
319        }
320    }
321
322    /// Wrap a [`TextAgent`] as a tool (shorthand for creating an agent tool entry).
323    ///
324    /// When invoked, the agent runs via `BaseLlm::generate()` and returns its
325    /// text output as the tool result. State is shared with the parent session.
326    pub fn agent(
327        name: impl Into<String>,
328        description: impl Into<String>,
329        agent: impl TextAgent + 'static,
330    ) -> ToolComposite {
331        ToolComposite {
332            entries: vec![ToolCompositeEntry::Agent {
333                name: name.into(),
334                description: description.into(),
335                agent: Arc::new(agent),
336            }],
337        }
338    }
339
340    /// Create an MCP (Model Context Protocol) toolset entry.
341    ///
342    /// `params` is the connection string (e.g. a URL or command) used to
343    /// establish the MCP session at runtime.
344    pub fn mcp(params: impl Into<String>) -> ToolComposite {
345        ToolComposite {
346            entries: vec![ToolCompositeEntry::Mcp {
347                params: params.into(),
348            }],
349        }
350    }
351
352    /// Create a mock tool that returns a fixed response.
353    ///
354    /// Useful for testing and prototyping without real tool implementations.
355    pub fn mock(
356        name: impl Into<String>,
357        description: impl Into<String>,
358        response: serde_json::Value,
359    ) -> ToolComposite {
360        ToolComposite {
361            entries: vec![ToolCompositeEntry::Mock {
362                name: name.into(),
363                description: description.into(),
364                response,
365            }],
366        }
367    }
368
369    /// Create a schema-defined tool (placeholder/marker).
370    ///
371    /// The tool's parameters are defined by the given JSON Schema value.
372    pub fn schema(name: impl Into<String>, schema: serde_json::Value) -> ToolComposite {
373        ToolComposite {
374            entries: vec![ToolCompositeEntry::Schema {
375                name: name.into(),
376                schema,
377            }],
378        }
379    }
380
381    /// Wrap each tool entry in a composite with a result transformer.
382    ///
383    /// The transformer function is applied to the tool's output value before
384    /// it is returned to the model.
385    pub fn transform<F, Fut>(tool: ToolComposite, f: F) -> ToolComposite
386    where
387        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
388        Fut: Future<Output = serde_json::Value> + Send + 'static,
389    {
390        let f: TransformFn = Arc::new(
391            move |v: serde_json::Value| -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> {
392                Box::pin(f(v))
393            },
394        );
395        ToolComposite {
396            entries: tool
397                .entries
398                .into_iter()
399                .map(|entry| ToolCompositeEntry::Transform {
400                    inner: Box::new(entry),
401                    transformer: Arc::clone(&f),
402                })
403                .collect(),
404        }
405    }
406}
407
408// ── Resolution ────────────────────────────────────────────────────────────
409
410/// A tool entry that needs asynchronous I/O (network or subprocess) to resolve,
411/// and is therefore resolved at connect time rather than when the composite is
412/// built. See [`crate::live::Live`] connection methods. A text
413/// [`AgentBuilder`](crate::builder::AgentBuilder) rejects these at `build`.
414#[derive(Clone, Debug)]
415pub enum DeferredTool {
416    /// MCP server connection — a stdio command line or an SSE/HTTP URL.
417    Mcp {
418        /// Connection string: an `http(s)://` URL (SSE) or a command line (stdio).
419        params: String,
420    },
421}
422
423/// The concrete outcome of classifying a single [`ToolCompositeEntry`].
424///
425/// This is the *single* exhaustive mapping from the composable tool algebra to
426/// the runtime; both [`crate::builder::AgentBuilder`] and [`crate::live::Live`]
427/// resolve through it, so no entry can be silently dropped.
428pub(crate) enum ToolResolution {
429    /// A runtime-executable tool function (register with a dispatcher).
430    Runtime(Arc<dyn ToolFunction>),
431    /// A built-in / declaration-only Gemini tool (add to the session config).
432    BuiltIn(Tool),
433    /// A text agent to expose as a tool (needs a shared session `State`).
434    Agent {
435        /// Tool name exposed to the model.
436        name: String,
437        /// Tool description exposed to the model.
438        description: String,
439        /// The text agent to invoke.
440        agent: Arc<dyn TextAgent>,
441    },
442    /// A tool that can only be resolved with async I/O at connect time.
443    Deferred(DeferredTool),
444}
445
446impl ToolCompositeEntry {
447    #[cfg(test)]
448    fn classify_name(self) -> String {
449        match self.classify() {
450            ToolResolution::Runtime(f) => f.name().to_string(),
451            _ => String::new(),
452        }
453    }
454
455    /// Classify this entry into its concrete [`ToolResolution`]. Exhaustive by
456    /// construction — adding a variant forces every consumer to handle it.
457    pub(crate) fn classify(self) -> ToolResolution {
458        match self {
459            ToolCompositeEntry::Function(f) => ToolResolution::Runtime(f),
460            ToolCompositeEntry::BuiltIn(t) => ToolResolution::BuiltIn(t),
461            ToolCompositeEntry::Agent {
462                name,
463                description,
464                agent,
465            } => ToolResolution::Agent {
466                name,
467                description,
468                agent,
469            },
470            ToolCompositeEntry::Mock {
471                name,
472                description,
473                response,
474            } => ToolResolution::Runtime(Arc::new(SimpleTool::new(
475                name,
476                description,
477                None,
478                move |_args| {
479                    let r = response.clone();
480                    async move { Ok(r) }
481                },
482            ))),
483            ToolCompositeEntry::Transform { inner, transformer } => match inner.classify() {
484                ToolResolution::Runtime(f) => ToolResolution::Runtime(Arc::new(TransformTool {
485                    inner: f,
486                    transformer,
487                })),
488                // A transformer only applies to a runtime function; for any other
489                // inner kind the transform is a no-op and the inner resolution
490                // passes through unchanged.
491                other => other,
492            },
493            ToolCompositeEntry::Schema { name, schema } => {
494                // A declaration-only tool: the model is told the function exists
495                // and the application services the call (e.g. via on_tool_call).
496                ToolResolution::BuiltIn(Tool::functions(vec![FunctionDeclaration {
497                    name,
498                    description: String::new(),
499                    parameters: Some(schema),
500                    behavior: None,
501                }]))
502            }
503            ToolCompositeEntry::Mcp { params } => {
504                ToolResolution::Deferred(DeferredTool::Mcp { params })
505            }
506        }
507    }
508}
509
510/// A [`ToolFunction`] that applies an async transformer to another tool's result.
511struct TransformTool {
512    inner: Arc<dyn ToolFunction>,
513    #[allow(clippy::type_complexity)]
514    transformer: Arc<
515        dyn Fn(serde_json::Value) -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>>
516            + Send
517            + Sync,
518    >,
519}
520
521#[async_trait::async_trait]
522impl ToolFunction for TransformTool {
523    fn name(&self) -> &str {
524        self.inner.name()
525    }
526
527    fn description(&self) -> &str {
528        self.inner.description()
529    }
530
531    fn parameters(&self) -> Option<serde_json::Value> {
532        self.inner.parameters()
533    }
534
535    async fn call(
536        &self,
537        args: serde_json::Value,
538    ) -> Result<serde_json::Value, gemini_adk_rs::error::ToolError> {
539        let result = self.inner.call(args).await?;
540        Ok((self.transformer)(result).await)
541    }
542
543    async fn call_with_context(
544        &self,
545        args: serde_json::Value,
546        ctx: gemini_adk_rs::tool::ToolContext,
547    ) -> Result<serde_json::Value, gemini_adk_rs::error::ToolError> {
548        let result = self.inner.call_with_context(args, ctx).await?;
549        Ok((self.transformer)(result).await)
550    }
551}
552
553#[cfg(test)]
554mod tests {
555    use super::*;
556
557    /// Classify the single entry of a one-element composite.
558    fn classify_one(c: ToolComposite) -> ToolResolution {
559        c.entries.into_iter().next().unwrap().classify()
560    }
561
562    #[test]
563    fn classify_maps_every_variant() {
564        // Synchronous, runtime-executable.
565        assert!(matches!(
566            classify_one(T::mock("m", "d", serde_json::json!({"ok": true}))),
567            ToolResolution::Runtime(_)
568        ));
569        assert!(matches!(
570            classify_one(T::simple("s", "d", |a| async move { Ok(a) })),
571            ToolResolution::Runtime(_)
572        ));
573        // Built-in / declaration-only.
574        assert!(matches!(
575            classify_one(T::google_search()),
576            ToolResolution::BuiltIn(_)
577        ));
578        assert!(matches!(
579            classify_one(T::schema("s", serde_json::json!({"type": "object"}))),
580            ToolResolution::BuiltIn(_)
581        ));
582        // Async, connect-time.
583        assert!(matches!(
584            classify_one(T::mcp("node ./server.js")),
585            ToolResolution::Deferred(DeferredTool::Mcp { .. })
586        ));
587    }
588
589    #[test]
590    fn a_single_tool_function_converts_into_a_composite() {
591        let composite: ToolComposite =
592            SimpleTool::new("one", "one", None, |_| async { Ok(serde_json::json!(1)) }).into();
593        assert_eq!(composite.len(), 1);
594        let arc: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new("two", "two", None, |_| async {
595            Ok(serde_json::json!(2))
596        }));
597        let composite: ToolComposite = arc.into();
598        assert_eq!(composite.len(), 1);
599        assert_eq!(composite.entries[0].clone().classify_name(), "two");
600    }
601
602    #[tokio::test]
603    async fn mock_resolves_to_callable_runtime_tool() {
604        let resolution = classify_one(T::mock(
605            "weather",
606            "Mock weather",
607            serde_json::json!({"temp": 22}),
608        ));
609        let ToolResolution::Runtime(tool) = resolution else {
610            panic!("mock should resolve to a runtime tool");
611        };
612        assert_eq!(tool.name(), "weather");
613        let out = tool.call(serde_json::json!({})).await.unwrap();
614        assert_eq!(out, serde_json::json!({"temp": 22}));
615    }
616
617    #[tokio::test]
618    async fn transform_wraps_inner_runtime_result() {
619        let composite = T::transform(
620            T::mock("base", "d", serde_json::json!({"n": 1})),
621            |mut v| async move {
622                v["doubled"] = serde_json::json!(true);
623                v
624            },
625        );
626        let ToolResolution::Runtime(tool) = classify_one(composite) else {
627            panic!("transform over a mock should resolve to a runtime tool");
628        };
629        assert_eq!(tool.name(), "base");
630        let out = tool.call(serde_json::json!({})).await.unwrap();
631        assert_eq!(out, serde_json::json!({"n": 1, "doubled": true}));
632    }
633
634    #[test]
635    fn google_search_creates_composite() {
636        let t = T::google_search();
637        assert_eq!(t.len(), 1);
638    }
639
640    #[test]
641    fn url_context_creates_composite() {
642        let t = T::url_context();
643        assert_eq!(t.len(), 1);
644    }
645
646    #[test]
647    fn code_execution_creates_composite() {
648        let t = T::code_execution();
649        assert_eq!(t.len(), 1);
650    }
651
652    #[test]
653    fn compose_with_bitor() {
654        let t = T::google_search() + T::url_context() + T::code_execution();
655        assert_eq!(t.len(), 3);
656    }
657
658    #[test]
659    fn simple_creates_tool() {
660        let t = T::simple("greet", "Greets the user", |_args| async {
661            Ok(serde_json::json!({"message": "hello"}))
662        });
663        assert_eq!(t.len(), 1);
664        match &t.entries[0] {
665            ToolCompositeEntry::Function(f) => assert_eq!(f.name(), "greet"),
666            _ => panic!("expected Function entry"),
667        }
668    }
669
670    #[tokio::test]
671    async fn timeout_modifier_enforces_timeout() {
672        use gemini_adk_rs::ToolError;
673        use std::time::Duration;
674
675        let t = T::timeout(
676            T::simple("slow", "slow tool", |_| async move {
677                tokio::time::sleep(Duration::from_secs(3600)).await;
678                Ok(serde_json::json!({"ok": true}))
679            }),
680            Duration::from_millis(50),
681        );
682        match &t.entries[0] {
683            ToolCompositeEntry::Function(f) => match f.call(serde_json::json!({})).await {
684                Err(ToolError::Timeout(d)) => assert_eq!(d, Duration::from_millis(50)),
685                other => panic!("expected Timeout, got {other:?}"),
686            },
687            _ => panic!("expected Function entry"),
688        }
689    }
690
691    #[tokio::test]
692    async fn cached_modifier_memoizes_results() {
693        use std::sync::atomic::{AtomicU32, Ordering};
694
695        let counter = Arc::new(AtomicU32::new(0));
696        let c = counter.clone();
697        let t = T::cached(T::simple("count", "counts calls", move |_| {
698            let c = c.clone();
699            async move {
700                let n = c.fetch_add(1, Ordering::SeqCst) + 1;
701                Ok(serde_json::json!({"n": n}))
702            }
703        }));
704        match &t.entries[0] {
705            ToolCompositeEntry::Function(f) => {
706                let first = f.call(serde_json::json!({"x": 1})).await.unwrap();
707                let second = f.call(serde_json::json!({"x": 1})).await.unwrap();
708                assert_eq!(first, second);
709                assert_eq!(first["n"], 1);
710                assert_eq!(counter.load(Ordering::SeqCst), 1);
711            }
712            _ => panic!("expected Function entry"),
713        }
714    }
715
716    #[test]
717    fn confirm_modifier_wraps_function() {
718        // confirm() wraps the function (preserving its name) so the policy flag
719        // travels to the runtime rather than being silently dropped.
720        let t = T::confirm(
721            T::simple("danger", "dangerous", |_| async move {
722                Ok(serde_json::json!({}))
723            }),
724            "are you sure?",
725        );
726        match &t.entries[0] {
727            ToolCompositeEntry::Function(f) => assert_eq!(f.name(), "danger"),
728            _ => panic!("expected Function entry"),
729        }
730    }
731
732    #[test]
733    fn toolset_combines_functions() {
734        let tool_a: Arc<dyn ToolFunction> =
735            Arc::new(SimpleTool::new("a", "tool a", None, |_| async {
736                Ok(serde_json::json!(null))
737            }));
738        let tool_b: Arc<dyn ToolFunction> =
739            Arc::new(SimpleTool::new("b", "tool b", None, |_| async {
740                Ok(serde_json::json!(null))
741            }));
742        let t = T::toolset(vec![tool_a, tool_b]);
743        assert_eq!(t.len(), 2);
744    }
745}