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);
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
118pub(crate) fn apply_schema(conn: &Connection) -> rusqlite::Result<()> {
121 conn.execute_batch(SCHEMA)?;
122 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 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
184impl SessionStore {
187 #[cfg(test)]
190 pub(crate) fn from_connection(conn: Connection) -> Self {
191 Self { db: conn }
192 }
193
194 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 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 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 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 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 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 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 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 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 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 pub fn save_cache_entry(&self, hash: &str, compressed: &CompressedContent) -> Result<()> {
418 self.save_cache_entry_with_original(hash, compressed, None)
419 }
420
421 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 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 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 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 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 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 pub fn get_cache_entry_by_prefix(
543 &self,
544 prefix: &str,
545 ) -> Result<Option<(String, CompressedContent)>> {
546 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 let _ = self.db.execute(
987 "DELETE FROM metadata WHERE key = 'last_compaction_at'", [],
988 );
989 Ok(count)
990 }
991
992 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 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 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 let _ = self.db.execute("VACUUM", []);
1029 Ok(())
1030 }
1031
1032 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 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 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#[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#[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#[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#[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 std::fs::write(&path, b"this is not a valid sqlite database").unwrap();
1336
1337 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 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 proptest! {
1378 #[test]
1384 fn prop_search_correctness(
1385 keyword in "[b-df-hj-np-tv-z]{5,8}",
1388 matching_suffixes in proptest::collection::vec("[a-z ]{4,20}", 1..=6usize),
1390 non_matching in proptest::collection::vec("[a-z ]{8,30}", 1..=6usize),
1392 ) {
1393 for s in &non_matching {
1395 prop_assume!(!s.contains(keyword.as_str()));
1396 }
1397
1398 let store = in_memory_store();
1399
1400 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 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 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 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 proptest! {
1449 #[test]
1455 fn prop_search_by_date_correctness(
1456 offsets in proptest::collection::vec(0i64..=86400i64 * 365, 2..=8usize),
1458 window_start_delta in 0i64..=3600i64,
1460 window_end_delta in 3600i64..=7200i64,
1461 ) {
1462 use chrono::TimeZone;
1463
1464 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 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 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}