Skip to main content

sqz_engine/
session_store.rs

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/// A lightweight summary of a session for search results. Doesn't include
12/// the full conversation — just enough to identify and filter sessions.
13#[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
22/// SQLite-backed persistent session and cache store with FTS5 full-text search.
23///
24/// Stores sessions, cache entries, compression logs, and known files in a
25/// single SQLite database. Uses WAL mode for concurrent read access.
26///
27/// ```rust,no_run
28/// use sqz_engine::SessionStore;
29/// use std::path::Path;
30///
31/// let store = SessionStore::open_or_create(Path::new("~/.sqz/sessions.db")).unwrap();
32/// let results = store.search("authentication refactor").unwrap();
33/// for session in &results {
34///     println!("{}: {}", session.id, session.compressed_summary);
35/// }
36/// ```
37pub struct SessionStore {
38    db: Connection,
39}
40
41// ── Schema ────────────────────────────────────────────────────────────────────
42
43const 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);
92
93CREATE TABLE IF NOT EXISTS compression_log (
94    id               INTEGER PRIMARY KEY AUTOINCREMENT,
95    tokens_original  INTEGER NOT NULL,
96    tokens_compressed INTEGER NOT NULL,
97    stages_applied   TEXT NOT NULL,
98    mode             TEXT NOT NULL DEFAULT 'auto',
99    created_at       TEXT NOT NULL
100);
101
102CREATE TABLE IF NOT EXISTS known_files (
103    path        TEXT PRIMARY KEY,
104    added_at    TEXT NOT NULL
105);
106
107-- Small key/value store for engine-wide state that needs to persist across
108-- short-lived sqz processes (each shell-hook invocation is a new process).
109-- Initially used only for the last_compaction_at marker: cache entries with
110-- `accessed_at < last_compaction_at` are treated as stale even if still
111-- within the normal TTL. See cache_manager.rs for the freshness model.
112CREATE TABLE IF NOT EXISTS metadata (
113    key         TEXT PRIMARY KEY,
114    value       TEXT NOT NULL
115);
116"#;
117
118// ── Helpers ───────────────────────────────────────────────────────────────────
119
120pub(crate) fn apply_schema(conn: &Connection) -> rusqlite::Result<()> {
121    conn.execute_batch(SCHEMA)?;
122    // Additive migration: add `original` BLOB column to cache_entries if
123    // it does not yet exist. Stores the raw pre-compression bytes so
124    // `sqz expand <prefix>` can return truly uncompressed content when
125    // an agent cannot parse `§ref:…§` tokens (reported on-list by
126    // SquireNed for GLM 5.1). `NULL` for rows written by older sqz
127    // versions — `expand` treats those as "original unavailable, fall
128    // back to the compressed blob" so users don't get spurious errors
129    // on pre-migration data.
130    //
131    // Using pragma_table_info rather than a version table because the
132    // rest of sqz does the same — this is the first additive migration.
133    let has_original: bool = conn
134        .prepare("SELECT 1 FROM pragma_table_info('cache_entries') WHERE name = 'original'")?
135        .query_row([], |_| Ok(()))
136        .is_ok();
137    if !has_original {
138        conn.execute("ALTER TABLE cache_entries ADD COLUMN original BLOB", [])?;
139    }
140
141    // Additive migration: add `project_dir` TEXT column to compression_log.
142    // Enables per-project stats filtering (issue #13). NULL for rows
143    // written by older sqz versions — aggregate queries treat those as
144    // "unknown project".
145    let has_project_dir: bool = conn
146        .prepare("SELECT 1 FROM pragma_table_info('compression_log') WHERE name = 'project_dir'")?
147        .query_row([], |_| Ok(()))
148        .is_ok();
149    if !has_project_dir {
150        conn.execute("ALTER TABLE compression_log ADD COLUMN project_dir TEXT", [])?;
151    }
152
153    Ok(())
154}
155
156fn open_connection(path: &Path) -> rusqlite::Result<Connection> {
157    let conn = Connection::open(path)?;
158    apply_schema(&conn)?;
159    Ok(conn)
160}
161
162fn row_to_summary(
163    id: String,
164    project_dir: String,
165    compressed_summary: String,
166    created_at: String,
167    updated_at: String,
168) -> Result<SessionSummary> {
169    let created_at = created_at
170        .parse::<DateTime<Utc>>()
171        .map_err(|e| SqzError::Other(format!("invalid created_at timestamp: {e}")))?;
172    let updated_at = updated_at
173        .parse::<DateTime<Utc>>()
174        .map_err(|e| SqzError::Other(format!("invalid updated_at timestamp: {e}")))?;
175    Ok(SessionSummary {
176        id,
177        project_dir: PathBuf::from(project_dir),
178        compressed_summary,
179        created_at,
180        updated_at,
181    })
182}
183
184// ── SessionStore ──────────────────────────────────────────────────────────────
185
186impl SessionStore {
187    /// Construct a `SessionStore` from an already-open `Connection`.
188    /// Intended for testing (e.g., in-memory databases).
189    #[cfg(test)]
190    pub(crate) fn from_connection(conn: Connection) -> Self {
191        Self { db: conn }
192    }
193
194    /// Open an existing database at `path`. Returns an error if the file does
195    /// not exist or cannot be opened.
196    pub fn open(path: &Path) -> Result<Self> {
197        let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)?;
198        apply_schema(&conn)?;
199        Ok(Self { db: conn })
200    }
201
202    /// Open the database at `path`, creating it if it does not exist.
203    /// If the database is corrupted, a fresh database is created at the same
204    /// path and a warning is logged to stderr.
205    pub fn open_or_create(path: &Path) -> Result<Self> {
206        match open_connection(path) {
207            Ok(conn) => Ok(Self { db: conn }),
208            Err(e) => {
209                eprintln!(
210                    "sqz warning: session store at '{}' is corrupted or inaccessible ({e}). \
211                     Creating a new database. Prior session data has been lost.",
212                    path.display()
213                );
214                // Remove the corrupted file so we can start fresh.
215                if path.exists() {
216                    let _ = std::fs::remove_file(path);
217                }
218                let conn = open_connection(path)
219                    .map_err(|e2| SqzError::Other(format!("failed to create new session store: {e2}")))?;
220                Ok(Self { db: conn })
221            }
222        }
223    }
224
225    // ── Session CRUD ──────────────────────────────────────────────────────────
226
227    /// Persist a session. Returns the session id.
228    pub fn save_session(&self, session: &SessionState) -> Result<SessionId> {
229        let data = serde_json::to_vec(session)?;
230        let project_dir = session.project_dir.to_string_lossy().to_string();
231        let created_at = session.created_at.to_rfc3339();
232        let updated_at = session.updated_at.to_rfc3339();
233
234        self.db.execute(
235            r#"INSERT INTO sessions (id, project_dir, compressed_summary, created_at, updated_at, data)
236               VALUES (?1, ?2, ?3, ?4, ?5, ?6)
237               ON CONFLICT(id) DO UPDATE SET
238                   project_dir        = excluded.project_dir,
239                   compressed_summary = excluded.compressed_summary,
240                   created_at         = excluded.created_at,
241                   updated_at         = excluded.updated_at,
242                   data               = excluded.data"#,
243            params![
244                session.id,
245                project_dir,
246                session.compressed_summary,
247                created_at,
248                updated_at,
249                data,
250            ],
251        )?;
252
253        Ok(session.id.clone())
254    }
255
256    /// Save a lightweight session from just an id and summary for dashboard use.
257    pub fn save_dashboard_session(&self, id: &str, summary: &str) -> Result<()> {
258        let now = Utc::now().to_rfc3339();
259        let project_dir = std::env::current_dir()
260            .unwrap_or_else(|_| std::path::PathBuf::from("."))
261            .to_string_lossy()
262            .to_string();
263        let empty_data = serde_json::to_vec(&serde_json::json!({"conversation":[]}))?;
264
265        self.db
266            .execute(
267                "INSERT INTO sessions (id, project_dir, compressed_summary, created_at, updated_at, data) \
268                 VALUES (?1, ?2, ?3, ?4, ?5, ?6) \
269                 ON CONFLICT(id) DO UPDATE SET \
270                 compressed_summary = excluded.compressed_summary, \
271                 updated_at = excluded.updated_at",
272                params![id, project_dir, summary, now, now, empty_data],
273            )
274            .map_err(SqzError::SessionStore)?;
275        Ok(())
276    }
277
278    /// Load a session by id.
279    pub fn load_session(&self, id: SessionId) -> Result<SessionState> {
280        let data: Vec<u8> = self.db.query_row(
281            "SELECT data FROM sessions WHERE id = ?1",
282            params![id],
283            |row| row.get(0),
284        )?;
285        let session: SessionState = serde_json::from_slice(&data)?;
286        Ok(session)
287    }
288
289    // ── Search ────────────────────────────────────────────────────────────────
290
291    /// Full-text search using FTS5 (porter stemmer, ASCII tokenizer).
292    pub fn search(&self, query: &str) -> Result<Vec<SessionSummary>> {
293        let mut stmt = self.db.prepare(
294            r#"SELECT s.id, s.project_dir, s.compressed_summary, s.created_at, s.updated_at
295               FROM sessions s
296               JOIN sessions_fts f ON s.rowid = f.rowid
297               WHERE sessions_fts MATCH ?1
298               ORDER BY rank"#,
299        )?;
300
301        let rows = stmt.query_map(params![query], |row| {
302            Ok((
303                row.get::<_, String>(0)?,
304                row.get::<_, String>(1)?,
305                row.get::<_, String>(2)?,
306                row.get::<_, String>(3)?,
307                row.get::<_, String>(4)?,
308            ))
309        })?;
310
311        let mut results = Vec::new();
312        for row in rows {
313            let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
314            results.push(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?);
315        }
316        Ok(results)
317    }
318
319    /// Query sessions whose `updated_at` falls within `[from, to]`.
320    pub fn search_by_date(
321        &self,
322        from: DateTime<Utc>,
323        to: DateTime<Utc>,
324    ) -> Result<Vec<SessionSummary>> {
325        let mut stmt = self.db.prepare(
326            r#"SELECT id, project_dir, compressed_summary, created_at, updated_at
327               FROM sessions
328               WHERE updated_at >= ?1 AND updated_at <= ?2
329               ORDER BY updated_at DESC"#,
330        )?;
331
332        let rows = stmt.query_map(params![from.to_rfc3339(), to.to_rfc3339()], |row| {
333            Ok((
334                row.get::<_, String>(0)?,
335                row.get::<_, String>(1)?,
336                row.get::<_, String>(2)?,
337                row.get::<_, String>(3)?,
338                row.get::<_, String>(4)?,
339            ))
340        })?;
341
342        let mut results = Vec::new();
343        for row in rows {
344            let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
345            results.push(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?);
346        }
347        Ok(results)
348    }
349
350    /// Return the most recently updated session, or `None` if no sessions exist.
351    pub fn latest_session(&self) -> Result<Option<SessionSummary>> {
352        let mut stmt = self.db.prepare(
353            r#"SELECT id, project_dir, compressed_summary, created_at, updated_at
354               FROM sessions
355               ORDER BY updated_at DESC
356               LIMIT 1"#,
357        ).map_err(SqzError::SessionStore)?;
358
359        let rows = stmt.query_map([], |row| {
360            Ok((
361                row.get::<_, String>(0)?,
362                row.get::<_, String>(1)?,
363                row.get::<_, String>(2)?,
364                row.get::<_, String>(3)?,
365                row.get::<_, String>(4)?,
366            ))
367        }).map_err(SqzError::SessionStore)?;
368
369        for row in rows {
370            let (id, project_dir, compressed_summary, created_at, updated_at) =
371                row.map_err(SqzError::SessionStore)?;
372            return Ok(Some(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?));
373        }
374        Ok(None)
375    }
376
377    /// Query sessions whose `project_dir` matches `dir` exactly.
378    pub fn search_by_project(&self, dir: &Path) -> Result<Vec<SessionSummary>> {
379        let dir_str = dir.to_string_lossy().to_string();
380        let mut stmt = self.db.prepare(
381            r#"SELECT id, project_dir, compressed_summary, created_at, updated_at
382               FROM sessions
383               WHERE project_dir = ?1
384               ORDER BY updated_at DESC"#,
385        )?;
386
387        let rows = stmt.query_map(params![dir_str], |row| {
388            Ok((
389                row.get::<_, String>(0)?,
390                row.get::<_, String>(1)?,
391                row.get::<_, String>(2)?,
392                row.get::<_, String>(3)?,
393                row.get::<_, String>(4)?,
394            ))
395        })?;
396
397        let mut results = Vec::new();
398        for row in rows {
399            let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
400            results.push(row_to_summary(id, project_dir, compressed_summary, created_at, updated_at)?);
401        }
402        Ok(results)
403    }
404
405    // ── Cache entries ─────────────────────────────────────────────────────────
406
407    /// Persist a cache entry keyed by content hash.
408    ///
409    /// Stores both the compressed JSON (`data`) used for dedup-hit
410    /// responses AND the raw uncompressed bytes (`original`) so that
411    /// `sqz expand <prefix>` can serve truly uncompressed content. See
412    /// [`save_cache_entry_with_original`] for the original-aware version;
413    /// this convenience wrapper exists for callers that do not (yet)
414    /// have the pre-compression bytes handy. Rows written through this
415    /// path leave `original` as `NULL`, and `expand` will degrade to
416    /// returning the compressed blob with a note.
417    pub fn save_cache_entry(&self, hash: &str, compressed: &CompressedContent) -> Result<()> {
418        self.save_cache_entry_with_original(hash, compressed, None)
419    }
420
421    /// Persist a cache entry with both compressed and original content.
422    ///
423    /// `original` must be the exact bytes that produced `compressed`, so
424    /// that `expand` is a true inverse of dedup. We store the raw bytes
425    /// (not the UTF-8 string) because command output may include
426    /// non-UTF-8 sequences — storing the text would lose them.
427    pub fn save_cache_entry_with_original(
428        &self,
429        hash: &str,
430        compressed: &CompressedContent,
431        original: Option<&[u8]>,
432    ) -> Result<()> {
433        let data = serde_json::to_string(compressed)?;
434        let now = Utc::now().to_rfc3339();
435        self.db.execute(
436            r#"INSERT INTO cache_entries (hash, data, accessed_at, original)
437               VALUES (?1, ?2, ?3, ?4)
438               ON CONFLICT(hash) DO UPDATE
439                   SET data = excluded.data,
440                       accessed_at = excluded.accessed_at,
441                       -- Don't overwrite a previously-stored `original`
442                       -- with NULL. Older callers (that go through
443                       -- save_cache_entry rather than the _with_original
444                       -- variant) shouldn't erase the expand-able bytes.
445                       original = COALESCE(excluded.original, original)"#,
446            params![hash, data, now, original],
447        )?;
448        Ok(())
449    }
450
451    /// Retrieve the stored original bytes for a cached hash, if the
452    /// caller populated them via `save_cache_entry_with_original`.
453    ///
454    /// Returns `Ok(None)` for missing entries AND for entries that were
455    /// saved by an older call site that did not pass `original`. The
456    /// caller should fall back to the compressed blob in the latter case
457    /// and surface a note to the user so they know this specific entry
458    /// wasn't round-trippable.
459    pub fn get_cache_entry_original(&self, hash: &str) -> Result<Option<Vec<u8>>> {
460        let result: rusqlite::Result<Option<Vec<u8>>> = self.db.query_row(
461            "SELECT original FROM cache_entries WHERE hash = ?1",
462            params![hash],
463            |row| row.get(0),
464        );
465        match result {
466            Ok(v) => Ok(v),
467            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
468            Err(e) => Err(SqzError::SessionStore(e)),
469        }
470    }
471
472    /// Delete a cache entry by content hash.
473    pub fn delete_cache_entry(&self, hash: &str) -> Result<()> {
474        self.db.execute(
475            "DELETE FROM cache_entries WHERE hash = ?1",
476            params![hash],
477        )?;
478        Ok(())
479    }
480
481    /// Return all cache entries ordered by `accessed_at` ASC (oldest first),
482    /// as `(hash, size_bytes)` pairs where `size_bytes` is the byte length of
483    /// the stored JSON data.
484    pub fn list_cache_entries_lru(&self) -> Result<Vec<(String, u64)>> {
485        let mut stmt = self.db.prepare(
486            "SELECT hash, length(data) FROM cache_entries ORDER BY accessed_at ASC",
487        )?;
488        let rows = stmt.query_map([], |row| {
489            Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
490        })?;
491        let mut entries = Vec::new();
492        for row in rows {
493            let (hash, size) = row?;
494            entries.push((hash, size as u64));
495        }
496        Ok(entries)
497    }
498
499    /// Retrieve a cache entry by content hash, updating `accessed_at`.
500    pub fn get_cache_entry(&self, hash: &str) -> Result<Option<CompressedContent>> {
501        let result: rusqlite::Result<String> = self.db.query_row(
502            "SELECT data FROM cache_entries WHERE hash = ?1",
503            params![hash],
504            |row| row.get(0),
505        );
506
507        match result {
508            Ok(data) => {
509                // Touch accessed_at for LRU tracking.
510                let now = Utc::now().to_rfc3339();
511                let _ = self.db.execute(
512                    "UPDATE cache_entries SET accessed_at = ?1 WHERE hash = ?2",
513                    params![now, hash],
514                );
515                let entry: CompressedContent = serde_json::from_str(&data)?;
516                Ok(Some(entry))
517            }
518            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
519            Err(e) => Err(SqzError::SessionStore(e)),
520        }
521    }
522
523    /// Look up a cache entry by a **prefix** of the content hash.
524    ///
525    /// The inline dedup refs we hand to the LLM carry only the first 16
526    /// hex chars of the SHA-256 (`§ref:<16-hex>§`), so when an agent runs
527    /// `sqz expand <prefix>` we need to resolve those 16 chars back to the
528    /// full 64-char key. Uses `LIKE 'prefix%'` with an index-friendly
529    /// anchored pattern so the query stays O(log n) on the primary key.
530    ///
531    /// Returns `Ok(Some((full_hash, entry)))` on unique match.
532    /// Returns `Ok(None)` if no entries match.
533    /// Returns `Err(_)` if the prefix is ambiguous (2+ matches) — the
534    /// caller should tell the user to use a longer prefix. 16-hex
535    /// collisions are astronomically unlikely (one in 2^64) but we
536    /// refuse to guess rather than quietly serve a surprise file.
537    ///
538    /// The prefix is validated as lowercase hex. Non-hex input returns
539    /// `None` without touching the database — this is how we handle the
540    /// common user-error case of pasting the ref with the `§` markers
541    /// still attached (they get rejected before we query SQLite).
542    pub fn get_cache_entry_by_prefix(
543        &self,
544        prefix: &str,
545    ) -> Result<Option<(String, CompressedContent)>> {
546        // Reject anything that isn't pure lowercase hex. The inline refs
547        // we emit are always lowercase so there's no reason to case-fold
548        // here; accidentally matching uppercase input would also match
549        // unrelated entries if someone hand-crafted a collision.
550        if prefix.is_empty() || !prefix.chars().all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()) {
551            return Ok(None);
552        }
553        let pattern = format!("{prefix}%");
554        let mut stmt = self
555            .db
556            .prepare("SELECT hash, data FROM cache_entries WHERE hash LIKE ?1 LIMIT 2")?;
557        let mut rows = stmt.query(params![pattern])?;
558
559        let first = match rows.next()? {
560            Some(r) => {
561                let hash: String = r.get(0)?;
562                let data: String = r.get(1)?;
563                (hash, data)
564            }
565            None => return Ok(None),
566        };
567
568        // Two or more hits — refuse. Prefix ambiguity is user-recoverable:
569        // rerun with a longer prefix. Matching one arbitrarily would be
570        // a silent data surprise.
571        if rows.next()?.is_some() {
572            return Err(SqzError::Other(format!(
573                "cache: prefix '{prefix}' matches multiple entries — use a longer prefix"
574            )));
575        }
576        drop(rows);
577        drop(stmt);
578
579        // Touch accessed_at for LRU tracking — symmetric with get_cache_entry.
580        let now = Utc::now().to_rfc3339();
581        let _ = self.db.execute(
582            "UPDATE cache_entries SET accessed_at = ?1 WHERE hash = ?2",
583            params![now, first.0],
584        );
585        let entry: CompressedContent = serde_json::from_str(&first.1)?;
586        Ok(Some((first.0, entry)))
587    }
588
589    /// Read the `accessed_at` timestamp for a cached hash without updating
590    /// it. Returns `None` if the hash is not cached.
591    ///
592    /// Used by the dedup freshness check: if `accessed_at` is recent, the
593    /// LLM likely still has the original content in its context window, so
594    /// returning a ref is safe. If it's old, re-send the full content.
595    pub fn get_cache_entry_accessed_at(&self, hash: &str) -> Result<Option<DateTime<Utc>>> {
596        let result: rusqlite::Result<String> = self.db.query_row(
597            "SELECT accessed_at FROM cache_entries WHERE hash = ?1",
598            params![hash],
599            |row| row.get(0),
600        );
601        match result {
602            Ok(s) => {
603                let ts = s
604                    .parse::<DateTime<Utc>>()
605                    .map_err(|e| SqzError::Other(format!("invalid accessed_at: {e}")))?;
606                Ok(Some(ts))
607            }
608            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
609            Err(e) => Err(SqzError::SessionStore(e)),
610        }
611    }
612
613    /// Check if a cache entry exists without updating `accessed_at`.
614    pub fn cache_entry_exists(&self, hash: &str) -> Result<bool> {
615        let result: rusqlite::Result<i64> = self.db.query_row(
616            "SELECT 1 FROM cache_entries WHERE hash = ?1",
617            params![hash],
618            |row| row.get(0),
619        );
620        match result {
621            Ok(_) => Ok(true),
622            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(false),
623            Err(e) => Err(SqzError::SessionStore(e)),
624        }
625    }
626
627    /// Update `accessed_at` for a cached hash to the current time. Called by
628    /// the cache manager when a ref is served so the next staleness check
629    /// sees the recent send.
630    pub fn touch_cache_entry(&self, hash: &str) -> Result<()> {
631        let now = Utc::now().to_rfc3339();
632        self.db.execute(
633            "UPDATE cache_entries SET accessed_at = ?1 WHERE hash = ?2",
634            params![now, hash],
635        )?;
636        Ok(())
637    }
638
639    /// Set a metadata key/value. Persists across sqz process boundaries
640    /// (each shell-hook invocation is a short-lived process).
641    pub fn set_metadata(&self, key: &str, value: &str) -> Result<()> {
642        self.db.execute(
643            "INSERT INTO metadata (key, value) VALUES (?1, ?2)
644             ON CONFLICT(key) DO UPDATE SET value = excluded.value",
645            params![key, value],
646        )?;
647        Ok(())
648    }
649
650    /// Get a metadata value. Returns `None` if the key has never been set.
651    pub fn get_metadata(&self, key: &str) -> Result<Option<String>> {
652        let result: rusqlite::Result<String> = self.db.query_row(
653            "SELECT value FROM metadata WHERE key = ?1",
654            params![key],
655            |row| row.get(0),
656        );
657        match result {
658            Ok(v) => Ok(Some(v)),
659            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
660            Err(e) => Err(SqzError::SessionStore(e)),
661        }
662    }
663
664    /// Log a compression event for cumulative stats tracking.
665    pub fn log_compression(
666        &self,
667        tokens_original: u32,
668        tokens_compressed: u32,
669        stages: &[String],
670        mode: &str,
671    ) -> Result<()> {
672        self.log_compression_with_project(tokens_original, tokens_compressed, stages, mode, None)
673    }
674
675    /// Log a compression event tagged with a project directory.
676    pub fn log_compression_with_project(
677        &self,
678        tokens_original: u32,
679        tokens_compressed: u32,
680        stages: &[String],
681        mode: &str,
682        project_dir: Option<&str>,
683    ) -> Result<()> {
684        let now = Utc::now().to_rfc3339();
685        let stages_str = stages.join(",");
686        self.db.execute(
687            "INSERT INTO compression_log (tokens_original, tokens_compressed, stages_applied, mode, created_at, project_dir) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
688            params![tokens_original, tokens_compressed, stages_str, mode, now, project_dir],
689        ).map_err(SqzError::SessionStore)?;
690        Ok(())
691    }
692
693    /// Get cumulative compression stats from the log.
694    pub fn compression_stats(&self) -> Result<CompressionStats> {
695        let mut stmt = self.db.prepare(
696            "SELECT COUNT(*), COALESCE(SUM(tokens_original), 0), COALESCE(SUM(tokens_compressed), 0) FROM compression_log",
697        ).map_err(SqzError::SessionStore)?;
698
699        let stats = stmt.query_row([], |row| {
700            Ok(CompressionStats {
701                total_compressions: row.get::<_, u32>(0)?,
702                total_tokens_in: row.get::<_, u64>(1)?,
703                total_tokens_out: row.get::<_, u64>(2)?,
704            })
705        }).map_err(SqzError::SessionStore)?;
706
707        Ok(stats)
708    }
709
710    /// Get daily compression gains for the last N days.
711    pub fn daily_gains(&self, days: u32) -> Result<Vec<DailyGain>> {
712        let mut stmt = self.db.prepare(
713            "SELECT date(created_at) as d, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
714             FROM compression_log \
715             WHERE created_at >= date('now', ?1) \
716             GROUP BY d ORDER BY d",
717        ).map_err(SqzError::SessionStore)?;
718
719        let offset = format!("-{days} days");
720        let rows = stmt.query_map(params![offset], |row| {
721            let tokens_in: u64 = row.get(2)?;
722            let tokens_out: u64 = row.get(3)?;
723            Ok(DailyGain {
724                date: row.get(0)?,
725                compressions: row.get(1)?,
726                tokens_in,
727                tokens_saved: tokens_in.saturating_sub(tokens_out),
728            })
729        }).map_err(SqzError::SessionStore)?;
730
731        let mut gains = Vec::new();
732        for row in rows {
733            gains.push(row.map_err(SqzError::SessionStore)?);
734        }
735        Ok(gains)
736    }
737
738    /// Get cumulative compression stats for a specific project directory.
739    pub fn compression_stats_for_project(&self, project_dir: &str) -> Result<CompressionStats> {
740        let mut stmt = self.db.prepare(
741            "SELECT COUNT(*), COALESCE(SUM(tokens_original), 0), COALESCE(SUM(tokens_compressed), 0) \
742             FROM compression_log WHERE project_dir = ?1",
743        ).map_err(SqzError::SessionStore)?;
744
745        let stats = stmt.query_row(params![project_dir], |row| {
746            Ok(CompressionStats {
747                total_compressions: row.get::<_, u32>(0)?,
748                total_tokens_in: row.get::<_, u64>(1)?,
749                total_tokens_out: row.get::<_, u64>(2)?,
750            })
751        }).map_err(SqzError::SessionStore)?;
752
753        Ok(stats)
754    }
755
756    /// Get daily compression gains for a specific project directory.
757    pub fn daily_gains_for_project(&self, days: u32, project_dir: &str) -> Result<Vec<DailyGain>> {
758        let mut stmt = self.db.prepare(
759            "SELECT date(created_at) as d, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
760             FROM compression_log \
761             WHERE created_at >= date('now', ?1) AND project_dir = ?2 \
762             GROUP BY d ORDER BY d",
763        ).map_err(SqzError::SessionStore)?;
764
765        let offset = format!("-{days} days");
766        let rows = stmt.query_map(params![offset, project_dir], |row| {
767            let tokens_in: u64 = row.get(2)?;
768            let tokens_out: u64 = row.get(3)?;
769            Ok(DailyGain {
770                date: row.get(0)?,
771                compressions: row.get(1)?,
772                tokens_in,
773                tokens_saved: tokens_in.saturating_sub(tokens_out),
774            })
775        }).map_err(SqzError::SessionStore)?;
776
777        let mut gains = Vec::new();
778        for row in rows {
779            gains.push(row.map_err(SqzError::SessionStore)?);
780        }
781        Ok(gains)
782    }
783
784    /// List distinct project directories that have compression data.
785    pub fn list_projects(&self) -> Result<Vec<(String, u32, u64)>> {
786        let mut stmt = self.db.prepare(
787            "SELECT project_dir, COUNT(*), COALESCE(SUM(tokens_original) - SUM(tokens_compressed), 0) \
788             FROM compression_log \
789             WHERE project_dir IS NOT NULL \
790             GROUP BY project_dir \
791             ORDER BY COUNT(*) DESC",
792        ).map_err(SqzError::SessionStore)?;
793
794        let rows = stmt.query_map([], |row| {
795            Ok((
796                row.get::<_, String>(0)?,
797                row.get::<_, u32>(1)?,
798                row.get::<_, u64>(2)?,
799            ))
800        }).map_err(SqzError::SessionStore)?;
801
802        let mut projects = Vec::new();
803        for row in rows {
804            projects.push(row.map_err(SqzError::SessionStore)?);
805        }
806        Ok(projects)
807    }
808
809    /// Per-command breakdown of token usage. Returns top N commands by
810    /// total tokens consumed, with compression stats for each.
811    /// Used by `sqz stats --breakdown` to show where tokens are going.
812    pub fn command_breakdown(&self, limit: u32) -> Result<Vec<CommandStats>> {
813        let mut stmt = self.db.prepare(
814            "SELECT mode, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
815             FROM compression_log \
816             GROUP BY mode \
817             ORDER BY SUM(tokens_original) DESC \
818             LIMIT ?1",
819        ).map_err(SqzError::SessionStore)?;
820
821        let rows = stmt.query_map(params![limit], |row| {
822            let tokens_in: u64 = row.get(2)?;
823            let tokens_out: u64 = row.get(3)?;
824            Ok(CommandStats {
825                command: row.get(0)?,
826                invocations: row.get(1)?,
827                tokens_in,
828                tokens_out,
829                tokens_saved: tokens_in.saturating_sub(tokens_out),
830            })
831        }).map_err(SqzError::SessionStore)?;
832
833        let mut results = Vec::new();
834        for row in rows {
835            results.push(row.map_err(SqzError::SessionStore)?);
836        }
837        Ok(results)
838    }
839
840    /// Per-command breakdown filtered to a specific project.
841    pub fn command_breakdown_for_project(&self, limit: u32, project_dir: &str) -> Result<Vec<CommandStats>> {
842        let mut stmt = self.db.prepare(
843            "SELECT mode, COUNT(*), SUM(tokens_original), SUM(tokens_compressed) \
844             FROM compression_log \
845             WHERE project_dir = ?1 \
846             GROUP BY mode \
847             ORDER BY SUM(tokens_original) DESC \
848             LIMIT ?2",
849        ).map_err(SqzError::SessionStore)?;
850
851        let rows = stmt.query_map(params![project_dir, limit], |row| {
852            let tokens_in: u64 = row.get(2)?;
853            let tokens_out: u64 = row.get(3)?;
854            Ok(CommandStats {
855                command: row.get(0)?,
856                invocations: row.get(1)?,
857                tokens_in,
858                tokens_out,
859                tokens_saved: tokens_in.saturating_sub(tokens_out),
860            })
861        }).map_err(SqzError::SessionStore)?;
862
863        let mut results = Vec::new();
864        for row in rows {
865            results.push(row.map_err(SqzError::SessionStore)?);
866        }
867        Ok(results)
868    }
869
870    /// Get the total tokens injected in the current session (since last
871    /// compaction or start). Used by the adaptive pressure system to
872    /// decide whether to escalate compression intensity.
873    pub fn session_pressure(&self, since_minutes: u32) -> Result<u64> {
874        let mut stmt = self.db.prepare(
875            "SELECT COALESCE(SUM(tokens_compressed), 0) \
876             FROM compression_log \
877             WHERE created_at >= datetime('now', ?1)",
878        ).map_err(SqzError::SessionStore)?;
879
880        let offset = format!("-{since_minutes} minutes");
881        let total: u64 = stmt.query_row(params![offset], |row| {
882            row.get(0)
883        }).map_err(SqzError::SessionStore)?;
884
885        Ok(total)
886    }
887
888    // ── Known files (persistent cross-command context tracking) ───────────
889
890    /// Record a file path as "known" (its content is in the dedup cache).
891    /// Used by cross-command context refs to annotate error messages.
892    pub fn add_known_file(&self, path: &str) -> Result<()> {
893        let now = Utc::now().to_rfc3339();
894        self.db.execute(
895            "INSERT OR REPLACE INTO known_files (path, added_at) VALUES (?1, ?2)",
896            params![path, now],
897        ).map_err(SqzError::SessionStore)?;
898        Ok(())
899    }
900
901    /// Load all known file paths from the persistent store.
902    pub fn known_files(&self) -> Result<Vec<String>> {
903        let mut stmt = self.db.prepare(
904            "SELECT path FROM known_files ORDER BY added_at DESC",
905        ).map_err(SqzError::SessionStore)?;
906
907        let rows = stmt.query_map([], |row| {
908            row.get::<_, String>(0)
909        }).map_err(SqzError::SessionStore)?;
910
911        let mut files = Vec::new();
912        for row in rows {
913            files.push(row.map_err(SqzError::SessionStore)?);
914        }
915        Ok(files)
916    }
917    /// Clear all known files (e.g. on session reset).
918    pub fn clear_known_files(&self) -> Result<()> {
919        self.db
920            .execute("DELETE FROM known_files", [])
921            .map_err(SqzError::SessionStore)?;
922        Ok(())
923    }
924
925    /// Number of entries currently in the dedup cache.
926    pub fn cache_entry_count(&self) -> Result<u64> {
927        self.db.query_row(
928            "SELECT COUNT(*) FROM cache_entries", [],
929            |row| row.get(0),
930        ).map_err(SqzError::SessionStore)
931    }
932
933    /// Increment the cache-hit counter (persisted in metadata).
934    pub fn record_cache_hit(&self) -> Result<()> {
935        let count: u64 = self.db.query_row(
936            "SELECT value FROM metadata WHERE key = 'cache_hits'", [],
937            |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
938        ).unwrap_or(0);
939        let new = count + 1;
940        self.db.execute(
941            "INSERT OR REPLACE INTO metadata (key, value) VALUES ('cache_hits', ?1)",
942            params![new.to_string()],
943        )?;
944        Ok(())
945    }
946
947    /// Increment the cache-miss counter (persisted in metadata).
948    pub fn record_cache_miss(&self) -> Result<()> {
949        let count: u64 = self.db.query_row(
950            "SELECT value FROM metadata WHERE key = 'cache_misses'", [],
951            |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
952        ).unwrap_or(0);
953        let new = count + 1;
954        self.db.execute(
955            "INSERT OR REPLACE INTO metadata (key, value) VALUES ('cache_misses', ?1)",
956            params![new.to_string()],
957        )?;
958        Ok(())
959    }
960
961    /// Return (cache_hits, cache_misses) from metadata counters.
962    pub fn cache_hit_miss_counts(&self) -> Result<(u64, u64)> {
963        let hits: u64 = self.db.query_row(
964            "SELECT value FROM metadata WHERE key = 'cache_hits'", [],
965            |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
966        ).unwrap_or(0);
967        let misses: u64 = self.db.query_row(
968            "SELECT value FROM metadata WHERE key = 'cache_misses'", [],
969            |row| row.get::<_, String>(0).map(|s| s.parse::<u64>().unwrap_or(0)),
970        ).unwrap_or(0);
971        Ok((hits, misses))
972    }
973
974    /// Clear the dedup cache. After this, all future reads will be treated
975    /// as cache misses and compressed fresh. Does NOT clear compression
976    /// stats or session history.
977    pub fn clear_cache(&self) -> Result<u64> {
978        let count: u64 = self.db.query_row(
979            "SELECT COUNT(*) FROM cache_entries", [],
980            |row| row.get(0),
981        ).unwrap_or(0);
982        self.db.execute("DELETE FROM cache_entries", [])
983            .map_err(SqzError::SessionStore)?;
984        // Also clear the compaction marker so freshness checks don't
985        // reference timestamps from the deleted entries.
986        let _ = self.db.execute(
987            "DELETE FROM metadata WHERE key = 'last_compaction_at'", [],
988        );
989        Ok(count)
990    }
991
992    /// Clear compression stats (the `compression_log` table). After this,
993    /// `sqz stats` and `sqz gain` will show zero until new compressions
994    /// happen.
995    pub fn clear_stats(&self) -> Result<u64> {
996        let count: u64 = self.db.query_row(
997            "SELECT COUNT(*) FROM compression_log", [],
998            |row| row.get(0),
999        ).unwrap_or(0);
1000        self.db.execute("DELETE FROM compression_log", [])
1001            .map_err(SqzError::SessionStore)?;
1002        Ok(count)
1003    }
1004
1005    /// Clear compression stats for a specific project directory only.
1006    pub fn clear_stats_for_project(&self, project_dir: &str) -> Result<u64> {
1007        let count: u64 = self.db.query_row(
1008            "SELECT COUNT(*) FROM compression_log WHERE project_dir = ?1",
1009            rusqlite::params![project_dir],
1010            |row| row.get(0),
1011        ).unwrap_or(0);
1012        self.db.execute(
1013            "DELETE FROM compression_log WHERE project_dir = ?1",
1014            rusqlite::params![project_dir],
1015        ).map_err(SqzError::SessionStore)?;
1016        Ok(count)
1017    }
1018
1019    /// Full reset: clear cache, stats, sessions, known files, and metadata.
1020    /// Equivalent to deleting sessions.db and letting sqz recreate it.
1021    pub fn reset_all(&self) -> Result<()> {
1022        self.db.execute("DELETE FROM cache_entries", []).map_err(SqzError::SessionStore)?;
1023        self.db.execute("DELETE FROM compression_log", []).map_err(SqzError::SessionStore)?;
1024        self.db.execute("DELETE FROM sessions", []).map_err(SqzError::SessionStore)?;
1025        self.db.execute("DELETE FROM known_files", []).map_err(SqzError::SessionStore)?;
1026        self.db.execute("DELETE FROM metadata", []).map_err(SqzError::SessionStore)?;
1027        // VACUUM to reclaim disk space.
1028        let _ = self.db.execute("VACUUM", []);
1029        Ok(())
1030    }
1031
1032    /// Return all sessions ordered by most recently updated, limited to `limit`.
1033    pub fn list_sessions(&self, limit: u32) -> Result<Vec<SessionSummary>> {
1034        let mut stmt = self
1035            .db
1036            .prepare(
1037                "SELECT id, project_dir, compressed_summary, created_at, updated_at \
1038             FROM sessions ORDER BY updated_at DESC LIMIT ?1",
1039            )
1040            .map_err(SqzError::SessionStore)?;
1041
1042        let rows = stmt
1043            .query_map(params![limit], |row| {
1044                Ok((
1045                    row.get::<_, String>(0)?,
1046                    row.get::<_, String>(1)?,
1047                    row.get::<_, String>(2)?,
1048                    row.get::<_, String>(3)?,
1049                    row.get::<_, String>(4)?,
1050                ))
1051            })
1052            .map_err(SqzError::SessionStore)?;
1053
1054        let mut results = Vec::new();
1055        for row in rows {
1056            let (id, project_dir, compressed_summary, created_at, updated_at) = row?;
1057            results.push(row_to_summary(
1058                id,
1059                project_dir,
1060                compressed_summary,
1061                created_at,
1062                updated_at,
1063            )?);
1064        }
1065        Ok(results)
1066    }
1067
1068    /// Per-tool breakdown: group compression_log by stages_applied (comma-split).
1069    pub fn per_tool_breakdown(&self) -> Result<Vec<(String, u64, u64, u32)>> {
1070        let mut stmt = self
1071            .db
1072            .prepare(
1073                "SELECT stages_applied, SUM(tokens_original), SUM(tokens_compressed), COUNT(*) \
1074             FROM compression_log GROUP BY stages_applied ORDER BY COUNT(*) DESC LIMIT 20",
1075            )
1076            .map_err(SqzError::SessionStore)?;
1077
1078        let rows = stmt
1079            .query_map([], |row| {
1080                Ok((
1081                    row.get::<_, String>(0)?,
1082                    row.get::<_, u64>(1)?,
1083                    row.get::<_, u64>(2)?,
1084                    row.get::<_, u32>(3)?,
1085                ))
1086            })
1087            .map_err(SqzError::SessionStore)?;
1088
1089        let mut tool_map: std::collections::HashMap<String, (u64, u64, u32)> =
1090            std::collections::HashMap::new();
1091        for row in rows {
1092            let (stages, tokens_in, tokens_out, count) = row?;
1093            for stage in stages.split(',') {
1094                let stage = stage.trim();
1095                if stage.is_empty() {
1096                    continue;
1097                }
1098                let entry = tool_map.entry(stage.to_string()).or_default();
1099                entry.0 += tokens_in;
1100                entry.1 += tokens_out;
1101                entry.2 += count;
1102            }
1103        }
1104
1105        let mut result: Vec<_> = tool_map.into_iter().collect();
1106        result.sort_by(|a, b| b.1 .2.cmp(&a.1 .2));
1107        Ok(result
1108            .into_iter()
1109            .map(|(name, (tin, tout, cnt))| (name, tin, tout, cnt))
1110            .collect())
1111    }
1112
1113    /// Per-command breakdown: group compression_log by mode.
1114    pub fn per_command_breakdown(&self) -> Result<Vec<(String, u64, u64, u32)>> {
1115        let mut stmt = self
1116            .db
1117            .prepare(
1118                "SELECT COALESCE(mode, 'auto'), SUM(tokens_original), SUM(tokens_compressed), COUNT(*) \
1119             FROM compression_log GROUP BY mode ORDER BY COUNT(*) DESC LIMIT 20",
1120            )
1121            .map_err(SqzError::SessionStore)?;
1122
1123        let rows = stmt
1124            .query_map([], |row| {
1125                Ok((
1126                    row.get::<_, String>(0)?,
1127                    row.get::<_, u64>(1)?,
1128                    row.get::<_, u64>(2)?,
1129                    row.get::<_, u32>(3)?,
1130                ))
1131            })
1132            .map_err(SqzError::SessionStore)?;
1133
1134        let mut result = Vec::new();
1135        for row in rows {
1136            result.push(row?);
1137        }
1138        Ok(result)
1139    }
1140}
1141
1142/// Cumulative compression statistics.
1143#[derive(Debug, Clone, Default)]
1144pub struct CompressionStats {
1145    pub total_compressions: u32,
1146    pub total_tokens_in: u64,
1147    pub total_tokens_out: u64,
1148}
1149
1150impl CompressionStats {
1151    pub fn tokens_saved(&self) -> u64 {
1152        self.total_tokens_in.saturating_sub(self.total_tokens_out)
1153    }
1154
1155    pub fn reduction_pct(&self) -> f64 {
1156        if self.total_tokens_in == 0 {
1157            0.0
1158        } else {
1159            (1.0 - self.total_tokens_out as f64 / self.total_tokens_in as f64) * 100.0
1160        }
1161    }
1162}
1163
1164/// A single day's compression gain.
1165#[derive(Debug, Clone)]
1166pub struct DailyGain {
1167    pub date: String,
1168    pub compressions: u32,
1169    pub tokens_saved: u64,
1170    pub tokens_in: u64,
1171}
1172
1173/// Per-command token usage breakdown.
1174#[derive(Debug, Clone)]
1175pub struct CommandStats {
1176    pub command: String,
1177    pub invocations: u32,
1178    pub tokens_in: u64,
1179    pub tokens_out: u64,
1180    pub tokens_saved: u64,
1181}
1182
1183impl CommandStats {
1184    pub fn reduction_pct(&self) -> f64 {
1185        if self.tokens_in == 0 {
1186            0.0
1187        } else {
1188            (1.0 - self.tokens_out as f64 / self.tokens_in as f64) * 100.0
1189        }
1190    }
1191}
1192
1193// ── Tests ─────────────────────────────────────────────────────────────────────
1194
1195#[cfg(test)]
1196mod tests {
1197    use super::*;
1198    use crate::types::{BudgetState, CorrectionLog, ModelFamily, SessionState};
1199    use chrono::Utc;
1200    use proptest::prelude::*;
1201    use std::path::PathBuf;
1202
1203    fn make_session(id: &str, project_dir: &str, summary: &str) -> SessionState {
1204        let now = Utc::now();
1205        SessionState {
1206            id: id.to_string(),
1207            project_dir: PathBuf::from(project_dir),
1208            conversation: vec![],
1209            corrections: CorrectionLog::default(),
1210            pins: vec![],
1211            learnings: vec![],
1212            compressed_summary: summary.to_string(),
1213            budget: BudgetState {
1214                window_size: 200_000,
1215                consumed: 0,
1216                pinned: 0,
1217                model_family: ModelFamily::AnthropicClaude,
1218            },
1219            tool_usage: vec![],
1220            created_at: now,
1221            updated_at: now,
1222        }
1223    }
1224
1225    fn in_memory_store() -> SessionStore {
1226        let conn = Connection::open_in_memory().unwrap();
1227        apply_schema(&conn).unwrap();
1228        SessionStore { db: conn }
1229    }
1230
1231    #[test]
1232    fn test_save_and_load_session() {
1233        let store = in_memory_store();
1234        let session = make_session("sess-1", "/home/user/project", "REST API refactor");
1235
1236        let id = store.save_session(&session).unwrap();
1237        assert_eq!(id, "sess-1");
1238
1239        let loaded = store.load_session("sess-1".to_string()).unwrap();
1240        assert_eq!(loaded.id, session.id);
1241        assert_eq!(loaded.compressed_summary, session.compressed_summary);
1242        assert_eq!(loaded.project_dir, session.project_dir);
1243    }
1244
1245    #[test]
1246    fn test_save_session_upsert() {
1247        let store = in_memory_store();
1248        let mut session = make_session("sess-2", "/proj", "initial summary");
1249        store.save_session(&session).unwrap();
1250
1251        session.compressed_summary = "updated summary".to_string();
1252        store.save_session(&session).unwrap();
1253
1254        let loaded = store.load_session("sess-2".to_string()).unwrap();
1255        assert_eq!(loaded.compressed_summary, "updated summary");
1256    }
1257
1258    #[test]
1259    fn test_load_nonexistent_session_errors() {
1260        let store = in_memory_store();
1261        let result = store.load_session("does-not-exist".to_string());
1262        assert!(result.is_err());
1263    }
1264
1265    #[test]
1266    fn test_search_fts() {
1267        let store = in_memory_store();
1268        store.save_session(&make_session("s1", "/proj", "REST API refactor with authentication")).unwrap();
1269        store.save_session(&make_session("s2", "/proj", "database migration postgres")).unwrap();
1270
1271        let results = store.search("authentication").unwrap();
1272        assert_eq!(results.len(), 1);
1273        assert_eq!(results[0].id, "s1");
1274    }
1275
1276    #[test]
1277    fn test_search_by_date() {
1278        let store = in_memory_store();
1279        let now = Utc::now();
1280        let past = now - chrono::Duration::hours(2);
1281        let future = now + chrono::Duration::hours(2);
1282
1283        store.save_session(&make_session("s1", "/proj", "recent session")).unwrap();
1284
1285        let results = store.search_by_date(past, future).unwrap();
1286        assert!(!results.is_empty());
1287        assert!(results.iter().any(|r| r.id == "s1"));
1288    }
1289
1290    #[test]
1291    fn test_search_by_project() {
1292        let store = in_memory_store();
1293        store.save_session(&make_session("s1", "/home/user/alpha", "alpha project")).unwrap();
1294        store.save_session(&make_session("s2", "/home/user/beta", "beta project")).unwrap();
1295
1296        let results = store.search_by_project(Path::new("/home/user/alpha")).unwrap();
1297        assert_eq!(results.len(), 1);
1298        assert_eq!(results[0].id, "s1");
1299    }
1300
1301    #[test]
1302    fn test_cache_entry_round_trip() {
1303        let store = in_memory_store();
1304        let entry = CompressedContent {
1305            data: "compressed data".to_string(),
1306            tokens_compressed: 10,
1307            tokens_original: 50,
1308            stages_applied: vec!["strip_nulls".to_string()],
1309            compression_ratio: 0.2,
1310            provenance: crate::types::Provenance::default(),
1311            verify: None,
1312        };
1313
1314        store.save_cache_entry("abc123", &entry).unwrap();
1315
1316        let loaded = store.get_cache_entry("abc123").unwrap().unwrap();
1317        assert_eq!(loaded.data, entry.data);
1318        assert_eq!(loaded.tokens_compressed, entry.tokens_compressed);
1319        assert_eq!(loaded.tokens_original, entry.tokens_original);
1320    }
1321
1322    #[test]
1323    fn test_get_cache_entry_missing_returns_none() {
1324        let store = in_memory_store();
1325        let result = store.get_cache_entry("nonexistent").unwrap();
1326        assert!(result.is_none());
1327    }
1328
1329    #[test]
1330    fn test_open_or_create_corrupted_db() {
1331        let dir = tempfile::tempdir().unwrap();
1332        let path = dir.path().join("store.db");
1333
1334        // Write garbage bytes to simulate a corrupted database.
1335        std::fs::write(&path, b"this is not a valid sqlite database").unwrap();
1336
1337        // Should succeed by creating a fresh database.
1338        let store = SessionStore::open_or_create(&path).unwrap();
1339        let session = make_session("s1", "/proj", "after corruption");
1340        store.save_session(&session).unwrap();
1341        let loaded = store.load_session("s1".to_string()).unwrap();
1342        assert_eq!(loaded.id, "s1");
1343    }
1344
1345    // ── Property-based tests ──────────────────────────────────────────────────
1346
1347    /// Build a `SessionState` with a specific `updated_at` timestamp.
1348    fn make_session_at(id: &str, summary: &str, updated_at: DateTime<Utc>) -> SessionState {
1349        let now = Utc::now();
1350        SessionState {
1351            id: id.to_string(),
1352            project_dir: PathBuf::from("/proj"),
1353            conversation: vec![],
1354            corrections: CorrectionLog::default(),
1355            pins: vec![],
1356            learnings: vec![],
1357            compressed_summary: summary.to_string(),
1358            budget: BudgetState {
1359                window_size: 200_000,
1360                consumed: 0,
1361                pinned: 0,
1362                model_family: ModelFamily::AnthropicClaude,
1363            },
1364            tool_usage: vec![],
1365            created_at: now,
1366            updated_at,
1367        }
1368    }
1369
1370    // ── Property 26: Session store search correctness ─────────────────────────
1371    // **Validates: Requirements 20.2, 20.3, 20.4**
1372    //
1373    // For any set of sessions saved to the store, a keyword search SHALL return
1374    // all sessions whose compressed_summary contains the keyword, and no
1375    // sessions that don't contain it.
1376
1377    proptest! {
1378        /// **Validates: Requirements 20.2, 20.3, 20.4**
1379        ///
1380        /// For any set of sessions saved to the store, a keyword search SHALL
1381        /// return all sessions whose `compressed_summary` contains the keyword,
1382        /// and no sessions that don't contain it.
1383        #[test]
1384        fn prop_search_correctness(
1385            // A simple ASCII keyword: 5-8 lowercase letters, no common English
1386            // words that the porter stemmer might conflate with other terms.
1387            keyword in "[b-df-hj-np-tv-z]{5,8}",
1388            // 1-6 summaries that embed the keyword
1389            matching_suffixes in proptest::collection::vec("[a-z ]{4,20}", 1..=6usize),
1390            // 1-6 summaries that do NOT contain the keyword
1391            non_matching in proptest::collection::vec("[a-z ]{8,30}", 1..=6usize),
1392        ) {
1393            // Ensure the keyword doesn't accidentally appear in non-matching summaries.
1394            for s in &non_matching {
1395                prop_assume!(!s.contains(keyword.as_str()));
1396            }
1397
1398            let store = in_memory_store();
1399
1400            // Save matching sessions (summary = "<suffix> <keyword> <suffix>")
1401            let mut matching_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1402            for (i, suffix) in matching_suffixes.iter().enumerate() {
1403                let id = format!("match-{i}");
1404                let summary = format!("{} {} end", suffix, keyword);
1405                store.save_session(&make_session(&id, "/proj", &summary)).unwrap();
1406                matching_ids.insert(id);
1407            }
1408
1409            // Save non-matching sessions
1410            let mut non_matching_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1411            for (i, summary) in non_matching.iter().enumerate() {
1412                let id = format!("nomatch-{i}");
1413                store.save_session(&make_session(&id, "/proj", summary)).unwrap();
1414                non_matching_ids.insert(id);
1415            }
1416
1417            let results = store.search(&keyword).unwrap();
1418            let result_ids: std::collections::HashSet<String> =
1419                results.iter().map(|r| r.id.clone()).collect();
1420
1421            // Every matching session must appear in results.
1422            for id in &matching_ids {
1423                prop_assert!(
1424                    result_ids.contains(id),
1425                    "matching session '{}' not found in search results for keyword '{}'",
1426                    id, keyword
1427                );
1428            }
1429
1430            // No non-matching session may appear in results.
1431            for id in &non_matching_ids {
1432                prop_assert!(
1433                    !result_ids.contains(id),
1434                    "non-matching session '{}' incorrectly appeared in search results for keyword '{}'",
1435                    id, keyword
1436                );
1437            }
1438        }
1439    }
1440
1441    // ── Property: search_by_date correctness ─────────────────────────────────
1442    // **Validates: Requirements 20.4**
1443    //
1444    // For any set of sessions with different timestamps, searching by a date
1445    // range SHALL return exactly the sessions whose `updated_at` falls within
1446    // [from, to], and no sessions outside that range.
1447
1448    proptest! {
1449        /// **Validates: Requirements 20.4**
1450        ///
1451        /// For any set of sessions with distinct timestamps, `search_by_date`
1452        /// SHALL return exactly the sessions whose `updated_at` is within
1453        /// `[from, to]`, and no sessions outside that range.
1454        #[test]
1455        fn prop_search_by_date_correctness(
1456            // Generate 2-8 offsets in seconds from epoch (spread over a wide range)
1457            offsets in proptest::collection::vec(0i64..=86400i64 * 365, 2..=8usize),
1458            // The search window: start and end offsets (relative to the minimum offset)
1459            window_start_delta in 0i64..=3600i64,
1460            window_end_delta   in 3600i64..=7200i64,
1461        ) {
1462            use chrono::TimeZone;
1463
1464            // Deduplicate offsets so each session has a unique timestamp.
1465            let mut unique_offsets: Vec<i64> = offsets.clone();
1466            unique_offsets.sort_unstable();
1467            unique_offsets.dedup();
1468            prop_assume!(unique_offsets.len() >= 2);
1469
1470            let base_offset = unique_offsets[0];
1471            let from_offset = base_offset + window_start_delta;
1472            let to_offset   = base_offset + window_end_delta;
1473
1474            let from = Utc.timestamp_opt(from_offset, 0).unwrap();
1475            let to   = Utc.timestamp_opt(to_offset,   0).unwrap();
1476
1477            let store = in_memory_store();
1478
1479            let mut in_range_ids:  std::collections::HashSet<String> = std::collections::HashSet::new();
1480            let mut out_range_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1481
1482            for (i, &offset) in unique_offsets.iter().enumerate() {
1483                let ts = Utc.timestamp_opt(offset, 0).unwrap();
1484                let id = format!("sess-{i}");
1485                let session = make_session_at(&id, "some summary", ts);
1486                store.save_session(&session).unwrap();
1487
1488                if ts >= from && ts <= to {
1489                    in_range_ids.insert(id);
1490                } else {
1491                    out_range_ids.insert(id);
1492                }
1493            }
1494
1495            let results = store.search_by_date(from, to).unwrap();
1496            let result_ids: std::collections::HashSet<String> =
1497                results.iter().map(|r| r.id.clone()).collect();
1498
1499            // Every in-range session must appear.
1500            for id in &in_range_ids {
1501                prop_assert!(
1502                    result_ids.contains(id),
1503                    "in-range session '{}' missing from search_by_date results",
1504                    id
1505                );
1506            }
1507
1508            // No out-of-range session may appear.
1509            for id in &out_range_ids {
1510                prop_assert!(
1511                    !result_ids.contains(id),
1512                    "out-of-range session '{}' incorrectly appeared in search_by_date results",
1513                    id
1514                );
1515            }
1516        }
1517    }
1518}