klieo-core 0.40.0

//! Inter-agent bus traits — Pubsub, RequestReply, KvStore, JobQueue.
//!
//! Trait shapes mirror NATS JetStream + KV semantics so the production
//! impl (`klieo-bus-nats`) is a thin adapter, and the in-process impl
//! (`klieo-bus-memory`) can faithfully simulate them. See the spec for
//! reliability invariants.

use crate::error::BusError;
use crate::ids::{DurableName, JobId};
use async_trait::async_trait;
use bytes::Bytes;
use futures_core::Stream;
use std::collections::HashMap;
use std::pin::Pin;
use std::time::Duration;

/// Opaque headers accompanying a bus message.
pub type Headers = HashMap<String, String>;

/// One message delivered to a subscriber.
pub struct Msg {
    /// Subject the message was published on.
    pub subject: String,
    /// Payload bytes.
    pub payload: Bytes,
    /// Headers.
    pub headers: Headers,
    /// Acknowledgement handle. The impl provides the underlying mechanism;
    /// callers must invoke exactly one of `ack` / `nak` / `term` per
    /// message or rely on visibility-timeout redelivery.
    pub ack: AckHandle,
}

impl std::fmt::Debug for Msg {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Msg")
            .field("subject", &self.subject)
            .field("payload_len", &self.payload.len())
            .field("headers", &self.headers)
            .finish()
    }
}

/// Acknowledgement handle attached to a delivered [`Msg`].
pub struct AckHandle(Box<dyn AckHandleImpl>);

impl AckHandle {
    /// Construct from an impl. Bus implementations call this when delivering messages.
    pub fn new(inner: Box<dyn AckHandleImpl>) -> Self {
        Self(inner)
    }
}

/// Implementor side of an ack handle. Exposed so impls can supply backend
/// behaviour without callers reaching into private types.
#[async_trait]
pub trait AckHandleImpl: Send + Sync {
    /// Acknowledge successful processing.
    async fn ack(self: Box<Self>) -> Result<(), BusError>;
    /// Negative-ack with optional redelivery delay. Triggers redelivery
    /// after `delay`.
    async fn nak(self: Box<Self>, delay: Duration) -> Result<(), BusError>;
    /// Terminate the message (do not redeliver).
    async fn term(self: Box<Self>) -> Result<(), BusError>;
}

impl AckHandle {
    /// Acknowledge.
    pub async fn ack(self) -> Result<(), BusError> {
        self.0.ack().await
    }
    /// Negative-ack.
    pub async fn nak(self, delay: Duration) -> Result<(), BusError> {
        self.0.nak(delay).await
    }
    /// Terminate.
    pub async fn term(self) -> Result<(), BusError> {
        self.0.term().await
    }
}

/// Stream of messages delivered to a subscriber.
pub type MsgStream = Pin<Box<dyn Stream<Item = Result<Msg, BusError>> + Send + 'static>>;

/// Pub/sub interface (subject-based, durable consumers).
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::{Headers, Pubsub};
/// use bytes::Bytes;
/// let (pubsub, _, _, _) = noop_bus();
/// pubsub.publish("subject.demo", Bytes::from_static(b"hi"), Headers::new())
///     .await.unwrap();
/// # });
/// ```
#[async_trait]
pub trait Pubsub: Send + Sync {
    /// Publish `payload` on `subject` with `headers`.
    async fn publish(
        &self,
        subject: &str,
        payload: Bytes,
        headers: Headers,
    ) -> Result<(), BusError>;

    /// Subscribe with a durable consumer name. Multiple calls with the
    /// same `durable` form a competing-consumer group sharing replays.
    async fn subscribe(&self, subject: &str, durable: DurableName) -> Result<MsgStream, BusError>;
}

/// Validate that `token` is safe to embed as a single NATS subject segment.
///
/// NATS reserves `.` (separator), `*` (single-token wildcard), and `>`
/// (greedy wildcard) as subject metacharacters. Whitespace, control
/// characters, and non-ASCII bytes are also rejected so wire-level
/// subjects remain printable and the segment cannot collapse the
/// caller-controlled subject namespace.
///
/// Use this at every site that embeds a caller-influenced identifier
/// (progressToken, task id, stream id, …) into a publish or subscribe
/// subject. Skipping validation is a cross-tenant data-leak vector
/// (CWE-74 subject injection): a single-character token like `>` or
/// `*` would subscribe to every other tenant's stream.
///
/// Returns `BusError::Invalid` on any rejected segment.
pub fn validate_subject_token(token: &str) -> Result<(), BusError> {
    if token.is_empty() {
        return Err(BusError::Invalid("subject segment is empty".into()));
    }
    for byte in token.bytes() {
        let forbidden = matches!(byte, b'.' | b'*' | b'>')
            || byte.is_ascii_whitespace()
            || byte.is_ascii_control()
            || !byte.is_ascii();
        if forbidden {
            return Err(BusError::Invalid(format!(
                "subject segment contains forbidden character (byte 0x{byte:02x})"
            )));
        }
    }
    Ok(())
}

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

    #[test]
    fn accepts_uuid_like_token() {
        validate_subject_token("550e8400-e29b-41d4-a716-446655440000").unwrap();
        validate_subject_token("task_42").unwrap();
        validate_subject_token("abc123").unwrap();
    }

    #[test]
    fn rejects_empty() {
        let e = validate_subject_token("").unwrap_err();
        assert!(matches!(e, BusError::Invalid(_)));
    }

    #[test]
    fn rejects_greedy_wildcard() {
        assert!(matches!(
            validate_subject_token(">"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_single_token_wildcard() {
        assert!(matches!(
            validate_subject_token("*"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_dot_separator() {
        assert!(matches!(
            validate_subject_token("a.b"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_whitespace_and_control() {
        assert!(matches!(
            validate_subject_token("a b"),
            Err(BusError::Invalid(_))
        ));
        assert!(matches!(
            validate_subject_token("a\nb"),
            Err(BusError::Invalid(_))
        ));
        assert!(matches!(
            validate_subject_token("a\tb"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_non_ascii() {
        assert!(matches!(
            validate_subject_token("café"),
            Err(BusError::Invalid(_))
        ));
    }
}

/// Synchronous request/response over the bus.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::{BusError, RequestReply};
/// use bytes::Bytes;
/// use std::time::Duration;
/// let (_, request_reply, _, _) = noop_bus();
/// let err = request_reply
///     .request("svc.add", Bytes::from_static(b"1"), Duration::from_secs(1))
///     .await
///     .unwrap_err();
/// assert!(matches!(err, BusError::NotFound(_)));
/// # });
/// ```
#[async_trait]
pub trait RequestReply: Send + Sync {
    /// Send a request and await one reply, bounded by `timeout`.
    async fn request(
        &self,
        subject: &str,
        payload: Bytes,
        timeout: Duration,
    ) -> Result<Bytes, BusError>;
}

/// CAS revision returned by KV writes.
pub type Revision = u64;

/// One KV entry.
#[derive(Debug, Clone)]
pub struct KvEntry {
    /// Stored value.
    pub value: Bytes,
    /// Revision number after the last write.
    pub revision: Revision,
}

/// Bucket-keyed durable KV store with CAS.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::KvStore;
/// use bytes::Bytes;
/// let (_, _, kv, _) = noop_bus();
/// let rev = kv.put("bucket", "key", Bytes::from_static(b"v")).await.unwrap();
/// assert_eq!(rev, 1);
/// # });
/// ```
#[async_trait]
pub trait KvStore: Send + Sync {
    /// Read a value.
    async fn get(&self, bucket: &str, key: &str) -> Result<Option<KvEntry>, BusError>;

    /// Unconditional write. Returns the new revision.
    async fn put(&self, bucket: &str, key: &str, value: Bytes) -> Result<Revision, BusError>;

    /// Compare-and-set. `expected = None` requires the key to be absent;
    /// `Some(rev)` requires the current revision to equal `rev`. Returns
    /// the new revision on success.
    async fn cas(
        &self,
        bucket: &str,
        key: &str,
        value: Bytes,
        expected: Option<Revision>,
    ) -> Result<Revision, BusError>;

    /// Delete a key.
    async fn delete(&self, bucket: &str, key: &str) -> Result<(), BusError>;

    /// Acquire an exclusive lease over `key` for `ttl`. Implementations
    /// should hold the lease as long as the returned [`Lease`] is live
    /// and call its `heartbeat` to extend.
    async fn lease(&self, bucket: &str, key: &str, ttl: Duration) -> Result<Lease, BusError>;

    /// Enumerate keys under `bucket`. Default impl returns
    /// `BusError::Unsupported` — backends that can scan (in-mem,
    /// NATS JetStream KV) override this. The resume-buffer sweeper
    /// (see `klieo-core::resume`) calls this to walk all buckets;
    /// backends that cannot enumerate skip proactive eviction and
    /// rely on access-time TTL checks.
    async fn keys(&self, bucket: &str) -> Result<Vec<String>, BusError> {
        let _ = bucket;
        Err(BusError::Unsupported(
            "keys() not implemented for this KvStore".into(),
        ))
    }
}

/// Lease handle. Heartbeat to extend; drop to release.
pub struct Lease(Box<dyn LeaseImpl>);

impl Lease {
    /// Construct from an impl. KvStore implementations call this when granting a lease.
    pub fn new(inner: Box<dyn LeaseImpl>) -> Self {
        Self(inner)
    }
}

/// Implementor side of a lease.
#[async_trait]
pub trait LeaseImpl: Send + Sync {
    /// Extend the TTL.
    async fn heartbeat(&self) -> Result<(), BusError>;
}

impl Lease {
    /// Extend the TTL.
    pub async fn heartbeat(&self) -> Result<(), BusError> {
        self.0.heartbeat().await
    }
}

/// Job enqueued for durable processing.
#[derive(Debug, Clone)]
pub struct Job {
    /// Job payload.
    pub payload: Bytes,
    /// Optional dedup key — if set, the impl writes a `dedup.<queue>`
    /// idempotency record before invoking the handler.
    pub dedup_key: Option<String>,
    /// Maximum redelivery attempts before routing to the DLQ subject.
    /// `None` = use queue default (5).
    pub max_attempts: Option<u32>,
}

impl Job {
    /// Build a job with default settings from raw bytes.
    pub fn new(payload: impl Into<Bytes>) -> Self {
        Self {
            payload: payload.into(),
            dedup_key: None,
            max_attempts: None,
        }
    }

    /// Start a fluent builder. Chain `.dedup(k)` and `.max_attempts(n)`
    /// then call `.build()`. Field-assignment on the struct still works
    /// for callers that prefer it.
    pub fn builder(payload: impl Into<Bytes>) -> JobBuilder {
        JobBuilder {
            payload: payload.into(),
            dedup_key: None,
            max_attempts: None,
        }
    }
}

/// Fluent builder for [`Job`].
pub struct JobBuilder {
    payload: Bytes,
    dedup_key: Option<String>,
    max_attempts: Option<u32>,
}

impl JobBuilder {
    /// Set the dedup key. The queue impl writes a `dedup.<queue>`
    /// idempotency record before invoking the handler.
    pub fn dedup(mut self, key: impl Into<String>) -> Self {
        self.dedup_key = Some(key.into());
        self
    }

    /// Set the maximum redelivery attempts before routing to the DLQ.
    pub fn max_attempts(mut self, n: u32) -> Self {
        self.max_attempts = Some(n);
        self
    }

    /// Finalise the builder into a [`Job`].
    pub fn build(self) -> Job {
        Job {
            payload: self.payload,
            dedup_key: self.dedup_key,
            max_attempts: self.max_attempts,
        }
    }
}

/// Job claimed by a worker.
pub struct ClaimedJob {
    /// Job id.
    pub id: JobId,
    /// Job payload.
    pub payload: Bytes,
    /// Lease handle. Caller heartbeats; on drop without ack/nak/dlq the
    /// lease expires and the impl redelivers.
    pub lease: Lease,
    /// Internal handle the impl uses to mark the claim resolved.
    pub claim: ClaimHandle,
}

impl std::fmt::Debug for ClaimedJob {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ClaimedJob")
            .field("id", &self.id)
            .field("payload_len", &self.payload.len())
            .finish()
    }
}

/// Handle the impl uses to resolve a claim.
pub struct ClaimHandle(Box<dyn ClaimHandleImpl>);

impl ClaimHandle {
    /// Construct from an impl. JobQueue implementations call this when claiming a job.
    pub fn new(inner: Box<dyn ClaimHandleImpl>) -> Self {
        Self(inner)
    }
}

/// Implementor side of a claim handle.
#[async_trait]
pub trait ClaimHandleImpl: Send + Sync {
    /// Successful completion. Releases the lease, marks the job done.
    async fn ack(self: Box<Self>) -> Result<(), BusError>;
    /// Retry with backoff.
    async fn nak(self: Box<Self>, delay: Duration) -> Result<(), BusError>;
    /// Send to DLQ.
    async fn dead_letter(self: Box<Self>, reason: &str) -> Result<(), BusError>;
}

impl ClaimedJob {
    /// Heartbeat the lease.
    pub async fn heartbeat(&self) -> Result<(), BusError> {
        self.lease.heartbeat().await
    }
    /// Acknowledge successful completion.
    pub async fn ack(self) -> Result<(), BusError> {
        self.claim.0.ack().await
    }
    /// Negative-ack; retry after `delay`.
    pub async fn nak(self, delay: Duration) -> Result<(), BusError> {
        self.claim.0.nak(delay).await
    }
    /// Route to DLQ.
    pub async fn dead_letter(self, reason: &str) -> Result<(), BusError> {
        self.claim.0.dead_letter(reason).await
    }
}

/// Durable competing-consumer job queue. **No ordering guarantee** —
/// workloads requiring order must use [`Pubsub`] on a partitioned subject
/// with one consumer per partition instead.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::{Job, JobQueue};
/// use bytes::Bytes;
/// let (_, _, _, jobs) = noop_bus();
/// let id = jobs.enqueue("queue.work", Job::new(Bytes::from_static(b"payload"))).await.unwrap();
/// assert_eq!(id.0, "noop-0");
/// # });
/// ```
#[async_trait]
pub trait JobQueue: Send + Sync {
    /// Enqueue a job. Returns a stable id.
    async fn enqueue(&self, queue: &str, job: Job) -> Result<JobId, BusError>;

    /// Claim the next available job. Returns `None` when the queue is
    /// empty. Implementations may long-poll up to a small bounded
    /// duration before returning `None`.
    async fn claim(
        &self,
        queue: &str,
        worker_id: &str,
        lease_ttl: Duration,
    ) -> Result<Option<ClaimedJob>, BusError>;
}

/// Resolved bundle of bus handles ready to drop into an
/// [`crate::agent::AgentContext`] or an `App`.
///
/// Impl crates (`klieo-bus-memory`, `klieo-bus-nats`) provide `From`
/// conversions; downstream code typically writes
/// `BusHandles::from(MemoryBus::new())` and never destructures the
/// four sub-handles by name.
#[derive(Clone)]
pub struct BusHandles {
    /// Pub/sub.
    pub pubsub: std::sync::Arc<dyn Pubsub>,
    /// KV store.
    pub kv: std::sync::Arc<dyn KvStore>,
    /// Synchronous request/reply.
    pub request_reply: std::sync::Arc<dyn RequestReply>,
    /// Durable job queue.
    pub jobs: std::sync::Arc<dyn JobQueue>,
}

impl BusHandles {
    /// Build directly from four already-`Arc`-wrapped handles.
    /// Most callers go through an impl crate's `From` instead.
    pub fn new(
        pubsub: std::sync::Arc<dyn Pubsub>,
        kv: std::sync::Arc<dyn KvStore>,
        request_reply: std::sync::Arc<dyn RequestReply>,
        jobs: std::sync::Arc<dyn JobQueue>,
    ) -> Self {
        Self {
            pubsub,
            kv,
            request_reply,
            jobs,
        }
    }
}

// ─── W3C tracecontext propagation helpers — cluster 0.23 ──────────────
//
// Bus messages crossing replica boundaries carry W3C tracecontext
// (W3C TR/trace-context) in the standard `traceparent` + `tracestate`
// `Headers` keys so OTEL consumers on the receiving side can stitch
// their spans as children of the publisher's span. No API break —
// `Headers` is already `HashMap<String, String>` on `Pubsub::publish`.
//
// No-op when no global OTEL trace provider is installed: the
// TraceContextPropagator silently emits no headers when the current
// context carries no span. Pre-0.23 deployments that ignore the
// headers continue to work unchanged.
//
// These helpers require the `otel` feature flag. Crates that need
// tracecontext propagation must declare:
//   klieo-core = { ..., features = ["otel"] }

#[cfg(feature = "otel")]
use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
#[cfg(feature = "otel")]
use opentelemetry_sdk::propagation::TraceContextPropagator;

/// Inject the current OpenTelemetry context's tracecontext into
/// `headers` under the standard W3C keys (`traceparent`, optionally
/// `tracestate`).
///
/// Call BEFORE `Pubsub::publish` from inside an active OTEL span
/// (typically reached via `tracing-opentelemetry`'s
/// `OpenTelemetrySpanExt::context` on `Span::current()`). Cluster
/// 0.23's impl crates do the bridging — T1 only supplies the
/// lower-level header injection given an explicit Context.
///
/// Requires feature `otel`.
#[cfg(feature = "otel")]
pub fn inject_traceparent(headers: &mut Headers, context: &opentelemetry::Context) {
    let propagator = TraceContextPropagator::new();
    let mut injector = HeaderMapInjector(headers);
    propagator.inject_context(context, &mut injector);
}

/// Extract a W3C tracecontext from `headers`. Returns the default
/// empty Context when no headers are set; callers should set the
/// result as the parent of subsequent spans only if they care about
/// cross-replica stitching.
///
/// Call AFTER `MsgStream::next` resolves.
///
/// Requires feature `otel`.
#[cfg(feature = "otel")]
pub fn extract_traceparent(headers: &Headers) -> opentelemetry::Context {
    let propagator = TraceContextPropagator::new();
    let extractor = HeaderMapExtractor(headers);
    propagator.extract(&extractor)
}

#[cfg(feature = "otel")]
struct HeaderMapInjector<'a>(&'a mut Headers);

#[cfg(feature = "otel")]
impl<'a> Injector for HeaderMapInjector<'a> {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_string(), value);
    }
}

#[cfg(feature = "otel")]
struct HeaderMapExtractor<'a>(&'a Headers);

#[cfg(feature = "otel")]
impl<'a> Extractor for HeaderMapExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).map(|s| s.as_str())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(|s| s.as_str()).collect()
    }
}

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

    #[allow(dead_code)]
    fn _assert_dyn_pubsub(_: &dyn Pubsub) {}
    #[allow(dead_code)]
    fn _assert_dyn_request(_: &dyn RequestReply) {}
    #[allow(dead_code)]
    fn _assert_dyn_kv(_: &dyn KvStore) {}
    #[allow(dead_code)]
    fn _assert_dyn_jobs(_: &dyn JobQueue) {}

    #[test]
    fn job_builder_defaults_match_job_new() {
        let by_builder = Job::builder(Bytes::from_static(b"x")).build();
        let by_new = Job::new(Bytes::from_static(b"x"));
        assert_eq!(by_builder.payload, by_new.payload);
        assert_eq!(by_builder.dedup_key, by_new.dedup_key);
        assert_eq!(by_builder.max_attempts, by_new.max_attempts);
    }

    #[test]
    fn job_builder_sets_dedup_and_max_attempts() {
        let job = Job::builder(Bytes::from_static(b"payload"))
            .dedup("idempotency-key-42")
            .max_attempts(7)
            .build();
        assert_eq!(job.payload, Bytes::from_static(b"payload"));
        assert_eq!(job.dedup_key.as_deref(), Some("idempotency-key-42"));
        assert_eq!(job.max_attempts, Some(7));
    }

    #[test]
    fn job_builder_accepts_string_dedup_via_into() {
        let owned = String::from("k");
        let job = Job::builder(Bytes::from_static(b"x")).dedup(owned).build();
        assert_eq!(job.dedup_key.as_deref(), Some("k"));
    }

    #[test]
    fn job_field_assignment_still_compiles() {
        let mut job = Job::new(Bytes::from_static(b"x"));
        job.dedup_key = Some("k".into());
        job.max_attempts = Some(3);
        assert_eq!(job.dedup_key.as_deref(), Some("k"));
        assert_eq!(job.max_attempts, Some(3));
    }

    #[test]
    fn bus_error_unsupported_renders_message() {
        let e = BusError::Unsupported("keys() not implemented".into());
        assert_eq!(
            e.to_string(),
            "unsupported operation: keys() not implemented"
        );
    }

    #[cfg(feature = "otel")]
    #[test]
    fn tracecontext_inject_then_extract_roundtrip() {
        use opentelemetry::trace::{
            SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState,
        };

        let trace_id = TraceId::from_hex("0123456789abcdef0123456789abcdef").unwrap();
        let span_id = SpanId::from_hex("0123456789abcdef").unwrap();
        let span_ctx = SpanContext::new(
            trace_id,
            span_id,
            TraceFlags::SAMPLED,
            true,
            TraceState::default(),
        );
        let cx = opentelemetry::Context::new().with_remote_span_context(span_ctx);

        let mut headers: Headers = HashMap::new();
        inject_traceparent(&mut headers, &cx);
        assert!(
            headers.contains_key("traceparent"),
            "traceparent header must be injected"
        );

        let extracted = extract_traceparent(&headers);
        let extracted_span_ctx = extracted.span().span_context().clone();
        assert_eq!(extracted_span_ctx.trace_id(), trace_id);
        assert_eq!(extracted_span_ctx.span_id(), span_id);
    }
}