klieo_a2a/

server.rs

//! `A2aServer` — listens on a NATS-style request/reply subject and
//! dispatches incoming JSON-RPC payloads to an [`A2aHandler`].
//!
//! # Transport contract
//!
//! - **Subject:** `<app_prefix>.a2a.<agent_id>.rpc`. Subscribe via
//!   [`klieo_core::Pubsub::subscribe`] with a fixed durable name.
//! - **Reply-to:** the inbound message MUST carry a `"Reply-To"` header
//!   naming the subject the client is listening on. The server publishes
//!   the [`JsonRpcResponse`] there.
//! - **Headers:** `Authorization`, `Content-Type`, `A2A-Version`,
//!   `A2A-Extensions` decoded via [`crate::envelope::A2aHeaders`].
//!
//! See module docs in `lib.rs` for why the server takes
//! `Arc<dyn Pubsub>` instead of `Arc<dyn RequestReply>`.

use crate::auth::RequestContext;
use crate::envelope::{
    codes, A2aHeaders, A2aMethod, JsonRpcError, JsonRpcRequest, JsonRpcResponse,
};
use crate::error::{A2aBuilderError, A2aError};
use crate::handler::A2aHandler;
use crate::types::{
    CancelTaskParams, DeletePushNotificationConfigParams, GetExtendedAgentCardParams,
    GetPushNotificationConfigParams, GetTaskParams, ListPushNotificationConfigsParams,
    ListTasksParams, Message, PushNotificationConfigParams, SendMessageParams, SendMessageResult,
    SubscribeToTaskParams, Task, TaskStatus,
};
use bytes::Bytes;
use klieo_auth_common::Authenticator;
use klieo_core::{DurableName, Headers, Msg, Pubsub};

const A2A_CANCEL_SUBJECT_PREFIX: &str = "klieo.a2a.cancel.";
const A2A_CANCEL_SUBJECT_PATTERN: &str = "klieo.a2a.cancel.>";
const A2A_CANCEL_LOG_TARGET: &str = "a2a.cancel";
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json::Value;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio_stream::StreamExt;
use tracing::{error, instrument, warn};
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

/// Default per-dispatcher cap on concurrent in-flight
/// [`TaskEventSink::send`] publishes. Bounds the runtime work the
/// best-effort fanout can pile up when a slow bus backend (e.g.
/// stalled NATS) lets publishes accumulate. Configurable via
/// [`A2aDispatcher::with_publish_concurrency`].
pub(crate) const DEFAULT_PUBLISH_PERMITS: usize = 64;

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

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

/// Ownership-key prefix for the A2A transport in the shared
/// `klieo-tenants` bucket. Same `a2a.{task_id}` shape as the
/// leader key so operators can correlate the two registries.
/// ADR-022.
const A2A_OWNERSHIP_KEY_PREFIX: &str = "a2a.";

fn ownership_key(task_id: &str) -> String {
    format!("{A2A_OWNERSHIP_KEY_PREFIX}{task_id}")
}

/// Pinned, boxed stream of [`TaskEvent`]s. Returned by
/// [`A2aDispatcher::dispatch_streaming`] and consumed by the HTTP
/// transport's SSE response builder.
pub type TaskEventStream = Pin<Box<dyn futures::Stream<Item = TaskEvent> + Send>>;

/// Lifecycle event emitted by [`crate::task_store::A2aTaskStore`] on each
/// mutating write. Serialised as the SSE `data:` payload for
/// `SubscribeToTask` / `SendStreamingMessage` responses in the HTTP
/// transport (the `http` cargo feature, not yet wired).
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
#[non_exhaustive]
pub struct TaskEvent {
    /// The task whose state transitioned.
    pub task_id: String,
    /// Current status post-transition.
    pub status: TaskStatus,
    /// Last message in the task history, if any.
    pub message: Option<Message>,
    /// `true` if `status` is terminal (Completed / Failed / Canceled /
    /// Rejected); stream consumers close after receiving it.
    pub final_event: bool,
    /// Monotonic per-task event id used by SSE `Last-Event-ID`
    /// resumption. Set by the HTTP transport when the event is yielded;
    /// defaults to 0 for events constructed outside the transport
    /// (e.g. synthetic initial replays from `task_store`).
    #[serde(default)]
    pub event_id: u64,
}

impl TaskEvent {
    /// Construct a `TaskEvent`. The only public constructor — required because
    /// `#[non_exhaustive]` prevents struct-literal construction outside this
    /// crate. `event_id` defaults to `0`; the HTTP transport overwrites it
    /// when the event is yielded into an SSE stream.
    pub fn new(
        task_id: impl Into<String>,
        status: TaskStatus,
        message: Option<Message>,
        final_event: bool,
    ) -> Self {
        Self {
            task_id: task_id.into(),
            status,
            message,
            final_event,
            event_id: 0,
        }
    }

    /// Override the monotonic event id assigned by the store or transport.
    ///
    /// Used by `A2aTaskStore::put` to stamp broadcast events, and by
    /// `dispatch_send_streaming` / `subscribe_snapshot_then_tail` to stamp
    /// synthetic initial events before they enter the SSE stream.
    #[must_use]
    pub fn with_event_id(mut self, event_id: u64) -> Self {
        self.event_id = event_id;
        self
    }

    /// Synthetic event emitted when a buffered event payload cannot
    /// be decoded back into a `TaskEvent`. Carries `event_id = 0` and
    /// only the task id; status is `Failed` and `final_event = true`
    /// so consumers don't hang.
    // Wired by the SSE resume branch (T8+); unused until then.
    #[allow(dead_code)]
    pub(crate) fn synthetic_corrupt(task_id: &str) -> Self {
        Self {
            task_id: task_id.into(),
            status: TaskStatus::Failed,
            message: None,
            final_event: true,
            event_id: 0,
        }
    }
}

/// Stream wrapper that owns a [`klieo_core::LeaderHandle`] and an
/// optional [`klieo_core::OwnershipHandle`] for the lifetime of an SSE
/// response body. Polling delegates to `inner`; `Drop` releases the
/// leader claim (heartbeat abort + best-effort KV delete inside
/// `LeaderHandle::drop`) and the ownership entry (best-effort
/// `kv.delete` spawned inside `OwnershipHandle::drop`) when the body
/// drops.
///
/// Both handle fields are `Option` so the same wrapper works when
/// leader election or tenant binding are disabled (`None` → behaves
/// as a transparent pass-through for that channel).
struct LeaderHoldStream<S> {
    inner: S,
    // Held for Drop side effect — never inspected.
    _leader: Option<klieo_core::LeaderHandle>,
    // Held for Drop side effect — never inspected. ADR-022.
    _ownership: Option<klieo_core::OwnershipHandle>,
}

impl<S: futures::Stream + Unpin> futures::Stream for LeaderHoldStream<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)
    }
}

/// Outcome of a leader-alive probe at `SubscribeToTask` 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,
}

/// Publish wrapper for cross-replica task event fanout. Replaces the
/// previous `tokio::sync::broadcast::Sender<TaskEvent>` returned by
/// [`A2aDispatcher::event_sink`].
///
/// The sink shares an [`Arc<Semaphore>`] with its parent
/// [`A2aDispatcher`] so [`Self::send`] can bound the number of
/// concurrent in-flight publishes. See [`Self::send`] for the
/// saturation-drop semantics.
#[derive(Clone)]
pub struct TaskEventSink {
    pubsub: Arc<dyn klieo_core::Pubsub>,
    permits: Arc<Semaphore>,
}

impl TaskEventSink {
    /// Wrap an `Arc<dyn Pubsub>` for publishing per-task events with a
    /// fresh, single-replica default semaphore
    /// (`DEFAULT_PUBLISH_PERMITS` permits). Multi-replica deployments
    /// or any wiring that wants to share the dispatcher's bound should
    /// go through [`A2aDispatcher::event_sink`] instead, which threads
    /// the dispatcher's semaphore through.
    pub fn new(pubsub: Arc<dyn klieo_core::Pubsub>) -> Self {
        Self::with_permits(pubsub, Arc::new(Semaphore::new(DEFAULT_PUBLISH_PERMITS)))
    }

    /// Wrap an `Arc<dyn Pubsub>` with an explicit publish-concurrency
    /// semaphore. Used by [`A2aDispatcher::event_sink`] to share the
    /// dispatcher's bound across every sink it hands out. The
    /// semaphore's `available_permits()` caps the number of concurrent
    /// publishes [`Self::send`] will execute; excess sends drop with
    /// `Ok(())` and a warn log.
    pub fn with_permits(pubsub: Arc<dyn klieo_core::Pubsub>, permits: Arc<Semaphore>) -> Self {
        Self { pubsub, permits }
    }

    /// Publish `event` on the per-task bus subject
    /// `klieo.a2a.task.{event.task_id}`.
    ///
    /// `task_id` is validated against
    /// [`klieo_core::validate_subject_token`] before subject
    /// construction; metacharacters (`.`, `*`, `>`, whitespace,
    /// non-ASCII) yield `A2aError::Bus(BusError::Invalid(_))`,
    /// preventing a caller-controlled task id from collapsing or
    /// wildcarding the subject namespace (CWE-74).
    ///
    /// # Backpressure
    ///
    /// Acquires a permit from the shared semaphore via
    /// [`Semaphore::try_acquire_owned`] BEFORE awaiting the publish.
    /// When the semaphore is saturated the event is dropped, a `warn`
    /// is logged at target `a2a.fanout` with
    /// `available_permits = 0`, and the method returns `Ok(())` — task
    /// event fanout is best-effort (matches the cluster-0.18 contract
    /// in [`crate::task_store::A2aTaskStore`] which already tolerates
    /// publish failures), and surfacing the saturation as an error
    /// would propagate into store writes that succeeded.
    ///
    /// Publish failures surface as `A2aError::Bus(_)`; server-side
    /// encode failures (effectively impossible for a server-built
    /// `TaskEvent`) surface as `A2aError::Internal { source }` so
    /// the wire envelope is `SERVER_ERROR` (`-32000`) rather than
    /// `PARSE_ERROR` (`-32700`). Both preserve `e.source()` for
    /// downstream tracing.
    #[instrument(
        skip_all,
        fields(
            messaging.system = "klieo-bus",
            messaging.destination = tracing::field::Empty,
            messaging.operation = "publish",
            messaging.message.payload_size_bytes = tracing::field::Empty,
            klieo.stream_id = %event.task_id,
        ),
        err,
    )]
    pub async fn send(&self, event: TaskEvent) -> Result<(), A2aError> {
        klieo_core::validate_subject_token(&event.task_id)?;
        let subject = format!("klieo.a2a.task.{}", event.task_id);
        tracing::Span::current().record("messaging.destination", subject.as_str());
        let bytes = serde_json::to_vec(&event).map_err(|e| A2aError::Internal {
            source: Box::new(e),
        })?;
        tracing::Span::current().record("messaging.message.payload_size_bytes", bytes.len());

        // Inject W3C tracecontext into bus headers so subscribers on
        // other replicas can stitch their decode span under the
        // publisher's span (cluster 0.23, ADR-023).
        let mut headers = klieo_core::Headers::default();
        let cx = tracing::Span::current().context();
        klieo_core::inject_traceparent(&mut headers, &cx);

        let permit = match self.permits.clone().try_acquire_owned() {
            Ok(p) => p,
            Err(_) => {
                tracing::warn!(
                    target: "a2a.fanout",
                    subject = %subject,
                    available_permits = self.permits.available_permits(),
                    "task event publish dropped: concurrency cap reached",
                );
                return Ok(()); // best-effort drop
            }
        };
        let result = self
            .pubsub
            .publish(&subject, bytes::Bytes::from(bytes), headers)
            .await;
        drop(permit);
        result?;
        Ok(())
    }
}

/// Builder for [`A2aDispatcher`] mirroring the shape of
/// `klieo_mcp_server::McpServerBuilder`:
/// accumulate `handler` / `authenticator` / `pubsub`, optionally
/// flag [`Self::with_cancel_subscription`], then finalise with
/// [`Self::build`] (unwrapped) or [`Self::build_arc`] (Arc, spawns
/// the wildcard cancel subscriber when the flag is set).
///
/// Symmetry with `McpServerBuilder` was carried as an architecture
/// gate medium for several rounds in cluster 0.18 — A2A previously
/// required a post-construction
/// [`A2aDispatcher::with_cancel_subscription`] call against an
/// already-wrapped [`Arc`], while MCP folded the same opt-in into
/// the builder. T5 closes the asymmetry; the pre-existing
/// [`A2aDispatcher::new`] shortcut and the post-construction method
/// remain for single-replica callers.
#[derive(Default)]
pub struct A2aDispatcherBuilder {
    handler: Option<Arc<dyn A2aHandler>>,
    authenticator: Option<Arc<dyn Authenticator>>,
    pubsub: Option<Arc<dyn klieo_core::Pubsub>>,
    publish_concurrency: Option<usize>,
    subscribe_cancels: bool,
    leader_kv: Option<Arc<dyn klieo_core::KvStore>>,
    tenant_kv: Option<Arc<dyn klieo_core::KvStore>>,
    tenant_strict: bool,
    leader_ttl: Option<std::time::Duration>,
    leader_heartbeat_interval: Option<std::time::Duration>,
    max_failover_attempts: Option<u32>,
    kv_reaper_interval: Option<std::time::Duration>,
    profile: klieo_core::DeploymentProfile,
}

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

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

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

    /// Wire an external [`klieo_core::Pubsub`] for cross-replica
    /// fanout (task events + cancel signals). Multi-replica
    /// deployments pass a shared NATS-backed pubsub.
    pub fn pubsub(mut self, pubsub: Arc<dyn klieo_core::Pubsub>) -> Self {
        self.pubsub = Some(pubsub);
        self
    }

    /// Convenience: wire a fresh in-process
    /// [`klieo_bus_memory::MemoryBus`] pubsub. Single-replica
    /// deployments only — multi-replica wiring MUST go through
    /// [`Self::pubsub`].
    #[must_use]
    pub fn with_in_process_pubsub(mut self) -> Self {
        self.pubsub = Some(klieo_bus_memory::MemoryBus::new().pubsub.clone());
        self
    }

    /// Override the per-dispatcher publish-concurrency cap (default
    /// `DEFAULT_PUBLISH_PERMITS`). Shared with every
    /// [`TaskEventSink`] handed out by [`A2aDispatcher::event_sink`]
    /// and with the drop-time cross-replica cancel publish.
    #[must_use]
    pub fn with_publish_concurrency(mut self, permits: usize) -> Self {
        self.publish_concurrency = Some(permits);
        self
    }

    /// Spawn the wildcard `klieo.a2a.cancel.>` background subscriber
    /// on [`Self::build_arc`]. Required for multi-replica deployments.
    /// The spawned task holds an [`Arc`] clone of the dispatcher, so
    /// this flag is only honoured by [`Self::build_arc`];
    /// [`Self::build`] panics if the flag is set.
    #[must_use]
    pub fn with_cancel_subscription(mut self) -> Self {
        self.subscribe_cancels = true;
        self
    }

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

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

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

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

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

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

    /// Opt in to the cluster-0.25 KV reaper. Spawns a background
    /// task scanning `klieo-leaders` (+ `klieo-tenants` when tenant
    /// binding is wired) every `interval`, evicting entries whose
    /// resume buffer is terminal. The scan task is spawned inside
    /// [`crate::http::A2aHttpServer::with_resume_buffer`] (which is
    /// where the resume buffer the reaper needs becomes available);
    /// calling this method on a dispatcher whose HTTP server never
    /// receives a non-noop resume buffer is a no-op.
    ///
    /// Recommended: 60s for production NATS deployments running
    /// `klieo-leaders` / `klieo-tenants` without bucket TTLs. Bucket
    /// TTLs are still the primary eviction mechanism — the reaper
    /// is a backstop for deployments where TTLs are not configured.
    #[must_use]
    pub fn with_kv_reaper(mut self, interval: std::time::Duration) -> Self {
        self.kv_reaper_interval = Some(interval);
        self
    }

    /// Finalise into a raw [`A2aDispatcher`].
    ///
    /// Returns [`A2aBuilderError::CancelRequiresArc`] when
    /// [`Self::with_cancel_subscription`] was set — the background
    /// subscriber requires an [`Arc`] clone; call [`Self::build_arc`] instead.
    /// Returns [`A2aBuilderError::MissingHandler`],
    /// [`A2aBuilderError::MissingAuthenticator`], or
    /// [`A2aBuilderError::MissingPubsub`] when the corresponding required
    /// field was not set on the builder.
    pub fn build(self) -> Result<A2aDispatcher, A2aBuilderError> {
        if self.subscribe_cancels {
            return Err(A2aBuilderError::CancelRequiresArc);
        }
        self.build_inner()
    }

    fn build_inner(self) -> Result<A2aDispatcher, A2aBuilderError> {
        let handler = self.handler.ok_or(A2aBuilderError::MissingHandler)?;
        let authenticator = self
            .authenticator
            .ok_or(A2aBuilderError::MissingAuthenticator)?;
        let pubsub = self.pubsub.ok_or(A2aBuilderError::MissingPubsub)?;
        let permits = self.publish_concurrency.unwrap_or(DEFAULT_PUBLISH_PERMITS);
        let leader_registry = self.leader_kv.map(|kv| {
            klieo_core::LeaderRegistry::new(
                kv,
                "klieo-leaders".into(),
                uuid::Uuid::new_v4().to_string(),
            )
        });
        let profile = self.profile;
        profile.validate(
            self.tenant_kv.is_some(),
            Some(authenticator.allows_anonymous()),
        )?;
        let tenant_strict = self.tenant_strict || profile.requires_strict_binding();
        let ownership_registry = self.tenant_kv.map(|kv| {
            let bucket = "klieo-tenants".into();
            if tenant_strict {
                klieo_core::OwnershipRegistry::new_strict(kv, bucket)
            } else {
                klieo_core::OwnershipRegistry::new(kv, bucket)
            }
        });
        if profile.requires_strict_binding() || profile.requires_named_principal() {
            tracing::warn!(
                target: "klieo.security",
                cwe = 639,
                "regulated multi-tenant profile active on this replica; \
                 cross-replica tenant isolation assumes ALL replicas run the \
                 same profile — a lenient peer reintroduces CWE-639. Fleet \
                 homogeneity is NOT verified by this replica."
            );
        }
        let leader_ttl = self.leader_ttl.unwrap_or(LEADER_TTL);
        let leader_heartbeat_interval = self.leader_heartbeat_interval.unwrap_or(leader_ttl / 2);
        let max_failover_attempts = self
            .max_failover_attempts
            .unwrap_or(klieo_core::FAILOVER_ATTEMPT_CAP);
        Ok(A2aDispatcher {
            handler,
            authenticator,
            pubsub,
            cancel_registry: klieo_core::CancelRegistry::new(),
            publish_permits: Arc::new(Semaphore::new(permits)),
            leader_registry,
            ownership_registry,
            leader_ttl,
            leader_heartbeat_interval,
            max_failover_attempts,
            kv_reaper_interval: self.kv_reaper_interval,
        })
    }

    /// Finalise into `Arc<A2aDispatcher>`. When
    /// [`Self::with_cancel_subscription`] was set, also spawns the
    /// wildcard cancel subscriber against the returned [`Arc`].
    pub fn build_arc(self) -> Result<Arc<A2aDispatcher>, A2aBuilderError> {
        let spawn_subscriber = self.subscribe_cancels;
        let dispatcher = Arc::new(self.build_inner()?);
        if spawn_subscriber {
            dispatcher.with_cancel_subscription();
        }
        Ok(dispatcher)
    }
}

/// Transport-agnostic dispatcher. Holds the handler and
/// authenticator. Both [`A2aServer`] (NATS) and the future
/// `A2aHttpServer` (HTTP, feature-gated) delegate to a shared
/// instance.
pub struct A2aDispatcher {
    handler: Arc<dyn A2aHandler>,
    authenticator: Arc<dyn Authenticator>,
    pubsub: Arc<dyn klieo_core::Pubsub>,
    cancel_registry: klieo_core::CancelRegistry<String>,
    publish_permits: Arc<Semaphore>,
    leader_registry: Option<klieo_core::LeaderRegistry>,
    ownership_registry: Option<klieo_core::OwnershipRegistry>,
    leader_ttl: std::time::Duration,
    leader_heartbeat_interval: std::time::Duration,
    max_failover_attempts: u32,
    kv_reaper_interval: Option<std::time::Duration>,
}

impl A2aDispatcher {
    /// Open an [`A2aDispatcherBuilder`] for accumulating handler /
    /// authenticator / pubsub and opting into cross-replica cancel
    /// subscription before finalising via
    /// [`A2aDispatcherBuilder::build`] or
    /// [`A2aDispatcherBuilder::build_arc`].
    ///
    /// Mirrors `klieo_mcp_server::McpServerBuilder`
    /// — see [`A2aDispatcherBuilder`] for the closed asymmetry.
    /// [`Self::new`] remains as a single-replica shortcut.
    pub fn builder() -> A2aDispatcherBuilder {
        A2aDispatcherBuilder::default()
    }

    /// Build a dispatcher with the given pubsub for cross-replica fanout.
    ///
    /// The dispatcher constructs an [`Arc<Semaphore>`] with
    /// `DEFAULT_PUBLISH_PERMITS` (64) permits to bound concurrent
    /// in-flight [`TaskEventSink::send`] publishes. Tune via
    /// [`Self::with_publish_concurrency`].
    pub fn new(
        handler: Arc<dyn A2aHandler>,
        authenticator: Arc<dyn Authenticator>,
        pubsub: Arc<dyn klieo_core::Pubsub>,
    ) -> Self {
        Self {
            handler,
            authenticator,
            pubsub,
            cancel_registry: klieo_core::CancelRegistry::new(),
            publish_permits: Arc::new(Semaphore::new(DEFAULT_PUBLISH_PERMITS)),
            leader_registry: None,
            ownership_registry: None,
            leader_ttl: LEADER_TTL,
            leader_heartbeat_interval: LEADER_TTL / 2,
            max_failover_attempts: klieo_core::FAILOVER_ATTEMPT_CAP,
            kv_reaper_interval: None,
        }
    }

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

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

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

    /// Configured KV-reaper interval, if any. `None` when
    /// [`A2aDispatcherBuilder::with_kv_reaper`] was not called.
    /// [`crate::http::A2aHttpServer::with_resume_buffer`] spawns
    /// the reaper task when this is `Some`.
    pub fn kv_reaper_interval(&self) -> Option<std::time::Duration> {
        self.kv_reaper_interval
    }

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

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

    /// Replace the dispatcher's publish-concurrency cap. Every
    /// [`TaskEventSink`] returned by [`Self::event_sink`] after this
    /// call shares the new semaphore; previously-issued sinks keep
    /// their old reference. Pass `0` to drop all fanout publishes
    /// (useful in tests that want to assert the saturation branch).
    /// Pass a larger value when the bus backend has demonstrated
    /// headroom and 64 in-flight publishes are a bottleneck.
    #[must_use]
    pub fn with_publish_concurrency(mut self, permits: usize) -> Self {
        self.publish_permits = Arc::new(Semaphore::new(permits));
        self
    }

    /// Convenience: builds with a fresh in-process `MemoryBus` pubsub.
    /// Single-replica deployments only. Multi-replica deployments wire a
    /// shared `Arc<dyn Pubsub>` via [`A2aDispatcher::new`].
    pub fn with_in_process_pubsub(
        handler: Arc<dyn A2aHandler>,
        authenticator: Arc<dyn Authenticator>,
    ) -> Self {
        let bus = klieo_bus_memory::MemoryBus::new();
        Self::new(handler, authenticator, bus.pubsub.clone())
    }

    /// Capability-shaped default for laptop-dev — wires
    /// [`klieo_auth_common::AllowAnonymous`] and a fresh in-process
    /// [`klieo_bus_memory::MemoryBus`] pubsub around the supplied
    /// handler.
    ///
    /// **TEST FIXTURE / DEMO ONLY.** Gated behind the `test-fixtures`
    /// feature (CWE-1188) so production builds cannot reach the
    /// permissive authenticator. Use [`Self::new`] or [`Self::builder`]
    /// for production wiring with a real
    /// [`klieo_auth_common::Authenticator`].
    #[cfg(feature = "test-fixtures")]
    pub fn local(handler: Arc<dyn A2aHandler>) -> Self {
        let bus = klieo_bus_memory::MemoryBus::new();
        Self::new(
            handler,
            Arc::new(klieo_auth_common::AllowAnonymous),
            bus.pubsub.clone(),
        )
    }

    /// Borrow the dispatcher's authenticator. Used by
    /// [`crate::http::A2aHttpServer`] to inspect
    /// [`crate::auth::Authenticator::allows_anonymous`] before binding.
    pub fn authenticator(&self) -> &Arc<dyn Authenticator> {
        &self.authenticator
    }

    /// Get a publish wrapper for hand-emitting TaskEvents (e.g. from
    /// `A2aTaskStore::with_event_sink`). The returned sink shares the
    /// dispatcher's publish-concurrency semaphore, so every sink
    /// handed out by one dispatcher is bounded by the same cap (see
    /// [`Self::with_publish_concurrency`]).
    pub fn event_sink(&self) -> TaskEventSink {
        TaskEventSink::with_permits(self.pubsub.clone(), self.publish_permits.clone())
    }

    /// Borrow the configured pubsub.
    pub fn pubsub(&self) -> &Arc<dyn klieo_core::Pubsub> {
        &self.pubsub
    }

    /// Borrow the dispatcher's publish-concurrency semaphore. Shared
    /// with every [`TaskEventSink`] this dispatcher hands out and
    /// threaded into `CancelOnDrop` (private wrapper in `crate::http`) so SSE-stream
    /// drop-time cross-replica cancel publishes bound under the same
    /// cap as task-event fanout. See [`Self::with_publish_concurrency`]
    /// for the dial.
    pub fn publish_permits(&self) -> &Arc<Semaphore> {
        &self.publish_permits
    }

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

    /// Publish a cross-replica cancel signal for `task_id` on
    /// `klieo.a2a.cancel.{task_id}`. Thin wrapper around
    /// [`klieo_core::cancel::publish_cancel_signal`] — see that
    /// helper for the full subject-validation + publish contract.
    /// `BusError` surfaces as `A2aError::Bus(_)` via the
    /// `#[from]` impl.
    ///
    /// # Security
    /// `task_id` is validated against
    /// [`klieo_core::validate_subject_token`] before subject
    /// construction; metacharacters (`.`, `*`, `>`, whitespace,
    /// non-ASCII) yield `A2aError::Bus(BusError::Invalid(_))`,
    /// preventing a caller-controlled task id from collapsing or
    /// wildcarding the subject namespace (CWE-74). Same guard
    /// shape as [`TaskEventSink::send`] on the per-task event
    /// subject.
    ///
    /// Cancel signals share the progressToken-as-credential threat
    /// model documented for resume in ADR-018 / ADR-019: knowledge
    /// of the task id grants the ability to cancel the invocation.
    /// Operators MUST mint unguessable task ids (UUID v4 or
    /// stronger) and gate per-tenant authorisation BEFORE the
    /// request reaches this dispatcher; without that, cross-tenant
    /// cancel becomes possible (CWE-639 IDOR).
    pub async fn publish_cancel(&self, task_id: &str) -> Result<(), A2aError> {
        klieo_core::cancel::publish_cancel_signal(&self.pubsub, A2A_CANCEL_SUBJECT_PREFIX, task_id)
            .await?;
        Ok(())
    }

    /// Spawn the wildcard cancel-subject background subscriber.
    ///
    /// Thin wrapper around
    /// [`klieo_core::cancel::spawn_wildcard_cancel_subscriber`] —
    /// see that helper for the loop body, ack policy, and
    /// startup-failure semantics. Required for multi-replica
    /// deployments; without it, only same-replica drop-cancel
    /// works.
    ///
    /// Idempotent at the call-site level: calling twice spawns
    /// two subscribers (not recommended; both will ack
    /// independently and the bus will balance delivery).
    pub fn with_cancel_subscription(self: &Arc<Self>) {
        klieo_core::cancel::spawn_wildcard_cancel_subscriber(
            self.pubsub.clone(),
            A2A_CANCEL_SUBJECT_PATTERN.to_string(),
            A2A_CANCEL_SUBJECT_PREFIX.to_string(),
            self.cancel_registry.clone(),
            A2A_CANCEL_LOG_TARGET,
        );
    }

    /// Dispatch one non-streaming JSON-RPC request.
    ///
    /// This is a flat dispatch table mapping each [`A2aMethod`] variant to
    /// its handler. Length exceeds the 40-line guideline by design —
    /// splitting the match into helper functions fragments the routing logic
    /// and makes it harder to see all method mappings at a glance.
    #[allow(clippy::too_many_lines)]
    #[instrument(
        skip_all,
        fields(rpc.system = "klieo-a2a", rpc.method = tracing::field::Empty),
    )]
    pub async fn dispatch(&self, ctx: &RequestContext, payload: &[u8]) -> JsonRpcResponse {
        let req: JsonRpcRequest = match serde_json::from_slice(payload) {
            Ok(r) => r,
            Err(e) => return A2aError::Json(e).to_json_rpc_error(Value::Null),
        };
        tracing::Span::current().record("rpc.method", req.method.as_str());
        let method = match A2aMethod::from_str(&req.method) {
            Ok(m) => m,
            Err(e) => return e.to_json_rpc_error(req.id.clone()),
        };
        let id = req.id.clone();
        if let Some(identity) = ctx.caller.as_ref() {
            if let Err(e) = self
                .authenticator
                .authorize_method(identity, &req.method)
                .await
            {
                warn!(
                    target: "security",
                    error = %e,
                    method = %req.method,
                    "a2a authorize_method rejected",
                );
                return A2aError::Unauthorized(e.to_string()).to_json_rpc_error(id);
            }
        } else if !self.authenticator.allows_anonymous() {
            warn!(
                target: "security",
                method = %req.method,
                "a2a anonymous caller rejected (authenticator requires an identity)",
            );
            return A2aError::Unauthorized("anonymous caller not permitted".into())
                .to_json_rpc_error(id);
        }
        let params_raw = req.params;
        let result: Result<Value, A2aError> = match method {
            A2aMethod::SendMessage => {
                run_method::<SendMessageParams, _, _>(params_raw, |p| async move {
                    let result = self.handler.send_message(ctx, p).await?;
                    // A non-streaming SendMessage creates a task that outlives the
                    // request and is later reached via GetTask/CancelTask/push-config
                    // — all of which gate on `enforce_owner`. Without a persistent
                    // ownership record an unclaimed task is Allowed for every tenant
                    // (CWE-639 IDOR). Claim ownership before returning, mirroring the
                    // streaming path's `try_claim_ownership`.
                    if let SendMessageResult::Task(task) = &result {
                        self.claim_owner_persistent(ctx, &task.id).await?;
                    }
                    Ok(result)
                })
                .await
            }
            // Streaming methods are only valid via dispatch_streaming;
            // the non-streaming dispatcher returns MethodNotFound so
            // NATS callers get a clear -32601 rather than a silent type error.
            A2aMethod::SendStreamingMessage => {
                Err(A2aError::MethodNotFound("SendStreamingMessage".into()))
            }
            A2aMethod::GetTask => {
                run_method::<GetTaskParams, _, _>(params_raw, |p| async move {
                    self.enforce_owner(ctx, &p.id).await?;
                    self.handler.get_task(ctx, p).await
                })
                .await
            }
            A2aMethod::ListTasks => {
                run_method::<ListTasksParams, _, _>(params_raw, |p| async move {
                    let mut result = self.handler.list_tasks(ctx, p).await?;
                    self.retain_owned_tasks(ctx, &mut result.tasks).await?;
                    Ok(result)
                })
                .await
            }
            A2aMethod::CancelTask => {
                run_method::<CancelTaskParams, _, _>(params_raw, |p| async move {
                    // Gate BEFORE the handler so a foreign cancel never reaches
                    // the cross-replica cancel publish.
                    self.enforce_owner(ctx, &p.id).await?;
                    self.handler.cancel_task(ctx, p).await
                })
                .await
            }
            A2aMethod::SubscribeToTask => Err(A2aError::MethodNotFound("SubscribeToTask".into())),
            A2aMethod::CreateTaskPushNotificationConfig => {
                run_method::<PushNotificationConfigParams, _, _>(params_raw, |p| async move {
                    self.enforce_owner(ctx, &p.taskId).await?;
                    self.handler
                        .create_task_push_notification_config(ctx, p)
                        .await
                })
                .await
            }
            A2aMethod::GetTaskPushNotificationConfig => {
                run_method::<GetPushNotificationConfigParams, _, _>(params_raw, |p| async move {
                    self.enforce_owner(ctx, &p.taskId).await?;
                    self.handler.get_task_push_notification_config(ctx, p).await
                })
                .await
            }
            A2aMethod::ListTaskPushNotificationConfigs => {
                run_method::<ListPushNotificationConfigsParams, _, _>(params_raw, |p| async move {
                    self.enforce_owner(ctx, &p.taskId).await?;
                    self.handler
                        .list_task_push_notification_configs(ctx, p)
                        .await
                })
                .await
            }
            A2aMethod::DeleteTaskPushNotificationConfig => {
                run_method::<DeletePushNotificationConfigParams, _, _>(params_raw, |p| async move {
                    self.enforce_owner(ctx, &p.taskId).await?;
                    self.handler
                        .delete_task_push_notification_config(ctx, p)
                        .await
                })
                .await
            }
            A2aMethod::GetExtendedAgentCard => {
                run_method::<GetExtendedAgentCardParams, _, _>(params_raw, |p| async move {
                    self.handler.get_extended_agent_card(ctx, p).await
                })
                .await
            }
        };
        match result {
            Ok(value) => JsonRpcResponse {
                jsonrpc: "2.0".into(),
                id,
                result: Some(value),
                error: None,
            },
            Err(e) => {
                log_internal_before_wire_seam(&e, &req.method);
                e.to_json_rpc_error(id)
            }
        }
    }

    /// Handle one unverified A2A JSON-RPC request: authenticate
    /// the caller, then dispatch on success.
    ///
    /// Returns a ready [`JsonRpcResponse`] in both the success and
    /// authentication-failure cases so callers never need to branch.
    ///
    /// # Security
    /// Authenticates before deserialising or dispatching
    /// (CWE-306). On auth failure the payload never reaches the
    /// handler; the caller receives a JSON-RPC `-32001`
    /// envelope.
    pub async fn handle_request(&self, headers: A2aHeaders, payload: &[u8]) -> JsonRpcResponse {
        match self.authenticator.authenticate(&headers, payload).await {
            Ok(identity) => {
                let ctx = RequestContext::new(headers, Some(identity));
                self.dispatch(&ctx, payload).await
            }
            Err(e) => {
                warn!(target: "security", error = %e, "a2a auth rejected");
                let req_id = serde_json::from_slice::<JsonRpcRequest>(payload)
                    .map(|r| r.id)
                    .unwrap_or(Value::Null);
                JsonRpcResponse {
                    jsonrpc: "2.0".into(),
                    id: req_id,
                    result: None,
                    error: Some(JsonRpcError {
                        code: codes::UNAUTHENTICATED,
                        message: "Unauthenticated".into(),
                        data: None,
                    }),
                }
            }
        }
    }

    /// Streaming-method analog of [`Self::handle_request`].
    /// Authenticates first (CWE-306), constructs a [`RequestContext`],
    /// then delegates to [`Self::dispatch_streaming`]. Used by the HTTP
    /// transport for `SendStreamingMessage` and `SubscribeToTask`.
    ///
    /// # Security
    /// Authenticates before deserialising or dispatching (CWE-306). On
    /// auth failure returns [`A2aError::Unauthorized`]; the payload never
    /// reaches the handler.
    pub async fn handle_streaming(
        &self,
        headers: A2aHeaders,
        payload: &[u8],
        task_store: &crate::task_store::A2aTaskStore,
        cancel: tokio_util::sync::CancellationToken,
        last_event_id: Option<u64>,
        resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    ) -> Result<TaskEventStream, A2aError> {
        match self.authenticator.authenticate(&headers, payload).await {
            Ok(identity) => {
                let ctx = RequestContext::new(headers, Some(identity)).with_cancel(cancel);
                self.dispatch_streaming(&ctx, payload, task_store, last_event_id, resume_buffer)
                    .await
            }
            Err(e) => {
                warn!(target: "security", error = %e, "a2a auth rejected (streaming path)");
                Err(A2aError::Unauthorized(e.to_string()))
            }
        }
    }

    /// Dispatch a streaming method (`SendStreamingMessage` or
    /// `SubscribeToTask`). Returns a [`TaskEventStream`] the HTTP
    /// transport serialises as SSE frames.
    ///
    /// The handler is tried first. On `Ok(stream)` the dispatcher
    /// returns it unchanged (handler override). On
    /// `Err(MethodNotFound)` the dispatcher synthesises a
    /// broadcast-derived stream:
    /// - `SendStreamingMessage`: converts params, calls
    ///   `send_message`, then tails the broadcast filtered by the
    ///   returned task id.
    /// - `SubscribeToTask`: reads current state from `task_store`,
    ///   emits a synthetic replay event, then tails the broadcast
    ///   until the first `final_event = true`.
    ///
    /// # Security
    /// Requires a pre-authenticated [`RequestContext`]. Callers
    /// must authenticate before invoking this method (the NATS
    /// path's [`A2aDispatcher::handle_request`] does this; an
    /// HTTP caller must authenticate at the request boundary).
    /// Passing an unauthenticated context violates CWE-306.
    #[instrument(
        skip_all,
        fields(rpc.system = "klieo-a2a", rpc.method = tracing::field::Empty),
        err,
    )]
    pub async fn dispatch_streaming(
        &self,
        ctx: &RequestContext,
        payload: &[u8],
        task_store: &crate::task_store::A2aTaskStore,
        last_event_id: Option<u64>,
        resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    ) -> Result<TaskEventStream, A2aError> {
        let req: JsonRpcRequest =
            serde_json::from_slice(payload).map_err(|e| A2aError::InvalidParams(e.to_string()))?;
        tracing::Span::current().record("rpc.method", req.method.as_str());
        let method = A2aMethod::from_str(&req.method)
            .map_err(|_| A2aError::MethodNotFound(req.method.clone()))?;
        if let Some(identity) = ctx.caller.as_ref() {
            if let Err(e) = self
                .authenticator
                .authorize_method(identity, &req.method)
                .await
            {
                warn!(
                    target: "security",
                    error = %e,
                    method = %req.method,
                    "a2a authorize_method rejected (streaming path)",
                );
                return Err(A2aError::Unauthorized(e.to_string()));
            }
        } else if !self.authenticator.allows_anonymous() {
            warn!(
                target: "security",
                method = %req.method,
                "a2a anonymous caller rejected (streaming path; authenticator requires an identity)",
            );
            return Err(A2aError::Unauthorized(
                "anonymous caller not permitted".into(),
            ));
        }
        let params_raw = req.params;
        match method {
            A2aMethod::SendStreamingMessage => {
                let params: SendMessageParams = serde_json::from_value(params_raw)
                    .map_err(|e| A2aError::InvalidParams(e.to_string()))?;
                // Resumption is subscribe-only in v1 — SendStreamingMessage
                // ignores last_event_id and resume_buffer.
                self.dispatch_send_streaming(ctx, params, task_store).await
            }
            A2aMethod::SubscribeToTask => {
                let params: SubscribeToTaskParams = serde_json::from_value(params_raw)
                    .map_err(|e| A2aError::InvalidParams(e.to_string()))?;
                self.dispatch_subscribe_to_task(
                    ctx,
                    params,
                    task_store,
                    last_event_id,
                    resume_buffer,
                )
                .await
            }
            _ => Err(A2aError::MethodNotFound(req.method)),
        }
    }

    #[instrument(
        skip_all,
        fields(
            rpc.system = "klieo-a2a",
            rpc.method = "SendStreamingMessage",
            klieo.stream_id = tracing::field::Empty,
        ),
        err,
    )]
    async fn dispatch_send_streaming(
        &self,
        ctx: &RequestContext,
        params: SendMessageParams,
        task_store: &crate::task_store::A2aTaskStore,
    ) -> Result<TaskEventStream, A2aError> {
        self.run_streaming_invoke(ctx, params, task_store, None)
            .await
    }

    /// Common body for `dispatch_send_streaming` and the cluster-0.24
    /// follower re-invoke path. When `existing_leader` is `Some`, skips
    /// the on-invoke leader claim and adopts the caller-provided handle
    /// (which already encodes attempt+1 and the cached payload via the
    /// follower's CAS claim). When `None`, captures the original params
    /// + principal and claims a fresh leader entry.
    async fn run_streaming_invoke(
        &self,
        ctx: &RequestContext,
        params: SendMessageParams,
        task_store: &crate::task_store::A2aTaskStore,
        existing_leader: Option<klieo_core::LeaderHandle>,
    ) -> Result<TaskEventStream, A2aError> {
        match self
            .handler
            .send_streaming_message(ctx, params.clone())
            .await
        {
            Ok(stream) => return Ok(stream),
            Err(A2aError::MethodNotFound(_)) => {}
            Err(other) => return Err(other),
        }
        // Capture the original params for cluster-0.24 failover
        // re-invoke BEFORE consuming `params` in the synthetic
        // fallback below. 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 payload_bytes_for_failover = match existing_leader {
            Some(_) => None,
            None => match serde_json::to_vec(&params) {
                Ok(v) => Some(Bytes::from(v)),
                Err(e) => {
                    warn!(
                        target: "a2a.failover",
                        error = %e,
                        "send-streaming params serialise for failover failed; \
                         proceeding without cached payload",
                    );
                    None
                }
            },
        };
        let principal_for_failover = match existing_leader {
            Some(_) => None,
            None => ctx
                .caller
                .as_ref()
                .filter(|id| !id.is_anonymous())
                .map(|id| id.as_str().to_string()),
        };
        let result = self.handler.send_message(ctx, params).await?;
        let task = match result {
            SendMessageResult::Task(t) => t,
            SendMessageResult::Message(_) => {
                // Non-task result — no task id to tail; return an
                // immediately-closed stream.
                return Ok(Box::pin(futures::stream::empty()));
            }
        };
        let task_id = task.id.clone();
        tracing::Span::current().record("klieo.stream_id", task_id.as_str());
        let terminal = task.status.is_terminal();
        let event_id = task_store.next_event_id(&task_id);
        let initial = TaskEvent::new(
            task_id.clone(),
            task.status,
            task.history.last().cloned(),
            terminal,
        )
        .with_event_id(event_id);
        if terminal {
            // Task already in a terminal state — emit the synthetic event
            // and close. No pubsub subscription needed, and no risk of
            // hanging when the handler (e.g. EchoHandler) doesn't write
            // through to a store.
            return Ok(Box::pin(futures::stream::once(futures::future::ready(
                initial,
            ))));
        }
        let initial_stream = futures::stream::once(futures::future::ready(initial));
        let tail = self.task_event_stream(task_id.clone()).await?;
        let chained = futures::StreamExt::chain(initial_stream, tail);
        // 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.
        let leader_handle = match existing_leader {
            Some(handle) => Some(handle),
            None => {
                self.try_claim_leader(&task_id, payload_bytes_for_failover, principal_for_failover)
                    .await
            }
        };
        let ownership_handle = self.try_claim_ownership(ctx, &task_id).await?;
        Ok(Box::pin(LeaderHoldStream {
            inner: chained,
            _leader: leader_handle,
            _ownership: ownership_handle,
        }))
    }

    /// Tenant-ownership claim for a streaming invoke. Returns `Ok(None)` (no
    /// claim, proceed) when no registry is wired, or the caller is anonymous
    /// (no principal to bind). On KV `put` failure the lenient registry returns
    /// `Ok(None)` (fail-open per ADR-022; the later `SubscribeToTask` lookup sees
    /// no entry and proceeds unbound), while a **strict** registry returns
    /// `Err` so the invoke is denied rather than started unprotected.
    ///
    /// The returned `OwnershipHandle` MUST be kept alive for the lifetime of the
    /// SSE response; the `LeaderHoldStream` wrapper carries it to drop-time so
    /// the `kv.delete` fires on stream end.
    async fn try_claim_ownership(
        &self,
        ctx: &RequestContext,
        task_id: &str,
    ) -> Result<Option<klieo_core::OwnershipHandle>, A2aError> {
        let Some(registry) = self.ownership_registry.as_ref() else {
            return Ok(None);
        };
        let Some(caller) = ctx.caller.as_ref() else {
            return Ok(None);
        };
        if caller.is_anonymous() {
            return Ok(None);
        }
        let key = ownership_key(task_id);
        let principal = caller.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(A2aError::Server(
                "ownership store unavailable; invoke denied (strict tenant binding)".into(),
            )),
            // `OwnershipClaim` is non_exhaustive: an unrecognised future outcome
            // is treated as deny — the fail-closed default for a security gate.
            _ => Err(A2aError::Server(
                "ownership claim inconclusive; invoke denied".into(),
            )),
        }
    }

    /// Record persistent tenant ownership for a task created by a non-streaming
    /// invoke. The streaming-path analogue [`Self::try_claim_ownership`] holds a
    /// drop-guard for the SSE lifetime; here the task outlives the request, so
    /// the entry must persist (reclaimed by the tenants-bucket reaper/TTL). No
    /// registry / no caller / anonymous caller → `Ok(())` (partial-deployment
    /// skip, ADR-022). On a strict-registry store error the invoke is denied so
    /// the task is never left readable by every tenant via `enforce_owner`.
    async fn claim_owner_persistent(
        &self,
        ctx: &RequestContext,
        task_id: &str,
    ) -> Result<(), A2aError> {
        let Some(registry) = self.ownership_registry.as_ref() else {
            return Ok(());
        };
        let Some(caller) = ctx.caller.as_ref() else {
            return Ok(());
        };
        if caller.is_anonymous() {
            return Ok(());
        }
        let key = ownership_key(task_id);
        match registry
            .record_owner(key, caller.as_str().to_string())
            .await
        {
            klieo_core::OwnershipRecord::Recorded | klieo_core::OwnershipRecord::Proceed => Ok(()),
            klieo_core::OwnershipRecord::Unavailable => Err(A2aError::Server(
                "ownership store unavailable; invoke denied (strict tenant binding)".into(),
            )),
            // `OwnershipRecord` is non_exhaustive: an unrecognised future outcome
            // is treated as deny — the fail-closed default for a security gate.
            _ => Err(A2aError::Server(
                "ownership record inconclusive; invoke denied".into(),
            )),
        }
    }

    /// SubscribeToTask tenant-binding gate. Looks up the owner in the
    /// `klieo-tenants` bucket and matches against the caller's
    /// principal.
    ///
    /// - Owner match → `Ok(())` (proceed).
    /// - Owner mismatch → `Err(A2aError::TaskNotFound)` — deny-as-
    ///   NotFound per OWASP IDOR best practice; same wire shape as a
    ///   nonexistent task 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(&self, ctx: &RequestContext, task_id: &str) -> Result<(), A2aError> {
        let Some(registry) = self.ownership_registry.as_ref() else {
            return Ok(());
        };
        let Some(caller) = ctx.caller.as_ref() else {
            return Ok(());
        };
        if caller.is_anonymous() {
            return Ok(());
        }
        let key = ownership_key(task_id);
        match registry.check_owner(&key, caller.as_str()).await {
            klieo_core::OwnershipCheck::Allowed => Ok(()),
            klieo_core::OwnershipCheck::Denied => {
                // Mismatch — deny-as-NotFound. No log of the claimed owner;
                // just the rejected principal so operators can spot IDOR
                // probes without leaking tenancy across the wire.
                warn!(
                    target: "a2a.tenants",
                    task_id = %task_id,
                    principal = %caller.as_str(),
                    "ownership mismatch on SubscribeToTask; denying as TaskNotFound",
                );
                Err(A2aError::TaskNotFound(task_id.to_string()))
            }
            klieo_core::OwnershipCheck::Unavailable => {
                warn!(
                    target: "a2a.tenants",
                    task_id = %task_id,
                    principal = %caller.as_str(),
                    "ownership store unavailable; request denied (strict tenant binding)",
                );
                Err(A2aError::OwnershipUnavailable)
            }
            other => {
                warn!(
                    target: "a2a.tenants",
                    task_id = %task_id,
                    principal = %caller.as_str(),
                    verdict = ?other,
                    "inconclusive ownership verdict; request denied",
                );
                Err(A2aError::OwnershipUnavailable)
            }
        }
    }

    /// Filter a `ListTasks` page down to the tasks the caller owns.
    ///
    /// Per-task analogue of [`Self::enforce_owner`]; same partial-deployment
    /// skip (no registry / no caller / anonymous → no filtering, ADR-022).
    /// A `Denied` task is dropped silently so the page cannot become a
    /// cross-tenant enumeration oracle (OWASP IDOR); an `Unavailable` verdict
    /// fails the whole list closed.
    ///
    /// [`klieo_core::KvStore`] has no multi-get, so the per-task checks are
    /// issued concurrently (one fan-out, ~one round-trip of latency) rather
    /// than serially — a page of N tasks must not stall on N sequential
    /// round-trips.
    async fn retain_owned_tasks(
        &self,
        ctx: &RequestContext,
        tasks: &mut Vec<Task>,
    ) -> Result<(), A2aError> {
        let Some(registry) = self.ownership_registry.as_ref() else {
            return Ok(());
        };
        let Some(caller) = ctx.caller.as_ref() else {
            return Ok(());
        };
        if caller.is_anonymous() {
            return Ok(());
        }
        let verdicts = futures::future::join_all(tasks.iter().map(|task| {
            let key = ownership_key(&task.id);
            async move { registry.check_owner(&key, caller.as_str()).await }
        }))
        .await;

        let mut owned = Vec::with_capacity(tasks.len());
        for (task, verdict) in std::mem::take(tasks).into_iter().zip(verdicts) {
            match verdict {
                klieo_core::OwnershipCheck::Allowed => owned.push(task),
                klieo_core::OwnershipCheck::Denied => {}
                klieo_core::OwnershipCheck::Unavailable => {
                    warn!(
                        target: "a2a.tenants",
                        principal = %caller.as_str(),
                        "ownership store unavailable; ListTasks denied (strict tenant binding)",
                    );
                    return Err(A2aError::OwnershipUnavailable);
                }
                other => {
                    warn!(
                        target: "a2a.tenants",
                        principal = %caller.as_str(),
                        verdict = ?other,
                        "inconclusive ownership verdict; ListTasks denied",
                    );
                    return Err(A2aError::OwnershipUnavailable);
                }
            }
        }
        *tasks = owned;
        Ok(())
    }

    /// Best-effort leader claim for a streaming invoke. Returns
    /// `None` when no registry is wired or when the KV `put` fails
    /// (fail-open per ADR-020). The returned `LeaderHandle` MUST be
    /// kept alive for the lifetime of the SSE response; the
    /// `LeaderHoldStream` wrapper carries it to drop-time.
    ///
    /// `payload` carries the original invoke params serialised to
    /// JSON bytes so a follower that detects an orphan can re-invoke
    /// an idempotent handler from the cached payload (cluster 0.24).
    /// `principal` records the authenticated subject (cluster 0.21)
    /// so the re-invoke runs under the same tenant binding as the
    /// original.
    async fn try_claim_leader(
        &self,
        task_id: &str,
        payload: Option<Bytes>,
        principal: Option<String>,
    ) -> Option<klieo_core::LeaderHandle> {
        let registry = self.leader_registry.as_ref()?;
        let key = format!("{A2A_LEADER_KEY_PREFIX}{task_id}");
        match registry
            .claim_with_heartbeat(
                key,
                self.leader_ttl,
                self.leader_heartbeat_interval,
                payload,
                principal,
            )
            .await
        {
            Ok(handle) => Some(handle),
            Err(e) => {
                warn!(
                    target: "a2a.leader",
                    task_id = %task_id,
                    error = %e,
                    "leader claim failed; degrading to no-claim (orphan detection \
                     disabled for this stream)",
                );
                None
            }
        }
    }

    /// 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 `Some(())` when the orphan write landed (caller raises
    /// [`A2aError::LeaderDied`] to terminate the current request);
    /// `None` 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 snapshot+tail 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(
        &self,
        task_id: &str,
        resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    ) -> Option<()> {
        let max_id = max_event_id(resume_buffer, task_id).await?;
        let next_id = max_id + 1;
        let frame = leader_died_sse_frame_bytes(task_id, next_id);
        if let Err(e) = resume_buffer.record(task_id, next_id, frame).await {
            warn!(
                target: "a2a.leader",
                task_id = %task_id,
                next_id,
                error = %e,
                "orphan terminal record failed; skipping orphan write",
            );
            return None;
        }
        if let Err(e) = resume_buffer.close(task_id).await {
            warn!(
                target: "a2a.leader",
                task_id = %task_id,
                error = %e,
                "orphan terminal close failed; resume buffer may retain stale stream",
            );
        }
        tracing::error!(
            target: "a2a.leader",
            task_id = %task_id,
            next_id,
            "stream leader died; emitted LEADER_DIED terminal frame at max+1",
        );
        Some(())
    }

    /// Cluster-0.24 follower-side response to a dead-leader probe.
    /// Returns:
    /// - `Ok(Some(stream))` when the follower re-invoked an idempotent
    ///   handler from the cached payload — the returned stream owns
    ///   the new leader claim via `LeaderHoldStream`.
    /// - `Ok(None)` when the follower wrote the terminal "leader died"
    ///   frame OR no orphan was detected (no resume-buffer entries) —
    ///   caller falls through to the regular snapshot+tail path on
    ///   `None`-from-no-orphan and the `LeaderDied` error path is
    ///   embedded as an `Err` returned from this method on terminate.
    /// - `Err(A2aError::LeaderDied { .. })` when the terminate-only
    ///   path fired (cluster-0.20 fallback).
    ///
    /// Public under `test-fixtures` ONLY so cluster-0.24 integration
    /// tests can drive the gate directly. Production callers reach
    /// this method via `dispatch_subscribe_to_task` after a
    /// `LeaderProbe::Dead` probe; tests bypass the probe because the
    /// probe and `lookup_entry_with_revision` both go through the
    /// same `kv.get` (an entry either exists or doesn't — there is
    /// no production window in which one probe sees None and the
    /// other sees Some, so an end-to-end HTTP-driven test that mocks
    /// "dead leader" via `kv.delete` collapses the entry visibility
    /// for the lookup too).
    #[cfg(not(feature = "test-fixtures"))]
    async fn handle_dead_leader_orphan(
        &self,
        ctx: &RequestContext,
        task_id: &str,
        task_store: &crate::task_store::A2aTaskStore,
        resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    ) -> Result<Option<TaskEventStream>, A2aError> {
        self.handle_dead_leader_orphan_impl(ctx, task_id, task_store, resume_buffer)
            .await
    }

    /// Test-only public alias for the cluster-0.24 orphan re-invoke
    /// gate. See the doc above for why integration tests reach the
    /// gate directly instead of driving through `is_alive`.
    #[cfg(feature = "test-fixtures")]
    pub async fn handle_dead_leader_orphan(
        &self,
        ctx: &RequestContext,
        task_id: &str,
        task_store: &crate::task_store::A2aTaskStore,
        resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    ) -> Result<Option<TaskEventStream>, A2aError> {
        self.handle_dead_leader_orphan_impl(ctx, task_id, task_store, resume_buffer)
            .await
    }

    async fn handle_dead_leader_orphan_impl(
        &self,
        ctx: &RequestContext,
        task_id: &str,
        task_store: &crate::task_store::A2aTaskStore,
        resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    ) -> Result<Option<TaskEventStream>, A2aError> {
        let registry = match self.leader_registry.as_ref() {
            Some(reg) => reg,
            None => return Ok(None),
        };
        let key = format!("{A2A_LEADER_KEY_PREFIX}{task_id}");
        let lookup = registry.lookup_entry_with_revision(&key).await;
        let Some((entry, prior_rev)) = lookup_ok_or_log(&key, lookup) else {
            // No readable entry (Ok(None) or KV error) — fall back to
            // terminate-only so the cluster-0.20 contract still holds.
            return self.terminate_orphan_a2a(task_id, resume_buffer).await;
        };
        if !self.handler.is_idempotent() {
            tracing::debug!(
                target: "a2a.failover",
                task_id,
                "handler not idempotent; emitting terminate frame",
            );
            return self.terminate_orphan_a2a(task_id, resume_buffer).await;
        }
        if entry.attempt >= self.max_failover_attempts {
            tracing::warn!(
                target: "a2a.failover",
                task_id,
                attempt = entry.attempt,
                cap = self.max_failover_attempts,
                "failover attempt cap reached; emitting terminate frame",
            );
            return self.terminate_orphan_a2a(task_id, resume_buffer).await;
        }
        let Some(payload_bytes) = entry.payload.clone() else {
            tracing::debug!(
                target: "a2a.failover",
                task_id,
                "no cached payload on leader entry; emitting terminate frame",
            );
            return self.terminate_orphan_a2a(task_id, resume_buffer).await;
        };
        let new_handle = match registry
            .claim_with_attempt_cas_and_heartbeat(
                key.clone(),
                self.leader_ttl,
                self.leader_heartbeat_interval,
                prior_rev,
                &entry,
            )
            .await
        {
            Ok(h) => h,
            Err(klieo_core::BusError::CasConflict { .. }) => {
                tracing::info!(
                    target: "a2a.failover",
                    task_id,
                    "another follower won the CAS race; emitting terminate frame",
                );
                return self.terminate_orphan_a2a(task_id, resume_buffer).await;
            }
            Err(e) => {
                tracing::warn!(
                    target: "a2a.failover",
                    task_id,
                    error = %e,
                    "CAS claim failed; emitting terminate frame",
                );
                return self.terminate_orphan_a2a(task_id, resume_buffer).await;
            }
        };
        self.record_failover_marker(task_id, &entry, &new_handle, resume_buffer)
            .await;
        let parsed_params: SendMessageParams = match serde_json::from_slice(&payload_bytes) {
            Ok(p) => p,
            Err(e) => {
                tracing::error!(
                    target: "a2a.failover",
                    task_id,
                    error = %e,
                    "cached payload parse failed; emitting terminate frame",
                );
                return self.terminate_orphan_a2a(task_id, resume_buffer).await;
            }
        };
        let reinvoke_ctx = self.fabricate_reinvoke_ctx(ctx, entry.principal.as_deref());
        let stream = self
            .run_streaming_invoke(&reinvoke_ctx, parsed_params, task_store, Some(new_handle))
            .await?;
        Ok(Some(stream))
    }

    /// Write the cluster-0.24 `failover-reinvoke` marker frame to the
    /// resume buffer at `max+1`. Best-effort: a record failure logs a
    /// warn + continues so the re-invoke still drives the production
    /// stream.
    async fn record_failover_marker(
        &self,
        task_id: &str,
        prior: &klieo_core::LeaderEntry,
        new_handle: &klieo_core::LeaderHandle,
        resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    ) {
        let Some(max) = max_event_id(resume_buffer, task_id).await else {
            // No retained events — re-invoke writes at event_id 1 onwards
            // via the regular task_store path; no marker frame needed.
            return;
        };
        let marker_id = max + 1;
        let frame = failover_reinvoke_sse_frame_bytes(
            task_id,
            marker_id,
            prior.attempt + 1,
            new_handle.replica_id(),
        );
        if let Err(e) = resume_buffer.record(task_id, marker_id, frame).await {
            warn!(
                target: "a2a.failover",
                task_id,
                marker_id,
                error = %e,
                "failover-reinvoke marker record failed; continuing without marker",
            );
        }
    }

    /// Build a fresh [`RequestContext`] for the failover re-invoke.
    /// The original `ctx.headers` are reused so handlers that read
    /// transport metadata see the resume request's framing; the
    /// caller identity is replaced with the cached principal so the
    /// re-invoke runs under the original tenant binding rather than
    /// the resume client's identity (cluster 0.22).
    fn fabricate_reinvoke_ctx(
        &self,
        ctx: &RequestContext,
        principal: Option<&str>,
    ) -> RequestContext {
        let caller = principal
            .map(|p| klieo_auth_common::Identity::new(p.to_string()))
            .or_else(|| Some(klieo_auth_common::Identity::anonymous()));
        RequestContext::new(ctx.headers.clone(), caller).with_cancel(ctx.cancel.clone())
    }

    /// Terminate-only orphan response (cluster 0.20 path). Writes the
    /// "leader died" frame to the resume buffer + raises
    /// [`A2aError::LeaderDied`]. Returns `Ok(None)` (caller falls
    /// through to the regular subscribe path) when no buffer entries
    /// exist — same shape as the inline pre-0.24 logic.
    async fn terminate_orphan_a2a(
        &self,
        task_id: &str,
        resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    ) -> Result<Option<TaskEventStream>, A2aError> {
        if self
            .write_orphan_terminal_frame(task_id, resume_buffer)
            .await
            .is_some()
        {
            return Err(A2aError::LeaderDied {
                stream_id: format!("{A2A_LEADER_KEY_PREFIX}{task_id}"),
            });
        }
        Ok(None)
    }

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

    /// # Authorization
    ///
    /// The synthetic fallback path (when the handler returns
    /// `MethodNotFound`) reads the task from [`crate::task_store::A2aTaskStore`]
    /// keyed by `params.id` without a per-task ACL check. Task IDs are
    /// UUID v4 (unguessable) by convention, but adopters that need
    /// per-task ownership enforcement should override
    /// [`A2aHandler::subscribe_to_task`] and perform the check before
    /// returning the stream.
    ///
    /// When `last_event_id` is set and the `resume_buffer` has retained
    /// events for the task, events since that id are replayed first, then
    /// the live broadcast tail is appended (with dedup by `event_id` to
    /// handle the subscribe-before-replay race).  Falls back to
    /// `subscribe_snapshot_then_tail` when the buffer returns `NotFound`
    /// or a backend error.  Returns [`A2aError::ResumeBufferExpired`] when
    /// the cursor predates the oldest retained event.
    #[instrument(
        skip_all,
        fields(
            rpc.system = "klieo-a2a",
            rpc.method = "SubscribeToTask",
            klieo.stream_id = %params.id,
        ),
        err,
    )]
    async fn dispatch_subscribe_to_task(
        &self,
        ctx: &RequestContext,
        params: SubscribeToTaskParams,
        task_store: &crate::task_store::A2aTaskStore,
        last_event_id: Option<u64>,
        resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
    ) -> Result<TaskEventStream, A2aError> {
        // Tenant binding gate (ADR-022). Runs BEFORE the leader probe
        // + task_store lookup so a cross-tenant probe cannot
        // distinguish "exists but owned by someone else" from
        // "doesn't exist" via timing or downstream errors. On
        // mismatch returns TaskNotFound — sanitised on the HTTP wire
        // by `dispatch_streaming` into the same `-32000 internal
        // server error` envelope as the nonexistent-task path so the
        // response leaks no existence info (OWASP IDOR).
        self.enforce_owner(ctx, &params.id).await?;
        // Orphan detection: leader probe at the SubscribeToTask 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 handler + attempt cap + payload available)
        // OR writes the terminal "leader died" SSE frame and raises
        // [`A2aError::LeaderDied`] (cluster 0.20 fallback).
        if let LeaderProbe::Dead = self.probe_leader(&params.id).await {
            if let Some(stream) = self
                .handle_dead_leader_orphan(ctx, &params.id, task_store, resume_buffer.as_ref())
                .await?
            {
                return Ok(stream);
            }
        }
        if let Some(since) = last_event_id {
            // Gate the resume branch through the handler's ACL: adopters that
            // override A2aHandler::subscribe_to_task to enforce per-task
            // ownership MUST run before any buffered events leave the server.
            //
            // - Ok(stream): handler owns resumption (and any ACL it enforced
            //   passed); return its stream verbatim, matching the
            //   `subscribe_snapshot_then_tail` handler-override semantic.
            // - Err(MethodNotFound): no override, fall through to our buffer
            //   replay (default impl returns this).
            // - Any other Err: hard reject (auth failure, invalid params, etc.).
            match self.handler.subscribe_to_task(ctx, params.clone()).await {
                Ok(stream) => return Ok(stream),
                Err(A2aError::MethodNotFound(_)) => {}
                Err(other) => return Err(other),
            }

            // Subscribe to the pubsub FIRST so no events emitted during
            // replay are lost between replay-end and tail-subscribe.
            let task_id_for_tail = params.id.clone();
            let tail_stream = self.task_event_stream(task_id_for_tail).await?;

            match resume_buffer.replay(&params.id, since).await {
                Ok(replay_stream) => {
                    use std::sync::atomic::{AtomicU64, Ordering};
                    use std::sync::Arc as StdArc;

                    let max_replayed = StdArc::new(AtomicU64::new(since));
                    let max_for_replay = max_replayed.clone();
                    let task_id_for_replay = params.id.clone();
                    let replay = replay_stream.map(move |(id, bytes)| {
                        max_for_replay.fetch_max(id, Ordering::SeqCst);
                        serde_json::from_slice::<TaskEvent>(&bytes)
                            .unwrap_or_else(|_| TaskEvent::synthetic_corrupt(&task_id_for_replay))
                    });
                    let max_for_tail = max_replayed.clone();
                    let tail = tail_stream.filter(move |ev: &TaskEvent| {
                        ev.event_id > max_for_tail.load(Ordering::SeqCst)
                    });
                    return Ok(Box::pin(futures::StreamExt::chain(replay, tail)));
                }
                Err(klieo_core::resume::ResumeError::Expired { since_id }) => {
                    return Err(A2aError::ResumeBufferExpired { since_id });
                }
                Err(klieo_core::resume::ResumeError::NotFound(_)) => { /* fall through */ }
                Err(klieo_core::resume::ResumeError::Backend(e)) => {
                    tracing::warn!(
                        target: "a2a.resume",
                        error = %e,
                        "resume buffer backend error; falling back to snapshot+tail"
                    );
                }
                // ResumeError is non_exhaustive; future variants fall through
                // to snapshot+tail like NotFound.
                Err(_) => {}
            }
        }
        self.subscribe_snapshot_then_tail(ctx, params, task_store)
            .await
    }

    /// Snapshot-then-tail fallback for `SubscribeToTask`. Reads current
    /// task state from the store, emits a synthetic replay event, then
    /// tails the broadcast until the first `final_event = true`.
    ///
    /// This is the pre-resumption path extracted verbatim from the
    /// original `dispatch_subscribe_to_task` body.
    #[instrument(
        skip_all,
        fields(klieo.stream_id = %params.id),
        level = "debug",
        err,
    )]
    async fn subscribe_snapshot_then_tail(
        &self,
        ctx: &RequestContext,
        params: SubscribeToTaskParams,
        task_store: &crate::task_store::A2aTaskStore,
    ) -> Result<TaskEventStream, A2aError> {
        match self.handler.subscribe_to_task(ctx, params.clone()).await {
            Ok(stream) => return Ok(stream),
            Err(A2aError::MethodNotFound(_)) => {}
            Err(other) => return Err(other),
        }
        let task = task_store
            .get(&params.id)
            .await?
            .ok_or_else(|| A2aError::TaskNotFound(params.id.clone()))?;
        let terminal = task.status.is_terminal();
        let event_id = task_store.next_event_id(&task.id);
        let initial = TaskEvent::new(
            task.id.clone(),
            task.status,
            task.history.last().cloned(),
            terminal,
        )
        .with_event_id(event_id);
        self.task_event_stream_with_initial(task.id, initial, terminal)
            .await
    }

    /// Subscribe to the per-task pubsub subject and return a stream of
    /// `TaskEvent`s, closing after the first event with `final_event = true`.
    ///
    /// Each received [`klieo_core::Msg`] is ack-ed immediately on receipt;
    /// ephemeral SSE consumers prefer duplicate delivery over redelivery
    /// storms on disconnect. Bus errors are logged and skipped.
    #[instrument(
        skip_all,
        fields(klieo.stream_id = %task_id),
        level = "debug",
        err,
    )]
    async fn task_event_stream(&self, task_id: String) -> Result<TaskEventStream, A2aError> {
        klieo_core::validate_subject_token(&task_id)?;
        let subject = format!("klieo.a2a.task.{task_id}");
        let durable = DurableName::new(format!("klieo-eph-{}", uuid::Uuid::new_v4()));
        let msg_stream = self.pubsub.subscribe(&subject, durable).await?;
        let stream = futures::stream::unfold((msg_stream, false), move |(mut ms, done)| {
            let task_id = task_id.clone();
            async move {
                if done {
                    return None;
                }
                loop {
                    match ms.next().await {
                        None => return None,
                        Some(Err(e)) => {
                            warn!(
                                target: "a2a",
                                task_id = %task_id,
                                error = %e,
                                "task event stream bus error; skipping",
                            );
                            continue;
                        }
                        Some(Ok(msg)) => {
                            // Ack immediately — ephemeral SSE consumer.
                            if let Err(ack_err) = msg.ack.ack().await {
                                tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
                            }
                            // 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!(
                                "task_event_decode",
                                messaging.system = "klieo-bus",
                                messaging.destination = %format!("klieo.a2a.task.{task_id}"),
                                messaging.operation = "receive",
                                klieo.stream_id = %task_id,
                            );
                            decode_span.set_parent(parent_cx);
                            let _enter = decode_span.enter();
                            match serde_json::from_slice::<TaskEvent>(&msg.payload) {
                                Err(e) => {
                                    warn!(
                                        target: "a2a",
                                        task_id = %task_id,
                                        error = %e,
                                        "task event decode error; skipping",
                                    );
                                    continue;
                                }
                                Ok(ev) => {
                                    let terminal = ev.final_event;
                                    return Some((ev, (ms, terminal)));
                                }
                            }
                        }
                    }
                }
            }
        });
        Ok(Box::pin(stream))
    }

    async fn task_event_stream_with_initial(
        &self,
        task_id: String,
        initial: TaskEvent,
        initial_terminal: bool,
    ) -> Result<TaskEventStream, A2aError> {
        use futures::StreamExt as FutExt;

        let initial_stream = futures::stream::once(futures::future::ready(initial));
        if initial_terminal {
            // Task already in a terminal state — emit the replay
            // event and close; no pubsub subscription needed.
            return Ok(Box::pin(initial_stream));
        }
        let tail = self.task_event_stream(task_id).await?;
        Ok(Box::pin(FutExt::chain(initial_stream, tail)))
    }
}

/// JSON-RPC dispatcher that listens on a NATS subject.
pub struct A2aServer {
    app_prefix: String,
    agent_id: String,
    dispatcher: Arc<A2aDispatcher>,
    pubsub: Arc<dyn Pubsub>,
}

impl A2aServer {
    /// Build a new server.
    ///
    /// **Authentication is mandatory** since 0.2.0. Wire a real
    /// [`Authenticator`] (e.g.
    /// [`BearerTokenAuthenticator`](crate::auth::BearerTokenAuthenticator))
    /// for production deployments. For tests and demos that intentionally
    /// run unauthenticated, pass
    /// [`AllowAnonymous`](crate::auth::AllowAnonymous) explicitly — that
    /// makes the choice visible at the wiring site.
    pub fn new(
        app_prefix: String,
        agent_id: String,
        handler: Arc<dyn A2aHandler>,
        authenticator: Arc<dyn Authenticator>,
        pubsub: Arc<dyn Pubsub>,
    ) -> Self {
        let dispatcher = Arc::new(A2aDispatcher::new(handler, authenticator, pubsub.clone()));
        Self {
            app_prefix,
            agent_id,
            dispatcher,
            pubsub,
        }
    }

    /// Borrow the inner dispatcher (e.g. to share with an HTTP bridge
    /// mounted in the same process).
    pub fn dispatcher(&self) -> &Arc<A2aDispatcher> {
        &self.dispatcher
    }

    /// Subject the server listens on.
    pub fn subject(&self) -> String {
        format!("{}.a2a.{}.rpc", self.app_prefix, self.agent_id)
    }

    /// Run the dispatch loop until the underlying subscription terminates.
    pub async fn run(self) -> Result<(), A2aError> {
        let subject = self.subject();
        let durable = DurableName::new(format!("a2a-server-{}", self.agent_id));
        let mut stream = self.pubsub.subscribe(&subject, durable).await?;
        while let Some(item) = stream.next().await {
            match item {
                Ok(msg) => {
                    self.handle_one(msg).await;
                }
                Err(e) => {
                    error!(error = %e, "a2a subscribe stream error; exiting loop");
                    return Err(A2aError::Bus(e));
                }
            }
        }
        Ok(())
    }

    async fn handle_one(&self, msg: Msg) {
        let reply_to = match msg.headers.get("Reply-To") {
            Some(s) => s.clone(),
            None => {
                warn!("a2a request missing Reply-To header; dropping");
                if let Err(ack_err) = msg.ack.ack().await {
                    tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
                }
                return;
            }
        };
        let headers = A2aHeaders::decode_from(&msg.headers);
        let response = self.dispatcher.handle_request(headers, &msg.payload).await;
        let bytes = match serde_json::to_vec(&response) {
            Ok(b) => Bytes::from(b),
            Err(e) => {
                error!(error = %e, "encode response");
                if let Err(ack_err) = msg.ack.ack().await {
                    tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
                }
                return;
            }
        };
        if let Err(e) = self.pubsub.publish(&reply_to, bytes, Headers::new()).await {
            error!(error = %e, "publish reply");
        }
        if let Err(ack_err) = msg.ack.ack().await {
            tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
        }
    }
}

/// 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(
    resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
    stream_id: &str,
) -> Option<u64> {
    let mut replay = match resume_buffer.replay(stream_id, 0).await {
        Ok(stream) => stream,
        Err(klieo_core::resume::ResumeError::NotFound(_)) => return None,
        Err(e) => {
            warn!(
                target: "a2a.leader",
                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
}

/// 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(_))`.
fn lookup_ok_or_log(
    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) => {
            warn!(
                target: "a2a.failover",
                key,
                error = %e,
                "leader entry lookup_with_revision failed; falling back to terminate",
            );
            None
        }
    }
}

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

/// Build the terminal SSE-frame payload bytes that get recorded in
/// the resume buffer at `max_id + 1` when an orphan is detected.
/// The frame is the raw JSON-RPC error envelope that resume clients
/// see verbatim when replaying past the orphan boundary.
fn leader_died_sse_frame_bytes(task_id: &str, event_id: u64) -> bytes::Bytes {
    let payload = serde_json::json!({
        "jsonrpc": "2.0",
        "id": serde_json::Value::Null,
        "error": {
            "code": codes::LEADER_DIED,
            "message": "stream leader died",
            "data": { "stream_id": format!("{A2A_LEADER_KEY_PREFIX}{task_id}") }
        },
        "event_id": event_id,
    });
    bytes::Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
}

/// Log an `A2aError` at the server-side seam BEFORE its wire envelope
/// is sent. The wire seam sanitises Display for internal classes
/// (`Bus`, `Server`, `Misconfigured`, `Internal`) — without this log
/// the source chain would never reach traces.
///
/// Logs at `warn` for client-class variants (`MethodNotFound`,
/// `InvalidParams`, `Unauthorized`, `TaskNotFound`,
/// `ResumeBufferExpired`, `Json`) and at `error` for internal-class
/// variants so operators can filter on the elevated level.
pub(crate) fn log_internal_before_wire_seam(err: &A2aError, method: &str) {
    let internal = matches!(
        err,
        A2aError::Bus(_)
            | A2aError::Server(_)
            | A2aError::Misconfigured(_)
            | A2aError::Internal { .. }
    );
    if internal {
        error!(
            target: "a2a.wire",
            method = %method,
            error = %err,
            source = ?std::error::Error::source(err),
            "internal error mapped to JSON-RPC wire envelope",
        );
    } else {
        warn!(
            target: "a2a.wire",
            method = %method,
            error = %err,
            "client-class error mapped to JSON-RPC wire envelope",
        );
    }
}

/// Helper: deserialise params from raw JSON, run the handler closure,
/// then serialise the typed result back to a JSON [`Value`].
async fn run_method<P, T, Fut>(raw: Value, run: impl FnOnce(P) -> Fut) -> Result<Value, A2aError>
where
    P: DeserializeOwned,
    T: Serialize,
    Fut: std::future::Future<Output = Result<T, A2aError>>,
{
    let params: P =
        serde_json::from_value(raw).map_err(|e| A2aError::InvalidParams(e.to_string()))?;
    let typed = run(params).await?;
    serde_json::to_value(typed).map_err(A2aError::from)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::envelope::A2aHeaders;
    use crate::handler::EchoHandler;
    use crate::task_store::A2aTaskStore;
    use crate::types::{Task, TaskStatus};
    use klieo_auth_common::{AllowAnonymous, Identity};
    use klieo_bus_memory::MemoryBus;

    fn anon_ctx() -> RequestContext {
        RequestContext::new(
            A2aHeaders::decode_from(&klieo_core::Headers::new()),
            Some(Identity::anonymous()),
        )
    }

    fn make_store_with_dispatcher(dispatcher: &A2aDispatcher) -> A2aTaskStore {
        A2aTaskStore::new(
            Arc::new(MemoryBus::new()).kv.clone(),
            crate::task_store::DEFAULT_BUCKET.into(),
        )
        .with_event_sink(dispatcher.event_sink())
    }

    fn make_task(id: &str, status: TaskStatus) -> Task {
        Task {
            id: id.into(),
            contextId: "ctx-1".into(),
            status,
            artifacts: vec![],
            history: vec![],
            metadata: None,
        }
    }

    fn noop_resume_buffer() -> Arc<dyn klieo_core::resume::ResumeBuffer> {
        Arc::new(klieo_core::resume::NoopResumeBuffer)
    }

    #[tokio::test]
    async fn local_wires_anonymous_auth_and_in_process_pubsub_and_dispatches() {
        let dispatcher = A2aDispatcher::local(Arc::new(EchoHandler::default()));
        let ctx = anon_ctx();
        let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
        let resp = dispatcher.dispatch(&ctx, payload).await;
        assert_eq!(resp.id, serde_json::json!(1));
        assert!(resp.result.is_some());
        assert!(resp.error.is_none());
    }

    #[tokio::test]
    async fn dispatcher_routes_send_message_to_handler() {
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            Arc::new(EchoHandler::default()),
            Arc::new(AllowAnonymous),
        );
        let ctx = anon_ctx();
        let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
        let resp = dispatcher.dispatch(&ctx, payload).await;
        assert_eq!(resp.id, serde_json::json!(1));
        assert!(resp.result.is_some());
        assert!(resp.error.is_none());
    }

    fn named_ctx(principal: &str) -> RequestContext {
        RequestContext::new(
            A2aHeaders::decode_from(&klieo_core::Headers::new()),
            Some(Identity::new(principal)),
        )
    }

    const SEND_HI: &[u8] = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;

    #[tokio::test]
    async fn non_streaming_send_message_claims_ownership_so_foreign_get_task_is_denied() {
        let kv = Arc::new(MemoryBus::new()).kv.clone();
        let dispatcher = A2aDispatcher::builder()
            .handler(Arc::new(EchoHandler::default()))
            .authenticator(Arc::new(AllowAnonymous))
            .with_in_process_pubsub()
            .with_tenant_binding(kv)
            .build()
            .unwrap();

        // alice creates a task via the non-streaming SendMessage path.
        let created = dispatcher.dispatch(&named_ctx("alice"), SEND_HI).await;
        let task_id = created.result.expect("SendMessage returns a task")["id"]
            .as_str()
            .expect("task id is a string")
            .to_string();
        let get = format!(
            r#"{{"jsonrpc":"2.0","id":2,"method":"GetTask","params":{{"id":"{task_id}"}}}}"#
        );

        // A foreign tenant must NOT be able to read alice's task (CWE-639 IDOR):
        // deny-as-TaskNotFound, surfaced as a JSON-RPC error.
        let foreign = dispatcher.dispatch(&named_ctx("bob"), get.as_bytes()).await;
        assert!(
            foreign.error.is_some(),
            "foreign GetTask must be denied, got result {:?}",
            foreign.result
        );

        // The owner can still read her own task.
        let owner = dispatcher
            .dispatch(&named_ctx("alice"), get.as_bytes())
            .await;
        assert!(
            owner.error.is_none(),
            "owner GetTask must succeed, got error {:?}",
            owner.error
        );
        assert!(owner.result.is_some());
    }

    #[tokio::test]
    async fn dispatcher_streaming_subscribe_to_task_replays_current_state() {
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            Arc::new(EchoHandler::default()),
            Arc::new(AllowAnonymous),
        );
        let store = make_store_with_dispatcher(&dispatcher);

        let task = make_task("t-1", TaskStatus::Working);
        store.put(&task).await.unwrap();

        let payload =
            br#"{"jsonrpc":"2.0","id":1,"method":"SubscribeToTask","params":{"id":"t-1"}}"#;
        let ctx = anon_ctx();
        let mut stream = dispatcher
            .dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
            .await
            .unwrap();

        let first = stream.next().await.expect("stream produced no first event");
        assert_eq!(first.task_id, "t-1");
        assert!(matches!(first.status, TaskStatus::Working));
        assert!(!first.final_event, "Working is not terminal");
    }

    #[tokio::test]
    async fn dispatcher_streaming_rejects_unknown_task() {
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            Arc::new(EchoHandler::default()),
            Arc::new(AllowAnonymous),
        );
        let store = make_store_with_dispatcher(&dispatcher);

        let payload =
            br#"{"jsonrpc":"2.0","id":1,"method":"SubscribeToTask","params":{"id":"nope"}}"#;
        let ctx = anon_ctx();
        let result = dispatcher
            .dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
            .await;
        assert!(matches!(result, Err(A2aError::TaskNotFound(_))));
    }

    struct NoAnonAuthn;

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

    #[tokio::test]
    async fn dispatcher_streaming_rejects_anonymous_under_named_authenticator() {
        // Mirror of the non-streaming anon-deny: a None caller on the streaming
        // path must be rejected when the authenticator forbids anonymous.
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            Arc::new(EchoHandler::default()),
            Arc::new(NoAnonAuthn),
        );
        let store = make_store_with_dispatcher(&dispatcher);
        let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendStreamingMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
        let ctx = RequestContext::new(A2aHeaders::decode_from(&klieo_core::Headers::new()), None);
        let result = dispatcher
            .dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
            .await;
        assert!(matches!(result, Err(A2aError::Unauthorized(_))));
    }

    #[tokio::test]
    async fn dispatcher_streaming_subscribe_terminal_task_closes_stream() {
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            Arc::new(EchoHandler::default()),
            Arc::new(AllowAnonymous),
        );
        let store = make_store_with_dispatcher(&dispatcher);

        let task = make_task("t-done", TaskStatus::Completed);
        store.put(&task).await.unwrap();

        let payload =
            br#"{"jsonrpc":"2.0","id":1,"method":"SubscribeToTask","params":{"id":"t-done"}}"#;
        let ctx = anon_ctx();
        let mut stream = dispatcher
            .dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
            .await
            .unwrap();

        let first = stream.next().await.expect("expected replay event");
        assert_eq!(first.task_id, "t-done");
        assert!(first.final_event, "Completed must set final_event=true");
        // Stream must close after the terminal replay event.
        assert!(stream.next().await.is_none(), "stream must be exhausted");
    }

    #[tokio::test]
    async fn dispatcher_streaming_send_streaming_message_returns_stream() {
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            Arc::new(EchoHandler::default()),
            Arc::new(AllowAnonymous),
        );
        let store = make_store_with_dispatcher(&dispatcher);

        let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendStreamingMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
        let ctx = anon_ctx();
        let result = dispatcher
            .dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
            .await;
        assert!(
            result.is_ok(),
            "dispatch_streaming should return Ok for SendStreamingMessage"
        );
    }

    #[tokio::test]
    async fn handle_streaming_threads_cancel_token_into_request_context() {
        use crate::auth::RequestContext;
        use crate::error::A2aError;
        use crate::handler::A2aHandler;
        use crate::server::TaskEventStream;
        use crate::types::SendMessageParams;
        use klieo_auth_common::{AllowAnonymous, Authenticator};
        use std::sync::{Arc, Mutex};
        use tokio_util::sync::CancellationToken;

        struct CancelSpy {
            observed: Mutex<Option<CancellationToken>>,
        }
        #[async_trait::async_trait]
        impl A2aHandler for CancelSpy {
            async fn send_streaming_message(
                &self,
                ctx: &RequestContext,
                _: SendMessageParams,
            ) -> Result<TaskEventStream, A2aError> {
                *self.observed.lock().unwrap() = Some(ctx.cancel.clone());
                Ok(Box::pin(futures::stream::empty()))
            }
        }

        let spy = Arc::new(CancelSpy {
            observed: Mutex::new(None),
        });
        let dispatcher = A2aDispatcher::with_in_process_pubsub(
            spy.clone() as Arc<dyn A2aHandler>,
            Arc::new(AllowAnonymous) as Arc<dyn Authenticator>,
        );
        let store = A2aTaskStore::new(
            Arc::new(klieo_bus_memory::MemoryBus::new()).kv.clone(),
            crate::task_store::DEFAULT_BUCKET.into(),
        )
        .with_event_sink(dispatcher.event_sink());

        let token = CancellationToken::new();
        let body = serde_json::to_vec(&serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "SendStreamingMessage",
            "params": {
                "message": {
                    "messageId": "m-spy",
                    "role": "user",
                    "parts": [],
                    "extensions": [],
                    "referenceTaskIds": []
                }
            }
        }))
        .unwrap();

        let _ = dispatcher
            .handle_streaming(
                A2aHeaders::decode_from(&klieo_core::Headers::new()),
                &body,
                &store,
                token.clone(),
                None,
                noop_resume_buffer(),
            )
            .await
            .unwrap();

        token.cancel();
        let observed = spy.observed.lock().unwrap().clone().unwrap();
        assert!(
            observed.is_cancelled(),
            "cancel must propagate from handle_streaming arg"
        );
    }

    #[tokio::test]
    async fn task_event_sink_publishes_to_per_task_subject() {
        let bus = klieo_bus_memory::MemoryBus::new();
        let pubsub: std::sync::Arc<dyn klieo_core::Pubsub> = bus.pubsub.clone();
        let sink = TaskEventSink::new(pubsub.clone());

        let durable = klieo_core::DurableName::new("test-eph");
        let mut stream = pubsub
            .subscribe("klieo.a2a.task.t-1", durable)
            .await
            .unwrap();

        let event = TaskEvent::new(
            "t-1".to_string(),
            crate::types::TaskStatus::Submitted,
            None,
            false,
        )
        .with_event_id(1);

        sink.send(event.clone()).await.unwrap();

        use tokio_stream::StreamExt as _;
        let msg = tokio::time::timeout(std::time::Duration::from_millis(500), stream.next())
            .await
            .expect("timeout")
            .expect("stream ended")
            .expect("subscribe err");
        let decoded: TaskEvent = serde_json::from_slice(&msg.payload).unwrap();
        assert_eq!(decoded.task_id, "t-1");
        assert_eq!(decoded.event_id, 1);
        msg.ack.ack().await.unwrap();
    }
}

#[cfg(test)]
mod profile_tests {
    use super::*;
    use crate::handler::EchoHandler;
    use klieo_auth_common::{AllowAnonymous, AuthError, Headers, Identity};
    use klieo_core::DeploymentProfile;

    struct NamedAuthn;

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

    fn builder_with(auth: Arc<dyn Authenticator>) -> A2aDispatcherBuilder {
        let bus = klieo_bus_memory::MemoryBus::new();
        A2aDispatcher::builder()
            .handler(Arc::new(EchoHandler::default()))
            .authenticator(auth)
            .pubsub(bus.pubsub.clone())
    }

    #[test]
    fn regulated_without_tenant_kv_fails_closed() {
        let err = builder_with(Arc::new(NamedAuthn))
            .profile(DeploymentProfile::RegulatedMultiTenant)
            .build()
            .err()
            .expect("expected RegulatedProfile error");
        assert!(matches!(
            err,
            A2aBuilderError::RegulatedProfile(klieo_core::ProfileViolation::MissingTenantKv)
        ));
    }

    #[test]
    fn regulated_with_anonymous_auth_fails_closed() {
        let bus = klieo_bus_memory::MemoryBus::new();
        let err = builder_with(Arc::new(AllowAnonymous))
            .with_tenant_binding(bus.kv.clone())
            .profile(DeploymentProfile::RegulatedMultiTenant)
            .build()
            .err()
            .expect("expected RegulatedProfile error");
        assert!(matches!(
            err,
            A2aBuilderError::RegulatedProfile(klieo_core::ProfileViolation::AnonymousAuth)
        ));
    }

    #[test]
    fn regulated_forces_strict_over_lenient_binding() {
        let bus = klieo_bus_memory::MemoryBus::new();
        let dispatcher = builder_with(Arc::new(NamedAuthn))
            .with_tenant_binding(bus.kv.clone())
            .profile(DeploymentProfile::RegulatedMultiTenant)
            .build()
            .expect("regulated build with named auth + kv must succeed");
        assert_eq!(
            dispatcher
                .ownership_registry
                .as_ref()
                .map(|r| r.is_strict()),
            Some(true),
            "regulated profile must force a strict registry even over lenient binding"
        );
    }

    #[test]
    fn unprofiled_keeps_lenient_binding() {
        let bus = klieo_bus_memory::MemoryBus::new();
        let dispatcher = builder_with(Arc::new(NamedAuthn))
            .with_tenant_binding(bus.kv.clone())
            .build()
            .expect("unprofiled build ok");
        assert_eq!(
            dispatcher
                .ownership_registry
                .as_ref()
                .map(|r| r.is_strict()),
            Some(false)
        );
    }

    /// Dispatcher with lenient tenant binding wired. `AllowAnonymous`
    /// authorizes every method, so the only gate exercised is the
    /// per-resource ownership check keyed off the caller in the context.
    fn tenant_bound_dispatcher() -> A2aDispatcher {
        let bus = klieo_bus_memory::MemoryBus::new();
        A2aDispatcher::builder()
            .handler(Arc::new(EchoHandler::default()))
            .authenticator(Arc::new(AllowAnonymous))
            .pubsub(bus.pubsub.clone())
            .with_tenant_binding(bus.kv.clone())
            .build()
            .expect("tenant-bound dispatcher builds")
    }

    fn named_ctx(principal: &str) -> RequestContext {
        RequestContext::new(
            A2aHeaders::decode_from(&klieo_core::Headers::new()),
            Some(Identity::new(principal)),
        )
    }

    const SEND_MESSAGE_PAYLOAD: &[u8] = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;

    #[tokio::test]
    async fn anonymous_caller_rejected_when_authenticator_requires_identity() {
        let dispatcher = builder_with(Arc::new(NamedAuthn)).build().expect("build");
        let ctx = RequestContext::new(A2aHeaders::decode_from(&klieo_core::Headers::new()), None);
        let resp = dispatcher.dispatch(&ctx, SEND_MESSAGE_PAYLOAD).await;
        let err = resp.error.expect("anonymous caller must be rejected");
        assert_eq!(err.code, codes::UNAUTHENTICATED);
        assert!(
            resp.result.is_none(),
            "rejected request must carry no result"
        );
    }

    /// Typed `SendMessageParams` for seeding a task on a handler directly,
    /// bypassing the (tenant-gated) dispatch path.
    fn hi_params() -> crate::types::SendMessageParams {
        serde_json::from_value(serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "user",
                "parts": [{"type": "text", "content": "hi"}],
                "extensions": [],
                "referenceTaskIds": []
            }
        }))
        .expect("valid SendMessageParams")
    }

    async fn create_task(dispatcher: &A2aDispatcher, ctx: &RequestContext) -> String {
        let resp = dispatcher.dispatch(ctx, SEND_MESSAGE_PAYLOAD).await;
        resp.result
            .expect("SendMessage result")
            .get("id")
            .and_then(|v| v.as_str())
            .expect("created task carries an id")
            .to_string()
    }

    #[tokio::test]
    async fn anonymous_send_message_records_no_ownership() {
        // The non-streaming claim skips anonymous callers (no principal to
        // bind, ADR-022). An anonymous SendMessage must leave no ownership
        // entry — not write one under an "anonymous" sentinel.
        let dispatcher = tenant_bound_dispatcher();
        let anon = RequestContext::new(
            A2aHeaders::decode_from(&klieo_core::Headers::new()),
            Some(Identity::anonymous()),
        );
        let created = dispatcher.dispatch(&anon, SEND_MESSAGE_PAYLOAD).await;
        let task_id = created.result.expect("SendMessage returns a task")["id"]
            .as_str()
            .expect("task id is a string")
            .to_string();
        let owner = dispatcher
            .ownership_registry()
            .expect("registry wired")
            .lookup(&format!("{A2A_OWNERSHIP_KEY_PREFIX}{task_id}"))
            .await
            .expect("lookup ok");
        assert!(
            owner.is_none(),
            "anonymous SendMessage must not write an ownership entry, got {owner:?}"
        );
    }

    /// Claim ownership and RETURN the handle — the caller MUST keep it
    /// alive, since dropping it deletes the KV entry and reopens the gate.
    #[must_use]
    async fn claim_for(
        dispatcher: &A2aDispatcher,
        task_id: &str,
        principal: &str,
    ) -> klieo_core::OwnershipHandle {
        dispatcher
            .ownership_registry()
            .expect("ownership registry wired")
            .claim(
                format!("{A2A_OWNERSHIP_KEY_PREFIX}{task_id}"),
                principal.into(),
            )
            .await
            .expect("claim ownership")
    }

    fn by_id_request(method: &str, task_id: &str) -> Vec<u8> {
        serde_json::to_vec(&serde_json::json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": method,
            "params": { "id": task_id },
        }))
        .expect("encode by-id request")
    }

    #[tokio::test]
    async fn get_task_owner_reads_but_foreign_principal_denied_as_not_found() {
        let dispatcher = tenant_bound_dispatcher();
        let alice = named_ctx("alice");
        let task_id = create_task(&dispatcher, &alice).await;
        let _ownership = claim_for(&dispatcher, &task_id, "alice").await;

        let request = by_id_request("GetTask", &task_id);

        let owner = dispatcher.dispatch(&alice, &request).await;
        assert!(owner.error.is_none(), "owner must read own task: {owner:?}");
        assert_eq!(
            owner
                .result
                .expect("owner result")
                .get("id")
                .and_then(|v| v.as_str()),
            Some(task_id.as_str()),
        );

        // Foreigner is denied — TaskNotFound (-32000), no body, never -32001.
        let bob = named_ctx("bob");
        let foreign = dispatcher.dispatch(&bob, &request).await;
        let err = foreign.error.expect("foreign principal must be denied");
        assert_eq!(err.code, codes::SERVER_ERROR);
        assert_ne!(
            err.code,
            codes::UNAUTHENTICATED,
            "deny must not surface -32001 (would leak existence info)",
        );
        assert!(
            err.message.contains("task not found"),
            "deny-as-not-found expected, got: {}",
            err.message,
        );
        assert!(foreign.result.is_none(), "must not leak the task body");
    }

    #[tokio::test]
    async fn cancel_task_foreign_principal_denied_before_mutation() {
        let dispatcher = tenant_bound_dispatcher();
        let alice = named_ctx("alice");
        let task_id = create_task(&dispatcher, &alice).await;
        let _ownership = claim_for(&dispatcher, &task_id, "alice").await;

        // Bob's cancel must be refused before the status flips.
        let bob = named_ctx("bob");
        let cancel = dispatcher
            .dispatch(&bob, &by_id_request("CancelTask", &task_id))
            .await;
        let err = cancel.error.expect("foreign cancel must be denied");
        assert_eq!(err.code, codes::SERVER_ERROR);
        assert!(err.message.contains("task not found"));

        // Alice still sees the task in its pre-cancel state — proving the
        // gate ran before the handler's mutation, not after.
        let owner_view = dispatcher
            .dispatch(&alice, &by_id_request("GetTask", &task_id))
            .await;
        let status = owner_view
            .result
            .expect("owner result")
            .get("status")
            .and_then(|v| v.as_str())
            .map(str::to_string);
        assert_ne!(
            status.as_deref(),
            Some("canceled"),
            "foreign cancel must not mutate the task",
        );
    }

    /// Every KV read fails, to drive the strict `Unavailable` verdict through
    /// `retain_owned_tasks`.
    struct UnavailableKv;

    #[async_trait::async_trait]
    impl klieo_core::KvStore for UnavailableKv {
        async fn get(
            &self,
            _: &str,
            _: &str,
        ) -> Result<Option<klieo_core::KvEntry>, klieo_core::BusError> {
            Err(klieo_core::BusError::Connection("kv down".into()))
        }
        async fn put(
            &self,
            _: &str,
            _: &str,
            _: Bytes,
        ) -> Result<klieo_core::Revision, klieo_core::BusError> {
            Err(klieo_core::BusError::Connection("kv down".into()))
        }
        async fn cas(
            &self,
            _: &str,
            _: &str,
            _: Bytes,
            _: Option<klieo_core::Revision>,
        ) -> Result<klieo_core::Revision, klieo_core::BusError> {
            Err(klieo_core::BusError::Connection("kv down".into()))
        }
        async fn delete(&self, _: &str, _: &str) -> Result<(), klieo_core::BusError> {
            Err(klieo_core::BusError::Connection("kv down".into()))
        }
        async fn lease(
            &self,
            _: &str,
            _: &str,
            _: std::time::Duration,
        ) -> Result<klieo_core::Lease, klieo_core::BusError> {
            Err(klieo_core::BusError::Connection("kv down".into()))
        }
        async fn keys(&self, _: &str) -> Result<Vec<String>, klieo_core::BusError> {
            Err(klieo_core::BusError::Connection("kv down".into()))
        }
    }

    #[tokio::test]
    async fn send_message_fails_closed_when_strict_ownership_store_unavailable() {
        // Under strict tenant binding a store-down on the ownership write must
        // deny the create rather than leave the task unprotected (an unrecorded
        // task is Allowed for every tenant by `enforce_owner`).
        let bus = klieo_bus_memory::MemoryBus::new();
        let dispatcher = A2aDispatcher::builder()
            .handler(Arc::new(EchoHandler::default()))
            .authenticator(Arc::new(AllowAnonymous))
            .pubsub(bus.pubsub.clone())
            .with_tenant_binding_strict(Arc::new(UnavailableKv))
            .build()
            .expect("strict tenant-bound dispatcher builds");

        let resp = dispatcher
            .dispatch(&named_ctx("alice"), SEND_MESSAGE_PAYLOAD)
            .await;
        let err = resp
            .error
            .expect("strict store-down must fail the create closed");
        assert_eq!(err.code, codes::SERVER_ERROR);
        assert!(resp.result.is_none(), "denied create must carry no task");
    }

    #[tokio::test]
    async fn list_tasks_fails_closed_when_ownership_store_unavailable() {
        // Seed a task directly on the handler: the strict store-down below would
        // also fail the SendMessage create closed, so the seed cannot go through
        // dispatch — this test isolates the ListTasks (`retain_owned_tasks`) gate.
        let handler = Arc::new(EchoHandler::default());
        handler
            .send_message(&named_ctx("alice"), hi_params())
            .await
            .expect("seed task on handler");
        let bus = klieo_bus_memory::MemoryBus::new();
        let dispatcher = A2aDispatcher::builder()
            .handler(handler)
            .authenticator(Arc::new(AllowAnonymous))
            .pubsub(bus.pubsub.clone())
            .with_tenant_binding_strict(Arc::new(UnavailableKv))
            .build()
            .expect("strict tenant-bound dispatcher builds");

        let alice = named_ctx("alice");
        let list_request = serde_json::to_vec(&serde_json::json!({
            "jsonrpc": "2.0", "id": 4, "method": "ListTasks", "params": {},
        }))
        .expect("encode ListTasks");
        let resp = dispatcher.dispatch(&alice, &list_request).await;
        let err = resp
            .error
            .expect("a strict store-down must fail the list closed");
        assert_eq!(err.code, codes::SERVER_ERROR);
        assert!(
            err.message.contains("unavailable") || err.message.contains("list denied"),
            "fail-closed message expected, got: {}",
            err.message,
        );
        assert!(resp.result.is_none(), "must not return a partial list");
    }

    #[tokio::test]
    async fn push_notification_config_arms_deny_foreign_principal() {
        let dispatcher = tenant_bound_dispatcher();
        let alice = named_ctx("alice");
        let task_id = create_task(&dispatcher, &alice).await;
        let _ownership = claim_for(&dispatcher, &task_id, "alice").await;
        let bob = named_ctx("bob");

        // Each push-config method is keyed by `taskId`; a foreign principal
        // must be gated before the handler (deny-as-NotFound), while the owner
        // passes the gate and reaches the handler default (MethodNotFound on
        // EchoHandler). The differing codes prove the gate fires for bob only.
        let requests = [
            serde_json::json!({"jsonrpc":"2.0","id":5,"method":"CreateTaskPushNotificationConfig","params":{"taskId":task_id,"url":"https://example.test/hook"}}),
            serde_json::json!({"jsonrpc":"2.0","id":6,"method":"GetTaskPushNotificationConfig","params":{"taskId":task_id,"id":"c1"}}),
            serde_json::json!({"jsonrpc":"2.0","id":7,"method":"ListTaskPushNotificationConfigs","params":{"taskId":task_id}}),
            serde_json::json!({"jsonrpc":"2.0","id":8,"method":"DeleteTaskPushNotificationConfig","params":{"taskId":task_id,"id":"c1"}}),
        ];

        for request in requests {
            let method = request["method"].as_str().expect("method").to_string();
            let body = serde_json::to_vec(&request).expect("encode push-config request");

            let foreign = dispatcher.dispatch(&bob, &body).await;
            let foreign_err = foreign
                .error
                .unwrap_or_else(|| panic!("{method}: foreign principal must be denied"));
            assert_eq!(
                foreign_err.code,
                codes::SERVER_ERROR,
                "{method}: foreign principal must be denied-as-not-found",
            );
            assert!(
                foreign_err.message.contains("task not found"),
                "{method}: expected deny-as-not-found, got {}",
                foreign_err.message,
            );
            assert!(foreign.result.is_none(), "{method}: must not leak a body");

            let owner = dispatcher.dispatch(&alice, &body).await;
            let owner_err = owner
                .error
                .unwrap_or_else(|| panic!("{method}: EchoHandler declines push-config"));
            assert_eq!(
                owner_err.code,
                codes::METHOD_NOT_FOUND,
                "{method}: owner must pass the gate and reach the handler default",
            );
        }
    }

    #[tokio::test]
    async fn list_tasks_drops_tasks_the_caller_does_not_own() {
        let dispatcher = tenant_bound_dispatcher();
        let alice = named_ctx("alice");
        let bob = named_ctx("bob");

        // Two tasks in one handler store, one owned by each principal.
        let alice_task = create_task(&dispatcher, &alice).await;
        let bob_task = create_task(&dispatcher, &bob).await;
        let _alice_owns = claim_for(&dispatcher, &alice_task, "alice").await;
        let _bob_owns = claim_for(&dispatcher, &bob_task, "bob").await;

        let list_request = serde_json::to_vec(&serde_json::json!({
            "jsonrpc": "2.0", "id": 3, "method": "ListTasks", "params": {},
        }))
        .expect("encode ListTasks");

        let resp = dispatcher.dispatch(&bob, &list_request).await;
        let ids: Vec<String> = resp
            .result
            .expect("list result")
            .get("tasks")
            .and_then(|v| v.as_array())
            .expect("tasks array")
            .iter()
            .filter_map(|t| t.get("id").and_then(|v| v.as_str()).map(str::to_string))
            .collect();

        assert!(ids.contains(&bob_task), "owner must see own task: {ids:?}");
        assert!(
            !ids.contains(&alice_task),
            "list must not leak a foreign-owned task: {ids:?}",
        );
    }
}