klieo-graph-projector 3.4.0

//! [`EpisodicMemory`] wrapper that broadcasts every recorded
//! `(RunId, Episode)` pair on a `tokio::sync::broadcast` channel.
//!
//! The wrapped store remains the single source of truth: persistence
//! happens **first** via the inner `record()` call. Only on success is
//! the pair sent on the broadcast channel. A failed `inner.record()`
//! never produces a broadcast — there is no event to project.
//!
//! Conversely a broadcast `send()` failure (no subscribers, lagging
//! receiver) never propagates to the caller. `record()` succeeds as
//! soon as the inner store has accepted the episode; the projection
//! pipeline is best-effort and its lag is observable via
//! [`klieo_memory_graph::RecallMetrics::record_projector_lag`] from
//! the subscriber side.

use async_trait::async_trait;
use klieo_core::error::MemoryError;
use klieo_core::ids::RunId;
use klieo_core::memory::{Episode, EpisodicMemory, RunFilter, RunSummary};
use std::sync::Arc;
use tokio::sync::broadcast;

/// Capacity of the internal broadcast channel. Sized for typical
/// agent-loop emission cadence (a handful of episodes per second)
/// with margin for a single slow subscriber to catch up. Subscribers
/// that fall further behind receive `RecvError::Lagged(n)` and the
/// subscriber implementation is expected to log + record the lag
/// rather than block the producer.
const CHANNEL_CAPACITY: usize = 256;

/// Drop-in wrapper around an [`EpisodicMemory`] that publishes each
/// recorded `(RunId, Episode)` pair on a `tokio::sync::broadcast`
/// channel.
///
/// Construction returns the wrapper together with the **first**
/// receiver. Additional subscribers can be obtained via
/// [`Self::subscribe`]. Each receiver gets a fan-out copy of every
/// subsequent record; subscribers added after a record was published
/// will not see that record (broadcast `subscribe` semantics).
#[non_exhaustive]
pub struct ProjectableEpisodic {
    inner: Arc<dyn EpisodicMemory>,
    tx: broadcast::Sender<(RunId, Episode)>,
}

impl ProjectableEpisodic {
    /// Wrap `inner` and return the wrapper plus the first subscriber
    /// receiver. Spawn an [`crate::EpisodeProjector`] against the
    /// returned receiver to start indexing into the graph.
    pub fn new(inner: Arc<dyn EpisodicMemory>) -> (Self, broadcast::Receiver<(RunId, Episode)>) {
        let (tx, rx) = broadcast::channel(CHANNEL_CAPACITY);
        (Self { inner, tx }, rx)
    }

    /// Open an additional subscription on the broadcast channel.
    ///
    /// Subscribers only receive episodes recorded after the
    /// `subscribe()` call. Use [`Self::new`] when you need to capture
    /// every record from the very first one.
    pub fn subscribe(&self) -> broadcast::Receiver<(RunId, Episode)> {
        self.tx.subscribe()
    }
}

#[async_trait]
impl EpisodicMemory for ProjectableEpisodic {
    async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError> {
        self.inner.record(run, event.clone()).await?;
        // best-effort fanout: send only fails when there are no live
        // subscribers (see module doc); never propagate to the caller.
        let _ = self.tx.send((run, event));
        Ok(())
    }

    async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError> {
        self.inner.replay(run).await
    }

    async fn list_runs(&self, filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError> {
        self.inner.list_runs(filter).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use klieo_core::test_utils::InMemoryEpisodic;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use tokio::sync::broadcast::error::TryRecvError;

    /// Stub episodic that always fails `record()` so we can prove the
    /// broadcast side-channel never fires on persistence failure.
    struct FailingEpisodic {
        record_calls: AtomicUsize,
    }

    impl FailingEpisodic {
        fn new() -> Self {
            Self {
                record_calls: AtomicUsize::new(0),
            }
        }
    }

    #[async_trait]
    impl EpisodicMemory for FailingEpisodic {
        async fn record(&self, _run: RunId, _event: Episode) -> Result<(), MemoryError> {
            self.record_calls.fetch_add(1, Ordering::SeqCst);
            Err(MemoryError::Store("simulated persistence failure".into()))
        }

        async fn replay(&self, _run: RunId) -> Result<Vec<Episode>, MemoryError> {
            Ok(Vec::new())
        }

        async fn list_runs(&self, _filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError> {
            Ok(Vec::new())
        }
    }

    #[tokio::test]
    async fn record_delegates_to_inner() {
        let inner: Arc<dyn EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
        let (wrapper, _rx) = ProjectableEpisodic::new(Arc::clone(&inner));
        let run = RunId::new();

        wrapper
            .record(
                run,
                Episode::Started {
                    agent: "alpha".into(),
                },
            )
            .await
            .expect("record succeeds");

        let stored = inner.replay(run).await.expect("replay");
        assert_eq!(stored.len(), 1, "inner should have one recorded episode");
        assert!(matches!(&stored[0], Episode::Started { agent } if agent == "alpha"));
    }

    #[tokio::test]
    async fn broadcasts_each_record() {
        let inner: Arc<dyn EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
        let (wrapper, mut rx) = ProjectableEpisodic::new(inner);
        let run = RunId::new();

        wrapper
            .record(
                run,
                Episode::Started {
                    agent: "alpha".into(),
                },
            )
            .await
            .unwrap();
        wrapper.record(run, Episode::Completed).await.unwrap();

        let (got_run, got_ep) = rx.recv().await.expect("first broadcast");
        assert_eq!(got_run, run);
        assert!(matches!(got_ep, Episode::Started { agent } if agent == "alpha"));

        let (got_run, got_ep) = rx.recv().await.expect("second broadcast");
        assert_eq!(got_run, run);
        assert!(matches!(got_ep, Episode::Completed));
    }

    #[tokio::test]
    async fn record_inner_failure_does_not_broadcast() {
        let failing = Arc::new(FailingEpisodic::new());
        let (wrapper, mut rx) =
            ProjectableEpisodic::new(Arc::clone(&failing) as Arc<dyn EpisodicMemory>);
        let run = RunId::new();

        let err = wrapper
            .record(
                run,
                Episode::Started {
                    agent: "alpha".into(),
                },
            )
            .await
            .expect_err("inner failure must propagate");
        assert!(matches!(err, MemoryError::Store(_)));
        assert_eq!(
            failing.record_calls.load(Ordering::SeqCst),
            1,
            "inner must be called exactly once"
        );

        // No broadcast on failure path.
        match rx.try_recv() {
            Err(TryRecvError::Empty) => {}
            other => panic!("expected Empty channel after failed record, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn replay_delegates_to_inner() {
        let inner = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        inner
            .record(
                run,
                Episode::Started {
                    agent: "alpha".into(),
                },
            )
            .await
            .unwrap();
        let (wrapper, _rx) =
            ProjectableEpisodic::new(Arc::clone(&inner) as Arc<dyn EpisodicMemory>);

        let episodes = wrapper.replay(run).await.expect("replay through wrapper");
        assert_eq!(episodes.len(), 1);
        assert!(matches!(&episodes[0], Episode::Started { agent } if agent == "alpha"));
    }

    #[tokio::test]
    async fn replay_propagates_inner_error() {
        struct FailingReplay;
        #[async_trait]
        impl EpisodicMemory for FailingReplay {
            async fn record(&self, _run: RunId, _event: Episode) -> Result<(), MemoryError> {
                Ok(())
            }
            async fn replay(&self, _run: RunId) -> Result<Vec<Episode>, MemoryError> {
                Err(MemoryError::Store("simulated replay failure".into()))
            }
            async fn list_runs(&self, _filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError> {
                Ok(Vec::new())
            }
        }
        let (wrapper, _rx) =
            ProjectableEpisodic::new(Arc::new(FailingReplay) as Arc<dyn EpisodicMemory>);
        let err = wrapper
            .replay(RunId::new())
            .await
            .expect_err("inner replay error must propagate");
        assert!(matches!(err, MemoryError::Store(s) if s.contains("simulated replay failure")));
    }

    #[tokio::test]
    async fn list_runs_delegates_to_inner() {
        let inner = Arc::new(InMemoryEpisodic::default());
        let run = RunId::new();
        inner.record(run, Episode::Completed).await.unwrap();
        let (wrapper, _rx) =
            ProjectableEpisodic::new(Arc::clone(&inner) as Arc<dyn EpisodicMemory>);

        let runs = wrapper
            .list_runs(RunFilter::default())
            .await
            .expect("list_runs through wrapper");
        assert_eq!(runs.len(), 1);
        assert_eq!(runs[0].run_id, run);
    }

    #[tokio::test]
    async fn subscribe_receives_events_published_after_subscription() {
        let inner: Arc<dyn EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
        let (wrapper, _rx_initial) = ProjectableEpisodic::new(inner);
        let mut rx_second = wrapper.subscribe();
        let run = RunId::new();

        wrapper
            .record(
                run,
                Episode::Started {
                    agent: "alpha".into(),
                },
            )
            .await
            .unwrap();

        let (got_run, got_ep) = rx_second.recv().await.expect("second subscriber receives");
        assert_eq!(got_run, run);
        assert!(matches!(got_ep, Episode::Started { agent } if agent == "alpha"));
    }
}