1use std::path::Path;
2use std::path::PathBuf;
3
4use chrono::{DateTime, Utc};
5use rusqlite::{params, Connection, OpenFlags};
6use serde::{Deserialize, Serialize};
7
8use crate::error::{Result, SqzError};
9use crate::types::{CompressedContent, SessionId, SessionState};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SessionSummary {
15 pub id: SessionId,
16 pub project_dir: PathBuf,
17 pub compressed_summary: String,
18 pub created_at: DateTime<Utc>,
19 pub updated_at: DateTime<Utc>,
20}
21
22pub struct SessionStore {
38 db: Connection,
39}
40
41const SCHEMA: &str = r#"
44PRAGMA journal_mode = WAL;
45
46CREATE TABLE IF NOT EXISTS sessions (
47 id TEXT PRIMARY KEY,
48 project_dir TEXT NOT NULL,
49 compressed_summary TEXT NOT NULL,
50 created_at TEXT NOT NULL,
51 updated_at TEXT NOT NULL,
52 data BLOB NOT NULL
53);
54
55CREATE VIRTUAL TABLE IF NOT EXISTS sessions_fts USING fts5(
56 id,
57 project_dir,
58 compressed_summary,
59 content='sessions',
60 content_rowid='rowid',
61 tokenize='porter ascii'
62);
63
64CREATE TRIGGER IF NOT EXISTS sessions_ai AFTER INSERT ON sessions BEGIN
65 INSERT INTO sessions_fts(rowid, id, project_dir, compressed_summary)
66 VALUES (new.rowid, new.id, new.project_dir, new.compressed_summary);
67END;
68
69CREATE TRIGGER IF NOT EXISTS sessions_ad AFTER DELETE ON sessions BEGIN
70 INSERT INTO sessions_fts(sessions_fts, rowid, id, project_dir, compressed_summary)
71 VALUES ('delete', old.rowid, old.id, old.project_dir, old.compressed_summary);
72END;
73
74CREATE TRIGGER IF NOT EXISTS sessions_au AFTER UPDATE ON sessions BEGIN
75 INSERT INTO sessions_fts(sessions_fts, rowid, id, project_dir, compressed_summary)
76 VALUES ('delete', old.rowid, old.id, old.project_dir, old.compressed_summary);
77 INSERT INTO sessions_fts(rowid, id, project_dir, compressed_summary)
78 VALUES (new.rowid, new.id, new.project_dir, new.compressed_summary);
79END;
80
81CREATE TABLE IF NOT EXISTS cache_entries (
82 hash TEXT PRIMARY KEY,
83 data TEXT NOT NULL,
84 accessed_at TEXT NOT NULL,
85 -- Raw pre-compression bytes so `sqz expand <prefix>` can serve
86 -- truly uncompressed content to agents that cannot parse
87 -- `§ref:…§` dedup tokens. Nullable because the column was added
88 -- in an additive migration; rows written before that migration
89 -- (or via callers that don't have the original bytes) have NULL.
90 original BLOB,
91 -- Pinned entries bypass the TTL freshness check. Used for stable
92 -- knowledge-base content (a wiki, project documentation, system
93 -- prompts) that's persistently re-injected into every agent and
94 -- therefore never gets compacted out of the LLM's context. See
95 -- `sqz pin` for the user-facing CLI.
96 pinned INTEGER NOT NULL DEFAULT 0
97);
98
99CREATE TABLE IF NOT EXISTS compression_log (
100 id INTEGER PRIMARY KEY AUTOINCREMENT,
101 tokens_original INTEGER NOT NULL,
102 tokens_compressed INTEGER NOT NULL,
103 stages_applied TEXT NOT NULL,
104 mode TEXT NOT NULL DEFAULT 'auto',
105 created_at TEXT NOT NULL
106);
107
108-- Compression-regret signals. A 'rerun' row means the agent re-produced
109-- byte-identical output within a short window of last seeing it (the
110-- repeat bought no new information); an 'expand' row means a §ref:…§
111-- token was expanded back to its original bytes (a recovery round trip).
112-- Both are proxies for "compression dropped something the agent needed",
113-- surfaced by `sqz stats` so aggressiveness can be judged per command.
114CREATE TABLE IF NOT EXISTS regret_log (
115 id INTEGER PRIMARY KEY AUTOINCREMENT,
116 kind TEXT NOT NULL,
117 cmd TEXT NOT NULL,
118 hash TEXT NOT NULL,
119 gap_secs REAL NOT NULL,
120 created_at TEXT NOT NULL
121);
122
123CREATE TABLE IF NOT EXISTS known_files (
124 path TEXT PRIMARY KEY,
125 added_at TEXT NOT NULL
126);
127
128-- Small key/value store for engine-wide state that needs to persist across
129-- short-lived sqz processes (each shell-hook invocation is a new process).
130-- Initially used only for the last_compaction_at marker: cache entries with
131-- `accessed_at < last_compaction_at` are treated as stale even if still
132-- within the normal TTL. See cache_manager.rs for the freshness model.
133CREATE TABLE IF NOT EXISTS metadata (
134 key TEXT PRIMARY KEY,
135 value TEXT NOT NULL
136);
137"#;
138
139pub(crate) fn apply_schema(conn: &Connection) -> rusqlite::Result<()> {
142 conn.execute_batch(SCHEMA)?;
143 let has_original: bool = conn
155 .prepare("SELECT 1 FROM pragma_table_info('cache_entries') WHERE name = 'original'")?
156 .query_row([], |_| Ok(()))
157 .is_ok();
158 if !has_original {
159 conn.execute("ALTER TABLE cache_entries ADD COLUMN original BLOB", [])?;
160 }
161
162 let has_project_dir: bool = conn
167 .prepare("SELECT 1 FROM pragma_table_info('compression_log') WHERE name = 'project_dir'")?
168 .query_row([], |_| Ok(()))
169 .is_ok();
170 if !has_project_dir {
171 conn.execute("ALTER TABLE compression_log ADD COLUMN project_dir TEXT", [])?;
172 }
173
174 let has_pinned: bool = conn
178 .prepare("SELECT 1 FROM pragma_table_info('cache_entries') WHERE name = 'pinned'")?
179 .query_row([], |_| Ok(()))
180 .is_ok();
181 if !has_pinned {
182 conn.execute(
183 "ALTER TABLE cache_entries ADD COLUMN pinned INTEGER NOT NULL DEFAULT 0",
184 [],
185 )?;
186 }
187
188 Ok(())
189}
190
191fn open_connection(path: &Path) -> rusqlite::Result<Connection> {
192 let conn = Connection::open(path)?;
193 apply_schema(&conn)?;
194 Ok(conn)
195}
196
197fn row_to_summary(
198 id: String,
199 project_dir: String,
200 compressed_summary: String,
201 created_at: String,
202 updated_at: String,
203) -> Result<SessionSummary> {
204 let created_at = created_at
205 .parse::<DateTime<Utc>>()
206 .map_err(|e| SqzError::Other(format!("invalid created_at timestamp: {e}")))?;
207 let updated_at = updated_at
208 .parse::<DateTime<Utc>>()
209 .map_err(|e| SqzError::Other(format!("invalid updated_at timestamp: {e}")))?;
210 Ok(SessionSummary {
211 id,
212 project_dir: PathBuf::from(project_dir),
213 compressed_summary,
214 created_at,
215 updated_at,
216 })
217}
218
219impl SessionStore {
222 #[cfg(test)]
225 pub(crate) fn from_connection(conn: Connection) -> Self {
226 Self { db: conn }
227 }
228
229 pub fn open(path: &Path) -> Result<Self> {
232 let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)?;
233 apply_schema(&conn)?;
234 Ok(Self { db: conn })
235 }
236
237 pub fn open_or_create(path: &Path) -> Result<Self> {
241 match open_connection(path) {
242 Ok(conn) => Ok(Self { db: conn }),
243 Err(e) => {
244 eprintln!(
245 "sqz warning: session store at '{}' is corrupted or inaccessible ({e}). \
246 Creating a new database. Prior session data has been lost.",
247 path.display()
248 );
249 if path.exists() {
251 let _ = std::fs::remove_file(path);
252 }
253 let conn = open_connection(path)
254 .map_err(|e2| SqzError::Other(format!("failed to create new session store: {e2}")))?;
255 Ok(Self { db: conn })
256 }
257 }
258 }
259
260 pub fn save_session(&self, session: &SessionState) -> Result<SessionId> {
264 let data = serde_json::to_vec(session)?;
265 let project_dir = session.project_dir.to_string_lossy().to_string();
266 let created_at = session.created_at.to_rfc3339();
267 let updated_at = session.updated_at.to_rfc3339();
268
269 self.db.execute(
270 r#"INSERT INTO sessions (id, project_dir, compressed_summary, created_at, updated_at, data)
271 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
272 ON CONFLICT(id) DO UPDATE SET
273 project_dir = excluded.project_dir,
274 compressed_summary = excluded.compressed_summary,
275 created_at = excluded.created_at,
276 updated_at = excluded.updated_at,
277 data = excluded.data"#,
278 params![
279 session.id,
280 project_dir,
281 session.compressed_summary,
282 created_at,
283 updated_at,
284 data,
285 ],
286 )?;
287
288 Ok(session.id.clone())
289 }
290
291 pub fn save_dashboard_session(&self, id: &str, summary: &str) -> Result<()> {
293 let now = Utc::now().to_rfc3339();
294 let project_dir = std::env::current_dir()
295 .unwrap_or_else(|_| std::path::PathBuf::from("."))
296 .to_string_lossy()
297 .to_string();
298 let empty_data = serde_json::to_vec(&serde_json::json!({"conversation":[]}))?;
299
300 self.db
301 .execute(
302 "INSERT INTO sessions (id, project_dir, compressed_summary, created_at, updated_at, data) \
303 VALUES (?1, ?2, ?3, ?4, ?5, ?6) \
304 ON CONFLICT(id) DO UPDATE SET \
305 compressed_summary = excluded.compressed_summary, \
306 updated_at = excluded.updated_at",
307 params![id, project_dir, summary, now, now, empty_data],
308 )
309 .map_err(SqzError::SessionStore)?;
310 Ok(())
311 }
312
313 pub fn load_session(&self, id: SessionId) -> Result<SessionState> {
315 let data: Vec<u8> = self.db.query_row(
316 "SELECT data FROM sessions WHERE id = ?1",
317 params![id],
318 |row| row.get(0),
319 )?;
320 let session: SessionState = serde_json::from_slice(&data)?;
321 Ok(session)
322 }
323
324 pub fn search(&self, query: &str) -> Result<Vec<SessionSummary>> {
328 let mut stmt = self.db.prepare(
329 r#"SELECT s.id, s.project_dir, s.compressed_summary, s.created_at, s.updated_at
330 FROM sessions s
331 JOIN sessions_fts f ON s.rowid = f.rowid
332 WHERE sessions_fts MATCH ?1
333 ORDER BY rank"#,
334 )?;
335
336 let rows = stmt.query_map(params![query], |row| {
337 Ok((
338 row.get::<_, String>(0)?,
339 row.get::<_, String>(1)?,
340 row.get::<_, String>(2)?,
341 row.get::<_, String>(3)?,
342 row.get::<_, String>(4)?,
343 ))
344 })?;
345
346 let mut results = Vec::new();
347 for row in rows {
348 let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
349 results.push(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?);
350 }
351 Ok(results)
352 }
353
354 pub fn search_by_date(
356 &self,
357 from: DateTime<Utc>,
358 to: DateTime<Utc>,
359 ) -> Result<Vec<SessionSummary>> {
360 let mut stmt = self.db.prepare(
361 r#"SELECT id, project_dir, compressed_summary, created_at, updated_at
362 FROM sessions
363 WHERE updated_at >= ?1 AND updated_at <= ?2
364 ORDER BY updated_at DESC"#,
365 )?;
366
367 let rows = stmt.query_map(params![from.to_rfc3339(), to.to_rfc3339()], |row| {
368 Ok((
369 row.get::<_, String>(0)?,
370 row.get::<_, String>(1)?,
371 row.get::<_, String>(2)?,
372 row.get::<_, String>(3)?,
373 row.get::<_, String>(4)?,
374 ))
375 })?;
376
377 let mut results = Vec::new();
378 for row in rows {
379 let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
380 results.push(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?);
381 }
382 Ok(results)
383 }
384
385 pub fn latest_session(&self) -> Result<Option<SessionSummary>> {
387 let mut stmt = self.db.prepare(
388 r#"SELECT id, project_dir, compressed_summary, created_at, updated_at
389 FROM sessions
390 ORDER BY updated_at DESC
391 LIMIT 1"#,
392 ).map_err(SqzError::SessionStore)?;
393
394 let rows = stmt.query_map([], |row| {
395 Ok((
396 row.get::<_, String>(0)?,
397 row.get::<_, String>(1)?,
398 row.get::<_, String>(2)?,
399 row.get::<_, String>(3)?,
400 row.get::<_, String>(4)?,
401 ))
402 }).map_err(SqzError::SessionStore)?;
403
404 for row in rows {
405 let (id, project_dir, compressed_summary, created_at, updated_at) =
406 row.map_err(SqzError::SessionStore)?;
407 return Ok(Some(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?));
408 }
409 Ok(None)
410 }
411
412 pub fn search_by_project(&self, dir: &Path) -> Result<Vec<SessionSummary>> {
414 let dir_str = dir.to_string_lossy().to_string();
415 let mut stmt = self.db.prepare(
416 r#"SELECT id, project_dir, compressed_summary, created_at, updated_at
417 FROM sessions
418 WHERE project_dir = ?1
419 ORDER BY updated_at DESC"#,
420 )?;
421
422 let rows = stmt.query_map(params![dir_str], |row| {
423 Ok((
424 row.get::<_, String>(0)?,
425 row.get::<_, String>(1)?,
426 row.get::<_, String>(2)?,
427 row.get::<_, String>(3)?,
428 row.get::<_, String>(4)?,
429 ))
430 })?;
431
432 let mut results = Vec::new();
433 for row in rows {
434 let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
435 results.push(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?);
436 }
437 Ok(results)
438 }
439
440 pub fn save_cache_entry(&self, hash: &str, compressed: &CompressedContent) -> Result<()> {
453 self.save_cache_entry_with_original(hash, compressed, None)
454 }
455
456 pub fn save_cache_entry_with_original(
463 &self,
464 hash: &str,
465 compressed: &CompressedContent,
466 original: Option<&[u8]>,
467 ) -> Result<()> {
468 let data = serde_json::to_string(compressed)?;
469 let now = Utc::now().to_rfc3339();
470 self.db.execute(
471 r#"INSERT INTO cache_entries (hash, data, accessed_at, original)
472 VALUES (?1, ?2, ?3, ?4)
473 ON CONFLICT(hash) DO UPDATE
474 SET data = excluded.data,
475 accessed_at = excluded.accessed_at,
476 -- Don't overwrite a previously-stored `original`
477 -- with NULL. Older callers (that go through
478 -- save_cache_entry rather than the _with_original
479 -- variant) shouldn't erase the expand-able bytes.
480 original = COALESCE(excluded.original, original)"#,
481 params![hash, data, now, original],
482 )?;
483 Ok(())
484 }
485
486 pub fn get_cache_entry_original(&self, hash: &str) -> Result<Option<Vec<u8>>> {
495 let result: rusqlite::Result<Option<Vec<u8>>> = self.db.query_row(
496 "SELECT original FROM cache_entries WHERE hash = ?1",
497 params![hash],
498 |row| row.get(0),
499 );
500 match result {
501 Ok(v) => Ok(v),
502 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
503 Err(e) => Err(SqzError::SessionStore(e)),
504 }
505 }
506
507 pub fn delete_cache_entry(&self, hash: &str) -> Result<()> {
509 self.db.execute(
510 "DELETE FROM cache_entries WHERE hash = ?1",
511 params![hash],
512 )?;
513 Ok(())
514 }
515
516 pub fn list_cache_entries_lru(&self) -> Result<Vec<(String, u64)>> {
520 let mut stmt = self.db.prepare(
521 "SELECT hash, length(data) FROM cache_entries ORDER BY accessed_at ASC",
522 )?;
523 let rows = stmt.query_map([], |row| {
524 Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
525 })?;
526 let mut entries = Vec::new();
527 for row in rows {
528 let (hash, size) = row?;
529 entries.push((hash, size as u64));
530 }
531 Ok(entries)
532 }
533
534 pub fn get_cache_entry(&self, hash: &str) -> Result<Option<CompressedContent>> {
536 let result: rusqlite::Result<String> = self.db.query_row(
537 "SELECT data FROM cache_entries WHERE hash = ?1",
538 params![hash],
539 |row| row.get(0),
540 );
541
542 match result {
543 Ok(data) => {
544 let now = Utc::now().to_rfc3339();
546 let _ = self.db.execute(
547 "UPDATE cache_entries SET accessed_at = ?1 WHERE hash = ?2",
548 params![now, hash],
549 );
550 let entry: CompressedContent = serde_json::from_str(&data)?;
551 Ok(Some(entry))
552 }
553 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
554 Err(e) => Err(SqzError::SessionStore(e)),
555 }
556 }
557
558 pub fn get_cache_entry_by_prefix(
578 &self,
579 prefix: &str,
580 ) -> Result<Option<(String, CompressedContent)>> {
581 if prefix.is_empty() || !prefix.chars().all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()) {
586 return Ok(None);
587 }
588 let pattern = format!("{prefix}%");
589 let mut stmt = self
590 .db
591 .prepare("SELECT hash, data FROM cache_entries WHERE hash LIKE ?1 LIMIT 2")?;
592 let mut rows = stmt.query(params![pattern])?;
593
594 let first = match rows.next()? {
595 Some(r) => {
596 let hash: String = r.get(0)?;
597 let data: String = r.get(1)?;
598 (hash, data)
599 }
600 None => return Ok(None),
601 };
602
603 if rows.next()?.is_some() {
607 return Err(SqzError::Other(format!(
608 "cache: prefix '{prefix}' matches multiple entries — use a longer prefix"
609 )));
610 }
611 drop(rows);
612 drop(stmt);
613
614 let now = Utc::now().to_rfc3339();
616 let _ = self.db.execute(
617 "UPDATE cache_entries SET accessed_at = ?1 WHERE hash = ?2",
618 params![now, first.0],
619 );
620 let entry: CompressedContent = serde_json::from_str(&first.1)?;
621 Ok(Some((first.0, entry)))
622 }
623
624 pub fn get_cache_entry_accessed_at(&self, hash: &str) -> Result<Option<DateTime<Utc>>> {
631 let result: rusqlite::Result<String> = self.db.query_row(
632 "SELECT accessed_at FROM cache_entries WHERE hash = ?1",
633 params![hash],
634 |row| row.get(0),
635 );
636 match result {
637 Ok(s) => {
638 let ts = s
639 .parse::<DateTime<Utc>>()
640 .map_err(|e| SqzError::Other(format!("invalid accessed_at: {e}")))?;
641 Ok(Some(ts))
642 }
643 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
644 Err(e) => Err(SqzError::SessionStore(e)),
645 }
646 }
647
648 pub fn cache_entry_exists(&self, hash: &str) -> Result<bool> {
650 let result: rusqlite::Result<i64> = self.db.query_row(
651 "SELECT 1 FROM cache_entries WHERE hash = ?1",
652 params![hash],
653 |row| row.get(0),
654 );
655 match result {
656 Ok(_) => Ok(true),
657 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(false),
658 Err(e) => Err(SqzError::SessionStore(e)),
659 }
660 }
661
662 pub fn touch_cache_entry(&self, hash: &str) -> Result<()> {
666 let now = Utc::now().to_rfc3339();
667 self.db.execute(
668 "UPDATE cache_entries SET accessed_at = ?1 WHERE hash = ?2",
669 params![now, hash],
670 )?;
671 Ok(())
672 }
673
674 pub fn resolve_cache_hash_by_prefix(&self, prefix: &str) -> Result<Option<String>> {
680 if prefix.is_empty()
681 || !prefix
682 .chars()
683 .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
684 {
685 return Ok(None);
686 }
687 let pattern = format!("{prefix}%");
688 let mut stmt = self
689 .db
690 .prepare("SELECT hash FROM cache_entries WHERE hash LIKE ?1 LIMIT 2")?;
691 let mut rows = stmt.query(params![pattern])?;
692 let first: Option<String> = match rows.next()? {
693 Some(r) => Some(r.get(0)?),
694 None => None,
695 };
696 if rows.next()?.is_some() {
697 return Err(SqzError::Other(format!(
698 "cache: prefix '{prefix}' matches multiple entries — use a longer prefix"
699 )));
700 }
701 Ok(first)
702 }
703
704 pub fn pin_cache_entry(&self, hash: &str) -> Result<bool> {
708 let n = self.db.execute(
709 "UPDATE cache_entries SET pinned = 1 WHERE hash = ?1",
710 params![hash],
711 )?;
712 Ok(n > 0)
713 }
714
715 pub fn unpin_cache_entry(&self, hash: &str) -> Result<bool> {
718 let n = self.db.execute(
719 "UPDATE cache_entries SET pinned = 0 WHERE hash = ?1",
720 params![hash],
721 )?;
722 Ok(n > 0)
723 }
724
725 pub fn unpin_all_cache_entries(&self) -> Result<usize> {
728 let n = self
729 .db
730 .execute("UPDATE cache_entries SET pinned = 0 WHERE pinned = 1", [])?;
731 Ok(n)
732 }
733
734 pub fn is_cache_entry_pinned(&self, hash: &str) -> Result<bool> {
737 let result: rusqlite::Result<i64> = self.db.query_row(
738 "SELECT pinned FROM cache_entries WHERE hash = ?1",
739 params![hash],
740 |row| row.get(0),
741 );
742 match result {
743 Ok(v) => Ok(v != 0),
744 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(false),
745 Err(e) => Err(SqzError::SessionStore(e)),
746 }
747 }
748
749 pub fn list_pinned_cache_entries(
752 &self,
753 ) -> Result<Vec<(String, DateTime<Utc>, u64)>> {
754 let mut stmt = self
755 .db
756 .prepare(
757 "SELECT hash, accessed_at, length(data) FROM cache_entries \
758 WHERE pinned = 1 ORDER BY accessed_at DESC",
759 )
760 .map_err(SqzError::SessionStore)?;
761 let rows = stmt
762 .query_map([], |row| {
763 Ok((
764 row.get::<_, String>(0)?,
765 row.get::<_, String>(1)?,
766 row.get::<_, i64>(2)? as u64,
767 ))
768 })
769 .map_err(SqzError::SessionStore)?;
770 let mut out = Vec::new();
771 for row in rows {
772 let (hash, accessed_at_str, size) = row?;
773 let accessed_at = accessed_at_str.parse::<DateTime<Utc>>().map_err(|e| {
774 SqzError::Other(format!("invalid accessed_at on pinned entry: {e}"))
775 })?;
776 out.push((hash, accessed_at, size));
777 }
778 Ok(out)
779 }
780
781 pub fn set_metadata(&self, key: &str, value: &str) -> Result<()> {
784 self.db.execute(
785 "INSERT INTO metadata (key, value) VALUES (?1, ?2)
786 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
787 params![key, value],
788 )?;
789 Ok(())
790 }
791
792 pub fn get_metadata(&self, key: &str) -> Result<Option<String>> {
794 let result: rusqlite::Result<String> = self.db.query_row(
795 "SELECT value FROM metadata WHERE key = ?1",
796 params![key],
797 |row| row.get(0),
798 );
799 match result {
800 Ok(v) => Ok(Some(v)),
801 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
802 Err(e) => Err(SqzError::SessionStore(e)),
803 }
804 }
805
806 pub fn log_compression(
808 &self,
809 tokens_original: u32,
810 tokens_compressed: u32,
811 stages: &[String],
812 mode: &str,
813 ) -> Result<()> {
814 self.log_compression_with_project(tokens_original, tokens_compressed, stages, mode, None)
815 }
816
817 pub fn log_compression_with_project(
819 &self,
820 tokens_original: u32,
821 tokens_compressed: u32,
822 stages: &[String],
823 mode: &str,
824 project_dir: Option<&str>,
825 ) -> Result<()> {
826 let now = Utc::now().to_rfc3339();
827 let stages_str = stages.join(",");
828 self.db.execute(
829 "INSERT INTO compression_log (tokens_original, tokens_compressed, stages_applied, mode, created_at, project_dir) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
830 params![tokens_original, tokens_compressed, stages_str, mode, now, project_dir],
831 ).map_err(SqzError::SessionStore)?;
832 Ok(())
833 }
834
835 pub fn log_regret(&self, kind: &str, cmd: &str, hash: &str, gap_secs: f64) -> Result<()> {
840 let now = Utc::now().to_rfc3339();
841 self.db.execute(
842 "INSERT INTO regret_log (kind, cmd, hash, gap_secs, created_at) VALUES (?1, ?2, ?3, ?4, ?5)",
843 params![kind, cmd, hash, gap_secs, now],
844 ).map_err(SqzError::SessionStore)?;
845 Ok(())
846 }
847
848 pub fn regret_stats(&self) -> Result<RegretStats> {
850 let mut stmt = self.db.prepare(
851 "SELECT \
852 COALESCE(SUM(CASE WHEN kind = 'rerun' THEN 1 ELSE 0 END), 0), \
853 COALESCE(SUM(CASE WHEN kind = 'expand' THEN 1 ELSE 0 END), 0), \
854 COALESCE(AVG(CASE WHEN kind = 'rerun' THEN gap_secs END), 0.0) \
855 FROM regret_log",
856 ).map_err(SqzError::SessionStore)?;
857 let stats = stmt.query_row([], |row| {
858 Ok(RegretStats {
859 reruns: row.get::<_, u32>(0)?,
860 expands: row.get::<_, u32>(1)?,
861 avg_rerun_gap_secs: row.get::<_, f64>(2)?,
862 })
863 }).map_err(SqzError::SessionStore)?;
864 Ok(stats)
865 }
866
867 pub fn regret_by_command(&self, limit: u32) -> Result<Vec<(String, u32)>> {
869 let mut stmt = self.db.prepare(
870 "SELECT cmd, COUNT(*) FROM regret_log \
871 WHERE kind = 'rerun' AND cmd != '' \
872 GROUP BY cmd ORDER BY COUNT(*) DESC LIMIT ?1",
873 ).map_err(SqzError::SessionStore)?;
874 let rows = stmt.query_map(params![limit], |row| {
875 Ok((row.get::<_, String>(0)?, row.get::<_, u32>(1)?))
876 }).map_err(SqzError::SessionStore)?;
877 let mut results = Vec::new();
878 for row in rows {
879 results.push(row.map_err(SqzError::SessionStore)?);
880 }
881 Ok(results)
882 }
883
884 pub fn compression_stats(&self) -> Result<CompressionStats> {
886 let mut stmt = self.db.prepare(
887 "SELECT COUNT(*), COALESCE(SUM(tokens_original), 0), COALESCE(SUM(tokens_compressed), 0) FROM compression_log",
888 ).map_err(SqzError::SessionStore)?;
889
890 let stats = stmt.query_row([], |row| {
891 Ok(CompressionStats {
892 total_compressions: row.get::<_, u32>(0)?,
893 total_tokens_in: row.get::<_, u64>(1)?,
894 total_tokens_out: row.get::<_, u64>(2)?,
895 })
896 }).map_err(SqzError::SessionStore)?;
897
898 Ok(stats)
899 }
900
901 pub fn daily_gains(&self, days: u32) -> Result<Vec<DailyGain>> {
903 let mut stmt = self.db.prepare(
904 "SELECT date(created_at) as d, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
905 FROM compression_log \
906 WHERE created_at >= date('now', ?1) \
907 GROUP BY d ORDER BY d",
908 ).map_err(SqzError::SessionStore)?;
909
910 let offset = format!("-{days} days");
911 let rows = stmt.query_map(params![offset], |row| {
912 let tokens_in: u64 = row.get(2)?;
913 let tokens_out: u64 = row.get(3)?;
914 Ok(DailyGain {
915 date: row.get(0)?,
916 compressions: row.get(1)?,
917 tokens_in,
918 tokens_saved: tokens_in.saturating_sub(tokens_out),
919 })
920 }).map_err(SqzError::SessionStore)?;
921
922 let mut gains = Vec::new();
923 for row in rows {
924 gains.push(row.map_err(SqzError::SessionStore)?);
925 }
926 Ok(gains)
927 }
928
929 pub fn compression_stats_for_project(&self, project_dir: &str) -> Result<CompressionStats> {
931 let mut stmt = self.db.prepare(
932 "SELECT COUNT(*), COALESCE(SUM(tokens_original), 0), COALESCE(SUM(tokens_compressed), 0) \
933 FROM compression_log WHERE project_dir = ?1",
934 ).map_err(SqzError::SessionStore)?;
935
936 let stats = stmt.query_row(params![project_dir], |row| {
937 Ok(CompressionStats {
938 total_compressions: row.get::<_, u32>(0)?,
939 total_tokens_in: row.get::<_, u64>(1)?,
940 total_tokens_out: row.get::<_, u64>(2)?,
941 })
942 }).map_err(SqzError::SessionStore)?;
943
944 Ok(stats)
945 }
946
947 pub fn daily_gains_for_project(&self, days: u32, project_dir: &str) -> Result<Vec<DailyGain>> {
949 let mut stmt = self.db.prepare(
950 "SELECT date(created_at) as d, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
951 FROM compression_log \
952 WHERE created_at >= date('now', ?1) AND project_dir = ?2 \
953 GROUP BY d ORDER BY d",
954 ).map_err(SqzError::SessionStore)?;
955
956 let offset = format!("-{days} days");
957 let rows = stmt.query_map(params![offset, project_dir], |row| {
958 let tokens_in: u64 = row.get(2)?;
959 let tokens_out: u64 = row.get(3)?;
960 Ok(DailyGain {
961 date: row.get(0)?,
962 compressions: row.get(1)?,
963 tokens_in,
964 tokens_saved: tokens_in.saturating_sub(tokens_out),
965 })
966 }).map_err(SqzError::SessionStore)?;
967
968 let mut gains = Vec::new();
969 for row in rows {
970 gains.push(row.map_err(SqzError::SessionStore)?);
971 }
972 Ok(gains)
973 }
974
975 pub fn list_projects(&self) -> Result<Vec<(String, u32, u64)>> {
977 let mut stmt = self.db.prepare(
978 "SELECT project_dir, COUNT(*), COALESCE(SUM(tokens_original) - SUM(tokens_compressed), 0) \
979 FROM compression_log \
980 WHERE project_dir IS NOT NULL \
981 GROUP BY project_dir \
982 ORDER BY COUNT(*) DESC",
983 ).map_err(SqzError::SessionStore)?;
984
985 let rows = stmt.query_map([], |row| {
986 Ok((
987 row.get::<_, String>(0)?,
988 row.get::<_, u32>(1)?,
989 row.get::<_, u64>(2)?,
990 ))
991 }).map_err(SqzError::SessionStore)?;
992
993 let mut projects = Vec::new();
994 for row in rows {
995 projects.push(row.map_err(SqzError::SessionStore)?);
996 }
997 Ok(projects)
998 }
999
1000 pub fn command_breakdown(&self, limit: u32) -> Result<Vec<CommandStats>> {
1004 let mut stmt = self.db.prepare(
1005 "SELECT mode, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
1006 FROM compression_log \
1007 GROUP BY mode \
1008 ORDER BY SUM(tokens_original) DESC \
1009 LIMIT ?1",
1010 ).map_err(SqzError::SessionStore)?;
1011
1012 let rows = stmt.query_map(params![limit], |row| {
1013 let tokens_in: u64 = row.get(2)?;
1014 let tokens_out: u64 = row.get(3)?;
1015 Ok(CommandStats {
1016 command: row.get(0)?,
1017 invocations: row.get(1)?,
1018 tokens_in,
1019 tokens_out,
1020 tokens_saved: tokens_in.saturating_sub(tokens_out),
1021 })
1022 }).map_err(SqzError::SessionStore)?;
1023
1024 let mut results = Vec::new();
1025 for row in rows {
1026 results.push(row.map_err(SqzError::SessionStore)?);
1027 }
1028 Ok(results)
1029 }
1030
1031 pub fn command_breakdown_for_project(&self, limit: u32, project_dir: &str) -> Result<Vec<CommandStats>> {
1033 let mut stmt = self.db.prepare(
1034 "SELECT mode, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
1035 FROM compression_log \
1036 WHERE project_dir = ?1 \
1037 GROUP BY mode \
1038 ORDER BY SUM(tokens_original) DESC \
1039 LIMIT ?2",
1040 ).map_err(SqzError::SessionStore)?;
1041
1042 let rows = stmt.query_map(params![project_dir, limit], |row| {
1043 let tokens_in: u64 = row.get(2)?;
1044 let tokens_out: u64 = row.get(3)?;
1045 Ok(CommandStats {
1046 command: row.get(0)?,
1047 invocations: row.get(1)?,
1048 tokens_in,
1049 tokens_out,
1050 tokens_saved: tokens_in.saturating_sub(tokens_out),
1051 })
1052 }).map_err(SqzError::SessionStore)?;
1053
1054 let mut results = Vec::new();
1055 for row in rows {
1056 results.push(row.map_err(SqzError::SessionStore)?);
1057 }
1058 Ok(results)
1059 }
1060
1061 pub fn session_pressure(&self, since_minutes: u32) -> Result<u64> {
1065 let mut stmt = self.db.prepare(
1066 "SELECT COALESCE(SUM(tokens_compressed), 0) \
1067 FROM compression_log \
1068 WHERE created_at >= datetime('now', ?1)",
1069 ).map_err(SqzError::SessionStore)?;
1070
1071 let offset = format!("-{since_minutes} minutes");
1072 let total: u64 = stmt.query_row(params![offset], |row| {
1073 row.get(0)
1074 }).map_err(SqzError::SessionStore)?;
1075
1076 Ok(total)
1077 }
1078
1079 pub fn add_known_file(&self, path: &str) -> Result<()> {
1084 let now = Utc::now().to_rfc3339();
1085 self.db.execute(
1086 "INSERT OR REPLACE INTO known_files (path, added_at) VALUES (?1, ?2)",
1087 params![path, now],
1088 ).map_err(SqzError::SessionStore)?;
1089 Ok(())
1090 }
1091
1092 pub fn known_files(&self) -> Result<Vec<String>> {
1094 let mut stmt = self.db.prepare(
1095 "SELECT path FROM known_files ORDER BY added_at DESC",
1096 ).map_err(SqzError::SessionStore)?;
1097
1098 let rows = stmt.query_map([], |row| {
1099 row.get::<_, String>(0)
1100 }).map_err(SqzError::SessionStore)?;
1101
1102 let mut files = Vec::new();
1103 for row in rows {
1104 files.push(row.map_err(SqzError::SessionStore)?);
1105 }
1106 Ok(files)
1107 }
1108 pub fn clear_known_files(&self) -> Result<()> {
1110 self.db
1111 .execute("DELETE FROM known_files", [])
1112 .map_err(SqzError::SessionStore)?;
1113 Ok(())
1114 }
1115
1116 pub fn cache_entry_count(&self) -> Result<u64> {
1118 self.db.query_row(
1119 "SELECT COUNT(*) FROM cache_entries", [],
1120 |row| row.get(0),
1121 ).map_err(SqzError::SessionStore)
1122 }
1123
1124 pub fn record_cache_hit(&self) -> Result<()> {
1126 let count: u64 = self.db.query_row(
1127 "SELECT value FROM metadata WHERE key = 'cache_hits'", [],
1128 |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
1129 ).unwrap_or(0);
1130 let new = count + 1;
1131 self.db.execute(
1132 "INSERT OR REPLACE INTO metadata (key, value) VALUES ('cache_hits', ?1)",
1133 params![new.to_string()],
1134 )?;
1135 Ok(())
1136 }
1137
1138 pub fn record_cache_miss(&self) -> Result<()> {
1140 let count: u64 = self.db.query_row(
1141 "SELECT value FROM metadata WHERE key = 'cache_misses'", [],
1142 |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
1143 ).unwrap_or(0);
1144 let new = count + 1;
1145 self.db.execute(
1146 "INSERT OR REPLACE INTO metadata (key, value) VALUES ('cache_misses', ?1)",
1147 params![new.to_string()],
1148 )?;
1149 Ok(())
1150 }
1151
1152 pub fn cache_hit_miss_counts(&self) -> Result<(u64, u64)> {
1154 let hits: u64 = self.db.query_row(
1155 "SELECT value FROM metadata WHERE key = 'cache_hits'", [],
1156 |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
1157 ).unwrap_or(0);
1158 let misses: u64 = self.db.query_row(
1159 "SELECT value FROM metadata WHERE key = 'cache_misses'", [],
1160 |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
1161 ).unwrap_or(0);
1162 Ok((hits, misses))
1163 }
1164
1165 pub fn clear_cache(&self) -> Result<u64> {
1169 let count: u64 = self.db.query_row(
1170 "SELECT COUNT(*) FROM cache_entries", [],
1171 |row| row.get(0),
1172 ).unwrap_or(0);
1173 self.db.execute("DELETE FROM cache_entries", [])
1174 .map_err(SqzError::SessionStore)?;
1175 let _ = self.db.execute(
1178 "DELETE FROM metadata WHERE key = 'last_compaction_at'", [],
1179 );
1180 Ok(count)
1181 }
1182
1183 pub fn clear_stats(&self) -> Result<u64> {
1187 let count: u64 = self.db.query_row(
1188 "SELECT COUNT(*) FROM compression_log", [],
1189 |row| row.get(0),
1190 ).unwrap_or(0);
1191 self.db.execute("DELETE FROM compression_log", [])
1192 .map_err(SqzError::SessionStore)?;
1193 Ok(count)
1194 }
1195
1196 pub fn clear_stats_for_project(&self, project_dir: &str) -> Result<u64> {
1198 let count: u64 = self.db.query_row(
1199 "SELECT COUNT(*) FROM compression_log WHERE project_dir = ?1",
1200 rusqlite::params![project_dir],
1201 |row| row.get(0),
1202 ).unwrap_or(0);
1203 self.db.execute(
1204 "DELETE FROM compression_log WHERE project_dir = ?1",
1205 rusqlite::params![project_dir],
1206 ).map_err(SqzError::SessionStore)?;
1207 Ok(count)
1208 }
1209
1210 pub fn reset_all(&self) -> Result<()> {
1213 self.db.execute("DELETE FROM cache_entries", []).map_err(SqzError::SessionStore)?;
1214 self.db.execute("DELETE FROM compression_log", []).map_err(SqzError::SessionStore)?;
1215 self.db.execute("DELETE FROM sessions", []).map_err(SqzError::SessionStore)?;
1216 self.db.execute("DELETE FROM known_files", []).map_err(SqzError::SessionStore)?;
1217 self.db.execute("DELETE FROM metadata", []).map_err(SqzError::SessionStore)?;
1218 let _ = self.db.execute("VACUUM", []);
1220 Ok(())
1221 }
1222
1223 pub fn list_sessions(&self, limit: u32) -> Result<Vec<SessionSummary>> {
1225 let mut stmt = self
1226 .db
1227 .prepare(
1228 "SELECT id, project_dir, compressed_summary, created_at, updated_at \
1229 FROM sessions ORDER BY updated_at DESC LIMIT ?1",
1230 )
1231 .map_err(SqzError::SessionStore)?;
1232
1233 let rows = stmt
1234 .query_map(params![limit], |row| {
1235 Ok((
1236 row.get::<_, String>(0)?,
1237 row.get::<_, String>(1)?,
1238 row.get::<_, String>(2)?,
1239 row.get::<_, String>(3)?,
1240 row.get::<_, String>(4)?,
1241 ))
1242 })
1243 .map_err(SqzError::SessionStore)?;
1244
1245 let mut results = Vec::new();
1246 for row in rows {
1247 let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
1248 results.push(row_to_summary(
1249 id,
1250 project_dir,
1251 compressed_summary,
1252 created_at,
1253 updated_at,
1254 )?);
1255 }
1256 Ok(results)
1257 }
1258
1259 pub fn per_tool_breakdown(&self) -> Result<Vec<(String, u64, u64, u32)>> {
1261 let mut stmt = self
1262 .db
1263 .prepare(
1264 "SELECT stages_applied, SUM(tokens_original), SUM(tokens_compressed), COUNT(*) \
1265 FROM compression_log GROUP BY stages_applied ORDER BY COUNT(*) DESC LIMIT 20",
1266 )
1267 .map_err(SqzError::SessionStore)?;
1268
1269 let rows = stmt
1270 .query_map([], |row| {
1271 Ok((
1272 row.get::<_, String>(0)?,
1273 row.get::<_, u64>(1)?,
1274 row.get::<_, u64>(2)?,
1275 row.get::<_, u32>(3)?,
1276 ))
1277 })
1278 .map_err(SqzError::SessionStore)?;
1279
1280 let mut tool_map: std::collections::HashMap<String, (u64, u64, u32)> =
1281 std::collections::HashMap::new();
1282 for row in rows {
1283 let (stages, tokens_in, tokens_out, count) = row?;
1284 for stage in stages.split(',') {
1285 let stage = stage.trim();
1286 if stage.is_empty() {
1287 continue;
1288 }
1289 let entry = tool_map.entry(stage.to_string()).or_default();
1290 entry.0 += tokens_in;
1291 entry.1 += tokens_out;
1292 entry.2 += count;
1293 }
1294 }
1295
1296 let mut result: Vec<_> = tool_map.into_iter().collect();
1297 result.sort_by(|a, b| b.1 .2.cmp(&a.1 .2));
1298 Ok(result
1299 .into_iter()
1300 .map(|(name, (tin, tout, cnt))| (name, tin, tout, cnt))
1301 .collect())
1302 }
1303
1304 pub fn per_command_breakdown(&self) -> Result<Vec<(String, u64, u64, u32)>> {
1306 let mut stmt = self
1307 .db
1308 .prepare(
1309 "SELECT COALESCE(mode, 'auto'), SUM(tokens_original), SUM(tokens_compressed), COUNT(*) \
1310 FROM compression_log GROUP BY mode ORDER BY COUNT(*) DESC LIMIT 20",
1311 )
1312 .map_err(SqzError::SessionStore)?;
1313
1314 let rows = stmt
1315 .query_map([], |row| {
1316 Ok((
1317 row.get::<_, String>(0)?,
1318 row.get::<_, u64>(1)?,
1319 row.get::<_, u64>(2)?,
1320 row.get::<_, u32>(3)?,
1321 ))
1322 })
1323 .map_err(SqzError::SessionStore)?;
1324
1325 let mut result = Vec::new();
1326 for row in rows {
1327 result.push(row?);
1328 }
1329 Ok(result)
1330 }
1331}
1332
1333#[derive(Debug, Clone, Default)]
1335pub struct CompressionStats {
1336 pub total_compressions: u32,
1337 pub total_tokens_in: u64,
1338 pub total_tokens_out: u64,
1339}
1340
1341impl CompressionStats {
1342 pub fn tokens_saved(&self) -> u64 {
1343 self.total_tokens_in.saturating_sub(self.total_tokens_out)
1344 }
1345
1346 pub fn reduction_pct(&self) -> f64 {
1347 if self.total_tokens_in == 0 {
1348 0.0
1349 } else {
1350 (1.0 - self.total_tokens_out as f64 / self.total_tokens_in as f64) * 100.0
1351 }
1352 }
1353}
1354
1355#[derive(Debug, Clone)]
1357pub struct DailyGain {
1358 pub date: String,
1359 pub compressions: u32,
1360 pub tokens_saved: u64,
1361 pub tokens_in: u64,
1362}
1363
1364#[derive(Debug, Clone)]
1366pub struct CommandStats {
1367 pub command: String,
1368 pub invocations: u32,
1369 pub tokens_in: u64,
1370 pub tokens_out: u64,
1371 pub tokens_saved: u64,
1372}
1373
1374#[derive(Debug, Clone, Default)]
1376pub struct RegretStats {
1377 pub reruns: u32,
1379 pub expands: u32,
1381 pub avg_rerun_gap_secs: f64,
1383}
1384
1385impl CommandStats {
1386 pub fn reduction_pct(&self) -> f64 {
1387 if self.tokens_in == 0 {
1388 0.0
1389 } else {
1390 (1.0 - self.tokens_out as f64 / self.tokens_in as f64) * 100.0
1391 }
1392 }
1393}
1394
1395#[cfg(test)]
1398mod tests {
1399 use super::*;
1400 use crate::types::{BudgetState, CorrectionLog, ModelFamily, SessionState};
1401 use chrono::Utc;
1402 use proptest::prelude::*;
1403 use std::path::PathBuf;
1404
1405 fn make_session(id: &str, project_dir: &str, summary: &str) -> SessionState {
1406 let now = Utc::now();
1407 SessionState {
1408 id: id.to_string(),
1409 project_dir: PathBuf::from(project_dir),
1410 conversation: vec![],
1411 corrections: CorrectionLog::default(),
1412 pins: vec![],
1413 learnings: vec![],
1414 compressed_summary: summary.to_string(),
1415 budget: BudgetState {
1416 window_size: 200_000,
1417 consumed: 0,
1418 pinned: 0,
1419 model_family: ModelFamily::AnthropicClaude,
1420 },
1421 tool_usage: vec![],
1422 created_at: now,
1423 updated_at: now,
1424 }
1425 }
1426
1427 fn in_memory_store() -> SessionStore {
1428 let conn = Connection::open_in_memory().unwrap();
1429 apply_schema(&conn).unwrap();
1430 SessionStore { db: conn }
1431 }
1432
1433 #[test]
1434 fn regret_log_roundtrip() {
1435 let store = in_memory_store();
1436 let empty = store.regret_stats().unwrap();
1438 assert_eq!(empty.reruns, 0);
1439 assert_eq!(empty.expands, 0);
1440 assert_eq!(empty.avg_rerun_gap_secs, 0.0);
1441
1442 store.log_regret("rerun", "git log", "aaaa", 10.0).unwrap();
1443 store.log_regret("rerun", "git log", "bbbb", 30.0).unwrap();
1444 store.log_regret("rerun", "cargo test", "cccc", 20.0).unwrap();
1445 store.log_regret("expand", "", "dddd", 0.0).unwrap();
1446
1447 let stats = store.regret_stats().unwrap();
1448 assert_eq!(stats.reruns, 3);
1449 assert_eq!(stats.expands, 1);
1450 assert!((stats.avg_rerun_gap_secs - 20.0).abs() < 1e-9);
1451
1452 let top = store.regret_by_command(10).unwrap();
1453 assert_eq!(top[0], ("git log".to_string(), 2));
1454 assert_eq!(top[1], ("cargo test".to_string(), 1));
1455 assert_eq!(top.len(), 2);
1457 }
1458
1459 #[test]
1460 fn test_save_and_load_session() {
1461 let store = in_memory_store();
1462 let session = make_session("sess-1", "/home/user/project", "REST API refactor");
1463
1464 let id = store.save_session(&session).unwrap();
1465 assert_eq!(id, "sess-1");
1466
1467 let loaded = store.load_session("sess-1".to_string()).unwrap();
1468 assert_eq!(loaded.id, session.id);
1469 assert_eq!(loaded.compressed_summary, session.compressed_summary);
1470 assert_eq!(loaded.project_dir, session.project_dir);
1471 }
1472
1473 #[test]
1474 fn test_save_session_upsert() {
1475 let store = in_memory_store();
1476 let mut session = make_session("sess-2", "/proj", "initial summary");
1477 store.save_session(&session).unwrap();
1478
1479 session.compressed_summary = "updated summary".to_string();
1480 store.save_session(&session).unwrap();
1481
1482 let loaded = store.load_session("sess-2".to_string()).unwrap();
1483 assert_eq!(loaded.compressed_summary, "updated summary");
1484 }
1485
1486 #[test]
1487 fn test_load_nonexistent_session_errors() {
1488 let store = in_memory_store();
1489 let result = store.load_session("does-not-exist".to_string());
1490 assert!(result.is_err());
1491 }
1492
1493 #[test]
1494 fn test_search_fts() {
1495 let store = in_memory_store();
1496 store.save_session(&make_session("s1", "/proj", "REST API refactor with authentication")).unwrap();
1497 store.save_session(&make_session("s2", "/proj", "database migration postgres")).unwrap();
1498
1499 let results = store.search("authentication").unwrap();
1500 assert_eq!(results.len(), 1);
1501 assert_eq!(results[0].id, "s1");
1502 }
1503
1504 #[test]
1505 fn test_search_by_date() {
1506 let store = in_memory_store();
1507 let now = Utc::now();
1508 let past = now - chrono::Duration::hours(2);
1509 let future = now + chrono::Duration::hours(2);
1510
1511 store.save_session(&make_session("s1", "/proj", "recent session")).unwrap();
1512
1513 let results = store.search_by_date(past, future).unwrap();
1514 assert!(!results.is_empty());
1515 assert!(results.iter().any(|r| r.id == "s1"));
1516 }
1517
1518 #[test]
1519 fn test_search_by_project() {
1520 let store = in_memory_store();
1521 store.save_session(&make_session("s1", "/home/user/alpha", "alpha project")).unwrap();
1522 store.save_session(&make_session("s2", "/home/user/beta", "beta project")).unwrap();
1523
1524 let results = store.search_by_project(Path::new("/home/user/alpha")).unwrap();
1525 assert_eq!(results.len(), 1);
1526 assert_eq!(results[0].id, "s1");
1527 }
1528
1529 #[test]
1530 fn test_cache_entry_round_trip() {
1531 let store = in_memory_store();
1532 let entry = CompressedContent {
1533 data: "compressed data".to_string(),
1534 tokens_compressed: 10,
1535 tokens_original: 50,
1536 stages_applied: vec!["strip_nulls".to_string()],
1537 compression_ratio: 0.2,
1538 provenance: crate::types::Provenance::default(),
1539 verify: None,
1540 };
1541
1542 store.save_cache_entry("abc123", &entry).unwrap();
1543
1544 let loaded = store.get_cache_entry("abc123").unwrap().unwrap();
1545 assert_eq!(loaded.data, entry.data);
1546 assert_eq!(loaded.tokens_compressed, entry.tokens_compressed);
1547 assert_eq!(loaded.tokens_original, entry.tokens_original);
1548 }
1549
1550 #[test]
1551 fn test_get_cache_entry_missing_returns_none() {
1552 let store = in_memory_store();
1553 let result = store.get_cache_entry("nonexistent").unwrap();
1554 assert!(result.is_none());
1555 }
1556
1557 #[test]
1558 fn test_open_or_create_corrupted_db() {
1559 let dir = tempfile::tempdir().unwrap();
1560 let path = dir.path().join("store.db");
1561
1562 std::fs::write(&path, b"this is not a valid sqlite database").unwrap();
1564
1565 let store = SessionStore::open_or_create(&path).unwrap();
1567 let session = make_session("s1", "/proj", "after corruption");
1568 store.save_session(&session).unwrap();
1569 let loaded = store.load_session("s1".to_string()).unwrap();
1570 assert_eq!(loaded.id, "s1");
1571 }
1572
1573 fn make_session_at(id: &str, summary: &str, updated_at: DateTime<Utc>) -> SessionState {
1577 let now = Utc::now();
1578 SessionState {
1579 id: id.to_string(),
1580 project_dir: PathBuf::from("/proj"),
1581 conversation: vec![],
1582 corrections: CorrectionLog::default(),
1583 pins: vec![],
1584 learnings: vec![],
1585 compressed_summary: summary.to_string(),
1586 budget: BudgetState {
1587 window_size: 200_000,
1588 consumed: 0,
1589 pinned: 0,
1590 model_family: ModelFamily::AnthropicClaude,
1591 },
1592 tool_usage: vec![],
1593 created_at: now,
1594 updated_at,
1595 }
1596 }
1597
1598 proptest! {
1606 #[test]
1612 fn prop_search_correctness(
1613 keyword in "[b-df-hj-np-tv-z]{5,8}",
1616 matching_suffixes in proptest::collection::vec("[a-z ]{4,20}", 1..=6usize),
1618 non_matching in proptest::collection::vec("[a-z ]{8,30}", 1..=6usize),
1620 ) {
1621 for s in &non_matching {
1623 prop_assume!(!s.contains(keyword.as_str()));
1624 }
1625
1626 let store = in_memory_store();
1627
1628 let mut matching_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1630 for (i, suffix) in matching_suffixes.iter().enumerate() {
1631 let id = format!("match-{i}");
1632 let summary = format!("{} {} end", suffix, keyword);
1633 store.save_session(&make_session(&id, "/proj", &summary)).unwrap();
1634 matching_ids.insert(id);
1635 }
1636
1637 let mut non_matching_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1639 for (i, summary) in non_matching.iter().enumerate() {
1640 let id = format!("nomatch-{i}");
1641 store.save_session(&make_session(&id, "/proj", summary)).unwrap();
1642 non_matching_ids.insert(id);
1643 }
1644
1645 let results = store.search(&keyword).unwrap();
1646 let result_ids: std::collections::HashSet<String> =
1647 results.iter().map(|r| r.id.clone()).collect();
1648
1649 for id in &matching_ids {
1651 prop_assert!(
1652 result_ids.contains(id),
1653 "matching session '{}' not found in search results for keyword '{}'",
1654 id, keyword
1655 );
1656 }
1657
1658 for id in &non_matching_ids {
1660 prop_assert!(
1661 !result_ids.contains(id),
1662 "non-matching session '{}' incorrectly appeared in search results for keyword '{}'",
1663 id, keyword
1664 );
1665 }
1666 }
1667 }
1668
1669 proptest! {
1677 #[test]
1683 fn prop_search_by_date_correctness(
1684 offsets in proptest::collection::vec(0i64..=86400i64 * 365, 2..=8usize),
1686 window_start_delta in 0i64..=3600i64,
1688 window_end_delta in 3600i64..=7200i64,
1689 ) {
1690 use chrono::TimeZone;
1691
1692 let mut unique_offsets: Vec<i64> = offsets.clone();
1694 unique_offsets.sort_unstable();
1695 unique_offsets.dedup();
1696 prop_assume!(unique_offsets.len() >= 2);
1697
1698 let base_offset = unique_offsets[0];
1699 let from_offset = base_offset + window_start_delta;
1700 let to_offset = base_offset + window_end_delta;
1701
1702 let from = Utc.timestamp_opt(from_offset, 0).unwrap();
1703 let to = Utc.timestamp_opt(to_offset, 0).unwrap();
1704
1705 let store = in_memory_store();
1706
1707 let mut in_range_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1708 let mut out_range_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1709
1710 for (i, &offset) in unique_offsets.iter().enumerate() {
1711 let ts = Utc.timestamp_opt(offset, 0).unwrap();
1712 let id = format!("sess-{i}");
1713 let session = make_session_at(&id, "some summary", ts);
1714 store.save_session(&session).unwrap();
1715
1716 if ts >= from && ts <= to {
1717 in_range_ids.insert(id);
1718 } else {
1719 out_range_ids.insert(id);
1720 }
1721 }
1722
1723 let results = store.search_by_date(from, to).unwrap();
1724 let result_ids: std::collections::HashSet<String> =
1725 results.iter().map(|r| r.id.clone()).collect();
1726
1727 for id in &in_range_ids {
1729 prop_assert!(
1730 result_ids.contains(id),
1731 "in-range session '{}' missing from search_by_date results",
1732 id
1733 );
1734 }
1735
1736 for id in &out_range_ids {
1738 prop_assert!(
1739 !result_ids.contains(id),
1740 "out-of-range session '{}' incorrectly appeared in search_by_date results",
1741 id
1742 );
1743 }
1744 }
1745 }
1746}