klieo-tools 3.2.0

//! `ChainedInvoker` — name-based tool dispatch with JSON-schema arg
//! validation and per-tool timeout enforcement.
//!
//! Each `invoke` call:
//! 1. Looks up the tool by name (returns [`klieo_core::ToolError::UnknownTool`] on miss).
//! 2. Validates the JSON arguments against the tool's declared `json_schema` via
//!    [`crate::validation::validate_args`] (returns [`klieo_core::ToolError::InvalidArgs`]
//!    on mismatch, [`klieo_core::ToolError::Permanent`] on bad schema).
//! 3. Runs `tool.invoke` inside `tokio::time::timeout` (returns
//!    [`klieo_core::ToolError::Timeout`] on overrun).

use async_trait::async_trait;
use klieo_core::error::ToolError;
use klieo_core::llm::ToolDef;
use klieo_core::tool::{Tool, ToolCtx, ToolInvoker};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

/// Errors returned by [`ChainedInvoker`] construction methods.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum InvokerError {
    /// Two tools with the same name were registered.
    #[error("duplicate tool name: {0:?}")]
    DuplicateTool(String),
}

const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(30);

/// Tool invoker that dispatches calls to a registered set of `Tool`
/// impls by name. Validates args against each tool's `json_schema`
/// before invocation, and enforces a per-invocation timeout.
pub struct ChainedInvoker {
    tools: HashMap<String, Arc<dyn Tool>>,
    catalogue: Vec<ToolDef>,
    default_timeout: Duration,
    per_tool_timeouts: HashMap<String, Duration>,
}

impl ChainedInvoker {
    /// Build an empty invoker with a 30-second default timeout.
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
            catalogue: Vec::new(),
            default_timeout: DEFAULT_TOOL_TIMEOUT,
            per_tool_timeouts: HashMap::new(),
        }
    }

    /// Override the default per-tool timeout. Default: 30 seconds.
    pub fn with_default_timeout(mut self, t: Duration) -> Self {
        self.default_timeout = t;
        self
    }

    /// Override the timeout for a specific tool by name. Repeated calls
    /// for the same `tool_name` keep the latest. The override applies
    /// regardless of registration order — the tool need not be added
    /// before the override.
    pub fn with_per_tool_timeout(mut self, tool_name: impl Into<String>, dur: Duration) -> Self {
        self.per_tool_timeouts.insert(tool_name.into(), dur);
        self
    }

    /// Register a tool by mutating the invoker in place. Useful when
    /// building from a dynamic source (e.g. iterating a `Vec`). Returns
    /// [`InvokerError::DuplicateTool`] when a tool with the same name is
    /// already registered.
    pub fn add_tool(&mut self, tool: Arc<dyn Tool>) -> Result<(), InvokerError> {
        let name = tool.name().to_string();
        let def = ToolDef::new(name.clone(), tool.description(), tool.json_schema().clone());
        if self.tools.insert(name.clone(), tool).is_some() {
            return Err(InvokerError::DuplicateTool(name));
        }
        self.catalogue.push(def);
        Ok(())
    }

    /// Register a tool, consuming and returning self. Convenient for
    /// inline chaining: `ChainedInvoker::new().with_tool(a)?.with_tool(b)?`.
    /// Delegates to [`Self::add_tool`].
    pub fn with_tool(mut self, tool: Arc<dyn Tool>) -> Result<Self, InvokerError> {
        self.add_tool(tool)?;
        Ok(self)
    }

    /// Returns [`InvokerError::DuplicateTool`] when a tool with the
    /// same name is already registered.
    pub fn with_tool_owned<T: Tool + 'static>(self, tool: T) -> Result<Self, InvokerError> {
        self.with_tool(Arc::new(tool))
    }
}

impl Default for ChainedInvoker {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for ChainedInvoker {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ChainedInvoker")
            .field(
                "tool_names",
                &self.catalogue.iter().map(|d| &d.name).collect::<Vec<_>>(),
            )
            .field("default_timeout", &self.default_timeout)
            .finish()
    }
}

#[async_trait]
impl ToolInvoker for ChainedInvoker {
    #[tracing::instrument(level = "debug", skip(self, args, ctx), fields(tool = %name))]
    async fn invoke(
        &self,
        name: &str,
        args: serde_json::Value,
        ctx: ToolCtx,
    ) -> Result<serde_json::Value, ToolError> {
        let tool = self
            .tools
            .get(name)
            .ok_or_else(|| ToolError::UnknownTool(name.to_string()))?;
        crate::validation::validate_args(tool.json_schema(), &args)?;
        let timeout = self
            .per_tool_timeouts
            .get(name)
            .copied()
            .unwrap_or(self.default_timeout);
        match tokio::time::timeout(timeout, tool.invoke(args, ctx)).await {
            Ok(res) => res,
            Err(_) => Err(ToolError::Timeout),
        }
    }

    fn catalogue(&self) -> Vec<ToolDef> {
        self.catalogue.clone()
    }

    fn tool_redacts_audit(&self, name: &str) -> bool {
        self.tools.get(name).is_some_and(|t| t.redacts_audit())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use klieo_bus_memory::MemoryBus;
    use klieo_core::error::ToolError;
    use klieo_core::tool::{Tool, ToolCtx, ToolInvoker};
    use std::sync::Arc;

    struct EchoTool;

    #[async_trait]
    impl Tool for EchoTool {
        fn name(&self) -> &str {
            "echo"
        }
        fn description(&self) -> &str {
            "echo args back"
        }
        fn json_schema(&self) -> &serde_json::Value {
            static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
            SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
        }
        async fn invoke(
            &self,
            args: serde_json::Value,
            _ctx: ToolCtx,
        ) -> Result<serde_json::Value, ToolError> {
            Ok(args)
        }
    }

    struct StrictTool;

    #[async_trait]
    impl Tool for StrictTool {
        fn name(&self) -> &str {
            "strict"
        }
        fn description(&self) -> &str {
            "requires query field"
        }
        fn json_schema(&self) -> &serde_json::Value {
            static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
            SCHEMA.get_or_init(|| {
                serde_json::json!({
                    "type": "object",
                    "properties": { "query": { "type": "string" } },
                    "required": ["query"],
                    "additionalProperties": false,
                })
            })
        }
        async fn invoke(
            &self,
            args: serde_json::Value,
            _ctx: ToolCtx,
        ) -> Result<serde_json::Value, ToolError> {
            Ok(args)
        }
    }

    struct PiiTool;

    #[async_trait]
    impl Tool for PiiTool {
        fn name(&self) -> &str {
            "claimant_lookup"
        }
        fn description(&self) -> &str {
            "handles claimant PII"
        }
        fn json_schema(&self) -> &serde_json::Value {
            static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
            SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
        }
        fn redacts_audit(&self) -> bool {
            true
        }
        async fn invoke(
            &self,
            args: serde_json::Value,
            _ctx: ToolCtx,
        ) -> Result<serde_json::Value, ToolError> {
            Ok(args)
        }
    }

    fn ctx() -> ToolCtx {
        let bus = MemoryBus::new();
        ToolCtx::new(bus.pubsub, bus.kv, bus.jobs)
    }

    #[tokio::test]
    async fn with_tool_owned_arc_wraps_internally_and_dispatches() {
        let inv = ChainedInvoker::new().with_tool_owned(EchoTool).unwrap();
        let out = inv
            .invoke("echo", serde_json::json!({"x": 7}), ctx())
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"x": 7}));
    }

    #[tokio::test]
    async fn unknown_tool_returns_unknown_tool_error() {
        let inv = ChainedInvoker::new();
        let err = inv
            .invoke("nope", serde_json::json!({}), ctx())
            .await
            .unwrap_err();
        assert!(matches!(err, ToolError::UnknownTool(name) if name == "nope"));
    }

    #[tokio::test]
    async fn registered_tool_is_dispatched() {
        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
        let out = inv
            .invoke("echo", serde_json::json!({"x": 1}), ctx())
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"x": 1}));
    }

    #[tokio::test]
    async fn catalogue_lists_registered_tools() {
        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
        let cat = inv.catalogue();
        assert_eq!(cat.len(), 1);
        assert_eq!(cat[0].name, "echo");
        assert_eq!(cat[0].description, "echo args back");
    }

    #[tokio::test]
    async fn pii_flagged_tool_makes_invoker_report_redacts_audit() {
        let inv = ChainedInvoker::new().with_tool(Arc::new(PiiTool)).unwrap();
        assert!(inv.tool_redacts_audit("claimant_lookup"));
    }

    #[tokio::test]
    async fn non_flagged_tool_does_not_report_redacts_audit() {
        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
        assert!(!inv.tool_redacts_audit("echo"));
    }

    #[tokio::test]
    async fn unregistered_name_does_not_report_redacts_audit() {
        let inv = ChainedInvoker::new().with_tool(Arc::new(PiiTool)).unwrap();
        assert!(!inv.tool_redacts_audit("never_registered"));
    }

    #[tokio::test]
    async fn invalid_args_rejected_before_invocation() {
        let inv = ChainedInvoker::new()
            .with_tool(Arc::new(StrictTool))
            .unwrap();
        // Missing required "query" field.
        let err = inv
            .invoke("strict", serde_json::json!({}), ctx())
            .await
            .unwrap_err();
        assert!(matches!(err, ToolError::InvalidArgs(_)));
    }

    #[tokio::test]
    async fn valid_args_pass_through_to_tool() {
        let inv = ChainedInvoker::new()
            .with_tool(Arc::new(StrictTool))
            .unwrap();
        let out = inv
            .invoke("strict", serde_json::json!({"query": "hello"}), ctx())
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"query": "hello"}));
    }

    #[tokio::test]
    async fn validation_short_circuits_invocation() {
        // A tool whose body increments a counter.
        struct CountingTool {
            counter: Arc<std::sync::atomic::AtomicU32>,
        }
        #[async_trait]
        impl Tool for CountingTool {
            fn name(&self) -> &str {
                "counting"
            }
            fn description(&self) -> &str {
                ""
            }
            fn json_schema(&self) -> &serde_json::Value {
                static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
                SCHEMA.get_or_init(|| {
                    serde_json::json!({
                        "type": "object",
                        "required": ["q"],
                    })
                })
            }
            async fn invoke(
                &self,
                _args: serde_json::Value,
                _ctx: ToolCtx,
            ) -> Result<serde_json::Value, ToolError> {
                self.counter
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                Ok(serde_json::Value::Null)
            }
        }
        let counter = Arc::new(std::sync::atomic::AtomicU32::new(0));
        let inv = ChainedInvoker::new()
            .with_tool(Arc::new(CountingTool {
                counter: counter.clone(),
            }))
            .unwrap();
        let _ = inv.invoke("counting", serde_json::json!({}), ctx()).await;
        assert_eq!(
            counter.load(std::sync::atomic::Ordering::Relaxed),
            0,
            "tool body must not run on validation failure"
        );
    }

    #[tokio::test]
    async fn slow_tool_times_out() {
        struct SlowTool;
        #[async_trait]
        impl Tool for SlowTool {
            fn name(&self) -> &str {
                "slow"
            }
            fn description(&self) -> &str {
                ""
            }
            fn json_schema(&self) -> &serde_json::Value {
                static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
                SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
            }
            async fn invoke(
                &self,
                _args: serde_json::Value,
                _ctx: ToolCtx,
            ) -> Result<serde_json::Value, ToolError> {
                tokio::time::sleep(std::time::Duration::from_secs(60)).await;
                Ok(serde_json::Value::Null)
            }
        }
        let inv = ChainedInvoker::new()
            .with_default_timeout(std::time::Duration::from_millis(20))
            .with_tool(Arc::new(SlowTool))
            .unwrap();
        let err = inv
            .invoke("slow", serde_json::json!({}), ctx())
            .await
            .unwrap_err();
        assert!(matches!(err, ToolError::Timeout));
    }

    #[tokio::test]
    async fn fast_tool_does_not_time_out() {
        let inv = ChainedInvoker::new()
            .with_default_timeout(std::time::Duration::from_millis(50))
            .with_tool(Arc::new(EchoTool))
            .unwrap();
        let out = inv
            .invoke("echo", serde_json::json!({"k": "v"}), ctx())
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"k": "v"}));
    }

    #[test]
    fn with_tool_returns_err_on_duplicate_name() {
        let err = ChainedInvoker::new()
            .with_tool(Arc::new(EchoTool))
            .unwrap()
            .with_tool(Arc::new(EchoTool))
            .unwrap_err();
        assert!(
            matches!(err, InvokerError::DuplicateTool(ref n) if n == "echo"),
            "expected DuplicateTool(\"echo\"), got {err:?}"
        );
    }

    #[test]
    fn add_tool_returns_err_on_duplicate_name() {
        let mut inv = ChainedInvoker::new();
        inv.add_tool(Arc::new(EchoTool)).unwrap();
        let err = inv.add_tool(Arc::new(EchoTool)).unwrap_err();
        assert!(
            matches!(err, InvokerError::DuplicateTool(ref n) if n == "echo"),
            "expected DuplicateTool(\"echo\"), got {err:?}"
        );
    }

    #[tokio::test]
    async fn catalogue_returns_independent_clones() {
        let inv = ChainedInvoker::new().with_tool(Arc::new(EchoTool)).unwrap();
        let mut cat = inv.catalogue();
        cat.clear();
        // Mutating the returned vec must not affect the invoker.
        let cat2 = inv.catalogue();
        assert_eq!(cat2.len(), 1);
    }

    #[tokio::test]
    async fn add_tool_mutates_in_place() {
        let mut inv = ChainedInvoker::new();
        inv.add_tool(Arc::new(EchoTool)).unwrap();
        let cat = inv.catalogue();
        assert_eq!(cat.len(), 1);
        assert_eq!(cat[0].name, "echo");
    }

    #[tokio::test]
    async fn per_tool_timeout_overrides_default_for_named_tool() {
        struct SlowTool;
        #[async_trait]
        impl Tool for SlowTool {
            fn name(&self) -> &str {
                "slow"
            }
            fn description(&self) -> &str {
                ""
            }
            fn json_schema(&self) -> &serde_json::Value {
                static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
                SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
            }
            async fn invoke(
                &self,
                _args: serde_json::Value,
                _ctx: ToolCtx,
            ) -> Result<serde_json::Value, ToolError> {
                tokio::time::sleep(std::time::Duration::from_millis(60)).await;
                Ok(serde_json::Value::Null)
            }
        }
        // Default 5s, but per-tool override of 10ms forces a timeout.
        let inv = ChainedInvoker::new()
            .with_default_timeout(std::time::Duration::from_secs(5))
            .with_per_tool_timeout("slow", std::time::Duration::from_millis(10))
            .with_tool(Arc::new(SlowTool))
            .unwrap();
        let err = inv
            .invoke("slow", serde_json::json!({}), ctx())
            .await
            .unwrap_err();
        assert!(matches!(err, ToolError::Timeout));
    }

    #[tokio::test]
    async fn per_tool_timeout_does_not_apply_to_other_tools() {
        // EchoTool keeps the generous default; per-tool override only on "slow".
        let inv = ChainedInvoker::new()
            .with_default_timeout(std::time::Duration::from_secs(5))
            .with_per_tool_timeout("slow", std::time::Duration::from_millis(1))
            .with_tool(Arc::new(EchoTool))
            .unwrap();
        let out = inv
            .invoke("echo", serde_json::json!({"x": 1}), ctx())
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"x": 1}));
    }

    #[tokio::test]
    async fn per_tool_timeout_repeated_call_keeps_latest() {
        let inv = ChainedInvoker::new()
            .with_default_timeout(std::time::Duration::from_secs(5))
            .with_per_tool_timeout("echo", std::time::Duration::from_millis(1))
            .with_per_tool_timeout("echo", std::time::Duration::from_secs(5))
            .with_tool(Arc::new(EchoTool))
            .unwrap();
        // The second override (5s) wins; echo runs cleanly.
        let out = inv
            .invoke("echo", serde_json::json!({"y": 2}), ctx())
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"y": 2}));
    }
}