klieo_tools/
invoker.rs

1//! `ChainedInvoker` — name-based tool dispatch with JSON-schema arg
2//! validation and per-tool timeout enforcement.
3//!
4//! Each `invoke` call:
5//! 1. Looks up the tool by name (returns [`klieo_core::ToolError::UnknownTool`] on miss).
6//! 2. Validates the JSON arguments against the tool's declared `json_schema` via
7//!    [`crate::validation::validate_args`] (returns [`klieo_core::ToolError::InvalidArgs`]
8//!    on mismatch, [`klieo_core::ToolError::Permanent`] on bad schema).
9//! 3. Runs `tool.invoke` inside `tokio::time::timeout` (returns
10//!    [`klieo_core::ToolError::Timeout`] on overrun).
11
12use async_trait::async_trait;
13use klieo_core::error::ToolError;
14use klieo_core::llm::ToolDef;
15use klieo_core::tool::{Tool, ToolCtx, ToolInvoker};
16use std::collections::HashMap;
17use std::sync::Arc;
18use std::time::Duration;
19
20/// Errors returned by [`ChainedInvoker`] construction methods.
21#[derive(Debug, thiserror::Error)]
22#[non_exhaustive]
23pub enum InvokerError {
24    /// Two tools with the same name were registered.
25    #[error("duplicate tool name: {0:?}")]
26    DuplicateTool(String),
27}
28
29const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(30);
30
31/// Tool invoker that dispatches calls to a registered set of `Tool`
32/// impls by name. Validates args against each tool's `json_schema`
33/// before invocation, and enforces a per-invocation timeout.
34pub struct ChainedInvoker {
35    tools: HashMap<String, Arc<dyn Tool>>,
36    catalogue: Vec<ToolDef>,
37    default_timeout: Duration,
38    per_tool_timeouts: HashMap<String, Duration>,
39}
40
41impl ChainedInvoker {
42    /// Build an empty invoker with a 30-second default timeout.
43    pub fn new() -> Self {
44        Self {
45            tools: HashMap::new(),
46            catalogue: Vec::new(),
47            default_timeout: DEFAULT_TOOL_TIMEOUT,
48            per_tool_timeouts: HashMap::new(),
49        }
50    }
51
52    /// Override the default per-tool timeout. Default: 30 seconds.
53    pub fn with_default_timeout(mut self, t: Duration) -> Self {
54        self.default_timeout = t;
55        self
56    }
57
58    /// Override the timeout for a specific tool by name. Repeated calls
59    /// for the same `tool_name` keep the latest. The override applies
60    /// regardless of registration order — the tool need not be added
61    /// before the override.
62    pub fn with_per_tool_timeout(mut self, tool_name: impl Into<String>, dur: Duration) -> Self {
63        self.per_tool_timeouts.insert(tool_name.into(), dur);
64        self
65    }
66
67    /// Register a tool by mutating the invoker in place. Useful when
68    /// building from a dynamic source (e.g. iterating a `Vec`). Returns
69    /// [`InvokerError::DuplicateTool`] when a tool with the same name is
70    /// already registered.
71    pub fn add_tool(&mut self, tool: Arc<dyn Tool>) -> Result<(), InvokerError> {
72        let name = tool.name().to_string();
73        let def = ToolDef::new(name.clone(), tool.description(), tool.json_schema().clone());
74        if self.tools.insert(name.clone(), tool).is_some() {
75            return Err(InvokerError::DuplicateTool(name));
76        }
77        self.catalogue.push(def);
78        Ok(())
79    }
80
81    /// Register a tool, consuming and returning self. Convenient for
82    /// inline chaining: `ChainedInvoker::new().with_tool(a)?.with_tool(b)?`.
83    /// Delegates to [`Self::add_tool`].
84    pub fn with_tool(mut self, tool: Arc<dyn Tool>) -> Result<Self, InvokerError> {
85        self.add_tool(tool)?;
86        Ok(self)
87    }
88
89    /// Returns [`InvokerError::DuplicateTool`] when a tool with the
90    /// same name is already registered.
91    pub fn with_tool_owned<T: Tool + 'static>(self, tool: T) -> Result<Self, InvokerError> {
92        self.with_tool(Arc::new(tool))
93    }
94}
95
96impl Default for ChainedInvoker {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102impl std::fmt::Debug for ChainedInvoker {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        f.debug_struct("ChainedInvoker")
105            .field(
106                "tool_names",
107                &self.catalogue.iter().map(|d| &d.name).collect::<Vec<_>>(),
108            )
109            .field("default_timeout", &self.default_timeout)
110            .finish()
111    }
112}
113
114#[async_trait]
115impl ToolInvoker for ChainedInvoker {
116    #[tracing::instrument(level = "debug", skip(self, args, ctx), fields(tool = %name))]
117    async fn invoke(
118        &self,
119        name: &str,
120        args: serde_json::Value,
121        ctx: ToolCtx,
122    ) -> Result<serde_json::Value, ToolError> {
123        let tool = self
124            .tools
125            .get(name)
126            .ok_or_else(|| ToolError::UnknownTool(name.to_string()))?;
127        crate::validation::validate_args(tool.json_schema(), &args)?;
128        let timeout = self
129            .per_tool_timeouts
130            .get(name)
131            .copied()
132            .unwrap_or(self.default_timeout);
133        match tokio::time::timeout(timeout, tool.invoke(args, ctx)).await {
134            Ok(res) => res,
135            Err(_) => Err(ToolError::Timeout),
136        }
137    }
138
139    fn catalogue(&self) -> Vec<ToolDef> {
140        self.catalogue.clone()
141    }
142
143    fn tool_redacts_audit(&self, name: &str) -> bool {
144        self.tools.get(name).is_some_and(|t| t.redacts_audit())
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use async_trait::async_trait;
152    use klieo_bus_memory::MemoryBus;
153    use klieo_core::error::ToolError;
154    use klieo_core::tool::{Tool, ToolCtx, ToolInvoker};
155    use std::sync::Arc;
156
157    struct EchoTool;
158
159    #[async_trait]
160    impl Tool for EchoTool {
161        fn name(&self) -> &str {
162            "echo"
163        }
164        fn description(&self) -> &str {
165            "echo args back"
166        }
167        fn json_schema(&self) -> &serde_json::Value {
168            static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
169            SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
170        }
171        async fn invoke(
172            &self,
173            args: serde_json::Value,
174            _ctx: ToolCtx,
175        ) -> Result<serde_json::Value, ToolError> {
176            Ok(args)
177        }
178    }
179
180    struct StrictTool;
181
182    #[async_trait]
183    impl Tool for StrictTool {
184        fn name(&self) -> &str {
185            "strict"
186        }
187        fn description(&self) -> &str {
188            "requires query field"
189        }
190        fn json_schema(&self) -> &serde_json::Value {
191            static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
192            SCHEMA.get_or_init(|| {
193                serde_json::json!({
194                    "type": "object",
195                    "properties": { "query": { "type": "string" } },
196                    "required": ["query"],
197                    "additionalProperties": false,
198                })
199            })
200        }
201        async fn invoke(
202            &self,
203            args: serde_json::Value,
204            _ctx: ToolCtx,
205        ) -> Result<serde_json::Value, ToolError> {
206            Ok(args)
207        }
208    }
209
210    struct PiiTool;
211
212    #[async_trait]
213    impl Tool for PiiTool {
214        fn name(&self) -> &str {
215            "claimant_lookup"
216        }
217        fn description(&self) -> &str {
218            "handles claimant PII"
219        }
220        fn json_schema(&self) -> &serde_json::Value {
221            static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
222            SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
223        }
224        fn redacts_audit(&self) -> bool {
225            true
226        }
227        async fn invoke(
228            &self,
229            args: serde_json::Value,
230            _ctx: ToolCtx,
231        ) -> Result<serde_json::Value, ToolError> {
232            Ok(args)
233        }
234    }
235
236    fn ctx() -> ToolCtx {
237        let bus = MemoryBus::new();
238        ToolCtx::new(bus.pubsub, bus.kv, bus.jobs)
239    }
240
241    #[tokio::test]
242    async fn with_tool_owned_arc_wraps_internally_and_dispatches() {
243        let inv = ChainedInvoker::new().with_tool_owned(EchoTool).unwrap();
244        let out = inv
245            .invoke("echo", serde_json::json!({"x": 7}), ctx())
246            .await
247            .unwrap();
248        assert_eq!(out, serde_json::json!({"x": 7}));
249    }
250
251    #[tokio::test]
252    async fn unknown_tool_returns_unknown_tool_error() {
253        let inv = ChainedInvoker::new();
254        let err = inv
255            .invoke("nope", serde_json::json!({}), ctx())
256            .await
257            .unwrap_err();
258        assert!(matches!(err, ToolError::UnknownTool(name) if name == "nope"));
259    }
260
261    #[tokio::test]
262    async fn registered_tool_is_dispatched() {
263        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
264        let out = inv
265            .invoke("echo", serde_json::json!({"x": 1}), ctx())
266            .await
267            .unwrap();
268        assert_eq!(out, serde_json::json!({"x": 1}));
269    }
270
271    #[tokio::test]
272    async fn catalogue_lists_registered_tools() {
273        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
274        let cat = inv.catalogue();
275        assert_eq!(cat.len(), 1);
276        assert_eq!(cat[0].name, "echo");
277        assert_eq!(cat[0].description, "echo args back");
278    }
279
280    #[tokio::test]
281    async fn pii_flagged_tool_makes_invoker_report_redacts_audit() {
282        let inv = ChainedInvoker::new().with_tool(Arc::new(PiiTool)).unwrap();
283        assert!(inv.tool_redacts_audit("claimant_lookup"));
284    }
285
286    #[tokio::test]
287    async fn non_flagged_tool_does_not_report_redacts_audit() {
288        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
289        assert!(!inv.tool_redacts_audit("echo"));
290    }
291
292    #[tokio::test]
293    async fn unregistered_name_does_not_report_redacts_audit() {
294        let inv = ChainedInvoker::new().with_tool(Arc::new(PiiTool)).unwrap();
295        assert!(!inv.tool_redacts_audit("never_registered"));
296    }
297
298    #[tokio::test]
299    async fn invalid_args_rejected_before_invocation() {
300        let inv = ChainedInvoker::new()
301            .with_tool(Arc::new(StrictTool))
302            .unwrap();
303        // Missing required "query" field.
304        let err = inv
305            .invoke("strict", serde_json::json!({}), ctx())
306            .await
307            .unwrap_err();
308        assert!(matches!(err, ToolError::InvalidArgs(_)));
309    }
310
311    #[tokio::test]
312    async fn valid_args_pass_through_to_tool() {
313        let inv = ChainedInvoker::new()
314            .with_tool(Arc::new(StrictTool))
315            .unwrap();
316        let out = inv
317            .invoke("strict", serde_json::json!({"query": "hello"}), ctx())
318            .await
319            .unwrap();
320        assert_eq!(out, serde_json::json!({"query": "hello"}));
321    }
322
323    #[tokio::test]
324    async fn validation_short_circuits_invocation() {
325        // A tool whose body increments a counter.
326        struct CountingTool {
327            counter: Arc<std::sync::atomic::AtomicU32>,
328        }
329        #[async_trait]
330        impl Tool for CountingTool {
331            fn name(&self) -> &str {
332                "counting"
333            }
334            fn description(&self) -> &str {
335                ""
336            }
337            fn json_schema(&self) -> &serde_json::Value {
338                static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
339                SCHEMA.get_or_init(|| {
340                    serde_json::json!({
341                        "type": "object",
342                        "required": ["q"],
343                    })
344                })
345            }
346            async fn invoke(
347                &self,
348                _args: serde_json::Value,
349                _ctx: ToolCtx,
350            ) -> Result<serde_json::Value, ToolError> {
351                self.counter
352                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
353                Ok(serde_json::Value::Null)
354            }
355        }
356        let counter = Arc::new(std::sync::atomic::AtomicU32::new(0));
357        let inv = ChainedInvoker::new()
358            .with_tool(Arc::new(CountingTool {
359                counter: counter.clone(),
360            }))
361            .unwrap();
362        let _ = inv.invoke("counting", serde_json::json!({}), ctx()).await;
363        assert_eq!(
364            counter.load(std::sync::atomic::Ordering::Relaxed),
365            0,
366            "tool body must not run on validation failure"
367        );
368    }
369
370    #[tokio::test]
371    async fn slow_tool_times_out() {
372        struct SlowTool;
373        #[async_trait]
374        impl Tool for SlowTool {
375            fn name(&self) -> &str {
376                "slow"
377            }
378            fn description(&self) -> &str {
379                ""
380            }
381            fn json_schema(&self) -> &serde_json::Value {
382                static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
383                SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
384            }
385            async fn invoke(
386                &self,
387                _args: serde_json::Value,
388                _ctx: ToolCtx,
389            ) -> Result<serde_json::Value, ToolError> {
390                tokio::time::sleep(std::time::Duration::from_secs(60)).await;
391                Ok(serde_json::Value::Null)
392            }
393        }
394        let inv = ChainedInvoker::new()
395            .with_default_timeout(std::time::Duration::from_millis(20))
396            .with_tool(Arc::new(SlowTool))
397            .unwrap();
398        let err = inv
399            .invoke("slow", serde_json::json!({}), ctx())
400            .await
401            .unwrap_err();
402        assert!(matches!(err, ToolError::Timeout));
403    }
404
405    #[tokio::test]
406    async fn fast_tool_does_not_time_out() {
407        let inv = ChainedInvoker::new()
408            .with_default_timeout(std::time::Duration::from_millis(50))
409            .with_tool(Arc::new(EchoTool))
410            .unwrap();
411        let out = inv
412            .invoke("echo", serde_json::json!({"k": "v"}), ctx())
413            .await
414            .unwrap();
415        assert_eq!(out, serde_json::json!({"k": "v"}));
416    }
417
418    #[test]
419    fn with_tool_returns_err_on_duplicate_name() {
420        let err = ChainedInvoker::new()
421            .with_tool(Arc::new(EchoTool))
422            .unwrap()
423            .with_tool(Arc::new(EchoTool))
424            .unwrap_err();
425        assert!(
426            matches!(err, InvokerError::DuplicateTool(ref n) if n == "echo"),
427            "expected DuplicateTool(\"echo\"), got {err:?}"
428        );
429    }
430
431    #[test]
432    fn add_tool_returns_err_on_duplicate_name() {
433        let mut inv = ChainedInvoker::new();
434        inv.add_tool(Arc::new(EchoTool)).unwrap();
435        let err = inv.add_tool(Arc::new(EchoTool)).unwrap_err();
436        assert!(
437            matches!(err, InvokerError::DuplicateTool(ref n) if n == "echo"),
438            "expected DuplicateTool(\"echo\"), got {err:?}"
439        );
440    }
441
442    #[tokio::test]
443    async fn catalogue_returns_independent_clones() {
444        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
445        let mut cat = inv.catalogue();
446        cat.clear();
447        // Mutating the returned vec must not affect the invoker.
448        let cat2 = inv.catalogue();
449        assert_eq!(cat2.len(), 1);
450    }
451
452    #[tokio::test]
453    async fn add_tool_mutates_in_place() {
454        let mut inv = ChainedInvoker::new();
455        inv.add_tool(Arc::new(EchoTool)).unwrap();
456        let cat = inv.catalogue();
457        assert_eq!(cat.len(), 1);
458        assert_eq!(cat[0].name, "echo");
459    }
460
461    #[tokio::test]
462    async fn per_tool_timeout_overrides_default_for_named_tool() {
463        struct SlowTool;
464        #[async_trait]
465        impl Tool for SlowTool {
466            fn name(&self) -> &str {
467                "slow"
468            }
469            fn description(&self) -> &str {
470                ""
471            }
472            fn json_schema(&self) -> &serde_json::Value {
473                static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
474                SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
475            }
476            async fn invoke(
477                &self,
478                _args: serde_json::Value,
479                _ctx: ToolCtx,
480            ) -> Result<serde_json::Value, ToolError> {
481                tokio::time::sleep(std::time::Duration::from_millis(60)).await;
482                Ok(serde_json::Value::Null)
483            }
484        }
485        // Default 5s, but per-tool override of 10ms forces a timeout.
486        let inv = ChainedInvoker::new()
487            .with_default_timeout(std::time::Duration::from_secs(5))
488            .with_per_tool_timeout("slow", std::time::Duration::from_millis(10))
489            .with_tool(Arc::new(SlowTool))
490            .unwrap();
491        let err = inv
492            .invoke("slow", serde_json::json!({}), ctx())
493            .await
494            .unwrap_err();
495        assert!(matches!(err, ToolError::Timeout));
496    }
497
498    #[tokio::test]
499    async fn per_tool_timeout_does_not_apply_to_other_tools() {
500        // EchoTool keeps the generous default; per-tool override only on "slow".
501        let inv = ChainedInvoker::new()
502            .with_default_timeout(std::time::Duration::from_secs(5))
503            .with_per_tool_timeout("slow", std::time::Duration::from_millis(1))
504            .with_tool(Arc::new(EchoTool))
505            .unwrap();
506        let out = inv
507            .invoke("echo", serde_json::json!({"x": 1}), ctx())
508            .await
509            .unwrap();
510        assert_eq!(out, serde_json::json!({"x": 1}));
511    }
512
513    #[tokio::test]
514    async fn per_tool_timeout_repeated_call_keeps_latest() {
515        let inv = ChainedInvoker::new()
516            .with_default_timeout(std::time::Duration::from_secs(5))
517            .with_per_tool_timeout("echo", std::time::Duration::from_millis(1))
518            .with_per_tool_timeout("echo", std::time::Duration::from_secs(5))
519            .with_tool(Arc::new(EchoTool))
520            .unwrap();
521        // The second override (5s) wins; echo runs cleanly.
522        let out = inv
523            .invoke("echo", serde_json::json!({"y": 2}), ctx())
524            .await
525            .unwrap();
526        assert_eq!(out, serde_json::json!({"y": 2}));
527    }
528}
klieo_tools/invoker.rs

klieo_tools/
invoker.rs
