klieo_memory_sqlite/

episodic.rs

//! `SqliteEpisodic` — `EpisodicMemory` over a SQLite table.

use crate::connection::DbHandle;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use klieo_core::error::MemoryError;
use klieo_core::ids::RunId;
use klieo_core::memory::{Episode, EpisodicMemory, RunFilter, RunSummary};

/// SQLite-backed episodic event log.
pub struct SqliteEpisodic {
    db: DbHandle,
}

impl SqliteEpisodic {
    pub(crate) fn new(db: DbHandle) -> Self {
        Self { db }
    }
}

fn episode_kind(ep: &Episode) -> &'static str {
    match ep {
        Episode::Started { .. } => "started",
        Episode::LlmCall { .. } => "llm_call",
        Episode::ToolCall { .. } => "tool_call",
        Episode::BusPublish { .. } => "bus_publish",
        Episode::BusReceive { .. } => "bus_receive",
        Episode::Completed => "completed",
        Episode::Failed { .. } => "failed",
        Episode::SummaryCheckpoint { .. } => "summary_checkpoint",
        // `Episode` is `#[non_exhaustive]`; future additive variants
        // get a generic kind label until explicitly handled.
        _ => "unknown",
    }
}

#[async_trait]
impl EpisodicMemory for SqliteEpisodic {
    async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError> {
        let kind = episode_kind(&event).to_string();
        let payload =
            serde_json::to_string(&event).map_err(|e| MemoryError::Serialization(e.to_string()))?;
        let now = Utc::now().to_rfc3339();
        let run_id = run.to_string();
        self.db
            .execute(move |conn| {
                // W2.A8: BEGIN IMMEDIATE so the write lock is acquired
                // before the MAX(seq) read. The default DEFERRED tx
                // promotes to a writer only at INSERT time, leaving a
                // race window where two concurrent `record()` calls
                // both read the same MAX, both try INSERT seq=N+1, and
                // the second fails on PRIMARY KEY (run_id, seq) once
                // the writer lock changes hands.
                let tx = conn.transaction_with_behavior(
                    rusqlite::TransactionBehavior::Immediate,
                )?;
                let next_seq: i64 = tx.query_row(
                    "SELECT COALESCE(MAX(seq), 0) + 1 FROM episodes WHERE run_id = ?1",
                    rusqlite::params![&run_id],
                    |r| r.get(0),
                )?;
                tx.execute(
                    "INSERT INTO episodes (run_id, seq, kind, payload, ts) VALUES (?1, ?2, ?3, ?4, ?5)",
                    rusqlite::params![&run_id, next_seq, &kind, &payload, &now],
                )?;
                tx.commit()?;
                Ok(())
            })
            .await
    }

    async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError> {
        let run_id = run.to_string();
        let payloads: Vec<String> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn
                    .prepare("SELECT payload FROM episodes WHERE run_id = ?1 ORDER BY seq ASC")?;
                let results = stmt
                    .query_map(rusqlite::params![&run_id], |r| r.get::<_, String>(0))?
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(results)
            })
            .await?;
        payloads
            .into_iter()
            .map(|p| {
                serde_json::from_str::<Episode>(&p)
                    .map_err(|e| MemoryError::Serialization(e.to_string()))
            })
            .collect()
    }

    /// Reads the `ts` column the append path has always written, so a projected
    /// `RunLog` can report a real duration instead of falling back to
    /// `Utc::now()` for both ends of the run.
    async fn replay_with_times(
        &self,
        run: RunId,
    ) -> Result<Vec<(Episode, Option<DateTime<Utc>>)>, MemoryError> {
        let run_id = run.to_string();
        let rows: Vec<(String, String)> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn.prepare(
                    "SELECT payload, ts FROM episodes WHERE run_id = ?1 ORDER BY seq ASC",
                )?;
                let results = stmt
                    .query_map(rusqlite::params![&run_id], |r| {
                        Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
                    })?
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(results)
            })
            .await?;
        rows.into_iter()
            .map(|(payload, ts)| {
                let episode = serde_json::from_str::<Episode>(&payload)
                    .map_err(|e| MemoryError::Serialization(e.to_string()))?;
                // A row whose ts cannot be parsed yields None rather than
                // failing the replay: the episode is the record of what happened,
                // the timestamp is metadata about when.
                let at = DateTime::parse_from_rfc3339(&ts)
                    .ok()
                    .map(|parsed| parsed.with_timezone(&Utc));
                Ok((episode, at))
            })
            .collect()
    }

    async fn replay_many(&self, runs: &[RunId]) -> Result<Vec<(RunId, Vec<Episode>)>, MemoryError> {
        if runs.is_empty() {
            return Ok(Vec::new());
        }
        let run_id_strs: Vec<String> = runs.iter().map(|run| run.to_string()).collect();
        let placeholders = vec!["?"; run_id_strs.len()].join(",");
        let sql = format!(
            "SELECT run_id, payload FROM episodes \
             WHERE run_id IN ({placeholders}) ORDER BY run_id ASC, seq ASC"
        );
        let bind = run_id_strs.clone();
        let rows: Vec<(String, String)> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn.prepare(&sql)?;
                let params = rusqlite::params_from_iter(bind.iter());
                let results = stmt
                    .query_map(params, |r| {
                        Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
                    })?
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(results)
            })
            .await?;

        let mut episodes_by_run: std::collections::HashMap<String, Vec<Episode>> =
            std::collections::HashMap::new();
        for (run_id_str, payload) in rows {
            let episode = serde_json::from_str::<Episode>(&payload)
                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
            episodes_by_run.entry(run_id_str).or_default().push(episode);
        }

        // Preserve the requested order; runs with no rows yield an empty Vec
        // so the caller's node set stays complete.
        let mut out = Vec::with_capacity(runs.len());
        for (run, run_id_str) in runs.iter().zip(run_id_strs) {
            let episodes = episodes_by_run.remove(&run_id_str).unwrap_or_default();
            out.push((*run, episodes));
        }
        Ok(out)
    }

    async fn list_runs(&self, filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError> {
        let RunFilter {
            agent: agent_filter,
            since,
            until,
            limit,
        } = filter;
        let limit_i64 = limit.map(|n| n as i64).unwrap_or(i64::MAX);
        let since_str = since.map(|t| t.to_rfc3339());
        let until_str = until.map(|t| t.to_rfc3339());

        // Push agent-substring filter into the HAVING clause so SQL's
        // LIMIT counts only matching rows, preserving the documented
        // "max results" semantics.
        let rows: Vec<(String, String, Option<String>, i64, Option<String>)> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn.prepare(
                    r#"
                    SELECT
                        run_id,
                        MIN(ts) AS started_at,
                        MAX(CASE WHEN kind IN ('completed', 'failed') THEN ts END) AS finished_ts,
                        COUNT(*) AS episode_count,
                        (SELECT json_extract(payload, '$.Started.agent')
                         FROM episodes e2
                         WHERE e2.run_id = episodes.run_id AND e2.kind = 'started'
                         ORDER BY seq ASC LIMIT 1) AS agent
                    FROM episodes
                    WHERE (?1 IS NULL OR ts >= ?1)
                      AND (?2 IS NULL OR ts <= ?2)
                    GROUP BY run_id
                    HAVING (?4 IS NULL OR (agent IS NOT NULL AND agent LIKE '%' || ?4 || '%'))
                    ORDER BY started_at DESC
                    LIMIT ?3
                    "#,
                )?;
                let iter = stmt.query_map(
                    rusqlite::params![&since_str, &until_str, limit_i64, &agent_filter],
                    |row| {
                        Ok((
                            row.get::<_, String>(0)?,
                            row.get::<_, String>(1)?,
                            row.get::<_, Option<String>>(2)?,
                            row.get::<_, i64>(3)?,
                            row.get::<_, Option<String>>(4)?,
                        ))
                    },
                )?;
                iter.collect::<Result<Vec<_>, _>>()
            })
            .await?;

        let mut out = Vec::new();
        for (run_id_str, started_at, finished_ts, count, agent) in rows {
            let agent_name = agent.unwrap_or_default();
            let started_at = started_at
                .parse::<DateTime<Utc>>()
                .map_err(|e| MemoryError::Serialization(format!("started_at: {e}")))?;
            // Only completed/failed runs report `finished_at` (matches the trait
            // doc + the Neo4j backend); an in-progress run leaves it `None`.
            let finished_at = finished_ts
                .map(|ts| ts.parse::<DateTime<Utc>>())
                .transpose()
                .map_err(|e| MemoryError::Serialization(format!("finished_at: {e}")))?;
            let run_id = run_id_str
                .parse::<ulid::Ulid>()
                .map(klieo_core::ids::RunId)
                .map_err(|e| MemoryError::Serialization(format!("run_id: {e}")))?;
            out.push(RunSummary {
                run_id,
                agent: agent_name,
                started_at,
                finished_at,
                episode_count: count as u32,
            });
        }
        Ok(out)
    }
}

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

    async fn fresh() -> SqliteEpisodic {
        let db = DbHandle::open(":memory:").await.unwrap();
        SqliteEpisodic::new(db)
    }

    use std::sync::Arc;

    #[tokio::test]
    async fn record_then_replay_round_trips() {
        let m = fresh().await;
        let run = RunId::new();
        m.record(
            run,
            Episode::Started {
                agent: "test".into(),
            },
        )
        .await
        .unwrap();
        m.record(run, Episode::llm_call(42, 10)).await.unwrap();
        m.record(run, Episode::Completed).await.unwrap();
        let replay = m.replay(run).await.unwrap();
        assert_eq!(replay.len(), 3);
        assert!(matches!(replay[0], Episode::Started { .. }));
        assert!(matches!(replay[1], Episode::LlmCall { tokens: 42, .. }));
        assert!(matches!(replay[2], Episode::Completed));
    }

    #[tokio::test]
    async fn replay_many_returns_episodes_for_each_run() {
        let m = fresh().await;
        let r1 = RunId::new();
        let r2 = RunId::new();
        let r3 = RunId::new();

        m.record(r1, Episode::Started { agent: "a1".into() })
            .await
            .unwrap();
        m.record(r1, Episode::Completed).await.unwrap();
        m.record(r2, Episode::Started { agent: "a2".into() })
            .await
            .unwrap();
        m.record(r2, Episode::llm_call(7, 3)).await.unwrap();
        m.record(r2, Episode::Completed).await.unwrap();
        // r3 has no episodes recorded.

        let loaded = m.replay_many(&[r1, r2, r3]).await.unwrap();
        assert_eq!(loaded.len(), 3, "one entry per requested run");

        // Requested order is preserved.
        assert_eq!(loaded[0].0, r1);
        assert_eq!(loaded[1].0, r2);
        assert_eq!(loaded[2].0, r3);

        // r1 episodes in sequence order.
        assert_eq!(loaded[0].1.len(), 2);
        assert!(matches!(loaded[0].1[0], Episode::Started { .. }));
        assert!(matches!(loaded[0].1[1], Episode::Completed));

        // r2 episodes in sequence order.
        assert_eq!(loaded[1].1.len(), 3);
        assert!(matches!(loaded[1].1[1], Episode::LlmCall { tokens: 7, .. }));

        // r3 has no episodes → empty Vec, not omitted.
        assert!(
            loaded[2].1.is_empty(),
            "run with no episodes returns empty Vec"
        );
    }

    #[tokio::test]
    async fn replay_many_with_no_runs_returns_empty() {
        let m = fresh().await;
        let loaded = m.replay_many(&[]).await.unwrap();
        assert!(loaded.is_empty());
    }

    #[tokio::test]
    async fn replay_unknown_run_returns_empty() {
        let m = fresh().await;
        let replay = m.replay(RunId::new()).await.unwrap();
        assert!(replay.is_empty());
    }

    #[tokio::test]
    async fn list_runs_returns_summary_per_run() {
        let m = fresh().await;
        let r1 = RunId::new();
        let r2 = RunId::new();
        m.record(
            r1,
            Episode::Started {
                agent: "alpha".into(),
            },
        )
        .await
        .unwrap();
        m.record(r1, Episode::Completed).await.unwrap();
        m.record(
            r2,
            Episode::Started {
                agent: "beta".into(),
            },
        )
        .await
        .unwrap();
        let summaries = m.list_runs(RunFilter::default()).await.unwrap();
        assert_eq!(summaries.len(), 2);
        let agents: Vec<_> = summaries.iter().map(|s| s.agent.as_str()).collect();
        assert!(agents.contains(&"alpha"));
        assert!(agents.contains(&"beta"));
        // alpha completed → finished_at set; beta is still in-progress → None.
        let alpha = summaries.iter().find(|s| s.agent == "alpha").unwrap();
        let beta = summaries.iter().find(|s| s.agent == "beta").unwrap();
        assert!(
            alpha.finished_at.is_some(),
            "completed run reports finished_at"
        );
        assert!(
            beta.finished_at.is_none(),
            "in-progress run leaves finished_at None"
        );
    }

    #[tokio::test]
    async fn list_runs_reports_finished_at_for_failed_run() {
        let m = fresh().await;
        let r = RunId::new();
        m.record(
            r,
            Episode::Started {
                agent: "gamma".into(),
            },
        )
        .await
        .unwrap();
        m.record(
            r,
            Episode::Failed {
                error: "boom".into(),
            },
        )
        .await
        .unwrap();
        let summaries = m.list_runs(RunFilter::default()).await.unwrap();
        let gamma = summaries.iter().find(|s| s.agent == "gamma").unwrap();
        assert!(
            gamma.finished_at.is_some(),
            "failed run also reports finished_at"
        );
    }

    #[tokio::test]
    async fn list_runs_filters_by_agent_substring() {
        let m = fresh().await;
        m.record(
            RunId::new(),
            Episode::Started {
                agent: "alpha-1".into(),
            },
        )
        .await
        .unwrap();
        m.record(
            RunId::new(),
            Episode::Started {
                agent: "beta-1".into(),
            },
        )
        .await
        .unwrap();
        let filter = RunFilter {
            agent: Some("alpha".into()),
            ..Default::default()
        };
        let summaries = m.list_runs(filter).await.unwrap();
        assert_eq!(summaries.len(), 1);
        assert_eq!(summaries[0].agent, "alpha-1");
    }

    #[tokio::test]
    async fn list_runs_respects_limit() {
        let m = fresh().await;
        for _ in 0..5 {
            m.record(RunId::new(), Episode::Started { agent: "x".into() })
                .await
                .unwrap();
        }
        let filter = RunFilter {
            limit: Some(2),
            ..Default::default()
        };
        let summaries = m.list_runs(filter).await.unwrap();
        assert_eq!(summaries.len(), 2);
    }

    #[tokio::test]
    async fn list_runs_limit_counts_post_filter() {
        let m = fresh().await;
        // 5 alpha runs + 5 beta runs.
        for _ in 0..5 {
            m.record(
                RunId::new(),
                Episode::Started {
                    agent: "alpha".into(),
                },
            )
            .await
            .unwrap();
        }
        for _ in 0..5 {
            m.record(
                RunId::new(),
                Episode::Started {
                    agent: "beta".into(),
                },
            )
            .await
            .unwrap();
        }
        let filter = RunFilter {
            agent: Some("alpha".into()),
            limit: Some(3),
            ..Default::default()
        };
        let summaries = m.list_runs(filter).await.unwrap();
        // Limit is "max post-filter", so we should get exactly 3 alpha
        // runs even though there are 5 matching.
        assert_eq!(summaries.len(), 3);
        assert!(summaries.iter().all(|s| s.agent.contains("alpha")));
    }

    /// W2.A8 / round-1 HIGH: concurrent `record()` calls into the same
    /// run must not race on the SELECT-MAX(seq)+INSERT path. Before
    /// switching to BEGIN IMMEDIATE this test was flaky — two DEFERRED
    /// transactions would both read the same MAX, both attempt INSERT
    /// seq=N+1, and one would fail on PRIMARY KEY (run_id, seq).
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn concurrent_record_into_same_run_does_not_collide() {
        let m = Arc::new(fresh().await);
        let run = RunId::new();
        let n = 100u32;
        let mut handles = Vec::with_capacity(n as usize);
        for _ in 0..n {
            let m = Arc::clone(&m);
            handles.push(tokio::spawn(async move {
                m.record(run, Episode::llm_call(100, 5)).await
            }));
        }
        for h in handles {
            h.await.expect("task did not panic").expect("record ok");
        }
        let replayed = m.replay(run).await.unwrap();
        assert_eq!(replayed.len(), n as usize, "every record must persist");
    }
}