Skip to main content

oxi/store/
memory_sqlite.rs

1//! SQLite-backed memory store implementing [`MemoryBackend`].
2//!
3//! Provides persistent storage with WAL journal mode for concurrent access
4//! safety. Embedding storage is reserved for a future cosine-search upgrade;
5//! search currently uses simple SQL `LIKE` matching.
6
7use oxi_agent::tools::{MemoryBackend, MemoryItem, ToolError};
8use rusqlite::{Connection, params};
9use std::path::Path;
10use std::pin::Pin;
11use tokio::sync::Mutex;
12
13/// SQLite-backed memory store.
14///
15/// Implements [`MemoryBackend`] for the `memory_*` agent tools.
16/// Each memory is stored with an auto-generated UUID, kind, content, and subject.
17/// The schema reserves an `embedding` column for future cosine-search support.
18#[derive(Debug)]
19pub struct SqliteMemoryStore {
20    db: Mutex<Connection>,
21}
22
23impl SqliteMemoryStore {
24    /// Open or create a SQLite memory store at `path`.
25    ///
26    /// Uses `:memory:` for an in-memory database. For persistent paths,
27    /// filesystem-aware journal-mode selection is applied (see
28    /// [`oxi_mnemopi::journal::JournalMode`]). On network filesystems the
29    /// engine falls back to `TRUNCATE` mode + per-host DB sibling to avoid
30    /// SIGBUS from mmap'd `-shm`.
31    pub fn open(path: &Path) -> Result<Self, rusqlite::Error> {
32        let is_memory = path == Path::new(":memory:");
33        // Durable primary store: detect journal mode (TRUNCATE on NFS for
34        // SIGBUS safety) but NEVER rewrite the path — user memories must
35        // stay coherent across hosts. See `JournalMode::effective_db_path`.
36        let mode = if is_memory {
37            oxi_mnemopi::journal::JournalMode::Wal
38        } else {
39            oxi_mnemopi::journal::JournalMode::for_db_path(path)
40        };
41        let conn = Connection::open(path)?;
42
43        conn.execute_batch("PRAGMA foreign_keys = ON;")?;
44        conn.execute_batch(&format!(
45            "PRAGMA busy_timeout = {};",
46            mode.busy_timeout_ms()
47        ))?;
48
49        if !is_memory {
50            // Apply the detected (or env-overridden) journal mode. Failures
51            // fall back to SQLite's default `delete` journal — still safe,
52            // just slower than WAL.
53            if let Err(e) = conn.execute_batch(&format!("PRAGMA journal_mode = {};", mode.as_str()))
54            {
55                tracing::warn!(
56                    mode = mode.as_str(),
57                    path = %path.display(),
58                    error = %e,
59                    "failed to set journal_mode; SQLite default will be used"
60                );
61            }
62        }
63
64        conn.execute_batch(
65            "CREATE TABLE IF NOT EXISTS memories (
66                id          TEXT PRIMARY KEY,
67                subject     TEXT NOT NULL,
68                kind        TEXT NOT NULL,
69                content     TEXT NOT NULL,
70                embedding   BLOB,
71                created_at  TEXT NOT NULL DEFAULT (datetime('now')),
72                updated_at  TEXT NOT NULL DEFAULT (datetime('now')),
73                metadata    TEXT
74            );",
75        )?;
76
77        Ok(Self {
78            db: Mutex::new(conn),
79        })
80    }
81}
82
83impl MemoryBackend for SqliteMemoryStore {
84    fn put<'a>(
85        &'a self,
86        content: &'a str,
87        kind: &'a str,
88        subject: &'a str,
89    ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
90        Box::pin(async move {
91            let id = uuid::Uuid::new_v4().to_string();
92            let db = self.db.lock().await;
93            db.execute(
94                "INSERT INTO memories (id, subject, kind, content)
95                 VALUES (?1, ?2, ?3, ?4)",
96                params![id, subject, kind, content],
97            )
98            .map_err(|e| format!("Failed to store memory: {e}"))?;
99            Ok(id)
100        })
101    }
102
103    fn search<'a>(
104        &'a self,
105        query: &'a str,
106        k: usize,
107    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
108        Box::pin(async move {
109            let db = self.db.lock().await;
110            let pattern = format!("%{}%", query.replace('%', "\\%").replace('_', "\\_"));
111            let mut stmt = db
112                .prepare(
113                    "SELECT id, kind, content, subject
114                     FROM memories
115                     WHERE content LIKE ?1 ESCAPE '\\'
116                     ORDER BY length(content) ASC
117                     LIMIT ?2",
118                )
119                .map_err(|e| format!("Failed to prepare search: {e}"))?;
120
121            let results: Vec<MemoryItem> = stmt
122                .query_map(params![pattern, k as i64], |row| {
123                    Ok(MemoryItem {
124                        id: row.get(0)?,
125                        kind: row.get(1)?,
126                        content: row.get(2)?,
127                        subject: row.get(3)?,
128                    })
129                })
130                .map_err(|e| format!("Failed to search memories: {e}"))?
131                .filter_map(|r| r.ok())
132                .collect();
133
134            Ok(results)
135        })
136    }
137
138    fn list<'a>(
139        &'a self,
140        subject: &'a str,
141    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
142        Box::pin(async move {
143            let db = self.db.lock().await;
144            let mut stmt = db
145                .prepare(
146                    "SELECT id, kind, content, subject
147                     FROM memories
148                     WHERE subject = ?1
149                     ORDER BY updated_at DESC",
150                )
151                .map_err(|e| format!("Failed to prepare list: {e}"))?;
152
153            let results: Vec<MemoryItem> = stmt
154                .query_map(params![subject], |row| {
155                    Ok(MemoryItem {
156                        id: row.get(0)?,
157                        kind: row.get(1)?,
158                        content: row.get(2)?,
159                        subject: row.get(3)?,
160                    })
161                })
162                .map_err(|e| format!("Failed to list memories: {e}"))?
163                .filter_map(|r| r.ok())
164                .collect();
165
166            Ok(results)
167        })
168    }
169
170    fn delete<'a>(
171        &'a self,
172        id: &'a str,
173    ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
174        Box::pin(async move {
175            let db = self.db.lock().await;
176            db.execute("DELETE FROM memories WHERE id = ?1", params![id])
177                .map_err(|e| format!("Failed to delete memory: {e}"))?;
178            Ok(())
179        })
180    }
181
182    /// Bulk delete: a single SQL statement, no per-row round-trip.
183    /// Returns the number of rows removed across all subjects.
184    fn clear_all<'a>(
185        &'a self,
186    ) -> Pin<Box<dyn Future<Output = Result<usize, ToolError>> + Send + 'a>> {
187        Box::pin(async move {
188            let db = self.db.lock().await;
189            let removed = db
190                .execute("DELETE FROM memories", [])
191                .map_err(|e| format!("Failed to clear memories: {e}"))?;
192            Ok(removed)
193        })
194    }
195
196    /// No background pipeline: surface an honest "not supported" rather
197    /// than claiming to have enqueued work that did not happen.
198    fn enqueue_consolidation<'a>(
199        &'a self,
200    ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
201        Box::pin(async move {
202            Err(
203                "enqueue_consolidation not supported by SqliteMemoryStore (no pipeline)"
204                    .to_string(),
205            )
206        })
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    fn tmp_store() -> (SqliteMemoryStore, tempfile::TempDir) {
215        let dir = tempfile::tempdir().expect("tempdir");
216        let store = SqliteMemoryStore::open(&dir.path().join("mem.db")).unwrap();
217        (store, dir)
218    }
219
220    #[tokio::test]
221    async fn put_list_delete_roundtrip() {
222        let (store, _dir) = tmp_store();
223        let id = store
224            .put("user prefers dark mode", "preference", "alice")
225            .await
226            .unwrap();
227        assert!(!id.is_empty());
228
229        let items = store.list("alice").await.unwrap();
230        assert_eq!(items.len(), 1);
231        assert_eq!(items[0].content, "user prefers dark mode");
232        assert_eq!(items[0].subject, "alice");
233        assert_eq!(items[0].kind, "preference");
234
235        store.delete(&id).await.unwrap();
236        assert!(store.list("alice").await.unwrap().is_empty());
237    }
238
239    #[tokio::test]
240    async fn list_scopes_per_subject() {
241        let (store, _dir) = tmp_store();
242        let a = store.put("a1", "fact", "alice").await.unwrap();
243        let b = store.put("b1", "fact", "bob").await.unwrap();
244
245        assert_eq!(store.list("alice").await.unwrap().len(), 1);
246        assert_eq!(store.list("bob").await.unwrap().len(), 1);
247        assert!(store.list("nobody").await.unwrap().is_empty());
248
249        store.delete(&a).await.unwrap();
250        store.delete(&b).await.unwrap();
251    }
252
253    #[tokio::test]
254    async fn search_finds_substring_match() {
255        let (store, _dir) = tmp_store();
256        store
257            .put("likes rust ownership", "fact", "alice")
258            .await
259            .unwrap();
260        store.put("prefers python", "fact", "alice").await.unwrap();
261
262        let results = store.search("rust", 5).await.unwrap();
263        assert!(!results.is_empty(), "LIKE search must match substring");
264        assert!(results.iter().any(|i| i.content.contains("rust")));
265    }
266
267    #[tokio::test]
268    async fn memory_info_returns_none_for_sqlite_backend() {
269        // Pins the contract — the SQLite backend has no bespoke
270        // diagnostics; future accidental override should be caught here.
271        let (store, _dir) = tmp_store();
272        assert!(store.memory_info().is_none());
273    }
274
275    #[tokio::test]
276    async fn clear_all_wipes_every_subject() {
277        let (store, _dir) = tmp_store();
278        store.put("a1", "fact", "alice").await.unwrap();
279        store.put("a2", "fact", "alice").await.unwrap();
280        store.put("b1", "fact", "bob").await.unwrap();
281
282        let removed = store.clear_all().await.unwrap();
283        assert_eq!(removed, 3, "all three rows removed in one DELETE");
284        assert!(store.list("alice").await.unwrap().is_empty());
285        assert!(store.list("bob").await.unwrap().is_empty());
286    }
287
288    #[tokio::test]
289    async fn clear_all_is_idempotent_on_empty_store() {
290        let (store, _dir) = tmp_store();
291        let removed = store.clear_all().await.unwrap();
292        assert_eq!(removed, 0);
293    }
294
295    #[tokio::test]
296    async fn enqueue_consolidation_surfaces_honest_unsupported() {
297        let (store, _dir) = tmp_store();
298        let err = store.enqueue_consolidation().await.unwrap_err();
299        assert!(
300            err.contains("not supported"),
301            "SqliteMemoryStore must surface the absence of a pipeline, got: {err}"
302        );
303    }
304}