klieo_memory_sqlite/
short_term.rs

1//! `SqliteShortTerm` — `ShortTermMemory` over a SQLite table.
2
3use crate::connection::DbHandle;
4use async_trait::async_trait;
5use chrono::Utc;
6use klieo_core::error::MemoryError;
7use klieo_core::ids::ThreadId;
8use klieo_core::llm::{Message, Role, ToolCall};
9use klieo_core::memory::ShortTermMemory;
10
11/// Extra rows fetched beyond the token budget so the in-Rust truncation always
12/// has the full kept-suffix available even when messages are tiny (~1 token).
13const LOAD_ROW_CAP_BUFFER: usize = 64;
14
15/// One stored row: its thread-local sequence number and its message columns.
16type StoredRow = (i64, String, String, String, Option<String>);
17
18/// SQLite-backed short-term conversation memory.
19pub struct SqliteShortTerm {
20    db: DbHandle,
21}
22
23impl SqliteShortTerm {
24    pub(crate) fn new(db: DbHandle) -> Self {
25        Self { db }
26    }
27
28    /// The newest rows that could survive the budget, plus the thread's head
29    /// row, in thread order.
30    ///
31    /// The head row is fetched in the SAME statement. `ORDER BY seq DESC LIMIT
32    /// n` returns the NEWEST rows, so on a thread longer than the cap the first
33    /// message is not in the result set at all — and no budget walk over that
34    /// set can honour `most_recent_within_budget`'s guarantee that the first
35    /// message survives. One statement rather than two because a second query
36    /// would race an interleaved `append`.
37    async fn newest_window_with_head(
38        &self,
39        thread: &ThreadId,
40        max_tokens: usize,
41    ) -> Result<Vec<StoredRow>, MemoryError> {
42        // Cap the row scan so a long-lived thread doesn't read every message
43        // into memory before the token-budget truncation runs. Each message
44        // costs >= 1 token, so at most `max_tokens` rows can survive the
45        // budget; a small buffer keeps the cap safe.
46        let row_cap =
47            i64::try_from(max_tokens.saturating_add(LOAD_ROW_CAP_BUFFER)).unwrap_or(i64::MAX);
48        let thread_id = thread.0.clone();
49        let mut rows: Vec<StoredRow> = self
50            .db
51            .execute(move |conn| {
52                let mut stmt = conn.prepare(
53                    "SELECT * FROM (\
54                       SELECT seq, role, content, tool_calls, tool_call_id \
55                       FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq DESC LIMIT ?2\
56                     ) \
57                     UNION \
58                     SELECT * FROM (\
59                       SELECT seq, role, content, tool_calls, tool_call_id \
60                       FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq ASC LIMIT 1\
61                     )",
62                )?;
63                let iter = stmt.query_map(rusqlite::params![&thread_id, row_cap], |row| {
64                    Ok((
65                        row.get::<_, i64>(0)?,
66                        row.get::<_, String>(1)?,
67                        row.get::<_, String>(2)?,
68                        row.get::<_, String>(3)?,
69                        row.get::<_, Option<String>>(4)?,
70                    ))
71                })?;
72                iter.collect::<Result<Vec<_>, _>>()
73            })
74            .await?;
75        // UNION makes no ordering promise; `seq` is the thread's own order.
76        rows.sort_by_key(|(seq, ..)| *seq);
77        Ok(rows)
78    }
79}
80
81/// How many rows the row cap left out between the head row and the newest
82/// window. `seq` is dense per thread (`MAX(seq) + 1` on append, and only
83/// `clear` removes rows), so a jump in it is exactly the rows not fetched.
84///
85/// Without this the gap is invisible: the walk downstream sees head and window
86/// side by side and reads them as one continuous conversation.
87fn elided_between_rows(rows: &[StoredRow]) -> usize {
88    rows.windows(2)
89        .map(|pair| {
90            let (previous, next) = (pair[0].0, pair[1].0);
91            usize::try_from((next - previous - 1).max(0)).unwrap_or(0)
92        })
93        .sum()
94}
95
96fn to_messages(rows: Vec<StoredRow>) -> Result<Vec<Message>, MemoryError> {
97    rows.into_iter()
98        .map(|(_, role, content, tool_calls_json, tool_call_id)| {
99            let tool_calls: Vec<ToolCall> = serde_json::from_str(&tool_calls_json)
100                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
101            Ok(Message {
102                role: role_from_str(&role)?,
103                content,
104                tool_calls,
105                tool_call_id,
106            })
107        })
108        .collect()
109}
110
111/// Marks a gap the row cap opened, when the budget walk has not already marked
112/// one. Two markers for one hole would misreport how much is missing.
113fn mark_row_cap_gap(
114    budgeted: klieo_core::memory::BudgetedHistory,
115    elided_by_row_cap: usize,
116) -> Vec<Message> {
117    let mut kept = budgeted.kept;
118    if elided_by_row_cap == 0 || budgeted.dropped > 0 || kept.len() < 2 {
119        return kept;
120    }
121    kept.insert(
122        1,
123        Message::system(format!(
124            "{}{elided_by_row_cap} earlier message(s) omitted to fit max_history_tokens; the \
125             first message and the most recent turns are intact]",
126            klieo_core::memory::ELIDED_MESSAGES_MARKER_PREFIX
127        )),
128    );
129    kept
130}
131
132fn role_to_str(r: Role) -> &'static str {
133    match r {
134        Role::System => "system",
135        Role::User => "user",
136        Role::Assistant => "assistant",
137        Role::Tool => "tool",
138        _ => "user",
139    }
140}
141
142fn role_from_str(s: &str) -> Result<Role, MemoryError> {
143    match s {
144        "system" => Ok(Role::System),
145        "user" => Ok(Role::User),
146        "assistant" => Ok(Role::Assistant),
147        "tool" => Ok(Role::Tool),
148        other => Err(MemoryError::Serialization(format!("unknown role: {other}"))),
149    }
150}
151
152#[async_trait]
153impl ShortTermMemory for SqliteShortTerm {
154    async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError> {
155        let tool_calls_json = serde_json::to_string(&msg.tool_calls)
156            .map_err(|e| MemoryError::Serialization(e.to_string()))?;
157        let role = role_to_str(msg.role);
158        let now = Utc::now().to_rfc3339();
159        self.db
160            .execute(move |conn| {
161                let tx = conn.transaction()?;
162                let next_seq: i64 = tx
163                    .query_row(
164                        "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
165                        rusqlite::params![&thread.0],
166                        |r| r.get(0),
167                    )?;
168                tx.execute(
169                    "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)",
170                    rusqlite::params![
171                        &thread.0,
172                        next_seq,
173                        role,
174                        &msg.content,
175                        &tool_calls_json,
176                        &msg.tool_call_id,
177                        &now,
178                    ],
179                )?;
180                tx.commit()?;
181                Ok(())
182            })
183            .await
184    }
185
186    async fn append_batch(
187        &self,
188        thread: ThreadId,
189        messages: Vec<Message>,
190    ) -> Result<(), MemoryError> {
191        if messages.is_empty() {
192            return Ok(());
193        }
194        // Serialize before opening the transaction so a bad message can't abort
195        // a half-written batch. The whole batch then lands in one transaction
196        // with a single MAX(seq) probe.
197        let mut rows: Vec<(&'static str, String, String, Option<String>)> =
198            Vec::with_capacity(messages.len());
199        for msg in &messages {
200            let tool_calls_json = serde_json::to_string(&msg.tool_calls)
201                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
202            rows.push((
203                role_to_str(msg.role),
204                msg.content.clone(),
205                tool_calls_json,
206                msg.tool_call_id.clone(),
207            ));
208        }
209        let now = Utc::now().to_rfc3339();
210        self.db
211            .execute(move |conn| {
212                let tx = conn.transaction()?;
213                let start_seq: i64 = tx.query_row(
214                    "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
215                    rusqlite::params![&thread.0],
216                    |r| r.get(0),
217                )?;
218                for (seq, (role, content, tool_calls_json, tool_call_id)) in (start_seq..).zip(&rows) {
219                    tx.execute(
220                        "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)",
221                        rusqlite::params![
222                            &thread.0, seq, role, content, tool_calls_json, tool_call_id, &now,
223                        ],
224                    )?;
225                }
226                tx.commit()?;
227                Ok(())
228            })
229            .await
230    }
231
232    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError> {
233        let rows = self.newest_window_with_head(&thread, max_tokens).await?;
234        let elided_by_row_cap = elided_between_rows(&rows);
235        let messages = to_messages(rows)?;
236
237        // The budget walk itself lives in `klieo-core` and is shared with every
238        // other backend. It used to be re-implemented here, which is why the
239        // same bug -- charging UTF-8 bytes instead of Unicode scalar values,
240        // so CJK, emoji and accented text cost roughly 3x and this store kept
241        // a third as much multibyte history as the others at the same budget
242        // -- had to be found and fixed separately in each copy. One walk means
243        // one place to be right, and it is the walk the conformance suite pins.
244        let budgeted =
245            klieo_core::memory::most_recent_within_budget_reporting(messages, max_tokens);
246        if budgeted.dropped > 0 || elided_by_row_cap > 0 {
247            tracing::warn!(
248                thread = %thread.0,
249                dropped_by_budget = budgeted.dropped,
250                elided_by_row_cap,
251                max_tokens,
252                "short-term load could not return the whole thread; the first message and the \
253                 newest turns were kept and the gap is marked in the history"
254            );
255        }
256        Ok(mark_row_cap_gap(budgeted, elided_by_row_cap))
257    }
258
259    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError> {
260        self.db
261            .execute(move |conn| {
262                conn.execute(
263                    "DELETE FROM short_term_messages WHERE thread_id = ?1",
264                    rusqlite::params![&thread.0],
265                )?;
266                Ok(())
267            })
268            .await
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use crate::connection::DbHandle;
276
277    fn user(text: &str) -> Message {
278        Message {
279            role: Role::User,
280            content: text.into(),
281            tool_calls: vec![],
282            tool_call_id: None,
283        }
284    }
285
286    async fn fresh() -> SqliteShortTerm {
287        let db = DbHandle::open(":memory:").await.unwrap();
288        SqliteShortTerm::new(db)
289    }
290
291    /// Holds this backend to the same contract as every other
292    /// `ShortTermMemory`, rather than only to the behaviours its own tests
293    /// happened to assert.
294    #[tokio::test]
295    async fn satisfies_short_term_conformance() {
296        klieo_core::conformance::short_term_memory(&fresh().await).await;
297    }
298
299    #[tokio::test]
300    async fn append_then_load_round_trips() {
301        let m = fresh().await;
302        let t = ThreadId::new("t1");
303        m.append(t.clone(), user("hello")).await.unwrap();
304        m.append(t.clone(), user("world")).await.unwrap();
305        let loaded = m.load(t, 10_000).await.unwrap();
306        assert_eq!(loaded.len(), 2);
307        assert_eq!(loaded[0].content, "hello");
308        assert_eq!(loaded[1].content, "world");
309    }
310
311    #[tokio::test]
312    async fn append_batch_preserves_order_and_continues_seq() {
313        let m = fresh().await;
314        let t = ThreadId::new("t1");
315        // A prior single append, then a batch — the batch must continue the
316        // sequence after it, in order.
317        m.append(t.clone(), user("first")).await.unwrap();
318        m.append_batch(t.clone(), vec![user("second"), user("third")])
319            .await
320            .unwrap();
321        let loaded = m.load(t, 10_000).await.unwrap();
322        let contents: Vec<&str> = loaded.iter().map(|msg| msg.content.as_str()).collect();
323        assert_eq!(contents, vec!["first", "second", "third"]);
324    }
325
326    #[tokio::test]
327    async fn append_batch_with_empty_input_is_a_noop() {
328        let m = fresh().await;
329        let t = ThreadId::new("t1");
330        m.append_batch(t.clone(), vec![]).await.unwrap();
331        assert!(m.load(t, 10_000).await.unwrap().is_empty());
332    }
333
334    #[tokio::test]
335    async fn load_caps_rows_at_token_budget_plus_buffer() {
336        let m = fresh().await;
337        let t = ThreadId::new("t1");
338        // 200 single-char messages (~1 token each); a budget of 5 must keep
339        // only the newest few, never load all 200.
340        let batch: Vec<Message> = (0..200).map(|i| user(&i.to_string())).collect();
341        m.append_batch(t.clone(), batch).await.unwrap();
342        let loaded = m.load(t, 5).await.unwrap();
343        assert!(
344            loaded.len() <= 5 + LOAD_ROW_CAP_BUFFER,
345            "load must bound the row scan; got {}",
346            loaded.len()
347        );
348        // Far below the 200 written — proves the LIMIT actually fired.
349        assert!(loaded.len() < 100, "cap must fire; got {}", loaded.len());
350        // The kept suffix is the newest messages, in order.
351        assert_eq!(loaded.last().unwrap().content, "199");
352    }
353
354    #[tokio::test]
355    async fn load_truncates_to_token_budget() {
356        let m = fresh().await;
357        let t = ThreadId::new("t1");
358        // Each ~40-char message ~= 10 tokens.
359        for i in 0..20 {
360            m.append(
361                t.clone(),
362                user(&format!("msg-{i:03}-padding-padding-padding")),
363            )
364            .await
365            .unwrap();
366        }
367        let loaded = m.load(t, 30).await.unwrap();
368        // A 30-token budget buys ~3 of these, and both ends are kept on top of
369        // that: head, gap marker, then the newest turns.
370        assert!(
371            loaded.len() <= 6 && !loaded.is_empty(),
372            "expected the budget to bound the middle, got {}",
373            loaded.len()
374        );
375        assert!(
376            loaded.first().unwrap().content.contains("msg-000"),
377            "the first message must survive truncation"
378        );
379        assert!(
380            loaded.last().unwrap().content.contains("msg-019"),
381            "newest message must survive truncation"
382        );
383        assert!(
384            loaded[1]
385                .content
386                .starts_with(klieo_core::memory::ELIDED_MESSAGES_MARKER_PREFIX),
387            "the omitted span must be marked, not closed silently: {:?}",
388            loaded[1].content
389        );
390    }
391
392    /// The defect that sent a factory Coder a system prompt and no task: a
393    /// brief larger than the whole budget loaded as an empty history, so the
394    /// agent had nothing to work from and invented its own task.
395    #[tokio::test]
396    async fn a_brief_larger_than_the_budget_still_loads() {
397        let m = fresh().await;
398        let t = ThreadId::new("t1");
399        m.append(t.clone(), user(&"x".repeat(40_000)))
400            .await
401            .unwrap();
402
403        let loaded = m.load(t, 8_000).await.unwrap();
404
405        assert_eq!(
406            loaded.len(),
407            1,
408            "a non-empty thread must never load as an empty history"
409        );
410    }
411
412    /// The row cap fetches the NEWEST rows, so the head row has to be fetched
413    /// explicitly or the first message is not even a candidate.
414    #[tokio::test]
415    async fn the_first_message_survives_a_thread_longer_than_the_row_cap() {
416        let m = fresh().await;
417        let t = ThreadId::new("t1");
418        m.append(t.clone(), user("the brief")).await.unwrap();
419        for i in 0..200 {
420            m.append(t.clone(), user(&format!("turn-{i}")))
421                .await
422                .unwrap();
423        }
424
425        let loaded = m.load(t, 10).await.unwrap();
426
427        assert_eq!(
428            loaded.first().unwrap().content,
429            "the brief",
430            "the head row must be fetched even when the row cap excludes it"
431        );
432        assert!(
433            loaded.last().unwrap().content.contains("turn-199"),
434            "and the newest turn must still be there"
435        );
436    }
437
438    #[tokio::test]
439    async fn clear_removes_thread() {
440        let m = fresh().await;
441        let t = ThreadId::new("t1");
442        m.append(t.clone(), user("hello")).await.unwrap();
443        m.clear(t.clone()).await.unwrap();
444        let loaded = m.load(t, 10_000).await.unwrap();
445        assert!(loaded.is_empty());
446    }
447
448    #[tokio::test]
449    async fn threads_are_isolated() {
450        let m = fresh().await;
451        m.append(ThreadId::new("a"), user("a-msg")).await.unwrap();
452        m.append(ThreadId::new("b"), user("b-msg")).await.unwrap();
453        let a = m.load(ThreadId::new("a"), 10_000).await.unwrap();
454        let b = m.load(ThreadId::new("b"), 10_000).await.unwrap();
455        assert_eq!(a.len(), 1);
456        assert_eq!(b.len(), 1);
457        assert_eq!(a[0].content, "a-msg");
458        assert_eq!(b[0].content, "b-msg");
459    }
460
461    #[tokio::test]
462    async fn role_round_trip_for_all_variants() {
463        let m = fresh().await;
464        let t = ThreadId::new("r");
465        for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
466            m.append(
467                t.clone(),
468                Message {
469                    role,
470                    content: "x".into(),
471                    tool_calls: vec![],
472                    tool_call_id: None,
473                },
474            )
475            .await
476            .unwrap();
477        }
478        let loaded = m.load(t, 10_000).await.unwrap();
479        assert_eq!(loaded.len(), 4);
480        assert_eq!(loaded[0].role, Role::System);
481        assert_eq!(loaded[3].role, Role::Tool);
482    }
483}
klieo_memory_sqlite/short_term.rs

klieo_memory_sqlite/
short_term.rs
