klieo-core 0.41.0

//! Per-replica registry of in-flight cancellation tokens.
//!
//! When klieo runs as multiple replicas behind a shared bus, a cancel
//! request published on a wildcard subject (such as
//! `klieo.a2a.cancel.{task_id}` or `klieo.mcp.cancel.{progressToken}`)
//! reaches every replica, but only one replica actually owns the
//! in-flight invocation for a given id. Each replica keeps a
//! [`CancelRegistry`] mapping ids it owns to their
//! [`CancellationToken`]; a background task subscribed to the wildcard
//! cancel subject calls [`CancelRegistry::cancel`] on every inbound
//! message. The replica that owns the id fires its token; replicas
//! that don't own it observe a no-op lookup miss.
//!
//! The expected lifecycle on the owning replica is:
//!
//! 1. Caller creates a fresh [`CancellationToken`] for the new
//!    invocation and calls [`CancelRegistry::register`] with the
//!    stream-specific id (A2A task id, MCP progress token, etc.).
//! 2. The invocation runs to completion or is aborted via the token.
//! 3. Caller calls [`CancelRegistry::deregister`] in a finally-style
//!    cleanup so the map does not leak entries past invocation end.

use std::collections::HashMap;
use std::hash::Hash;
use std::pin::Pin;
use std::sync::{Arc, RwLock};
use std::task::{Context, Poll};

use bytes::Bytes;
use futures_core::Stream;
use tokio_util::sync::CancellationToken;

use crate::bus::{validate_subject_token, Headers, Pubsub};
use crate::error::BusError;

/// Per-replica registry mapping in-flight stream ids to their
/// [`CancellationToken`]. Cloning the registry yields a handle that
/// shares the same underlying map, so the wildcard cancel-subject
/// subscription task and the invoke-handling task can both hold one.
///
/// The key type is generic so the same shape serves both A2A (task id
/// as `String`) and MCP (progress token as `String`) without coupling
/// `klieo-core` to either transport's id type.
#[derive(Clone, Debug, Default)]
pub struct CancelRegistry<K>
where
    K: Eq + Hash + Clone,
{
    inner: Arc<RwLock<HashMap<K, CancellationToken>>>,
}

impl<K> CancelRegistry<K>
where
    K: Eq + Hash + Clone + std::fmt::Debug,
{
    /// Create an empty registry. Equivalent to
    /// [`CancelRegistry::default`].
    pub fn new() -> Self {
        Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Record that this replica owns the in-flight invocation
    /// identified by `id`, with the supplied token. Call this at
    /// invoke start, before the invocation can plausibly receive a
    /// cross-replica cancel for that id.
    ///
    /// An existing entry for `id` is overwritten; the caller is
    /// expected to keep ids unique per invocation.
    #[tracing::instrument(skip_all, fields(klieo.cancel.key = ?id), level = "debug")]
    pub fn register(&self, id: K, token: CancellationToken) {
        self.inner
            .write()
            .expect("CancelRegistry lock poisoned")
            .insert(id, token);
    }

    /// Remove the entry for `id`. Call this when the invocation
    /// finishes (success or failure) so the map does not retain
    /// tokens past invocation end. A missing entry is silently
    /// ignored — deregister-after-cancel is a normal interleaving.
    #[tracing::instrument(skip_all, fields(klieo.cancel.key = ?id), level = "debug")]
    pub fn deregister(&self, id: &K) {
        self.inner
            .write()
            .expect("CancelRegistry lock poisoned")
            .remove(id);
    }

    /// Fire the token registered for `id` and return `true`. When no
    /// entry exists for `id` return `false` without panicking — this
    /// is the expected path on replicas that don't own the cancelled
    /// invocation, and on cancel messages that arrive after the
    /// invocation already finished.
    #[tracing::instrument(
        skip_all,
        fields(klieo.cancel.key = ?id, klieo.cancel.hit = tracing::field::Empty),
        level = "debug",
    )]
    pub fn cancel(&self, id: &K) -> bool {
        let guard = self.inner.read().expect("CancelRegistry lock poisoned");
        let hit = match guard.get(id) {
            Some(token) => {
                token.cancel();
                true
            }
            None => false,
        };
        tracing::Span::current().record("klieo.cancel.hit", hit);
        hit
    }
}

/// Publish a cross-replica cancel signal for `id` on the subject
/// `format!("{subject_prefix}{id}")`. The body is empty — subscribers
/// only need the subject to identify which in-flight invocation to
/// cancel via their local [`CancelRegistry`].
///
/// Callers pass the transport-specific subject prefix, e.g.
/// `"klieo.a2a.cancel."` or `"klieo.mcp.cancel."`. `id` is validated
/// against [`validate_subject_token`] before subject construction;
/// metacharacters (`.`, `*`, `>`, whitespace, non-ASCII) yield
/// `BusError::Invalid(_)`, preventing a caller-controlled id from
/// collapsing or wildcarding the subject namespace (CWE-74).
///
/// # Security
///
/// Cancel signals share the progressToken-as-credential threat model
/// documented for resume in ADR-018. The cancel subject embeds an
/// opaque caller-supplied id (`task_id` for A2A, `progressToken` for
/// MCP); any caller who knows the id can publish a cancel that the
/// owning replica will honour. Operators MUST mint unguessable ids
/// (UUID v4 or stronger) and gate per-tenant authorisation BEFORE
/// the request reaches the dispatcher — without that, knowledge of
/// the id grants the ability to cancel the invocation across tenant
/// boundaries (CWE-639 IDOR). See ADR-018 and ADR-019.
pub async fn publish_cancel_signal(
    pubsub: &Arc<dyn Pubsub>,
    subject_prefix: &str,
    id: &str,
) -> Result<(), BusError> {
    validate_subject_token(id)?;
    let subject = format!("{subject_prefix}{id}");
    pubsub
        .publish(&subject, Bytes::new(), Headers::default())
        .await
}

/// Spawn the wildcard cancel-subject background subscriber for one
/// transport (A2A or MCP).
///
/// Required for multi-replica deployments — without it, only same-
/// replica drop-cancel works. Subscribes to `subject_pattern` (e.g.
/// `"klieo.a2a.cancel.>"`) with an ephemeral durable name
/// `klieo-cancel-{uuid}`. On each delivered [`crate::Msg`]: parse
/// the per-invoke id by stripping `subject_prefix` from
/// `msg.subject`, ack the message, then look up the id in
/// `registry` and fire the matching token (or no-op on lookup miss
/// — the originating replica is elsewhere).
///
/// Subscribe failure at startup logs `error` at `log_target` and
/// the background task exits — replica falls back to single-replica
/// semantics. Per-message bus errors log `warn` and the loop
/// continues. Ack failures log `warn` (and the message will be
/// redelivered by the bus); the loop continues so a transient ack
/// error does not stall the cancel pipeline.
pub fn spawn_wildcard_cancel_subscriber(
    pubsub: Arc<dyn Pubsub>,
    subject_pattern: String,
    subject_prefix: String,
    registry: CancelRegistry<String>,
    log_target: &'static str,
) {
    // `tracing` macros require `target` as a compile-time literal,
    // so the runtime `log_target` is surfaced as a structured field
    // (`transport`) on each event instead. Operators filtering on
    // `klieo.cancel` see both transports unified; field-level
    // filtering on `transport = "a2a.cancel"` / `"mcp.cancel"`
    // recovers the per-transport view the inline loops previously
    // provided.
    tokio::spawn(async move {
        let durable =
            crate::ids::DurableName::new(format!("klieo-cancel-{}", uuid::Uuid::new_v4()));
        let mut stream = match pubsub.subscribe(&subject_pattern, durable).await {
            Ok(s) => s,
            Err(e) => {
                tracing::error!(
                    target: "klieo.cancel",
                    transport = %log_target,
                    pattern = %subject_pattern,
                    error = %e,
                    "wildcard cancel subscription failed; \
                     replica falls back to single-replica semantics",
                );
                return;
            }
        };
        while let Some(item) = tokio_stream::StreamExt::next(&mut stream).await {
            match item {
                Ok(msg) => {
                    let id = msg
                        .subject
                        .strip_prefix(subject_prefix.as_str())
                        .unwrap_or_default()
                        .to_string();
                    let subject = msg.subject.clone();
                    let ack = msg.ack;
                    let span = tracing::info_span!(
                        "cancel_dispatch",
                        messaging.system = "klieo-bus",
                        messaging.destination = %subject,
                        messaging.operation = "receive",
                        transport = %log_target,
                    );
                    // Extract publisher's W3C tracecontext from bus
                    // headers so the dispatch span stitches under the
                    // originating replica's span (cluster 0.23,
                    // ADR-023). Requires the `otel` feature.
                    #[cfg(feature = "otel")]
                    {
                        use tracing_opentelemetry::OpenTelemetrySpanExt as _;
                        let parent_cx = crate::bus::extract_traceparent(&msg.headers);
                        span.set_parent(parent_cx);
                    }
                    use tracing::Instrument as _;
                    let registry = registry.clone();
                    async move {
                        if let Err(e) = ack.ack().await {
                            tracing::warn!(
                                target: "klieo.cancel",
                                transport = %log_target,
                                subject = %subject,
                                error = %e,
                                "cancel ack failed; expect redelivery",
                            );
                        }
                        if !id.is_empty() {
                            let hit = registry.cancel(&id);
                            tracing::debug!(
                                target: "klieo.cancel",
                                transport = %log_target,
                                id = %id,
                                hit = hit,
                                "cancel signal dispatched",
                            );
                        }
                    }
                    .instrument(span)
                    .await;
                }
                Err(e) => {
                    tracing::warn!(
                        target: "klieo.cancel",
                        transport = %log_target,
                        error = %e,
                        "cancel subscription stream error; continuing",
                    );
                }
            }
        }
    });
}

/// Drop-time helper that publishes a cross-replica cancel signal on
/// the supplied subject. Returns immediately; the publish is best-
/// effort via `tokio::spawn`. If no tokio runtime is active (e.g.
/// synchronous shutdown), the publish is skipped and a `warn` is
/// emitted instead of panicking.
///
/// Caller is expected to invoke this from within `Drop::drop` AFTER
/// the local `DropGuard` has fired — the bus publish merely fans the
/// signal out to peer replicas. `subject` is moved in (use
/// [`std::mem::take`] on the field). Empty `subject` short-circuits
/// to a no-op so call-sites do not have to guard separately.
///
/// `log_target` is surfaced as a structured `transport` field on
/// `tracing` events rather than the macro `target:` slot, because
/// `tracing` requires the target to be a compile-time literal. All
/// events fire under the unified `klieo.cancel` target — operators
/// filter on `transport = "a2a.cancel"` / `"mcp.cancel"` for the
/// per-transport view. Matches the precedent set by
/// [`spawn_wildcard_cancel_subscriber`] in this module.
///
/// `permits` optionally bounds the concurrency of in-flight drop
/// publishes per transport. When `Some(sem)` the call uses
/// `try_acquire_owned` — on saturation the publish is dropped + a
/// `warn` fires (preferred over flooding the tokio runtime with
/// unbounded spawns under a slow bus backend). The acquired permit
/// is moved into the spawned task and released on task end. When
/// `None` the historical unbounded behaviour is preserved so
/// pre-existing call sites keep their semantics until they wire a
/// semaphore in.
pub fn spawn_drop_publish(
    pubsub: Arc<dyn Pubsub>,
    subject: String,
    log_target: &'static str,
    permits: Option<Arc<tokio::sync::Semaphore>>,
    trace_headers: Headers,
) {
    if subject.is_empty() {
        return;
    }
    let permit = match permits.as_ref() {
        Some(sem) => match sem.clone().try_acquire_owned() {
            Ok(p) => Some(p),
            Err(_) => {
                tracing::warn!(
                    target: "klieo.cancel",
                    transport = log_target,
                    subject = %subject,
                    available_permits = sem.available_permits(),
                    "cancel publish dropped: concurrency cap reached",
                );
                return;
            }
        },
        None => None,
    };
    match tokio::runtime::Handle::try_current() {
        Ok(handle) => {
            handle.spawn(async move {
                let _permit = permit;
                if let Err(e) = pubsub.publish(&subject, Bytes::new(), trace_headers).await {
                    tracing::warn!(
                        target: "klieo.cancel",
                        transport = %log_target,
                        subject = %subject,
                        error = %e,
                        "cancel publish failed; cross-replica cancel degraded",
                    );
                }
            });
        }
        Err(_) => {
            tracing::warn!(
                target: "klieo.cancel",
                transport = %log_target,
                subject = %subject,
                "cancel publish skipped: no tokio runtime active during Drop",
            );
        }
    }
}

/// Spawn a best-effort bus publish bounded by `permits`. If the
/// semaphore has no available permits the publish is dropped + a
/// `tracing::warn!` fires; the call returns immediately and never
/// blocks the caller. Used by SSE yield-publish hot paths
/// (`stream_tools_call`, `TaskEventSink::send`) where dropping a
/// non-critical fanout event is preferable to flooding the runtime
/// with unbounded tasks.
///
/// `log_target` is a structured `transport` field on a
/// `target: "klieo.cancel"` event (tracing requires const-literal
/// targets).
pub fn spawn_publish_bounded(
    pubsub: Arc<dyn Pubsub>,
    subject: String,
    payload: Bytes,
    log_target: &'static str,
    permits: Arc<tokio::sync::Semaphore>,
    trace_headers: Headers,
) {
    let permit = match permits.clone().try_acquire_owned() {
        Ok(p) => p,
        Err(_) => {
            tracing::warn!(
                target: "klieo.cancel",
                transport = log_target,
                subject = %subject,
                available_permits = permits.available_permits(),
                "publish dropped: concurrency cap reached",
            );
            return;
        }
    };
    match tokio::runtime::Handle::try_current() {
        Ok(handle) => {
            handle.spawn(async move {
                let _permit = permit;
                if let Err(e) = pubsub.publish(&subject, payload, trace_headers).await {
                    tracing::warn!(
                        target: "klieo.cancel",
                        transport = log_target,
                        subject = %subject,
                        error = %e,
                        "publish failed",
                    );
                }
            });
        }
        Err(_) => {
            tracing::warn!(
                target: "klieo.cancel",
                transport = log_target,
                subject = %subject,
                "publish skipped: no tokio runtime active",
            );
        }
    }
}

/// Stream wrapper that deregisters an entry keyed by `id` from
/// `registry` on `Drop`. Use in HTTP-layer streaming paths to ensure
/// the per-invoke registry entry installed at invoke start is
/// removed when the response body terminates (final event, server
/// cancel, or client disconnect).
///
/// `Stream` is implemented by delegating `poll_next` to the inner
/// stream; this is a transparent wrapper at the polling layer.
///
/// `K` is fixed at `String` because both current call sites (A2A
/// `task_id`, MCP `progressToken`) use that key type. Generify
/// further when a non-`String` key appears.
///
/// An empty `id` is treated as a no-op on drop so callers that lack
/// a meaningful id (e.g. A2A `SendStreamingMessage` paths where
/// `task_id` is not known at SSE-build time) can wrap unconditionally.
pub struct RegistryDeregisterOnDrop<S> {
    inner: S,
    registry: CancelRegistry<String>,
    id: String,
}

impl<S> RegistryDeregisterOnDrop<S> {
    /// Wrap `inner` so that on `Drop` the entry keyed by `id` is
    /// removed from `registry`. Empty `id` is allowed and yields a
    /// no-op on drop.
    pub fn new(inner: S, registry: CancelRegistry<String>, id: String) -> Self {
        Self {
            inner,
            registry,
            id,
        }
    }
}

impl<S: Stream + Unpin> Stream for RegistryDeregisterOnDrop<S> {
    type Item = S::Item;
    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.inner).poll_next(cx)
    }
}

impl<S> Drop for RegistryDeregisterOnDrop<S> {
    fn drop(&mut self) {
        if !self.id.is_empty() {
            self.registry.deregister(&self.id);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bus::MsgStream;
    use crate::ids::DurableName;
    use async_trait::async_trait;
    use std::time::Duration;
    use tokio::sync::{mpsc, Semaphore};

    /// Recording `Pubsub` used to assert publish side-effects without
    /// pulling in `klieo-bus-memory` as a dev-dep (which would form a
    /// cyclic compile graph). Publishes push `(subject, payload)` onto
    /// an mpsc; `subscribe` is unused by these tests and yields an
    /// `Unimplemented` error so a misuse is loud.
    struct RecordingPubsub {
        tx: mpsc::UnboundedSender<(String, Bytes)>,
    }

    #[async_trait]
    impl Pubsub for RecordingPubsub {
        async fn publish(
            &self,
            subject: &str,
            payload: Bytes,
            _headers: Headers,
        ) -> Result<(), crate::error::BusError> {
            let _ = self.tx.send((subject.to_string(), payload));
            Ok(())
        }

        async fn subscribe(
            &self,
            _subject: &str,
            _durable: DurableName,
        ) -> Result<MsgStream, crate::error::BusError> {
            Err(crate::error::BusError::Unsupported(
                "RecordingPubsub::subscribe is not used in tests".into(),
            ))
        }
    }

    fn recording_pubsub() -> (Arc<dyn Pubsub>, mpsc::UnboundedReceiver<(String, Bytes)>) {
        let (tx, rx) = mpsc::unbounded_channel();
        (Arc::new(RecordingPubsub { tx }), rx)
    }

    #[test]
    fn register_then_cancel_fires_token() {
        let registry: CancelRegistry<String> = CancelRegistry::new();
        let token = CancellationToken::new();
        registry.register("task-1".to_string(), token.clone());

        assert!(!token.is_cancelled());

        let hit = registry.cancel(&"task-1".to_string());

        assert!(hit, "cancel on registered id returns true");
        assert!(token.is_cancelled(), "registered token fires on cancel");
    }

    #[test]
    fn deregister_removes_entry() {
        let registry: CancelRegistry<String> = CancelRegistry::new();
        let token = CancellationToken::new();
        registry.register("task-2".to_string(), token.clone());

        registry.deregister(&"task-2".to_string());

        let hit = registry.cancel(&"task-2".to_string());

        assert!(!hit, "cancel after deregister returns false");
        assert!(
            !token.is_cancelled(),
            "deregistered token is not fired by later cancel",
        );
    }

    #[test]
    fn cancel_on_missing_key_is_noop() {
        let registry: CancelRegistry<String> = CancelRegistry::new();

        let hit = registry.cancel(&"never".to_string());

        assert!(!hit, "cancel on absent id returns false");
    }

    #[tokio::test(flavor = "current_thread")]
    async fn spawn_publish_bounded_drops_on_saturation() {
        let (pubsub, mut rx) = recording_pubsub();
        // Keep a second sender alive so the channel stays open while
        // the bounded helper drops its pubsub clone on the saturation
        // path; otherwise `rx.recv()` returns `Ok(None)` immediately on
        // sender-drop and masquerades as a delivered message.
        let _keepalive = pubsub.clone();
        let permits = Arc::new(Semaphore::new(0));

        spawn_publish_bounded(
            pubsub,
            "subj.saturated".to_string(),
            Bytes::from_static(b"x"),
            "test.cancel",
            permits,
            Headers::default(),
        );

        let recv = tokio::time::timeout(Duration::from_millis(100), rx.recv()).await;
        assert!(
            recv.is_err(),
            "no publish should reach recorder under zero-permit saturation"
        );
    }

    #[tokio::test(flavor = "current_thread")]
    async fn spawn_publish_bounded_publishes_with_available_permit() {
        let (pubsub, mut rx) = recording_pubsub();
        let permits = Arc::new(Semaphore::new(1));

        spawn_publish_bounded(
            pubsub,
            "subj.ok".to_string(),
            Bytes::from_static(b"hello"),
            "test.cancel",
            permits,
            Headers::default(),
        );

        let recv = tokio::time::timeout(Duration::from_millis(500), rx.recv())
            .await
            .expect("publish should arrive within 500ms when a permit is available");
        let (subject, payload) = recv.expect("recorder channel closed unexpectedly");
        assert_eq!(subject, "subj.ok");
        assert_eq!(payload.as_ref(), b"hello");
    }

    #[tokio::test(flavor = "current_thread")]
    async fn spawn_drop_publish_with_permits_drops_on_saturation() {
        let (pubsub, mut rx) = recording_pubsub();
        // See `spawn_publish_bounded_drops_on_saturation` — keep the
        // sender side of the recorder alive across the saturation path
        // so `rx.recv()` blocks for the full timeout instead of
        // returning `Ok(None)` on the inner Arc drop.
        let _keepalive = pubsub.clone();
        let permits = Arc::new(Semaphore::new(0));

        spawn_drop_publish(
            pubsub,
            "subj.drop.saturated".to_string(),
            "test.cancel",
            Some(permits),
            Headers::default(),
        );

        let recv = tokio::time::timeout(Duration::from_millis(100), rx.recv()).await;
        assert!(
            recv.is_err(),
            "spawn_drop_publish with zero permits must not publish",
        );
    }
}