klieo_core/

tool.rs

//! 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;

/// Verified caller principal carried on [`ToolCtx`] for server-side
/// authorisation (resume-ticket binding) and non-PII attribution.
///
/// A newtype rather than a bare `String` so the value cannot be passed
/// where arbitrary text is expected, and — deliberately NOT `Serialize`
/// — it cannot be folded into an [`crate::Episode`] or any other
/// persisted/LLM-visible structure by accident. The raw value reaches
/// only server-side authz/tracing; attribution into the audit trail goes
/// through [`crate::principal_hash`].
#[derive(Clone, Debug)]
pub struct CallerPrincipal(String);

impl CallerPrincipal {
    /// Wrap a verified principal supplied by the transport layer.
    pub fn new(value: impl Into<String>) -> Self {
        Self(value.into())
    }

    /// Borrow the raw principal for an authz compare or hashing. The
    /// caller is responsible for keeping it off LLM-visible surfaces.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// Opaque cross-hop provenance anchor carried on [`ToolCtx`].
///
/// Same rationale as [`CallerPrincipal`]: a distinct, non-`Serialize`
/// type so a caller-supplied anchor cannot leak into a persisted
/// structure except through the explicit run-origin recording path,
/// which copies its [`Self::as_str`] verbatim into a fresh `Episode`.
#[derive(Clone, Debug)]
pub struct ParentAnchor(String);

impl ParentAnchor {
    /// Wrap a validated cross-hop anchor supplied by the transport layer.
    pub fn new(value: impl Into<String>) -> Self {
        Self(value.into())
    }

    /// Borrow the raw anchor for recording as the run origin.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// 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>>,
    /// Verified caller principal for server-side authorisation
    /// (e.g. binding a `klieo/run/resume` ticket to its issuer).
    /// Server-side metadata ONLY: implementations MUST NOT append
    /// this value to short-term/episodic memory, conversation history,
    /// or any LLM-visible message — those surfaces are operator-
    /// reviewed and PII-redacted on their own contract. `None` when
    /// no authenticator is wired, when the transport carries no
    /// principal (stdio, unauthenticated HTTP), or when the request
    /// is anonymous.
    pub caller_principal: Option<CallerPrincipal>,
    /// Opaque cross-hop provenance anchor supplied by an authenticated
    /// caller — by convention the caller's own provenance chain-entry id
    /// (or its run's episodic-root hash). When set, the run records it as
    /// its origin so klieo→klieo composition can be stitched across
    /// deployment boundaries in the audit trail. Audit metadata ONLY:
    /// implementations MUST NOT append it to short-term/episodic memory,
    /// conversation history, or any LLM-visible message. The value is
    /// recorded verbatim (stitching requires it to equal the caller's
    /// chain-entry id), is a caller-ASSERTED, UNVERIFIED claim, and must
    /// be treated as such downstream. `None` when no anchor was supplied
    /// or the caller is unauthenticated.
    pub parent_anchor: Option<ParentAnchor>,
}

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,
            caller_principal: None,
            parent_anchor: 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
    }

    /// Stamp the verified caller principal so server-side
    /// authorisation helpers (e.g. resume-ticket issuance) can bind
    /// follow-up requests to the original caller. See the field
    /// docstring for the prohibition on leaking this value into
    /// LLM-visible memory.
    #[must_use]
    pub fn with_caller_principal(mut self, principal: String) -> Self {
        self.caller_principal = Some(CallerPrincipal::new(principal));
        self
    }

    /// Stamp the cross-hop provenance anchor (the caller's chain-entry
    /// id) so the run can record its origin. See the field docstring for
    /// the verbatim-record / unverified-claim / no-LLM-leak contract.
    #[must_use]
    pub fn with_parent_anchor(mut self, anchor: String) -> Self {
        self.parent_anchor = Some(ParentAnchor::new(anchor));
        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>;

    /// Advisory metadata: whether this tool causes an irreversible side effect
    /// (payment, mutation of external state, sent message, etc.). Defaults to
    /// `false` (read-only); side-effecting tools should override to `true`.
    ///
    /// This flag does NOT, by itself, cause gating. Gating is opt-in: wrap a
    /// tool in `klieo_ops::GatedTool` to enforce a pre-effect approval boundary
    /// (it gates regardless of this flag). The flag is surfaced for callers and
    /// audit metadata, not consumed by the runtime to auto-gate.
    fn is_effectful(&self) -> bool {
        false
    }

    /// Whether dispatch must record a redacted projection of this tool's
    /// input, output, and error to [`crate::EpisodicMemory`] in place of
    /// the raw values. Tools handling claimant PII MUST override to
    /// return `true`; dispatch then never persists raw args/results for
    /// the call, falling back to an opaque digest when no
    /// [`crate::AuditRedactor`] is wired. Defaults to `false` (raw).
    /// Hash-chain provenance is computed over the raw values upstream
    /// and is unaffected by this flag.
    fn redacts_audit(&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
    }

    /// Mirrors [`Tool::redacts_audit`] across the invoker boundary
    /// (dispatch holds a [`ToolInvoker`], not the concrete [`Tool`]).
    /// Per-name granularity so one invoker can mix PII-handling and
    /// non-PII tools. Default false (raw).
    fn tool_redacts_audit(&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());
    }

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

    #[test]
    fn with_caller_principal_records_value() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let ctx = ToolCtx::new(pubsub, kv, jobs).with_caller_principal("alice@x".into());
        assert_eq!(
            ctx.caller_principal.as_ref().map(CallerPrincipal::as_str),
            Some("alice@x")
        );
    }

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

    #[test]
    fn with_parent_anchor_records_value() {
        let (pubsub, _, kv, jobs) = noop_bus();
        let ctx = ToolCtx::new(pubsub, kv, jobs).with_parent_anchor("sha256:abc123".into());
        assert_eq!(
            ctx.parent_anchor.as_ref().map(ParentAnchor::as_str),
            Some("sha256:abc123")
        );
    }
}