klieo-a2a 3.4.0

//! HTTP/SSE transport for `A2aDispatcher`. Streamable HTTP shape
//! per A2A v1.0 §9.4.2: single `POST /a2a`; server picks JSON vs
//! SSE response based on JSON-RPC method.
//!
//! Gated behind the `http` cargo feature. See ADR-013 for the design
//! rationale and out-of-scope items (SSE resumption, multi-replica
//! fanout, HTTP-layer auth).

use crate::auth::RequestContext;
use crate::envelope::{codes, A2aHeaders, A2aMethod, JsonRpcError, JsonRpcResponse};
use crate::error::{A2aBuilderError, A2aError};
use crate::handler::A2aHandler;
use crate::server::{A2aDispatcher, TaskEvent, TaskEventStream};
use crate::task_store::{A2aTaskStore, DEFAULT_BUCKET};
use axum::{
    body::Bytes,
    extract::{DefaultBodyLimit, State},
    http::{header, HeaderMap, StatusCode},
    response::{sse::Event, IntoResponse, Response, Sse},
    routing::post,
    Json, Router,
};
use futures::StreamExt;
use klieo_auth_common::Authenticator;
use klieo_core::{KvStore, Pubsub};
use serde_json::Value;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

const MAX_BODY_BYTES: usize = 1 << 20; // 1 MiB
const CANCEL_SUBJECT_PREFIX: &str = "klieo.a2a.cancel.";

/// HTTP/SSE server wrapping an [`A2aDispatcher`]. Owns the task store
/// used by `SubscribeToTask` (for current-state replay) and a parent
/// cancel token for graceful shutdown.
pub struct A2aHttpServer {
    pub(crate) dispatcher: Arc<A2aDispatcher>,
    pub(crate) task_store: Arc<A2aTaskStore>,
    pub(crate) parent_cancel: CancellationToken,
    pub(crate) resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    allow_public_bind: bool,
    // Held for Drop side effect — Drop aborts the background scan
    // task. `None` when [`A2aDispatcherBuilder::with_kv_reaper`] was
    // not called or [`Self::with_resume_buffer`] was never wired.
    _kv_reaper: Option<klieo_core::KvReaperHandle>,
}

impl A2aHttpServer {
    /// Open an [`A2aHttpServerBuilder`].
    ///
    /// Prefer this over hand-assembling [`A2aDispatcher`] +
    /// [`A2aTaskStore`] + [`Self::new`]: the builder wires the task
    /// store's event sink to the dispatcher's pubsub unconditionally,
    /// so the silent-streaming trap (a task store built without
    /// `with_event_sink` whose `emit` no-ops, hanging every
    /// `SubscribeToTask` subscriber forever) is impossible to hit.
    pub fn builder() -> A2aHttpServerBuilder {
        A2aHttpServerBuilder::default()
    }

    /// Build a new HTTP/SSE server.
    ///
    /// Defaults to [`klieo_core::resume::NoopResumeBuffer`] — no SSE
    /// resumption. Call [`Self::with_resume_buffer`] to opt in.
    pub fn new(
        dispatcher: Arc<A2aDispatcher>,
        task_store: Arc<A2aTaskStore>,
        parent_cancel: CancellationToken,
    ) -> Self {
        Self {
            dispatcher,
            task_store,
            parent_cancel,
            resume_buffer: Arc::new(klieo_core::resume::NoopResumeBuffer),
            allow_public_bind: false,
            _kv_reaper: None,
        }
    }

    /// Opt in to SSE resumption via `Last-Event-ID` / `klieo/tools/resume`.
    ///
    /// Default is [`klieo_core::resume::NoopResumeBuffer`] (zero-cost no-op).
    /// Pass a [`klieo_core::resume::KvResumeBuffer`] (or any custom impl)
    /// to enable live resumption.
    ///
    /// Also wires the cluster-0.25 KV reaper when the dispatcher was
    /// built with [`crate::server::A2aDispatcherBuilder::with_kv_reaper`].
    /// The reaper scans the `klieo-leaders` (+ `klieo-tenants` when
    /// tenant binding is wired) buckets at the configured interval
    /// and evicts entries whose corresponding resume buffer reports
    /// terminal. See ADR-025.
    #[must_use]
    pub fn with_resume_buffer(mut self, buffer: Arc<dyn klieo_core::resume::ResumeBuffer>) -> Self {
        self.resume_buffer = buffer.clone();
        self._kv_reaper = self.spawn_kv_reaper_if_configured(buffer);
        self
    }

    fn spawn_kv_reaper_if_configured(
        &self,
        buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    ) -> Option<klieo_core::KvReaperHandle> {
        let interval = self.dispatcher.kv_reaper_interval()?;
        let leader_registry = self.dispatcher.leader_registry()?;
        let kv = leader_registry.kv().clone();
        let mut buckets = vec![leader_registry.bucket().to_string()];
        if let Some(ownership) = self.dispatcher.ownership_registry() {
            buckets.push(ownership.bucket().to_string());
        }
        Some(klieo_core::spawn_kv_reaper(kv, buffer, buckets, interval))
    }

    /// Opt in to binding non-loopback addresses when the dispatcher's
    /// `Authenticator` returns `allows_anonymous() == true`.
    ///
    /// Without this, [`Self::serve_http`] returns
    /// [`crate::error::A2aError::Misconfigured`] when asked to bind a
    /// non-loopback address with an anonymous authenticator.
    ///
    /// Use ONLY when the dispatcher's authenticator is a real auth mechanism
    /// (e.g. `BearerTokenAuthenticator`) AND the service is
    /// fronted by an auth-enforcing reverse proxy. Calling this with
    /// `AllowAnonymous` wired and no reverse proxy is a
    /// security bug.
    pub fn allow_public_bind(mut self) -> Self {
        self.allow_public_bind = true;
        self
    }

    /// Borrow the task store (for test helpers that need to pre-populate
    /// tasks before issuing `SubscribeToTask` requests).
    pub fn task_store(&self) -> &Arc<A2aTaskStore> {
        &self.task_store
    }

    /// Build the axum [`Router`]. `POST /a2a` route with JSON-or-SSE
    /// response negotiation, plus `GET /.well-known/agent-card.json`
    /// for unauthenticated discovery (A2A spec §5 / IANA well-known-URI
    /// registration). 1 MiB body cap + content-type guard applied to
    /// `/a2a`; caller owns CORS / TLS / rate-limiting via additional
    /// layers.
    ///
    /// # Security
    /// `/a2a` adds no authentication beyond what the configured
    /// `Authenticator` on the dispatcher enforces. Bind
    /// `127.0.0.1` unless fronted by an auth-enforcing reverse proxy.
    /// `/.well-known/agent-card.json` is UNAUTHENTICATED BY PROTOCOL
    /// DESIGN — see [`crate::handler::A2aHandler::get_agent_card`] for
    /// why it must never return auth-gated detail.
    pub fn router(self: &Arc<Self>) -> Router {
        Router::new()
            .route("/a2a", post(post_a2a))
            .layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
            .route(
                "/.well-known/agent-card.json",
                axum::routing::get(get_agent_card),
            )
            .with_state(self.clone())
    }

    /// Bind to `addr` and serve until `parent_cancel` fires or the
    /// listener errors out.
    ///
    /// # Security
    /// No TLS — terminate at a reverse proxy. Bind `127.0.0.1` unless
    /// the listener is fronted by an auth-enforcing proxy.
    pub async fn serve_http(self: Arc<Self>, addr: SocketAddr) -> Result<(), A2aError> {
        if !addr.ip().is_loopback()
            && self.dispatcher.authenticator().allows_anonymous()
            && !self.allow_public_bind
        {
            tracing::error!(
                target: "a2a",
                %addr,
                "refusing non-loopback bind with anonymous authenticator",
            );
            return Err(A2aError::Misconfigured(format!(
                "refusing to bind non-loopback address {addr} with anonymous authenticator; \
                 call `allow_public_bind()` to override \
                 (only safe behind an auth-enforcing reverse proxy)",
            )));
        }
        let cancel = self.parent_cancel.clone();
        let listener = tokio::net::TcpListener::bind(addr)
            .await
            .map_err(|e| A2aError::Server(e.to_string()))?;
        let router = self.router();
        axum::serve(listener, router)
            .with_graceful_shutdown(async move { cancel.cancelled().await })
            .await
            .map_err(|e| A2aError::Server(e.to_string()))?;
        Ok(())
    }
}

/// Builder for [`A2aHttpServer`] that closes the event-sink wiring trap.
///
/// Assembling the server by hand means constructing the dispatcher, the
/// task store, and remembering `A2aTaskStore::with_event_sink(dispatcher
/// .event_sink())`. Omit that one call and the store's `emit` silently
/// no-ops, so `SubscribeToTask` streaming subscribers hang forever with
/// no error. This builder wires the sink from the same dispatcher it
/// builds, unconditionally, so the trap cannot be hit. Setter naming
/// mirrors [`crate::server::A2aDispatcherBuilder`].
#[derive(Default)]
pub struct A2aHttpServerBuilder {
    handler: Option<Arc<dyn A2aHandler>>,
    authenticator: Option<Arc<dyn Authenticator>>,
    kv: Option<Arc<dyn KvStore>>,
    bucket: Option<String>,
    cancel: Option<CancellationToken>,
    pubsub: Option<Arc<dyn Pubsub>>,
}

impl A2aHttpServerBuilder {
    /// Set the [`A2aHandler`] the dispatcher delegates requests to.
    pub fn handler(mut self, handler: Arc<dyn A2aHandler>) -> Self {
        self.handler = Some(handler);
        self
    }

    /// Set the [`Authenticator`] enforced at every request boundary.
    pub fn authenticator(mut self, authenticator: Arc<dyn Authenticator>) -> Self {
        self.authenticator = Some(authenticator);
        self
    }

    /// Set the [`KvStore`] backing the task store. Required — there is
    /// no sensible default backend for durable task persistence.
    pub fn kv(mut self, kv: Arc<dyn KvStore>) -> Self {
        self.kv = Some(kv);
        self
    }

    /// Override the task-store KV bucket. Defaults to
    /// [`DEFAULT_BUCKET`] when unset.
    pub fn bucket(mut self, bucket: String) -> Self {
        self.bucket = Some(bucket);
        self
    }

    /// Override the parent cancel token used for graceful shutdown.
    /// Defaults to a fresh [`CancellationToken`] when unset.
    pub fn cancel(mut self, cancel: CancellationToken) -> Self {
        self.cancel = Some(cancel);
        self
    }

    /// Wire a shared [`Pubsub`] for cross-replica fanout. Multi-replica
    /// deployments pass a NATS-backed pubsub here so task events fan out
    /// across replicas; tests pass a known bus to observe emitted events.
    /// Defaults to a fresh in-process bus when unset.
    pub fn pubsub(mut self, pubsub: Arc<dyn Pubsub>) -> Self {
        self.pubsub = Some(pubsub);
        self
    }

    /// Finalise into an [`A2aHttpServer`] with the task-store event sink
    /// wired to the built dispatcher's pubsub by construction.
    ///
    /// Returns [`A2aBuilderError::MissingHandler`],
    /// [`A2aBuilderError::MissingAuthenticator`], or
    /// [`A2aBuilderError::MissingKv`] when the corresponding required
    /// field was not set.
    pub fn build(self) -> Result<A2aHttpServer, A2aBuilderError> {
        let handler = self.handler.ok_or(A2aBuilderError::MissingHandler)?;
        let authenticator = self
            .authenticator
            .ok_or(A2aBuilderError::MissingAuthenticator)?;
        let kv = self.kv.ok_or(A2aBuilderError::MissingKv)?;

        let dispatcher_builder = A2aDispatcher::builder()
            .handler(handler)
            .authenticator(authenticator);
        let dispatcher_builder = match self.pubsub {
            Some(pubsub) => dispatcher_builder.pubsub(pubsub),
            None => dispatcher_builder.with_in_process_pubsub(),
        };
        let dispatcher = dispatcher_builder.build_arc()?;

        let bucket = self.bucket.unwrap_or_else(|| DEFAULT_BUCKET.to_string());
        let task_store =
            Arc::new(A2aTaskStore::new(kv, bucket).with_event_sink(dispatcher.event_sink()));

        let cancel = self.cancel.unwrap_or_default();
        Ok(A2aHttpServer::new(dispatcher, task_store, cancel))
    }
}

#[tracing::instrument(
    skip_all,
    fields(
        rpc.system = "klieo-a2a",
        rpc.method = tracing::field::Empty,
        http.request.method = "POST",
    ),
)]
async fn post_a2a(
    State(server): State<Arc<A2aHttpServer>>,
    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 shutdown_response();
    }

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

    if headers.get_all(header::AUTHORIZATION).iter().count() >= 2 {
        warn!(target: "a2a", "rejected request with duplicate Authorization header");
        return (
            StatusCode::BAD_REQUEST,
            Json(error_envelope(
                serde_json::Value::Null,
                codes::INVALID_REQUEST,
                "duplicate Authorization header",
            )),
        )
            .into_response();
    }

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

    let method = raw.get("method").and_then(|m| m.as_str()).unwrap_or("");
    let req_id = raw.get("id").cloned().unwrap_or(Value::Null);
    tracing::Span::current().record("rpc.method", method);

    if is_streaming_method(method) {
        dispatch_streaming(&server, headers, body, req_id, method.to_owned()).await
    } else {
        dispatch_json(&server, headers, body).await
    }
}

/// `GET /.well-known/agent-card.json` — A2A spec §5 discovery endpoint.
///
/// # Security
/// UNAUTHENTICATED BY PROTOCOL DESIGN: a client fetches this before it
/// knows which auth scheme the agent requires, so it is served with
/// `ctx.caller = None` regardless of the dispatcher's configured
/// `Authenticator` — that authenticator gates `/a2a`
/// only. [`crate::handler::A2aHandler::get_agent_card`] implementations
/// MUST NOT return anything here that isn't safe for an anonymous
/// caller (auth-gated detail belongs behind `GetExtendedAgentCard`
/// instead, which IS authenticated).
///
/// A handler that hasn't implemented [`A2aHandler::get_agent_card`]
/// (the default `MethodNotFound`) yields `404 Not Found` — a plain
/// HTTP status, not a JSON-RPC envelope, since this route is a REST
/// resource rather than a JSON-RPC method.
async fn get_agent_card(State(server): State<Arc<A2aHttpServer>>, headers: HeaderMap) -> Response {
    let a2a_headers = axum_headers_to_a2a(&headers);
    let ctx = RequestContext::new(a2a_headers, None);
    match server.dispatcher.handler().get_agent_card(&ctx).await {
        Ok(card) => (StatusCode::OK, Json(card)).into_response(),
        Err(A2aError::MethodNotFound(_)) => StatusCode::NOT_FOUND.into_response(),
        Err(err) => {
            warn!(target: "a2a", error = %err, "agent-card discovery handler failed");
            StatusCode::INTERNAL_SERVER_ERROR.into_response()
        }
    }
}

async fn dispatch_json(server: &Arc<A2aHttpServer>, headers: HeaderMap, body: Bytes) -> Response {
    let a2a_headers = axum_headers_to_a2a(&headers);
    let resp = server.dispatcher.handle_request(a2a_headers, &body).await;
    (StatusCode::OK, Json(resp)).into_response()
}

async fn dispatch_streaming(
    server: &Arc<A2aHttpServer>,
    headers: HeaderMap,
    body: Bytes,
    req_id: Value,
    method: String,
) -> Response {
    let a2a_headers = axum_headers_to_a2a(&headers);
    let last_event_id = last_event_id_from(&headers);
    let request_cancel = server.parent_cancel.child_token();
    // Extract the task id for SubscribeToTask so we can register the
    // per-invoke cancel token + publish a cross-replica cancel on
    // body drop. SendStreamingMessage does not yet know the task id
    // at this point (the handler mints it); the registry hook there
    // is deferred until the synthesised stream knows the id.
    let task_id = task_id_from_body(&method, &body);
    match server
        .dispatcher
        .handle_streaming(
            a2a_headers,
            &body,
            &server.task_store,
            request_cancel.clone(),
            last_event_id,
            server.resume_buffer.clone(),
        )
        .await
    {
        Ok(stream) => build_sse_response(
            stream,
            req_id,
            request_cancel,
            server.dispatcher.clone(),
            server.task_store.clone(),
            server.resume_buffer.clone(),
            task_id,
        ),
        Err(A2aError::ResumeBufferExpired { since_id }) => (
            StatusCode::OK,
            Json(error_envelope(
                req_id,
                codes::RESUME_BUFFER_EXPIRED,
                &format!("resume window expired (since_id={since_id})"),
            )),
        )
            .into_response(),
        Err(A2aError::Unauthorized(_)) => {
            // Auth failure must not open an SSE body — return JSON envelope.
            (
                StatusCode::OK,
                Json(error_envelope(
                    req_id,
                    codes::UNAUTHENTICATED,
                    "Authentication required",
                )),
            )
                .into_response()
        }
        Err(err) => {
            // Log server-side BEFORE the wire envelope sanitises Display —
            // mirrors the NATS path in A2aDispatcher::dispatch so both
            // seams get the same error-class-aware severity split
            // (Bus/Server/Misconfigured/Internal -> error!,
            // client-class -> warn!).
            crate::server::log_internal_before_wire_seam(&err, &method);
            if let A2aError::LeaderDied { stream_id } = &err {
                return (
                    StatusCode::OK,
                    Json(leader_died_envelope(req_id, &err.to_string(), stream_id)),
                )
                    .into_response();
            }
            let (code, msg) = match &err {
                A2aError::InvalidParams(m) => (codes::INVALID_PARAMS, m.clone()),
                A2aError::MethodNotFound(m) => {
                    (codes::METHOD_NOT_FOUND, format!("method not found: {m}"))
                }
                // All other variants: sanitise on wire so A2aError::Bus /
                // Server payloads cannot leak NATS URLs, bucket names, or
                // OS-level I/O details to the peer.
                _ => (codes::SERVER_ERROR, "internal server error".into()),
            };
            (StatusCode::OK, Json(error_envelope(req_id, code, &msg))).into_response()
        }
    }
}

/// Stream wrapper that holds a `DropGuard` whose `Drop` impl fires the
/// request-scoped `CancellationToken`. When axum/hyper drops the
/// response body (client TCP close), the wrapped stream drops and the
/// guard cancels the token, propagating through `RequestContext.cancel`.
///
/// 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 task id 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 task id can cause a cancel on the owning replica. Operators
/// MUST mint unguessable task ids and gate per-tenant authorisation
/// BEFORE the request reaches the dispatcher; 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 the dispatcher's
    /// [`crate::server::A2aDispatcher::publish_permits`] so the
    /// drop-time cross-replica cancel publish bounds under the same
    /// cap as `TaskEventSink::send` fanout. Threaded through
    /// [`klieo_core::cancel::spawn_drop_publish`] which `try_acquire`s
    /// before spawning; saturation drops the publish with a `warn`
    /// (local cancel still fires via the conventional drop-of-`_guard`
    /// path).
    permits: Arc<tokio::sync::Semaphore>,
}

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),
            "a2a.cancel",
            Some(self.permits.clone()),
            trace_headers,
        );
    }
}

/// Extract the task id from a streaming JSON-RPC request body so the
/// HTTP layer can register the per-invoke cancel token before the
/// stream is built.  Returns the empty string when the method does
/// not embed a task id at this stage (e.g. `SendStreamingMessage`)
/// or when parsing fails — callers treat empty as "skip registration".
fn task_id_from_body(method: &str, body: &Bytes) -> String {
    // Alias-aware for the same reason as `is_streaming_method`: a
    // `tasks/resubscribe` request must have its task id extracted too.
    if !matches!(A2aMethod::from_str(method), Ok(A2aMethod::SubscribeToTask)) {
        return String::new();
    }
    let parsed: Result<Value, _> = serde_json::from_slice(body);
    parsed
        .ok()
        .and_then(|v| v.get("params")?.get("id")?.as_str().map(str::to_owned))
        .unwrap_or_default()
}

fn build_sse_response(
    stream: TaskEventStream,
    req_id: Value,
    cancel: CancellationToken,
    dispatcher: Arc<A2aDispatcher>,
    task_store: Arc<A2aTaskStore>,
    resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    task_id: String,
) -> Response {
    // Register the per-invoke cancel token under task_id so the
    // wildcard cancel-subject subscription can fire it on inbound
    // klieo.a2a.cancel.{task_id} messages. Skip on empty task_id.
    if !task_id.is_empty() {
        dispatcher
            .cancel_registry()
            .register(task_id.clone(), cancel.clone());
    }
    let registry_handle = dispatcher.cancel_registry().clone();
    let pubsub = dispatcher.pubsub().clone();
    let permits = dispatcher.publish_permits().clone();
    let cancel_subject = if task_id.is_empty() {
        String::new()
    } else {
        format!("{CANCEL_SUBJECT_PREFIX}{task_id}")
    };
    let deregistered =
        klieo_core::cancel::RegistryDeregisterOnDrop::new(stream, registry_handle, task_id);
    let mapped = deregistered.map(move |mut event: TaskEvent| {
        // Pre-stamped events (e.g. from task_store::put broadcast path)
        // carry a non-zero id — preserve them to maintain monotonicity in
        // the resume replay branch. Mint a fresh id only for events that
        // arrive without one (e.g. handler-supplied streams).
        if event.event_id == 0 {
            event.event_id = task_store.next_event_id(&event.task_id);
        }
        let id = event.event_id;
        let final_event = event.final_event;
        let task_id = event.task_id.clone();
        let payload = bytes::Bytes::from(serde_json::to_vec(&event).unwrap_or_default());
        let buffer = resume_buffer.clone();
        tokio::spawn(async move {
            if let Err(e) = buffer.record(&task_id, id, payload).await {
                warn!(
                    target: "a2a.resume",
                    task_id = %task_id,
                    id,
                    error = %e,
                    "resume buffer record failed",
                );
            }
            if final_event {
                if let Err(e) = buffer.close(&task_id).await {
                    warn!(
                        target: "a2a.resume",
                        task_id = %task_id,
                        error = %e,
                        "resume buffer close failed",
                    );
                }
            }
        });
        task_event_to_sse_frame(event, &req_id, id)
    });
    let guarded = CancelOnDrop {
        inner: mapped,
        _guard: cancel.drop_guard(),
        pubsub,
        cancel_subject,
        permits,
    };
    (StatusCode::OK, Sse::new(guarded)).into_response()
}

fn task_event_to_sse_frame(
    event: TaskEvent,
    req_id: &Value,
    seq: u64,
) -> Result<Event, Infallible> {
    let payload = serde_json::json!({
        "jsonrpc": "2.0",
        "id": req_id,
        "result": {
            "task_id": event.task_id,
            "status": event.status,
            "message": event.message,
            "final": event.final_event,
        },
    });
    let data = match serde_json::to_string(&payload) {
        Ok(s) => s,
        Err(e) => {
            warn!(
                target: "a2a",
                task_id = %event.task_id,
                error = %e,
                "sse frame serialise failed; emitting comment",
            );
            // Emit a comment-only event; SSE clients ignore comments.
            return Ok(Event::default().comment("serialise-fail"));
        }
    };
    Ok(Event::default()
        .event("task-update")
        .id(seq.to_string())
        .data(data))
}

fn shutdown_response() -> Response {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(JsonRpcResponse {
            jsonrpc: "2.0".into(),
            id: Value::Null,
            result: None,
            error: Some(JsonRpcError {
                code: codes::SERVER_ERROR,
                message: "server shutting down".into(),
                data: None,
            }),
        }),
    )
        .into_response()
}

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)
}

/// Routes on the RESOLVED [`crate::envelope::A2aMethod`] rather than the
/// raw request string, so a legacy slash-form alias (`message/stream`,
/// `tasks/resubscribe`) is recognised as streaming exactly like its
/// current CamelCase form. Matching the raw string here would silently
/// route an aliased streaming request to the non-streaming dispatcher,
/// which rejects both `SendStreamingMessage` and `SubscribeToTask` with
/// `MethodNotFound` by design (see `A2aDispatcher::dispatch`) — so an
/// un-normalised check would make every alias streaming call fail.
fn is_streaming_method(method: &str) -> bool {
    matches!(
        A2aMethod::from_str(method),
        Ok(A2aMethod::SendStreamingMessage) | Ok(A2aMethod::SubscribeToTask)
    )
}

/// Parse `Last-Event-ID` header as a monotonic event counter.
///
/// Returns `None` when the header is absent, non-UTF-8, or non-numeric.
/// Invalid values are silently ignored (client starts fresh).
fn last_event_id_from(headers: &HeaderMap) -> Option<u64> {
    headers
        .get("last-event-id")
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.trim().parse::<u64>().ok())
}

/// 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 axum_headers_to_a2a(headers: &HeaderMap) -> A2aHeaders {
    let mut kheaders = klieo_core::Headers::default();
    for (name, value) in headers.iter() {
        if let Ok(v) = value.to_str() {
            kheaders.insert(name.as_str().to_string(), v.to_string());
        }
    }
    A2aHeaders::decode_from(&kheaders)
}

fn error_envelope(id: Value, code: i32, message: &str) -> Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": { "code": code, "message": message },
    })
}

/// JSON-RPC error envelope for [`A2aError::LeaderDied`]. Distinct
/// from [`error_envelope`] because the `data` field carries the
/// leader-registry stream id so the client can correlate the
/// orphan signal with operator-side bus telemetry. ADR-020.
fn leader_died_envelope(id: Value, message: &str, stream_id: &str) -> Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": {
            "code": codes::LEADER_DIED,
            "message": message,
            "data": { "stream_id": stream_id },
        },
    })
}