klieo-core 0.40.0

//! Per-stream SSE event buffer used by HTTP/SSE transports for
//! resumption via `Last-Event-ID` / `klieo/tools/resume`.
//!
//! The [`ResumeBuffer`] trait is opaque on `stream_id` (`klieo-a2a`
//! passes task IDs; `klieo-mcp-server` passes stringified
//! progressTokens). [`NoopResumeBuffer`] is the zero-cost default for
//! adopters who do not enable resumption. `KvResumeBuffer` (added in
//! T5) is the default `KvStore`-backed impl.

use async_trait::async_trait;
use bytes::Bytes;
use futures_core::Stream;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use thiserror::Error;

/// Maximum events retained per stream before the oldest is evicted.
pub const RESUME_RING_CAP: usize = 256;

/// Idle TTL — buckets with no writes for this duration are evicted by
/// the sweeper.
pub const RESUME_IDLE_TTL: Duration = Duration::from_secs(5 * 60);

/// How often the sweeper walks buckets.
pub const RESUME_SWEEP_INTERVAL: Duration = Duration::from_secs(60);

/// Grace period between a stream being marked terminal and its bucket
/// being evicted by the sweeper.
pub const RESUME_TERMINAL_GRACE: Duration = Duration::from_secs(30);

/// Per-stream SSE event buffer. `stream_id` is opaque to the buffer —
/// `klieo-a2a` passes task IDs; `klieo-mcp-server` passes stringified
/// progressTokens.
#[async_trait]
pub trait ResumeBuffer: Send + Sync {
    /// Append one event under `stream_id` with the given `event_id`.
    /// Implementations evict the oldest event when the per-stream
    /// ring is full (cap = [`RESUME_RING_CAP`]).
    ///
    /// **Caller contract: `event_id` must be strictly monotonically
    /// increasing per `stream_id`.** Implementations are free to drop,
    /// reorder, or silently overwrite out-of-order writes; A2A's
    /// per-task counter and MCP's per-progressToken counter both
    /// satisfy this contract by construction.
    async fn record(
        &self,
        stream_id: &str,
        event_id: u64,
        payload: Bytes,
    ) -> Result<(), ResumeError>;

    /// Replay events with id > since_id, in monotonic order. Returns
    /// [`ResumeError::Expired`] if the cursor is before the oldest
    /// retained id; [`ResumeError::NotFound`] if no buffer exists for
    /// the stream.
    async fn replay(
        &self,
        stream_id: &str,
        since_id: u64,
    ) -> Result<Pin<Box<dyn Stream<Item = (u64, Bytes)> + Send>>, ResumeError>;

    /// Mark the stream terminal. Buffer is eligible for sweeper
    /// eviction after [`RESUME_TERMINAL_GRACE`].
    async fn close(&self, stream_id: &str) -> Result<(), ResumeError>;

    /// Return `true` if the stream has been marked terminal via `close`.
    /// Implementations that do not track terminal state return `false`
    /// (conservative: assume live). Used by `stream_resume` to decide
    /// whether to attach a live-tail subscription after replay drains.
    async fn is_terminal(&self, _stream_id: &str) -> Result<bool, ResumeError> {
        Ok(false)
    }
}

/// Failure modes for [`ResumeBuffer`].
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ResumeError {
    /// The requested cursor is before the oldest retained id.
    #[error("resume window expired (oldest retained id > requested {since_id})")]
    Expired {
        /// Cursor the caller supplied.
        since_id: u64,
    },
    /// No buffer exists for this stream id.
    #[error("stream not found: {0}")]
    NotFound(String),
    /// Backend error (KV store failure, serialisation error, etc).
    #[error("buffer backend: {0}")]
    Backend(#[source] Box<dyn std::error::Error + Send + Sync>),
}

/// Zero-cost default — `record` is a no-op, `replay` returns
/// [`ResumeError::NotFound`], `close` is a no-op. Used by builders
/// when resumption is not enabled.
pub struct NoopResumeBuffer;

#[async_trait]
impl ResumeBuffer for NoopResumeBuffer {
    async fn record(
        &self,
        _stream_id: &str,
        _event_id: u64,
        _payload: Bytes,
    ) -> Result<(), ResumeError> {
        Ok(())
    }

    async fn replay(
        &self,
        stream_id: &str,
        _since_id: u64,
    ) -> Result<Pin<Box<dyn Stream<Item = (u64, Bytes)> + Send>>, ResumeError> {
        Err(ResumeError::NotFound(stream_id.to_string()))
    }

    async fn close(&self, _stream_id: &str) -> Result<(), ResumeError> {
        Ok(())
    }
}

use crate::bus::KvStore;
use crate::error::BusError;

/// Prefix for resume-buffer KV buckets: `klieo.resume.{stream_id}`.
const RESUME_BUCKET_PREFIX: &str = "klieo.resume.";
/// Directory bucket listing all active stream IDs for the sweeper.
const RESUME_DIRECTORY_BUCKET: &str = "klieo.resume.index";
const META_KEY: &str = "meta";
const EVENT_KEY_PREFIX: &str = "event/";

/// Sweeper interval tuning. Default values match [`RESUME_IDLE_TTL`],
/// [`RESUME_SWEEP_INTERVAL`], and [`RESUME_TERMINAL_GRACE`]. Tests
/// override with short values to exercise the sweeper without sleeping
/// minutes.
#[derive(Debug, Clone)]
pub struct SweeperTuning {
    /// Buckets with no writes for this duration are eligible for eviction.
    pub idle_ttl: Duration,
    /// How often the sweeper walks the directory bucket.
    pub sweep_interval: Duration,
    /// Grace period between a stream being marked terminal and its eviction.
    pub terminal_grace: Duration,
}

impl Default for SweeperTuning {
    fn default() -> Self {
        Self {
            idle_ttl: RESUME_IDLE_TTL,
            sweep_interval: RESUME_SWEEP_INTERVAL,
            terminal_grace: RESUME_TERMINAL_GRACE,
        }
    }
}

/// `KvStore`-backed [`ResumeBuffer`]. Layout per stream:
///
/// - `klieo.resume.{stream_id}` / `meta` → JSON [`BufferMeta`].
/// - `klieo.resume.{stream_id}` / `event/{id}` → raw event payload bytes.
/// - `klieo.resume.index` / `{stream_id}` → `"1"` (directory for sweeper).
pub struct KvResumeBuffer {
    kv: Arc<dyn KvStore>,
    tuning: SweeperTuning,
}

/// Per-stream metadata persisted as JSON under the `meta` key.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
pub(crate) struct BufferMeta {
    /// Next id to assign. `next_id - 1` is the highest retained id.
    pub next_id: u64,
    /// Oldest retained id. `event/{oldest_id}` exists until ring
    /// eviction bumps it past the requested id.
    pub oldest_id: u64,
    /// Unix milliseconds of the most recent `record` call. Stored in
    /// milliseconds so the sweeper can use sub-second TTLs in tests.
    pub last_write_unix_ms: u64,
    /// `true` once `close` has been called. Used by the sweeper for
    /// terminal-grace eviction.
    pub terminal: bool,
}

impl KvResumeBuffer {
    /// Build from any `Arc<dyn KvStore>` with default sweeper tuning.
    pub fn new(kv: Arc<dyn KvStore>) -> Self {
        Self {
            kv,
            tuning: SweeperTuning::default(),
        }
    }

    /// Build with custom sweeper tuning. Use in tests with short TTLs.
    pub fn with_tuning(kv: Arc<dyn KvStore>, tuning: SweeperTuning) -> Self {
        Self { kv, tuning }
    }

    fn bucket(stream_id: &str) -> String {
        format!("{RESUME_BUCKET_PREFIX}{stream_id}")
    }

    fn event_key(id: u64) -> String {
        format!("{EVENT_KEY_PREFIX}{id}")
    }

    fn now_unix_ms() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_millis() as u64)
            .unwrap_or(0)
    }

    pub(crate) async fn read_meta(&self, stream_id: &str) -> Result<BufferMeta, ResumeError> {
        let bucket = Self::bucket(stream_id);
        let entry = self
            .kv
            .get(&bucket, META_KEY)
            .await
            .map_err(|e| ResumeError::Backend(Box::new(e)))?;
        match entry {
            Some(entry) => serde_json::from_slice::<BufferMeta>(&entry.value)
                .map_err(|e| ResumeError::Backend(Box::new(e))),
            None => Err(ResumeError::NotFound(stream_id.to_string())),
        }
    }

    async fn write_meta(&self, stream_id: &str, meta: &BufferMeta) -> Result<(), ResumeError> {
        let bucket = Self::bucket(stream_id);
        let bytes = serde_json::to_vec(meta).map_err(|e| ResumeError::Backend(Box::new(e)))?;
        self.kv
            .put(&bucket, META_KEY, Bytes::from(bytes))
            .await
            .map_err(|e| ResumeError::Backend(Box::new(e)))?;
        Ok(())
    }
}

#[async_trait]
impl ResumeBuffer for KvResumeBuffer {
    #[tracing::instrument(
        skip_all,
        fields(
            db.system = "klieo-kv",
            db.namespace = "klieo-resume",
            db.operation = "record",
            klieo.resume.stream_id = %stream_id,
            klieo.resume.event_id = event_id,
        ),
        level = "debug",
        err,
    )]
    async fn record(
        &self,
        stream_id: &str,
        event_id: u64,
        payload: Bytes,
    ) -> Result<(), ResumeError> {
        let bucket = Self::bucket(stream_id);
        let mut meta = match self.read_meta(stream_id).await {
            Ok(m) => m,
            Err(ResumeError::NotFound(_)) => BufferMeta {
                next_id: event_id,
                oldest_id: event_id,
                last_write_unix_ms: Self::now_unix_ms(),
                terminal: false,
            },
            Err(other) => return Err(other),
        };

        // Write event payload before bumping next_id so a partial
        // failure here leaves meta untouched.
        self.kv
            .put(&bucket, &Self::event_key(event_id), payload)
            .await
            .map_err(|e| ResumeError::Backend(Box::new(e)))?;

        // Maintain the directory index so the sweeper can enumerate
        // active streams. Best-effort; missing-index just disables
        // proactive eviction for this stream.
        if let Err(e) = self
            .kv
            .put(RESUME_DIRECTORY_BUCKET, stream_id, Bytes::from_static(b"1"))
            .await
        {
            tracing::warn!(
                target: "klieo.resume",
                stream_id,
                error = %e,
                "resume directory index write failed; sweeper may miss this stream"
            );
        }

        // Keep the highest seen id monotonically.
        if event_id >= meta.next_id {
            meta.next_id = event_id + 1;
        }

        // Evict oldest slots while ring exceeds capacity.
        while meta.next_id.saturating_sub(meta.oldest_id) > RESUME_RING_CAP as u64 {
            let to_evict = meta.oldest_id;
            // Best-effort delete; benign if already absent.
            let _ = self.kv.delete(&bucket, &Self::event_key(to_evict)).await;
            meta.oldest_id = to_evict + 1;
        }

        meta.last_write_unix_ms = Self::now_unix_ms();
        self.write_meta(stream_id, &meta).await
    }

    #[tracing::instrument(
        skip_all,
        fields(
            db.system = "klieo-kv",
            db.namespace = "klieo-resume",
            db.operation = "replay",
            klieo.resume.stream_id = %stream_id,
            klieo.resume.since_id = since_id,
        ),
        level = "debug",
        err,
    )]
    async fn replay(
        &self,
        stream_id: &str,
        since_id: u64,
    ) -> Result<Pin<Box<dyn Stream<Item = (u64, Bytes)> + Send>>, ResumeError> {
        let meta = self.read_meta(stream_id).await?;
        // Allow since_id + 1 == oldest_id (caller asks for events
        // starting at oldest_id; we still have them). Anything older
        // is expired.
        if since_id + 1 < meta.oldest_id {
            return Err(ResumeError::Expired { since_id });
        }
        let start = since_id + 1;
        let end = meta.next_id; // exclusive

        let bucket = Self::bucket(stream_id);
        // Eagerly collect retained events so we don't need the `futures` crate.
        let mut events: Vec<(u64, Bytes)> = Vec::new();
        for id in start..end {
            let key = Self::event_key(id);
            match self.kv.get(&bucket, &key).await {
                Ok(Some(entry)) => events.push((id, entry.value)),
                // Race: event evicted between meta read and get. Skip.
                Ok(None) => {}
                Err(e) => return Err(ResumeError::Backend(Box::new(e))),
            }
        }
        Ok(Box::pin(tokio_stream::iter(events)))
    }

    #[tracing::instrument(
        skip_all,
        fields(
            db.system = "klieo-kv",
            db.namespace = "klieo-resume",
            db.operation = "close",
            klieo.resume.stream_id = %stream_id,
        ),
        level = "debug",
        err,
    )]
    async fn close(&self, stream_id: &str) -> Result<(), ResumeError> {
        let mut meta = self.read_meta(stream_id).await?;
        meta.terminal = true;
        self.write_meta(stream_id, &meta).await
    }

    #[tracing::instrument(
        skip_all,
        fields(
            db.system = "klieo-kv",
            db.namespace = "klieo-resume",
            db.operation = "is_terminal",
            klieo.resume.stream_id = %stream_id,
        ),
        level = "debug",
        err,
    )]
    async fn is_terminal(&self, stream_id: &str) -> Result<bool, ResumeError> {
        match self.read_meta(stream_id).await {
            Ok(meta) => Ok(meta.terminal),
            Err(ResumeError::NotFound(_)) => Ok(false),
            Err(e) => Err(e),
        }
    }
}

impl KvResumeBuffer {
    /// Spawn a background task that periodically walks the resume
    /// directory bucket and evicts buckets idle past `tuning.idle_ttl`
    /// or terminal-past-grace (`tuning.terminal_grace`).
    ///
    /// The returned `JoinHandle` is owned by the caller; the task
    /// terminates when the supplied `cancel` token fires. For production,
    /// wire the token to the server's parent cancel so server shutdown
    /// cleanly stops the sweeper.
    ///
    /// Backends without [`KvStore::keys`] enumeration support return
    /// `BusError::Unsupported`; the sweeper logs once and becomes inert —
    /// record/replay still work, but proactive eviction is disabled.
    pub fn spawn_sweeper(
        self: Arc<Self>,
        cancel: tokio_util::sync::CancellationToken,
    ) -> tokio::task::JoinHandle<()> {
        tokio::spawn(async move {
            let mut tick = tokio::time::interval(self.tuning.sweep_interval);
            // First tick fires immediately; skip it so the sweeper does
            // not run before any data has landed.
            tick.tick().await;
            loop {
                tokio::select! {
                    biased;
                    _ = cancel.cancelled() => break,
                    _ = tick.tick() => {
                        if let Err(e) = self.sweep_once().await {
                            tracing::warn!(
                                target: "klieo.resume.sweeper",
                                error = %e,
                                "resume buffer sweep failed; will retry next interval"
                            );
                        }
                    }
                }
            }
        })
    }

    async fn sweep_once(&self) -> Result<(), BusError> {
        let stream_ids = match self.kv.keys(RESUME_DIRECTORY_BUCKET).await {
            Ok(ids) => ids,
            Err(BusError::Unsupported(_)) => {
                tracing::debug!(
                    target: "klieo.resume.sweeper",
                    "KvStore backend does not support keys(); proactive eviction disabled"
                );
                return Ok(());
            }
            Err(e) => return Err(e),
        };

        let now_ms = Self::now_unix_ms();
        for stream_id in stream_ids {
            let meta = match self.read_meta(&stream_id).await {
                Ok(m) => m,
                // Genuinely stale entry: meta gone, directory entry leaked.
                Err(ResumeError::NotFound(_)) => {
                    let _ = self.kv.delete(RESUME_DIRECTORY_BUCKET, &stream_id).await;
                    continue;
                }
                // Backend error reading meta — DO NOT delete the directory
                // entry. A transient KV error would otherwise look like a
                // stale entry and orphan live event data with no visibility.
                Err(other) => {
                    tracing::warn!(
                        target: "klieo.resume.sweeper",
                        stream_id,
                        error = %other,
                        "read_meta failed during sweep; skipping this stream this interval"
                    );
                    continue;
                }
            };
            let age_ms = now_ms.saturating_sub(meta.last_write_unix_ms);
            let evict = if meta.terminal {
                age_ms >= self.tuning.terminal_grace.as_millis() as u64
            } else {
                age_ms >= self.tuning.idle_ttl.as_millis() as u64
            };
            if evict {
                self.evict_bucket(&stream_id).await?;
            }
        }
        Ok(())
    }

    async fn evict_bucket(&self, stream_id: &str) -> Result<(), BusError> {
        let bucket = Self::bucket(stream_id);
        let keys = match self.kv.keys(&bucket).await {
            Ok(k) => k,
            Err(BusError::Unsupported(_)) => return Ok(()),
            Err(e) => return Err(e),
        };
        let mut delete_failures: usize = 0;
        for k in keys {
            if let Err(e) = self.kv.delete(&bucket, &k).await {
                delete_failures += 1;
                tracing::warn!(
                    target: "klieo.resume.sweeper",
                    stream_id,
                    key = %k,
                    error = %e,
                    "evict delete failed; will retry next interval"
                );
            }
        }
        // Only purge the directory entry when every key delete succeeded.
        // Otherwise a partial eviction would orphan the surviving event/*
        // data with no sweeper visibility (data loss with no recovery
        // path next interval).
        if delete_failures == 0 {
            let _ = self.kv.delete(RESUME_DIRECTORY_BUCKET, stream_id).await;
        } else {
            tracing::warn!(
                target: "klieo.resume.sweeper",
                stream_id,
                failures = delete_failures,
                "directory entry retained for retry; partial eviction"
            );
        }
        Ok(())
    }
}

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

    #[tokio::test]
    async fn noop_record_is_ok() {
        let buf = NoopResumeBuffer;
        buf.record("s1", 1, Bytes::from_static(b"hi"))
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn noop_replay_returns_not_found() {
        let buf = NoopResumeBuffer;
        let result = buf.replay("s1", 0).await;
        match result {
            Err(ResumeError::NotFound(s)) => assert_eq!(s, "s1"),
            Err(other) => panic!("expected NotFound, got {other:?}"),
            Ok(_) => panic!("expected Err, got Ok"),
        }
    }

    #[tokio::test]
    async fn noop_close_is_ok() {
        let buf = NoopResumeBuffer;
        buf.close("s1").await.unwrap();
    }

    #[test]
    fn backend_error_preserves_source_chain() {
        use std::error::Error as _;
        let inner = std::io::Error::other("disk gone");
        let e = ResumeError::Backend(Box::new(inner));
        assert_eq!(e.to_string(), "buffer backend: disk gone");
        let src = e
            .source()
            .expect("Backend must preserve source for e.source()");
        assert_eq!(src.to_string(), "disk gone");
    }

    #[test]
    fn expired_renders_since_id() {
        let e = ResumeError::Expired { since_id: 42 };
        assert_eq!(
            e.to_string(),
            "resume window expired (oldest retained id > requested 42)"
        );
    }

    // ---- KvResumeBuffer tests ----

    fn in_mem_buffer() -> KvResumeBuffer {
        use crate::test_utils::fake_kv;
        KvResumeBuffer::new(fake_kv())
    }

    #[tokio::test]
    async fn kv_record_and_replay_in_order() {
        let buf = in_mem_buffer();
        for i in 1..=5u64 {
            buf.record("s1", i, Bytes::from(format!("e{i}")))
                .await
                .unwrap();
        }
        let stream = buf.replay("s1", 2).await.expect("replay");
        use tokio_stream::StreamExt as _;
        let mut stream = stream;
        let mut got = Vec::new();
        while let Some((id, bytes)) = stream.next().await {
            got.push((id, String::from_utf8(bytes.to_vec()).unwrap()));
        }
        assert_eq!(
            got,
            vec![(3, "e3".into()), (4, "e4".into()), (5, "e5".into())]
        );
    }

    #[tokio::test]
    async fn kv_ring_evicts_oldest_when_full() {
        let buf = in_mem_buffer();
        for i in 1..=300u64 {
            buf.record("s1", i, Bytes::from(format!("e{i}")))
                .await
                .unwrap();
        }
        // Cap is 256, so oldest_id == 300 - 256 + 1 == 45.
        // Replaying since 44 must succeed (next event is 45).
        let _ = buf.replay("s1", 44).await.expect("replay since 44");
        // Replaying since 5 (well before oldest=45) must return Expired.
        match buf.replay("s1", 5).await {
            Err(ResumeError::Expired { since_id }) => assert_eq!(since_id, 5),
            Err(other) => panic!("expected Expired, got Err({other:?})"),
            Ok(_) => panic!("expected Expired, got Ok(stream)"),
        }
    }

    #[tokio::test]
    async fn kv_replay_unknown_stream_returns_not_found() {
        let buf = in_mem_buffer();
        match buf.replay("missing", 0).await {
            Err(ResumeError::NotFound(s)) => assert_eq!(s, "missing"),
            Err(other) => panic!("expected NotFound, got Err({other:?})"),
            Ok(_) => panic!("expected NotFound, got Ok(stream)"),
        }
    }

    #[tokio::test]
    async fn kv_close_marks_terminal_in_meta() {
        let buf = in_mem_buffer();
        buf.record("s1", 1, Bytes::from_static(b"x")).await.unwrap();
        buf.close("s1").await.unwrap();
        let meta = buf.read_meta("s1").await.unwrap();
        assert!(meta.terminal);
    }

    #[tokio::test]
    async fn is_terminal_false_before_close() {
        let buf = in_mem_buffer();
        buf.record("s-1", 1, Bytes::from_static(b"hi"))
            .await
            .unwrap();
        assert!(!buf.is_terminal("s-1").await.unwrap());
    }

    #[tokio::test]
    async fn is_terminal_true_after_close() {
        let buf = in_mem_buffer();
        buf.record("s-1", 1, Bytes::from_static(b"hi"))
            .await
            .unwrap();
        buf.close("s-1").await.unwrap();
        assert!(buf.is_terminal("s-1").await.unwrap());
    }

    #[tokio::test]
    async fn is_terminal_false_on_unknown_id() {
        let buf = in_mem_buffer();
        assert!(!buf.is_terminal("never").await.unwrap());
    }

    // ---- Sweeper tests ----

    fn fast_sweeper_buffer() -> KvResumeBuffer {
        use crate::test_utils::fake_kv;
        KvResumeBuffer::with_tuning(
            fake_kv(),
            SweeperTuning {
                idle_ttl: Duration::from_millis(100),
                sweep_interval: Duration::from_millis(20),
                terminal_grace: Duration::from_millis(50),
            },
        )
    }

    #[tokio::test]
    async fn sweeper_evicts_idle_buckets_after_ttl() {
        let buf = Arc::new(fast_sweeper_buffer());
        buf.record("s1", 1, Bytes::from_static(b"x")).await.unwrap();
        assert!(buf.read_meta("s1").await.is_ok());

        let cancel = tokio_util::sync::CancellationToken::new();
        let sweep_handle = buf.clone().spawn_sweeper(cancel.clone());

        // idle_ttl=100ms + sweep_interval=20ms + headroom.
        tokio::time::sleep(Duration::from_millis(300)).await;

        match buf.read_meta("s1").await {
            Err(ResumeError::NotFound(_)) => {}
            Ok(_) => panic!("sweeper did not evict idle bucket"),
            Err(other) => panic!("unexpected error: {other:?}"),
        }

        cancel.cancel();
        let _ = tokio::time::timeout(Duration::from_secs(1), sweep_handle).await;
    }

    #[tokio::test]
    async fn sweeper_evicts_terminal_buckets_after_grace() {
        let buf = Arc::new(fast_sweeper_buffer());
        buf.record("s1", 1, Bytes::from_static(b"x")).await.unwrap();
        buf.close("s1").await.unwrap();

        let cancel = tokio_util::sync::CancellationToken::new();
        let sweep_handle = buf.clone().spawn_sweeper(cancel.clone());

        // terminal_grace=50ms + sweep_interval=20ms + headroom.
        tokio::time::sleep(Duration::from_millis(200)).await;

        assert!(matches!(
            buf.read_meta("s1").await,
            Err(ResumeError::NotFound(_))
        ));
        cancel.cancel();
        let _ = tokio::time::timeout(Duration::from_secs(1), sweep_handle).await;
    }
}