klieo-memory-sqlite 0.4.0

//! `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()
    }

    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, String, i64, Option<String>)> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn.prepare(
                    r#"
                    SELECT
                        run_id,
                        MIN(ts) AS started_at,
                        MAX(ts) AS last_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::<_, 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, last_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}")))?;
            let last_ts = last_ts
                .parse::<DateTime<Utc>>()
                .map_err(|e| MemoryError::Serialization(format!("last_ts: {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: Some(last_ts),
                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)
    }

    #[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::LlmCall {
                tokens: 42,
                latency_ms: 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_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"));
    }

    #[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);
            let run = run.clone();
            handles.push(tokio::spawn(async move {
                m.record(
                    run,
                    Episode::LlmCall {
                        tokens: 100,
                        latency_ms: 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");
    }
}

// `Arc` is needed only in the concurrent test above.
#[cfg(test)]
use std::sync::Arc;