klieo-mcp-server 2.2.0

//! HTTP transport for `McpServer`. Streamable HTTP per the MCP
//! 2025-03-26 spec. Single-shot JSON responses for all methods except
//! `tools/call` with `_meta.progressToken`, which upgrades to an SSE
//! stream per MCP spec §6.4.
//!
//! Gated behind the `http` cargo feature. See the crate-level
//! docs for out-of-scope items (sessions, auth).
//!
//! ## Lock acquisition order
//!
//! Across the multi-session paths:
//!   1. [`McpServer::sessions`](crate::McpServer) (write or read)
//!   2. [`McpServer::principal_counts`](crate::McpServer) (write)
//!
//! Either lock is taken alone where the other is not needed. When
//! both are needed (the `initialize` admission check), `sessions` is
//! taken FIRST and `principal_counts` SECOND. Eviction paths
//! decrement `principal_counts` AFTER releasing the `sessions`
//! guard, preserving the order.

use crate::{
    rpc_error, tool_error_to_envelope, McpServer, McpServerError, JSONRPC_INVALID_PARAMS,
    JSONRPC_LEADER_DIED, JSONRPC_PARSE_ERROR, JSONRPC_SERVER_ERROR, JSONRPC_UNAUTHENTICATED,
    MCP_LEADER_KEY_PREFIX,
};
use axum::{
    body::Bytes,
    extract::{DefaultBodyLimit, State},
    http::{header, HeaderMap, StatusCode},
    response::{IntoResponse, Response},
    routing::post,
    Json, Router,
};
use futures::{Stream, StreamExt as _};
use klieo_auth_common::Identity;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
use tracing::{error, instrument, warn};
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

const MAX_BODY_BYTES: usize = 1 << 20; // 1 MiB
const CANCEL_SUBJECT_PREFIX: &str = "klieo.mcp.cancel.";
/// Key-prefix in the `klieo-tenants` KV bucket for MCP streams.
/// Mirrors `A2A_OWNERSHIP_KEY_PREFIX` on the A2A side so a single
/// bucket can host both transports without collisions.
const MCP_OWNERSHIP_KEY_PREFIX: &str = "mcp.";
/// Wire header carrying the streamable-HTTP session identifier.
/// Set on the `initialize` response, echoed by the client on every
/// subsequent POST + GET. ADR-028.
const MCP_SESSION_ID_HEADER: &str = "Mcp-Session-Id";

/// Standard SSE-resumption header. Clients echo the last `id:` they
/// observed; the GET /mcp handler parses it as a `u64` and snapshots
/// the session's resume buffer for any frames strictly newer than the
/// supplied id.
const LAST_EVENT_ID_HEADER: &str = "Last-Event-Id";

/// Initial capacity hint for the per-frame SSE body buffer.
///
/// Sized to cover the `id: <N>\ndata: <payload>\n\n` framing
/// plus a typical small JSON payload without forcing a realloc.
/// Larger payloads grow the `BytesMut` once.
pub(crate) const SSE_FRAME_HINT_BYTES: usize = 512;

/// Per-batch JSON-RPC element cap. Combined with the 1 MiB body cap
/// this bounds worst-case batch work.
const MAX_BATCH_ITEMS: usize = 100;

/// Construct the structured-logged 500 returned when a session-mint
/// `OnceCell::set` lost a race. Distinguishes a real server invariant
/// violation from the legitimate `.get().is_some()` 409 path: the
/// `.get()` branch fires whenever a second `initialize` races past the
/// guard with a session already minted; hitting THIS branch means two
/// requests both raced past the `.get()` check, which is an internal
/// invariant violation worth surfacing to operators rather than
/// presenting as a benign client-side conflict.
fn mint_session_race_500(
    raw_id: Option<&serde_json::Value>,
    field: &'static str,
    attempted_session_id: Option<uuid::Uuid>,
) -> Response {
    error!(
        target: "klieo::mcp::session",
        field,
        attempted_session_id = ?attempted_session_id,
        "session-mint race (concurrent initialize past guard)"
    );
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        Json(rpc_error(
            raw_id.cloned(),
            JSONRPC_SERVER_ERROR,
            "internal: session mint race",
        )),
    )
        .into_response()
}

impl McpServer {
    /// Build the axum [`Router`] that serves `POST /mcp`. The
    /// caller may mount this under their own application or hand
    /// it to [`Self::serve_http`].
    ///
    /// No middleware is applied beyond a 1 MiB body limit and a
    /// `Content-Type: application/json` guard. CORS, tracing,
    /// auth, and rate-limiting are caller-owned — wrap the
    /// returned `Router` with whatever `tower` layers you need.
    pub fn router(self: &Arc<Self>) -> Router {
        router_impl(self.clone())
    }

    /// Bind to `addr` and serve until `parent_cancel` fires or
    /// the listener errors out.
    ///
    /// # Security
    /// No authentication is performed. Bind to `127.0.0.1`
    /// unless the listener is fronted by a reverse proxy that
    /// enforces auth. Plain HTTP only — terminate TLS at the
    /// proxy.
    pub async fn serve_http(self: Arc<Self>, addr: SocketAddr) -> Result<(), McpServerError> {
        let cancel = self.parent_cancel.clone();
        let listener = tokio::net::TcpListener::bind(addr).await?;
        let router = self.router();
        axum::serve(listener, router)
            .with_graceful_shutdown(async move { cancel.cancelled().await })
            .await?;
        Ok(())
    }
}

fn router_impl(server: Arc<McpServer>) -> Router {
    Router::new()
        .route("/mcp", post(post_mcp).get(get_mcp).delete(delete_mcp))
        .layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
        .with_state(server)
}

#[instrument(
    skip_all,
    fields(
        rpc.system = "klieo-mcp",
        rpc.method = tracing::field::Empty,
        http.request.method = "POST",
    ),
)]
async fn post_mcp(
    State(server): State<Arc<McpServer>>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    // Cluster 0.23: stitch the current span under any upstream W3C
    // tracecontext carried in `traceparent` / `tracestate` headers so
    // load-balancer / sibling-service traces fold into klieo's tree.
    // Empty / malformed headers yield an empty Context; behaviour
    // pre-0.23 (no parent) is preserved when callers send neither.
    let parent_cx = klieo_core::extract_traceparent(&klieo_headers_from_axum(&headers));
    tracing::Span::current().set_parent(parent_cx);

    if server.parent_cancel.is_cancelled() {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(rpc_error(
                None,
                JSONRPC_SERVER_ERROR,
                "server shutting down",
            )),
        )
            .into_response();
    }

    if !content_type_is_json(&headers) {
        return StatusCode::UNSUPPORTED_MEDIA_TYPE.into_response();
    }

    let raw: serde_json::Value = match serde_json::from_slice(&body) {
        Ok(v) => v,
        Err(e) => {
            warn!(error = %e, "rejected malformed JSON-RPC body");
            return (
                StatusCode::BAD_REQUEST,
                Json(rpc_error(
                    None,
                    JSONRPC_PARSE_ERROR,
                    "malformed JSON-RPC body",
                )),
            )
                .into_response();
        }
    };

    if let Some(method) = raw.get("method").and_then(|m| m.as_str()) {
        tracing::Span::current().record("rpc.method", method);
    }

    // OAuth / bearer / arbitrary-Authenticator gate. Runs BEFORE
    // `try_sse_upgrade` so an unauthenticated tools/call with a
    // progressToken returns a JSON -32001 envelope rather than an
    // SSE stream (the stream would otherwise leak the start of an
    // invocation an unauthorised principal must not see). Skipped
    // entirely when no authenticator is wired, preserving the
    // pre-0.21 caller-owned-middleware contract.
    let identity = match enforce_authenticator(&server, &headers, &body, &raw).await {
        Ok(identity) => identity,
        Err(rejection) => return rejection,
    };

    // Streamable-HTTP session lifecycle. `initialize` mints the session
    // id + echoes it in the response header. Every subsequent non-
    // `initialize` request must echo the same id back — mismatched ids
    // are rejected with 409, missing ids with 400. ADR-028.
    let method = raw.get("method").and_then(|m| m.as_str()).unwrap_or("");
    if method == "initialize" {
        return handle_initialize_post(&server, raw, identity).await;
    }

    // Non-initialize POSTs require a session header once any session
    // has been minted on this server. The registry-emptiness guard
    // lets a non-initialize POST on a fresh server (no sessions yet)
    // fall through to `dispatch` without a session, so HTTP servers
    // that never `initialize` still serve in-process unit-test
    // traffic. Once the registry holds at least one entry, the
    // header is mandatory. `identity` (when present) is threaded into
    // `require_session` so the per-session principal binding is
    // enforced on every subsequent request — CWE-639 cross-tenant
    // session reuse returns 404.
    let session = if server.sessions.read().await.is_empty() {
        None
    } else {
        match require_session(&headers, &server, identity.as_ref()).await {
            Ok(s) => Some(s),
            Err(rejection) => return rejection,
        }
    };

    // HTTP outbound responses arrive as POST bodies with shape
    // {id, result|error, ...} and no `method`. Route them to the
    // session's outbound correlation table rather than handle_jsonrpc.
    // Only single-object bodies qualify — batch arrays fall through
    // to `dispatch` so their per-element method semantics survive.
    // ADR-028.
    if raw.is_object() && raw.get("method").is_none() {
        let Some(session) = session.as_ref() else {
            warn!(
                target: "mcp.http",
                "POST /mcp body has no method and no session header; rejecting",
            );
            return (
                StatusCode::BAD_REQUEST,
                Json(rpc_error(None, JSONRPC_PARSE_ERROR, "missing method")),
            )
                .into_response();
        };
        return route_outbound_response(&server, session, raw).await;
    }

    if let Some(session) = session.as_ref() {
        touch_last_activity(&server, session);
    }

    // SSE upgrade: `tools/call` with `_meta.progressToken` switches to
    // an SSE stream carrying progress frames + a terminal result frame.
    // `klieo/tools/resume` replays buffered events for a progressToken.
    // The verified `Identity` (when present) is threaded into the
    // upgrade path so the cluster-0.22 tenant-binding gate can claim
    // ownership on `tools/call` and check it on `klieo/tools/resume`.
    if let Some(sse) = try_sse_upgrade(server.clone(), raw.clone(), identity).await {
        return sse;
    }

    let resp = dispatch(&server, raw, session.as_ref()).await;
    (StatusCode::OK, Json(resp)).into_response()
}

/// Bump the per-session activity timestamp on the supplied session.
/// Called on every successful POST and on every outbound frame send.
/// The idle reaper compares this against `session_idle_timeout` to
/// detect idle sessions. ADR-028.
///
/// Lock-free relaxed atomic store; synchronous. `server` supplies the
/// reference instant used to encode the timestamp.
fn touch_last_activity(server: &Arc<McpServer>, session: &Arc<crate::session::Session>) {
    session.mark_active(server.server_start);
}

/// Route a POST body that has no `method` (a JSON-RPC response shape:
/// `{id, result|error, ...}`) to the supplied session's outbound
/// correlation table. Returns `202 Accepted` once `complete_pending`
/// has fired, or `400 Bad Request` when the body is missing a
/// numeric `id` OR no outbound primitive is wired on this session
/// (the SSE-stream open path wires it on `GET /mcp`). ADR-028.
async fn route_outbound_response(
    server: &Arc<McpServer>,
    session: &Arc<crate::session::Session>,
    raw: serde_json::Value,
) -> Response {
    let id = raw.get("id").and_then(|v| v.as_i64());
    let outbound = session.outbound.get();
    if let (Some(id), Some(outbound)) = (id, outbound) {
        outbound.complete_pending(id, raw).await;
        touch_last_activity(server, session);
        return StatusCode::ACCEPTED.into_response();
    }
    warn!(
        target: "mcp.http",
        "POST /mcp body has no method and no routable outbound; rejecting",
    );
    (
        StatusCode::BAD_REQUEST,
        Json(rpc_error(None, JSONRPC_PARSE_ERROR, "missing method")),
    )
        .into_response()
}

/// Streamable-HTTP outbound stream entry. Validates the session,
/// authenticates the caller, mints a per-session mpsc whose receiver
/// becomes the response body, and returns an SSE response carrying
/// newline-delimited JSON-RPC frames. ADR-028.
///
/// Single-shot per session: a second concurrent `GET /mcp` for the
/// same active session returns `409 Conflict`. The receiver's
/// lifetime equals the response body's lifetime.
#[instrument(
    skip_all,
    fields(
        rpc.system = "klieo-mcp",
        http.request.method = "GET",
    ),
)]
async fn get_mcp(State(server): State<Arc<McpServer>>, headers: HeaderMap) -> Response {
    let parent_cx = klieo_core::extract_traceparent(&klieo_headers_from_axum(&headers));
    tracing::Span::current().set_parent(parent_cx);

    if server.parent_cancel.is_cancelled() {
        return StatusCode::SERVICE_UNAVAILABLE.into_response();
    }

    // Authenticator runs BEFORE the SSE stream is minted so a 401
    // never leaks a stream byte to an unauthenticated peer. GET
    // carries no JSON-RPC body and no `method` slot, so only the
    // `authenticate` half runs; `authorize_method` is deliberately
    // skipped — scope on the GET observer is granted transitively
    // through the session-id check from the (already-authorized)
    // `initialize`. The verified `Identity` is threaded into
    // `require_session` so the per-session principal binding is
    // enforced before the SSE upgrade starts. ADR-028.
    let identity = match enforce_authenticator_for_get(&server, &headers).await {
        Ok(identity) => identity,
        Err(rejection) => return rejection,
    };

    let session = match require_session(&headers, &server, identity.as_ref()).await {
        Ok(s) => s,
        Err(rejection) => return rejection,
    };
    let session_id = session.id.expect("HTTP session always carries Some(uuid)");

    let replay = match compute_replay_window(&headers, &session, &server) {
        Ok(slice) => slice,
        Err(rejection) => return rejection,
    };

    let (tx, rx) = crate::outbound_ring::bounded_ring::<(u64, std::sync::Arc<serde_json::Value>)>(
        crate::outbound_sink::OUTBOUND_QUEUE_CAPACITY,
    );

    // The outbound primitive's `OnceCell` is a single-writer slot per
    // session: only the first `GET /mcp` for a session may set it.
    // Hitting `Err` here means a concurrent GET raced past
    // `require_session` for the same session id — a server invariant
    // violation, surfaced via `mint_session_race_500` rather than
    // collapsed to a generic 4xx.
    if session.outbound_tx.set(tx.clone()).is_err() {
        return mint_session_race_500(None, "session.outbound_tx", session.id);
    }

    if let Err(rejection) = wire_session_outbound(&server, &session, tx) {
        return rejection;
    }

    let (cleanup_tx, cleanup_rx) = tokio::sync::oneshot::channel::<()>();
    spawn_session_cleanup(server.clone(), session_id, cleanup_rx);
    build_outbound_sse_response(session_id, replay, rx, cleanup_tx)
}

/// Snapshots the per-session resume buffer for the slice with event id
/// strictly greater than the client's `Last-Event-Id`. Returns
/// `Ok(Vec::new())` when the header is absent. Returns
/// `Err(Response)` for the 400 / 410 / 501 wire-status edges:
///
/// - `400 Bad Request` — header is non-ASCII, not parseable as `u64`,
///   or ahead of the server's monotonic event sequence.
/// - `410 Gone` (rpc_error envelope, message `"resume gap; reconnect
///   with fresh initialize"`) — id older than the oldest retained
///   buffer entry, so the replay window cannot be honoured without
///   forging a gap.
/// - `501 Not Implemented` — server built with
///   `with_sse_replay_capacity(0)`.
///
/// The buffer Mutex is dropped before the returned `Vec` leaves; the
/// caller never holds the lock across the SSE body's lifetime. The
/// function is synchronous because every step — header parse, atomic
/// load, `parking_lot::Mutex` guard acquisition, slice collection —
/// completes without yielding.
#[allow(clippy::result_large_err)]
fn compute_replay_window(
    headers: &HeaderMap,
    session: &crate::session::Session,
    server: &McpServer,
) -> Result<Vec<(u64, std::sync::Arc<serde_json::Value>)>, Response> {
    let Some(raw) = headers.get(LAST_EVENT_ID_HEADER) else {
        return Ok(Vec::new());
    };
    let Ok(header_str) = raw.to_str() else {
        return Err((StatusCode::BAD_REQUEST, "Last-Event-Id must be ASCII u64").into_response());
    };
    let Ok(last_id) = header_str.parse::<u64>() else {
        return Err((StatusCode::BAD_REQUEST, "Last-Event-Id is not a valid u64").into_response());
    };
    if !server.sse_replay_enabled() {
        return Err((
            StatusCode::NOT_IMPLEMENTED,
            "resume buffer disabled (with_sse_replay_capacity(0))",
        )
            .into_response());
    }
    let current_head = session
        .next_event_id
        .load(std::sync::atomic::Ordering::Relaxed);
    if last_id >= current_head {
        return Err((
            StatusCode::BAD_REQUEST,
            "Last-Event-Id is ahead of the server's event sequence",
        )
            .into_response());
    }
    let buffer = session.sse_replay_buffer.lock();
    let oldest_retained = buffer.front().map(|(id, _)| *id).unwrap_or(current_head);
    if last_id < oldest_retained.saturating_sub(1) {
        return Err((
            StatusCode::GONE,
            Json(rpc_error(
                None,
                JSONRPC_SERVER_ERROR,
                "resume gap; reconnect with fresh initialize",
            )),
        )
            .into_response());
    }
    Ok(buffer
        .iter()
        .filter(|(id, _)| *id > last_id)
        .cloned()
        .collect())
}

/// Handle `DELETE /mcp`: client-initiated session shutdown.
///
/// Authenticates the caller, parses the `Mcp-Session-Id` header,
/// removes the matching session from the registry under a single
/// write-lock acquisition, marks it closed, drains any pending
/// outbound oneshots so awaiting callers wake with
/// `TransportClosed`, emits
/// `klieo_mcp_session_deleted_total{reason="client_delete"}` for
/// operator visibility, and returns 204 No Content.
///
/// Status codes:
/// - `204 No Content` — session removed.
/// - `400 Bad Request` — missing or malformed `Mcp-Session-Id`.
/// - `401 Unauthorized` — authenticator rejected the caller.
/// - `404 Not Found` — id absent from the registry (never minted,
///   already deleted, or evicted by the idle reaper).
///
/// Idempotency: the registry remove + close marker are both
/// one-shot, so a repeated DELETE for the same id yields 404 rather
/// than re-running drain logic.
async fn delete_mcp(State(server): State<Arc<McpServer>>, headers: HeaderMap) -> Response {
    let identity = match enforce_authenticator_for_delete(&server, &headers).await {
        Ok(identity) => identity,
        Err(rejection) => return rejection,
    };
    let id = match extract_session_id(&headers) {
        Some(id) => id,
        None => {
            return (StatusCode::BAD_REQUEST, "missing or invalid Mcp-Session-Id").into_response();
        }
    };
    // Two-step lookup-then-remove: peek first under the read-lock so a
    // principal mismatch returns 404 without taking the registry
    // write-lock OR mutating state. Only proceed to the remove when
    // the caller's principal matches the binding established at
    // `initialize`. Mismatch yields 404 (same status as unknown id)
    // so a peer cannot probe the registry for existence — CWE-639.
    {
        let sessions = server.sessions.read().await;
        let Some(session) = sessions.get(&id) else {
            return StatusCode::NOT_FOUND.into_response();
        };
        if !principal_matches(identity.as_ref(), session.principal.as_deref()) {
            return StatusCode::NOT_FOUND.into_response();
        }
    }
    let session = {
        let mut sessions = server.sessions.write().await;
        sessions.remove(&id)
    };
    let Some(session) = session else {
        return StatusCode::NOT_FOUND.into_response();
    };
    let principal = session.principal.clone();
    session.close_and_drain().await;
    server.decrement_principal_count(principal.as_deref()).await;
    metrics::counter!(
        "klieo_mcp_session_deleted_total",
        "reason" => "client_delete"
    )
    .increment(1);
    tracing::info!(
        target: "klieo::mcp::session",
        session_id = %id,
        "session deleted by client"
    );
    StatusCode::NO_CONTENT.into_response()
}

/// Spawn the disconnect-cleanup task tied to the SSE body.
///
/// The task waits on a `oneshot::Receiver` that fires when
/// [`GuardedSseStream`] is dropped (TCP close, axum shutdown, client
/// disconnect). On wake it looks the session up in the
/// [`McpServer::sessions`] registry by `session_id` — the entry may
/// already be gone (DELETE /mcp, idle reaper) and the task tolerates
/// that as a no-op. When present, the session is marked closed, any
/// pending outbound oneshots drain as
/// [`klieo_core::ServerOutboundError::TransportClosed`] so awaiting
/// callers wake immediately, and the entry is removed from the
/// registry so the slot is reclaimed for a fresh `initialize`. The
/// `Sender` side travels with the guard wrapping the response body,
/// so the task can only fire after the body is gone. ADR-028.
fn spawn_session_cleanup(
    server: Arc<McpServer>,
    session_id: uuid::Uuid,
    cleanup_rx: tokio::sync::oneshot::Receiver<()>,
) {
    tokio::spawn(async move {
        let _ = cleanup_rx.await;
        let session = {
            let mut sessions = server.sessions.write().await;
            sessions.remove(&session_id)
        };
        let Some(session) = session else {
            return;
        };
        let principal = session.principal.clone();
        session.close_and_drain().await;
        server.decrement_principal_count(principal.as_deref()).await;
    });
}

/// Wire the per-session outbound primitive (and the roots cache, when
/// the server declared `sampling`) on top of the freshly-minted GET
/// mpsc tx. Mirrors stdio's `ensure_outbound_and_roots`. The caller
/// (`get_mcp`) passes the [`crate::session::Session`] resolved by
/// `require_session`, so the outbound + roots primitives land on the
/// exact session this GET is wiring. Returns `Err(500)` via
/// [`mint_session_race_500`] only if the outbound `OnceCell` was
/// already populated — a server invariant violation (concurrent GET
/// past `require_session` for the same id), not a benign client
/// conflict. The `roots_cache` set is best-effort; a stale cache on
/// the same `Session` survives by design.
#[allow(clippy::result_large_err)]
fn wire_session_outbound(
    server: &Arc<McpServer>,
    session: &Arc<crate::session::Session>,
    tx: crate::outbound_ring::RingSender<(u64, std::sync::Arc<serde_json::Value>)>,
) -> Result<(), Response> {
    use crate::outbound::OutboundRequests;
    use crate::outbound_sink::HttpFrameSink;
    use klieo_core::ServerOutbound;

    let sink: Arc<dyn crate::OutboundFrameSink> = Arc::new(HttpFrameSink::new(
        Arc::downgrade(session),
        tx,
        server.sse_replay_capacity,
    ));
    let outbound = Arc::new(OutboundRequests::new(sink));

    if session.outbound.set(outbound.clone()).is_err() {
        return Err(mint_session_race_500(None, "session.outbound", session.id));
    }
    if server.declare_sampling {
        let as_trait: Arc<dyn ServerOutbound> = outbound.clone();
        let _ = session
            .roots_cache
            .set(Arc::new(crate::roots::RootsCache::new(as_trait)));
    }
    Ok(())
}

/// Encodes one SSE frame into a `Bytes` buffer: `id: <N>\ndata: <json>\n\n`.
///
/// Writes the framing prefix, then streams the JSON payload directly
/// into the same [`bytes::BytesMut`] via [`serde_json::to_writer`],
/// then appends the SSE record terminator. The initial capacity is
/// `SSE_FRAME_HINT_BYTES`; oversized payloads grow the buffer once.
///
/// Returns `None` and logs at `target = "klieo::mcp::sse"` with
/// `event_id` + `session_id` when the JSON payload fails to serialise,
/// so the caller can skip the frame without aborting the stream.
pub fn encode_sse_frame(
    event_id: u64,
    frame: &serde_json::Value,
    session_id: uuid::Uuid,
) -> Option<Bytes> {
    use bytes::BufMut as _;
    use std::fmt::Write as _;
    let mut buf = bytes::BytesMut::with_capacity(SSE_FRAME_HINT_BYTES);
    write!(&mut buf, "id: {event_id}\ndata: ").expect("BytesMut write is infallible");
    match serde_json::to_writer((&mut buf).writer(), frame) {
        Ok(()) => {
            buf.extend_from_slice(b"\n\n");
            Some(buf.freeze())
        }
        Err(err) => {
            tracing::error!(
                target: "klieo::mcp::sse",
                event_id,
                session_id = %session_id,
                error = %err,
                "outbound frame serialisation failed; skipping id",
            );
            None
        }
    }
}

/// Wrap the per-session outbound receiver in an SSE-framed body and
/// build the `text/event-stream` response. `replay` carries the
/// snapshot of buffered `(event_id, frame)` pairs taken by
/// [`compute_replay_window`] before the live ring is attached;
/// fresh streams (no `Last-Event-Id` header) pass `Vec::new()`. The
/// body emits every replay frame in order, then attaches to the
/// live ring receiver.
///
/// Each yielded frame carries the monotonic per-session `event_id`
/// allocated by [`crate::outbound_sink::HttpFrameSink`] as the SSE
/// `id:` field, followed by the JSON-RPC payload as a single-line
/// `data:` field and the SSE record terminator (blank line).
/// Compliant clients track `id:` automatically and replay it via
/// `Last-Event-Id` on reconnect; non-tracking clients see the
/// `data:` payload unchanged. JSON serialisation failures log the
/// offending `event_id` + `session_id` at error severity (target
/// `klieo::mcp::sse`) and skip the frame rather than terminate the
/// stream — the next frame still gets through.
///
/// The stream is wrapped in a [`GuardedSseStream`] so dropping the
/// body (TCP close, axum shutdown, client disconnect) fires
/// `cleanup_tx`, waking the session-cleanup task that drains pending
/// outbound oneshots. Constructing the response cannot fail in
/// practice (status, headers, and body type are all well-formed); the
/// fallback path exists only to keep the handler total.
fn build_outbound_sse_response(
    session_id: uuid::Uuid,
    replay: Vec<(u64, std::sync::Arc<serde_json::Value>)>,
    mut rx: crate::outbound_ring::RingReceiver<(u64, std::sync::Arc<serde_json::Value>)>,
    cleanup_tx: tokio::sync::oneshot::Sender<()>,
) -> Response {
    let stream = async_stream::stream! {
        for (event_id, value) in replay {
            if let Some(frame) = encode_sse_frame(event_id, &value, session_id) {
                yield Ok::<Bytes, std::io::Error>(frame);
            }
        }
        while let Some((event_id, value)) = rx.recv().await {
            if let Some(frame) = encode_sse_frame(event_id, &value, session_id) {
                yield Ok::<Bytes, std::io::Error>(frame);
            }
        }
    };
    let guarded = GuardedSseStream::new(Box::pin(stream), cleanup_tx);
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/event-stream")
        .header(header::CACHE_CONTROL, "no-cache")
        .header(MCP_SESSION_ID_HEADER, session_id.to_string())
        .body(axum::body::Body::from_stream(guarded))
        .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
}

/// Wraps the outbound SSE body stream so that when the stream is
/// dropped (TCP close, axum shutdown, client disconnect) the server
/// clears the per-session state for the session this stream serves.
/// On drop the contained [`tokio::sync::oneshot::Sender`] fires,
/// waking the cleanup task spawned by [`spawn_session_cleanup`],
/// which removes the session from the registry, marks it closed,
/// and drains pending outbound oneshots. ADR-028.
struct GuardedSseStream<S> {
    inner: S,
    cleanup: Option<tokio::sync::oneshot::Sender<()>>,
}

impl<S> GuardedSseStream<S> {
    fn new(inner: S, cleanup: tokio::sync::oneshot::Sender<()>) -> Self {
        Self {
            inner,
            cleanup: Some(cleanup),
        }
    }
}

impl<S: Stream + Unpin> Stream for GuardedSseStream<S> {
    type Item = S::Item;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        Pin::new(&mut self.inner).poll_next(cx)
    }
}

impl<S> Drop for GuardedSseStream<S> {
    fn drop(&mut self) {
        if let Some(sender) = self.cleanup.take() {
            // Receiver may already be gone if the cleanup task was
            // cancelled (server shutdown); the send-error is the
            // intended no-op signal in that case.
            let _ = sender.send(());
        }
    }
}

/// Apply the configured [`klieo_auth_common::Authenticator`] to the
/// inbound request. Returns the verified [`Identity`] on success
/// (or `None` when no authenticator is wired), or a JSON-RPC
/// `-32001` envelope on failure. Authentication failure and
/// authorisation failure both yield the same `-32001` envelope
/// shape; the differing branches are logged so operators can
/// distinguish them in `tracing` output without leaking the
/// distinction onto the wire.
///
/// The `Identity` is returned (rather than consumed inline) so
/// the SSE upgrade path (`stream_tools_call` / `stream_resume`)
/// can thread it into the cluster-0.22 tenant-binding gate.
async fn enforce_authenticator(
    server: &Arc<McpServer>,
    headers: &HeaderMap,
    body: &Bytes,
    raw: &serde_json::Value,
) -> Result<Option<Identity>, Response> {
    let Some(auth) = server.authenticator() else {
        return Ok(None);
    };
    let adapter = AxumHeaders(headers);
    let identity = match auth.authenticate(&adapter, body).await {
        Ok(identity) => identity,
        Err(e) => {
            warn!(target: "mcp.auth", error = ?e, "authenticate failed");
            return Err(unauthenticated_response(raw, "authentication required"));
        }
    };
    let method = raw.get("method").and_then(|m| m.as_str()).unwrap_or("");
    if let Err(e) = auth.authorize_method(&identity, method).await {
        warn!(target: "mcp.auth", method = %method, error = ?e, "authorize_method failed");
        return Err(unauthenticated_response(
            raw,
            "method not authorized for principal",
        ));
    }
    Ok(Some(identity))
}

fn unauthenticated_response(raw: &serde_json::Value, message: &str) -> Response {
    (
        StatusCode::OK,
        Json(rpc_error(
            raw.get("id").cloned(),
            JSONRPC_UNAUTHENTICATED,
            message,
        )),
    )
        .into_response()
}

/// Apply the configured authenticator to a `GET /mcp` request. Only
/// the credential check (`authenticate`) runs; the POST path's
/// per-method scope check (`authorize_method`) is skipped because
/// GET carries no JSON-RPC `method` slot and the observer's scope
/// is granted transitively from the `initialize` POST that minted
/// the session id (the session-id requirement on subsequent calls
/// binds the GET observer to that same principal). Returns `None`
/// when no authenticator is wired, the verified [`Identity`] when
/// it succeeds, or a 401 response with no JSON-RPC envelope on
/// failure. ADR-028.
#[allow(clippy::result_large_err)]
async fn enforce_authenticator_for_get(
    server: &Arc<McpServer>,
    headers: &HeaderMap,
) -> Result<Option<Identity>, Response> {
    let Some(auth) = server.authenticator() else {
        return Ok(None);
    };
    let adapter = AxumHeaders(headers);
    match auth.authenticate(&adapter, &Bytes::new()).await {
        Ok(identity) => Ok(Some(identity)),
        Err(e) => {
            warn!(target: "mcp.auth", error = ?e, "authenticate failed on GET /mcp");
            Err(unauthenticated_response_for_get())
        }
    }
}

/// Build the 401 response returned from `enforce_authenticator_for_get`.
/// Plain status + short reason — GET /mcp has no JSON-RPC envelope to
/// embed an error in, unlike POST. The body is fixed text to avoid
/// echoing any caller-supplied data.
fn unauthenticated_response_for_get() -> Response {
    (StatusCode::UNAUTHORIZED, "authentication required").into_response()
}

/// Authentication gate for `DELETE /mcp`. Mirrors the GET dispatch:
/// runs only the `authenticate` half (no JSON-RPC body, no `method`
/// slot to `authorize_method` against), and yields `Err(rejection)`
/// when the configured authenticator rejects the request — otherwise
/// `Ok(_)`. Kept as a thin wrapper around
/// [`enforce_authenticator_for_get`] so a future divergence (e.g. a
/// distinct DELETE scope policy) has a single edit point.
#[allow(clippy::result_large_err)]
async fn enforce_authenticator_for_delete(
    server: &Arc<McpServer>,
    headers: &HeaderMap,
) -> Result<Option<Identity>, Response> {
    enforce_authenticator_for_get(server, headers).await
}

/// Adapter that exposes axum's [`HeaderMap`] as a
/// [`klieo_auth_common::Headers`] bag so the shared `Authenticator`
/// trait can read the same case-insensitive header surface on both
/// transports. `HeaderName`'s internal representation is already
/// lowercase per RFC 7230 §3.2, so we lowercase the lookup name and
/// pass it through `HeaderMap::get` (which does its own case-
/// insensitive match).
struct AxumHeaders<'a>(&'a HeaderMap);

impl<'a> klieo_auth_common::Headers for AxumHeaders<'a> {
    fn get(&self, name: &str) -> Option<&str> {
        self.0.get(name).and_then(|v| v.to_str().ok())
    }
}

/// Parse the `Mcp-Session-Id` header into a UUID. Returns `None` when
/// the header is absent, non-UTF-8, or not a valid UUID — callers map
/// each to the appropriate HTTP status (400 Bad Request).
fn extract_session_id(headers: &HeaderMap) -> Option<uuid::Uuid> {
    let raw = headers.get(MCP_SESSION_ID_HEADER)?.to_str().ok()?;
    uuid::Uuid::parse_str(raw).ok()
}

/// Parse the `Mcp-Session-Id` header, look up the corresponding
/// session in the registry, and verify the caller's principal matches
/// the binding established at `initialize`.
///
/// `caller` is the [`Identity`] returned by the configured
/// authenticator on this request (`None` when no authenticator is
/// wired — in that mode every session's stored principal is also
/// `None` and the equality check passes trivially).
///
/// A principal mismatch returns the same 404 envelope as an unknown
/// session id: a peer with a stolen session UUID must not be able to
/// distinguish "id never minted" from "id minted by someone else"
/// (CWE-639 IDOR existence-leakage).
///
/// Returns:
/// - `Err((400, "missing or invalid Mcp-Session-Id"))` on missing
///   or malformed header.
/// - `Err((404, JSONRPC_SERVER_ERROR + "unknown session id"))` when
///   the id is not in the registry OR the caller's principal does
///   not match the session's recorded principal.
/// - `Ok(Arc<Session>)` on a successful lookup with matching
///   principal.
async fn require_session(
    headers: &HeaderMap,
    server: &McpServer,
    caller: Option<&Identity>,
) -> Result<std::sync::Arc<crate::session::Session>, Response> {
    let id = extract_session_id(headers).ok_or_else(|| {
        (StatusCode::BAD_REQUEST, "missing or invalid Mcp-Session-Id").into_response()
    })?;
    let unknown_session = || -> Response {
        (
            StatusCode::NOT_FOUND,
            Json(rpc_error(None, JSONRPC_SERVER_ERROR, "unknown session id")),
        )
            .into_response()
    };
    let sessions = server.sessions.read().await;
    let session = sessions.get(&id).cloned().ok_or_else(unknown_session)?;
    if !principal_matches(caller, session.principal.as_deref()) {
        return Err(unknown_session());
    }
    Ok(session)
}

/// Compare a caller's verified principal against the principal
/// recorded on a session at `initialize` time.
///
/// `caller` is `None` when no authenticator is wired on this request;
/// `session_principal` is `None` when the session was minted on a
/// server with no authenticator. Both-None passes (auth-disabled
/// deployment, backwards compatible). Any other combination requires
/// string equality on the principal value (typically the OAuth `sub`
/// claim returned by `Identity::as_str`).
fn principal_matches(caller: Option<&Identity>, session_principal: Option<&str>) -> bool {
    match (caller, session_principal) {
        (None, None) => true,
        (Some(identity), Some(stored)) => identity.as_str() == stored,
        _ => false,
    }
}

/// Handle the `initialize`-shaped POST: mint a fresh UUID v4, build
/// a [`crate::session::Session`] bound to the verified caller
/// principal, insert it atomically into the server's `sessions`
/// registry under a single write-lock acquisition (cap check + insert
/// co-located so two concurrent inserts cannot both pass the cap),
/// and return the dispatch result with the session id echoed in the
/// `Mcp-Session-Id` response header. Called from [`post_mcp`].
///
/// `caller` is the [`Identity`] verified by `enforce_authenticator`
/// for this request (`None` when no authenticator is wired). It is
/// recorded on the [`crate::session::Session`] and enforced by
/// `require_session` on every subsequent POST / GET / DELETE for the
/// session — a different principal presenting the same session UUID
/// is rejected as 404 (CWE-639).
///
/// Cap enforcement: when `sessions.len() >= server.max_sessions` the
/// request is rejected with `503 Service Unavailable` and a
/// `klieo_mcp_session_cap_rejected_total` counter increment for
/// operator visibility. Successful inserts lazy-start the idle reaper
/// task via [`ensure_idle_reaper`].
async fn handle_initialize_post(
    server: &Arc<McpServer>,
    raw: serde_json::Value,
    caller: Option<Identity>,
) -> Response {
    let session_id = uuid::Uuid::new_v4();
    let principal = caller.as_ref().map(|id| id.as_str().to_string());
    let session = std::sync::Arc::new(crate::session::Session::new_http(
        session_id,
        principal,
        server.server_start,
    ));

    {
        let mut sessions = server.sessions.write().await;
        let mut principal_counts = server.principal_counts.write().await;

        if sessions.len() >= server.max_sessions {
            drop(principal_counts);
            drop(sessions);
            tracing::warn!(
                target: "klieo::mcp::session",
                scope = "global",
                cap = server.max_sessions,
                "session cap reached; rejecting initialize"
            );
            metrics::counter!(
                "klieo_mcp_session_cap_rejected_total",
                "scope" => "global"
            )
            .increment(1);
            return (
                StatusCode::SERVICE_UNAVAILABLE,
                Json(rpc_error(
                    raw.get("id").cloned(),
                    JSONRPC_SERVER_ERROR,
                    "session cap reached",
                )),
            )
                .into_response();
        }

        if let Some(p) = session.principal.as_deref() {
            let current = principal_counts.get(p).copied().unwrap_or(0);
            if current >= server.max_sessions_per_principal {
                drop(principal_counts);
                drop(sessions);
                tracing::warn!(
                    target: "klieo::mcp::session",
                    scope = "per_principal",
                    principal = p,
                    cap = server.max_sessions_per_principal,
                    "per-principal session cap reached; rejecting initialize"
                );
                metrics::counter!(
                    "klieo_mcp_session_cap_rejected_total",
                    "scope" => "per_principal"
                )
                .increment(1);
                return (
                    StatusCode::SERVICE_UNAVAILABLE,
                    Json(rpc_error(
                        raw.get("id").cloned(),
                        JSONRPC_SERVER_ERROR,
                        "per-principal session cap reached",
                    )),
                )
                    .into_response();
            }
            *principal_counts.entry(p.to_string()).or_insert(0) += 1;
        }

        sessions.insert(session_id, session.clone());
    }

    ensure_idle_reaper(server).await;

    // Successful initialize counts as the session's first activity:
    // touch the per-session clock before the OK envelope leaves so
    // the idle reaper does not race the response and evict an
    // otherwise-live session.
    let resp = dispatch(server, raw, Some(&session)).await;
    touch_last_activity(server, &session);
    (
        StatusCode::OK,
        [(MCP_SESSION_ID_HEADER, session_id.to_string())],
        Json(resp),
    )
        .into_response()
}

/// Idempotently spawn the idle-session reaper task for this server.
/// The first call populates [`McpServer::idle_reaper_started`] and
/// spawns [`idle_reaper_loop`]; subsequent calls are no-ops, so
/// duplicate `initialize` POSTs do not spawn duplicate reapers. The
/// spawned task holds an `Arc<McpServer>` and exits cleanly when the
/// runtime drops it (typically server shutdown).
async fn ensure_idle_reaper(server: &Arc<McpServer>) {
    let server_for_task = server.clone();
    let _ = server
        .idle_reaper_started
        .get_or_init(|| async move {
            tokio::spawn(idle_reaper_loop(server_for_task));
        })
        .await;
}

/// Drive periodic idle-session reaping for the HTTP session
/// registry. Holds an `Arc<McpServer>` for the task's lifetime and
/// runs until the runtime drops the spawned task (typically server
/// shutdown).
///
/// Wakes on the cadence stored in [`McpServer::idle_reaper_tick`]
/// (defaults to 10 seconds; configurable via
/// [`McpServerBuilder::with_idle_reaper_tick`] in test builds). On
/// each tick, when `session_idle_timeout` is non-zero, snapshots
/// `(id, Arc<Session>)` pairs under a read-lock then drops the
/// lock; per-session probes of `last_activity_millis` run lock-free
/// outside the registry lock so writers (initialize POST, DELETE,
/// SSE disconnect) never queue behind a scan. Entries that probe
/// idle are then removed under a write-lock.
/// After the registry lock is released, each evicted session is
/// marked closed, its pending outbound oneshots are drained as
/// `TransportClosed`, a
/// `klieo_mcp_session_deleted_total{reason="idle_timeout"}` counter
/// is incremented, and an info-level eviction record is emitted.
///
/// `MissedTickBehavior::Skip` prevents tick accumulation if the
/// runtime stalls under load; missed ticks coalesce into a single
/// scan. A `session_idle_timeout` of `Duration::ZERO` disables the
/// scan entirely while leaving the task alive.
async fn idle_reaper_loop(server: Arc<McpServer>) {
    let tick = server.idle_reaper_tick;
    let mut interval = tokio::time::interval(tick);
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
    loop {
        interval.tick().await;
        let timeout = server.session_idle_timeout;
        if timeout.is_zero() {
            continue;
        }
        // Snapshot Arcs under the read-lock; probe `last_activity_millis`
        // lock-free outside it so registry writers never queue behind a
        // scan.
        let snapshot: Vec<(uuid::Uuid, Arc<crate::session::Session>)> = {
            let sessions = server.sessions.read().await;
            sessions.iter().map(|(id, s)| (*id, s.clone())).collect()
        };

        let now_millis = server.server_start.elapsed().as_millis() as u64;
        let timeout_millis = timeout.as_millis() as u64;
        let to_evict: Vec<(uuid::Uuid, Arc<crate::session::Session>)> = snapshot
            .into_iter()
            .filter(|(_, session)| {
                let last = session
                    .last_activity_millis
                    .load(std::sync::atomic::Ordering::Relaxed);
                now_millis.saturating_sub(last) > timeout_millis
            })
            .collect();

        if to_evict.is_empty() {
            continue;
        }

        {
            let mut sessions = server.sessions.write().await;
            for (id, _) in &to_evict {
                sessions.remove(id);
            }
        }

        for (id, session) in to_evict {
            let principal = session.principal.clone();
            session.close_and_drain().await;
            server.decrement_principal_count(principal.as_deref()).await;
            metrics::counter!(
                "klieo_mcp_session_deleted_total",
                "reason" => "idle_timeout"
            )
            .increment(1);
            tracing::info!(
                target: "klieo::mcp::session",
                session_id = %id,
                "session evicted by idle reaper"
            );
        }
    }
}

/// Route `tools/call` (SSE upgrade) and `klieo/tools/resume` (replay-only
/// SSE) out of the normal JSON dispatch path. Returns `Some(Response)` when
/// either branch fires, `None` when the request should continue to
/// `dispatch`.
///
/// `identity` is the [`Identity`] verified by `enforce_authenticator` —
/// `None` when no authenticator is wired. Threaded through to the
/// tenant-binding gate in both SSE paths (ADR-022).
#[instrument(skip_all, level = "debug")]
async fn try_sse_upgrade(
    server: Arc<McpServer>,
    raw: serde_json::Value,
    identity: Option<Identity>,
) -> Option<Response> {
    let method = raw.get("method").and_then(|m| m.as_str())?;
    match method {
        "tools/call" => {
            let token_val = raw.pointer("/params/_meta/progressToken")?;
            match token_val {
                serde_json::Value::String(_) | serde_json::Value::Number(_) => {
                    Some(stream_tools_call(server, raw.clone(), token_val.clone(), identity).await)
                }
                _ => {
                    let id = raw.get("id").cloned();
                    Some(
                        (
                            StatusCode::BAD_REQUEST,
                            Json(rpc_error(
                                id,
                                JSONRPC_INVALID_PARAMS,
                                "_meta.progressToken must be a string or number",
                            )),
                        )
                            .into_response(),
                    )
                }
            }
        }
        "klieo/run/resume" => Some(handle_run_resume(server, raw, identity).await),
        "klieo/tools/resume" => {
            let params = match raw.get("params").cloned() {
                Some(p) => p,
                None => {
                    let id = raw.get("id").cloned();
                    return Some(
                        (
                            StatusCode::BAD_REQUEST,
                            Json(rpc_error(
                                id,
                                JSONRPC_INVALID_PARAMS,
                                "klieo/tools/resume: missing params",
                            )),
                        )
                            .into_response(),
                    );
                }
            };
            let parsed: ResumeParams = match serde_json::from_value(params) {
                Ok(p) => p,
                Err(e) => {
                    let id = raw.get("id").cloned();
                    return Some(
                        (
                            StatusCode::BAD_REQUEST,
                            Json(rpc_error(
                                id,
                                JSONRPC_INVALID_PARAMS,
                                &format!("klieo/tools/resume: invalid params: {e}"),
                            )),
                        )
                            .into_response(),
                    );
                }
            };
            Some(stream_resume(server, raw, parsed, identity).await)
        }
        _ => None,
    }
}

#[derive(serde::Deserialize)]
struct ResumeParams {
    #[serde(rename = "progressToken")]
    progress_token: serde_json::Value,
    #[serde(rename = "lastEventId")]
    last_event_id: u64,
}

/// Wire shape of `klieo/run/resume` params. The ticket is the only
/// opaque handle the peer holds; the decision carries the operator's
/// approve/reject verdict plus an optional rejection reason that
/// gets fed back to the model on resume (ADR-045).
#[derive(serde::Deserialize)]
struct RunResumeParams {
    ticket: String,
    decision: RunResumeDecision,
}

#[derive(serde::Deserialize)]
struct RunResumeDecision {
    approved: bool,
    #[serde(default)]
    reason: Option<String>,
}

/// Stable wire message used by `klieo/run/resume` for every
/// fail-closed branch: unknown ticket, foreign-principal ticket, race-
/// loser ticket. Identical bytes in all three cases so the peer cannot
/// distinguish them (IDOR mitigation, CWE-639). The reason a request
/// failed lives in the server log only.
const RUN_RESUME_DENY_MESSAGE: &str = "resume ticket invalid";

/// Stable wire message used when the server has no checkpoint KV
/// wired or no workflow registered under the ticket's workflow name —
/// caller-side config, not a tenant secret. Distinct from the deny
/// shape because the peer cannot derive ticket presence from it.
const RUN_RESUME_UNAVAILABLE_MESSAGE: &str = "resume unavailable";

/// JSON-RPC handler for `klieo/run/resume`. Every error branch
/// surfaces the sanitised stable string only (CWE-209); the internal
/// cause lives in the server log.
async fn handle_run_resume(
    server: Arc<McpServer>,
    raw: serde_json::Value,
    identity: Option<Identity>,
) -> Response {
    let req_id = raw.get("id").cloned();
    let parsed = match parse_run_resume_params(&raw, req_id.clone()) {
        Ok(p) => p,
        Err(resp) => return resp,
    };
    let claimed = match claim_resume_record(&server, &parsed.ticket, identity, req_id.clone()).await
    {
        Ok(rec) => rec,
        Err(resp) => return resp,
    };
    drive_resume(&server, parsed.decision, claimed, req_id).await
}

#[allow(clippy::result_large_err)]
fn parse_run_resume_params(
    raw: &serde_json::Value,
    req_id: Option<serde_json::Value>,
) -> Result<RunResumeParams, Response> {
    let Some(params_value) = raw.get("params").cloned() else {
        return Err(run_resume_invalid_params(req_id, "missing params"));
    };
    serde_json::from_value::<RunResumeParams>(params_value)
        .map_err(|_| run_resume_invalid_params(req_id, "invalid params"))
}

/// Run the lookup-then-authz-then-claim sequence. Authz runs BEFORE
/// the atomic claim so a foreign principal cannot ever consume
/// another caller's ticket; the same opaque deny string surfaces for
/// every failure branch so the peer cannot tell unknown from
/// foreign-owned.
#[allow(clippy::result_large_err)]
async fn claim_resume_record(
    server: &McpServer,
    ticket: &str,
    identity: Option<Identity>,
    req_id: Option<serde_json::Value>,
) -> Result<crate::resume_ticket::ResumeTicketRecord, Response> {
    let Some(store) = server.resume_ticket_store.as_ref() else {
        return Err(run_resume_unavailable(req_id));
    };
    let caller = identity
        .as_ref()
        .filter(|id| !id.is_anonymous())
        .ok_or_else(|| run_resume_denied(req_id.clone(), "anonymous caller"))?;

    let peeked = store
        .peek(ticket)
        .await
        .map_err(|err| log_and_deny(err, req_id.clone(), "peek failure"))?;
    let record = peeked.ok_or_else(|| run_resume_denied(req_id.clone(), "unknown ticket"))?;
    if caller.as_str() != record.principal {
        return Err(run_resume_denied(req_id, "principal mismatch"));
    }
    let claimed = store
        .claim(ticket)
        .await
        .map_err(|err| log_and_deny(err, req_id.clone(), "claim failure"))?;
    claimed.ok_or_else(|| run_resume_denied(req_id, "claim lost race"))
}

async fn drive_resume(
    server: &McpServer,
    decision: RunResumeDecision,
    record: crate::resume_ticket::ResumeTicketRecord,
    req_id: Option<serde_json::Value>,
) -> Response {
    let Some(handle) = server.workflow_resume_handles.get(&record.workflow_name) else {
        tracing::warn!(
            target: "klieo.mcp.resume",
            workflow = %record.workflow_name,
            "claimed ticket references an unregistered workflow",
        );
        return run_resume_unavailable(req_id);
    };
    let approval = if decision.approved {
        klieo_core::checkpoint::ApprovalDecision::Approved
    } else {
        klieo_core::checkpoint::ApprovalDecision::Rejected {
            reason: decision.reason.unwrap_or_default(),
        }
    };
    // Bind the resumed run to the ticket's principal (hashed, non-PII) so
    // the continuation is attributed + governed like the original call —
    // resume must not bypass the inbound per-tenant LLM budget.
    let tenant_label = klieo_core::principal_hash(&record.principal);
    match handle
        .resume(record.checkpoint, approval, tenant_label)
        .await
    {
        Ok(result) => (StatusCode::OK, Json(crate::rpc_ok(req_id, result))).into_response(),
        Err(err) => {
            tracing::warn!(
                target: "klieo.mcp.resume",
                workflow = %record.workflow_name,
                error = %err,
                "workflow resume failed",
            );
            run_resume_server_error(req_id)
        }
    }
}

fn log_and_deny(
    err: crate::resume_ticket::TicketStoreError,
    req_id: Option<serde_json::Value>,
    log_reason: &str,
) -> Response {
    tracing::warn!(
        target: "klieo.mcp.resume",
        rpc_id = ?req_id,
        error = %err,
        reason = log_reason,
        "resume ticket-store op failed; denying fail-closed",
    );
    run_resume_denied(req_id, log_reason)
}

fn run_resume_denied(id: Option<serde_json::Value>, log_reason: &str) -> Response {
    tracing::info!(
        target: "klieo.mcp.resume",
        rpc_id = ?id,
        reason = log_reason,
        "klieo/run/resume denied",
    );
    (
        StatusCode::OK,
        Json(rpc_error(id, JSONRPC_INVALID_PARAMS, RUN_RESUME_DENY_MESSAGE)),
    )
        .into_response()
}

fn run_resume_unavailable(id: Option<serde_json::Value>) -> Response {
    (
        StatusCode::OK,
        Json(rpc_error(
            id,
            JSONRPC_SERVER_ERROR,
            RUN_RESUME_UNAVAILABLE_MESSAGE,
        )),
    )
        .into_response()
}

fn run_resume_invalid_params(id: Option<serde_json::Value>, log_reason: &str) -> Response {
    tracing::warn!(
        target: "klieo.mcp.resume",
        rpc_id = ?id,
        reason = log_reason,
        "klieo/run/resume params invalid",
    );
    (
        StatusCode::BAD_REQUEST,
        Json(rpc_error(id, JSONRPC_INVALID_PARAMS, RUN_RESUME_DENY_MESSAGE)),
    )
        .into_response()
}

fn run_resume_server_error(id: Option<serde_json::Value>) -> Response {
    (
        StatusCode::OK,
        Json(rpc_error(id, JSONRPC_SERVER_ERROR, "resume execution failed")),
    )
        .into_response()
}

/// Cluster-0.23 helper: copy the W3C tracecontext headers from an
/// axum [`HeaderMap`] into a [`klieo_core::Headers`] bag so
/// [`klieo_core::extract_traceparent`] can lift the upstream
/// `opentelemetry::Context` and parent the entry span under it.
/// Only `traceparent` + `tracestate` are forwarded — everything else
/// stays on the inbound `HeaderMap` for the dispatch path.
fn klieo_headers_from_axum(headers: &HeaderMap) -> klieo_core::Headers {
    let mut out = klieo_core::Headers::default();
    if let Some(value) = headers.get("traceparent").and_then(|v| v.to_str().ok()) {
        out.insert("traceparent".into(), value.to_string());
    }
    if let Some(value) = headers.get("tracestate").and_then(|v| v.to_str().ok()) {
        out.insert("tracestate".into(), value.to_string());
    }
    out
}

fn content_type_is_json(headers: &HeaderMap) -> bool {
    headers
        .get(header::CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .map(|s| {
            s.split(';')
                .next()
                .unwrap_or("")
                .trim()
                .eq_ignore_ascii_case("application/json")
        })
        .unwrap_or(false)
}

/// Broadcast channel capacity: enough headroom for a burst of events
/// before the SSE loop drains them. Lagged receivers emit a single
/// synthetic `lagged` notification rather than hard-failing.
const PROGRESS_CHANNEL_CAP: usize = 64;

/// Stringify a `progressToken` JSON value for use as a `ResumeBuffer` stream id.
fn progress_token_to_string(token: &serde_json::Value) -> String {
    match token {
        serde_json::Value::String(s) => s.clone(),
        serde_json::Value::Number(n) => n.to_string(),
        other => other.to_string(),
    }
}

/// Prepend `id: {id}\n` to an existing SSE frame so reconnecting clients
/// can supply `Last-Event-ID`.
fn id_prefixed_frame(id: u64, frame_bytes: Bytes) -> Bytes {
    let mut prefixed = format!("id: {id}\n").into_bytes();
    prefixed.extend_from_slice(&frame_bytes);
    Bytes::from(prefixed)
}

/// Extract the `id: N` line from a raw SSE frame payload (as written by
/// `id_prefixed_frame`). Returns `None` if the prefix is missing or
/// the value is not a u64.
fn parse_id_prefix(bytes: &Bytes) -> Option<u64> {
    let text = std::str::from_utf8(bytes).ok()?;
    let id_line = text.lines().next()?;
    id_line.strip_prefix("id: ")?.parse::<u64>().ok()
}

/// Spawn a best-effort write-through to the resume buffer. A record
/// failure is logged but never propagates to the SSE caller. When
/// `close_after` is `true`, `buffer.close()` is called after the record
/// succeeds (or fails) so the stream is marked terminal.
fn spawn_record(
    buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    stream_id: String,
    id: u64,
    payload: Bytes,
    close_after: bool,
) {
    tokio::spawn(async move {
        if let Err(e) = buffer.record(&stream_id, id, payload.clone()).await {
            tracing::warn!(
                target: "mcp.resume",
                stream_id,
                id,
                error = %e,
                "resume buffer record failed"
            );
        }
        if close_after {
            if let Err(e) = buffer.close(&stream_id).await {
                tracing::warn!(
                    target: "mcp.resume",
                    stream_id,
                    error = %e,
                    "resume buffer close failed"
                );
            }
        }
    });
}

/// Spawn a best-effort cross-replica publish of one SSE frame to the
/// per-progressToken bus subject, bounded by the server's shared
/// publish-concurrency semaphore. Saturation drops the publish + emits
/// a `warn`; cross-replica subscribers fall back to the resume buffer
/// for gap-fill on the next `klieo/tools/resume` request. Failures
/// inside the spawned task log + continue. Thin wrapper over
/// [`klieo_core::cancel::spawn_publish_bounded`] so the same
/// concurrency cap covers both this hot path and the drop-time cancel
/// publish in [`CancelOnDrop`].
fn spawn_publish(
    pubsub: Arc<dyn klieo_core::Pubsub>,
    subject: String,
    payload: Bytes,
    permits: Arc<tokio::sync::Semaphore>,
) {
    // Capture the publisher's OTEL context via the tracing →
    // opentelemetry bridge so the bus headers carry a W3C `traceparent`.
    // Subscribers on other replicas extract it and stitch their decode
    // span under this span (cluster 0.23, ADR-023).
    let mut trace_headers = klieo_core::Headers::default();
    klieo_core::inject_traceparent(&mut trace_headers, &tracing::Span::current().context());
    klieo_core::cancel::spawn_publish_bounded(
        pubsub,
        subject,
        payload,
        "mcp.fanout",
        permits,
        trace_headers,
    );
}

/// Wrap an SSE body stream with the cancel-fanout sandwich both
/// streaming paths require: `RegistryDeregisterOnDrop` to remove the
/// per-invoke cancel-registry entry when the stream ends, plus
/// `CancelOnDrop` to fire the request-scoped token and publish a
/// best-effort cross-replica cancel on body drop.
///
/// `leader_handle` is held inside `CancelOnDrop` for the lifetime of
/// the SSE body; on drop the held `klieo_core::LeaderHandle` releases
/// its KV claim so a follower's orphan probe observes a TTL-expired
/// entry (ADR-020). `None` when leader election is disabled.
///
/// `ownership_handle` is held inside `CancelOnDrop` for the lifetime
/// of the SSE body; on drop the held `klieo_core::OwnershipHandle`
/// fires a best-effort `kv.delete` so a follower's `lookup` no longer
/// returns the stale principal (ADR-022). `None` when tenant binding
/// is disabled, when the request is anonymous, or when the claim
/// itself failed.
///
/// Callers MUST have already registered `request_cancel` under
/// `stream_id` in `server.cancel_registry()`. The wrapper handles
/// deregistration on drop.
fn wrap_with_cancel_fanout<S>(
    server: &Arc<McpServer>,
    inner: S,
    stream_id: String,
    request_cancel: tokio_util::sync::CancellationToken,
    leader_handle: Option<klieo_core::LeaderHandle>,
    ownership_handle: Option<klieo_core::OwnershipHandle>,
) -> CancelOnDrop<klieo_core::cancel::RegistryDeregisterOnDrop<S>>
where
    S: Stream<Item = Result<Bytes, std::convert::Infallible>> + Send + Unpin + 'static,
{
    let cancel_subject = CANCEL_SUBJECT_PREFIX.to_string() + &stream_id;
    let deregistered = klieo_core::cancel::RegistryDeregisterOnDrop::new(
        inner,
        server.cancel_registry().clone(),
        stream_id,
    );
    CancelOnDrop {
        inner: deregistered,
        _guard: request_cancel.drop_guard(),
        pubsub: server.pubsub.clone(),
        cancel_subject,
        permits: server.publish_permits.clone(),
        _leader: leader_handle,
        _ownership: ownership_handle,
    }
}

/// Best-effort leader claim for a streaming invoke. Returns `None`
/// when no registry is wired (single-replica deployment) OR when the
/// KV `put` fails (fail-open per ADR-020). The returned
/// `LeaderHandle` MUST be held for the lifetime of the SSE response;
/// `CancelOnDrop::_leader` carries it to drop-time so the KV entry
/// is released and a follower's orphan probe observes TTL expiry on
/// replica failure.
///
/// `payload` carries the original JSON-RPC request body serialised
/// to bytes so a follower that detects an orphan can re-invoke an
/// idempotent tool from the cached payload (cluster 0.24).
/// `principal` records the authenticated subject (cluster 0.22) so
/// the re-invoke runs under the same tenant binding as the original.
async fn try_claim_leader(
    server: &Arc<McpServer>,
    stream_id: &str,
    payload: Option<Bytes>,
    principal: Option<String>,
) -> Option<klieo_core::LeaderHandle> {
    let registry = server.leader_registry()?;
    let key = format!("{MCP_LEADER_KEY_PREFIX}{stream_id}");
    match registry
        .claim_with_heartbeat(
            key,
            server.leader_ttl(),
            server.leader_heartbeat_interval(),
            payload,
            principal,
        )
        .await
    {
        Ok(handle) => Some(handle),
        Err(e) => {
            tracing::warn!(
                target: "mcp.leader",
                stream_id = %stream_id,
                error = %e,
                "leader claim failed; degrading to no-claim (orphan detection \
                 disabled for this stream)",
            );
            None
        }
    }
}

/// Tenant-binding claim for a streaming invoke. Returns `Ok(None)` (proceed,
/// no claim) when no `OwnershipRegistry` is wired or the caller is
/// unauthenticated/anonymous. On KV `put` failure a lenient registry returns
/// `Ok(None)` (fail-open per ADR-022), while a **strict** registry returns
/// `Err(Response)` so the invoke is denied rather than started unprotected. The
/// returned `OwnershipHandle` MUST be held for the lifetime of the SSE response;
/// `CancelOnDrop::_ownership` carries it to drop-time so the entry is removed on
/// stream end.
async fn try_claim_ownership(
    server: &Arc<McpServer>,
    stream_id: &str,
    identity: &Option<Identity>,
) -> Result<Option<klieo_core::OwnershipHandle>, Response> {
    let Some(registry) = server.ownership_registry() else {
        return Ok(None);
    };
    let Some(identity) = identity.as_ref() else {
        return Ok(None);
    };
    if identity.is_anonymous() {
        return Ok(None);
    }
    let key = format!("{MCP_OWNERSHIP_KEY_PREFIX}{stream_id}");
    let principal = identity.as_str().to_string();
    match registry.claim_guarded(key, principal).await {
        klieo_core::OwnershipClaim::Claimed(handle) => Ok(Some(handle)),
        klieo_core::OwnershipClaim::Proceed => Ok(None),
        klieo_core::OwnershipClaim::Unavailable => {
            Err(stream_unavailable_response(serde_json::Value::Null))
        }
        // Non_exhaustive: an unrecognised future outcome denies (fail-closed).
        _ => Err(stream_unavailable_response(serde_json::Value::Null)),
    }
}

/// `klieo/tools/resume` tenant-binding gate. Looks up the owner of
/// `stream_id` in the `klieo-tenants` bucket and matches against the
/// caller's principal.
///
/// - Owner match → `Ok(())` (proceed).
/// - Owner mismatch → `Err(stream_not_found_response)` — deny-as-
///   NotFound per OWASP IDOR best practice; same wire shape as a
///   nonexistent progressToken so the response leaks no existence
///   info.
/// - No registry / no caller / anonymous caller → `Ok(())`
///   (partial-deployment skip, ADR-022).
/// - Missing entry (`Ok(None)`) → `Ok(())` (legacy pre-0.22 stream
///   or no auth wired at invoke).
/// - KV lookup error → `Ok(())` (fail-open; same posture as the
///   leader is_alive probe in ADR-020).
async fn enforce_owner(
    server: &Arc<McpServer>,
    stream_id: &str,
    identity: &Option<Identity>,
    req_id: &serde_json::Value,
) -> Result<(), Response> {
    let Some(registry) = server.ownership_registry() else {
        return Ok(());
    };
    let Some(identity) = identity.as_ref() else {
        return Ok(());
    };
    if identity.is_anonymous() {
        return Ok(());
    }
    let key = format!("{MCP_OWNERSHIP_KEY_PREFIX}{stream_id}");
    match registry.check_owner(&key, identity.as_str()).await {
        klieo_core::OwnershipCheck::Allowed => Ok(()),
        klieo_core::OwnershipCheck::Denied => {
            tracing::warn!(
                target: "mcp.tenants",
                stream_id = %stream_id,
                principal = %identity.as_str(),
                "ownership mismatch on klieo/tools/resume; denying as stream-not-found",
            );
            Err(stream_not_found_response(req_id.clone()))
        }
        // Strict mode, store unreachable: `check_owner` already logged the
        // cause; deny rather than risk a cross-tenant resume.
        klieo_core::OwnershipCheck::Unavailable => Err(stream_unavailable_response(req_id.clone())),
        // Non_exhaustive: an unrecognised future verdict denies (fail-closed).
        _ => Err(stream_unavailable_response(req_id.clone())),
    }
}

/// Build the JSON-RPC envelope the resume gate returns on owner
/// mismatch. Wire-shape MUST match the `ResumeError::NotFound` arm
/// in [`stream_resume`] byte-for-byte (modulo request id) so the
/// response is indistinguishable from a request for a nonexistent
/// progressToken — IDOR mitigation per OWASP.
fn stream_not_found_response(req_id: serde_json::Value) -> Response {
    (
        StatusCode::OK,
        Json(rpc_error(
            Some(req_id),
            crate::JSONRPC_RESUME_BUFFER_NOT_FOUND,
            "no buffered stream for progressToken",
        )),
    )
        .into_response()
}

/// Build the envelope returned when **strict** tenant binding cannot reach the
/// ownership store. 503 + `JSONRPC_SERVER_ERROR` — a retryable infrastructure
/// failure, deliberately distinct from the leak-safe not-found used on an owner
/// mismatch: the gate fails closed rather than risk a cross-tenant resume.
fn stream_unavailable_response(req_id: serde_json::Value) -> Response {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(rpc_error(
            Some(req_id),
            crate::JSONRPC_SERVER_ERROR,
            "ownership store unavailable; denied (strict tenant binding)",
        )),
    )
        .into_response()
}

/// Outcome of a leader-alive probe at `stream_resume` entry.
enum LeaderProbe {
    /// No registry wired — single-replica deployment, orphan detection skipped.
    NoRegistry,
    /// Probe returned `Ok(true)` OR the KV errored (fail-open per ADR-020).
    Alive,
    /// Probe returned `Ok(false)` — the leader's claim has lapsed or never existed.
    Dead,
}

/// Best-effort leader-alive probe for a `stream_resume` entry.
///
/// Returns `LeaderProbe::NoRegistry` when leader election is not
/// wired; otherwise probes the KV. `Err(_)` from the registry
/// fails open (treat as alive) per ADR-020: a transient KV blip
/// must NOT trip an orphan write that would terminate a live
/// stream prematurely.
async fn probe_leader(server: &Arc<McpServer>, stream_id: &str) -> LeaderProbe {
    let Some(registry) = server.leader_registry() else {
        return LeaderProbe::NoRegistry;
    };
    let key = format!("{MCP_LEADER_KEY_PREFIX}{stream_id}");
    match registry.is_alive(&key).await {
        Ok(true) => LeaderProbe::Alive,
        Ok(false) => LeaderProbe::Dead,
        Err(e) => {
            tracing::warn!(
                target: "mcp.leader",
                stream_id = %stream_id,
                error = %e,
                "is_alive probe failed; treating as alive (fail-open per ADR-020)",
            );
            LeaderProbe::Alive
        }
    }
}

/// Probe a [`klieo_core::resume::ResumeBuffer`] for the highest
/// retained event id. Returns `None` when the buffer has no
/// retained events for `stream_id` (`NotFound`) OR when the
/// backend errors (best-effort: log warn + skip orphan write).
async fn max_event_id(
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    stream_id: &str,
) -> Option<u64> {
    let mut replay = match buffer.replay(stream_id, 0).await {
        Ok(stream) => stream,
        Err(klieo_core::resume::ResumeError::NotFound(_)) => return None,
        Err(e) => {
            tracing::warn!(
                target: "mcp.leader",
                stream_id = %stream_id,
                error = %e,
                "max_event_id replay failed; skipping orphan terminal write",
            );
            return None;
        }
    };
    let mut highest: Option<u64> = None;
    while let Some((id, _)) = tokio_stream::StreamExt::next(&mut replay).await {
        highest = Some(match highest {
            Some(current) => current.max(id),
            None => id,
        });
    }
    highest
}

/// JSON-RPC error envelope for the LEADER_DIED single-shot
/// response returned from `stream_resume` on orphan detection.
/// Carries `data.stream_id` so the client can correlate the orphan
/// signal with operator-side bus telemetry. Mirrors A2A's
/// `leader_died_envelope`. ADR-020.
fn leader_died_envelope(id: serde_json::Value, stream_id: &str) -> serde_json::Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": {
            "code": JSONRPC_LEADER_DIED,
            "message": "stream leader died",
            "data": { "stream_id": stream_id },
        }
    })
}

/// Build the terminal SSE-frame bytes recorded in the resume buffer
/// at `max_id + 1` when an orphan is detected on `stream_resume`.
/// Replay restamps with `id_prefixed_frame` on read, so the stored
/// payload is the bare `event: error\ndata: {...}\n\n` SSE block.
fn leader_died_sse_frame_bytes(stream_id: &str) -> Bytes {
    let envelope = serde_json::json!({
        "jsonrpc": "2.0",
        "id": serde_json::Value::Null,
        "error": {
            "code": JSONRPC_LEADER_DIED,
            "message": "stream leader died",
            "data": { "stream_id": stream_id },
        }
    });
    Bytes::from(format!("event: error\ndata: {}\n\n", envelope))
}

/// Orphan recovery: on a dead leader probe, write a terminal
/// "leader died" SSE frame at `max_id + 1` to the resume buffer and
/// mark the buffer terminal. Future resume clients replaying past
/// `max_id` see the synthetic frame and close.
///
/// Returns `true` when the orphan write landed (caller raises the
/// single-shot JSON-RPC `LEADER_DIED` envelope); `false` when:
/// - The buffer has no retained events (leader never claimed +
///   never wrote → not an orphan, just a non-streaming client).
///   Caller falls through to the regular replay path.
/// - The terminal-frame `record` failed (best-effort: log warn +
///   fall through to the regular subscribe path so a transient
///   KV blip does not surface as a hard-fail to the client).
async fn write_orphan_terminal_frame(
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    stream_id: &str,
) -> bool {
    let Some(max_id) = max_event_id(buffer, stream_id).await else {
        return false;
    };
    let next_id = max_id + 1;
    let frame = leader_died_sse_frame_bytes(stream_id);
    if let Err(e) = buffer.record(stream_id, next_id, frame).await {
        tracing::warn!(
            target: "mcp.leader",
            stream_id = %stream_id,
            next_id,
            error = %e,
            "orphan terminal record failed; skipping orphan write",
        );
        return false;
    }
    if let Err(e) = buffer.close(stream_id).await {
        tracing::warn!(
            target: "mcp.leader",
            stream_id = %stream_id,
            error = %e,
            "orphan terminal close failed; resume buffer may retain stale stream",
        );
    }
    tracing::error!(
        target: "mcp.leader",
        stream_id = %stream_id,
        next_id,
        "stream leader died; emitted LEADER_DIED terminal frame at max+1",
    );
    true
}

/// Terminate-only orphan response (cluster-0.20 path). Writes the
/// "leader died" SSE frame to the resume buffer at `max + 1` + marks
/// the buffer terminal, then returns the single-shot JSON-RPC
/// `LEADER_DIED` (-32099) envelope. When the resume buffer has no
/// retained events the orphan write is skipped (not an orphan, just
/// a no-stream resume client) and the caller falls through to the
/// regular replay path; the [`OrphanOutcome::Passthrough`] variant
/// signals that fall-through.
async fn terminate_orphan_mcp(
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    req_id: &serde_json::Value,
    stream_id: &str,
) -> OrphanOutcome {
    if write_orphan_terminal_frame(buffer, stream_id).await {
        let resp = (
            StatusCode::OK,
            Json(leader_died_envelope(req_id.clone(), stream_id)),
        )
            .into_response();
        return OrphanOutcome::Terminated(resp);
    }
    OrphanOutcome::Passthrough
}

/// Outcome of the cluster-0.24 follower-side orphan gate.
///
/// - [`OrphanOutcome::Reinvoked`] — the gate CAS-claimed a new leader
///   entry with `attempt + 1` and is driving a fresh `tools/call`
///   stream from the cached payload. The carried [`Response`] is the
///   SSE stream the resume client should consume.
/// - [`OrphanOutcome::Terminated`] — cluster-0.20 fall-back fired
///   (handler not idempotent, attempt cap reached, no cached payload,
///   CAS conflict, or no resume-buffer events to anchor on). The
///   carried [`Response`] is the single-shot LEADER_DIED envelope.
/// - [`OrphanOutcome::Passthrough`] — no leader registry wired OR
///   the orphan write found no buffer entries. Caller falls through
///   to the regular replay path.
#[non_exhaustive]
pub enum OrphanOutcome {
    /// Follower CAS-claimed a fresh leader entry and is driving a new
    /// `tools/call` stream; the carried response is the SSE stream
    /// the resume client should consume.
    Reinvoked(Response),
    /// Cluster-0.20 fall-back fired (handler not idempotent, attempt
    /// cap reached, no cached payload, CAS conflict, or no
    /// resume-buffer events). The carried response is the single-shot
    /// LEADER_DIED envelope.
    Terminated(Response),
    /// No leader registry wired or the orphan write found no buffer
    /// entries; caller falls through to the regular replay path.
    Passthrough,
}

/// Cluster-0.24 follower-side response to a dead-leader probe on
/// `klieo/tools/resume`.
///
/// Branches on the cached [`klieo_core::LeaderEntry`]:
/// 1. No registry wired → [`OrphanOutcome::Passthrough`].
/// 2. Lookup yields no entry / KV error → terminate-only.
/// 3. Cached payload missing / unparseable → terminate-only.
/// 4. `tool_invoker.is_tool_idempotent(name) == false` → terminate-
///    only.
/// 5. `entry.attempt >= FAILOVER_ATTEMPT_CAP` → terminate-only.
/// 6. CAS-claim with `attempt + 1` succeeds → record a
///    `failover-reinvoke` marker frame at `max + 1`, fabricate an
///    [`Identity`] from the cached principal, and drive a fresh
///    `tools/call` invocation via `run_tools_call` (private helper) with the new
///    handle adopted as `existing_leader`. The re-invoke's events
///    append to the existing resume buffer at `max + 2` onwards.
/// 7. CAS conflict (another follower won) / CAS error → terminate-
///    only.
///
/// `pub` under `test-fixtures` so integration tests can drive the
/// gate directly with a pre-seeded `LeaderEntry` in the KV bucket —
/// the production path's `is_alive` probe and
/// `lookup_entry_with_revision` both go through `kv.get`, so any HTTP-
/// driven test that mocks "dead leader" via `kv.delete` collapses the
/// entry visibility for the lookup too. Same shape as A2A's
/// `handle_dead_leader_orphan`.
#[cfg(feature = "test-fixtures")]
pub async fn handle_dead_leader_orphan_mcp(
    server: &Arc<McpServer>,
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    req_id: &serde_json::Value,
    stream_id: &str,
) -> OrphanOutcome {
    handle_dead_leader_orphan_mcp_impl(server, buffer, req_id, stream_id).await
}

#[cfg(not(feature = "test-fixtures"))]
async fn handle_dead_leader_orphan_mcp(
    server: &Arc<McpServer>,
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    req_id: &serde_json::Value,
    stream_id: &str,
) -> OrphanOutcome {
    handle_dead_leader_orphan_mcp_impl(server, buffer, req_id, stream_id).await
}

async fn handle_dead_leader_orphan_mcp_impl(
    server: &Arc<McpServer>,
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    req_id: &serde_json::Value,
    stream_id: &str,
) -> OrphanOutcome {
    let Some(registry) = server.leader_registry() else {
        return OrphanOutcome::Passthrough;
    };
    let key = format!("{MCP_LEADER_KEY_PREFIX}{stream_id}");
    let lookup = registry.lookup_entry_with_revision(&key).await;
    let Some((entry, prior_rev)) = lookup_ok_or_log_mcp(&key, lookup) else {
        return terminate_orphan_mcp(buffer, req_id, stream_id).await;
    };
    let Some(payload_bytes) = entry.payload.clone() else {
        tracing::debug!(
            target: "mcp.failover",
            stream_id = %stream_id,
            "no cached payload on leader entry; emitting terminate frame",
        );
        return terminate_orphan_mcp(buffer, req_id, stream_id).await;
    };
    let parsed_body: serde_json::Value = match serde_json::from_slice(&payload_bytes) {
        Ok(v) => v,
        Err(e) => {
            tracing::error!(
                target: "mcp.failover",
                stream_id = %stream_id,
                error = %e,
                "cached payload parse failed; emitting terminate frame",
            );
            return terminate_orphan_mcp(buffer, req_id, stream_id).await;
        }
    };
    let tool_name = parsed_body
        .pointer("/params/name")
        .and_then(|v| v.as_str())
        .unwrap_or("");
    if !server.invoker.is_tool_idempotent(tool_name) {
        tracing::debug!(
            target: "mcp.failover",
            stream_id = %stream_id,
            tool = %tool_name,
            "tool not idempotent; emitting terminate frame",
        );
        return terminate_orphan_mcp(buffer, req_id, stream_id).await;
    }
    if entry.attempt >= server.max_failover_attempts() {
        tracing::warn!(
            target: "mcp.failover",
            stream_id = %stream_id,
            attempt = entry.attempt,
            cap = server.max_failover_attempts(),
            "failover attempt cap reached; emitting terminate frame",
        );
        return terminate_orphan_mcp(buffer, req_id, stream_id).await;
    }
    let new_handle = match registry
        .claim_with_attempt_cas_and_heartbeat(
            key.clone(),
            server.leader_ttl(),
            server.leader_heartbeat_interval(),
            prior_rev,
            &entry,
        )
        .await
    {
        Ok(h) => h,
        Err(klieo_core::BusError::CasConflict { .. }) => {
            tracing::info!(
                target: "mcp.failover",
                stream_id = %stream_id,
                "another follower won the CAS race; emitting terminate frame",
            );
            return terminate_orphan_mcp(buffer, req_id, stream_id).await;
        }
        Err(e) => {
            tracing::warn!(
                target: "mcp.failover",
                stream_id = %stream_id,
                error = %e,
                "CAS claim failed; emitting terminate frame",
            );
            return terminate_orphan_mcp(buffer, req_id, stream_id).await;
        }
    };
    record_failover_marker_mcp(buffer, stream_id, &entry, &new_handle).await;
    let progress_token = parsed_body
        .pointer("/params/_meta/progressToken")
        .cloned()
        .unwrap_or(serde_json::Value::Null);
    let reinvoke_identity = entry.principal.as_ref().map(|p| Identity::new(p.clone()));
    let resp = run_tools_call(
        server.clone(),
        parsed_body,
        progress_token,
        reinvoke_identity,
        Some(new_handle),
    )
    .await;
    OrphanOutcome::Reinvoked(resp)
}

/// Unwrap a `lookup_entry_with_revision` result, logging KV-layer
/// errors at `warn` before falling back to terminate. Returns
/// `Some((entry, rev))` only on `Ok(Some(_))`. Mirrors A2A's
/// `lookup_ok_or_log`.
fn lookup_ok_or_log_mcp(
    key: &str,
    lookup: Result<Option<(klieo_core::LeaderEntry, klieo_core::Revision)>, klieo_core::BusError>,
) -> Option<(klieo_core::LeaderEntry, klieo_core::Revision)> {
    match lookup {
        Ok(Some(pair)) => Some(pair),
        Ok(None) => None,
        Err(e) => {
            tracing::warn!(
                target: "mcp.failover",
                key,
                error = %e,
                "leader entry lookup_with_revision failed; falling back to terminate",
            );
            None
        }
    }
}

/// Write the cluster-0.24 `failover-reinvoke` marker frame to the
/// resume buffer at `max + 1` so resume clients see a clear boundary
/// between the dead leader's last event and the re-invoke's fresh
/// sequence (which appends at `max + 2` onwards). Best-effort: a
/// record failure logs a `warn` + continues so the re-invoke still
/// drives the production stream.
async fn record_failover_marker_mcp(
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    stream_id: &str,
    prior: &klieo_core::LeaderEntry,
    new_handle: &klieo_core::LeaderHandle,
) {
    let Some(max) = max_event_id(buffer, stream_id).await else {
        return;
    };
    let marker_id = max + 1;
    let frame = failover_reinvoke_sse_frame_bytes_mcp(
        stream_id,
        marker_id,
        prior.attempt + 1,
        new_handle.replica_id(),
    );
    if let Err(e) = buffer.record(stream_id, marker_id, frame).await {
        tracing::warn!(
            target: "mcp.failover",
            stream_id = %stream_id,
            marker_id,
            error = %e,
            "failover-reinvoke marker record failed; continuing without marker",
        );
    }
}

/// Build the cluster-0.24 `failover-reinvoke` SSE marker frame
/// recorded in the resume buffer at `max + 1` when a follower
/// re-invokes an idempotent tool from the cached payload. The
/// re-invoke's own events append at `max + 2` onwards via
/// [`emit_progress_frame`], so resume clients see the marker between
/// the original leader's last event and the re-invoke's fresh
/// sequence. Mirrors A2A's `failover_reinvoke_sse_frame_bytes`.
fn failover_reinvoke_sse_frame_bytes_mcp(
    stream_id: &str,
    event_id: u64,
    attempt: u32,
    new_replica_id: &str,
) -> Bytes {
    let payload = serde_json::json!({
        "jsonrpc": "2.0",
        "id": serde_json::Value::Null,
        "event": "failover-reinvoke",
        "data": {
            "stream_id": format!("{MCP_LEADER_KEY_PREFIX}{stream_id}"),
            "attempt": attempt,
            "by_replica": new_replica_id,
        },
        "event_id": event_id,
    });
    Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
}

/// Combine the resume buffer's `replay` stream with an optional live
/// bus tail into one `Result<Bytes, Infallible>` stream of SSE frames.
///
/// `max_replayed` tracks the highest id observed during replay so the
/// tail-side `filter_map` can drop frames already delivered by replay
/// (preserves monotonic ordering across the boundary). Tail frames
/// missing a parseable `id:` prefix are warn-logged and dropped.
fn combine_replay_and_tail(
    replay: Pin<Box<dyn Stream<Item = (u64, Bytes)> + Send>>,
    live_tail: Option<klieo_core::MsgStream>,
    max_replayed: Arc<AtomicU64>,
    stream_id: String,
) -> Pin<Box<dyn Stream<Item = Result<Bytes, std::convert::Infallible>> + Send>> {
    let max_for_replay = max_replayed.clone();
    let replay = replay.map(move |(id, payload)| {
        max_for_replay.fetch_max(id, std::sync::atomic::Ordering::SeqCst);
        Ok::<_, std::convert::Infallible>(id_prefixed_frame(id, payload))
    });
    let Some(tail_msgs) = live_tail else {
        return replay.boxed();
    };
    let max_for_tail = max_replayed;
    let stream_id_for_log = stream_id;
    let tail = tail_msgs.filter_map(move |msg_result| {
        let stream_id = stream_id_for_log.clone();
        let max_for_tail = max_for_tail.clone();
        async move {
            let msg = match msg_result {
                Ok(m) => m,
                Err(e) => {
                    tracing::warn!(
                        target: "mcp.fanout",
                        stream_id = %stream_id,
                        error = %e,
                        "live tail subscription stream error; skipping"
                    );
                    return None;
                }
            };
            // Extract W3C tracecontext from bus headers and parent the
            // decode span under the publisher's span so cross-replica
            // traces stitch into one tree (cluster 0.23, ADR-023).
            let parent_cx = klieo_core::extract_traceparent(&msg.headers);
            let decode_span = tracing::info_span!(
                "tail_frame_decode",
                messaging.system = "klieo-bus",
                messaging.destination = %format!("klieo.mcp.progress.{stream_id}"),
                messaging.operation = "receive",
                klieo.stream_id = %stream_id,
            );
            decode_span.set_parent(parent_cx);
            let _enter = decode_span.enter();
            let payload = msg.payload.clone();
            if let Err(e) = msg.ack.ack().await {
                tracing::warn!(
                    target: "mcp.fanout",
                    error = %e,
                    "ack failed; ephemeral consumer continues"
                );
            }
            match parse_id_prefix(&payload) {
                Some(id) if id > max_for_tail.load(std::sync::atomic::Ordering::SeqCst) => {
                    Some(Ok::<_, std::convert::Infallible>(payload))
                }
                Some(_) => None,
                None => {
                    tracing::warn!(
                        target: "mcp.fanout",
                        payload_len = payload.len(),
                        "tail frame missing or unparseable id prefix; dropping",
                    );
                    None
                }
            }
        }
    });
    replay.chain(tail).boxed()
}

/// Stamp one progress/lagged/result frame with the next monotonic
/// `id:` prefix, spawn the best-effort resume-buffer record + bus
/// publish, and return the stamped frame for the SSE yield. Collapses
/// the inline 6-line block repeated 9 times in `stream_tools_call`.
///
/// Frame routing preserves the pre-refactor split: the resume buffer
/// stores the UNSTAMPED frame (replay restamps via
/// `id_prefixed_frame` on read), while the cross-replica publish + the
/// SSE yield carry the STAMPED bytes.
#[allow(clippy::too_many_arguments)]
fn emit_progress_frame(
    buffer: &Arc<dyn klieo_core::resume::ResumeBuffer>,
    pubsub: &Arc<dyn klieo_core::Pubsub>,
    permits: &Arc<tokio::sync::Semaphore>,
    publish_subject: &str,
    stream_id: &str,
    next_id: &AtomicU64,
    frame: Bytes,
    terminal: bool,
) -> Bytes {
    let id = next_id.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
    spawn_record(
        buffer.clone(),
        stream_id.to_string(),
        id,
        frame.clone(),
        terminal,
    );
    let stamped = id_prefixed_frame(id, frame);
    spawn_publish(
        pubsub.clone(),
        publish_subject.to_string(),
        stamped.clone(),
        permits.clone(),
    );
    stamped
}

/// Lower bound on a cross-hop provenance anchor — short enough to admit a
/// truncated hash, long enough to reject stray single characters.
const MIN_PARENT_ANCHOR_BYTES: usize = 8;
/// Upper bound — a hash token, not a payload. Caps the audit-trail string
/// a caller can write per run.
const MAX_PARENT_ANCHOR_BYTES: usize = 256;

/// `true` for the hash-token charset accepted in a parent-chain anchor:
/// base64url (`-`/`_`), hex, `=` padding, and `:` algorithm namespacing
/// (`sha256:…`). Excludes whitespace, control chars, and `@`/`.` so most
/// freeform/PII strings are rejected at the boundary.
fn is_parent_anchor_byte(byte: u8) -> bool {
    byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'=' | b':' | b'-')
}

/// Validate the optional cross-hop provenance anchor at `params._meta.
/// parentAnchor`. Absent → `Ok(None)`. Present but malformed →
/// `Err(stable message)`: a caller that supplied the field asked for a
/// cross-hop link, so silently dropping it would leave an invisible gap
/// in the audit chain — reject loud instead. Validation is generous; a
/// conformant hash token never trips it.
fn parse_parent_anchor(raw: &serde_json::Value) -> Result<Option<String>, &'static str> {
    let Some(value) = raw.pointer("/params/_meta/parentAnchor") else {
        return Ok(None);
    };
    let anchor = value
        .as_str()
        .ok_or("_meta.parentAnchor must be a string")?;
    if anchor.len() < MIN_PARENT_ANCHOR_BYTES || anchor.len() > MAX_PARENT_ANCHOR_BYTES {
        return Err("_meta.parentAnchor length out of bounds");
    }
    if !anchor.bytes().all(is_parent_anchor_byte) {
        return Err("_meta.parentAnchor must be a hash token");
    }
    Ok(Some(anchor.to_string()))
}

#[cfg(test)]
mod parent_anchor_tests {
    use super::parse_parent_anchor;
    use serde_json::json;

    fn req_with_anchor(anchor: serde_json::Value) -> serde_json::Value {
        json!({ "params": { "name": "t", "_meta": { "parentAnchor": anchor } } })
    }

    #[test]
    fn absent_anchor_is_ok_none() {
        let req = json!({ "params": { "name": "t" } });
        assert_eq!(parse_parent_anchor(&req), Ok(None));
        // Even an empty `_meta` object yields None.
        let req = json!({ "params": { "name": "t", "_meta": {} } });
        assert_eq!(parse_parent_anchor(&req), Ok(None));
    }

    #[test]
    fn valid_hash_token_is_accepted_verbatim() {
        let req = req_with_anchor(json!("sha256:0123abcd_ef-ABCD="));
        assert_eq!(
            parse_parent_anchor(&req),
            Ok(Some("sha256:0123abcd_ef-ABCD=".to_string()))
        );
    }

    #[test]
    fn non_string_anchor_is_rejected() {
        assert!(parse_parent_anchor(&req_with_anchor(json!(42))).is_err());
        assert!(parse_parent_anchor(&req_with_anchor(json!(["a"]))).is_err());
    }

    #[test]
    fn empty_or_short_anchor_is_rejected() {
        assert!(parse_parent_anchor(&req_with_anchor(json!(""))).is_err());
        assert!(parse_parent_anchor(&req_with_anchor(json!("abc"))).is_err());
    }

    #[test]
    fn oversize_anchor_is_rejected() {
        let huge = "a".repeat(257);
        assert!(parse_parent_anchor(&req_with_anchor(json!(huge))).is_err());
    }

    #[test]
    fn freeform_or_pii_shaped_anchor_is_rejected() {
        // An email (`@`, `.`) and whitespace are outside the hash-token
        // charset, so most freeform/PII strings fail at the boundary.
        assert!(parse_parent_anchor(&req_with_anchor(json!("alice@example.com"))).is_err());
        assert!(parse_parent_anchor(&req_with_anchor(json!("hello world"))).is_err());
    }
}

/// Extract params and spawn the invoke task for SSE-upgrade `tools/call`.
/// Returns the request ID, broadcast channel, and task handle.
///
/// `parent_anchor` is the validated cross-hop provenance anchor; it is
/// honored only for an authenticated (non-anonymous) caller so every
/// recorded `Episode::RunOrigin` is co-attributable to a principal.
fn spawn_progress_stream_task(
    server: &Arc<McpServer>,
    raw: &serde_json::Value,
    cancel: tokio_util::sync::CancellationToken,
    identity: Option<&Identity>,
    parent_anchor: Option<String>,
) -> (
    serde_json::Value,
    tokio::sync::broadcast::Receiver<klieo_core::AgentEvent>,
    tokio::task::JoinHandle<Result<serde_json::Value, klieo_core::error::ToolError>>,
) {
    let req_id = raw.get("id").cloned().unwrap_or(serde_json::Value::Null);
    let params = raw
        .get("params")
        .cloned()
        .unwrap_or(serde_json::Value::Null);

    let name = params
        .get("name")
        .and_then(|n| n.as_str())
        .unwrap_or("")
        .to_string();
    let args = params
        .get("arguments")
        .cloned()
        .unwrap_or(serde_json::Value::Null);

    let (tx, rx) = tokio::sync::broadcast::channel::<klieo_core::AgentEvent>(PROGRESS_CHANNEL_CAP);
    // Anonymous identities are excluded so the workflow's authz path
    // never binds a ticket to the "anonymous" principal — a deny-by-
    // default the slice-1 no-ticket envelope already covers.
    let principal = identity
        .filter(|id| !id.is_anonymous())
        .map(|id| id.as_str().to_string());
    // Cross-hop anchor honored only for authenticated callers: an
    // anonymous caller has no principal to attribute the (unverified)
    // parent claim to, so its anchor is dropped rather than recorded
    // as an orphan `RunOrigin`.
    let parent_anchor = if principal.is_some() {
        parent_anchor
    } else {
        None
    };
    let tool_ctx = server.tool_ctx_with_progress(tx, cancel, principal, parent_anchor);

    let invoker = server.invoker.clone();
    let invoke_handle = tokio::spawn(async move { invoker.invoke(&name, args, tool_ctx).await });

    (req_id, rx, invoke_handle)
}

/// Replay buffered events for `params.progress_token` since
/// `params.last_event_id`, then tail live events published to
/// `klieo.mcp.progress.{token}` so cross-replica clients see events
/// emitted on other replicas.
///
/// Ordering guarantee: subscribe-before-replay — the bus subscription
/// is established before the replay stream is consumed, so events
/// published during the replay drain are not lost.
///
/// Error mapping:
/// - `ResumeError::Expired`  → JSON-RPC -32011.
/// - `ResumeError::NotFound` → JSON-RPC -32012.
/// - `ResumeError::Backend`  → JSON-RPC -32603 (logged server-side only).
/// - Subscribe failure       → replay-only (logged warning).
#[instrument(
    skip_all,
    fields(
        rpc.system = "klieo-mcp",
        rpc.method = "klieo/tools/resume",
        klieo.stream_id = tracing::field::Empty,
    ),
)]
async fn stream_resume(
    server: Arc<McpServer>,
    raw: serde_json::Value,
    params: ResumeParams,
    identity: Option<Identity>,
) -> Response {
    let req_id = raw.get("id").cloned().unwrap_or(serde_json::Value::Null);
    let stream_id = progress_token_to_string(&params.progress_token);
    if let Err(e) = klieo_core::validate_subject_token(&stream_id) {
        tracing::warn!(
            target: "mcp.resume",
            error = %e,
            "rejected progressToken: invalid subject segment",
        );
        return (
            StatusCode::OK,
            Json(rpc_error(
                Some(req_id),
                JSONRPC_INVALID_PARAMS,
                "progressToken contains reserved bus-subject metacharacters",
            )),
        )
            .into_response();
    }
    tracing::Span::current().record("klieo.stream_id", stream_id.as_str());
    let request_cancel = server.parent_cancel.child_token();
    let buffer = server.resume_buffer.clone();
    let pubsub = server.pubsub.clone();
    let since = params.last_event_id;

    // Orphan detection: leader probe at the resume entry. When the
    // leader is dead AND the resume buffer has retained events (the
    // leader must have written some before dying), the follower
    // either re-invokes from the cached payload (cluster 0.24,
    // idempotent tool + attempt cap + payload available) OR writes
    // the terminal "leader died" SSE frame and returns a single-shot
    // JSON-RPC error envelope (code -32099) so the client sees clean
    // termination + can retry. ADR-020 / ADR-024.
    if let LeaderProbe::Dead = probe_leader(&server, &stream_id).await {
        match handle_dead_leader_orphan_mcp(&server, &buffer, &req_id, &stream_id).await {
            OrphanOutcome::Reinvoked(resp) | OrphanOutcome::Terminated(resp) => return resp,
            OrphanOutcome::Passthrough => {}
        }
    }

    // Tenant-binding gate (ADR-022). Runs BEFORE the resume buffer
    // `replay` so a cross-tenant probe cannot infer existence via
    // backend side effects (NotFound vs Expired surfaces differ on
    // some backends). Owner mismatch returns the same envelope as
    // the buffer's `ResumeError::NotFound` arm below — wire-shape
    // identical (modulo request id) per OWASP IDOR mitigation.
    if let Err(resp) = enforce_owner(&server, &stream_id, &identity, &req_id).await {
        return resp;
    }

    let replay_stream = match buffer.replay(&stream_id, since).await {
        Ok(s) => s,
        Err(klieo_core::resume::ResumeError::Expired { since_id }) => {
            return (
                StatusCode::OK,
                Json(rpc_error(
                    Some(req_id),
                    crate::JSONRPC_RESUME_BUFFER_EXPIRED,
                    &format!("resume window expired (since_id={since_id})"),
                )),
            )
                .into_response();
        }
        Err(klieo_core::resume::ResumeError::NotFound(_)) => {
            return (
                StatusCode::OK,
                Json(rpc_error(
                    Some(req_id),
                    crate::JSONRPC_RESUME_BUFFER_NOT_FOUND,
                    "no buffered stream for progressToken",
                )),
            )
                .into_response();
        }
        Err(klieo_core::resume::ResumeError::Backend(e)) => {
            tracing::warn!(
                target: "mcp.resume",
                stream_id = %stream_id,
                error = %e,
                "resume backend error"
            );
            return (
                StatusCode::OK,
                Json(rpc_error(
                    Some(req_id),
                    JSONRPC_SERVER_ERROR,
                    "resume backend unavailable",
                )),
            )
                .into_response();
        }
        // ResumeError is #[non_exhaustive]; future variants fall back to a
        // generic server-error envelope.
        Err(_) => {
            return (
                StatusCode::OK,
                Json(rpc_error(
                    Some(req_id),
                    JSONRPC_SERVER_ERROR,
                    "resume backend error",
                )),
            )
                .into_response();
        }
    };

    // Subscribe BEFORE consuming replay so events published during the
    // replay drain are not lost (subscribe-before-replay race fix).
    // Skip the tail when the buffer is already terminal — no further
    // events will be published, so an infinite subscription would
    // prevent the SSE response from closing after replay drains.
    let already_closed = match buffer.is_terminal(&stream_id).await {
        Ok(v) => v,
        Err(e) => {
            tracing::warn!(
                target: "mcp.resume",
                stream_id = %stream_id,
                error = %e,
                "is_terminal check failed; assuming live"
            );
            false
        }
    };
    let subject = format!("klieo.mcp.progress.{}", stream_id);
    let live_tail = if already_closed {
        None
    } else {
        let durable = klieo_core::DurableName::new(format!("klieo-eph-{}", uuid::Uuid::new_v4()));
        match pubsub.subscribe(&subject, durable).await {
            Ok(s) => Some(s),
            Err(e) => {
                tracing::warn!(
                    target: "mcp.resume",
                    subject = %subject,
                    error = %e,
                    "live tail subscribe failed; returning replay-only response"
                );
                None
            }
        }
    };

    let max_replayed = Arc::new(AtomicU64::new(since));
    let combined =
        combine_replay_and_tail(replay_stream, live_tail, max_replayed, stream_id.clone());

    // Register the per-invoke cancel token under stream_id so the
    // wildcard cancel-subject subscription can fire it on inbound
    // klieo.mcp.cancel.{progressToken} messages. The DeregisterOnDrop
    // wrapper inside `wrap_with_cancel_fanout` removes the entry on
    // stream end or client disconnect.
    server
        .cancel_registry()
        .register(stream_id.clone(), request_cancel.clone());
    // stream_resume holds no leader claim — claims belong to the
    // originating invoke (`stream_tools_call`). Resume is read-side.
    // Ownership handle is `None` for the same reason: resume must
    // never `kv.delete` on drop, only the invoke owner does.
    let guarded = wrap_with_cancel_fanout(&server, combined, stream_id, request_cancel, None, None);
    axum::response::Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/event-stream")
        .header("Cache-Control", "no-cache")
        .body(axum::body::Body::from_stream(guarded))
        .unwrap()
}

/// Serve a `tools/call` request as an SSE stream.
///
/// Flow:
/// 1. Extract tool name and arguments, spawn the invoke task.
/// 2. Build the SSE stream pump: interleaves progress events with
///    the terminal result frame. The pump is one cohesive select-loop
///    that cannot be split further without fragmenting the state machine.
/// 3. Each yielded frame receives a monotonically-increasing `id:` line
///    (per-progressToken counter) and is recorded to the resume buffer.
/// 4. Wrap the body stream in `CancelOnDrop` so dropping the response
///    body (client TCP close) fires the request-scoped cancel token,
///    which the spawned invoke task observes via `ToolCtx.cancel`.
/// 5. Return the SSE response with correct headers.
#[allow(clippy::too_many_lines)]
#[instrument(
    skip_all,
    fields(
        rpc.system = "klieo-mcp",
        rpc.method = "tools/call",
        klieo.stream_id = tracing::field::Empty,
    ),
)]
async fn stream_tools_call(
    server: Arc<McpServer>,
    raw: serde_json::Value,
    progress_token: serde_json::Value,
    identity: Option<Identity>,
) -> Response {
    run_tools_call(server, raw, progress_token, identity, None).await
}

/// Shared body of the `tools/call` SSE-upgrade path.
///
/// Driven from two sites:
/// - [`stream_tools_call`] passes `existing_leader = None` and the
///   path performs the cluster-0.20 [`try_claim_leader`] claim,
///   capturing the original JSON-RPC body + authenticated principal
///   inside the [`klieo_core::LeaderEntry`] for cluster-0.24
///   follower re-invoke.
/// - [`handle_dead_leader_orphan_mcp`] passes `existing_leader =
///   Some(handle)` (already CAS-claimed against the dead leader's
///   revision with `attempt + 1`); this path SKIPS the on-invoke
///   claim and adopts the supplied handle. The cached payload-+-
///   principal pair came from the dead leader's entry, so the
///   re-invoke runs under the same tenant binding as the original.
#[allow(clippy::too_many_lines)]
async fn run_tools_call(
    server: Arc<McpServer>,
    raw: serde_json::Value,
    progress_token: serde_json::Value,
    identity: Option<Identity>,
    existing_leader: Option<klieo_core::LeaderHandle>,
) -> Response {
    let stream_id = progress_token_to_string(&progress_token);
    if let Err(e) = klieo_core::validate_subject_token(&stream_id) {
        tracing::warn!(
            target: "mcp.tools_call",
            error = %e,
            "rejected progressToken: invalid subject segment",
        );
        let req_id = raw.get("id").cloned().unwrap_or(serde_json::Value::Null);
        return (
            StatusCode::OK,
            Json(rpc_error(
                Some(req_id),
                JSONRPC_INVALID_PARAMS,
                "progressToken contains reserved bus-subject metacharacters",
            )),
        )
            .into_response();
    }
    // Validate the optional cross-hop provenance anchor before claiming
    // leadership/ownership — a malformed anchor fails the call loud
    // rather than silently dropping a requested audit link.
    let parent_anchor = match parse_parent_anchor(&raw) {
        Ok(anchor) => anchor,
        Err(message) => {
            tracing::warn!(
                target: "mcp.tools_call",
                reason = message,
                "rejected tools/call: malformed _meta.parentAnchor",
            );
            let req_id = raw.get("id").cloned().unwrap_or(serde_json::Value::Null);
            return (
                StatusCode::OK,
                Json(rpc_error(Some(req_id), JSONRPC_INVALID_PARAMS, message)),
            )
                .into_response();
        }
    };
    tracing::Span::current().record("klieo.stream_id", stream_id.as_str());
    let request_cancel = server.parent_cancel.child_token();
    // Register the per-invoke cancel token under stream_id so the
    // wildcard cancel-subject subscription can fire it on inbound
    // klieo.mcp.cancel.{progressToken} messages. The matching
    // deregister fires from `wrap_with_cancel_fanout` on stream drop.
    server
        .cancel_registry()
        .register(stream_id.clone(), request_cancel.clone());
    // Adopt the follower-supplied handle on cluster-0.24 re-invoke;
    // otherwise claim fresh leadership for the lifetime of the SSE
    // response (multi-replica orphan detection, ADR-020). KV-layer
    // failures on a fresh claim degrade silently: log warn + proceed
    // without a claim (a follower stream_resume will treat the
    // buffer as orphaned only if it ALSO has buffered events;
    // otherwise the regular replay path runs unchanged).
    //
    // Cluster 0.24: capture the original JSON-RPC body + the
    // authenticated principal so a follower can re-invoke an
    // idempotent tool from the cached payload under the same
    // tenant binding. A serialisation failure here is non-fatal:
    // the invoke proceeds with `payload = None` and any follower
    // that detects an orphan will terminate (no payload to
    // re-invoke from).
    let leader_handle = match existing_leader {
        Some(handle) => Some(handle),
        None => {
            let payload_bytes_for_failover = match serde_json::to_vec(&raw) {
                Ok(v) => Some(Bytes::from(v)),
                Err(e) => {
                    tracing::warn!(
                        target: "mcp.failover",
                        stream_id = %stream_id,
                        error = %e,
                        "tools/call body serialise for failover failed; \
                         proceeding without cached payload",
                    );
                    None
                }
            };
            let principal_for_failover = identity
                .as_ref()
                .filter(|id| !id.is_anonymous())
                .map(|id| id.as_str().to_string());
            try_claim_leader(
                &server,
                &stream_id,
                payload_bytes_for_failover,
                principal_for_failover,
            )
            .await
        }
    };
    // Tenant-binding claim (ADR-022). Held alongside the leader
    // handle inside `CancelOnDrop` so both KV entries fire their
    // best-effort delete on body drop in a single drop pass.
    let ownership_handle = match try_claim_ownership(&server, &stream_id, &identity).await {
        Ok(handle) => handle,
        Err(resp) => return resp,
    };
    let (req_id, rx, invoke_handle) = spawn_progress_stream_task(
        &server,
        &raw,
        request_cancel.clone(),
        identity.as_ref(),
        parent_anchor,
    );

    let next_id = Arc::new(AtomicU64::new(0));
    let buffer = server.resume_buffer.clone();
    let pubsub = server.pubsub.clone();
    let publish_subject = format!("klieo.mcp.progress.{stream_id}");
    let publish_permits = server.publish_permits.clone();

    // `async_stream::stream!` expands to `async move { ... }` and captures
    // referenced locals by move. Clone the locals the stream needs so the
    // originals remain available for the `wrap_with_cancel_fanout` wiring below.
    let request_cancel_for_stream = request_cancel.clone();
    let stream_id_for_stream = stream_id.clone();

    let body_stream = async_stream::stream! {
        let mut rx = rx;
        let invoke_handle = invoke_handle;
        tokio::pin!(invoke_handle);
        tokio::task::yield_now().await;

        loop {
            tokio::select! {
                biased;
                _ = request_cancel_for_stream.cancelled() => {
                    // Body about to drop; stop yielding to avoid writing into a
                    // closed body. Spawned task observes the same token via ToolCtx.
                    break;
                }
                recv_result = rx.recv() => {
                    match recv_result {
                        Ok(event) => {
                            let frame = match progress_event(&progress_token, &event) {
                                Ok(f) => f,
                                Err(e) => match e {},
                            };
                            yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                                &buffer, &pubsub, &publish_permits,
                                &publish_subject, &stream_id_for_stream,
                                &next_id, frame, false,
                            ));
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                            let frame = match lagged_event(&progress_token, n) {
                                Ok(f) => f,
                                Err(e) => match e {},
                            };
                            yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                                &buffer, &pubsub, &publish_permits,
                                &publish_subject, &stream_id_for_stream,
                                &next_id, frame, false,
                            ));
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                            break;
                        }
                    }
                }
                result = &mut invoke_handle => {
                    loop {
                        match rx.try_recv() {
                            Ok(event) => {
                                let frame = match progress_event(&progress_token, &event) {
                                    Ok(f) => f,
                                    Err(e) => match e {},
                                };
                                yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                                    &buffer, &pubsub, &publish_permits,
                                    &publish_subject, &stream_id_for_stream,
                                    &next_id, frame, false,
                                ));
                            }
                            Err(tokio::sync::broadcast::error::TryRecvError::Lagged(n)) => {
                                let frame = match lagged_event(&progress_token, n) {
                                    Ok(f) => f,
                                    Err(e) => match e {},
                                };
                                yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                                    &buffer, &pubsub, &publish_permits,
                                    &publish_subject, &stream_id_for_stream,
                                    &next_id, frame, false,
                                ));
                            }
                            Err(_) => break,
                        }
                    }
                    let outcome = result
                        .unwrap_or_else(|_| Err(klieo_core::error::ToolError::Permanent(
                            "invoke task panicked".into()
                        )));
                    let req_id_opt = if req_id == serde_json::Value::Null {
                        None
                    } else {
                        Some(req_id.clone())
                    };
                    let frame = match result_event(req_id_opt, outcome) {
                        Ok(f) => f,
                        Err(e) => match e {},
                    };
                    yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                        &buffer, &pubsub, &publish_permits,
                        &publish_subject, &stream_id_for_stream,
                        &next_id, frame, true,
                    ));
                    return;
                }
            }
        }

        let result = invoke_handle.await
            .unwrap_or_else(|_| Err(klieo_core::error::ToolError::Permanent(
                "invoke task panicked".into()
            )));
        loop {
            match rx.try_recv() {
                Ok(event) => {
                    let frame = match progress_event(&progress_token, &event) {
                        Ok(f) => f,
                        Err(e) => match e {},
                    };
                    yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                        &buffer, &pubsub, &publish_permits,
                        &publish_subject, &stream_id_for_stream,
                        &next_id, frame, false,
                    ));
                }
                Err(tokio::sync::broadcast::error::TryRecvError::Lagged(n)) => {
                    let frame = match lagged_event(&progress_token, n) {
                        Ok(f) => f,
                        Err(e) => match e {},
                    };
                    yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
                        &buffer, &pubsub, &publish_permits,
                        &publish_subject, &stream_id_for_stream,
                        &next_id, frame, false,
                    ));
                }
                Err(_) => break,
            }
        }
        let req_id_opt = if req_id == serde_json::Value::Null {
            None
        } else {
            Some(req_id)
        };
        let frame = match result_event(req_id_opt, result) {
            Ok(f) => f,
            Err(e) => match e {},
        };
        yield Ok::<_, std::convert::Infallible>(emit_progress_frame(
            &buffer, &pubsub, &publish_permits,
            &publish_subject, &stream_id_for_stream,
            &next_id, frame, true,
        ));
    };

    let guarded = wrap_with_cancel_fanout(
        &server,
        body_stream.boxed(),
        stream_id,
        request_cancel,
        leader_handle,
        ownership_handle,
    );

    axum::response::Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/event-stream")
        .header("Cache-Control", "no-cache")
        .body(axum::body::Body::from_stream(guarded))
        .unwrap()
}

/// Format one `AgentEvent` as an SSE `event: progress` frame carrying a
/// `notifications/progress` JSON-RPC notification envelope.
fn progress_event(
    token: &serde_json::Value,
    event: &klieo_core::AgentEvent,
) -> Result<axum::body::Bytes, std::convert::Infallible> {
    let payload = serde_json::json!({
        "jsonrpc": "2.0",
        "method": "notifications/progress",
        "params": {
            "progressToken": token,
            "data": event,
        }
    });
    Ok(axum::body::Bytes::from(format!(
        "event: progress\ndata: {}\n\n",
        payload
    )))
}

/// Synthetic progress frame emitted when the broadcast channel lags.
fn lagged_event(
    token: &serde_json::Value,
    skipped: u64,
) -> Result<axum::body::Bytes, std::convert::Infallible> {
    let payload = serde_json::json!({
        "jsonrpc": "2.0",
        "method": "notifications/progress",
        "params": {
            "progressToken": token,
            "data": { "kind": "lagged", "message": format!("lagged: skipped {} events", skipped) }
        }
    });
    Ok(axum::body::Bytes::from(format!(
        "event: progress\ndata: {}\n\n",
        payload
    )))
}

/// Terminal SSE frame: carries the JSON-RPC `tools/call` result or a
/// redacted error envelope.
fn result_event(
    id: Option<serde_json::Value>,
    outcome: Result<serde_json::Value, klieo_core::error::ToolError>,
) -> Result<axum::body::Bytes, std::convert::Infallible> {
    let envelope = match outcome {
        Ok(v) => {
            let result = serde_json::json!({
                "content": [{ "type": "text", "text": v.to_string() }]
            });
            serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": result })
        }
        Err(e) => tool_error_to_envelope(id, e),
    };
    Ok(axum::body::Bytes::from(format!(
        "event: result\ndata: {}\n\n",
        envelope
    )))
}

/// Stream wrapper that fires a `DropGuard` on `Drop`. When axum/hyper
/// drops the response body after a client TCP close, the held guard
/// cancels the request-scoped token, propagating to `ToolCtx.cancel`
/// inside the spawned invoke task.
///
/// Drop ordering: the explicit `Drop::drop` body runs FIRST and only
/// spawns a best-effort cross-replica publish (no `.await`). The
/// `_guard` field then drops as part of normal field-drop, firing
/// the local request CancellationToken. Because the body merely
/// spawns, awaiting tasks see the local cancel before the bus
/// publish completes.
///
/// # Security
/// The cross-replica cancel signal published from this `Drop` body
/// embeds the caller-supplied progressToken in its subject. Cancel
/// signals share the progressToken-as-credential threat model
/// documented for resume in ADR-018 / ADR-019: any caller who knows
/// the progressToken can cause a cancel on the owning replica.
/// Operators MUST mint unguessable progressTokens and gate
/// per-tenant authorisation BEFORE the request reaches the server;
/// otherwise cross-tenant cancel becomes possible (CWE-639 IDOR).
struct CancelOnDrop<S> {
    inner: S,
    _guard: tokio_util::sync::DropGuard,
    pubsub: Arc<dyn klieo_core::Pubsub>,
    cancel_subject: String,
    /// Shared with [`McpServer::publish_permits`] so the drop-time
    /// cross-replica cancel publish bounds under the same cap as the
    /// per-stream SSE-frame publishes. Passed through to
    /// [`klieo_core::cancel::spawn_drop_publish`] which `try_acquire`s
    /// before spawning; saturation drops the cancel publish with a
    /// `warn` (cross-replica subscribers still see the same task's
    /// local cancel fire in the conventional drop-of-`_guard` path).
    permits: Arc<tokio::sync::Semaphore>,
    /// Held for the lifetime of the SSE response body. `Drop` releases
    /// the leader claim (heartbeat abort + best-effort KV delete inside
    /// [`klieo_core::LeaderHandle::drop`]) when the body drops, so a
    /// replica failure manifests as a TTL-expired KV entry that a
    /// follower's `stream_resume` orphan probe observes. `None` when
    /// leader election is disabled (no `with_leader_election`) or when
    /// the claim itself failed (fail-open per ADR-020).
    _leader: Option<klieo_core::LeaderHandle>,
    /// Held for the lifetime of the SSE response body. `Drop` fires a
    /// best-effort `kv.delete` inside
    /// [`klieo_core::OwnershipHandle::drop`] so the `klieo-tenants`
    /// entry no longer survives the stream. `None` when tenant
    /// binding is disabled (no `with_tenant_binding`), when no
    /// authenticator produced a non-anonymous principal, or when the
    /// claim itself failed (fail-open per ADR-022). Resume-side
    /// streams never hold this — claims belong to the originating
    /// invoke.
    _ownership: Option<klieo_core::OwnershipHandle>,
}

impl<S: futures::Stream + Unpin> futures::Stream for CancelOnDrop<S> {
    type Item = S::Item;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<S::Item>> {
        std::pin::Pin::new(&mut self.inner).poll_next(cx)
    }
}

impl<S> Drop for CancelOnDrop<S> {
    fn drop(&mut self) {
        // This body runs FIRST; `_guard` drops after the body returns
        // (normal field-drop order) and fires the local request token.
        // Cross-replica fan-out is delegated to
        // [`klieo_core::cancel::spawn_drop_publish`], which handles
        // the empty-subject short-circuit and the no-runtime guard.
        let mut trace_headers = klieo_core::Headers::default();
        klieo_core::inject_traceparent(&mut trace_headers, &opentelemetry::Context::current());
        klieo_core::cancel::spawn_drop_publish(
            self.pubsub.clone(),
            std::mem::take(&mut self.cancel_subject),
            "mcp.cancel",
            Some(self.permits.clone()),
            trace_headers,
        );
    }
}

/// Dispatch one parsed JSON-RPC value, supporting batch arrays.
#[instrument(
    skip_all,
    fields(
        rpc.system = "klieo-mcp",
        rpc.method = tracing::field::Empty,
    ),
)]
async fn dispatch(
    server: &McpServer,
    raw: serde_json::Value,
    session: Option<&std::sync::Arc<crate::session::Session>>,
) -> serde_json::Value {
    if let Some(method) = raw.get("method").and_then(|m| m.as_str()) {
        tracing::Span::current().record("rpc.method", method);
    }
    match raw {
        serde_json::Value::Array(items) => {
            if items.len() > MAX_BATCH_ITEMS {
                warn!(
                    items = items.len(),
                    max = MAX_BATCH_ITEMS,
                    "rejected oversized JSON-RPC batch"
                );
                return rpc_error(None, JSONRPC_SERVER_ERROR, "batch size exceeds limit");
            }
            let mut out = Vec::with_capacity(items.len());
            for item in items {
                if server.parent_cancel.is_cancelled() {
                    out.push(rpc_error(
                        item.get("id").cloned(),
                        JSONRPC_SERVER_ERROR,
                        "server shutting down",
                    ));
                    continue;
                }
                out.push(server.handle_jsonrpc(item, session).await);
            }
            serde_json::Value::Array(out)
        }
        single => server.handle_jsonrpc(single, session).await,
    }
}

#[cfg(test)]
mod post_body_classification_tests {
    //! Unit tests for the POST /mcp body-classification arm + the
    //! `touch_last_activity` helper introduced for HTTP outbound
    //! responses. Driven through the axum [`Router`] via
    //! `tower::ServiceExt::oneshot` so the wire-side semantics
    //! (status codes, headers) are exercised end to end without
    //! binding a port.
    use super::*;
    use crate::outbound::OutboundRequests;
    use crate::{OutboundFrameSink, OutboundSinkError};
    use async_trait::async_trait;
    use axum::body::{to_bytes, Body};
    use axum::http::{header, Method, Request, StatusCode};
    use klieo_core::error::ToolError;
    use klieo_core::llm::ToolDef;
    use klieo_core::tool::{ToolCtx, ToolInvoker};
    use serde_json::{json, Value};
    use tokio::sync::Mutex as AsyncMutex;
    use tower::ServiceExt;

    struct NoopInvoker;

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

    fn server() -> Arc<McpServer> {
        McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .build_arc()
            .unwrap()
    }

    /// Capturing frame sink so the second test can wire an
    /// `OutboundRequests` primitive onto the server without spinning
    /// a transport. Mirrors the test-only sink used by
    /// `outbound::tests::CapturingSink`.
    struct LocalCapturingSink {
        frames: AsyncMutex<Vec<Value>>,
    }

    #[async_trait]
    impl OutboundFrameSink for LocalCapturingSink {
        async fn send_frame(&self, frame: std::sync::Arc<Value>) -> Result<(), OutboundSinkError> {
            self.frames.lock().await.push((*frame).clone());
            Ok(())
        }
    }

    async fn post_init(server: &Arc<McpServer>) -> String {
        let body = json!({
            "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}
        });
        let req = Request::builder()
            .method(Method::POST)
            .uri("/mcp")
            .header(header::CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_vec(&body).unwrap()))
            .unwrap();
        let resp = server.router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        resp.headers()
            .get(MCP_SESSION_ID_HEADER)
            .expect("initialize must echo Mcp-Session-Id")
            .to_str()
            .unwrap()
            .to_string()
    }

    fn outbound_response_request(session_id: &str, id: i64) -> Request<Body> {
        let body = json!({"jsonrpc": "2.0", "id": id, "result": {"x": 1}});
        Request::builder()
            .method(Method::POST)
            .uri("/mcp")
            .header(header::CONTENT_TYPE, "application/json")
            .header(MCP_SESSION_ID_HEADER, session_id)
            .body(Body::from(serde_json::to_vec(&body).unwrap()))
            .unwrap()
    }

    /// Without an outbound primitive wired on the server (the SSE-
    /// stream open path wires it in a follow-on task), a POST body
    /// shaped `{id, result}` must fall through to the rejection arm
    /// and surface as 400 Bad Request with a JSON-RPC parse-error
    /// envelope. Pins the no-outbound branch behaviour.
    #[tokio::test]
    async fn post_outbound_response_with_no_outbound_wired_returns_400() {
        let server = server();
        let session_id = post_init(&server).await;
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        {
            let sessions = server.sessions.read().await;
            let session = sessions
                .get(&session_uuid)
                .expect("post_init inserts the session into the registry");
            assert!(
                session.outbound.get().is_none(),
                "outbound must be unset on a plain HTTP server"
            );
        }

        let resp = server
            .router()
            .oneshot(outbound_response_request(&session_id, 42))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let bytes = to_bytes(resp.into_body(), 1 << 16).await.unwrap();
        let envelope: Value = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(envelope["error"]["code"], JSONRPC_PARSE_ERROR);
    }

    /// With the outbound primitive wired (mocked by directly
    /// populating the `OnceCell` over a capturing sink), a POST body
    /// shaped `{id, result}` must route into `complete_pending` and
    /// return 202 Accepted. Drives the production write path via
    /// `outbound_request` so the awaited future resolves with the
    /// supplied result — the integration ledger that the body-shape
    /// branch + the correlation table cooperate end to end.
    #[tokio::test]
    async fn post_outbound_response_routes_when_outbound_wired() {
        let server = server();
        let session_id = post_init(&server).await;

        let sink: Arc<dyn OutboundFrameSink> = Arc::new(LocalCapturingSink {
            frames: AsyncMutex::new(Vec::new()),
        });
        let outbound = Arc::new(OutboundRequests::new(sink));
        // Production reads route through the per-session registry;
        // seed the mock onto the Session cell so route_outbound_response
        // observes it.
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        let session = {
            let sessions = server.sessions.read().await;
            sessions
                .get(&session_uuid)
                .expect("post_init inserts the session into the registry")
                .clone()
        };
        if session.outbound.set(outbound.clone()).is_err() {
            panic!("session.outbound OnceCell must be empty for this test");
        }

        let call_handle = {
            let outbound = outbound.clone();
            tokio::spawn(async move {
                use klieo_core::ServerOutbound;
                outbound
                    .outbound_request(
                        "custom/method",
                        Value::Null,
                        std::time::Duration::from_secs(2),
                    )
                    .await
            })
        };

        // Yield + sleep so the spawned task registers its pending entry
        // (id=1, the AtomicI64 seed) before we deliver the response.
        // Without this settle the inbound POST may arrive before
        // `outbound_request` has inserted the oneshot sender, and
        // `complete_pending` would log+drop as unknown id.
        tokio::time::sleep(std::time::Duration::from_millis(30)).await;

        let resp = server
            .router()
            .oneshot(outbound_response_request(&session_id, 1))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::ACCEPTED);

        let result = tokio::time::timeout(std::time::Duration::from_secs(2), call_handle)
            .await
            .expect("outbound_request did not resolve")
            .expect("task panicked")
            .expect("outbound_request returned error");
        assert_eq!(result["x"], 1);
    }

    /// Every successful POST must bump `session_last_activity` so the
    /// idle-timeout watchdog sees fresh traffic. Captures the clock
    /// before + after a tools/list POST that passes session
    /// validation and asserts strict monotonicity.
    #[tokio::test]
    async fn post_updates_last_activity() {
        let server = server();
        let session_id = post_init(&server).await;
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        let session = {
            let sessions = server.sessions.read().await;
            sessions
                .get(&session_uuid)
                .expect("post_init inserts the session into the registry")
                .clone()
        };

        let before = session
            .last_activity_millis
            .load(std::sync::atomic::Ordering::Relaxed);
        // Ensure the monotonic clock advances at least one tick.
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        let body = json!({"jsonrpc": "2.0", "id": 2, "method": "tools/list"});
        let req = Request::builder()
            .method(Method::POST)
            .uri("/mcp")
            .header(header::CONTENT_TYPE, "application/json")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::from(serde_json::to_vec(&body).unwrap()))
            .unwrap();
        let resp = server.router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        let after = session
            .last_activity_millis
            .load(std::sync::atomic::Ordering::Relaxed);
        assert!(
            after > before,
            "session last_activity must advance on successful POST"
        );
    }
}

#[cfg(test)]
mod get_outbound_wiring_tests {
    //! Unit tests for the `GET /mcp` outbound + roots-cache wiring.
    //! A successful GET MUST populate the active session's `outbound`
    //! cell over the same per-session mpsc tx the SSE body drains,
    //! and MUST populate `session.roots_cache` iff the server was
    //! built with `with_client_sampling()`.
    use super::*;
    use axum::body::Body;
    use axum::http::{header, Method, Request, StatusCode};
    use klieo_core::error::ToolError;
    use klieo_core::llm::ToolDef;
    use klieo_core::tool::{ToolCtx, ToolInvoker};
    use serde_json::{json, Value};
    use tower::ServiceExt;

    struct NoopInvoker;

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

    fn server_with_sampling() -> Arc<McpServer> {
        McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .with_client_sampling()
            .build_arc()
            .unwrap()
    }

    fn server_without_sampling() -> Arc<McpServer> {
        McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .build_arc()
            .unwrap()
    }

    async fn post_init(server: &Arc<McpServer>) -> String {
        let body = json!({
            "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}
        });
        let req = Request::builder()
            .method(Method::POST)
            .uri("/mcp")
            .header(header::CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_vec(&body).unwrap()))
            .unwrap();
        let resp = server.router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        resp.headers()
            .get(MCP_SESSION_ID_HEADER)
            .expect("initialize echoes session id")
            .to_str()
            .unwrap()
            .to_string()
    }

    async fn get_mcp_with_session(server: &Arc<McpServer>, session_id: &str) -> StatusCode {
        let req = Request::builder()
            .method(Method::GET)
            .uri("/mcp")
            .header(MCP_SESSION_ID_HEADER, session_id)
            .body(Body::empty())
            .unwrap();
        let resp = server.router().oneshot(req).await.unwrap();
        resp.status()
    }

    #[tokio::test]
    async fn http_get_wires_outbound_primitive() {
        let server = server_with_sampling();
        let session_id = post_init(&server).await;
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        let session = {
            let sessions = server.sessions.read().await;
            sessions
                .get(&session_uuid)
                .expect("post_init inserts the session into the registry")
                .clone()
        };
        assert!(
            session.outbound.get().is_none(),
            "outbound must be unset before GET"
        );
        let status = get_mcp_with_session(&server, &session_id).await;
        assert_eq!(status, StatusCode::OK);
        assert!(
            session.outbound.get().is_some(),
            "GET must populate the outbound primitive"
        );
    }

    #[tokio::test]
    async fn http_get_wires_roots_cache_when_sampling_declared() {
        let server = server_with_sampling();
        let session_id = post_init(&server).await;
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        let session = {
            let sessions = server.sessions.read().await;
            sessions
                .get(&session_uuid)
                .expect("post_init inserts the session into the registry")
                .clone()
        };
        assert!(session.roots_cache.get().is_none());
        let status = get_mcp_with_session(&server, &session_id).await;
        assert_eq!(status, StatusCode::OK);
        assert!(
            session.roots_cache.get().is_some(),
            "with_client_sampling + GET must populate roots_cache"
        );
    }

    #[tokio::test]
    async fn http_get_skips_roots_cache_when_sampling_absent() {
        let server = server_without_sampling();
        let session_id = post_init(&server).await;
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        let session = {
            let sessions = server.sessions.read().await;
            sessions
                .get(&session_uuid)
                .expect("post_init inserts the session into the registry")
                .clone()
        };
        let status = get_mcp_with_session(&server, &session_id).await;
        assert_eq!(status, StatusCode::OK);
        assert!(
            session.outbound.get().is_some(),
            "outbound is wired regardless of sampling"
        );
        assert!(
            session.roots_cache.get().is_none(),
            "roots_cache is gated on declare_sampling"
        );
    }

    /// Dropping the SSE body must fire the [`GuardedSseStream`] cleanup
    /// path: pending outbound oneshots are drained so awaiting callers
    /// resolve with [`klieo_core::ServerOutboundError::TransportClosed`].
    /// Pins ADR-028's "single-session reaping" failure-semantics
    /// guarantee.
    #[tokio::test]
    async fn disconnect_drains_pending_outbound() {
        use klieo_core::{ServerOutbound, ServerOutboundError};

        let server = server_with_sampling();
        let session_id = post_init(&server).await;

        // Open the SSE stream; the response body owns the GuardedSseStream
        // whose Drop fires the cleanup task.
        let req = Request::builder()
            .method(Method::GET)
            .uri("/mcp")
            .header(MCP_SESSION_ID_HEADER, &session_id)
            .body(Body::empty())
            .unwrap();
        let resp = server.router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let session_uuid = uuid::Uuid::parse_str(&session_id).expect("session id is a uuid");
        let outbound = {
            let sessions = server.sessions.read().await;
            sessions
                .get(&session_uuid)
                .and_then(|s| s.outbound.get().cloned())
                .expect("GET must populate outbound primitive")
        };

        // Spawn an outbound caller that parks on its receiver; wait
        // until the pending entry is registered so the drain sees it.
        let call_handle = {
            let outbound = outbound.clone();
            tokio::spawn(async move {
                outbound
                    .outbound_request(
                        "custom/method",
                        Value::Null,
                        std::time::Duration::from_secs(5),
                    )
                    .await
            })
        };
        tokio::time::sleep(std::time::Duration::from_millis(30)).await;

        // Dropping the response drops the body → GuardedSseStream::drop
        // fires cleanup_tx → spawn_session_cleanup awaits + marks the
        // session closed + drains pending outbound oneshots.
        drop(resp);

        let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), call_handle)
            .await
            .expect("outbound_request did not resolve within 2s of disconnect")
            .expect("task panicked");
        assert!(
            matches!(outcome, Err(ServerOutboundError::TransportClosed)),
            "disconnect must surface as TransportClosed; got {outcome:?}"
        );
    }
}

#[cfg(test)]
mod idle_watchdog_tests {
    //! Unit tests for the idle-reaper loop. The reaper is armed on
    //! the first successful `initialize` POST and evicts a session
    //! once `session_last_activity` has been idle longer than the
    //! configured deadline. ADR-028.
    use super::*;
    use axum::body::Body;
    use axum::http::{header, Method, Request, StatusCode};
    use klieo_core::error::ToolError;
    use klieo_core::llm::ToolDef;
    use klieo_core::tool::{ToolCtx, ToolInvoker};
    use serde_json::{json, Value};
    use tower::ServiceExt;

    struct NoopInvoker;

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

    /// POST `initialize` against the router and return the minted
    /// session id parsed out of the `Mcp-Session-Id` response header.
    async fn post_init(server: &Arc<McpServer>) -> uuid::Uuid {
        let body = json!({
            "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}
        });
        let req = Request::builder()
            .method(Method::POST)
            .uri("/mcp")
            .header(header::CONTENT_TYPE, "application/json")
            .body(Body::from(serde_json::to_vec(&body).unwrap()))
            .unwrap();
        let resp = server.router().oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let raw = resp
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .and_then(|v| v.to_str().ok())
            .expect("Mcp-Session-Id header on initialize");
        uuid::Uuid::parse_str(raw).expect("Mcp-Session-Id parses as UUID")
    }

    /// Poll [`McpServer::is_session_closed_by_id`] until it reports
    /// the id has been evicted (returns `None`) or the deadline
    /// elapses. The reaper removes idle sessions from the registry
    /// before marking them closed, so eviction surfaces as `None`.
    async fn wait_for_session_evicted(server: &Arc<McpServer>, id: uuid::Uuid) {
        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
        loop {
            if server.is_session_closed_by_id(id).await.is_none() {
                return;
            }
            if std::time::Instant::now() >= deadline {
                panic!("session {id} never evicted by idle reaper");
            }
            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        }
    }

    /// A session left idle past `session_idle_timeout` MUST observe
    /// the idle reaper evicting it from the registry and draining
    /// pending outbound oneshots. The short tick override drives the
    /// reaper fast enough that the test stays cheap on CI.
    #[tokio::test]
    async fn idle_timeout_fires_after_inactivity() {
        let server = McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .with_session_idle_timeout(std::time::Duration::from_millis(150))
            .with_idle_reaper_tick(std::time::Duration::from_millis(50))
            .build_arc()
            .unwrap();
        let id = post_init(&server).await;
        assert_eq!(server.is_session_closed_by_id(id).await, Some(false));

        wait_for_session_evicted(&server, id).await;
    }

    /// `Duration::ZERO` disables the reaper scan. After an idle
    /// stretch the session must still be live — the reaper loop wakes
    /// but the zero-timeout branch skips the scan.
    #[tokio::test]
    async fn zero_timeout_disables_watchdog() {
        let server = McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .with_session_idle_timeout(std::time::Duration::ZERO)
            .with_idle_reaper_tick(std::time::Duration::from_millis(50))
            .build_arc()
            .unwrap();
        let id = post_init(&server).await;
        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        assert_eq!(
            server.is_session_closed_by_id(id).await,
            Some(false),
            "Duration::ZERO must skip eviction"
        );
    }
}

#[cfg(test)]
mod mint_race_tests {
    //! Direct unit tests for [`mint_session_race_500`]. The
    //! concurrent-initialize integration test
    //! (`tests/http_outbound.rs::concurrent_initialize_yields_one_200_and_rest_409_or_500`)
    //! asserts the UNION `{409, 500}` of race-loser counts but never
    //! proves the 500 branch is reachable. If the helper became dead
    //! code (e.g. `.set()` always loses to `.get()` in the scheduler),
    //! that test would still pass with every loser surfacing as 409.
    //!
    //! These tests pin the helper's wire contract independently of
    //! race timing: status, JSON-RPC code, the stable operator-facing
    //! message, and the id-echo behaviour. A future refactor that
    //! drifts any of these breaks the suite — by design.
    use super::*;
    use axum::body::to_bytes;
    use axum::http::StatusCode;

    async fn extract_envelope(resp: Response) -> (StatusCode, serde_json::Value) {
        let status = resp.status();
        let body_bytes = to_bytes(resp.into_body(), usize::MAX)
            .await
            .expect("response body collects");
        let envelope: serde_json::Value =
            serde_json::from_slice(&body_bytes).expect("response body is JSON-RPC envelope");
        (status, envelope)
    }

    #[tokio::test]
    async fn mint_session_race_500_returns_500_with_stable_wire_envelope() {
        let raw_id = serde_json::json!(7);
        let attempted = uuid::Uuid::from_u128(0x0123_4567_89ab_cdef_0123_4567_89ab_cdef);
        let resp = mint_session_race_500(Some(&raw_id), "active_session", Some(attempted));

        let (status, envelope) = extract_envelope(resp).await;
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(envelope["jsonrpc"], "2.0");
        assert_eq!(envelope["id"], 7);
        assert_eq!(envelope["error"]["code"], JSONRPC_SERVER_ERROR);
        assert_eq!(
            envelope["error"]["message"], "internal: session mint race",
            "stable wire message must not drift — operators key alerting on this string"
        );
    }

    #[tokio::test]
    async fn mint_session_race_500_handles_none_attempted_session_id() {
        // wire_session_outbound passes Option<Uuid> directly from
        // active_session_id() — when None, the helper must still
        // return 500 + the stable envelope (operator log shape
        // differs but the wire shape is invariant).
        let resp = mint_session_race_500(None, "outbound", None);
        let (status, envelope) = extract_envelope(resp).await;
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(envelope["error"]["code"], JSONRPC_SERVER_ERROR);
        assert_eq!(envelope["error"]["message"], "internal: session mint race");
        // raw_id == None → JSON-RPC `id: null` (the rpc_error builder
        // emits Null when no id was provided).
        assert_eq!(envelope["id"], serde_json::Value::Null);
    }

    #[tokio::test]
    async fn mint_session_race_500_distinct_message_from_session_already_active() {
        // Negative fixture: confirm the 500-helper's wire message is
        // NOT the same as the legitimate 409 "session already
        // active" path. If a future refactor accidentally collapses
        // both paths to the same string, operators lose the ability
        // to distinguish client mistake from server invariant
        // violation.
        let resp = mint_session_race_500(None, "active_session", None);
        let (_, envelope) = extract_envelope(resp).await;
        assert_ne!(
            envelope["error"]["message"], "session already active",
            "race-500 path must remain distinct from the .get().is_some() 409 path"
        );
    }
}

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

    #[test]
    fn both_none_passes() {
        // Auth-disabled deployment: no authenticator wired on the
        // request, no principal recorded on the session. Every
        // caller can reach every session — the backwards-compatible
        // shape for servers that never opted into auth.
        assert!(principal_matches(None, None));
    }

    #[test]
    fn matching_some_passes() {
        // Verified caller equals the principal stamped on the session
        // at initialize time — the only `Some, Some` shape that may
        // pass.
        let caller = Identity::new("alice");
        assert!(principal_matches(Some(&caller), Some("alice")));
    }

    #[test]
    fn mismatched_some_fails() {
        // Cross-tenant attempt: a verified peer presenting another
        // principal's session id. CWE-639 boundary — must reject so
        // the caller path returns 404 rather than honour the request.
        let caller = Identity::new("alice");
        assert!(!principal_matches(Some(&caller), Some("bob")));
    }

    #[test]
    fn caller_authenticated_session_anonymous_fails() {
        // Authenticator wired AFTER an anonymous session was minted.
        // The session carries no principal; rejecting denies the
        // verified caller without forcing operators to drop every
        // pre-auth session when policy flips on.
        let caller = Identity::new("alice");
        assert!(!principal_matches(Some(&caller), None));
    }

    #[test]
    fn caller_anonymous_session_authenticated_fails() {
        // Authenticator removed AFTER an authenticated session was
        // minted. An anonymous caller cannot reach a session bound
        // to a specific principal — the binding outlives the policy
        // change.
        assert!(!principal_matches(None, Some("alice")));
    }
}

#[cfg(test)]
mod run_resume_authz_tests {
    //! Direct tests for the `klieo/run/resume` authorization seam
    //! (`claim_resume_record`): anonymous, foreign-principal, and
    //! unknown-ticket callers all deny with the same opaque envelope and
    //! do NOT consume the ticket (IDOR / CWE-639); the rightful owner
    //! then claims exactly once, and a replayed claim loses the race.
    use super::*;
    use crate::resume_ticket::{ResumeTicketRecord, ResumeTicketStore};
    use async_trait::async_trait;
    use axum::body::to_bytes;
    use klieo_bus_memory::MemoryKv;
    use klieo_core::error::ToolError;
    use klieo_core::llm::ToolDef;
    use klieo_core::tool::{ToolCtx, ToolInvoker};
    use serde_json::Value;

    struct NoopInvoker;

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

    fn seeded_record(principal: &str) -> ResumeTicketRecord {
        let cp = serde_json::json!({
            "run_id": klieo_core::ids::RunId::new(),
            "step_index": 1,
            "thread_id": "t-authz",
            "messages": [],
            "pending_tool_calls": null,
            "created_at": "2026-06-18T00:00:00Z",
        });
        ResumeTicketRecord {
            principal: principal.into(),
            workflow_name: "wf".into(),
            checkpoint: serde_json::from_value(cp).unwrap(),
            created_at: chrono::Utc::now(),
        }
    }

    /// Build a server whose resume-ticket store shares `kv`, so a ticket
    /// seeded through a sibling store over the same KV is visible to the
    /// handler.
    fn server_over(kv: Arc<MemoryKv>) -> Arc<McpServer> {
        McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .with_checkpoint_kv(kv)
            .build_arc()
            .unwrap()
    }

    async fn deny_message(resp: Response) -> String {
        let body = to_bytes(resp.into_body(), usize::MAX)
            .await
            .expect("deny body collects");
        let json: serde_json::Value =
            serde_json::from_slice(&body).expect("deny body is JSON-RPC");
        json["error"]["message"]
            .as_str()
            .expect("deny envelope carries an error message")
            .to_string()
    }

    #[tokio::test]
    async fn anonymous_caller_is_denied_and_ticket_survives() {
        let kv = Arc::new(MemoryKv::new());
        let server = server_over(kv.clone());
        let store = ResumeTicketStore::new(kv);
        let token = ResumeTicketStore::mint_token();
        store
            .persist(&token, &seeded_record("alice@x"))
            .await
            .unwrap();

        let resp = claim_resume_record(&server, &token, Some(Identity::anonymous()), None)
            .await
            .expect_err("anonymous caller must be denied");
        assert_eq!(deny_message(resp).await, RUN_RESUME_DENY_MESSAGE);

        let owner =
            claim_resume_record(&server, &token, Some(Identity::new("alice@x")), None).await;
        assert!(
            owner.is_ok(),
            "an anonymous denial must not consume the ticket"
        );
    }

    #[tokio::test]
    async fn foreign_principal_is_denied_and_owner_still_claims() {
        let kv = Arc::new(MemoryKv::new());
        let server = server_over(kv.clone());
        let store = ResumeTicketStore::new(kv);
        let token = ResumeTicketStore::mint_token();
        store
            .persist(&token, &seeded_record("alice@x"))
            .await
            .unwrap();

        let resp = claim_resume_record(&server, &token, Some(Identity::new("mallory@x")), None)
            .await
            .expect_err("foreign principal must be denied (IDOR)");
        assert_eq!(deny_message(resp).await, RUN_RESUME_DENY_MESSAGE);

        let owner = claim_resume_record(&server, &token, Some(Identity::new("alice@x")), None)
            .await
            .expect("rightful owner claims after the foreign denial");
        assert_eq!(owner.principal, "alice@x");
    }

    #[tokio::test]
    async fn unknown_ticket_is_denied_with_opaque_message() {
        let kv = Arc::new(MemoryKv::new());
        let server = server_over(kv);
        let resp =
            claim_resume_record(&server, "no-such-token", Some(Identity::new("alice@x")), None)
                .await
                .expect_err("unknown ticket must be denied");
        assert_eq!(deny_message(resp).await, RUN_RESUME_DENY_MESSAGE);
    }

    #[tokio::test]
    async fn owner_claims_exactly_once_replay_loses_race() {
        let kv = Arc::new(MemoryKv::new());
        let server = server_over(kv.clone());
        let store = ResumeTicketStore::new(kv);
        let token = ResumeTicketStore::mint_token();
        store
            .persist(&token, &seeded_record("alice@x"))
            .await
            .unwrap();

        let first =
            claim_resume_record(&server, &token, Some(Identity::new("alice@x")), None).await;
        assert!(first.is_ok(), "the first claim by the owner succeeds");

        let replay = claim_resume_record(&server, &token, Some(Identity::new("alice@x")), None)
            .await
            .expect_err("a replayed claim must lose the race");
        assert_eq!(deny_message(replay).await, RUN_RESUME_DENY_MESSAGE);
    }

    /// A `WorkflowResumeHandle` stub whose outcome the test controls,
    /// so `drive_resume`'s success vs. failure routing is exercised
    /// without standing up a real workflow.
    struct StubResumeHandle {
        fail: bool,
    }

    #[async_trait]
    impl crate::workflow::WorkflowResumeHandle for StubResumeHandle {
        async fn resume(
            &self,
            _checkpoint: klieo_core::checkpoint::RunCheckpoint,
            _decision: klieo_core::checkpoint::ApprovalDecision,
            _tenant_label: String,
        ) -> Result<Value, ToolError> {
            if self.fail {
                Err(ToolError::Permanent("resume blew up".into()))
            } else {
                Ok(serde_json::json!({ "resumed": true }))
            }
        }
    }

    fn server_with_handle(fail: bool) -> McpServer {
        let mut server = McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .build()
            .unwrap();
        server.workflow_resume_handles.insert(
            "wf".to_string(),
            Arc::new(StubResumeHandle { fail })
                as Arc<dyn crate::workflow::WorkflowResumeHandle>,
        );
        server
    }

    async fn result_value(resp: Response) -> serde_json::Value {
        let body = to_bytes(resp.into_body(), usize::MAX)
            .await
            .expect("body collects");
        serde_json::from_slice(&body).expect("body is JSON-RPC")
    }

    #[tokio::test]
    async fn drive_resume_unregistered_workflow_is_unavailable() {
        let server = McpServer::builder()
            .add_tools(Arc::new(NoopInvoker))
            .build()
            .unwrap();
        let decision = RunResumeDecision {
            approved: true,
            reason: None,
        };
        // `seeded_record` names workflow "wf"; the server has no handles.
        let resp = drive_resume(&server, decision, seeded_record("alice@x"), None).await;
        assert_eq!(deny_message(resp).await, RUN_RESUME_UNAVAILABLE_MESSAGE);
    }

    #[tokio::test]
    async fn drive_resume_dispatches_to_registered_handle() {
        let server = server_with_handle(false);
        let decision = RunResumeDecision {
            approved: true,
            reason: None,
        };
        let resp = drive_resume(&server, decision, seeded_record("alice@x"), None).await;
        let body = result_value(resp).await;
        assert_eq!(
            body["result"]["resumed"],
            serde_json::Value::Bool(true),
            "approved resume returns the handle's result envelope; got {body}"
        );
    }

    #[tokio::test]
    async fn drive_resume_handle_error_yields_server_error() {
        let server = server_with_handle(true);
        let decision = RunResumeDecision {
            approved: false,
            reason: Some("operator rejected".into()),
        };
        let resp = drive_resume(&server, decision, seeded_record("alice@x"), None).await;
        // Server error envelope, NOT the opaque deny string — the
        // failure is post-authz execution, not an authz refusal.
        assert_eq!(deny_message(resp).await, "resume execution failed");
    }
}