Skip to main content

context_forge/storage/
mod.rs

1use std::path::Path;
2use std::sync::Arc;
3
4/// Shared in-memory tantivy FTS index.
5pub(crate) mod fts_index;
6/// Turso-backed async FTS searcher.
7pub mod turso_searcher;
8/// Turso-backed async storage implementation.
9pub mod turso_storage;
10
11pub use turso_searcher::TursoSearcher;
12pub use turso_storage::TursoStorage;
13
14/// Create a paired storage + searcher backed by turso.
15///
16/// # Errors
17///
18/// Returns an error if the database cannot be opened or migrations fail.
19pub async fn open_storage(
20    db_path: &Path,
21    max_entries: usize,
22) -> crate::Result<(TursoStorage, TursoSearcher)> {
23    let storage = TursoStorage::open(db_path, max_entries).await?;
24    let db = Arc::clone(&storage.db);
25    let fts = Arc::clone(&storage.fts);
26    let searcher = TursoSearcher::new(db, fts);
27    Ok((storage, searcher))
28}
29
30#[cfg(test)]
31mod tests {
32    use std::path::Path;
33
34    use crate::engine::MATCH_ALL_QUERY;
35    use crate::entry::{kind, ContextEntry};
36    use crate::storage::open_storage;
37    use crate::traits::{ContextStorage, Searcher};
38
39    fn make_entry(id: &str, content: &str, timestamp: i64, kind: &str) -> ContextEntry {
40        ContextEntry {
41            id: id.into(),
42            content: content.into(),
43            timestamp,
44            kind: kind.to_owned(),
45            scope: None,
46            session_id: None,
47            token_count: None,
48            metadata: None,
49        }
50    }
51
52    fn make_scoped_entry(id: &str, content: &str, timestamp: i64, scope: &str) -> ContextEntry {
53        ContextEntry {
54            id: id.into(),
55            content: content.into(),
56            timestamp,
57            kind: kind::MANUAL.to_owned(),
58            scope: Some(scope.to_owned()),
59            session_id: None,
60            token_count: None,
61            metadata: None,
62        }
63    }
64
65    #[tokio::test]
66    async fn test_save_and_count() {
67        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
68        let entry = make_entry("e1", "hello world", 1000, kind::MANUAL);
69        storage.save(&entry).await.unwrap();
70        assert_eq!(storage.count().await.unwrap(), 1);
71    }
72
73    #[tokio::test]
74    async fn test_save_and_get_top_k() {
75        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
76        storage
77            .save(&make_entry("e1", "first", 100, kind::MANUAL))
78            .await
79            .unwrap();
80        storage
81            .save(&make_entry("e2", "second", 200, kind::SNAPSHOT))
82            .await
83            .unwrap();
84        storage
85            .save(&make_entry("e3", "third", 300, kind::SUMMARY))
86            .await
87            .unwrap();
88
89        let top2 = storage.get_top_k(2).await.unwrap();
90        assert_eq!(top2.len(), 2);
91        assert_eq!(top2[0].id, "e3"); // most recent
92        assert_eq!(top2[1].id, "e2");
93    }
94
95    #[tokio::test]
96    async fn test_save_and_get_all() {
97        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
98        storage
99            .save(&make_entry("e1", "first", 100, kind::MANUAL))
100            .await
101            .unwrap();
102        storage
103            .save(&make_entry("e2", "second", 200, kind::MANUAL))
104            .await
105            .unwrap();
106        storage
107            .save(&make_entry("e3", "third", 300, kind::MANUAL))
108            .await
109            .unwrap();
110
111        let all = storage.get_all().await.unwrap();
112        assert_eq!(all.len(), 3);
113        assert_eq!(all[0].id, "e3");
114        assert_eq!(all[1].id, "e2");
115        assert_eq!(all[2].id, "e1");
116    }
117
118    #[tokio::test]
119    async fn test_delete() {
120        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
121        storage
122            .save(&make_entry("e1", "hello", 1000, kind::MANUAL))
123            .await
124            .unwrap();
125
126        assert!(storage.delete("e1").await.unwrap());
127        assert!(!storage.delete("nonexistent").await.unwrap());
128        assert_eq!(storage.count().await.unwrap(), 0);
129    }
130
131    #[tokio::test]
132    async fn test_clear() {
133        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
134        storage
135            .save(&make_entry("e1", "a", 100, kind::MANUAL))
136            .await
137            .unwrap();
138        storage
139            .save(&make_entry("e2", "b", 200, kind::MANUAL))
140            .await
141            .unwrap();
142        storage
143            .save(&make_entry("e3", "c", 300, kind::MANUAL))
144            .await
145            .unwrap();
146
147        let cleared = storage.clear().await.unwrap();
148        assert_eq!(cleared, 3);
149        assert_eq!(storage.count().await.unwrap(), 0);
150    }
151
152    #[tokio::test]
153    async fn test_clear_scope() {
154        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
155        storage
156            .save(&make_scoped_entry("e1", "a", 100, "scope-a"))
157            .await
158            .unwrap();
159        storage
160            .save(&make_scoped_entry("e2", "b", 200, "scope-a"))
161            .await
162            .unwrap();
163        storage
164            .save(&make_scoped_entry("e3", "c", 300, "scope-b"))
165            .await
166            .unwrap();
167
168        let cleared = storage.clear_scope("scope-a").await.unwrap();
169        assert_eq!(cleared, 2);
170        assert_eq!(storage.count().await.unwrap(), 1);
171
172        let all = storage.get_all().await.unwrap();
173        assert_eq!(all[0].id, "e3");
174    }
175
176    #[tokio::test]
177    async fn test_clear_scope_no_match() {
178        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
179        storage
180            .save(&make_scoped_entry("e1", "a", 100, "scope-a"))
181            .await
182            .unwrap();
183
184        let cleared = storage.clear_scope("scope-z").await.unwrap();
185        assert_eq!(cleared, 0);
186        assert_eq!(storage.count().await.unwrap(), 1);
187    }
188
189    #[tokio::test]
190    async fn test_lru_eviction() {
191        let (storage, _) = open_storage(Path::new(":memory:"), 2).await.unwrap();
192        storage
193            .save(&make_entry("e1", "oldest", 100, kind::MANUAL))
194            .await
195            .unwrap();
196        storage
197            .save(&make_entry("e2", "middle", 200, kind::MANUAL))
198            .await
199            .unwrap();
200        storage
201            .save(&make_entry("e3", "newest", 300, kind::MANUAL))
202            .await
203            .unwrap();
204
205        assert_eq!(storage.count().await.unwrap(), 2);
206
207        let all = storage.get_all().await.unwrap();
208        let ids: Vec<&str> = all.iter().map(|e| e.id.as_str()).collect();
209        assert!(
210            !ids.contains(&"e1"),
211            "oldest entry should have been evicted"
212        );
213        assert!(ids.contains(&"e2"));
214        assert!(ids.contains(&"e3"));
215    }
216
217    #[tokio::test]
218    async fn test_save_batch_evicts_oldest_over_capacity() {
219        // A single batch that overflows capacity must evict the oldest rows once,
220        // keeping the newest `max_entries` — the batched-eviction path.
221        let (storage, _) = open_storage(Path::new(":memory:"), 3).await.unwrap();
222        let batch = vec![
223            make_entry("e1", "oldest", 100, kind::MANUAL),
224            make_entry("e2", "second", 200, kind::MANUAL),
225            make_entry("e3", "third", 300, kind::MANUAL),
226            make_entry("e4", "fourth", 400, kind::MANUAL),
227            make_entry("e5", "newest", 500, kind::MANUAL),
228        ];
229        storage.save_batch(&batch).await.unwrap();
230
231        assert_eq!(storage.count().await.unwrap(), 3);
232        let all = storage.get_all().await.unwrap();
233        let ids: Vec<&str> = all.iter().map(|e| e.id.as_str()).collect();
234        assert!(!ids.contains(&"e1"), "oldest should be evicted");
235        assert!(!ids.contains(&"e2"), "second-oldest should be evicted");
236        assert!(ids.contains(&"e3"));
237        assert!(ids.contains(&"e4"));
238        assert!(ids.contains(&"e5"));
239    }
240
241    #[tokio::test]
242    async fn test_save_batch_persists_and_indexes() {
243        // save_batch must persist every entry and commit the FTS index so the
244        // batch is immediately searchable — equivalent to a sequence of save().
245        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
246        let batch = vec![
247            make_entry("b1", "rust programming language", 100, kind::MANUAL),
248            make_entry("b2", "python scripting tools", 200, kind::MANUAL),
249        ];
250        storage.save_batch(&batch).await.unwrap();
251
252        assert_eq!(storage.count().await.unwrap(), 2);
253        let hits = searcher.search("rust", None, 10).await.unwrap();
254        assert!(
255            hits.iter().any(|e| e.entry.id == "b1"),
256            "batch-saved entry should be FTS-searchable"
257        );
258    }
259
260    #[tokio::test]
261    async fn test_fts_search() {
262        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
263        storage
264            .save(&make_entry(
265                "e1",
266                "rust programming language",
267                100,
268                kind::MANUAL,
269            ))
270            .await
271            .unwrap();
272        storage
273            .save(&make_entry("e2", "python scripting", 200, kind::MANUAL))
274            .await
275            .unwrap();
276        storage
277            .save(&make_entry("e3", "rust borrow checker", 300, kind::MANUAL))
278            .await
279            .unwrap();
280
281        let results = searcher.search("rust", None, 5).await.unwrap();
282        assert_eq!(results.len(), 2);
283        assert!(
284            results[0].score >= results[1].score,
285            "results should be ordered by descending score"
286        );
287    }
288
289    #[tokio::test]
290    async fn test_fts_search_no_results() {
291        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
292        storage
293            .save(&make_entry("e1", "hello world", 100, kind::MANUAL))
294            .await
295            .unwrap();
296
297        let results = searcher.search("nonexistent", None, 5).await.unwrap();
298        assert!(results.is_empty());
299    }
300
301    #[tokio::test]
302    async fn test_fts_search_scoped() {
303        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
304        storage
305            .save(&make_scoped_entry("e1", "rust programming", 100, "a"))
306            .await
307            .unwrap();
308        storage
309            .save(&make_scoped_entry("e2", "rust borrow checker", 200, "b"))
310            .await
311            .unwrap();
312
313        let results_a = searcher.search("rust", Some("a"), 5).await.unwrap();
314        assert_eq!(results_a.len(), 1);
315        assert_eq!(results_a[0].entry.id, "e1");
316
317        let results_b = searcher.search("rust", Some("b"), 5).await.unwrap();
318        assert_eq!(results_b.len(), 1);
319        assert_eq!(results_b[0].entry.id, "e2");
320
321        let results_all = searcher.search("rust", None, 5).await.unwrap();
322        assert_eq!(results_all.len(), 2);
323    }
324
325    #[tokio::test]
326    async fn test_new_entry_with_scope_and_metadata() {
327        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
328
329        let metadata = serde_json::json!({"runtime": "codex", "model": "gpt-5.3-codex"});
330        let entry = ContextEntry {
331            id: "v3-entry".into(),
332            content: "entry with scope and metadata".into(),
333            timestamp: 1_700_000_100,
334            kind: kind::SUMMARY.to_owned(),
335            scope: Some("project:homelab-rs".into()),
336            session_id: Some("session-123".into()),
337            token_count: Some(6),
338            metadata: Some(metadata.clone()),
339        };
340
341        storage.save(&entry).await.unwrap();
342
343        let all = storage.get_all().await.unwrap();
344        assert_eq!(all.len(), 1);
345        let got = &all[0];
346        assert_eq!(got.id, entry.id);
347        assert_eq!(got.content, entry.content);
348        assert_eq!(got.timestamp, entry.timestamp);
349        assert_eq!(got.kind, entry.kind);
350        assert_eq!(got.scope, entry.scope);
351        assert_eq!(got.session_id, entry.session_id);
352        assert_eq!(got.token_count, entry.token_count);
353        assert_eq!(got.metadata, Some(metadata));
354    }
355
356    #[tokio::test]
357    async fn test_insert_or_replace() {
358        let (storage, _) = open_storage(Path::new(":memory:"), 100).await.unwrap();
359        storage
360            .save(&make_entry("e1", "original content", 100, kind::MANUAL))
361            .await
362            .unwrap();
363        storage
364            .save(&make_entry("e1", "updated content", 200, kind::SUMMARY))
365            .await
366            .unwrap();
367
368        assert_eq!(storage.count().await.unwrap(), 1);
369
370        let all = storage.get_all().await.unwrap();
371        assert_eq!(all[0].content, "updated content");
372    }
373
374    #[tokio::test]
375    async fn test_search_match_all_query() {
376        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
377        storage
378            .save(&make_entry("e1", "first entry", 100, kind::MANUAL))
379            .await
380            .unwrap();
381        storage
382            .save(&make_entry("e2", "second entry", 200, kind::SNAPSHOT))
383            .await
384            .unwrap();
385        storage
386            .save(&make_entry("e3", "third entry", 300, kind::SUMMARY))
387            .await
388            .unwrap();
389
390        let results = searcher.search(MATCH_ALL_QUERY, None, 10).await.unwrap();
391        assert_eq!(results.len(), 3);
392
393        // Ordered by timestamp descending (newest first).
394        assert_eq!(results[0].entry.id, "e3");
395        assert_eq!(results[1].entry.id, "e2");
396        assert_eq!(results[2].entry.id, "e1");
397
398        // Match-all results all share the same fixed score.
399        for r in &results {
400            assert!((r.score - 1.0).abs() < f64::EPSILON);
401        }
402    }
403
404    #[tokio::test]
405    async fn test_search_match_all_query_scoped() {
406        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
407        storage
408            .save(&make_scoped_entry("e1", "first entry", 100, "a"))
409            .await
410            .unwrap();
411        storage
412            .save(&make_scoped_entry("e2", "second entry", 200, "b"))
413            .await
414            .unwrap();
415
416        let results_a = searcher
417            .search(MATCH_ALL_QUERY, Some("a"), 10)
418            .await
419            .unwrap();
420        assert_eq!(results_a.len(), 1);
421        assert_eq!(results_a[0].entry.id, "e1");
422
423        let results_all = searcher.search(MATCH_ALL_QUERY, None, 10).await.unwrap();
424        assert_eq!(results_all.len(), 2);
425    }
426
427    #[tokio::test]
428    async fn search_with_fts5_operator_characters_does_not_error() {
429        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
430        storage
431            .save(&make_entry("e1", "marco polo", 100, kind::MANUAL))
432            .await
433            .unwrap();
434
435        // Production bug: a query containing `"` and `.` previously caused
436        // `fts5: syntax error`. It must now succeed and find the entry.
437        let results = searcher
438            .search(r#"if I say "marco"."#, None, 10)
439            .await
440            .unwrap();
441
442        assert!(
443            results.iter().any(|r| r.entry.id == "e1"),
444            "expected 'marco polo' entry to be found despite FTS5 operator characters in the query"
445        );
446    }
447
448    #[tokio::test]
449    async fn test_fts_search_scores_are_nonzero() {
450        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
451        storage
452            .save(&make_entry(
453                "e1",
454                "rust programming language",
455                100,
456                kind::MANUAL,
457            ))
458            .await
459            .unwrap();
460        storage
461            .save(&make_entry(
462                "e2",
463                "python scripting tools",
464                200,
465                kind::MANUAL,
466            ))
467            .await
468            .unwrap();
469
470        let results = searcher.search("rust", None, 5).await.unwrap();
471        assert!(!results.is_empty());
472        assert!(
473            results[0].score > 0.0,
474            "expected non-zero BM25 score from tantivy, got {}",
475            results[0].score
476        );
477    }
478
479    #[tokio::test]
480    async fn search_or_joins_terms_for_message_length_queries() {
481        let (storage, searcher) = open_storage(Path::new(":memory:"), 100).await.unwrap();
482        storage
483            .save(&make_entry("e1", "alpha one", 100, kind::MANUAL))
484            .await
485            .unwrap();
486        storage
487            .save(&make_entry("e2", "beta two", 200, kind::MANUAL))
488            .await
489            .unwrap();
490
491        // Under implicit AND, no entry contains all of "alpha", "beta", and
492        // "gamma", so this would return zero results. OR-join must return both.
493        let results = searcher.search("alpha beta gamma", None, 10).await.unwrap();
494
495        let ids: Vec<&str> = results.iter().map(|r| r.entry.id.as_str()).collect();
496        assert!(ids.contains(&"e1"), "expected 'alpha one' entry in results");
497        assert!(ids.contains(&"e2"), "expected 'beta two' entry in results");
498    }
499}