klieo-memory-sqlite 3.16.0

//! `SqliteShortTerm` — `ShortTermMemory` over a SQLite table.

use crate::connection::DbHandle;
use async_trait::async_trait;
use chrono::Utc;
use klieo_core::error::MemoryError;
use klieo_core::ids::ThreadId;
use klieo_core::llm::{Message, Role, ToolCall};
use klieo_core::memory::ShortTermMemory;

/// Extra rows fetched beyond the token budget so the in-Rust truncation always
/// has the full kept-suffix available even when messages are tiny (~1 token).
const LOAD_ROW_CAP_BUFFER: usize = 64;

/// One stored row: its thread-local sequence number and its message columns.
type StoredRow = (i64, String, String, String, Option<String>);

/// SQLite-backed short-term conversation memory.
pub struct SqliteShortTerm {
    db: DbHandle,
}

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

    /// The newest rows that could survive the budget, plus the thread's head
    /// row, in thread order.
    ///
    /// The head row is fetched in the SAME statement. `ORDER BY seq DESC LIMIT
    /// n` returns the NEWEST rows, so on a thread longer than the cap the first
    /// message is not in the result set at all — and no budget walk over that
    /// set can honour `most_recent_within_budget`'s guarantee that the first
    /// message survives. One statement rather than two because a second query
    /// would race an interleaved `append`.
    async fn newest_window_with_head(
        &self,
        thread: &ThreadId,
        max_tokens: usize,
    ) -> Result<Vec<StoredRow>, MemoryError> {
        // Cap the row scan so a long-lived thread doesn't read every message
        // into memory before the token-budget truncation runs. Each message
        // costs >= 1 token, so at most `max_tokens` rows can survive the
        // budget; a small buffer keeps the cap safe.
        let row_cap =
            i64::try_from(max_tokens.saturating_add(LOAD_ROW_CAP_BUFFER)).unwrap_or(i64::MAX);
        let thread_id = thread.0.clone();
        let mut rows: Vec<StoredRow> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn.prepare(
                    "SELECT * FROM (\
                       SELECT seq, role, content, tool_calls, tool_call_id \
                       FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq DESC LIMIT ?2\
                     ) \
                     UNION \
                     SELECT * FROM (\
                       SELECT seq, role, content, tool_calls, tool_call_id \
                       FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq ASC LIMIT 1\
                     )",
                )?;
                let iter = stmt.query_map(rusqlite::params![&thread_id, row_cap], |row| {
                    Ok((
                        row.get::<_, i64>(0)?,
                        row.get::<_, String>(1)?,
                        row.get::<_, String>(2)?,
                        row.get::<_, String>(3)?,
                        row.get::<_, Option<String>>(4)?,
                    ))
                })?;
                iter.collect::<Result<Vec<_>, _>>()
            })
            .await?;
        // UNION makes no ordering promise; `seq` is the thread's own order.
        rows.sort_by_key(|(seq, ..)| *seq);
        Ok(rows)
    }
}

/// How many rows the row cap left out between the head row and the newest
/// window. `seq` is dense per thread (`MAX(seq) + 1` on append, and only
/// `clear` removes rows), so a jump in it is exactly the rows not fetched.
///
/// Without this the gap is invisible: the walk downstream sees head and window
/// side by side and reads them as one continuous conversation.
fn elided_between_rows(rows: &[StoredRow]) -> usize {
    rows.windows(2)
        .map(|pair| {
            let (previous, next) = (pair[0].0, pair[1].0);
            usize::try_from((next - previous - 1).max(0)).unwrap_or(0)
        })
        .sum()
}

fn to_messages(rows: Vec<StoredRow>) -> Result<Vec<Message>, MemoryError> {
    rows.into_iter()
        .map(|(_, role, content, tool_calls_json, tool_call_id)| {
            let tool_calls: Vec<ToolCall> = serde_json::from_str(&tool_calls_json)
                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
            Ok(Message {
                role: role_from_str(&role)?,
                content,
                tool_calls,
                tool_call_id,
            })
        })
        .collect()
}

/// Marks a gap the row cap opened, when the budget walk has not already marked
/// one. Two markers for one hole would misreport how much is missing.
fn mark_row_cap_gap(
    budgeted: klieo_core::memory::BudgetedHistory,
    elided_by_row_cap: usize,
) -> Vec<Message> {
    let mut kept = budgeted.kept;
    if elided_by_row_cap == 0 || budgeted.dropped > 0 || kept.len() < 2 {
        return kept;
    }
    kept.insert(
        1,
        Message::system(format!(
            "{}{elided_by_row_cap} earlier message(s) omitted to fit max_history_tokens; the \
             first message and the most recent turns are intact]",
            klieo_core::memory::ELIDED_MESSAGES_MARKER_PREFIX
        )),
    );
    kept
}

fn role_to_str(r: Role) -> &'static str {
    match r {
        Role::System => "system",
        Role::User => "user",
        Role::Assistant => "assistant",
        Role::Tool => "tool",
        _ => "user",
    }
}

fn role_from_str(s: &str) -> Result<Role, MemoryError> {
    match s {
        "system" => Ok(Role::System),
        "user" => Ok(Role::User),
        "assistant" => Ok(Role::Assistant),
        "tool" => Ok(Role::Tool),
        other => Err(MemoryError::Serialization(format!("unknown role: {other}"))),
    }
}

#[async_trait]
impl ShortTermMemory for SqliteShortTerm {
    async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError> {
        let tool_calls_json = serde_json::to_string(&msg.tool_calls)
            .map_err(|e| MemoryError::Serialization(e.to_string()))?;
        let role = role_to_str(msg.role);
        let now = Utc::now().to_rfc3339();
        self.db
            .execute(move |conn| {
                let tx = conn.transaction()?;
                let next_seq: i64 = tx
                    .query_row(
                        "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
                        rusqlite::params![&thread.0],
                        |r| r.get(0),
                    )?;
                tx.execute(
                    "INSERT INTO short_term_messages (thread_id, seq, role, content, tool_calls, tool_call_id, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
                    rusqlite::params![
                        &thread.0,
                        next_seq,
                        role,
                        &msg.content,
                        &tool_calls_json,
                        &msg.tool_call_id,
                        &now,
                    ],
                )?;
                tx.commit()?;
                Ok(())
            })
            .await
    }

    async fn append_batch(
        &self,
        thread: ThreadId,
        messages: Vec<Message>,
    ) -> Result<(), MemoryError> {
        if messages.is_empty() {
            return Ok(());
        }
        // Serialize before opening the transaction so a bad message can't abort
        // a half-written batch. The whole batch then lands in one transaction
        // with a single MAX(seq) probe.
        let mut rows: Vec<(&'static str, String, String, Option<String>)> =
            Vec::with_capacity(messages.len());
        for msg in &messages {
            let tool_calls_json = serde_json::to_string(&msg.tool_calls)
                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
            rows.push((
                role_to_str(msg.role),
                msg.content.clone(),
                tool_calls_json,
                msg.tool_call_id.clone(),
            ));
        }
        let now = Utc::now().to_rfc3339();
        self.db
            .execute(move |conn| {
                let tx = conn.transaction()?;
                let start_seq: i64 = tx.query_row(
                    "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
                    rusqlite::params![&thread.0],
                    |r| r.get(0),
                )?;
                for (seq, (role, content, tool_calls_json, tool_call_id)) in (start_seq..).zip(&rows) {
                    tx.execute(
                        "INSERT INTO short_term_messages (thread_id, seq, role, content, tool_calls, tool_call_id, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
                        rusqlite::params![
                            &thread.0, seq, role, content, tool_calls_json, tool_call_id, &now,
                        ],
                    )?;
                }
                tx.commit()?;
                Ok(())
            })
            .await
    }

    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError> {
        let rows = self.newest_window_with_head(&thread, max_tokens).await?;
        let elided_by_row_cap = elided_between_rows(&rows);
        let messages = to_messages(rows)?;

        // The budget walk itself lives in `klieo-core` and is shared with every
        // other backend. It used to be re-implemented here, which is why the
        // same bug -- charging UTF-8 bytes instead of Unicode scalar values,
        // so CJK, emoji and accented text cost roughly 3x and this store kept
        // a third as much multibyte history as the others at the same budget
        // -- had to be found and fixed separately in each copy. One walk means
        // one place to be right, and it is the walk the conformance suite pins.
        let budgeted =
            klieo_core::memory::most_recent_within_budget_reporting(messages, max_tokens);
        if budgeted.dropped > 0 || elided_by_row_cap > 0 {
            tracing::warn!(
                thread = %thread.0,
                dropped_by_budget = budgeted.dropped,
                elided_by_row_cap,
                max_tokens,
                "short-term load could not return the whole thread; the first message and the \
                 newest turns were kept and the gap is marked in the history"
            );
        }
        Ok(mark_row_cap_gap(budgeted, elided_by_row_cap))
    }

    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError> {
        self.db
            .execute(move |conn| {
                conn.execute(
                    "DELETE FROM short_term_messages WHERE thread_id = ?1",
                    rusqlite::params![&thread.0],
                )?;
                Ok(())
            })
            .await
    }
}

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

    fn user(text: &str) -> Message {
        Message {
            role: Role::User,
            content: text.into(),
            tool_calls: vec![],
            tool_call_id: None,
        }
    }

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

    /// Holds this backend to the same contract as every other
    /// `ShortTermMemory`, rather than only to the behaviours its own tests
    /// happened to assert.
    #[tokio::test]
    async fn satisfies_short_term_conformance() {
        klieo_core::conformance::short_term_memory(&fresh().await).await;
    }

    #[tokio::test]
    async fn append_then_load_round_trips() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append(t.clone(), user("hello")).await.unwrap();
        m.append(t.clone(), user("world")).await.unwrap();
        let loaded = m.load(t, 10_000).await.unwrap();
        assert_eq!(loaded.len(), 2);
        assert_eq!(loaded[0].content, "hello");
        assert_eq!(loaded[1].content, "world");
    }

    #[tokio::test]
    async fn append_batch_preserves_order_and_continues_seq() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        // A prior single append, then a batch — the batch must continue the
        // sequence after it, in order.
        m.append(t.clone(), user("first")).await.unwrap();
        m.append_batch(t.clone(), vec![user("second"), user("third")])
            .await
            .unwrap();
        let loaded = m.load(t, 10_000).await.unwrap();
        let contents: Vec<&str> = loaded.iter().map(|msg| msg.content.as_str()).collect();
        assert_eq!(contents, vec!["first", "second", "third"]);
    }

    #[tokio::test]
    async fn append_batch_with_empty_input_is_a_noop() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append_batch(t.clone(), vec![]).await.unwrap();
        assert!(m.load(t, 10_000).await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn load_caps_rows_at_token_budget_plus_buffer() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        // 200 single-char messages (~1 token each); a budget of 5 must keep
        // only the newest few, never load all 200.
        let batch: Vec<Message> = (0..200).map(|i| user(&i.to_string())).collect();
        m.append_batch(t.clone(), batch).await.unwrap();
        let loaded = m.load(t, 5).await.unwrap();
        assert!(
            loaded.len() <= 5 + LOAD_ROW_CAP_BUFFER,
            "load must bound the row scan; got {}",
            loaded.len()
        );
        // Far below the 200 written — proves the LIMIT actually fired.
        assert!(loaded.len() < 100, "cap must fire; got {}", loaded.len());
        // The kept suffix is the newest messages, in order.
        assert_eq!(loaded.last().unwrap().content, "199");
    }

    #[tokio::test]
    async fn load_truncates_to_token_budget() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        // Each ~40-char message ~= 10 tokens.
        for i in 0..20 {
            m.append(
                t.clone(),
                user(&format!("msg-{i:03}-padding-padding-padding")),
            )
            .await
            .unwrap();
        }
        let loaded = m.load(t, 30).await.unwrap();
        // A 30-token budget buys ~3 of these, and both ends are kept on top of
        // that: head, gap marker, then the newest turns.
        assert!(
            loaded.len() <= 6 && !loaded.is_empty(),
            "expected the budget to bound the middle, got {}",
            loaded.len()
        );
        assert!(
            loaded.first().unwrap().content.contains("msg-000"),
            "the first message must survive truncation"
        );
        assert!(
            loaded.last().unwrap().content.contains("msg-019"),
            "newest message must survive truncation"
        );
        assert!(
            loaded[1]
                .content
                .starts_with(klieo_core::memory::ELIDED_MESSAGES_MARKER_PREFIX),
            "the omitted span must be marked, not closed silently: {:?}",
            loaded[1].content
        );
    }

    /// The defect that sent a factory Coder a system prompt and no task: a
    /// brief larger than the whole budget loaded as an empty history, so the
    /// agent had nothing to work from and invented its own task.
    #[tokio::test]
    async fn a_brief_larger_than_the_budget_still_loads() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append(t.clone(), user(&"x".repeat(40_000)))
            .await
            .unwrap();

        let loaded = m.load(t, 8_000).await.unwrap();

        assert_eq!(
            loaded.len(),
            1,
            "a non-empty thread must never load as an empty history"
        );
    }

    /// The row cap fetches the NEWEST rows, so the head row has to be fetched
    /// explicitly or the first message is not even a candidate.
    #[tokio::test]
    async fn the_first_message_survives_a_thread_longer_than_the_row_cap() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append(t.clone(), user("the brief")).await.unwrap();
        for i in 0..200 {
            m.append(t.clone(), user(&format!("turn-{i}")))
                .await
                .unwrap();
        }

        let loaded = m.load(t, 10).await.unwrap();

        assert_eq!(
            loaded.first().unwrap().content,
            "the brief",
            "the head row must be fetched even when the row cap excludes it"
        );
        assert!(
            loaded.last().unwrap().content.contains("turn-199"),
            "and the newest turn must still be there"
        );
    }

    #[tokio::test]
    async fn clear_removes_thread() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append(t.clone(), user("hello")).await.unwrap();
        m.clear(t.clone()).await.unwrap();
        let loaded = m.load(t, 10_000).await.unwrap();
        assert!(loaded.is_empty());
    }

    #[tokio::test]
    async fn threads_are_isolated() {
        let m = fresh().await;
        m.append(ThreadId::new("a"), user("a-msg")).await.unwrap();
        m.append(ThreadId::new("b"), user("b-msg")).await.unwrap();
        let a = m.load(ThreadId::new("a"), 10_000).await.unwrap();
        let b = m.load(ThreadId::new("b"), 10_000).await.unwrap();
        assert_eq!(a.len(), 1);
        assert_eq!(b.len(), 1);
        assert_eq!(a[0].content, "a-msg");
        assert_eq!(b[0].content, "b-msg");
    }

    #[tokio::test]
    async fn role_round_trip_for_all_variants() {
        let m = fresh().await;
        let t = ThreadId::new("r");
        for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
            m.append(
                t.clone(),
                Message {
                    role,
                    content: "x".into(),
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await
            .unwrap();
        }
        let loaded = m.load(t, 10_000).await.unwrap();
        assert_eq!(loaded.len(), 4);
        assert_eq!(loaded[0].role, Role::System);
        assert_eq!(loaded[3].role, Role::Tool);
    }
}