klieo_mcp_server/

lib.rs

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
//! Expose any klieo [`ToolInvoker`] as an MCP server over stdio or HTTP (see the `http` cargo feature).
//!
//! This crate is the inverse of `klieo-tools-mcp`: instead of a Rust
//! host calling external MCP servers, it lets external MCP hosts
//! (Claude Desktop, Continue, LangGraph, OpenAI Agents SDK) call
//! into a Rust-built klieo `ToolInvoker`.
//!
//! Polyglot strategy: klieo exposes three wire-level contracts so
//! non-Rust hosts can integrate without depending on klieo crates —
//! the NATS bus (durable pipelines), this MCP server (stdio/HTTP
//! tool calls), and the A2A peer protocol (JSON-RPC). This crate
//! covers the MCP-server contract.
//!
//! ## Scope
//!
//! - **stdio transport** — newline-delimited JSON-RPC 2.0 frames on
//!   stdin/stdout, matching the MCP reference spec.
//! - **`tools/list`** — derived from the invoker's catalogue.
//! - **`tools/call`** — dispatches to `ToolInvoker::invoke` with a
//!   `ToolCtx` minted per request from the
//!   [`McpServerBuilder::with_tool_ctx_factory`]-supplied factory.
//!   Default factory is the in-memory noop wiring (`Pubsub`/`KvStore`/
//!   `JobQueue` from `klieo-bus-memory`, fresh per call). Override
//!   the factory to expose tools that need real bus access.
//! - **`initialize` / `shutdown`** — standard MCP handshake.
//!
//! Out-of-scope:
//! - SSE streaming responses (HTTP transport returns JSON only).
//! - Multi-tenant auth (wrap `router()` with a tower auth layer, or front with an auth-enforcing reverse proxy).
//!
//! ## Agent exposure (0.9, ADR-010)
//!
//! `McpServer::expose_agent_with_schema(agent, schema, ctx_factory)`
//! wraps any [`klieo_core::Agent`] as a single MCP tool whose
//! `inputSchema` is the caller-supplied JSON Schema. The
//! `ctx_factory` closure is called per `tools/call` to mint a
//! fresh [`klieo_core::agent::AgentContext`] (so each invocation gets its own RunId).
//!
//! Behind the `schemars` cargo feature, `expose_agent::<A>(agent,
//! ctx_factory)` derives the schema automatically via
//! `schema_for!(A::Input)`. See ADR-010 for the trade-off.

#[cfg(feature = "http")]
pub(crate) mod http;

pub(crate) mod outbound;

#[cfg(feature = "http")]
pub(crate) mod outbound_ring;

pub(crate) mod session;

pub(crate) mod workflow;

// Inbound per-tenant LLM budget governor. Wired through
// `McpServerBuilder::with_governor`; gated on `feature = "governor"`
// so adopters that bring their own rate-limiting layer keep the lean
// default dep graph.
#[cfg(feature = "governor")]
pub(crate) mod governor;

// The HTTP `klieo/run/resume` handler is the only production consumer;
// stdio builds still need the module compiled so the tests + workflow
// suspend path can mint tickets, but the read-side API
// (`peek` / `claim` / decode-error variant) only fires under
// `--features http`.
#[cfg_attr(not(feature = "http"), allow(dead_code))]
pub(crate) mod resume_ticket;

pub mod outbound_sink;
pub use outbound_sink::{OutboundFrameSink, OutboundSinkError};

pub mod sampling;
pub use sampling::{
    ModelHint, ModelPreferences, SamplingContent, SamplingMessage, SamplingRequest,
    SamplingResponse,
};

pub mod roots;
pub use roots::Root;

pub mod outbound_ext;
pub use outbound_ext::McpOutboundExt;

/// Re-export the cluster-0.24 follower-side orphan re-invoke gate
/// for integration tests. Gated on `test-fixtures` so production
/// callers cannot reach the helper. See [`crate::http::OrphanOutcome`]
/// + [`crate::http::handle_dead_leader_orphan_mcp`] for the contract.
#[cfg(all(feature = "http", feature = "test-fixtures"))]
pub use http::{handle_dead_leader_orphan_mcp, OrphanOutcome};

#[cfg(all(feature = "http", feature = "bench"))]
pub use http::encode_sse_frame;
#[cfg(feature = "bench")]
pub use outbound_sink::bench_stdio_sink;

/// Scans a replay-buffer snapshot and returns entries with `event_id >
/// since_id`. Approximates the per-frame scan cost of the SSE replay
/// path without HTTP header parsing or lock acquisition overhead.
/// Available only under the `bench` feature.
#[cfg(all(feature = "http", feature = "bench"))]
pub fn bench_filter_replay(
    entries: &[(u64, std::sync::Arc<serde_json::Value>)],
    since_id: u64,
) -> Vec<(u64, std::sync::Arc<serde_json::Value>)> {
    entries
        .iter()
        .filter(|(id, _)| *id > since_id)
        .cloned()
        .collect()
}

use async_trait::async_trait;
use klieo_core::agent::Agent;
use klieo_core::error::ToolError;
use klieo_core::llm::ToolDef;
use klieo_core::tool::{ToolCtx, ToolInvoker};
use std::sync::Arc;
use thiserror::Error;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use tracing::warn;

/// JSON-RPC 2.0 standard error codes used in handler dispatch.
/// See https://www.jsonrpc.org/specification#error_object.
pub(crate) const JSONRPC_PARSE_ERROR: i64 = -32700;
const JSONRPC_METHOD_NOT_FOUND: i64 = -32601;
#[cfg(feature = "http")]
pub(crate) const JSONRPC_INVALID_PARAMS: i64 = -32602;
pub(crate) const JSONRPC_SERVER_ERROR: i64 = -32000;
/// Klieo-specific: resume requested with an expired window.
/// Maps to JSON-RPC application code `-32011`. Used by `stream_resume`
/// in the HTTP transport.
#[cfg(feature = "http")]
pub(crate) const JSONRPC_RESUME_BUFFER_EXPIRED: i64 = -32011;
/// Klieo-specific: resume requested for an unknown progressToken.
/// Maps to JSON-RPC application code `-32012`. Used by `stream_resume`
/// in the HTTP transport.
#[cfg(feature = "http")]
pub(crate) const JSONRPC_RESUME_BUFFER_NOT_FOUND: i64 = -32012;
/// Klieo-specific: stream leader (originating replica running the
/// invoke) died or never claimed leadership. Follower observed an
/// orphaned resume buffer + wrote a terminal frame so the client sees
/// clean termination + can retry. Mirrors `codes::LEADER_DIED` on
/// the A2A side. ADR-020.
#[cfg(feature = "http")]
pub(crate) const JSONRPC_LEADER_DIED: i64 = -32099;
/// Klieo-specific: a second `initialize` arrived while the HTTP
/// transport already owns an active session. Returned in the JSON-RPC
/// error envelope alongside a `409 Conflict` HTTP status. Renumbered
/// from `-32099` to `-32002` in 0.28 to resolve the collision with
/// [`JSONRPC_LEADER_DIED`]; LEADER_DIED kept the older slot because
/// it has the longer external visibility. ADR-028, ADR-029.
#[cfg(feature = "http")]
pub(crate) const JSONRPC_SESSION_CONFLICT: i64 = -32002;
/// JSON-RPC application code returned when an
/// [`klieo_auth_common::Authenticator`] wired via
/// [`McpServerBuilder::with_authenticator`] rejects a request —
/// either `authenticate(headers)` failed or `authorize_method`
/// denied the principal for the requested method. Mirrors
/// `klieo_auth_common UNAUTHENTICATED (-32001)` so both
/// transports surface auth failures with the same wire code. ADR-021.
#[cfg(feature = "http")]
pub(crate) const JSONRPC_UNAUTHENTICATED: i64 = -32001;

// Compile-time uniqueness: every JSONRPC_* code listed here must be
// distinct. Adding a new code with a duplicate value is a compile
// error.
#[cfg(feature = "http")]
const _: () = {
    let codes: [i64; 9] = [
        JSONRPC_PARSE_ERROR,
        JSONRPC_METHOD_NOT_FOUND,
        JSONRPC_INVALID_PARAMS,
        JSONRPC_SERVER_ERROR,
        JSONRPC_UNAUTHENTICATED,
        JSONRPC_RESUME_BUFFER_EXPIRED,
        JSONRPC_RESUME_BUFFER_NOT_FOUND,
        JSONRPC_LEADER_DIED,
        JSONRPC_SESSION_CONFLICT,
    ];
    let mut i = 0;
    while i < codes.len() {
        let mut j = i + 1;
        while j < codes.len() {
            assert!(codes[i] != codes[j], "JSONRPC_* code collision");
            j += 1;
        }
        i += 1;
    }
};

/// Leader-claim TTL for the `klieo-leaders` KV bucket used by
/// `stream_tools_call` (claim on invoke) and `stream_resume` (orphan
/// check on resume). Operators MUST configure the JetStream KV bucket
/// with `max_age` equal to this value so a dead replica's leader
/// entry evicts automatically; the registry's heartbeat runs every
/// `TTL / 2`. See ADR-020.
pub const LEADER_TTL: std::time::Duration = std::time::Duration::from_secs(5);

/// Leader-key prefix for the MCP transport in the shared
/// `klieo-leaders` bucket. Mirrors `a2a.<task_id>` on the A2A side.
pub const MCP_LEADER_KEY_PREFIX: &str = "mcp.";

/// MCP protocol revision advertised on every `initialize` response.
///
/// Pinned to the 2025-03-26 Streamable HTTP revision the HTTP
/// transport implements (`Mcp-Session-Id` lifecycle + SSE upgrade).
/// The single source for the `protocolVersion` field — the
/// conformance test fails if any handshake site diverges from it.
pub const MCP_PROTOCOL_VERSION: &str = "2025-03-26";

/// Top-level error from the MCP server loop.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum McpServerError {
    /// stdin/stdout I/O failure.
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
    /// JSON-RPC payload could not be decoded.
    #[error("json decode error: {0}")]
    Json(#[from] serde_json::Error),
    /// Klieo-specific: resume requested with an expired window.
    /// Maps to JSON-RPC application code `-32011`.
    #[error("resume window expired (since_id={since_id})")]
    ResumeBufferExpired {
        /// Cursor the caller supplied.
        since_id: u64,
    },

    /// Klieo-specific: resume requested for an unknown progressToken.
    /// Maps to JSON-RPC application code `-32012`.
    #[error("no buffered stream for progressToken: {0}")]
    ResumeBufferNotFound(String),

    /// Underlying tool invocation failed.
    #[error("tool error: {0}")]
    Tool(#[from] ToolError),

    /// Caller-supplied subject segment (progressToken, task id, …)
    /// failed [`klieo_core::validate_subject_token`] — contains a
    /// reserved metacharacter (`.`, `*`, `>`), whitespace, or
    /// non-ASCII byte. Split out of [`McpServerError::Bus`] so
    /// callers can distinguish caller-input validation failures
    /// (permanent — retry would fail identically) from genuine
    /// bus transport failures.
    ///
    /// The typed cause is reachable via `Error::source()` —
    /// `BusError::Invalid(_)`.
    #[error("invalid subject token: {0}")]
    InvalidSubject(#[source] klieo_core::BusError),

    /// Bus publish or subscribe transport-class failure (connection,
    /// timeout, retryable, permanent, …). The wire envelope maps
    /// this to JSON-RPC `SERVER_ERROR` (`-32000`); the typed cause
    /// stays on `e.source()` for downstream tracing.
    ///
    /// Caller-input validation failures
    /// ([`klieo_core::BusError::Invalid`]) are surfaced as
    /// [`McpServerError::InvalidSubject`] instead.
    #[error("bus error: {0}")]
    Bus(#[source] klieo_core::BusError),

    /// Outbound request timed out (peer didn't respond within the
    /// configured deadline).
    #[error("outbound request timed out")]
    OutboundTimeout,

    /// Peer returned an MCP-spec error envelope on an outbound
    /// request.
    #[error("client returned error: code={code} message={message}")]
    ClientReturnedError {
        /// JSON-RPC error code returned by the peer.
        code: i64,
        /// JSON-RPC error message returned by the peer.
        message: String,
    },

    /// Transport closed while an outbound request was in flight.
    #[error("transport closed")]
    TransportClosed,

    /// The transport carrying this server does not support
    /// server-initiated outbound requests (e.g. HTTP today).
    #[error("outbound channel unsupported on this transport")]
    OutboundUnsupported,

    /// The outbound request frame could not be serialised. Indicates an
    /// internal invariant violation rather than a transport or capability failure.
    #[error("outbound serialisation failed: {0}")]
    OutboundSerialisation(#[source] serde_json::Error),

    /// Failed to serialise a sampling request to JSON.
    #[error("failed to serialise sampling request: {0}")]
    SamplingSerialise(serde_json::Error),

    /// Failed to deserialise the client's sampling response.
    #[error("failed to deserialise sampling response: {0}")]
    SamplingDeserialise(serde_json::Error),
}

impl From<klieo_core::ServerOutboundError> for McpServerError {
    fn from(e: klieo_core::ServerOutboundError) -> Self {
        use klieo_core::ServerOutboundError as E;
        match e {
            E::Timeout => McpServerError::OutboundTimeout,
            E::PeerError { code, message } => McpServerError::ClientReturnedError { code, message },
            E::TransportClosed => McpServerError::TransportClosed,
            E::Unsupported => McpServerError::OutboundUnsupported,
            E::Serialisation(err) => McpServerError::OutboundSerialisation(err),
            // `ServerOutboundError` is `#[non_exhaustive]`. Future
            // variants surface as `OutboundUnsupported` until an
            // explicit arm is added here.
            _ => McpServerError::OutboundUnsupported,
        }
    }
}

impl From<klieo_core::BusError> for McpServerError {
    fn from(e: klieo_core::BusError) -> Self {
        match e {
            klieo_core::BusError::Invalid(_) => Self::InvalidSubject(e),
            other => Self::Bus(other),
        }
    }
}

/// Errors returned by [`McpServerBuilder::build`] and [`McpServerBuilder::build_arc`].
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum McpBuildError {
    /// `with_cancel_subscription` was set but `build()` (not `build_arc()`) was called.
    #[error("with_cancel_subscription requires build_arc()")]
    CancelRequiresArc,
    /// No invokers registered — call `add_tools`, `add_agent`, or `add_agent_with_schema` first.
    #[error("at least one invoker required; call add_tools or add_agent before build")]
    NoInvokers,
    /// Two agents/invokers share a tool name — rename or split the colliding agent.
    #[error("duplicate tool name {0:?} across registered invokers")]
    DuplicateTool(String),
    /// The builder's `profile(..)` requirements were not met (e.g. regulated
    /// profile without tenant binding or with an anonymous authenticator).
    #[error(transparent)]
    RegulatedProfile(#[from] klieo_core::ProfileViolation),
    /// A workflow was registered via `add_workflow_with_schema` /
    /// `add_workflow` without first calling
    /// [`McpServerBuilder::with_hitl`]. The workflow path drives
    /// [`klieo_hitl::run_with_hitl`], which requires a shared HITL
    /// client + config; without them every suspension would fail with
    /// no compliance endpoint to submit to.
    #[error("workflow registered without with_hitl(..); call with_hitl before build")]
    WorkflowWithoutHitl,
    /// A workflow was registered without a wired
    /// [`McpServerBuilder::with_governor`]. The build hard-gates
    /// workflow exposure on a per-tenant LLM budget: an inbound MCP
    /// caller must NOT drive unbounded paid LLM spend. The error is
    /// emitted regardless of the `governor` cargo feature — the build
    /// rejects the unsafe configuration both ways so adopters can't
    /// silently fall through by toggling a feature flag.
    #[error("workflow registered without with_governor(..); call with_governor before build")]
    WorkflowWithoutGovernor,
}

/// MCP server that exposes a `klieo-core` `ToolInvoker` as a stdio
/// MCP server.
///
/// ## Construction
///
/// ```no_run
/// # use std::sync::Arc;
/// # use klieo_core::tool::ToolInvoker;
/// # async fn _ex(invoker: Arc<dyn ToolInvoker>) {
/// use klieo_mcp_server::McpServer;
/// let server = Arc::new(McpServer::expose_tools(invoker));
/// server.serve_stdio().await.expect("server loop");
/// # }
/// ```
pub struct McpServer {
    pub(crate) invoker: Arc<dyn ToolInvoker>,
    tool_ctx_factory: ToolCtxFactory,
    pub(crate) parent_cancel: CancellationToken,
    pub(crate) resume_buffer: std::sync::Arc<dyn klieo_core::resume::ResumeBuffer>,
    pub(crate) pubsub: std::sync::Arc<dyn klieo_core::Pubsub>,
    pub(crate) cancel_registry: klieo_core::CancelRegistry<String>,
    /// Bounds in-flight SSE-frame fanout publishes on the per-
    /// progressToken bus subject (`klieo.mcp.progress.{token}`) and
    /// drop-time cross-replica cancel publishes. Shared with
    /// [`crate::http::spawn_publish`] and the `CancelOnDrop` body
    /// so every replica caps fanout work at the same value;
    /// saturation drops the publish + emits a `warn`. Default
    /// [`DEFAULT_PUBLISH_PERMITS`] (64); configurable via
    /// [`McpServerBuilder::with_publish_concurrency`].
    #[cfg(feature = "http")]
    pub(crate) publish_permits: std::sync::Arc<tokio::sync::Semaphore>,
    pub(crate) leader_registry: Option<klieo_core::LeaderRegistry>,
    pub(crate) ownership_registry: Option<klieo_core::OwnershipRegistry>,
    /// One-shot ticket store keyed under `klieo.mcp.resume-tickets`,
    /// populated when the builder was wired with
    /// [`McpServerBuilder::with_checkpoint_kv`]. `None` keeps the
    /// slice-1 no-ticket suspend behaviour. Read only by the
    /// HTTP `klieo/run/resume` handler; gated dead-code accordingly.
    #[cfg_attr(not(feature = "http"), allow(dead_code))]
    pub(crate) resume_ticket_store: Option<Arc<crate::resume_ticket::ResumeTicketStore>>,
    /// Resume-side workflow handles keyed by `workflow_name`.
    /// Populated only when at least one workflow was registered AND
    /// the builder was wired with both `with_hitl` + the workflow
    /// registration path. Read by the HTTP `klieo/run/resume`
    /// handler after the ticket-bound authz gate has cleared.
    #[cfg_attr(not(feature = "http"), allow(dead_code))]
    pub(crate) workflow_resume_handles:
        std::collections::HashMap<String, Arc<dyn crate::workflow::WorkflowResumeHandle>>,
    /// Optional [`klieo_auth_common::Authenticator`] wired into the
    /// HTTP entry path via
    /// [`McpServerBuilder::with_authenticator`]. `post_mcp` calls
    /// `authenticate(headers)` + `authorize_method(&identity, method)`
    /// before dispatch / SSE upgrade when this is `Some`. None
    /// preserves the pre-0.21 caller-owned-middleware contract
    /// (no auth applied by the server).
    pub(crate) authenticator: Option<Arc<dyn klieo_auth_common::Authenticator>>,
    pub(crate) leader_ttl: std::time::Duration,
    pub(crate) leader_heartbeat_interval: std::time::Duration,
    pub(crate) max_failover_attempts: u32,
    pub(crate) kv_reaper_interval: Option<std::time::Duration>,
    /// Held for Drop semantics — Drop aborts the background scan
    /// task spawned by [`McpServerBuilder::with_kv_reaper`]. `None`
    /// when the builder was not asked to spawn the reaper or when
    /// neither leader-election nor tenant-binding was wired (no
    /// bucket to scan).
    _kv_reaper: Option<klieo_core::KvReaperHandle>,
    /// Stdio transport's per-process session container. Populated by
    /// [`Self::ensure_outbound_and_roots`] on stdio entry and holds
    /// the same `Session` shape as the HTTP registry entries.
    /// Remains empty for HTTP-only servers; stdio servers populate
    /// once per process lifetime.
    pub(crate) stdio_session: tokio::sync::OnceCell<std::sync::Arc<crate::session::Session>>,
    /// HTTP session registry, keyed by minted `Mcp-Session-Id` UUID.
    /// Populated by `handle_initialize_post` when a fresh session
    /// passes the cap check; evicted by `delete_mcp` on explicit
    /// client teardown or by the idle reaper when the per-session
    /// watchdog fires. Stdio transports leave this empty.
    #[cfg(feature = "http")]
    pub(crate) sessions: std::sync::Arc<
        tokio::sync::RwLock<
            std::collections::HashMap<uuid::Uuid, std::sync::Arc<crate::session::Session>>,
        >,
    >,
    /// Hard cap on concurrent HTTP sessions held in
    /// [`Self::sessions`]. Read by `handle_initialize_post` to
    /// decide whether a fresh `initialize` POST proceeds or yields
    /// 503. Configured via
    /// [`McpServerBuilder::with_max_sessions`]; defaults to
    /// [`DEFAULT_MAX_SESSIONS`].
    #[cfg(feature = "http")]
    pub(crate) max_sessions: usize,
    /// Hard cap on concurrent HTTP sessions per authenticated
    /// principal. Resolved at build time from
    /// [`McpServerBuilder::with_max_sessions_per_principal`] or the
    /// default [`default_max_sessions_per_principal`] applied to
    /// `max_sessions`. Read by `handle_initialize_post` to decide
    /// whether a fresh `initialize` POST from an already-tracked
    /// principal proceeds or yields 503.
    #[cfg(feature = "http")]
    pub(crate) max_sessions_per_principal: usize,
    /// Per-session SSE replay buffer capacity in frames. Resolved at
    /// build time from the builder; default
    /// [`DEFAULT_SSE_REPLAY_CAPACITY`]. Setting to 0 disables SSE
    /// resumption.
    #[cfg(feature = "http")]
    pub(crate) sse_replay_capacity: usize,
    /// Per-principal session count cache. Keyed by the verified
    /// [`klieo_auth_common::Identity::as_str`] value of the principal
    /// that minted the session; auth-disabled deployments leave this
    /// map empty.
    ///
    /// Lock acquisition order: [`Self::sessions`] write FIRST, then
    /// `principal_counts` write. Never the reverse. Either may be
    /// taken alone. Eviction paths decrement OUTSIDE the `sessions`
    /// critical section to preserve the invariant.
    #[cfg(feature = "http")]
    pub(crate) principal_counts:
        std::sync::Arc<tokio::sync::RwLock<std::collections::HashMap<String, usize>>>,
    /// Idempotent spawn sentinel for the idle-session reaper task.
    /// The first call to `ensure_idle_reaper` populates this cell
    /// and spawns the background scan; subsequent calls are no-ops
    /// so duplicate `initialize` POSTs do not spawn duplicate
    /// reapers.
    #[cfg(feature = "http")]
    pub(crate) idle_reaper_started: tokio::sync::OnceCell<()>,
    /// Set by [`McpServerBuilder::with_client_sampling`]. Gates the
    /// `capabilities.sampling = {}` field on the initialize response
    /// and signals to the stdio loop that outbound correlation should
    /// be wired when the transport mints its shared writer.
    pub(crate) declare_sampling: bool,
    /// Shared stdout writer for the stdio transport. Primed by
    /// [`Self::serve_with_streams`] on the first call so
    /// [`Self::ensure_outbound_and_roots`] hands the same `Arc` to the
    /// outbound primitive — preventing interleaved writes between the
    /// inbound reply path and the outbound request path on a single
    /// underlying stream.
    pub(crate) stdout_writer: tokio::sync::OnceCell<crate::outbound::SharedWriter>,
    /// Per-server snapshot of client-declared capabilities, parsed from
    /// the `initialize` request's `params.capabilities` payload. Read
    /// by the `notifications/initialized` arm to decide whether to
    /// drive the initial `roots/list` fetch.
    pub(crate) client_caps: tokio::sync::Mutex<ClientCaps>,
    /// Idle deadline for the active session. Zero disables the
    /// watchdog entirely. Read by the GET handler and the watchdog
    /// task spawned on `initialize`. ADR-028.
    #[cfg(feature = "http")]
    pub(crate) session_idle_timeout: std::time::Duration,
    /// Tick cadence for the idle-reaper background task. Defaults to
    /// 10 seconds in production builds; test builds can shorten this
    /// through [`McpServerBuilder::with_idle_reaper_tick`] so short
    /// `session_idle_timeout` values trip within the test deadline.
    #[cfg(feature = "http")]
    pub(crate) idle_reaper_tick: std::time::Duration,
    /// Reference instant captured at server build time. Each HTTP
    /// session's `last_activity_millis` is encoded as millis since
    /// this anchor; the idle reaper compares `server_start.elapsed()`
    /// against per-session values to decide eviction. Monotonic
    /// (`Instant`) so wall-clock jumps cannot corrupt the comparison.
    #[cfg(feature = "http")]
    pub(crate) server_start: std::time::Instant,
}

/// Per-server snapshot of client-declared MCP capabilities parsed
/// from the `initialize` request. Held under
/// [`McpServer::client_caps`] (a `tokio::sync::Mutex`) so the
/// `initialize` write and the `notifications/initialized` read
/// cannot race.
#[derive(Default, Debug)]
pub(crate) struct ClientCaps {
    /// True iff the client advertised `capabilities.roots` (any
    /// non-null payload) on `initialize`. Gates the spawned
    /// initial-`roots/list` fetch on `notifications/initialized`.
    pub roots_supported: bool,
}

/// Factory for fresh [`klieo_core::agent::AgentContext`] values,
/// called once per `tools/call` so each invocation gets its own
/// `RunId` + cancel token. Caller-owned to keep `klieo-mcp-server`
/// free of opinions about which memory / bus / llm backends an
/// agent runs against.
pub type AgentContextFactory =
    Arc<dyn Fn() -> klieo_core::agent::AgentContext + Send + Sync + 'static>;

/// Factory for fresh [`klieo_core::tool::ToolCtx`] values, called
/// once per `tools/call` so each invocation gets its own
/// `Pubsub` / `KvStore` / `JobQueue` (or a shared one, caller's
/// choice). Default = `default_tool_ctx_factory` which mints
/// the same per-request in-memory wiring `serve_stdio` has used
/// since 0.8. Override via
/// [`McpServerBuilder::with_tool_ctx_factory`] to expose
/// bus-using tools through this server.
pub type ToolCtxFactory = Arc<dyn Fn() -> klieo_core::tool::ToolCtx + Send + Sync + 'static>;

/// Returns the default [`ToolCtxFactory`]: a per-request in-memory
/// `Pubsub` / `KvStore` / `JobQueue`, equivalent to what
/// `serve_stdio` used before the factory was introduced.
fn default_tool_ctx_factory() -> ToolCtxFactory {
    Arc::new(noop_ctx)
}

/// Builder for [`McpServer`]. Collects one or more invokers (raw
/// `ToolInvoker`s or `Agent`s wrapped as single-tool invokers) and
/// configures server-level concerns such as a parent
/// [`CancellationToken`] that propagates into every minted
/// `AgentContext`.
///
/// ## Why a builder
///
/// The original 0.9.0 surface offered four constructors along two
/// orthogonal axes (schema-explicit vs schemars-derive, cancellable
/// vs non-cancellable). Adding a third axis (multi-agent / multi-
/// tool) would have produced eight constructors. The builder
/// collapses those axes into a single fluent API.
///
/// ## Multi-invoker dispatch
///
/// Repeated `add_agent_with_schema` / `add_agent` / `add_tools` calls
/// accumulate invokers; [`Self::build`] wraps them in a private
/// `MergedInvoker` that merges their catalogues and routes each
/// `tools/call` to the inner invoker that claims the named tool.
/// Tool-name collisions across invokers are rejected at `build` time
/// (caller bug — fix by renaming or splitting the colliding agent).
/// Default per-server cap on concurrent in-flight SSE-frame fanout
/// publishes and drop-time cross-replica cancel publishes. Bounds the
/// runtime work `stream_tools_call` can pile up when a slow bus
/// backend (e.g. stalled NATS) lets publishes accumulate. Configurable
/// via [`McpServerBuilder::with_publish_concurrency`]; mirrors the
/// same default applied to `A2aDispatcher` so both transports share
/// the cap.
#[cfg(feature = "http")]
const DEFAULT_PUBLISH_PERMITS: usize = 64;

/// Default idle deadline for an HTTP streamable session before the
/// per-session watchdog forcibly closes it. Configurable via
/// [`McpServerBuilder::with_session_idle_timeout`]. ADR-028.
#[cfg(feature = "http")]
pub(crate) const DEFAULT_SESSION_IDLE_TIMEOUT: std::time::Duration =
    std::time::Duration::from_secs(300);

/// Hard cap on concurrent HTTP sessions held in
/// `McpServer::sessions`. Configurable via
/// [`McpServerBuilder::with_max_sessions`]. Reaching the cap returns
/// 503 on subsequent `initialize` POSTs until a session is evicted
/// (via `DELETE /mcp` or the idle reaper).
#[cfg(feature = "http")]
pub(crate) const DEFAULT_MAX_SESSIONS: usize = 1024;

/// Divisor for the default per-principal sub-cap relative to
/// [`DEFAULT_MAX_SESSIONS`]. With `DEFAULT_MAX_SESSIONS = 1024` and
/// divisor 16, the default per-principal cap is 64.
#[cfg(feature = "http")]
pub(crate) const DEFAULT_MAX_SESSIONS_PER_PRINCIPAL_DIVISOR: usize = 16;

/// Default per-session SSE replay buffer capacity (frame count).
///
/// Servers retain the most recent N outbound frames per session
/// so a disconnected client can reconnect with `Last-Event-Id`
/// and replay the strictly-newer slice.
#[cfg(feature = "http")]
pub(crate) const DEFAULT_SSE_REPLAY_CAPACITY: usize = 256;

/// Compute the default per-principal session sub-cap from
/// `max_sessions`. Floors at 1 so tiny `max_sessions` configurations
/// never produce a sub-cap of zero (which would reject every
/// authenticated `initialize` POST with no recovery).
#[cfg(feature = "http")]
pub(crate) fn default_max_sessions_per_principal(max_sessions: usize) -> usize {
    (max_sessions / DEFAULT_MAX_SESSIONS_PER_PRINCIPAL_DIVISOR).max(1)
}

/// Default idle-reaper tick cadence. The reaper wakes on this period
/// to scan the session registry for idle entries. Sized for
/// production: a 10-second tick keeps overhead negligible against
/// the default 5-minute idle timeout. Tests override via
/// [`McpServerBuilder::with_idle_reaper_tick`].
#[cfg(feature = "http")]
pub(crate) const DEFAULT_IDLE_REAPER_TICK: std::time::Duration = std::time::Duration::from_secs(10);

/// Builder for [`McpServer`]. Collects one or more invokers (raw
/// `ToolInvoker`s or `Agent`s wrapped as single-tool invokers) and
/// configures server-level concerns such as transports, parent
/// cancel, leader election, tenant binding, and authentication.
///
/// See [`McpServer::builder`] for the entry point and
/// [`Self::build`] / [`Self::build_arc`] for the terminal calls.
pub struct McpServerBuilder {
    invokers: Vec<Arc<dyn ToolInvoker>>,
    parent_cancel: CancellationToken,
    tool_ctx_factory: ToolCtxFactory,
    resume_buffer: Option<std::sync::Arc<dyn klieo_core::resume::ResumeBuffer>>,
    pubsub: Option<std::sync::Arc<dyn klieo_core::Pubsub>>,
    subscribe_cancels: bool,
    #[cfg(feature = "http")]
    publish_permits: Option<usize>,
    leader_kv: Option<Arc<dyn klieo_core::KvStore>>,
    tenant_kv: Option<Arc<dyn klieo_core::KvStore>>,
    /// Durable store for suspended-workflow resume tickets
    /// (ADR-045). Kept separate from `tenant_kv`
    /// (tenant binding is per `tools/call` invocation, scoped under
    /// `klieo-tenants`) — tickets live under their own bucket
    /// (`klieo.mcp.resume-tickets`) and their lifecycle is bound to
    /// the suspended run, not the request. Sharing one underlying
    /// backend is fine and recommended, but the wiring is explicit
    /// so a deployment can opt in to one without the other.
    checkpoint_kv: Option<Arc<dyn klieo_core::KvStore>>,
    tenant_strict: bool,
    profile: klieo_core::DeploymentProfile,
    authenticator: Option<Arc<dyn klieo_auth_common::Authenticator>>,
    leader_ttl: Option<std::time::Duration>,
    leader_heartbeat_interval: Option<std::time::Duration>,
    max_failover_attempts: Option<u32>,
    kv_reaper_interval: Option<std::time::Duration>,
    declare_sampling: bool,
    #[cfg(feature = "http")]
    session_idle_timeout: Option<std::time::Duration>,
    #[cfg(feature = "http")]
    max_sessions: Option<usize>,
    #[cfg(feature = "http")]
    max_sessions_per_principal: Option<usize>,
    #[cfg(feature = "http")]
    sse_replay_capacity: Option<usize>,
    /// Test-only override for the idle-reaper tick cadence. Production
    /// callers cannot reach this; integration + unit tests use it to
    /// drive the reaper fast enough that short idle deadlines (tens to
    /// hundreds of milliseconds) trip within the test timeout. See
    /// [`McpServerBuilder::with_idle_reaper_tick`].
    #[cfg(all(feature = "http", any(test, feature = "test-fixtures")))]
    idle_reaper_tick: Option<std::time::Duration>,
    /// Shared HITL wiring for every workflow registered via
    /// [`Self::add_workflow_with_schema`] / [`Self::add_workflow`].
    /// `None` until the caller invokes [`Self::with_hitl`].
    hitl: Option<crate::workflow::HitlBundle>,
    /// Pending workflow registrations. Materialised into
    /// [`crate::workflow::WorkflowAsToolInvoker`] at `build()` time
    /// once the shared HITL bundle is known. Kept as type-erased
    /// closures so the builder doesn't need a `WorkflowAsToolInvoker<A>`
    /// generic at storage time.
    pending_workflows: Vec<crate::workflow::WorkflowRegistration>,
    /// Shared per-tenant LLM budget governor wired via
    /// [`Self::with_governor`]. `None` keeps the legacy non-governed
    /// path; workflow exposure refuses to build without it. Held
    /// regardless of the `governor` cargo feature so the
    /// `WorkflowWithoutGovernor` hard gate fires uniformly — only the
    /// runtime wire-in lives behind the feature flag.
    governor_bundle: Option<GovernorBundleHolder>,
}

/// Feature-gated payload carried by the builder. The outer
/// [`Option<GovernorBundleHolder>`] field stays present in every build
/// so the hard gate fires uniformly; the inner type erases to a unit
/// struct when the `governor` feature is off (the build still rejects
/// the unsafe configuration, just without a governor to wire).
#[cfg(feature = "governor")]
pub(crate) type GovernorBundleHolder = crate::governor::GovernorBundle;

/// Zero-sized stand-in used when the `governor` feature is disabled.
/// `with_governor` is then unreachable, but the
/// [`McpBuildError::WorkflowWithoutGovernor`] hard gate still fires
/// — the no-feature build rejects workflows outright.
#[cfg(not(feature = "governor"))]
#[derive(Clone)]
pub(crate) struct GovernorBundleHolder;

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

/// Spawn the cluster-0.25 KV reaper iff the builder was asked for it
/// AND at least one bucket (leader or ownership) is wired AND a
/// non-noop resume buffer is available. Returns `None` when any
/// precondition fails — the reaper has nothing useful to do without
/// a bucket to scan or a buffer to query for terminal liveness.
fn spawn_reaper_if_configured(
    interval: Option<std::time::Duration>,
    leader_registry: Option<&klieo_core::LeaderRegistry>,
    ownership_registry: Option<&klieo_core::OwnershipRegistry>,
    resume_buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
) -> Option<klieo_core::KvReaperHandle> {
    let interval = interval?;
    let mut buckets: Vec<String> = Vec::new();
    let kv = if let Some(reg) = leader_registry {
        buckets.push(reg.bucket().to_string());
        reg.kv().clone()
    } else if let Some(reg) = ownership_registry {
        buckets.push(reg.bucket().to_string());
        reg.kv().clone()
    } else {
        return None;
    };
    if let (Some(_), Some(ownership)) = (leader_registry, ownership_registry) {
        buckets.push(ownership.bucket().to_string());
    }
    Some(klieo_core::spawn_kv_reaper(
        kv,
        resume_buffer.clone(),
        buckets,
        interval,
    ))
}

impl McpServerBuilder {
    /// Open a fresh builder with an empty invoker list and a
    /// never-cancelled parent token.
    pub fn new() -> Self {
        Self {
            invokers: Vec::new(),
            parent_cancel: CancellationToken::new(),
            tool_ctx_factory: default_tool_ctx_factory(),
            resume_buffer: None,
            pubsub: None,
            subscribe_cancels: false,
            #[cfg(feature = "http")]
            publish_permits: None,
            leader_kv: None,
            tenant_kv: None,
            checkpoint_kv: None,
            tenant_strict: false,
            profile: klieo_core::DeploymentProfile::Unprofiled,
            authenticator: None,
            leader_ttl: None,
            leader_heartbeat_interval: None,
            max_failover_attempts: None,
            kv_reaper_interval: None,
            declare_sampling: false,
            #[cfg(feature = "http")]
            session_idle_timeout: None,
            #[cfg(feature = "http")]
            max_sessions: None,
            #[cfg(feature = "http")]
            max_sessions_per_principal: None,
            #[cfg(feature = "http")]
            sse_replay_capacity: None,
            #[cfg(all(feature = "http", any(test, feature = "test-fixtures")))]
            idle_reaper_tick: None,
            hitl: None,
            pending_workflows: Vec::new(),
            governor_bundle: None,
        }
    }

    /// Configure the parent [`CancellationToken`] applied to every
    /// agent registered via [`Self::add_agent_with_schema`] or
    /// [`Self::add_agent`]. Cancelling this token propagates into
    /// all in-flight agent runs by overriding the `cancel` field of
    /// each freshly-minted `AgentContext` with a child token.
    ///
    /// Suitable for SIGINT-driven graceful shutdown of the stdio
    /// loop. Raw `ToolInvoker`s registered via [`Self::add_tools`]
    /// are unaffected (they manage their own concurrency).
    pub fn with_parent_cancel(mut self, parent_cancel: CancellationToken) -> Self {
        self.parent_cancel = parent_cancel;
        self
    }

    /// Override the per-request `ToolCtx`. Default = noop in-memory
    /// `Pubsub`/`KvStore`/`JobQueue` (the same wiring `serve_stdio`
    /// used in 0.8/0.9). Override to wire a shared bus so bus-using
    /// tools can be exposed through this server.
    ///
    /// Applies to all transports (stdio + http).
    pub fn with_tool_ctx_factory(mut self, factory: ToolCtxFactory) -> Self {
        self.tool_ctx_factory = factory;
        self
    }

    /// Opt in to SSE resumption via `klieo/tools/resume`. Default is
    /// [`klieo_core::resume::NoopResumeBuffer`] (zero-cost no-op).
    #[must_use]
    pub fn with_resume_buffer(
        mut self,
        buffer: std::sync::Arc<dyn klieo_core::resume::ResumeBuffer>,
    ) -> Self {
        self.resume_buffer = Some(buffer);
        self
    }

    /// Opt in to cross-replica progress-event fanout. Default is a
    /// fresh in-process `MemoryBus` pubsub (single-replica only).
    /// Multi-replica deployments wire a shared `Arc<dyn Pubsub>`
    /// (NATS-backed) via this builder method.
    ///
    /// # Security — progressToken IS the ownership credential
    ///
    /// Per the MCP spec, `_meta.progressToken` is caller-supplied and
    /// klieo treats it as an opaque session identifier: anyone who can
    /// present a progressToken to `klieo/tools/resume` receives the
    /// associated stream's buffered + live events. With cross-replica
    /// fanout, that exposure spans every replica reading from the same
    /// shared pubsub.
    ///
    /// **Operators MUST ensure progressTokens are unguessable** (UUID v4
    /// or equivalent) and authorise their issuance at the tenant
    /// boundary BEFORE the `POST /mcp` request reaches `McpServer`.
    /// `klieo-mcp-server` performs NO progressToken→tenant binding —
    /// any auth proxy in front cannot infer it either. Sharing one
    /// `Arc<dyn Pubsub>` across replicas serving multiple tenants without
    /// an upstream layer that mints + validates per-tenant unguessable
    /// progressTokens is a cross-tenant data-leak vector (CWE-639).
    ///
    /// See ADR-018 for the full threat model.
    #[must_use]
    pub fn with_pubsub(mut self, pubsub: std::sync::Arc<dyn klieo_core::Pubsub>) -> Self {
        self.pubsub = Some(pubsub);
        self
    }

    /// Spawn the wildcard cancel-subject background subscriber on
    /// [`Self::build_arc`]. Required for multi-replica deployments —
    /// without it, only same-replica drop-cancel works. The
    /// background task subscribes to `klieo.mcp.cancel.>` and
    /// dispatches each inbound cancel through the local
    /// `cancel_registry`. Subscribe failure at startup is logged
    /// at `error` and the task exits; replica falls back to
    /// single-replica semantics.
    ///
    /// The spawned task holds an [`Arc`] clone of the server, so
    /// this flag is only honoured by [`Self::build_arc`].
    /// [`Self::build`] panics if this flag is set (the unwrapped
    /// `McpServer` shape cannot keep the background task alive).
    #[must_use]
    pub fn with_cancel_subscription(mut self) -> Self {
        self.subscribe_cancels = true;
        self
    }

    /// Replace the server's publish-concurrency cap. The semaphore is
    /// shared with every SSE-yield publish in `stream_tools_call` (private
    /// helper in `crate::http`) and with the drop-time cross-replica
    /// cancel publish in `CancelOnDrop`, so a single dial bounds all
    /// best-effort fanout work this server performs.
    ///
    /// Pass `0` to drop all fanout publishes (useful in tests that
    /// want to assert the saturation branch — see
    /// `tests/publish_backpressure.rs`). Pass a larger value when the
    /// bus backend has demonstrated headroom and 64 in-flight publishes
    /// are a bottleneck.
    #[cfg(feature = "http")]
    #[must_use]
    pub fn with_publish_concurrency(mut self, permits: usize) -> Self {
        self.publish_permits = Some(permits);
        self
    }

    /// Opt in to leader election for multi-replica orphan
    /// detection. On `tools/call` start the server claims
    /// leadership in the `klieo-leaders` KV bucket; on
    /// `klieo/tools/resume` the server checks the leader is
    /// alive and writes a terminal "leader died" SSE frame to
    /// the resume buffer if not.
    ///
    /// Operators MUST configure the `klieo-leaders` JetStream KV
    /// bucket with `max_age = 5s` (matches the cluster's fixed
    /// TTL). See ADR-020.
    #[must_use]
    pub fn with_leader_election(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
        self.leader_kv = Some(kv);
        self
    }

    /// Opt in to tenant binding. On `tools/call` start the
    /// server claims ownership in the `klieo-tenants` KV
    /// bucket keyed by the authenticated `Identity.principal`;
    /// on `klieo/tools/resume` the server rejects mismatched
    /// principals as `JSONRPC_STREAM_NOT_FOUND` (-32004; deny-
    /// as-NotFound per OWASP IDOR best practice).
    ///
    /// Operators MUST also wire an `Authenticator` via
    /// `with_authenticator` that produces non-anonymous
    /// identities (typically OAuth via cluster 0.21); without
    /// it no entries are written and no checks run. See
    /// ADR-022.
    #[must_use]
    pub fn with_tenant_binding(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
        self.tenant_kv = Some(kv);
        self.tenant_strict = false;
        self
    }

    /// Like [`with_tenant_binding`](Self::with_tenant_binding) but **fail-closed**
    /// on store error: when the `klieo-tenants` KV is unreachable, an invoke
    /// claim and a `klieo/tools/resume` ownership check both DENY (503) rather
    /// than proceed, closing the transient-KV-blip cross-tenant-resume window
    /// for regulated multi-tenant deployments (at the cost of availability
    /// during a KV outage). An unclaimed key still proceeds. See ADR-022.
    #[must_use]
    pub fn with_tenant_binding_strict(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
        self.tenant_kv = Some(kv);
        self.tenant_strict = true;
        self
    }

    /// Opt in to principal-scoped resume tickets (ADR-045). When a
    /// workflow registered via
    /// [`Self::add_workflow_with_schema`] suspends on a `ReviewPolicy`,
    /// the server persists the checkpoint to `kv` under an opaque
    /// 256-bit token bound to the verified caller principal and
    /// returns the token in the wire envelope. A later
    /// `klieo/run/resume` request authorises the token against the
    /// caller before atomically consuming it.
    ///
    /// Distinct from [`Self::with_tenant_binding`] (per-`tools/call`
    /// ownership claim in the `klieo-tenants` bucket) — tickets live
    /// under `klieo.mcp.resume-tickets` and carry the checkpoint
    /// payload across the operator-review window. Sharing one KV
    /// backend across both is fine.
    ///
    /// Leaving this unwired keeps the slice-1 no-ticket suspend
    /// envelope so unprincipaled / unauthenticated deployments stay
    /// backward-compatible.
    #[must_use]
    pub fn with_checkpoint_kv(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
        self.checkpoint_kv = Some(kv);
        self
    }

    /// Wire an [`klieo_auth_common::Authenticator`] into the HTTP
    /// entry path. `post_mcp` calls `authenticate(headers)` +
    /// `authorize_method(&identity, method)` before dispatch /
    /// SSE upgrade. Auth failures return a JSON-RPC -32001
    /// envelope (no SSE upgrade on failure).
    ///
    /// Designed for `klieo_auth_oauth::OAuthAuthenticator` but
    /// any `Authenticator` impl works (e.g. an in-house bearer
    /// validator, or `klieo_auth_common::AllowAnonymous` for
    /// explicit no-auth wiring).
    #[must_use]
    pub fn with_authenticator(
        mut self,
        authenticator: Arc<dyn klieo_auth_common::Authenticator>,
    ) -> Self {
        self.authenticator = Some(authenticator);
        self
    }

    /// Apply a [`DeploymentProfile`](klieo_core::DeploymentProfile). Regulated
    /// profiles force strict tenant binding + require a non-anonymous
    /// authenticator; `build()` fails closed if a prerequisite is missing.
    pub fn profile(mut self, profile: klieo_core::DeploymentProfile) -> Self {
        self.profile = profile;
        self
    }

    /// Override the cluster-0.20 leader TTL (default [`LEADER_TTL`],
    /// 5s). Wider TTL tolerates network blips at the cost of slower
    /// dead-leader detection; tighter TTL detects deaths faster but
    /// risks spurious orphan probes under load. The heartbeat
    /// interval defaults to `ttl / 2` — override independently via
    /// [`Self::with_leader_heartbeat_interval`]. Cluster 0.25.
    #[must_use]
    pub fn with_leader_ttl(mut self, ttl: std::time::Duration) -> Self {
        self.leader_ttl = Some(ttl);
        self
    }

    /// Override the leader heartbeat interval. Default is half the
    /// configured TTL (or [`LEADER_TTL`] / 2 when the TTL itself is
    /// at its default). Tighten when KV write latency is high and
    /// you want extra headroom before TTL expires. Cluster 0.25.
    #[must_use]
    pub fn with_leader_heartbeat_interval(mut self, interval: std::time::Duration) -> Self {
        self.leader_heartbeat_interval = Some(interval);
        self
    }

    /// Override the cluster-0.24 failover-attempt cap (default
    /// [`klieo_core::FAILOVER_ATTEMPT_CAP`], 3). Higher cap
    /// accommodates flaky downstream where bad-state tools are
    /// rare; lower cap exits failover loops faster. Cluster 0.25.
    #[must_use]
    pub fn with_max_failover_attempts(mut self, cap: u32) -> Self {
        self.max_failover_attempts = Some(cap);
        self
    }

    /// Opt in to the cluster-0.25 KV reaper. Spawns a background
    /// task scanning `klieo-leaders` (+ `klieo-tenants` when tenant
    /// binding is wired) every `interval`, evicting entries whose
    /// resume buffer is terminal. The scan task is spawned at
    /// [`Self::build`] / [`Self::build_arc`] time (the `McpServer`
    /// holds both leader-KV / ownership-KV via its registries AND
    /// the resume buffer directly); calling this method without
    /// [`Self::with_resume_buffer`] AND at least one of
    /// [`Self::with_leader_election`] / [`Self::with_tenant_binding`]
    /// is a no-op.
    ///
    /// Recommended: 60s for production NATS deployments running
    /// `klieo-leaders` / `klieo-tenants` without bucket TTLs. Bucket
    /// TTLs are still the primary eviction mechanism — the reaper
    /// is a backstop for deployments where TTLs are not configured.
    #[must_use]
    pub fn with_kv_reaper(mut self, interval: std::time::Duration) -> Self {
        self.kv_reaper_interval = Some(interval);
        self
    }

    /// Declare that this server may issue `sampling/createMessage`
    /// outbound requests to the client. Sets
    /// `capabilities.sampling = {}` on the initialize response so
    /// MCP-conformant clients know to listen.
    ///
    /// This is the precondition for constructing the outbound
    /// correlation primitive: until the flag is set, the stdio read
    /// loop drops inbound JSON-RPC responses (no table to route them
    /// into). Tools-only deployments leave the flag unset and pay
    /// zero outbound cost.
    #[must_use]
    pub fn with_client_sampling(mut self) -> Self {
        self.declare_sampling = true;
        self
    }

    /// Configure the HTTP streamable-session idle timeout. Default:
    /// 5 minutes (see `DEFAULT_SESSION_IDLE_TIMEOUT`). Pass
    /// [`std::time::Duration::ZERO`] to disable the per-session
    /// watchdog entirely (sessions then live until the client SSE
    /// disconnects). ADR-028.
    #[cfg(feature = "http")]
    #[must_use]
    pub fn with_session_idle_timeout(mut self, ttl: std::time::Duration) -> Self {
        self.session_idle_timeout = Some(ttl);
        self
    }

    /// Override the `DEFAULT_MAX_SESSIONS` cap on concurrent HTTP
    /// sessions held in the server's private `sessions` registry.
    /// Reaching the cap causes `handle_initialize_post` to return 503
    /// until a session is evicted via `DELETE /mcp` or the idle reaper.
    ///
    /// Panics when `cap == 0` — a zero cap would deadlock the
    /// `initialize` path with no recovery, which is a caller bug.
    #[cfg(feature = "http")]
    #[must_use]
    pub fn with_max_sessions(mut self, cap: usize) -> Self {
        assert!(cap > 0, "max_sessions must be > 0");
        self.max_sessions = Some(cap);
        self
    }

    /// Override the default per-principal session sub-cap. The
    /// default is `default_max_sessions_per_principal` applied to
    /// the effective `max_sessions` —
    /// `max_sessions / DEFAULT_MAX_SESSIONS_PER_PRINCIPAL_DIVISOR`
    /// floored at 1. Reaching the sub-cap causes
    /// `handle_initialize_post` to return 503 on subsequent
    /// `initialize` POSTs from the same authenticated principal until
    /// one of that principal's sessions is evicted (via `DELETE /mcp`
    /// or the idle reaper).
    ///
    /// Panics when `cap == 0` — a zero sub-cap would reject every
    /// authenticated `initialize` with no recovery, which is a caller
    /// bug.
    #[cfg(feature = "http")]
    #[must_use]
    pub fn with_max_sessions_per_principal(mut self, cap: usize) -> Self {
        assert!(cap > 0, "max_sessions_per_principal must be > 0");
        self.max_sessions_per_principal = Some(cap);
        self
    }

    /// Cap on the per-session SSE replay buffer (frame count).
    ///
    /// Defaults to `DEFAULT_SSE_REPLAY_CAPACITY` (256). A value of
    /// `0` disables resumption; reconnects with `Last-Event-Id` then
    /// return 501 Not Implemented.
    #[cfg(feature = "http")]
    #[must_use]
    pub fn with_sse_replay_capacity(mut self, capacity: usize) -> Self {
        self.sse_replay_capacity = Some(capacity);
        self
    }

    /// Override the idle-reaper tick cadence used by the per-server
    /// background task. Production callers leave this unset — the
    /// reaper wakes every 10 seconds, which is the right cadence for
    /// minute-scale `session_idle_timeout` values. Tests configure
    /// short idle deadlines (tens of milliseconds) and rely on a
    /// matching short tick so the reaper fires within the test
    /// timeout.
    ///
    /// Gated on `test-fixtures` (or `#[cfg(test)]` inside the crate)
    /// so the production wire surface stays fixed.
    #[cfg(all(feature = "http", any(test, feature = "test-fixtures")))]
    #[must_use]
    pub fn with_idle_reaper_tick(mut self, tick: std::time::Duration) -> Self {
        self.idle_reaper_tick = Some(tick);
        self
    }

    /// Append a raw [`ToolInvoker`]. The invoker's catalogue will
    /// be merged into the server's `tools/list` response.
    pub fn add_tools(mut self, invoker: Arc<dyn ToolInvoker>) -> Self {
        self.invokers.push(invoker);
        self
    }

    /// Register an [`Agent`] as a single MCP tool with the
    /// caller-supplied JSON Schema. See [`McpServer::expose_agent_with_schema`]
    /// for the dispatch semantics.
    pub fn add_agent_with_schema<A>(
        mut self,
        agent: A,
        input_schema: serde_json::Value,
        ctx_factory: AgentContextFactory,
    ) -> Self
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + Send + 'static,
        A::Output: serde::Serialize + Send + 'static,
    {
        let name = agent.name().to_string();
        let invoker: Arc<dyn ToolInvoker> = Arc::new(AgentAsToolInvoker {
            agent: Arc::new(agent),
            name,
            input_schema,
            ctx_factory,
            // Governed agents opt in iff `with_governor` wired the
            // bundle at registration time. Cloning here is
            // cheap (Arc + ProviderId) and binds the invoker to the
            // governor visible at registration — later `with_governor`
            // calls do not retroactively govern earlier-registered
            // agents.
            #[cfg(feature = "governor")]
            governor: self.governor_bundle.clone(),
        });
        self.invokers.push(invoker);
        self
    }

    /// Auto-derive variant of [`Self::add_agent_with_schema`].
    /// Requires `A::Input: schemars::JsonSchema` and the `schemars`
    /// cargo feature on `klieo-mcp-server`.
    #[cfg(feature = "schemars")]
    pub fn add_agent<A>(self, agent: A, ctx_factory: AgentContextFactory) -> Self
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
        A::Output: serde::Serialize + Send + 'static,
    {
        let schema = serde_json::to_value(schemars::schema_for!(A::Input))
            .expect("schemars::Schema serialises to JSON via #[derive(Serialize)]");
        self.add_agent_with_schema(agent, schema, ctx_factory)
    }

    /// Wire the inbound per-tenant LLM budget governor.
    ///
    /// `governor` enforces the RPS cap; `provider` is the
    /// [`klieo_ops::ProviderId`] every inbound LLM call is scoped to.
    /// Workflows registered via [`Self::add_workflow_with_schema`] /
    /// [`Self::add_workflow`] are hard-gated on this wiring — the
    /// build refuses to ship a workflow-exposing server without a
    /// governor and surfaces [`McpBuildError::WorkflowWithoutGovernor`].
    /// Governed agents (registered via [`Self::add_agent`]) opt in
    /// automatically: the governor wraps their `ctx.llm` whenever
    /// one is wired, otherwise they keep the legacy ungoverned path.
    ///
    /// Available only with the `governor` cargo feature. Without
    /// the feature, the build still refuses workflows (the hard gate
    /// fires regardless), so adopters always know exposure of a paid
    /// LLM path through MCP must be explicit.
    #[cfg(feature = "governor")]
    #[must_use]
    pub fn with_governor(
        mut self,
        governor: Arc<dyn klieo_ops::governor::Governor>,
        provider: klieo_ops::ProviderId,
    ) -> Self {
        self.governor_bundle = Some(crate::governor::GovernorBundle { governor, provider });
        self
    }

    /// Configure the shared HITL client + config consumed by every
    /// workflow registered via [`Self::add_workflow_with_schema`] /
    /// [`Self::add_workflow`]. Building with a workflow but no HITL
    /// wiring fails with [`McpBuildError::WorkflowWithoutHitl`] — the
    /// workflow path has no compliance endpoint to submit to without
    /// it (ADR-045).
    pub fn with_hitl(
        mut self,
        client: Arc<klieo_hitl_client::HitlClient>,
        cfg: Arc<klieo_hitl::HitlConfig>,
    ) -> Self {
        self.hitl = Some(crate::workflow::HitlBundle { client, cfg });
        self
    }

    /// Register an [`Agent`] as a single MCP tool whose `tools/call`
    /// drives [`klieo_hitl::run_with_hitl`] instead of the bare
    /// [`Agent::run`]. `system_prompt` is the prompt the underlying
    /// `run_steps` loop uses; `run_options` carries the
    /// [`klieo_core::runtime::ReviewPolicy`] that decides when to
    /// suspend. On suspend the response is
    /// `{"status":"suspended","reason":...}` (no checkpoint / ticket
    /// when no resume KV is wired; ADR-045).
    ///
    /// Building this server without first calling [`Self::with_hitl`]
    /// fails with [`McpBuildError::WorkflowWithoutHitl`].
    pub fn add_workflow_with_schema<A>(
        mut self,
        agent: A,
        system_prompt: impl Into<String>,
        input_schema: serde_json::Value,
        run_options: klieo_core::runtime::RunOptions,
        ctx_factory: AgentContextFactory,
    ) -> Self
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + Send + 'static,
    {
        let name = agent.name().to_string();
        let prompt = system_prompt.into();
        // The workflow body never calls `A::run`; the agent value is
        // captured purely to anchor the generic `A` (which carries
        // `A::Input`'s DeserializeOwned + Send bounds through the
        // type-erased registration). PhantomData on the invoker
        // does the same job once materialised.
        drop(agent);
        let materialise: crate::workflow::WorkflowMaterialiser = Box::new(
            move |bundle: crate::workflow::HitlBundle,
                  ticket_store: Option<Arc<crate::resume_ticket::ResumeTicketStore>>,
                  governor_bundle: Option<crate::GovernorBundleHolder>| {
                #[cfg(not(feature = "governor"))]
                let _ = governor_bundle;
                let invoker = Arc::new(crate::workflow::WorkflowAsToolInvoker::<A>::new(
                    name.clone(),
                    prompt.clone(),
                    input_schema.clone(),
                    ctx_factory.clone(),
                    run_options.clone(),
                    bundle,
                    ticket_store,
                    #[cfg(feature = "governor")]
                    governor_bundle,
                ));
                crate::workflow::WorkflowMaterialisation {
                    name: name.clone(),
                    resume_handle: invoker.clone()
                        as Arc<dyn crate::workflow::WorkflowResumeHandle>,
                    invoker: invoker as Arc<dyn ToolInvoker>,
                }
            },
        );
        self.pending_workflows
            .push(crate::workflow::WorkflowRegistration { materialise });
        self
    }

    /// Auto-derive variant of [`Self::add_workflow_with_schema`].
    /// Requires `A::Input: schemars::JsonSchema` and the `schemars`
    /// cargo feature on `klieo-mcp-server`.
    #[cfg(feature = "schemars")]
    pub fn add_workflow<A>(
        self,
        agent: A,
        system_prompt: impl Into<String>,
        run_options: klieo_core::runtime::RunOptions,
        ctx_factory: AgentContextFactory,
    ) -> Self
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
    {
        let schema = serde_json::to_value(schemars::schema_for!(A::Input))
            .expect("schemars::Schema serialises to JSON via #[derive(Serialize)]");
        self.add_workflow_with_schema(agent, system_prompt, schema, run_options, ctx_factory)
    }

    /// Finalise the builder into a runnable [`McpServer`].
    ///
    /// - **0 invokers** → [`McpBuildError::NoInvokers`].
    /// - **1 invoker** → wrap directly with no merge overhead.
    /// - **≥2 invokers** → wrap in `MergedInvoker`. Returns
    ///   [`McpBuildError::DuplicateTool`] on the first duplicate tool
    ///   name across the merged catalogues.
    /// - [`Self::with_cancel_subscription`] set → [`McpBuildError::CancelRequiresArc`].
    pub fn build(self) -> Result<McpServer, McpBuildError> {
        if self.subscribe_cancels {
            return Err(McpBuildError::CancelRequiresArc);
        }
        self.build_inner()
    }

    fn build_inner(mut self) -> Result<McpServer, McpBuildError> {
        let ticket_store = self
            .checkpoint_kv
            .clone()
            .map(|kv| Arc::new(crate::resume_ticket::ResumeTicketStore::new(kv)));
        let mut workflow_resume_handles: std::collections::HashMap<
            String,
            Arc<dyn crate::workflow::WorkflowResumeHandle>,
        > = std::collections::HashMap::new();
        if !self.pending_workflows.is_empty() {
            let bundle = self
                .hitl
                .clone()
                .ok_or(McpBuildError::WorkflowWithoutHitl)?;
            // Hard gate: workflow exposure is illegal without a
            // governor — an inbound MCP caller must not be
            // able to drive unbounded paid LLM spend. Mirrors
            // `WorkflowWithoutHitl`. The gate fires only when the
            // `governor` cargo feature is enabled; the off-feature
            // build retains the legacy ungoverned path because the
            // adopter explicitly opted out of inbound budget
            // enforcement at compile time.
            #[cfg(feature = "governor")]
            if self.governor_bundle.is_none() {
                return Err(McpBuildError::WorkflowWithoutGovernor);
            }
            // The off-feature build path is unreachable because the
            // hard gate above already errored; we still propagate the
            // (always-`None`) holder so the materialiser signature
            // stays feature-uniform.
            let governor_bundle = self.governor_bundle.clone();
            let pending = std::mem::take(&mut self.pending_workflows);
            for reg in pending {
                let mat = (reg.materialise)(
                    bundle.clone(),
                    ticket_store.clone(),
                    governor_bundle.clone(),
                );
                if workflow_resume_handles
                    .insert(mat.name.clone(), mat.resume_handle)
                    .is_some()
                {
                    return Err(McpBuildError::DuplicateTool(mat.name));
                }
                self.invokers.push(mat.invoker);
            }
        }
        let invoker_count = self.invokers.len();
        if invoker_count == 0 {
            return Err(McpBuildError::NoInvokers);
        }
        let invoker: Arc<dyn ToolInvoker> = if invoker_count == 1 {
            self.invokers.into_iter().next().unwrap() // safe: len == 1
        } else {
            Arc::new(MergedInvoker::new(self.invokers)?)
        };
        #[cfg(feature = "http")]
        let permits = self.publish_permits.unwrap_or(DEFAULT_PUBLISH_PERMITS);
        #[cfg(feature = "http")]
        let max_sessions = self.max_sessions.unwrap_or(DEFAULT_MAX_SESSIONS);
        #[cfg(feature = "http")]
        let max_sessions_per_principal = self
            .max_sessions_per_principal
            .unwrap_or_else(|| default_max_sessions_per_principal(max_sessions));
        let leader_registry = self.leader_kv.map(|kv| {
            klieo_core::LeaderRegistry::new(
                kv,
                "klieo-leaders".into(),
                uuid::Uuid::new_v4().to_string(),
            )
        });
        let profile = self.profile;
        profile.validate(
            self.tenant_kv.is_some(),
            self.authenticator.as_ref().map(|a| a.allows_anonymous()),
        )?;
        let tenant_strict = self.tenant_strict || profile.requires_strict_binding();
        let ownership_registry = self.tenant_kv.map(|kv| {
            let bucket = "klieo-tenants".into();
            if tenant_strict {
                klieo_core::OwnershipRegistry::new_strict(kv, bucket)
            } else {
                klieo_core::OwnershipRegistry::new(kv, bucket)
            }
        });
        if profile.requires_strict_binding() || profile.requires_named_principal() {
            tracing::warn!(
                target: "klieo.security",
                cwe = 639,
                "regulated multi-tenant profile active on this replica; \
                 cross-replica tenant isolation assumes ALL replicas run the \
                 same profile — a lenient peer reintroduces CWE-639. Fleet \
                 homogeneity is NOT verified by this replica."
            );
        }
        let resume_buffer = self
            .resume_buffer
            .unwrap_or_else(|| std::sync::Arc::new(klieo_core::resume::NoopResumeBuffer));
        let leader_ttl = self.leader_ttl.unwrap_or(LEADER_TTL);
        let leader_heartbeat_interval = self.leader_heartbeat_interval.unwrap_or(leader_ttl / 2);
        let max_failover_attempts = self
            .max_failover_attempts
            .unwrap_or(klieo_core::FAILOVER_ATTEMPT_CAP);
        let kv_reaper = spawn_reaper_if_configured(
            self.kv_reaper_interval,
            leader_registry.as_ref(),
            ownership_registry.as_ref(),
            &resume_buffer,
        );
        Ok(McpServer {
            invoker,
            tool_ctx_factory: self.tool_ctx_factory,
            parent_cancel: self.parent_cancel,
            resume_buffer,
            pubsub: self
                .pubsub
                .unwrap_or_else(|| klieo_bus_memory::MemoryBus::new().pubsub.clone()),
            cancel_registry: klieo_core::CancelRegistry::new(),
            #[cfg(feature = "http")]
            publish_permits: std::sync::Arc::new(tokio::sync::Semaphore::new(permits)),
            leader_registry,
            ownership_registry,
            resume_ticket_store: ticket_store,
            workflow_resume_handles,
            authenticator: self.authenticator,
            leader_ttl,
            leader_heartbeat_interval,
            max_failover_attempts,
            kv_reaper_interval: self.kv_reaper_interval,
            _kv_reaper: kv_reaper,
            // Stdio session container starts empty. The stdio entry
            // path populates it on first inbound request; HTTP
            // transports never touch it.
            stdio_session: tokio::sync::OnceCell::new(),
            // HTTP session registry starts empty. The initialize POST
            // handler inserts a fresh entry when a client mints a
            // session; the reaper and DELETE handler remove entries.
            #[cfg(feature = "http")]
            sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
                std::collections::HashMap::new(),
            )),
            #[cfg(feature = "http")]
            max_sessions,
            #[cfg(feature = "http")]
            max_sessions_per_principal,
            #[cfg(feature = "http")]
            sse_replay_capacity: self
                .sse_replay_capacity
                .unwrap_or(DEFAULT_SSE_REPLAY_CAPACITY),
            #[cfg(feature = "http")]
            principal_counts: std::sync::Arc::new(tokio::sync::RwLock::new(
                std::collections::HashMap::new(),
            )),
            // Idle-reaper spawn sentinel starts empty. The first
            // initialize POST that needs the reaper populates it via
            // `ensure_idle_reaper`; subsequent populates are no-ops.
            #[cfg(feature = "http")]
            idle_reaper_started: tokio::sync::OnceCell::new(),
            declare_sampling: self.declare_sampling,
            // Shared stdout writer is minted lazily by `serve_stdio`
            // on first call (build-time stays sync; the writer needs
            // an async runtime context).
            stdout_writer: tokio::sync::OnceCell::new(),
            // Client capabilities populated by the `initialize` arm.
            client_caps: tokio::sync::Mutex::new(ClientCaps::default()),
            #[cfg(feature = "http")]
            session_idle_timeout: self
                .session_idle_timeout
                .unwrap_or(DEFAULT_SESSION_IDLE_TIMEOUT),
            #[cfg(all(feature = "http", any(test, feature = "test-fixtures")))]
            idle_reaper_tick: self.idle_reaper_tick.unwrap_or(DEFAULT_IDLE_REAPER_TICK),
            #[cfg(all(feature = "http", not(any(test, feature = "test-fixtures"))))]
            idle_reaper_tick: DEFAULT_IDLE_REAPER_TICK,
            // Reference instant for per-session `last_activity_millis`.
            // Captured once at build time and never mutated; lives for
            // the lifetime of the server.
            #[cfg(feature = "http")]
            server_start: std::time::Instant::now(),
        })
    }

    /// Identical to [`Self::build`] but returns `Arc<McpServer>`,
    /// which is required by the HTTP transport (it shares state
    /// across handler invocations).
    ///
    /// If [`Self::with_cancel_subscription`] was called, this also
    /// spawns the wildcard `klieo.mcp.cancel.>` background task
    /// against an [`Arc`] clone of the returned server.
    pub fn build_arc(self) -> Result<std::sync::Arc<McpServer>, McpBuildError> {
        let spawn_subscriber = self.subscribe_cancels;
        let server = std::sync::Arc::new(self.build_inner()?);
        if spawn_subscriber {
            klieo_core::cancel::spawn_wildcard_cancel_subscriber(
                server.pubsub.clone(),
                "klieo.mcp.cancel.>".to_string(),
                "klieo.mcp.cancel.".to_string(),
                server.cancel_registry.clone(),
                "mcp.cancel",
            );
        }
        Ok(server)
    }
}

impl McpServer {
    /// Open a [`McpServerBuilder`] for accumulating tools / agents
    /// and configuring server-level concerns such as a parent
    /// [`CancellationToken`].
    ///
    /// The three shorthand constructors below
    /// ([`Self::expose_tools`], [`Self::expose_agent_with_schema`],
    /// [`Self::expose_agent`]) build on this — they are zero-config
    /// shims for the single-tool / single-agent cases.
    pub fn builder() -> McpServerBuilder {
        McpServerBuilder::new()
    }

    /// Borrow the leader registry if configured via the builder's
    /// `with_leader_election`. None when running single-replica
    /// without orphan detection.
    pub fn leader_registry(&self) -> Option<&klieo_core::LeaderRegistry> {
        self.leader_registry.as_ref()
    }

    /// Borrow the ownership registry if configured via the
    /// builder's `with_tenant_binding`. None when running
    /// without tenant binding (pre-0.22 deployments).
    pub fn ownership_registry(&self) -> Option<&klieo_core::OwnershipRegistry> {
        self.ownership_registry.as_ref()
    }

    /// Borrow the configured authenticator if any.
    pub fn authenticator(&self) -> Option<&Arc<dyn klieo_auth_common::Authenticator>> {
        self.authenticator.as_ref()
    }

    /// Configured leader TTL (cluster-0.20 + cluster-0.25 tunable).
    /// Default [`LEADER_TTL`] when
    /// [`McpServerBuilder::with_leader_ttl`] was not called.
    pub fn leader_ttl(&self) -> std::time::Duration {
        self.leader_ttl
    }

    /// Configured leader heartbeat interval. Default is half the
    /// configured TTL when
    /// [`McpServerBuilder::with_leader_heartbeat_interval`] was not
    /// called.
    pub fn leader_heartbeat_interval(&self) -> std::time::Duration {
        self.leader_heartbeat_interval
    }

    /// Configured failover-attempt cap (cluster-0.24 + cluster-0.25
    /// tunable). Default [`klieo_core::FAILOVER_ATTEMPT_CAP`] when
    /// [`McpServerBuilder::with_max_failover_attempts`] was not
    /// called.
    pub fn max_failover_attempts(&self) -> u32 {
        self.max_failover_attempts
    }

    /// Configured KV-reaper interval, if any. `None` when
    /// [`McpServerBuilder::with_kv_reaper`] was not called or when
    /// no leader/ownership KV was wired (no bucket to scan).
    pub fn kv_reaper_interval(&self) -> Option<std::time::Duration> {
        self.kv_reaper_interval
    }

    /// Snapshot of every live HTTP session id, taken under the
    /// registry read-lock. The lock is released before this returns,
    /// so the snapshot's relationship with concurrent `DELETE /mcp`
    /// requests or idle-reaper evictions depends on observation
    /// timing. Order is undefined. Returns an empty vector on stdio
    /// servers and on HTTP servers with no live sessions.
    #[cfg(feature = "http")]
    pub async fn session_ids(&self) -> Vec<uuid::Uuid> {
        self.sessions.read().await.keys().copied().collect()
    }

    /// Close state of a specific HTTP session, looked up by minted
    /// `Mcp-Session-Id`. Returns `None` when the id is unknown — the
    /// session either never existed or has already been evicted via
    /// `DELETE /mcp` or the idle reaper. Returns `Some(true)` after
    /// the session's drain path has run; `Some(false)` while the
    /// session is still serving traffic.
    #[cfg(feature = "http")]
    pub async fn is_session_closed_by_id(&self, id: uuid::Uuid) -> Option<bool> {
        self.sessions.read().await.get(&id).map(|s| s.is_closed())
    }

    /// Returns `true` when SSE resumption is enabled
    /// (`sse_replay_capacity > 0`). Read by the SSE producer to
    /// gate buffer writes and by the `GET /mcp` resume branch to
    /// reject `Last-Event-Id` with 501 when disabled.
    #[cfg(feature = "http")]
    pub(crate) fn sse_replay_enabled(&self) -> bool {
        self.sse_replay_capacity > 0
    }

    /// Decrement the per-principal session count for `principal`.
    /// Removes the entry when the count reaches zero so the HashMap
    /// does not grow unbounded under principal churn. No-op when
    /// `principal` is `None` (stdio session or auth-disabled
    /// deployment, neither of which populates the count cache).
    ///
    /// Callers MUST release any [`Self::sessions`] write guard before
    /// invoking this method to preserve the documented lock-acquisition
    /// order.
    #[cfg(feature = "http")]
    pub(crate) async fn decrement_principal_count(&self, principal: Option<&str>) {
        let Some(principal) = principal else { return };
        let mut counts = self.principal_counts.write().await;
        if let Some(entry) = counts.get_mut(principal) {
            *entry = entry.saturating_sub(1);
            if *entry == 0 {
                counts.remove(principal);
            }
        }
    }

    /// Borrow the stdio transport's server-initiated outbound channel
    /// — `Some` once the stdio loop has wired it (i.e. the builder
    /// opted in via [`McpServerBuilder::with_client_sampling`] AND
    /// [`Self::serve_with_streams`] has run at least once), `None`
    /// otherwise. Surfaces an `Arc<dyn ServerOutbound>` so callers
    /// thread it onto `ToolCtx::server_outbound` in their own
    /// [`McpServerBuilder::with_tool_ctx_factory`] closures.
    ///
    /// Scoped to the stdio transport: HTTP servers carry one outbound
    /// per session, accessible via the per-session registry rather
    /// than through this accessor.
    pub fn outbound(&self) -> Option<Arc<dyn klieo_core::ServerOutbound>> {
        self.stdio_session
            .get()
            .and_then(|session| session.outbound.get())
            .map(|o| o.clone() as Arc<dyn klieo_core::ServerOutbound>)
    }

    /// Push a notification-shaped JSON-RPC frame through the named
    /// HTTP session's wired outbound primitive. Used by the
    /// slow-SSE-consumer integration test to drive the per-session
    /// ring past [`crate::outbound_sink::OUTBOUND_QUEUE_CAPACITY`]
    /// without fabricating pending-response entries — notifications
    /// carry no `id`, so the table stays empty while the ring fills.
    /// The `payload_bytes` parameter inflates each frame so the
    /// kernel TCP send buffer saturates before the ring does —
    /// otherwise the receiver dequeues fast enough that drop-oldest
    /// never fires.
    ///
    /// Returns `Ok(())` when the frame was enqueued (including the
    /// drop-oldest path, which the sink reports out-of-band via
    /// metrics and a `warn` log). Returns `Err(())` when the session
    /// id is unknown OR when no outbound primitive has been wired
    /// (caller must run `POST /mcp initialize` then `GET /mcp` first).
    ///
    /// Gated on `test-fixtures` so production callers cannot reach it.
    /// Test-only snapshot of the per-HTTP-session roots cache, keyed
    /// by minted `Mcp-Session-Id`. Returns `None` when the id is
    /// unknown OR the roots cache has not been wired (the server was
    /// not built with `with_client_sampling`). Gated on
    /// `test-fixtures` so production callers cannot reach it.
    #[cfg(all(feature = "http", feature = "test-fixtures"))]
    pub async fn client_roots_for_session(&self, session_id: uuid::Uuid) -> Option<Vec<Root>> {
        let sessions = self.sessions.read().await;
        let session = sessions.get(&session_id)?;
        let cache = session.roots_cache.get()?;
        Some(cache.snapshot())
    }

    /// Test-only lookup of a session's wired outbound primitive,
    /// keyed by minted `Mcp-Session-Id`. Returns `Some` once the SSE
    /// `GET /mcp` for the session has installed the primitive and the
    /// session is still in the registry; `None` for unknown ids or
    /// before the SSE wiring completes. Gated on `test-fixtures` so
    /// production callers cannot reach it — see [`Self::session_ids`]
    /// for the production-facing accessor.
    #[cfg(all(feature = "http", feature = "test-fixtures"))]
    pub async fn outbound_for_session(
        &self,
        session_id: uuid::Uuid,
    ) -> Option<Arc<dyn klieo_core::ServerOutbound>> {
        let sessions = self.sessions.read().await;
        let session = sessions.get(&session_id)?;
        let outbound = session.outbound.get()?.clone();
        Some(outbound as Arc<dyn klieo_core::ServerOutbound>)
    }

    /// Test-only read of the resolved per-session SSE replay buffer
    /// capacity. Returns the value handed in via
    /// [`McpServerBuilder::with_sse_replay_capacity`] or the builder
    /// default when the knob was not set. Gated on `test-fixtures`
    /// so production callers cannot reach it.
    #[cfg(all(feature = "http", feature = "test-fixtures"))]
    pub fn sse_replay_capacity(&self) -> usize {
        self.sse_replay_capacity
    }

    /// Test-only snapshot of a session's SSE replay buffer as a list
    /// of `(event_id, frame)` pairs in delivery order. Returns
    /// `None` when the session id is absent from the registry. Gated
    /// on `test-fixtures` so production callers cannot reach it.
    #[cfg(all(feature = "http", feature = "test-fixtures"))]
    pub async fn sse_replay_snapshot(
        &self,
        session_id: uuid::Uuid,
    ) -> Option<Vec<(u64, std::sync::Arc<serde_json::Value>)>> {
        let sessions = self.sessions.read().await;
        let session = sessions.get(&session_id)?;
        let buffer = session.sse_replay_buffer.lock();
        Some(buffer.iter().cloned().collect())
    }

    /// Push a notification-shaped JSON-RPC frame through the named
    /// HTTP session's wired outbound primitive — used by the slow-SSE
    /// backpressure integration test to drive the per-session ring past
    /// [`crate::outbound_sink::OUTBOUND_QUEUE_CAPACITY`]. Each frame is
    /// inflated to `payload_bytes` so the kernel TCP send buffer
    /// saturates before the ring does. Returns `Err(())` when the
    /// session id is unknown or no outbound primitive is wired.
    ///
    /// Gated on `test-fixtures` so production callers cannot reach it.
    #[cfg(all(feature = "http", feature = "test-fixtures"))]
    pub async fn emit_test_notification(
        &self,
        session_id: uuid::Uuid,
        method: &str,
        payload_bytes: usize,
    ) -> Result<(), ()> {
        let session = {
            let sessions = self.sessions.read().await;
            sessions.get(&session_id).cloned()
        };
        let Some(outbound) = session.as_ref().and_then(|s| s.outbound.get()) else {
            return Err(());
        };
        outbound
            .send_notification_frame(method, payload_bytes)
            .await
            .map_err(|_| ())
    }

    /// Return the latest cached client roots for the stdio transport.
    /// Returns an empty vec when the client did not advertise roots
    /// support, when the outbound primitive was not enabled at build
    /// time (no cache to query), or when the server is HTTP-only.
    /// Callers that want to react to changes should use
    /// [`Self::subscribe_root_changes`] instead of polling this.
    ///
    /// Scoped to the stdio transport: HTTP servers carry one roots
    /// cache per session.
    pub fn client_roots(&self) -> Vec<Root> {
        self.stdio_session
            .get()
            .and_then(|session| session.roots_cache.get())
            .map(|c| c.snapshot())
            .unwrap_or_default()
    }

    /// Subscribe to roots-changed notifications for the stdio
    /// transport. Returns `None` when the cache is not wired (no
    /// outbound primitive) or when the server is HTTP-only. When
    /// `Some`, the receiver fires every time a fresh `roots/list`
    /// response lands; consumers read the current value via
    /// [`tokio::sync::watch::Receiver::borrow`].
    ///
    /// Scoped to the stdio transport: HTTP servers carry one roots
    /// cache per session.
    pub fn subscribe_root_changes(&self) -> Option<tokio::sync::watch::Receiver<Vec<Root>>> {
        self.stdio_session
            .get()
            .and_then(|session| session.roots_cache.get())
            .map(|c| c.subscribe())
    }

    /// Shim for [`McpServerBuilder::add_tools`] + [`McpServerBuilder::build`].
    /// The invoker's catalogue becomes the MCP `tools/list` response;
    /// `tools/call` dispatches via `ToolInvoker::invoke`.
    pub fn expose_tools(invoker: Arc<dyn ToolInvoker>) -> Self {
        Self::builder()
            .add_tools(invoker)
            .build()
            .expect("single-invoker build cannot fail")
    }

    /// Shim for [`McpServerBuilder::add_agent_with_schema`] +
    /// [`McpServerBuilder::build`]. Wraps an [`Agent`] as a single
    /// MCP tool. `tools/list` reports one entry named `agent.name()`
    /// whose `inputSchema` is the caller-supplied `input_schema` (an
    /// arbitrary JSON Schema). `tools/call` decodes the `arguments`
    /// blob into `A::Input`, mints a fresh
    /// [`AgentContext`](klieo_core::agent::AgentContext) via
    /// `ctx_factory`, runs `agent.run(ctx, input).await`, and returns
    /// the JSON-serialised `A::Output`.
    ///
    /// Use this when the agent's `Input` cannot easily derive
    /// `schemars::JsonSchema`. For the auto-derive path enable the
    /// `schemars` cargo feature and use [`Self::expose_agent`].
    ///
    /// For graceful-shutdown — propagating a parent
    /// [`CancellationToken`] into every minted `AgentContext` — open
    /// the builder explicitly:
    /// `McpServer::builder().with_parent_cancel(tok).add_agent_with_schema(..).build()`.
    /// See ADR-010 for the design trade-off.
    pub fn expose_agent_with_schema<A>(
        agent: A,
        input_schema: serde_json::Value,
        ctx_factory: AgentContextFactory,
    ) -> Self
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + Send + 'static,
        A::Output: serde::Serialize + Send + 'static,
    {
        Self::builder()
            .add_agent_with_schema(agent, input_schema, ctx_factory)
            .build()
            .expect("single-invoker build cannot fail")
    }

    /// Shim for [`McpServerBuilder::add_agent`] + [`McpServerBuilder::build`].
    /// Auto-derive variant of [`Self::expose_agent_with_schema`].
    /// Requires `A::Input: schemars::JsonSchema` and the `schemars`
    /// cargo feature on `klieo-mcp-server`.
    #[cfg(feature = "schemars")]
    pub fn expose_agent<A>(agent: A, ctx_factory: AgentContextFactory) -> Self
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
        A::Output: serde::Serialize + Send + 'static,
    {
        Self::builder()
            .add_agent(agent, ctx_factory)
            .build()
            .expect("single-invoker build cannot fail")
    }

    /// Shim for [`McpServerBuilder::with_hitl`] +
    /// [`McpServerBuilder::add_workflow_with_schema`] +
    /// [`McpServerBuilder::build`]. Wraps an [`Agent`] as a single MCP
    /// tool whose `tools/call` drives [`klieo_hitl::run_with_hitl`]
    /// against the supplied HITL client + config. Suspensions surface
    /// as `{"status":"suspended","reason":...}` (no checkpoint / ticket
    /// when no resume KV is wired; ADR-045).
    ///
    /// Ungoverned-build convenience constructor. With the `governor`
    /// feature on, the workflow hard gate requires `with_governor`, so use
    /// the builder path (`with_hitl` + `with_governor` + `add_workflow_*` +
    /// `build`) instead — this shorthand is compiled out there.
    #[cfg(not(feature = "governor"))]
    pub fn expose_workflow_with_schema<A>(
        agent: A,
        system_prompt: impl Into<String>,
        input_schema: serde_json::Value,
        run_options: klieo_core::runtime::RunOptions,
        hitl_client: Arc<klieo_hitl_client::HitlClient>,
        hitl_cfg: Arc<klieo_hitl::HitlConfig>,
        ctx_factory: AgentContextFactory,
    ) -> Result<Self, McpBuildError>
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + Send + 'static,
    {
        Self::builder()
            .with_hitl(hitl_client, hitl_cfg)
            .add_workflow_with_schema(agent, system_prompt, input_schema, run_options, ctx_factory)
            .build()
    }

    /// Auto-derive variant of [`Self::expose_workflow_with_schema`].
    /// Requires `A::Input: schemars::JsonSchema` and the `schemars`
    /// cargo feature. Ungoverned-build only (see the sibling shorthand);
    /// governed builds use the `with_governor` builder path.
    #[cfg(all(feature = "schemars", not(feature = "governor")))]
    pub fn expose_workflow<A>(
        agent: A,
        system_prompt: impl Into<String>,
        run_options: klieo_core::runtime::RunOptions,
        hitl_client: Arc<klieo_hitl_client::HitlClient>,
        hitl_cfg: Arc<klieo_hitl::HitlConfig>,
        ctx_factory: AgentContextFactory,
    ) -> Result<Self, McpBuildError>
    where
        A: Agent + 'static,
        A::Input: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
    {
        Self::builder()
            .with_hitl(hitl_client, hitl_cfg)
            .add_workflow(agent, system_prompt, run_options, ctx_factory)
            .build()
    }

    /// Borrow the merged [`ToolInvoker`] the server dispatches against.
    /// Test helper for asserting per-invocation behaviour (governor
    /// wrap, tenant isolation, error mapping) without standing up the
    /// full HTTP transport. Production callers route through
    /// [`Self::serve_stdio`] / [`Self::serve_http`] instead.
    pub fn invoker(&self) -> &std::sync::Arc<dyn ToolInvoker> {
        &self.invoker
    }

    /// Borrow the configured resume buffer (test helper).
    pub fn resume_buffer(&self) -> &std::sync::Arc<dyn klieo_core::resume::ResumeBuffer> {
        &self.resume_buffer
    }

    /// Borrow the configured pubsub (test helper + cross-replica
    /// fanout consumers).
    pub fn pubsub(&self) -> &std::sync::Arc<dyn klieo_core::Pubsub> {
        &self.pubsub
    }

    /// Borrow the per-server [`klieo_core::CancelRegistry`] keyed by
    /// progressToken. The cancel-subject subscription task and the
    /// `tools/call` dispatch path share a single registry instance:
    /// `tools/call` registers a
    /// [`tokio_util::sync::CancellationToken`] under the progressToken
    /// at invocation start, the subscription task fires it on inbound
    /// `klieo.mcp.cancel.{progressToken}` messages, and `tools/call`
    /// deregisters in a finally-style cleanup once the invocation
    /// terminates.
    pub fn cancel_registry(&self) -> &klieo_core::CancelRegistry<String> {
        &self.cancel_registry
    }

    /// Publish a cross-replica cancel signal for `progress_token` on
    /// `klieo.mcp.cancel.{progress_token}`. Thin wrapper around
    /// [`klieo_core::cancel::publish_cancel_signal`] with explicit
    /// mapping of [`klieo_core::BusError::Invalid`] onto
    /// [`McpServerError::InvalidSubject`] so callers can distinguish
    /// caller-input validation failure from transport failure
    /// without source-chain downcasting.
    ///
    /// # Security
    /// `progress_token` is validated against
    /// [`klieo_core::validate_subject_token`] before subject
    /// construction; metacharacters (`.`, `*`, `>`, whitespace,
    /// non-ASCII) yield [`McpServerError::InvalidSubject`],
    /// preventing a caller-controlled progressToken from collapsing
    /// or wildcarding the subject namespace (CWE-74). Same guard
    /// shape as the per-progressToken event subject published from
    /// the HTTP transport's streaming `tools/call` path.
    ///
    /// Cancel signals share the progressToken-as-credential threat
    /// model documented for resume in ADR-018 / ADR-019: knowledge
    /// of the progressToken grants the ability to cancel the
    /// invocation. Operators MUST mint unguessable progressTokens
    /// (UUID v4 or stronger) and gate per-tenant authorisation
    /// BEFORE the request reaches this server; without that,
    /// cross-tenant cancel becomes possible (CWE-639 IDOR).
    pub async fn publish_cancel(&self, progress_token: &str) -> Result<(), McpServerError> {
        // `?` routes via `From<BusError> for McpServerError`, which
        // maps `BusError::Invalid` → `McpServerError::InvalidSubject`
        // and every other variant → `McpServerError::Bus`. Keep the
        // two-arm dispatch in the `From` impl so this thin wrapper
        // does not duplicate the mapping.
        klieo_core::cancel::publish_cancel_signal(
            &self.pubsub,
            "klieo.mcp.cancel.",
            progress_token,
        )
        .await?;
        Ok(())
    }

    /// Mint a [`ToolCtx`] for the streaming path, overlaying the supplied
    /// progress sender and request-scoped cancel token. Used by the
    /// HTTP transport's `stream_tools_call` so the broadcast is in
    /// place and disconnect-cancel is observable before `invoke` runs.
    #[cfg(feature = "http")]
    pub(crate) fn tool_ctx_with_progress(
        &self,
        progress: tokio::sync::broadcast::Sender<klieo_core::AgentEvent>,
        cancel: tokio_util::sync::CancellationToken,
        caller_principal: Option<String>,
        parent_anchor: Option<String>,
    ) -> klieo_core::tool::ToolCtx {
        let mut ctx = (self.tool_ctx_factory)()
            .with_progress(progress)
            .with_cancel(cancel);
        if let Some(principal) = caller_principal {
            ctx = ctx.with_caller_principal(principal);
        }
        if let Some(anchor) = parent_anchor {
            ctx = ctx.with_parent_anchor(anchor);
        }
        ctx
    }

    /// Drive the stdio loop against the real process stdin/stdout. Thin
    /// wrapper around [`Self::serve_with_streams`] that mints a writer
    /// over [`tokio::io::stdout`] and feeds [`tokio::io::stdin`] as the
    /// reader; all loop semantics live in `serve_with_streams` so the
    /// integration tests drive the same code path over
    /// [`tokio::io::duplex`] pairs.
    ///
    /// Returns when stdin closes (peer disconnect) or on fatal I/O
    /// error.
    pub async fn serve_stdio(self: Arc<Self>) -> Result<(), McpServerError> {
        let stdin = tokio::io::stdin();
        let stdout: outbound::SharedWriter = Arc::new(Mutex::new(tokio::io::stdout()));
        self.serve_with_streams(stdin, stdout).await
    }

    /// Drive the stdio loop against caller-supplied I/O streams. Reads
    /// newline-delimited JSON-RPC frames from `reader`, classifies each
    /// by shape, and either dispatches a response onto `writer` or
    /// routes a server-initiated response into the outbound correlation
    /// table. The same `writer` is primed into the private
    /// `stdout_writer` slot so the outbound primitive shares it —
    /// preventing interleaved writes between inbound replies and
    /// outbound requests on a single underlying stream.
    ///
    /// ## Classification
    ///
    /// - **Request** (`method` + `id`): dispatched via the private
    ///   `handle_jsonrpc` helper; the wire envelope is written onto
    ///   `writer`.
    /// - **Notification** (`method`, no `id`): dispatched for side
    ///   effects only; nothing is written back per JSON-RPC §4.1.
    /// - **Outbound response** (`id` + `result`/`error`, no `method`):
    ///   routed into `outbound::OutboundRequests::complete_pending`
    ///   when the server is wired with outbound support; logged and
    ///   dropped otherwise.
    /// - **Unparseable** (no `method`, no `id`): logged at `warn`.
    pub async fn serve_with_streams<R>(
        self: Arc<Self>,
        reader: R,
        writer: outbound::SharedWriter,
    ) -> Result<(), McpServerError>
    where
        R: tokio::io::AsyncRead + Unpin,
    {
        let stdout = self.ensure_stdout_writer_with(writer).await;
        self.ensure_outbound_and_roots().await;
        let mut lines = BufReader::new(reader).lines();

        while let Some(line) = lines.next_line().await? {
            if line.trim().is_empty() {
                continue;
            }
            self.dispatch_stdio_line(line, stdout.clone());
        }
        Ok(())
    }

    /// Dispatch one inbound frame on a spawned task so the read loop
    /// keeps draining stdin while the handler runs. Spawning is
    /// essential for the sampling path: a `tools/call` whose handler
    /// issues `sampling/createMessage` over `ctx.server_outbound`
    /// blocks until the matching response arrives — and that response
    /// can only arrive once the read loop reads the next frame. Without
    /// spawning, the loop and the handler deadlock.
    fn dispatch_stdio_line(self: &Arc<Self>, line: String, writer: outbound::SharedWriter) {
        let server = self.clone();
        tokio::spawn(async move {
            if let Err(error) = server.process_stdio_line(&line, &writer).await {
                warn!(error = ?error, "stdio dispatch task failed");
            }
        });
    }

    /// Prime [`Self::stdout_writer`] with `writer` iff the cell is
    /// empty, then return a clone of the cell's value.
    /// [`tokio::sync::OnceCell::set`] is no-op-on-populated, so the
    /// first caller's writer wins; subsequent [`Self::serve_with_streams`]
    /// / [`Self::serve_stdio`] calls on the same server reuse the cached
    /// handle so the outbound primitive never observes a writer swap
    /// mid-flight.
    async fn ensure_stdout_writer_with(
        &self,
        writer: outbound::SharedWriter,
    ) -> outbound::SharedWriter {
        // `set` returns `Err` when the cell is already populated — that
        // is the expected fast path on every subsequent invocation. We
        // ignore the result and read back via `get().expect(..)` because
        // we've just guaranteed the cell is non-empty.
        let _ = self.stdout_writer.set(writer);
        self.stdout_writer
            .get()
            .expect("stdout_writer populated above")
            .clone()
    }

    /// Construct the outbound correlation primitive + roots cache when
    /// the builder opted in via [`McpServerBuilder::with_client_sampling`].
    /// No-op when the flag is unset (tools-only deployments pay zero cost)
    /// or when the cells are already populated (subsequent
    /// [`Self::serve_stdio`] calls share the same primitives).
    async fn ensure_outbound_and_roots(&self) {
        if !self.declare_sampling {
            return;
        }
        // Cloning the cached handle keeps inbound reply writes and
        // outbound request writes serialised on the same `Mutex`.
        let writer = self
            .stdout_writer
            .get()
            .expect("serve_with_streams primes stdout_writer before ensure_outbound_and_roots")
            .clone();
        let session = self
            .stdio_session
            .get_or_init(|| async { std::sync::Arc::new(crate::session::Session::new_stdio()) })
            .await
            .clone();
        let outbound = session
            .outbound
            .get_or_init(|| async {
                let sink: Arc<dyn OutboundFrameSink> =
                    Arc::new(crate::outbound_sink::StdioFrameSink::new(writer.clone()));
                Arc::new(crate::outbound::OutboundRequests::new(sink))
            })
            .await
            .clone();
        let _ = session
            .roots_cache
            .get_or_init(|| async {
                let outbound: Arc<dyn klieo_core::ServerOutbound> = outbound.clone();
                Arc::new(crate::roots::RootsCache::new(outbound))
            })
            .await;
    }

    /// Classify one inbound JSON-RPC frame and dispatch to the request,
    /// notification, or outbound-response path.
    async fn process_stdio_line(
        &self,
        line: &str,
        writer: &outbound::SharedWriter,
    ) -> Result<(), McpServerError> {
        let parsed: serde_json::Value = match serde_json::from_str(line) {
            Ok(value) => value,
            Err(error) => {
                // Sanitise: log the serde error, return a stable message so attacker bytes never echo back.
                warn!(error = %error, "rejected malformed JSON-RPC frame");
                let envelope = rpc_error(None, JSONRPC_PARSE_ERROR, "malformed JSON-RPC frame");
                return write_frame(writer, &envelope).await;
            }
        };
        let stdio_session = self.stdio_session.get();
        match classify_inbound(&parsed) {
            InboundKind::Request => {
                let envelope = self.handle_jsonrpc(parsed, stdio_session).await;
                write_frame(writer, &envelope).await
            }
            InboundKind::Notification => {
                // JSON-RPC §4.1 — notifications carry no response.
                self.handle_jsonrpc(parsed, stdio_session).await;
                Ok(())
            }
            InboundKind::OutboundResponse(id) => {
                self.route_outbound_response(id, parsed).await;
                Ok(())
            }
            InboundKind::Unparseable => {
                warn!("rejected inbound frame: no method and no id");
                Ok(())
            }
        }
    }

    /// Route a server-initiated response into the stdio session's
    /// correlation table. When the table is absent (the builder did
    /// not opt in to a transport that carries reverse-direction
    /// JSON-RPC) the frame is logged at `warn` and dropped — the peer
    /// is misbehaving or the server is misconfigured, but neither
    /// warrants tearing down the stdio loop.
    async fn route_outbound_response(&self, id: i64, frame: serde_json::Value) {
        if let Some(outbound) = self
            .stdio_session
            .get()
            .and_then(|session| session.outbound.get())
        {
            outbound.complete_pending(id, frame).await;
        } else {
            warn!(
                rpc_id = id,
                "outbound response received but server has no outbound table wired"
            );
        }
    }

    /// Parse + dispatch a single newline-delimited JSON-RPC frame.
    /// Retained from the pre-T2 surface for unit tests that feed a
    /// request line straight into the dispatcher and expect the wire
    /// envelope back. Notifications + outbound responses go through
    /// [`Self::process_stdio_line`] instead, which classifies the
    /// frame by shape before dispatch.
    #[cfg(test)]
    async fn handle_line(&self, line: &str) -> serde_json::Value {
        let req: serde_json::Value = match serde_json::from_str(line) {
            Ok(v) => v,
            Err(e) => {
                warn!(error = %e, "rejected malformed JSON-RPC frame");
                return rpc_error(None, JSONRPC_PARSE_ERROR, "malformed JSON-RPC frame");
            }
        };
        self.handle_jsonrpc(req, self.stdio_session.get()).await
    }

    /// Dispatch one already-parsed JSON-RPC request value to the
    /// matching handler. Returns a fully-formed JSON-RPC response
    /// envelope (`result` on success, `error` on failure). Shared
    /// by all transports.
    ///
    /// `session` carries the per-transport [`crate::session::Session`]
    /// that owns the outbound primitive + roots cache for this
    /// dispatch. Stdio callers pass `self.stdio_session.get()`; HTTP
    /// callers pass the session minted by `handle_initialize_post` or
    /// resolved by `require_session`. Notification arms that drive
    /// per-session side effects (roots seed + refresh) read from this
    /// parameter; method arms that touch only catalogue state ignore
    /// it. `None` degrades the notification arms to a silent no-op.
    pub(crate) async fn handle_jsonrpc(
        &self,
        req: serde_json::Value,
        session: Option<&std::sync::Arc<crate::session::Session>>,
    ) -> serde_json::Value {
        let id = req.get("id").cloned();
        let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");

        match method {
            "initialize" => rpc_ok(id, self.handle_initialize(&req).await),
            "notifications/initialized" => {
                self.handle_initialized_notification(session).await;
                serde_json::Value::Null
            }
            "notifications/roots/list_changed" => {
                self.handle_roots_list_changed_notification(session);
                serde_json::Value::Null
            }
            "shutdown" => rpc_ok(id, serde_json::Value::Null),
            "tools/list" => rpc_ok(id, self.tools_list()),
            "tools/call" => match self.tools_call(req.get("params")).await {
                Ok(v) => rpc_ok(id, v),
                Err(e) => tool_error_to_envelope(id, e),
            },
            other => {
                warn!(rpc_id = ?id, method = other, "method not found");
                rpc_error(
                    id,
                    JSONRPC_METHOD_NOT_FOUND,
                    &format!("method not found: {other}"),
                )
            }
        }
    }

    /// Parse the client-declared capabilities from the `initialize`
    /// request, store the flags this server cares about under
    /// [`Self::client_caps`], and return the result payload for the
    /// outgoing JSON-RPC response. Splits the "absorb caps + mint
    /// response" pair out of [`Self::handle_jsonrpc`] so the dispatch
    /// arm stays a single-statement readable line.
    async fn handle_initialize(&self, req: &serde_json::Value) -> serde_json::Value {
        let roots_supported = req.pointer("/params/capabilities/roots").is_some();
        {
            let mut caps = self.client_caps.lock().await;
            caps.roots_supported = roots_supported;
        }
        if self.declare_sampling {
            initialize_result_with_sampling()
        } else {
            initialize_result_without_sampling()
        }
    }

    /// Handle `notifications/initialized`. When the client advertised
    /// `capabilities.roots` on `initialize` AND the dispatch session
    /// carries a wired outbound primitive (i.e. the builder opted in
    /// via [`McpServerBuilder::with_client_sampling`]), spawn a
    /// one-shot `roots/list` fetch to seed the cache. The fetch runs
    /// on a detached task so the notification dispatch does not block
    /// on peer round-trip latency.
    ///
    /// `session` resolves to the per-transport session that owns the
    /// cache; `None` (no session in dispatch context) or a missing
    /// cache both degrade silently.
    async fn handle_initialized_notification(
        &self,
        session: Option<&std::sync::Arc<crate::session::Session>>,
    ) {
        let roots_supported = self.client_caps.lock().await.roots_supported;
        if !roots_supported {
            return;
        }
        let Some(cache) = session
            .and_then(|session| session.roots_cache.get())
            .cloned()
        else {
            return;
        };
        tokio::spawn(async move {
            if let Err(error) = cache.refresh().await {
                warn!(error = ?error, "initial roots/list fetch failed");
            }
        });
    }

    /// Handle `notifications/roots/list_changed`. The client signals
    /// that its declared roots have changed; the server reacts by
    /// re-issuing `roots/list` to refresh the cached snapshot. The
    /// refresh runs on a detached task so the notification dispatch
    /// does not block on peer round-trip latency.
    ///
    /// `session` resolves to the per-transport session that owns the
    /// cache. Acts as a no-op when `None` (no session in dispatch
    /// context) or when the cache is not wired — tools-only
    /// deployments that never opted in to the outbound primitive
    /// have nothing to refresh, and a misbehaving peer sending the
    /// notification anyway should not panic the server.
    fn handle_roots_list_changed_notification(
        &self,
        session: Option<&std::sync::Arc<crate::session::Session>>,
    ) {
        let Some(cache) = session
            .and_then(|session| session.roots_cache.get())
            .cloned()
        else {
            return;
        };
        tokio::spawn(async move {
            if let Err(error) = cache.refresh().await {
                warn!(error = ?error, "roots list_changed re-fetch failed");
            }
        });
    }

    fn tools_list(&self) -> serde_json::Value {
        let tools: Vec<serde_json::Value> = self
            .invoker
            .catalogue()
            .iter()
            .map(tool_def_to_mcp_descriptor)
            .collect();
        serde_json::json!({ "tools": tools })
    }

    async fn tools_call(
        &self,
        params: Option<&serde_json::Value>,
    ) -> Result<serde_json::Value, ToolError> {
        let params = params.ok_or_else(|| ToolError::InvalidArgs("missing params".into()))?;
        let name = params
            .get("name")
            .and_then(|n| n.as_str())
            .ok_or_else(|| ToolError::InvalidArgs("missing tool name".into()))?;
        let args = params
            .get("arguments")
            .cloned()
            .unwrap_or(serde_json::Value::Null);
        let ctx = (self.tool_ctx_factory)().with_cancel(self.parent_cancel.child_token());
        let out = self.invoker.invoke(name, args, ctx).await?;
        Ok(serde_json::json!({
            "content": [
                { "type": "text", "text": out.to_string() }
            ]
        }))
    }
}

/// Internal adapter: presents one [`Agent`] as a single-tool
/// [`ToolInvoker`] so the MCP dispatch loop need not learn about
/// agents. The wrapper mints an [`AgentContext`] per request via
/// `ctx_factory` and derives a child cancel token from the incoming
/// [`ToolCtx`], propagating both server-level and per-request
/// cancellation into the agent run.
struct AgentAsToolInvoker<A>
where
    A: Agent + 'static,
    A::Input: serde::de::DeserializeOwned + Send + 'static,
    A::Output: serde::Serialize + Send + 'static,
{
    agent: Arc<A>,
    name: String,
    input_schema: serde_json::Value,
    ctx_factory: AgentContextFactory,
    /// When the builder was wired with a governor, every agent run
    /// wraps `ctx.llm` with a [`GovernedLlmClient`] keyed off the
    /// tenant label. `None` keeps the legacy
    /// ungoverned path so adopters that govern out-of-band stay
    /// unaffected.
    #[cfg(feature = "governor")]
    governor: Option<crate::governor::GovernorBundle>,
}

#[async_trait]
impl<A> ToolInvoker for AgentAsToolInvoker<A>
where
    A: Agent + 'static,
    A::Input: serde::de::DeserializeOwned + Send + 'static,
    A::Output: serde::Serialize + Send + 'static,
{
    fn catalogue(&self) -> Vec<ToolDef> {
        vec![ToolDef::new(
            self.name.clone(),
            format!("klieo agent: {}", self.name),
            self.input_schema.clone(),
        )]
    }

    async fn invoke(
        &self,
        name: &str,
        args: serde_json::Value,
        tool_ctx: ToolCtx,
    ) -> Result<serde_json::Value, ToolError> {
        if name != self.name {
            return Err(ToolError::UnknownTool(name.into()));
        }
        // Sanitise wire-bound error messages: keep stable strings on
        // the response payload; log the full third-party Display
        // server-side. Mirrors the parse-error sanitisation already
        // applied in `handle_line` (CWE-209: no internal error
        // text — type names, file paths, response bodies — over the
        // transport boundary).
        let input: A::Input = serde_json::from_value(args).map_err(|e| {
            warn!(agent = %self.name, error = %e, "decode of MCP tools/call args failed");
            ToolError::InvalidArgs("arguments do not match inputSchema".into())
        })?;
        let mut ctx = (self.ctx_factory)();
        // Child token so both server-level cancel (carried into `tool_ctx`
        // by `tools_call`) and per-request cancel (carried in by the HTTP
        // streaming path's `DropGuard`) reach the agent without the agent
        // needing to know which origin triggered the disconnect.
        ctx.cancel = tool_ctx.cancel.child_token();
        ctx.progress = tool_ctx.progress.clone();
        // The verified raw principal stays server-side; hash it into a
        // non-PII label so the audit trail records attribution via
        // `Episode::RunAttributed` without admitting the principal into
        // agent memory or any LLM-visible surface (ADR-045).
        if let Some(principal) = tool_ctx.caller_principal.as_ref() {
            ctx = ctx.with_tenant_label(klieo_core::principal_hash(principal.as_str()));
        }
        // The HTTP boundary threads the cross-hop anchor only for
        // authenticated callers, so a stamped `Episode::RunOrigin` is
        // always co-attributable to the `RunAttributed` principal above.
        // Recorded verbatim; never reaches an LLM-visible surface.
        if let Some(anchor) = tool_ctx.parent_anchor.as_ref() {
            ctx = ctx.with_parent_anchor(anchor.as_str().to_string());
        }
        // Agents are governed only when a bundle was wired at
        // registration; ungoverned agents keep the legacy path (no hard
        // gate at `add_agent` time, unlike workflows).
        #[cfg(feature = "governor")]
        if let Some(bundle) = self.governor.as_ref() {
            ctx = crate::governor::wrap_ctx_with_governor(ctx, bundle);
        }
        let output = self.agent.run(ctx, input).await.map_err(|e| {
            warn!(agent = %self.name, error = %e, "exposed agent execution failed");
            ToolError::Permanent("agent execution failed".into())
        })?;
        serde_json::to_value(output).map_err(|e| {
            warn!(agent = %self.name, error = %e, "encode of agent output failed");
            ToolError::Permanent("agent output not serialisable".into())
        })
    }
}

/// Internal multi-invoker: merges the catalogues of multiple
/// [`ToolInvoker`]s and routes each `invoke(name, ..)` to the inner
/// invoker that claims the named tool.
///
/// Construction walks every inner catalogue once and builds two
/// memoised structures:
/// - `routes: HashMap<tool_name, inner_index>` for O(1) dispatch.
/// - `merged_catalogue: Vec<ToolDef>` cached so `catalogue()` does
///   not re-walk the inner invokers per `tools/list` request.
///
/// Tool-name collisions are rejected at construction time —
/// [`Self::new`] panics on the first duplicate so routing stays
/// unambiguous (silent first-match-wins is fragile when two agents
/// share a name).
struct MergedInvoker {
    inner: Vec<Arc<dyn ToolInvoker>>,
    routes: std::collections::HashMap<String, usize>,
    merged_catalogue: Vec<ToolDef>,
}

impl MergedInvoker {
    fn new(inner: Vec<Arc<dyn ToolInvoker>>) -> Result<Self, McpBuildError> {
        let mut routes: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
        let mut merged_catalogue: Vec<ToolDef> = Vec::new();
        for (index, invoker) in inner.iter().enumerate() {
            for tool in invoker.catalogue() {
                if routes.insert(tool.name.clone(), index).is_some() {
                    return Err(McpBuildError::DuplicateTool(tool.name));
                }
                merged_catalogue.push(tool);
            }
        }
        Ok(Self {
            inner,
            routes,
            merged_catalogue,
        })
    }
}

#[async_trait]
impl ToolInvoker for MergedInvoker {
    fn catalogue(&self) -> Vec<ToolDef> {
        self.merged_catalogue.clone()
    }

    async fn invoke(
        &self,
        name: &str,
        args: serde_json::Value,
        ctx: ToolCtx,
    ) -> Result<serde_json::Value, ToolError> {
        match self.routes.get(name) {
            Some(&index) => self.inner[index].invoke(name, args, ctx).await,
            None => Err(ToolError::UnknownTool(name.into())),
        }
    }

    /// Forward to the inner invoker that owns `name` so per-tool
    /// idempotency overrides survive the merge. Tools whose owning
    /// invoker is not registered fall through to the default `false`.
    /// Cluster-0.24 follower re-invoke reads this on the orphan
    /// resume path; see [`crate::http::handle_dead_leader_orphan_mcp`].
    fn is_tool_idempotent(&self, name: &str) -> bool {
        match self.routes.get(name) {
            Some(&index) => self.inner[index].is_tool_idempotent(name),
            None => false,
        }
    }

    /// Forward to the inner invoker that owns `name` so a PII-flagged
    /// tool's audit-redaction requirement survives the merge. Without
    /// this the decorator's default `false` would make dispatch record
    /// raw PII for any tool reached through a merged catalogue.
    fn tool_redacts_audit(&self, name: &str) -> bool {
        match self.routes.get(name) {
            Some(&index) => self.inner[index].tool_redacts_audit(name),
            None => false,
        }
    }
}

/// Map a [`ToolError`] to a JSON-RPC error envelope. Applies the
/// standard CWE-209 redaction policy shared by both the JSON path
/// (in `handle_jsonrpc`) and the SSE streaming path.
///
/// - `UnknownTool(name)` — safe to echo; surfaces tool name verbatim.
/// - `InvalidArgs(reason)` — sanitised by the invoker layer; echo as-is.
/// - All other variants — log full error server-side, surface only the
///   stable `"tool invocation failed"` string.
pub(crate) fn tool_error_to_envelope(
    id: Option<serde_json::Value>,
    e: ToolError,
) -> serde_json::Value {
    let stable_msg = match &e {
        ToolError::UnknownTool(name) => {
            warn!(rpc_id = ?id, tool = %name, "tools/call: unknown tool");
            format!("unknown tool: {name}")
        }
        ToolError::InvalidArgs(reason) => {
            warn!(rpc_id = ?id, reason = %reason, "tools/call: invalid args");
            reason.clone()
        }
        _ => {
            warn!(rpc_id = ?id, error = %e, "tools/call failed");
            "tool invocation failed".into()
        }
    };
    rpc_error(id, JSONRPC_SERVER_ERROR, &stable_msg)
}

fn tool_def_to_mcp_descriptor(def: &ToolDef) -> serde_json::Value {
    serde_json::json!({
        "name": def.name,
        "description": def.description,
        "inputSchema": def.json_schema,
    })
}

fn initialize_result_with_sampling() -> serde_json::Value {
    initialize_result_inner(true)
}

fn initialize_result_without_sampling() -> serde_json::Value {
    initialize_result_inner(false)
}

/// Build the JSON payload of an `initialize` response.
///
/// `with_sampling` gates the `capabilities.sampling = {}` field so
/// MCP-conformant clients listen for `sampling/createMessage` outbound
/// requests only when the server actually wires the outbound primitive.
fn initialize_result_inner(with_sampling: bool) -> serde_json::Value {
    let mut capabilities = serde_json::json!({ "tools": {} });
    if with_sampling {
        capabilities["sampling"] = serde_json::json!({});
    }
    serde_json::json!({
        "protocolVersion": MCP_PROTOCOL_VERSION,
        "capabilities": capabilities,
        "serverInfo": { "name": "klieo-mcp-server", "version": env!("CARGO_PKG_VERSION") }
    })
}

pub(crate) fn rpc_ok(id: Option<serde_json::Value>, result: serde_json::Value) -> serde_json::Value {
    serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": result })
}

/// Shape of an inbound JSON-RPC frame as classified by
/// [`classify_inbound`]. The stdio read loop maps each variant to a
/// distinct dispatch path (request vs notification vs server-initiated
/// response vs reject).
#[derive(Debug)]
enum InboundKind {
    /// `method` + `id`: peer-initiated request — dispatch + write the
    /// response.
    Request,
    /// `method`, no `id`: peer-initiated notification — dispatch for
    /// side effects only; no reply per JSON-RPC §4.1.
    Notification,
    /// `id` + (`result` or `error`), no `method`: response to a
    /// previous server-initiated outbound request.
    OutboundResponse(i64),
    /// Neither `method` nor a usable `id` payload — peer is
    /// misbehaving; log + drop.
    Unparseable,
}

/// Classify an already-parsed JSON-RPC frame by shape. Pure function:
/// no side effects, fully unit-testable.
fn classify_inbound(value: &serde_json::Value) -> InboundKind {
    let has_method = value.get("method").is_some();
    let id = value.get("id");
    if has_method {
        return if id.is_some() {
            InboundKind::Request
        } else {
            InboundKind::Notification
        };
    }
    let has_payload = value.get("result").is_some() || value.get("error").is_some();
    match (id.and_then(serde_json::Value::as_i64), has_payload) {
        (Some(id), true) => InboundKind::OutboundResponse(id),
        _ => InboundKind::Unparseable,
    }
}

/// Serialise + flush a single JSON-RPC envelope onto the shared
/// transport writer. Holds the writer mutex across write + flush so
/// two frames cannot interleave on a single underlying file
/// descriptor (stdout in production, in-memory duplex in tests).
async fn write_frame(
    writer: &outbound::SharedWriter,
    envelope: &serde_json::Value,
) -> Result<(), McpServerError> {
    let bytes = serde_json::to_vec(envelope)?;
    let mut guard = writer.lock().await;
    guard.write_all(&bytes).await?;
    guard.write_all(b"\n").await?;
    guard.flush().await?;
    Ok(())
}

pub(crate) fn rpc_error(
    id: Option<serde_json::Value>,
    code: i64,
    message: &str,
) -> serde_json::Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": { "code": code, "message": message }
    })
}

/// Construct a default-shape `ToolCtx` for use in downstream
/// integration tests. Hidden from rustdoc — not part of the
/// public API contract.
#[doc(hidden)]
pub fn __test_noop_ctx() -> klieo_core::tool::ToolCtx {
    noop_ctx()
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use klieo_core::tool::Tool;
    use std::sync::OnceLock;

    struct EmptyInvoker;

    #[async_trait]
    impl ToolInvoker for EmptyInvoker {
        fn catalogue(&self) -> Vec<ToolDef> {
            Vec::new()
        }
        async fn invoke(
            &self,
            name: &str,
            _args: serde_json::Value,
            _ctx: ToolCtx,
        ) -> Result<serde_json::Value, ToolError> {
            Err(ToolError::UnknownTool(name.into()))
        }
    }

    struct Echo;

    #[async_trait]
    impl Tool for Echo {
        fn name(&self) -> &str {
            "echo"
        }
        fn description(&self) -> &str {
            "echoes back its args"
        }
        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)
        }
    }

    struct OneToolInvoker;

    #[async_trait]
    impl ToolInvoker for OneToolInvoker {
        fn catalogue(&self) -> Vec<ToolDef> {
            vec![ToolDef::new(
                "echo",
                "echoes back its args",
                serde_json::json!({"type": "object"}),
            )]
        }
        async fn invoke(
            &self,
            name: &str,
            args: serde_json::Value,
            ctx: ToolCtx,
        ) -> Result<serde_json::Value, ToolError> {
            if name == "echo" {
                Echo.invoke(args, ctx).await
            } else {
                Err(ToolError::UnknownTool(name.into()))
            }
        }
    }

    #[tokio::test]
    async fn initialize_returns_server_info() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server
            .handle_line(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#)
            .await;
        let info = resp["result"]["serverInfo"]["name"].as_str().unwrap();
        assert_eq!(info, "klieo-mcp-server");
    }

    #[tokio::test]
    async fn tools_list_surfaces_invoker_catalogue() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server
            .handle_line(r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#)
            .await;
        let tools = resp["result"]["tools"].as_array().unwrap();
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0]["name"], "echo");
    }

    #[tokio::test]
    async fn tools_call_dispatches_to_invoker() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server
            .handle_line(
                r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"echo","arguments":{"hello":"world"}}}"#,
            )
            .await;
        let text = resp["result"]["content"][0]["text"].as_str().unwrap();
        assert!(text.contains("hello"));
        assert!(text.contains("world"));
    }

    #[tokio::test]
    async fn unknown_method_returns_method_not_found() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server
            .handle_line(r#"{"jsonrpc":"2.0","id":4,"method":"nope"}"#)
            .await;
        assert_eq!(resp["error"]["code"], JSONRPC_METHOD_NOT_FOUND);
    }

    #[tokio::test]
    async fn tools_call_without_params_returns_server_error() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server
            .handle_line(r#"{"jsonrpc":"2.0","id":5,"method":"tools/call"}"#)
            .await;
        assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
        assert!(resp["error"]["message"]
            .as_str()
            .unwrap()
            .contains("missing params"));
    }

    #[tokio::test]
    async fn tools_call_unknown_tool_surfaces_invoker_error() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server
            .handle_line(
                r#"{"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"name":"does-not-exist","arguments":{}}}"#,
            )
            .await;
        assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
        assert!(resp["error"]["message"]
            .as_str()
            .unwrap()
            .contains("does-not-exist"));
    }

    #[tokio::test]
    async fn malformed_frame_returns_sanitised_parse_error() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let resp = server.handle_line("not json").await;
        assert_eq!(resp["error"]["code"], JSONRPC_PARSE_ERROR);
        // The peer must NOT see the underlying serde_json error text;
        // only the stable sanitised message.
        let msg = resp["error"]["message"].as_str().unwrap();
        assert_eq!(msg, "malformed JSON-RPC frame");
    }

    #[tokio::test]
    async fn handle_jsonrpc_dispatches_initialize() {
        let server = McpServer::builder()
            .add_tools(Arc::new(EmptyInvoker))
            .build()
            .unwrap();
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {}
        });
        let resp = server.handle_jsonrpc(req, None).await;
        assert_eq!(resp["jsonrpc"], "2.0");
        assert_eq!(resp["id"], 1);
        assert!(resp["result"].is_object());
    }

    #[tokio::test]
    async fn handle_jsonrpc_returns_method_not_found_for_unknown() {
        let server = McpServer::builder()
            .add_tools(Arc::new(EmptyInvoker))
            .build()
            .unwrap();
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 7,
            "method": "no_such_method"
        });
        let resp = server.handle_jsonrpc(req, None).await;
        assert_eq!(resp["error"]["code"], JSONRPC_METHOD_NOT_FOUND);
        assert_eq!(resp["id"], 7);
    }

    #[tokio::test]
    async fn classify_inbound_recognises_request_shape() {
        let frame = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/list"
        });
        assert!(matches!(classify_inbound(&frame), InboundKind::Request));
    }

    #[tokio::test]
    async fn classify_inbound_recognises_notification_shape() {
        let frame = serde_json::json!({
            "jsonrpc": "2.0",
            "method": "notifications/initialized"
        });
        assert!(matches!(
            classify_inbound(&frame),
            InboundKind::Notification
        ));
    }

    #[tokio::test]
    async fn classify_inbound_recognises_outbound_result_shape() {
        let frame = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 42,
            "result": {"role": "assistant"}
        });
        match classify_inbound(&frame) {
            InboundKind::OutboundResponse(id) => assert_eq!(id, 42),
            other => panic!("expected OutboundResponse(42), got {other:?}"),
        }
    }

    #[tokio::test]
    async fn classify_inbound_recognises_outbound_error_shape() {
        let frame = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 7,
            "error": {"code": -32601, "message": "Method not found"}
        });
        match classify_inbound(&frame) {
            InboundKind::OutboundResponse(id) => assert_eq!(id, 7),
            other => panic!("expected OutboundResponse(7), got {other:?}"),
        }
    }

    #[tokio::test]
    async fn classify_inbound_rejects_no_method_no_id() {
        let frame = serde_json::json!({"jsonrpc": "2.0"});
        assert!(matches!(classify_inbound(&frame), InboundKind::Unparseable));
    }

    #[tokio::test]
    async fn classify_inbound_rejects_bare_id_without_payload() {
        // `id` present, but neither `result` nor `error` and no
        // `method`. JSON-RPC does not define this shape; treat as
        // unparseable rather than guess the peer's intent.
        let frame = serde_json::json!({"jsonrpc": "2.0", "id": 9});
        assert!(matches!(classify_inbound(&frame), InboundKind::Unparseable));
    }

    /// Shared in-memory byte buffer used by [`BufferSink`]. Wrapped in
    /// `std::sync::Mutex` so `AsyncWrite::poll_write` can append
    /// synchronously without yielding to the runtime.
    type CapturedBytes = std::sync::Arc<std::sync::Mutex<Vec<u8>>>;

    /// Mint a [`outbound::SharedWriter`] backed by an in-memory buffer
    /// and return the buffer handle so the test can read back what
    /// the loop wrote.
    fn duplex_writer() -> (outbound::SharedWriter, CapturedBytes) {
        let buffer: CapturedBytes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
        let shared: outbound::SharedWriter = Arc::new(Mutex::new(BufferSink(buffer.clone())));
        (shared, buffer)
    }

    /// In-memory [`tokio::io::AsyncWrite`] sink that appends every
    /// byte into a shared `Vec<u8>` so the test can read back what the
    /// loop wrote out.
    struct BufferSink(CapturedBytes);

    impl tokio::io::AsyncWrite for BufferSink {
        fn poll_write(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
            buf: &[u8],
        ) -> std::task::Poll<std::io::Result<usize>> {
            self.0
                .lock()
                .expect("BufferSink mutex poisoned in test")
                .extend_from_slice(buf);
            std::task::Poll::Ready(Ok(buf.len()))
        }

        fn poll_flush(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            std::task::Poll::Ready(Ok(()))
        }

        fn poll_shutdown(
            self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            std::task::Poll::Ready(Ok(()))
        }
    }

    fn captured_bytes(buffer: &CapturedBytes) -> Vec<u8> {
        buffer
            .lock()
            .expect("captured-bytes mutex poisoned in test")
            .clone()
    }

    #[tokio::test]
    async fn process_stdio_line_writes_response_for_request() {
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let (writer, buffer) = duplex_writer();
        let request = r#"{"jsonrpc":"2.0","id":11,"method":"tools/list"}"#;
        server
            .process_stdio_line(request, &writer)
            .await
            .expect("stdio dispatch must not fail");
        let bytes = captured_bytes(&buffer);
        assert!(bytes.ends_with(b"\n"), "frames are newline-delimited");
        let envelope: serde_json::Value =
            serde_json::from_slice(bytes.trim_ascii_end()).expect("written frame must be JSON");
        assert_eq!(envelope["id"], 11);
        assert!(envelope["result"]["tools"].is_array());
    }

    #[tokio::test]
    async fn process_stdio_line_drops_outbound_response_when_table_absent() {
        // Server built without the future `with_client_sampling` opt-in
        // has `outbound = None`. An incoming response shape must be
        // logged and dropped, NOT written back to the peer (that would
        // echo the same id back and confuse the client).
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        assert!(
            server
                .stdio_session
                .get()
                .and_then(|s| s.outbound.get())
                .is_none(),
            "default-built server must not wire an outbound table"
        );
        let (writer, buffer) = duplex_writer();
        let stray = r#"{"jsonrpc":"2.0","id":99,"result":{"role":"assistant"}}"#;
        server
            .process_stdio_line(stray, &writer)
            .await
            .expect("stray response must not break the loop");
        assert!(
            captured_bytes(&buffer).is_empty(),
            "stray outbound responses must never produce wire output"
        );
    }

    #[tokio::test]
    async fn process_stdio_line_drops_notification_without_writing() {
        // Per JSON-RPC §4.1, a notification (method + no id) must not
        // produce a wire response — even when `handle_jsonrpc`'s
        // method-not-found path generates an envelope internally.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let (writer, buffer) = duplex_writer();
        let notification = r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#;
        server
            .process_stdio_line(notification, &writer)
            .await
            .expect("notification dispatch must not fail");
        assert!(
            captured_bytes(&buffer).is_empty(),
            "notifications must not produce wire output"
        );
    }

    #[tokio::test]
    async fn process_stdio_line_drops_unparseable_frame() {
        // Frame missing both `method` and a usable `id` payload is
        // unparseable per `classify_inbound`. Must log + drop without
        // writing.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let (writer, buffer) = duplex_writer();
        let unparseable = r#"{"jsonrpc":"2.0"}"#;
        server
            .process_stdio_line(unparseable, &writer)
            .await
            .expect("unparseable frame must not break the loop");
        assert!(
            captured_bytes(&buffer).is_empty(),
            "unparseable frames must not produce wire output"
        );
    }

    #[tokio::test]
    async fn process_stdio_line_writes_parse_error_for_malformed_json() {
        // Malformed bytes still produce a wire-visible parse error so
        // the peer can react. The detailed serde error stays server-
        // side; only the sanitised message reaches the wire.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let (writer, buffer) = duplex_writer();
        server
            .process_stdio_line("not json", &writer)
            .await
            .expect("parse-error path must not fail the loop");
        let bytes = captured_bytes(&buffer);
        let envelope: serde_json::Value =
            serde_json::from_slice(bytes.trim_ascii_end()).expect("parse-error envelope is JSON");
        assert_eq!(envelope["error"]["code"], JSONRPC_PARSE_ERROR);
        assert_eq!(envelope["error"]["message"], "malformed JSON-RPC frame");
    }

    #[tokio::test]
    async fn initialize_arm_records_roots_capability_when_advertised() {
        // Client advertised `capabilities.roots = {}` → server records
        // `roots_supported = true` so the subsequent
        // `notifications/initialized` arm can drive the seed fetch.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 400,
            "method": "initialize",
            "params": { "capabilities": { "roots": {} } }
        });
        server.handle_jsonrpc(req, None).await;
        assert!(
            server.client_caps.lock().await.roots_supported,
            "initialize must record advertised roots capability"
        );
    }

    #[tokio::test]
    async fn initialize_arm_defaults_roots_unsupported_when_absent() {
        // No `capabilities.roots` on the initialize payload → server
        // must NOT spawn the seed fetch on `notifications/initialized`.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 401,
            "method": "initialize",
            "params": { "capabilities": {} }
        });
        server.handle_jsonrpc(req, None).await;
        assert!(
            !server.client_caps.lock().await.roots_supported,
            "initialize must leave roots_supported=false when absent"
        );
    }

    #[tokio::test]
    async fn initialize_result_includes_sampling_when_flag_set() {
        let payload = super::initialize_result_with_sampling();
        assert!(
            payload["capabilities"]["sampling"].is_object(),
            "initialize_result_with_sampling must surface capabilities.sampling; got: {payload}"
        );
    }

    #[tokio::test]
    async fn initialize_result_omits_sampling_when_flag_unset() {
        let payload = super::initialize_result_without_sampling();
        assert!(
            payload["capabilities"].get("sampling").is_none(),
            "initialize_result_without_sampling must omit capabilities.sampling; got: {payload}"
        );
    }

    #[tokio::test]
    async fn initialized_notification_returns_null_value() {
        // Per JSON-RPC §4.1 the dispatcher returns `Value::Null` for a
        // notification so the stdio loop's classifier drops the frame
        // without writing a reply. Asserts the response shape directly;
        // the spawned roots/list fetch (when roots_supported) is covered
        // by T9's integration suite.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "method": "notifications/initialized"
        });
        let resp = server.handle_jsonrpc(req, None).await;
        assert!(
            resp.is_null(),
            "notifications/initialized must yield a Null sentinel; got: {resp}"
        );
    }

    #[tokio::test]
    async fn list_changed_notification_returns_null_value() {
        // `notifications/roots/list_changed` is a JSON-RPC notification —
        // the dispatcher returns `Value::Null` so the stdio loop's
        // classifier discards the frame without writing a reply.
        // Cache-refresh side effect is asserted end-to-end by T9.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "method": "notifications/roots/list_changed"
        });
        let resp = server.handle_jsonrpc(req, None).await;
        assert!(
            resp.is_null(),
            "notifications/roots/list_changed must yield a Null sentinel; got: {resp}"
        );
    }

    #[tokio::test]
    async fn list_changed_when_cache_absent_is_noop() {
        // Default-built server (no `with_client_sampling`) never wires
        // the outbound primitive or the roots cache. The dispatch arm
        // must degrade silently — no panic, still returns the Null
        // sentinel so the stdio classifier drops the frame.
        let server = McpServer::expose_tools(Arc::new(OneToolInvoker));
        assert!(
            server
                .stdio_session
                .get()
                .and_then(|s| s.roots_cache.get())
                .is_none(),
            "default-built server must not wire a roots cache"
        );
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "method": "notifications/roots/list_changed"
        });
        let resp = server.handle_jsonrpc(req, None).await;
        assert!(
            resp.is_null(),
            "cache-absent list_changed must still yield Null; got: {resp}"
        );
    }

    #[cfg(feature = "http")]
    #[tokio::test]
    async fn tool_ctx_with_progress_threads_cancel() {
        use std::sync::Arc;
        use tokio_util::sync::CancellationToken;
        let server = Arc::new(
            McpServer::builder()
                .add_tools(Arc::new(EmptyInvoker))
                .build()
                .unwrap(),
        );
        let (tx, _rx) = tokio::sync::broadcast::channel::<klieo_core::AgentEvent>(8);
        let token = CancellationToken::new();
        let ctx = server.tool_ctx_with_progress(tx, token.clone(), None, None);
        token.cancel();
        assert!(ctx.cancel.is_cancelled());
    }

    mod expose_agent_tests {
        use super::*;
        use async_trait::async_trait;
        use klieo_core::agent::{Agent, AgentContext};
        use klieo_core::error::Error as KlieoError;
        use klieo_core::llm::ToolDef;
        use klieo_core::test_utils::fake_context;
        use serde::{Deserialize, Serialize};

        #[derive(Debug, Clone, Deserialize, Serialize)]
        struct GreetIn {
            who: String,
        }

        #[derive(Debug, Clone, Serialize)]
        struct GreetOut {
            greeting: String,
        }

        struct Greeter;

        #[async_trait]
        impl Agent for Greeter {
            type Input = GreetIn;
            type Output = GreetOut;
            type Error = KlieoError;

            fn name(&self) -> &str {
                "greeter"
            }
            fn system_prompt(&self) -> &str {
                ""
            }
            fn tools(&self) -> &[ToolDef] {
                &[]
            }
            async fn run(
                &self,
                _ctx: AgentContext,
                input: GreetIn,
            ) -> Result<GreetOut, KlieoError> {
                Ok(GreetOut {
                    greeting: format!("hello {}", input.who),
                })
            }
        }

        // Used by both the propagation and default-token tests so
        // each cancel scenario refers to the same Agent shape.

        #[derive(Debug, Clone, Serialize)]
        struct CancelObserveOut {
            state: String,
        }

        struct CancelObserver;

        #[async_trait]
        impl Agent for CancelObserver {
            type Input = serde_json::Value;
            type Output = CancelObserveOut;
            type Error = KlieoError;
            fn name(&self) -> &str {
                "cancel-observer"
            }
            fn system_prompt(&self) -> &str {
                ""
            }
            fn tools(&self) -> &[ToolDef] {
                &[]
            }
            async fn run(
                &self,
                ctx: AgentContext,
                _input: serde_json::Value,
            ) -> Result<CancelObserveOut, KlieoError> {
                let state = if ctx.cancel.is_cancelled() {
                    "cancelled".into()
                } else {
                    "ran".into()
                };
                Ok(CancelObserveOut { state })
            }
        }

        fn fresh_ctx() -> AgentContext {
            fake_context("greeter")
        }

        fn one_object_schema() -> serde_json::Value {
            serde_json::json!({
                "type": "object",
                "properties": {"who": {"type": "string"}},
                "required": ["who"]
            })
        }

        #[tokio::test]
        async fn expose_agent_with_schema_lists_agent_as_single_tool() {
            let server = McpServer::expose_agent_with_schema(
                Greeter,
                one_object_schema(),
                Arc::new(fresh_ctx),
            );
            let resp = server
                .handle_line(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#)
                .await;
            let tools = resp["result"]["tools"].as_array().unwrap();
            assert_eq!(tools.len(), 1);
            assert_eq!(tools[0]["name"], "greeter");
            assert_eq!(tools[0]["inputSchema"]["type"], "object");
        }

        #[tokio::test]
        async fn expose_agent_with_schema_dispatches_tools_call_through_agent() {
            let server = McpServer::expose_agent_with_schema(
                Greeter,
                one_object_schema(),
                Arc::new(fresh_ctx),
            );
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"greeter","arguments":{"who":"world"}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains(r#""greeting":"hello world""#),
                "tools/call must return serialised agent output; got: {text}"
            );
        }

        #[tokio::test]
        async fn expose_agent_with_schema_rejects_unknown_tool_name() {
            let server = McpServer::expose_agent_with_schema(
                Greeter,
                one_object_schema(),
                Arc::new(fresh_ctx),
            );
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"not-greeter","arguments":{}}}"#,
                )
                .await;
            assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
            assert!(resp["error"]["message"]
                .as_str()
                .unwrap()
                .contains("not-greeter"));
        }

        #[tokio::test]
        async fn expose_agent_with_schema_rejects_malformed_args() {
            let server = McpServer::expose_agent_with_schema(
                Greeter,
                one_object_schema(),
                Arc::new(fresh_ctx),
            );
            // `who` field missing — A::Input decode fails.
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"greeter","arguments":{}}}"#,
                )
                .await;
            assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
            let msg = resp["error"]["message"].as_str().unwrap();
            assert!(
                msg.contains("arguments do not match inputSchema"),
                "wire message must be the sanitised string; got: {msg}"
            );
            // The serde decode internals (struct names, paths) MUST NOT
            // surface to the peer — verify they are NOT present.
            assert!(
                !msg.contains("GreetIn") && !msg.contains("missing field"),
                "internal decode detail must not leak: {msg}"
            );
        }

        /// Agent execution errors are sanitised over the wire — peer
        /// sees the stable `"agent execution failed"` message and the
        /// full chain logs server-side via tracing::warn! (CWE-209).
        /// Asserts both directions: stable string present + internal
        /// secret-bearing content absent.
        #[tokio::test]
        async fn expose_agent_sanitises_run_error_on_wire() {
            struct Failing;
            #[async_trait]
            impl Agent for Failing {
                type Input = serde_json::Value;
                type Output = serde_json::Value;
                type Error = KlieoError;
                fn name(&self) -> &str {
                    "failing"
                }
                fn system_prompt(&self) -> &str {
                    ""
                }
                fn tools(&self) -> &[ToolDef] {
                    &[]
                }
                async fn run(
                    &self,
                    _ctx: AgentContext,
                    _input: serde_json::Value,
                ) -> Result<serde_json::Value, KlieoError> {
                    Err(KlieoError::BadResponse(
                        "internal: token=secret-abc upstream=https://provider/url".into(),
                    ))
                }
            }
            let server = McpServer::expose_agent_with_schema(
                Failing,
                serde_json::json!({}),
                Arc::new(fresh_ctx),
            );
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":99,"method":"tools/call","params":{"name":"failing","arguments":{}}}"#,
                )
                .await;
            assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
            let msg = resp["error"]["message"].as_str().unwrap();
            assert!(
                msg.contains("tool invocation failed"),
                "wire message must contain the sanitised stable string; got: {msg}"
            );
            // The agent's internal error payload (secret-like content,
            // URLs) MUST NOT reach the peer.
            assert!(
                !msg.contains("secret-abc") && !msg.contains("https://"),
                "internal error detail must not leak: {msg}"
            );
        }

        /// Cancellation propagation via the builder: cancelling the
        /// parent token mid-flight cancels every in-flight agent's
        /// `ctx.cancel`. Agent observes `ctx.cancel.is_cancelled()`
        /// and reports `"ran"` or `"cancelled"` in its Output so the
        /// test can distinguish before and after `parent.cancel()`.
        #[tokio::test]
        async fn builder_propagates_parent_cancel_into_ctx() {
            let parent = CancellationToken::new();
            let server = McpServer::builder()
                .with_parent_cancel(parent.clone())
                .add_agent_with_schema(CancelObserver, serde_json::json!({}), Arc::new(fresh_ctx))
                .build()
                .unwrap();

            // Uncancelled parent → agent sees a live cancel token.
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":200,"method":"tools/call","params":{"name":"cancel-observer","arguments":{}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains(r#""state":"ran""#),
                "live parent token must produce live ctx.cancel; got: {text}"
            );

            // Cancel the parent; new invocations must observe a
            // cancelled token via the child-token override.
            parent.cancel();
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":201,"method":"tools/call","params":{"name":"cancel-observer","arguments":{}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains(r#""state":"cancelled""#),
                "cancelled parent must propagate into ctx.cancel via child_token; got: {text}"
            );
        }

        /// Streaming path: `tool_ctx_with_progress(tx, token)` threads
        /// the request-scoped token into the `ToolCtx`; `AgentAsToolInvoker`
        /// derives a child and places it on the minted `AgentContext`. The
        /// `CancelObserver` agent reports its `ctx.cancel` state so the
        /// test can assert end-to-end propagation without needing SSE wiring.
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn tool_ctx_with_progress_cancel_cascades_into_agent_context() {
            let request_cancel = CancellationToken::new();
            let server = Arc::new(
                McpServer::builder()
                    .add_agent_with_schema(
                        CancelObserver,
                        serde_json::json!({}),
                        Arc::new(fresh_ctx),
                    )
                    .build()
                    .unwrap(),
            );
            let (tx, _rx) = tokio::sync::broadcast::channel::<klieo_core::AgentEvent>(8);
            request_cancel.cancel();
            let tool_ctx = server.tool_ctx_with_progress(tx, request_cancel, None, None);
            let result = server
                .invoker
                .invoke("cancel-observer", serde_json::json!({}), tool_ctx)
                .await
                .unwrap();
            let text = result.to_string();
            assert!(
                text.contains(r#""state":"cancelled""#),
                "cancelled request token must cascade into AgentContext.cancel; got: {text}"
            );
        }

        /// The shorthand `expose_agent_with_schema` ctor — which
        /// delegates to the builder without ever calling
        /// `with_parent_cancel` — yields a fresh, never-cancelled
        /// default token. Agents see `ctx.cancel.is_cancelled() ==
        /// false`. Guards against accidental wiring of an external
        /// token into the default path.
        #[tokio::test]
        async fn shim_ctor_uses_default_uncancelled_parent_token() {
            let server = McpServer::expose_agent_with_schema(
                CancelObserver,
                serde_json::json!({}),
                Arc::new(fresh_ctx),
            );
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":202,"method":"tools/call","params":{"name":"cancel-observer","arguments":{}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains(r#""state":"ran""#),
                "shim ctor must default to a never-cancelled parent token; got: {text}"
            );
        }

        /// An `AgentAsToolInvoker` receiving a
        /// `ToolCtx` with `caller_principal=Some(p)` mints an
        /// `AgentContext` whose `tenant_label=principal_hash(p)`. The
        /// run loop then records exactly one `Episode::RunAttributed`
        /// carrying that label, and the raw principal never reaches
        /// episodic memory or short-term memory.
        #[tokio::test]
        async fn agent_as_tool_invoker_installs_tenant_label_from_caller_principal() {
            use klieo_core::test_utils::{noop_bus, FakeLlmClient, FakeLlmStep};
            const PRINCIPAL: &str = "alice@example.com";

            struct EchoLoopAgent;

            #[async_trait]
            impl Agent for EchoLoopAgent {
                type Input = serde_json::Value;
                type Output = serde_json::Value;
                type Error = KlieoError;
                fn name(&self) -> &str {
                    "echo-loop"
                }
                fn system_prompt(&self) -> &str {
                    ""
                }
                fn tools(&self) -> &[ToolDef] {
                    &[]
                }
                async fn run(
                    &self,
                    ctx: AgentContext,
                    _input: serde_json::Value,
                ) -> Result<serde_json::Value, KlieoError> {
                    let out = klieo_core::runtime::run_steps(
                        &ctx,
                        "",
                        klieo_core::ids::ThreadId::new("echo-loop-thread"),
                        klieo_core::runtime::RunOptions::default(),
                    )
                    .await?;
                    Ok(serde_json::Value::String(out))
                }
            }

            let mut ctx_seed = fake_context("echo-loop");
            ctx_seed.llm = Arc::new(
                FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]),
            );
            let episodic_for_probe = ctx_seed.episodic.clone();
            let short_term_for_probe = ctx_seed.short_term.clone();
            let run_id_for_probe = ctx_seed.run_id;

            let slot = Arc::new(std::sync::Mutex::new(Some(ctx_seed)));
            let ctx_factory: AgentContextFactory = Arc::new(move || {
                slot.lock()
                    .unwrap()
                    .take()
                    .expect("ctx_factory called more than once")
            });
            let server = McpServer::builder()
                .add_agent_with_schema(
                    EchoLoopAgent,
                    serde_json::json!({"type": "object"}),
                    ctx_factory,
                )
                .build()
                .unwrap();
            let (pubsub, _, kv, jobs) = noop_bus();
            let tool_ctx = klieo_core::tool::ToolCtx::new(pubsub, kv, jobs)
                .with_caller_principal(PRINCIPAL.into());
            let _ = server
                .invoker
                .invoke(
                    "echo-loop",
                    serde_json::json!({}),
                    tool_ctx,
                )
                .await
                .unwrap();

            let expected = klieo_core::principal_hash(PRINCIPAL);
            let episodes = episodic_for_probe.replay(run_id_for_probe).await.unwrap();
            let labels: Vec<&str> = episodes
                .iter()
                .filter_map(|e| match e {
                    klieo_core::Episode::RunAttributed { tenant_label } => {
                        Some(tenant_label.as_str())
                    }
                    _ => None,
                })
                .collect();
            assert_eq!(
                labels,
                vec![expected.as_str()],
                "exactly one RunAttributed carrying principal_hash; got {episodes:?}",
            );
            for ep in &episodes {
                let payload = serde_json::to_string(ep).unwrap();
                assert!(
                    !payload.contains(PRINCIPAL),
                    "raw principal leaked into recorded episode: {payload}",
                );
            }
            let history = short_term_for_probe
                .load(klieo_core::ids::ThreadId::new("echo-loop-thread"), 8192)
                .await
                .unwrap_or_default();
            for msg in &history {
                assert!(
                    !msg.content.contains(PRINCIPAL),
                    "principal leaked into short-term memory: {}",
                    msg.content
                );
            }
        }

        /// An `AgentAsToolInvoker` receiving a `ToolCtx`
        /// with `parent_anchor=Some(a)` records exactly one
        /// `Episode::RunOrigin` carrying the anchor **verbatim** (not
        /// hashed/rewritten), co-emitted with the `RunAttributed`
        /// attribution for the authenticated caller, and never admits the
        /// anchor into short-term (LLM-visible) memory.
        #[tokio::test]
        async fn agent_as_tool_invoker_records_run_origin_from_parent_anchor() {
            use klieo_core::test_utils::{noop_bus, FakeLlmClient, FakeLlmStep};
            const PRINCIPAL: &str = "alice@example.com";
            const ANCHOR: &str = "sha256:deadbeefcafe0123";

            struct EchoLoopAgent;

            #[async_trait]
            impl Agent for EchoLoopAgent {
                type Input = serde_json::Value;
                type Output = serde_json::Value;
                type Error = KlieoError;
                fn name(&self) -> &str {
                    "echo-origin"
                }
                fn system_prompt(&self) -> &str {
                    ""
                }
                fn tools(&self) -> &[ToolDef] {
                    &[]
                }
                async fn run(
                    &self,
                    ctx: AgentContext,
                    _input: serde_json::Value,
                ) -> Result<serde_json::Value, KlieoError> {
                    let out = klieo_core::runtime::run_steps(
                        &ctx,
                        "",
                        klieo_core::ids::ThreadId::new("echo-origin-thread"),
                        klieo_core::runtime::RunOptions::default(),
                    )
                    .await?;
                    Ok(serde_json::Value::String(out))
                }
            }

            let mut ctx_seed = fake_context("echo-origin");
            ctx_seed.llm = Arc::new(
                FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]),
            );
            let episodic_for_probe = ctx_seed.episodic.clone();
            let short_term_for_probe = ctx_seed.short_term.clone();
            let run_id_for_probe = ctx_seed.run_id;

            let slot = Arc::new(std::sync::Mutex::new(Some(ctx_seed)));
            let ctx_factory: AgentContextFactory = Arc::new(move || {
                slot.lock()
                    .unwrap()
                    .take()
                    .expect("ctx_factory called more than once")
            });
            let server = McpServer::builder()
                .add_agent_with_schema(
                    EchoLoopAgent,
                    serde_json::json!({"type": "object"}),
                    ctx_factory,
                )
                .build()
                .unwrap();
            let (pubsub, _, kv, jobs) = noop_bus();
            let tool_ctx = klieo_core::tool::ToolCtx::new(pubsub, kv, jobs)
                .with_caller_principal(PRINCIPAL.into())
                .with_parent_anchor(ANCHOR.into());
            let _ = server
                .invoker
                .invoke("echo-origin", serde_json::json!({}), tool_ctx)
                .await
                .unwrap();

            let episodes = episodic_for_probe.replay(run_id_for_probe).await.unwrap();
            let anchors: Vec<&str> = episodes
                .iter()
                .filter_map(|e| match e {
                    klieo_core::Episode::RunOrigin { parent_anchor } => {
                        Some(parent_anchor.as_str())
                    }
                    _ => None,
                })
                .collect();
            assert_eq!(
                anchors,
                vec![ANCHOR],
                "exactly one RunOrigin carrying the verbatim anchor; got {episodes:?}",
            );
            // Co-attribution: RunOrigin never stands alone — the
            // authenticated principal is recorded in the same run so the
            // unverified parent claim is attributable.
            let attributed = episodes
                .iter()
                .filter(|e| matches!(e, klieo_core::Episode::RunAttributed { .. }))
                .count();
            assert_eq!(
                attributed, 1,
                "RunOrigin co-emitted with exactly one RunAttributed; got {episodes:?}",
            );
            let history = short_term_for_probe
                .load(klieo_core::ids::ThreadId::new("echo-origin-thread"), 8192)
                .await
                .unwrap_or_default();
            for msg in &history {
                assert!(
                    !msg.content.contains(ANCHOR),
                    "anchor leaked into short-term memory: {}",
                    msg.content
                );
            }
        }

        /// `build()` returns `NoInvokers` when called with zero
        /// invokers — the server has nothing to serve. Documents
        /// the typed-error guard.
        #[tokio::test]
        async fn builder_with_no_invokers_returns_no_invokers_error() {
            let result = McpServer::builder().build();
            assert!(matches!(result, Err(McpBuildError::NoInvokers)));
        }

        /// Opting into client sampling via the builder threads the
        /// flag onto the built `McpServer`. The capability-negotiation
        /// arm reads it to gate `capabilities.sampling = {}` on the
        /// initialize response.
        #[tokio::test]
        async fn with_client_sampling_sets_capability_flag() {
            let server = McpServer::builder()
                .with_client_sampling()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .build()
                .unwrap();
            assert!(
                server.declare_sampling,
                "with_client_sampling() must set declare_sampling=true on the built server"
            );
        }

        /// A builder that never calls `with_client_sampling()` produces
        /// a server with the flag unset. Initialize must NOT advertise
        /// `capabilities.sampling` for tools-only deployments.
        #[tokio::test]
        async fn default_builder_does_not_declare_sampling() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .build()
                .unwrap();
            assert!(
                !server.declare_sampling,
                "default builder must leave declare_sampling=false"
            );
        }

        /// A builder that never calls `with_session_idle_timeout`
        /// produces a server whose idle deadline matches
        /// `DEFAULT_SESSION_IDLE_TIMEOUT` (5 minutes per ADR-028).
        /// Also pins that the per-session SSE tx + activity clock
        /// fields are initialised to their pre-session sentinels.
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn default_session_idle_timeout_is_5min() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .build()
                .unwrap();
            assert_eq!(
                server.session_idle_timeout,
                std::time::Duration::from_secs(300),
                "default session idle timeout must be 5 minutes"
            );
            assert!(
                server.sessions.read().await.is_empty(),
                "HTTP server must hold zero sessions before any initialize POST"
            );
        }

        /// `with_session_idle_timeout(ttl)` overrides the default
        /// idle deadline on the built server. Pins the builder
        /// method threads its argument through `build_inner`.
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn with_session_idle_timeout_overrides_default() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .with_session_idle_timeout(std::time::Duration::from_secs(42))
                .build()
                .unwrap();
            assert_eq!(
                server.session_idle_timeout,
                std::time::Duration::from_secs(42),
                "with_session_idle_timeout must override the default"
            );
        }

        /// `with_session_idle_timeout(Duration::ZERO)` records the
        /// disabled-watchdog sentinel. The watchdog task spawned on
        /// `initialize` reads this and returns immediately when zero.
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn zero_duration_records_disabled_watchdog() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .with_session_idle_timeout(std::time::Duration::ZERO)
                .build()
                .unwrap();
            assert_eq!(
                server.session_idle_timeout,
                std::time::Duration::ZERO,
                "Duration::ZERO must thread through to record disabled-watchdog intent"
            );
        }

        /// `DEFAULT_MAX_SESSIONS` pins the concurrent-HTTP-session cap
        /// at 1024. Operators tuning the knob downward are calling
        /// `with_max_sessions` against this baseline; the value must
        /// not drift without a CHANGELOG entry.
        #[cfg(feature = "http")]
        #[test]
        fn default_max_sessions_is_1024() {
            assert_eq!(DEFAULT_MAX_SESSIONS, 1024);
        }

        /// `with_max_sessions(cap)` overrides the
        /// [`DEFAULT_MAX_SESSIONS`] baseline on the built server.
        /// Pins that the builder method threads its argument through
        /// `build_inner` into the `max_sessions` field read by
        /// `handle_initialize_post`.
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn with_max_sessions_overrides_default() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .with_max_sessions(64)
                .build()
                .unwrap();
            assert_eq!(
                server.max_sessions, 64,
                "with_max_sessions must override the default cap"
            );
        }

        /// `with_max_sessions(0)` panics — a zero cap would deadlock
        /// the `initialize` POST path with no recovery, so the
        /// builder rejects it eagerly rather than producing a
        /// permanently 503ing server.
        #[cfg(feature = "http")]
        #[test]
        #[should_panic(expected = "max_sessions must be > 0")]
        fn with_max_sessions_panics_on_zero() {
            let _ = McpServer::builder().with_max_sessions(0);
        }

        /// `DEFAULT_MAX_SESSIONS_PER_PRINCIPAL_DIVISOR` pins the
        /// divisor used by [`default_max_sessions_per_principal`] at
        /// 16. Operators see the default sub-cap as
        /// `max_sessions / 16`; drifting the constant changes the
        /// implicit cap on every untouched deployment.
        #[cfg(feature = "http")]
        #[test]
        fn default_divisor_is_sixteen() {
            assert_eq!(DEFAULT_MAX_SESSIONS_PER_PRINCIPAL_DIVISOR, 16);
        }

        /// `default_max_sessions_per_principal` derives the sub-cap
        /// by dividing `max_sessions` by
        /// [`DEFAULT_MAX_SESSIONS_PER_PRINCIPAL_DIVISOR`]. Pins the
        /// production-baseline pair (1024 → 64) plus a smaller
        /// configuration (32 → 2) to lock in the linear scaling.
        #[cfg(feature = "http")]
        #[test]
        fn default_per_principal_derives_from_max_sessions() {
            assert_eq!(default_max_sessions_per_principal(1024), 64);
            assert_eq!(default_max_sessions_per_principal(32), 2);
        }

        /// `default_max_sessions_per_principal` floors at 1 when
        /// integer division would yield 0. Guarantees the default
        /// sub-cap is always reachable so authenticated `initialize`
        /// POSTs never face a permanent 503 under small
        /// `max_sessions` values.
        #[cfg(feature = "http")]
        #[test]
        fn default_per_principal_floors_at_one() {
            assert_eq!(default_max_sessions_per_principal(0), 1);
            assert_eq!(default_max_sessions_per_principal(15), 1);
            assert_eq!(default_max_sessions_per_principal(16), 1);
        }

        /// `with_max_sessions_per_principal(cap)` overrides the
        /// default sub-cap on the built server. Pins that the builder
        /// method threads its argument through `build_inner` into the
        /// `max_sessions_per_principal` field read by the admission
        /// branch on `handle_initialize_post`.
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn with_max_sessions_per_principal_overrides_default() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .with_max_sessions(1024)
                .with_max_sessions_per_principal(8)
                .build()
                .unwrap();
            assert_eq!(
                server.max_sessions_per_principal, 8,
                "with_max_sessions_per_principal must override the default sub-cap"
            );
        }

        /// `with_max_sessions_per_principal(0)` panics — a zero
        /// sub-cap rejects every authenticated `initialize` with no
        /// recovery, so the builder rejects it eagerly rather than
        /// producing a permanently 503ing server for any
        /// authenticated principal.
        #[cfg(feature = "http")]
        #[test]
        #[should_panic(expected = "max_sessions_per_principal must be > 0")]
        fn with_max_sessions_per_principal_panics_on_zero() {
            let _ = McpServer::builder().with_max_sessions_per_principal(0);
        }

        /// `DEFAULT_SSE_REPLAY_CAPACITY` pins the production-baseline
        /// per-session SSE replay buffer capacity at 256 frames.
        /// Operators observing the implicit ring size depend on this
        /// number; a drift would silently change memory pressure on
        /// every untouched deployment.
        #[cfg(feature = "http")]
        #[test]
        fn default_sse_replay_capacity_is_256() {
            assert_eq!(DEFAULT_SSE_REPLAY_CAPACITY, 256);
        }

        /// `with_sse_replay_capacity(capacity)` overrides the default
        /// on the built server. Pins that the builder method threads
        /// its argument through `build_inner` into the
        /// `sse_replay_capacity` field, and that `sse_replay_enabled()`
        /// returns `true` for positive capacities and `false` for 0
        /// (the disable knob).
        #[cfg(feature = "http")]
        #[tokio::test]
        async fn with_sse_replay_capacity_overrides_default() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .with_sse_replay_capacity(8)
                .build()
                .unwrap();
            assert_eq!(server.sse_replay_capacity, 8);
            assert!(server.sse_replay_enabled());

            let off = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .with_sse_replay_capacity(0)
                .build()
                .unwrap();
            assert_eq!(off.sse_replay_capacity, 0);
            assert!(!off.sse_replay_enabled());
        }

        /// Two agents registered via repeated `add_agent_with_schema`
        /// both show up in `tools/list` and each `tools/call` routes
        /// to the correct agent through `MergedInvoker`.
        #[tokio::test]
        async fn builder_supports_multi_agent_dispatch() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .add_agent_with_schema(CancelObserver, serde_json::json!({}), Arc::new(fresh_ctx))
                .build()
                .unwrap();

            // tools/list surfaces both agents.
            let resp = server
                .handle_line(r#"{"jsonrpc":"2.0","id":300,"method":"tools/list"}"#)
                .await;
            let tools = resp["result"]["tools"].as_array().unwrap();
            let names: Vec<&str> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect();
            assert_eq!(tools.len(), 2);
            assert!(names.contains(&"greeter"));
            assert!(names.contains(&"cancel-observer"));

            // tools/call to greeter dispatches to Greeter.
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":301,"method":"tools/call","params":{"name":"greeter","arguments":{"who":"multi"}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(text.contains(r#""greeting":"hello multi""#));

            // tools/call to cancel-observer dispatches to CancelObserver.
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":302,"method":"tools/call","params":{"name":"cancel-observer","arguments":{}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(text.contains(r#""state":"ran""#));
        }

        /// Parent-cancel token registered ONCE on the builder
        /// propagates into every agent the builder produces. Pinning
        /// this prevents a future refactor from giving each
        /// `add_agent_*` call its own copy that doesn't share the
        /// cancellation chain.
        #[tokio::test]
        async fn builder_parent_cancel_propagates_into_every_agent() {
            let parent = CancellationToken::new();
            let server = McpServer::builder()
                .with_parent_cancel(parent.clone())
                .add_agent_with_schema(CancelObserver, serde_json::json!({}), Arc::new(fresh_ctx))
                .add_tools(Arc::new(super::OneToolInvoker))
                .build()
                .unwrap();

            parent.cancel();
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":303,"method":"tools/call","params":{"name":"cancel-observer","arguments":{}}}"#,
                )
                .await;
            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains(r#""state":"cancelled""#),
                "builder-level parent_cancel must reach every add_agent_* invoker; got: {text}"
            );
        }

        /// On a multi-invoker server, `tools/call` for a name that
        /// matches no inner catalogue maps to `ToolError::UnknownTool`,
        /// surfaced as a JSON-RPC server error containing the
        /// unknown name. Covers `MergedInvoker::invoke`'s `None`
        /// branch which the dispatch tests do not exercise.
        #[tokio::test]
        async fn builder_multi_agent_unknown_tool_returns_error() {
            let server = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .add_agent_with_schema(CancelObserver, serde_json::json!({}), Arc::new(fresh_ctx))
                .build()
                .unwrap();

            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":304,"method":"tools/call","params":{"name":"no-such-agent","arguments":{}}}"#,
                )
                .await;
            assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
            assert!(
                resp["error"]["message"]
                    .as_str()
                    .unwrap()
                    .contains("no-such-agent"),
                "UnknownTool error must reference the requested name"
            );
        }

        /// Two agents that report the same `name()` are a caller bug.
        /// `build` with ≥2 invokers returns `DuplicateTool` on the first
        /// duplicate so routing stays unambiguous.
        #[tokio::test]
        async fn builder_build_returns_duplicate_tool_error() {
            let result = McpServer::builder()
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .add_agent_with_schema(Greeter, one_object_schema(), Arc::new(fresh_ctx))
                .build();
            let Err(McpBuildError::DuplicateTool(ref name)) = result else {
                panic!("expected DuplicateTool error");
            };
            assert!(
                !name.is_empty(),
                "DuplicateTool must carry the colliding tool name"
            );
            assert_eq!(
                name, "greeter",
                "DuplicateTool must name the colliding tool"
            );
        }

        /// `MergedInvoker::tool_redacts_audit` forwards to the inner
        /// invoker that owns the tool, so a PII-flagged tool keeps its
        /// audit-redaction requirement after the merge — without the
        /// forwarding, the decorator default `false` would make dispatch
        /// record raw PII for any tool reached through a merged catalogue.
        /// Covers the routed-true, routed-false, and unrouted-None arms.
        #[test]
        fn merged_invoker_forwards_tool_redacts_audit_per_owner() {
            use klieo_core::test_utils::FakeToolInvoker;

            let pii_owner: Arc<dyn ToolInvoker> = Arc::new(
                FakeToolInvoker::new().with_redacting_tool("claimant_lookup", "handles PII", Ok),
            );
            let plain_owner: Arc<dyn ToolInvoker> =
                Arc::new(FakeToolInvoker::new().with_tool("echo", "plain", Ok));

            let merged = MergedInvoker::new(vec![pii_owner, plain_owner])
                .expect("distinct tool names must merge without DuplicateTool");

            assert!(
                merged.tool_redacts_audit("claimant_lookup"),
                "a PII-flagged tool's redaction must survive the merge; \
                 default-false here would record raw PII (fail-open)"
            );
            assert!(
                !merged.tool_redacts_audit("echo"),
                "an unflagged tool must not be reported as redacting"
            );
            assert!(
                !merged.tool_redacts_audit("no-such-tool"),
                "an unrouted name hits the None arm and must default to false"
            );
        }

        /// `with_cancel_subscription` + `build()` (not `build_arc()`) must
        /// return `CancelRequiresArc` — the unwrapped server shape cannot
        /// keep a background subscriber task alive.
        #[test]
        fn mcp_builder_build_returns_cancel_requires_arc_error() {
            let result = McpServer::builder()
                .add_tools(Arc::new(OneToolInvoker))
                .with_cancel_subscription()
                .build();
            assert!(
                matches!(result, Err(McpBuildError::CancelRequiresArc)),
                "build() with cancel subscription must return CancelRequiresArc",
            );
        }

        /// Output-encode failure is sanitised on the wire the same way
        /// as run-error and decode-error. An Agent whose Output type's
        /// Serialize impl returns Err must NOT leak the inner serde
        /// error Display (which can embed arbitrary internal state
        /// chosen by the impl author) over the JSON-RPC boundary.
        #[tokio::test]
        async fn expose_agent_sanitises_encode_error_on_wire() {
            struct NonSerialisable;

            impl serde::Serialize for NonSerialisable {
                fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
                where
                    S: serde::Serializer,
                {
                    Err(serde::ser::Error::custom(
                        "internal: token=secret-encode-abc upstream=https://provider/encode",
                    ))
                }
            }

            struct EncodeFailing;
            #[async_trait]
            impl Agent for EncodeFailing {
                type Input = serde_json::Value;
                type Output = NonSerialisable;
                type Error = KlieoError;
                fn name(&self) -> &str {
                    "encode-failing"
                }
                fn system_prompt(&self) -> &str {
                    ""
                }
                fn tools(&self) -> &[ToolDef] {
                    &[]
                }
                async fn run(
                    &self,
                    _ctx: AgentContext,
                    _input: serde_json::Value,
                ) -> Result<NonSerialisable, KlieoError> {
                    Ok(NonSerialisable)
                }
            }

            let server = McpServer::expose_agent_with_schema(
                EncodeFailing,
                serde_json::json!({}),
                Arc::new(fresh_ctx),
            );
            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":100,"method":"tools/call","params":{"name":"encode-failing","arguments":{}}}"#,
                )
                .await;
            assert_eq!(resp["error"]["code"], JSONRPC_SERVER_ERROR);
            let msg = resp["error"]["message"].as_str().unwrap();
            assert!(
                msg.contains("tool invocation failed"),
                "wire message must contain the sanitised stable string; got: {msg}"
            );
            assert!(
                !msg.contains("secret-encode-abc") && !msg.contains("https://"),
                "internal encode-error detail must not leak: {msg}"
            );
        }

        #[tokio::test]
        async fn tool_ctx_factory_invoked_per_request() {
            use std::sync::atomic::{AtomicUsize, Ordering};
            let counter = Arc::new(AtomicUsize::new(0));
            let c2 = counter.clone();
            let factory: ToolCtxFactory = Arc::new(move || {
                c2.fetch_add(1, Ordering::SeqCst);
                default_tool_ctx_factory()()
            });

            let server = McpServer::builder()
                .with_tool_ctx_factory(factory)
                .add_tools(Arc::new(super::OneToolInvoker))
                .build()
                .unwrap();

            let req = serde_json::json!({
                "jsonrpc": "2.0", "id": 1, "method": "tools/call",
                "params": { "name": "echo", "arguments": {"x": 1} }
            });
            server.handle_jsonrpc(req.clone(), None).await;
            server.handle_jsonrpc(req, None).await;

            assert_eq!(counter.load(Ordering::SeqCst), 2);
        }

        /// `AgentAsToolInvoker::invoke` overlays `ToolCtx::progress`
        /// onto the `AgentContext` it mints via `ctx_factory`. A
        /// transport that injects a `broadcast::Sender<AgentEvent>`
        /// into the `ToolCtx` must see it arrive inside `Agent::run`.
        #[tokio::test]
        async fn agent_as_tool_invoker_propagates_progress_to_agent_context() {
            use klieo_core::AgentEvent;
            use std::sync::Mutex;
            use tokio::sync::broadcast;

            struct CapturingAgent {
                captured: Arc<Mutex<Option<Option<broadcast::Sender<AgentEvent>>>>>,
            }

            #[async_trait]
            impl Agent for CapturingAgent {
                type Input = serde_json::Value;
                type Output = serde_json::Value;
                type Error = KlieoError;

                fn name(&self) -> &str {
                    "capturing"
                }
                fn system_prompt(&self) -> &str {
                    ""
                }
                fn tools(&self) -> &[ToolDef] {
                    &[]
                }
                async fn run(
                    &self,
                    ctx: AgentContext,
                    _input: serde_json::Value,
                ) -> Result<serde_json::Value, KlieoError> {
                    *self.captured.lock().unwrap() = Some(ctx.progress.clone());
                    Ok(serde_json::json!({}))
                }
            }

            let captured = Arc::new(Mutex::new(None::<Option<broadcast::Sender<AgentEvent>>>));
            let agent = CapturingAgent {
                captured: captured.clone(),
            };

            let (tx, _rx) = broadcast::channel::<AgentEvent>(16);
            let tx_for_factory = tx.clone();
            let factory: ToolCtxFactory = Arc::new(move || {
                let bus = klieo_bus_memory::MemoryBus::new();
                klieo_core::tool::ToolCtx::new(bus.pubsub, bus.kv, bus.jobs)
                    .with_progress(tx_for_factory.clone())
            });

            let server = McpServer::builder()
                .with_tool_ctx_factory(factory)
                .add_agent_with_schema(agent, serde_json::json!({}), Arc::new(fresh_ctx))
                .build()
                .unwrap();

            let req = serde_json::json!({
                "jsonrpc": "2.0", "id": 1, "method": "tools/call",
                "params": { "name": "capturing", "arguments": {} }
            });
            let resp = server.handle_jsonrpc(req, None).await;
            assert!(
                resp["result"].is_object(),
                "tools/call must succeed; got: {resp}"
            );

            let captured_progress = captured
                .lock()
                .unwrap()
                .clone()
                .expect("Agent::run was never invoked");
            assert!(
                captured_progress.is_some(),
                "AgentContext.progress was None despite ToolCtx.progress=Some"
            );
        }

        #[cfg(feature = "schemars")]
        mod auto_derive {
            use super::*;
            use schemars::JsonSchema;

            #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)]
            struct DerivedIn {
                who: String,
            }

            #[derive(Debug, Clone, Serialize)]
            struct DerivedOut {
                greeting: String,
            }

            struct DerivedGreeter;

            #[async_trait]
            impl Agent for DerivedGreeter {
                type Input = DerivedIn;
                type Output = DerivedOut;
                type Error = KlieoError;

                fn name(&self) -> &str {
                    "derived-greeter"
                }
                fn system_prompt(&self) -> &str {
                    ""
                }
                fn tools(&self) -> &[ToolDef] {
                    &[]
                }
                async fn run(
                    &self,
                    _ctx: AgentContext,
                    input: DerivedIn,
                ) -> Result<DerivedOut, KlieoError> {
                    Ok(DerivedOut {
                        greeting: format!("hi {}", input.who),
                    })
                }
            }

            #[tokio::test]
            async fn expose_agent_auto_derives_schema_via_schemars() {
                let server = McpServer::expose_agent(DerivedGreeter, Arc::new(fresh_ctx));
                let resp = server
                    .handle_line(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#)
                    .await;
                let schema = &resp["result"]["tools"][0]["inputSchema"];
                // schemars emits a JSON Schema with a "properties" map.
                assert!(
                    schema["properties"]["who"].is_object(),
                    "derived schema must include the `who` field; got: {schema}"
                );
            }
        }
    }


    // These tests exercise the no-governor convenience shim. Under
    // `--features governor` they would hit the
    // `WorkflowWithoutGovernor` hard gate; that's the intended
    // behaviour (workflows must wire a governor when the feature is
    // on). `governor_inbound.rs` covers the same contracts on the
    // builder + `with_governor` path.
    #[cfg(not(feature = "governor"))]
    mod expose_workflow_tests {
        use super::*;
        use async_trait::async_trait;
        use chrono::Utc;
        use klieo_core::agent::{Agent, AgentContext};
        use klieo_core::error::Error as KlieoError;
        use klieo_core::llm::Message;
        use klieo_core::runtime::{ReviewPolicy, RunOptions};
        use klieo_core::test_utils::{fake_context, fake_kv, FakeLlmClient, FakeLlmStep};
        use klieo_core::ToolDef;
        use klieo_hitl::HitlConfig;
        use klieo_hitl_client::HitlClient;
        use secrecy::SecretString;
        use serde::{Deserialize, Serialize};
        use serde_json::json;
        use std::time::Duration;
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        const CHECKPOINT_BUCKET: &str = "klieo.run-checkpoints";
        const WORKSPACE_ID: &str = "ws-test";
        const PLANTED_SENTINEL: &str = "PLANT-SENTINEL-9F7C";

        #[derive(Debug, Clone, Deserialize, Serialize)]
        struct WorkflowIn {
            #[allow(dead_code)]
            payload: String,
        }

        #[derive(Debug, Clone, Serialize)]
        struct UnusedOut;

        struct WorkflowAgent {
            name: &'static str,
        }

        #[async_trait]
        impl Agent for WorkflowAgent {
            type Input = WorkflowIn;
            type Output = UnusedOut;
            type Error = KlieoError;

            fn name(&self) -> &str {
                self.name
            }
            fn system_prompt(&self) -> &str {
                ""
            }
            fn tools(&self) -> &[ToolDef] {
                &[]
            }
            async fn run(
                &self,
                _ctx: AgentContext,
                _input: WorkflowIn,
            ) -> Result<UnusedOut, KlieoError> {
                Err(KlieoError::BadResponse(
                    "workflow path must not call Agent::run".into(),
                ))
            }
        }

        struct PauseOnce(std::sync::atomic::AtomicBool);

        impl PauseOnce {
            fn new() -> Self {
                Self(std::sync::atomic::AtomicBool::new(false))
            }
        }

        #[async_trait]
        impl ReviewPolicy for PauseOnce {
            async fn should_pause_for_approval(
                &self,
                _step: u32,
                _message: &Message,
            ) -> Result<Option<String>, KlieoError> {
                if self.0.swap(true, std::sync::atomic::Ordering::SeqCst) {
                    Ok(None)
                } else {
                    Ok(Some("policy reason that MUST NOT leak to peer".into()))
                }
            }
        }

        fn workflow_ctx_with(steps: Vec<FakeLlmStep>) -> AgentContext {
            let mut ctx = fake_context("workflow-test");
            ctx.llm = Arc::new(FakeLlmClient::new("fake").with_steps(steps));
            ctx.kv = fake_kv();
            ctx
        }

        fn item_json(id: &str, state: &str) -> serde_json::Value {
            json!({
                "id": id, "workspace_id": WORKSPACE_ID, "state": state, "version": 1,
                "escalation_count": 0,
                "decision_context": {"subject_ref":"x","run_id":"r","payload_hash_hex":"h"},
                "reviewer": null, "updated_at": "2026-06-18T00:00:00Z"
            })
        }

        fn hitl_cfg(poll_timeout: Duration) -> HitlConfig {
            HitlConfig::new(
                WORKSPACE_ID,
                CHECKPOINT_BUCKET,
                Duration::from_millis(1),
                poll_timeout,
            )
        }

        fn gated_run_options() -> RunOptions {
            RunOptions::default()
                .with_review_policy(Arc::new(PauseOnce::new()))
                .with_checkpoint_bucket(CHECKPOINT_BUCKET)
        }

        fn plain_run_options() -> RunOptions {
            RunOptions::default()
        }

        fn one_shot_ctx_factory(ctx: AgentContext) -> AgentContextFactory {
            let slot = Arc::new(std::sync::Mutex::new(Some(ctx)));
            Arc::new(move || {
                slot.lock()
                    .unwrap()
                    .take()
                    .expect("ctx_factory called more than once")
            })
        }

        fn input_schema() -> serde_json::Value {
            json!({
                "type": "object",
                "properties": {"payload": {"type": "string"}},
                "required": ["payload"]
            })
        }

        /// Happy path: NeverReview-equivalent runs (no policy installed)
        /// drive `run_with_hitl` straight through. The HITL endpoint is
        /// never hit, and the response carries the final LLM text.
        #[tokio::test]
        async fn invoke_happy_path_returns_text_without_hitl_traffic() {
            let mock = MockServer::start().await;
            // Any HITL traffic at all signals an incorrect suspension.
            Mock::given(method("POST"))
                .and(path("/api/v1/hitl/items"))
                .respond_with(ResponseTemplate::new(500))
                .expect(0)
                .mount(&mock)
                .await;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text("workflow done".into())]);
            let client = Arc::new(HitlClient::new(
                mock.uri(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_secs(1)));

            let server = McpServer::expose_workflow_with_schema(
                WorkflowAgent { name: "wf-happy" },
                "you are a workflow",
                input_schema(),
                plain_run_options(),
                client,
                cfg,
                one_shot_ctx_factory(ctx),
            )
            .unwrap();

            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"wf-happy","arguments":{"payload":"hi"}}}"#,
                )
                .await;

            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains("workflow done"),
                "tools/call must return run_with_hitl text body; got: {text}"
            );
        }

        /// Suspension path: a pausing ReviewPolicy + near-zero
        /// poll_timeout drives `run_with_hitl` into `Suspended`. The
        /// wire response must use the SAFE reason string, MUST NOT echo
        /// the raw policy reason, and MUST NOT contain the planted
        /// sentinel (proves checkpoint/conversation bytes are dropped).
        #[tokio::test]
        async fn invoke_suspend_path_redacts_reason_and_drops_checkpoint() {
            let mock = MockServer::start().await;
            Mock::given(method("POST"))
                .and(path("/api/v1/hitl/items"))
                .respond_with(
                    ResponseTemplate::new(201)
                        .set_body_json(item_json("item-suspend", "awaiting")),
                )
                .mount(&mock)
                .await;
            Mock::given(method("GET"))
                .and(path("/api/v1/hitl/items/item-suspend"))
                .respond_with(
                    ResponseTemplate::new(200)
                        .set_body_json(item_json("item-suspend", "awaiting")),
                )
                .mount(&mock)
                .await;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text(PLANTED_SENTINEL.into())]);
            let client = Arc::new(HitlClient::new(
                mock.uri(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_millis(5)));

            let server = McpServer::expose_workflow_with_schema(
                WorkflowAgent {
                    name: "wf-suspend",
                },
                "you are a suspending workflow",
                input_schema(),
                gated_run_options(),
                client,
                cfg,
                one_shot_ctx_factory(ctx),
            )
            .unwrap();

            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"wf-suspend","arguments":{"payload":"please-approve"}}}"#,
                )
                .await;

            let text = resp["result"]["content"][0]["text"].as_str().unwrap();
            assert!(
                text.contains(r#""status":"suspended""#),
                "suspend response must carry status=suspended; got: {text}"
            );
            assert!(
                text.contains("workflow suspended for human review"),
                "suspend response must carry the safe wire reason; got: {text}"
            );
            assert!(
                !text.contains("policy reason that MUST NOT leak"),
                "raw ReviewPolicy reason leaked to peer: {text}"
            );
            assert!(
                !text.contains(PLANTED_SENTINEL),
                "checkpoint/conversation bytes leaked to peer: {text}"
            );
        }

        /// Error path: HITL submit fails (compliance endpoint 403). The
        /// wire response must be the sanitised stable error envelope;
        /// the underlying `HitlClientError` body MUST NOT leak.
        #[tokio::test]
        async fn invoke_hitl_submit_failure_maps_to_sanitised_tool_error() {
            let mock = MockServer::start().await;
            Mock::given(method("POST"))
                .and(path("/api/v1/hitl/items"))
                .respond_with(ResponseTemplate::new(403).set_body_string("forbidden: token=xyz"))
                .mount(&mock)
                .await;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text("never reached".into())]);
            let client = Arc::new(HitlClient::new(
                mock.uri(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_secs(1)));

            let server = McpServer::expose_workflow_with_schema(
                WorkflowAgent { name: "wf-err" },
                "",
                input_schema(),
                gated_run_options(),
                client,
                cfg,
                one_shot_ctx_factory(ctx),
            )
            .unwrap();

            let resp = server
                .handle_line(
                    r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"wf-err","arguments":{"payload":"x"}}}"#,
                )
                .await;

            assert!(
                resp.get("error").is_some(),
                "submit failure must surface as JSON-RPC error; got: {resp}"
            );
            let msg = resp["error"]["message"].as_str().unwrap();
            assert!(
                msg.contains("tool invocation failed"),
                "wire message must be the sanitised stable string; got: {msg}"
            );
            assert!(
                !msg.contains("token=xyz") && !msg.contains("forbidden"),
                "internal HitlClientError detail leaked: {msg}"
            );
        }


        /// Drive a workflow run through the public `ToolInvoker`
        /// surface with `caller_principal=Some("alice@x")` and assert
        /// no recorded message contains the principal — server-side
        /// authz metadata must never leak into agent memory (D3).
        #[tokio::test]
        async fn caller_principal_does_not_enter_run_state() {
            use klieo_core::test_utils::noop_bus;
            const PRINCIPAL: &str = "alice-NEVER-IN-MEMORY@x";

            let mock = MockServer::start().await;
            Mock::given(method("POST"))
                .and(path("/api/v1/hitl/items"))
                .respond_with(ResponseTemplate::new(500))
                .expect(0)
                .mount(&mock)
                .await;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text("done".into())]);
            let short_term_for_probe = ctx.short_term.clone();
            let episodic_for_probe = ctx.episodic.clone();
            let run_id_for_probe = ctx.run_id;
            let client = Arc::new(HitlClient::new(
                mock.uri(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_secs(1)));

            let server = Arc::new(
                McpServer::builder()
                    .with_hitl(client, cfg)
                    .add_workflow_with_schema(
                        WorkflowAgent { name: "wf-no-leak" },
                        "you are a workflow",
                        input_schema(),
                        plain_run_options(),
                        one_shot_ctx_factory(ctx),
                    )
                    .build()
                    .unwrap(),
            );

            let (pubsub, _, kv, jobs) = noop_bus();
            let tool_ctx = klieo_core::tool::ToolCtx::new(pubsub, kv, jobs)
                .with_caller_principal(PRINCIPAL.into());
            let _result = server
                .invoker
                .invoke(
                    "wf-no-leak",
                    json!({"payload": "hi"}),
                    tool_ctx,
                )
                .await
                .unwrap();

            let history = short_term_for_probe
                .load(klieo_core::ids::ThreadId::new("wf-no-leak:any"), 8192)
                .await
                .unwrap_or_default();
            for msg in &history {
                assert!(
                    !msg.content.contains(PRINCIPAL),
                    "principal leaked into short-term memory: {}",
                    msg.content
                );
            }

            // Hashed attribution label MUST land on the audit trail,
            // raw principal MUST NOT.
            let episodes = episodic_for_probe
                .replay(run_id_for_probe)
                .await
                .expect("episodic replay must succeed");
            let expected_label = klieo_core::principal_hash(PRINCIPAL);
            let attributed_labels: Vec<&str> = episodes
                .iter()
                .filter_map(|e| match e {
                    klieo_core::Episode::RunAttributed { tenant_label } => {
                        Some(tenant_label.as_str())
                    }
                    _ => None,
                })
                .collect();
            assert_eq!(
                attributed_labels,
                vec![expected_label.as_str()],
                "exactly one RunAttributed carrying principal_hash; got {episodes:?}",
            );
            for ep in &episodes {
                let payload = serde_json::to_string(ep).expect("episode serialises");
                assert!(
                    !payload.contains(PRINCIPAL),
                    "raw principal leaked into recorded episode: {payload}",
                );
            }
        }

        /// Negative path: no `caller_principal` on the inbound
        /// `ToolCtx` (non-MCP / unauthenticated transport) records
        /// zero `RunAttributed` episodes.
        #[tokio::test]
        async fn no_principal_yields_no_run_attributed_episode() {
            use klieo_core::test_utils::noop_bus;
            let mock = MockServer::start().await;
            Mock::given(method("POST"))
                .and(path("/api/v1/hitl/items"))
                .respond_with(ResponseTemplate::new(500))
                .expect(0)
                .mount(&mock)
                .await;
            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text("done".into())]);
            let episodic_for_probe = ctx.episodic.clone();
            let run_id_for_probe = ctx.run_id;
            let client = Arc::new(HitlClient::new(
                mock.uri(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_secs(1)));
            let server = Arc::new(
                McpServer::builder()
                    .with_hitl(client, cfg)
                    .add_workflow_with_schema(
                        WorkflowAgent { name: "wf-anon" },
                        "you are a workflow",
                        input_schema(),
                        plain_run_options(),
                        one_shot_ctx_factory(ctx),
                    )
                    .build()
                    .unwrap(),
            );
            let (pubsub, _, kv, jobs) = noop_bus();
            let tool_ctx = klieo_core::tool::ToolCtx::new(pubsub, kv, jobs);
            let _ = server
                .invoker
                .invoke("wf-anon", json!({"payload": "hi"}), tool_ctx)
                .await
                .unwrap();
            let episodes = episodic_for_probe
                .replay(run_id_for_probe)
                .await
                .expect("episodic replay must succeed");
            let attributed_count = episodes
                .iter()
                .filter(|e| matches!(e, klieo_core::Episode::RunAttributed { .. }))
                .count();
            assert_eq!(
                attributed_count, 0,
                "RunAttributed must not appear without a caller_principal; got {episodes:?}",
            );
        }

        /// Ticket-issuance gate: suspend with both `with_checkpoint_kv`
        /// AND a principal yields `{status,ticket,reason}`; suspend
        /// missing either falls back to the slice-1 envelope. In both
        /// paths the planted PLANTED_SENTINEL bytes (carried inside
        /// the dropped checkpoint) MUST NOT cross the wire.
        async fn run_suspend_with(
            with_kv: bool,
            with_principal: bool,
        ) -> serde_json::Value {
            use klieo_core::test_utils::noop_bus;
            let mock = MockServer::start().await;
            Mock::given(method("POST"))
                .and(path("/api/v1/hitl/items"))
                .respond_with(
                    ResponseTemplate::new(201)
                        .set_body_json(item_json("item-suspend", "awaiting")),
                )
                .mount(&mock)
                .await;
            Mock::given(method("GET"))
                .and(path("/api/v1/hitl/items/item-suspend"))
                .respond_with(
                    ResponseTemplate::new(200)
                        .set_body_json(item_json("item-suspend", "awaiting")),
                )
                .mount(&mock)
                .await;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text(PLANTED_SENTINEL.into())]);
            let client = Arc::new(HitlClient::new(
                mock.uri(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_millis(5)));

            let mut builder = McpServer::builder()
                .with_hitl(client, cfg)
                .add_workflow_with_schema(
                    WorkflowAgent { name: "wf-suspend" },
                    "",
                    input_schema(),
                    gated_run_options(),
                    one_shot_ctx_factory(ctx),
                );
            if with_kv {
                builder = builder.with_checkpoint_kv(fake_kv());
            }
            let server = Arc::new(builder.build().unwrap());

            let (pubsub, _, kv, jobs) = noop_bus();
            let mut tool_ctx = klieo_core::tool::ToolCtx::new(pubsub, kv, jobs);
            if with_principal {
                tool_ctx = tool_ctx.with_caller_principal("alice@x".into());
            }
            server
                .invoker
                .invoke(
                    "wf-suspend",
                    json!({"payload": "please-approve"}),
                    tool_ctx,
                )
                .await
                .unwrap()
        }

        #[tokio::test]
        async fn suspend_with_kv_and_principal_issues_ticket() {
            let envelope = run_suspend_with(true, true).await;
            let body = envelope.to_string();
            assert_eq!(envelope["status"], "suspended");
            let ticket = envelope["ticket"].as_str().expect("ticket present");
            assert!(!ticket.is_empty(), "issued ticket must be non-empty");
            assert!(
                !body.contains(PLANTED_SENTINEL),
                "checkpoint bytes leaked: {body}"
            );
            assert!(
                !body.contains("policy reason that MUST NOT leak"),
                "raw policy reason leaked: {body}"
            );
        }

        #[tokio::test]
        async fn suspend_without_kv_falls_back_to_no_ticket_envelope() {
            let envelope = run_suspend_with(false, true).await;
            assert_eq!(envelope["status"], "suspended");
            assert!(
                envelope.get("ticket").is_none(),
                "no checkpoint KV must yield slice-1 envelope (no ticket field)",
            );
            let body = envelope.to_string();
            assert!(!body.contains(PLANTED_SENTINEL));
        }

        #[tokio::test]
        async fn suspend_without_principal_falls_back_to_no_ticket_envelope() {
            let envelope = run_suspend_with(true, false).await;
            assert_eq!(envelope["status"], "suspended");
            assert!(
                envelope.get("ticket").is_none(),
                "no caller principal must yield slice-1 envelope (no ticket field)",
            );
            let body = envelope.to_string();
            assert!(!body.contains(PLANTED_SENTINEL));
        }

        /// IDOR mandatory test: principal B resumes principal A's
        /// ticket → fail-closed and the ticket stays consumable by A.
        /// Drives the lookup-then-authz-then-claim sequence directly
        /// against the ticket store (the HTTP method is a thin wrapper
        /// around it).
        #[tokio::test]
        async fn principal_b_cannot_consume_principal_a_ticket() {
            use crate::resume_ticket::{ResumeTicketRecord, ResumeTicketStore};
            let store = ResumeTicketStore::new(fake_kv());
            let token = ResumeTicketStore::mint_token();
            let cp_json = serde_json::json!({
                "run_id": klieo_core::ids::RunId::new(),
                "step_index": 1,
                "thread_id": "t-idor",
                "messages": [],
                "pending_tool_calls": null,
                "created_at": "2026-06-18T00:00:00Z",
            });
            let checkpoint = serde_json::from_value(cp_json).unwrap();
            let record = ResumeTicketRecord {
                principal: "alice@x".into(),
                workflow_name: "wf".into(),
                checkpoint,
                created_at: Utc::now(),
            };
            store.persist(&token, &record).await.unwrap();

            // Principal B follows the lookup→authz→claim sequence.
            let peeked = store.peek(&token).await.unwrap().expect("ticket present");
            let principal_b = "mallory@x";
            assert_ne!(
                peeked.principal, principal_b,
                "fixture must seed a distinct principal so the authz arm engages"
            );
            // The handler refuses BEFORE the claim — record this:
            // the ticket must therefore stay consumable by Alice.
            let consumed_by_alice = store.claim(&token).await.unwrap();
            assert!(
                consumed_by_alice.is_some(),
                "after a foreign-principal denial the rightful owner can still resume"
            );
            // And a second claim now returns None (the previous succeeded).
            let after = store.claim(&token).await.unwrap();
            assert!(after.is_none(), "the now-consumed ticket cannot be reused");
        }

        /// Concurrent double-resume of one ticket: exactly one wins,
        /// the other observes None (and would fail-closed at the
        /// handler seam).
        #[tokio::test]
        async fn concurrent_resume_runs_exactly_once() {
            use crate::resume_ticket::{ResumeTicketRecord, ResumeTicketStore};
            let store = Arc::new(ResumeTicketStore::new(fake_kv()));
            let token = ResumeTicketStore::mint_token();
            let cp_json = serde_json::json!({
                "run_id": klieo_core::ids::RunId::new(),
                "step_index": 1,
                "thread_id": "t-conc",
                "messages": [],
                "pending_tool_calls": null,
                "created_at": "2026-06-18T00:00:00Z",
            });
            let checkpoint = serde_json::from_value(cp_json).unwrap();
            let record = ResumeTicketRecord {
                principal: "alice@x".into(),
                workflow_name: "wf".into(),
                checkpoint,
                created_at: Utc::now(),
            };
            store.persist(&token, &record).await.unwrap();
            let racers: Vec<_> = (0..8)
                .map(|_| {
                    let store = store.clone();
                    let token = token.clone();
                    tokio::spawn(async move { store.claim(&token).await })
                })
                .collect();
            let mut winners = 0usize;
            for handle in racers {
                if handle.await.unwrap().unwrap().is_some() {
                    winners += 1;
                }
            }
            assert_eq!(
                winners, 1,
                "concurrent ticket consumption must run exactly once; got {winners}"
            );
        }

        /// Happy approve resume: a `WorkflowAsToolInvoker` resume
        /// handle drives `resume_from_checkpoint` to completion when
        /// the ctx_factory mints a real bus + kv. Validates D8 (the
        /// resume-time ctx shares the same KV semantics as suspend
        /// — here we exercise the bucketless branch, the latch test
        /// in klieo-core covers the bucket variant).
        #[tokio::test]
        async fn approve_resume_drives_run_to_completion() {
            use klieo_core::checkpoint::ApprovalDecision;
            use std::sync::Mutex;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text("approved".into())]);
            let cp_json = serde_json::json!({
                "run_id": ctx.run_id,
                "step_index": 1,
                "thread_id": "t-resume-approve",
                "messages": [],
                "pending_tool_calls": null,
                "created_at": "2026-06-18T00:00:00Z",
            });
            let checkpoint: klieo_core::checkpoint::RunCheckpoint =
                serde_json::from_value(cp_json).unwrap();

            let client = Arc::new(HitlClient::new(
                "http://unused".to_string(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_secs(1)));
            let ctx_holder = Arc::new(Mutex::new(Some(ctx)));
            let ctx_factory: AgentContextFactory = Arc::new(move || {
                ctx_holder
                    .lock()
                    .unwrap()
                    .take()
                    .expect("ctx_factory drained")
            });

            let invoker = Arc::new(crate::workflow::WorkflowAsToolInvoker::<WorkflowAgent>::new(
                "wf-resume".into(),
                "".into(),
                input_schema(),
                ctx_factory,
                plain_run_options(),
                crate::workflow::HitlBundle { client, cfg },
                None,
                #[cfg(feature = "governor")]
                None,
            ));

            let handle: Arc<dyn crate::workflow::WorkflowResumeHandle> = invoker;
            let result = handle
                .resume(checkpoint, ApprovalDecision::Approved, "hashed-tenant".into())
                .await
                .unwrap();
            assert_eq!(result, serde_json::Value::String("approved".into()));
        }

        /// Reject resume: the rejection reason is appended to short-
        /// term memory so the LLM sees the operator's verdict, then
        /// the run completes.
        #[tokio::test]
        async fn reject_resume_feeds_reason_back_to_model() {
            use klieo_core::checkpoint::ApprovalDecision;
            use klieo_core::llm::Role;
            use std::sync::Mutex;

            let ctx = workflow_ctx_with(vec![FakeLlmStep::Text("acknowledged".into())]);
            let short_term_for_probe = ctx.short_term.clone();
            let cp_json = serde_json::json!({
                "run_id": ctx.run_id,
                "step_index": 1,
                "thread_id": "t-resume-reject",
                "messages": [],
                "pending_tool_calls": null,
                "created_at": "2026-06-18T00:00:00Z",
            });
            let checkpoint: klieo_core::checkpoint::RunCheckpoint =
                serde_json::from_value(cp_json).unwrap();

            let client = Arc::new(HitlClient::new(
                "http://unused".to_string(),
                SecretString::from("tok".to_string()),
            ));
            let cfg = Arc::new(hitl_cfg(Duration::from_secs(1)));
            let ctx_holder = Arc::new(Mutex::new(Some(ctx)));
            let ctx_factory: AgentContextFactory = Arc::new(move || {
                ctx_holder
                    .lock()
                    .unwrap()
                    .take()
                    .expect("ctx_factory drained")
            });

            let invoker = Arc::new(crate::workflow::WorkflowAsToolInvoker::<WorkflowAgent>::new(
                "wf-reject".into(),
                "".into(),
                input_schema(),
                ctx_factory,
                plain_run_options(),
                crate::workflow::HitlBundle { client, cfg },
                None,
                #[cfg(feature = "governor")]
                None,
            ));

            let handle: Arc<dyn crate::workflow::WorkflowResumeHandle> = invoker;
            let _ = handle
                .resume(
                    checkpoint,
                    ApprovalDecision::Rejected {
                        reason: "BAD-IDEA-XYZ".into(),
                    },
                    "hashed-tenant".into(),
                )
                .await
                .unwrap();

            let history = short_term_for_probe
                .load(
                    klieo_core::ids::ThreadId::new("t-resume-reject"),
                    8192,
                )
                .await
                .unwrap();
            let rejection_seen = history
                .iter()
                .any(|m| m.role == Role::Tool && m.content.contains("BAD-IDEA-XYZ"));
            assert!(
                rejection_seen,
                "the model must see the operator's rejection reason on resume"
            );
        }

        /// Builder guard: workflow registered without `with_hitl` must
        /// fail `build()` with the typed `WorkflowWithoutHitl` variant.
        #[test]
        fn builder_rejects_workflow_without_hitl() {
            let ctx_factory: AgentContextFactory = Arc::new(|| fake_context("guard-test"));
            let err = McpServer::builder()
                .add_workflow_with_schema(
                    WorkflowAgent { name: "wf-guard" },
                    "",
                    input_schema(),
                    plain_run_options(),
                    ctx_factory,
                )
                .build()
                .err()
                .expect("workflow without with_hitl must fail build");
            assert!(
                matches!(err, McpBuildError::WorkflowWithoutHitl),
                "expected WorkflowWithoutHitl, got: {err:?}"
            );
        }
    }

    #[test]
    fn resume_errors_render_messages() {
        let a = McpServerError::ResumeBufferExpired { since_id: 7 };
        assert_eq!(a.to_string(), "resume window expired (since_id=7)");
        let b = McpServerError::ResumeBufferNotFound("tok".into());
        assert_eq!(b.to_string(), "no buffered stream for progressToken: tok");
    }

    #[test]
    fn from_server_outbound_serialisation_maps_to_outbound_serialisation() {
        use klieo_core::ServerOutboundError;
        let serde_err = serde_json::from_str::<serde_json::Value>("{invalid}").unwrap_err();
        let mcp_err = McpServerError::from(ServerOutboundError::Serialisation(serde_err));
        assert!(
            matches!(mcp_err, McpServerError::OutboundSerialisation(_)),
            "ServerOutboundError::Serialisation must map to McpServerError::OutboundSerialisation; got {mcp_err:?}"
        );
        use std::error::Error;
        assert!(
            mcp_err.source().is_some(),
            "McpServerError::OutboundSerialisation must expose source via #[source]"
        );
    }

    struct NamedAuthn;

    #[async_trait]
    impl klieo_auth_common::Authenticator for NamedAuthn {
        async fn authenticate(
            &self,
            _headers: &dyn klieo_auth_common::Headers,
            _payload: &[u8],
        ) -> Result<klieo_auth_common::Identity, klieo_auth_common::AuthError> {
            Ok(klieo_auth_common::Identity::new("alice"))
        }
    }

    #[test]
    fn regulated_without_tenant_kv_fails_closed() {
        let err = McpServer::builder()
            .add_tools(Arc::new(OneToolInvoker))
            .with_authenticator(Arc::new(NamedAuthn))
            .profile(klieo_core::DeploymentProfile::RegulatedMultiTenant)
            .build()
            .err()
            .expect("must fail closed");
        assert!(matches!(
            err,
            McpBuildError::RegulatedProfile(klieo_core::ProfileViolation::MissingTenantKv)
        ));
    }

    #[test]
    fn regulated_without_authenticator_fails_closed() {
        let kv = klieo_bus_memory::MemoryBus::new().kv;
        let err = McpServer::builder()
            .add_tools(Arc::new(OneToolInvoker))
            .with_tenant_binding(kv)
            .profile(klieo_core::DeploymentProfile::RegulatedMultiTenant)
            .build()
            .err()
            .expect("must fail closed");
        assert!(matches!(
            err,
            McpBuildError::RegulatedProfile(klieo_core::ProfileViolation::AnonymousAuth)
        ));
    }

    struct AnonAuthn;

    #[async_trait]
    impl klieo_auth_common::Authenticator for AnonAuthn {
        async fn authenticate(
            &self,
            _headers: &dyn klieo_auth_common::Headers,
            _payload: &[u8],
        ) -> Result<klieo_auth_common::Identity, klieo_auth_common::AuthError> {
            Ok(klieo_auth_common::Identity::anonymous())
        }

        fn allows_anonymous(&self) -> bool {
            true
        }
    }

    #[test]
    fn regulated_with_anonymous_authenticator_fails_closed() {
        let kv = klieo_bus_memory::MemoryBus::new().kv;
        let err = McpServer::builder()
            .add_tools(Arc::new(OneToolInvoker))
            .with_tenant_binding(kv)
            .with_authenticator(Arc::new(AnonAuthn))
            .profile(klieo_core::DeploymentProfile::RegulatedMultiTenant)
            .build()
            .err()
            .expect("must fail closed");
        assert!(matches!(
            err,
            McpBuildError::RegulatedProfile(klieo_core::ProfileViolation::AnonymousAuth)
        ));
    }

    #[test]
    fn regulated_forces_strict_over_lenient_binding() {
        let kv = klieo_bus_memory::MemoryBus::new().kv;
        let server = McpServer::builder()
            .add_tools(Arc::new(OneToolInvoker))
            .with_tenant_binding(kv) // lenient — profile must upgrade
            .with_authenticator(Arc::new(NamedAuthn))
            .profile(klieo_core::DeploymentProfile::RegulatedMultiTenant)
            .build()
            .expect("regulated build with named auth + kv must succeed");
        assert_eq!(
            server.ownership_registry.as_ref().map(|r| r.is_strict()),
            Some(true)
        );
    }

    #[test]
    fn unprofiled_keeps_lenient_binding() {
        let kv = klieo_bus_memory::MemoryBus::new().kv;
        let server = McpServer::builder()
            .add_tools(Arc::new(OneToolInvoker))
            .with_tenant_binding(kv)
            .with_authenticator(Arc::new(NamedAuthn))
            .build()
            .expect("unprofiled build ok");
        assert_eq!(
            server.ownership_registry.as_ref().map(|r| r.is_strict()),
            Some(false)
        );
    }
}

#[cfg(test)]
#[cfg(feature = "http")]
mod jsonrpc_const_tests {
    use super::*;

    #[test]
    fn jsonrpc_constants_are_i64_and_unique_at_runtime() {
        // Documented runtime mirror of the const _ uniqueness block.
        // The block itself enforces the contract at compile time;
        // this test makes it grep-able and includes a negative-fixture
        // assertion that the dedup logic actually catches duplicates.
        let codes: [i64; 9] = [
            JSONRPC_PARSE_ERROR,
            JSONRPC_METHOD_NOT_FOUND,
            JSONRPC_INVALID_PARAMS,
            JSONRPC_SERVER_ERROR,
            JSONRPC_UNAUTHENTICATED,
            JSONRPC_RESUME_BUFFER_EXPIRED,
            JSONRPC_RESUME_BUFFER_NOT_FOUND,
            JSONRPC_LEADER_DIED,
            JSONRPC_SESSION_CONFLICT,
        ];
        let mut seen = std::collections::HashSet::new();
        let mut duplicates: Vec<i64> = Vec::new();
        for code in codes {
            if !seen.insert(code) {
                duplicates.push(code);
            }
        }
        assert!(
            duplicates.is_empty(),
            "JSONRPC_* codes must be unique; found duplicates: {duplicates:?}"
        );

        // Negative fixture — the dedup logic actually rejects a known
        // duplicate (regression guard against the test itself drifting
        // into vacuous truth).
        let mut local = seen.clone();
        assert!(
            !local.insert(JSONRPC_PARSE_ERROR),
            "duplicate detection logic broken"
        );
    }
}