klieo-core 0.41.2

//! Tool trait, invoker, and tool context.
//!
//! Tools are I/O leaves. They receive a narrowed [`ToolCtx`] holding only
//! bus + KV access — they cannot reach into LLM or memory directly.
//! Memory updates triggered by tools must flow through bus events that
//! agents subscribe to, keeping the dependency graph one-directional.

use crate::bus::{JobQueue, KvStore, Pubsub};
use crate::error::ToolError;
use async_trait::async_trait;
use std::sync::Arc;

/// Narrowed slice of `AgentContext` (see `agent.rs`) passed to tools.
#[derive(Clone)]
#[non_exhaustive]
pub struct ToolCtx {
    /// Pub/sub bus.
    pub pubsub: Arc<dyn Pubsub>,
    /// KV store.
    pub kv: Arc<dyn KvStore>,
    /// Job queue.
    pub jobs: Arc<dyn JobQueue>,
    /// Optional progress sender. Tools that wrap an [`crate::Agent`]
    /// (e.g. `AgentAsToolInvoker` in klieo-mcp-server) overlay this
    /// onto the minted `AgentContext`. Non-agent invokers ignore the
    /// field.
    pub progress: Option<tokio::sync::broadcast::Sender<crate::AgentEvent>>,
    /// Cooperative cancellation token. HTTP/SSE transports pass a
    /// request-scoped child token derived from the server's parent
    /// cancel; tools should `select!` against `cancel.cancelled()` to
    /// honor client disconnects mid-call.
    pub cancel: tokio_util::sync::CancellationToken,
    /// Optional outbound channel for server-initiated JSON-RPC
    /// (e.g. MCP `sampling/createMessage`). Tools requiring the
    /// outbound surface return `Err(Unsupported)` when `None`.
    /// Default: `None`.
    pub server_outbound: Option<Arc<dyn crate::ServerOutbound>>,
}

impl ToolCtx {
    /// Construct a `ToolCtx` with default `progress = None`, no
    /// outbound channel, and a fresh uncancelled `CancellationToken`.
    /// Builders below override.
    pub fn new(pubsub: Arc<dyn Pubsub>, kv: Arc<dyn KvStore>, jobs: Arc<dyn JobQueue>) -> Self {
        Self {
            pubsub,
            kv,
            jobs,
            progress: None,
            cancel: tokio_util::sync::CancellationToken::new(),
            server_outbound: None,
        }
    }

    /// Override the cancellation token (HTTP/SSE callers pass a
    /// request-scoped child token).
    #[must_use]
    pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
        self.cancel = cancel;
        self
    }

    /// Override the progress sender (HTTP/SSE callers pass the
    /// per-request broadcast for `AgentEvent` forwarding).
    #[must_use]
    pub fn with_progress(
        mut self,
        progress: tokio::sync::broadcast::Sender<crate::AgentEvent>,
    ) -> Self {
        self.progress = Some(progress);
        self
    }

    /// Attach a server-initiated outbound channel so tools can issue
    /// reverse-direction JSON-RPC requests to the connected peer.
    #[must_use]
    pub fn with_server_outbound(mut self, outbound: Arc<dyn crate::ServerOutbound>) -> Self {
        self.server_outbound = Some(outbound);
        self
    }
}

/// One executable tool.
///
/// Implementations live wherever they make sense. The
/// `klieo_macros::tool` proc-macro generates one for you from a
/// plain async function.
///
/// ```
/// # tokio_test::block_on(async {
/// use async_trait::async_trait;
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::error::ToolError;
/// use klieo_core::tool::{Tool, ToolCtx};
/// use std::sync::OnceLock;
///
/// struct Echo;
///
/// #[async_trait]
/// impl Tool for Echo {
///     fn name(&self) -> &str { "echo" }
///     fn description(&self) -> &str { "echoes input" }
///     fn json_schema(&self) -> &serde_json::Value {
///         static S: OnceLock<serde_json::Value> = OnceLock::new();
///         S.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)
///     }
/// }
///
/// let (pubsub, _, kv, jobs) = noop_bus();
/// let ctx = ToolCtx::new(pubsub, kv, jobs);
/// let out = Echo.invoke(serde_json::json!({"y": 2}), ctx).await.unwrap();
/// assert_eq!(out, serde_json::json!({"y": 2}));
/// # });
/// ```
#[async_trait]
pub trait Tool: Send + Sync {
    /// Tool name shown to the LLM. Must be unique within a catalogue.
    fn name(&self) -> &str;

    /// Human-readable description shown to the LLM.
    fn description(&self) -> &str;

    /// JSON-schema for the tool's arguments.
    fn json_schema(&self) -> &serde_json::Value;

    /// Invoke the tool. Args are pre-validated against `json_schema`.
    async fn invoke(
        &self,
        args: serde_json::Value,
        ctx: ToolCtx,
    ) -> Result<serde_json::Value, ToolError>;

    /// Whether this tool causes an irreversible side effect (payment,
    /// mutation of external state, sent message, etc.). Defaults to
    /// `false` (read-only). Tools that perform side effects MUST override
    /// to return `true` — `klieo-ops` Gates evaluate at the pre-effect
    /// boundary on every `is_effectful() == true` tool.
    fn is_effectful(&self) -> bool {
        false
    }
}

/// Dispatches tool calls by name.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::{noop_bus, FakeToolInvoker};
/// use klieo_core::{ToolCtx, ToolInvoker};
/// let (pubsub, _, kv, jobs) = noop_bus();
/// let inv = FakeToolInvoker::new()
///     .with_tool("echo", "echoes back", |args| Ok(args));
/// let ctx = ToolCtx::new(pubsub, kv, jobs);
/// let out = inv.invoke("echo", serde_json::json!({"x": 1}), ctx).await.unwrap();
/// assert_eq!(out, serde_json::json!({"x": 1}));
/// assert_eq!(inv.catalogue().len(), 1);
/// # });
/// ```
#[async_trait]
pub trait ToolInvoker: Send + Sync {
    /// Invoke `name` with `args` against the supplied context.
    async fn invoke(
        &self,
        name: &str,
        args: serde_json::Value,
        ctx: ToolCtx,
    ) -> Result<serde_json::Value, ToolError>;

    /// Return the catalogue this invoker exposes.
    fn catalogue(&self) -> Vec<crate::llm::ToolDef>;

    /// Returns true if `invoke(name, args, ctx)` is safe to
    /// re-execute from the original arguments — i.e. read-only,
    /// deterministic, or duplicate execution has no business
    /// effect.
    ///
    /// Per-tool granularity so an invoker can expose a mix of
    /// idempotent + state-mutating tools. Default false (safe).
    /// Tool authors override per-name only when the tool is
    /// provably safe to re-execute.
    ///
    /// Cluster 0.24: when true for a given tool, the follower
    /// replica detecting an orphaned stream re-invokes
    /// `tools/call` from the cached args instead of writing a
    /// terminal "leader died" frame. See ADR-024.
    fn is_tool_idempotent(&self, _name: &str) -> bool {
        false
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[allow(dead_code)]
    fn _assert_dyn_invoker(_: &dyn ToolInvoker) {}

    #[allow(dead_code)]
    fn _assert_dyn_tool(_: &dyn Tool) {}

    struct StubInvoker;
    #[async_trait::async_trait]
    impl ToolInvoker for StubInvoker {
        fn catalogue(&self) -> Vec<crate::llm::ToolDef> {
            vec![]
        }
        async fn invoke(
            &self,
            _name: &str,
            _args: serde_json::Value,
            _ctx: ToolCtx,
        ) -> Result<serde_json::Value, ToolError> {
            Ok(serde_json::Value::Null)
        }
    }

    #[test]
    fn default_is_tool_idempotent_returns_false() {
        let inv = StubInvoker;
        assert!(!inv.is_tool_idempotent("anything"));
        assert!(!inv.is_tool_idempotent(""));
    }
}

#[cfg(test)]
mod ctx_tests {
    use super::*;
    use crate::server_outbound::{ServerOutbound, ServerOutboundError};
    use crate::test_utils::noop_bus;
    use std::time::Duration;
    use tokio_util::sync::CancellationToken;

    #[test]
    fn new_defaults_progress_none_and_cancel_uncancelled() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let ctx = ToolCtx::new(pubsub, kv, jobs);
        assert!(ctx.progress.is_none());
        assert!(!ctx.cancel.is_cancelled());
    }

    #[test]
    fn with_cancel_overrides_default_token() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let token = CancellationToken::new();
        let ctx = ToolCtx::new(pubsub, kv, jobs).with_cancel(token.clone());
        token.cancel();
        assert!(ctx.cancel.is_cancelled());
    }

    #[test]
    fn with_progress_sets_sender() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let (tx, _rx) = tokio::sync::broadcast::channel::<crate::AgentEvent>(8);
        let ctx = ToolCtx::new(pubsub, kv, jobs).with_progress(tx);
        assert!(ctx.progress.is_some());
    }

    struct StubOutbound;

    #[async_trait::async_trait]
    impl ServerOutbound for StubOutbound {
        async fn outbound_request(
            &self,
            _method: &str,
            _params: serde_json::Value,
            _timeout: Duration,
        ) -> Result<serde_json::Value, ServerOutboundError> {
            Err(ServerOutboundError::Unsupported)
        }
    }

    #[test]
    fn new_defaults_server_outbound_none() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let ctx = ToolCtx::new(pubsub, kv, jobs);
        assert!(ctx.server_outbound.is_none());
    }

    #[test]
    fn with_server_outbound_attaches_channel() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let outbound: Arc<dyn ServerOutbound> = Arc::new(StubOutbound);
        let ctx = ToolCtx::new(pubsub, kv, jobs).with_server_outbound(outbound);
        assert!(ctx.server_outbound.is_some());
    }
}