1#![allow(
5 clippy::cast_possible_truncation,
6 clippy::cast_possible_wrap,
7 clippy::cast_precision_loss,
8 clippy::cast_sign_loss
9)]
10use anyhow::{Context, Result};
11use chrono::Utc;
12use rusqlite::{Connection, OptionalExtension};
13use uuid::Uuid;
14
15use crate::provider::ModelPricing;
16
17#[derive(Debug, Clone)]
19pub struct ModelPref {
20 pub id: String,
21 pub favorite: bool,
22 pub last_used: Option<String>,
23 pub reasoning: Option<String>,
24}
25
26#[derive(Debug, Clone)]
29pub struct Space {
30 pub id: String,
31 pub name: String,
32 pub created_at: String,
33}
34
35#[derive(Debug, Clone)]
37pub struct Session {
38 pub id: String,
39 pub title: String,
40 pub model: String,
41 pub slug: Option<String>,
44 pub created_at: String,
45 pub compact_summary: Option<String>,
48 pub compact_through: i64,
52 pub web_mode: bool,
54 pub swarm_mode: bool,
57 pub kind: String,
59 pub research_parent_id: Option<String>,
62}
63
64#[derive(Debug, Clone)]
66pub struct Persona {
67 pub name: String,
68 pub model: String,
69 pub blurb: String,
70}
71
72#[derive(Debug, Clone)]
74pub struct Watch {
75 pub id: String,
76 pub space_id: String,
77 pub topic: String,
78 pub interval_hours: i64,
79 pub session_id: String,
80 pub last_run_at: Option<String>,
81}
82
83#[derive(Debug, Clone)]
89pub struct FileRow {
90 pub id: String,
91 pub name: String,
92 pub hash: String,
93 pub size: i64,
94 pub status: String,
95 pub mtime: i64,
98}
99
100#[derive(Debug, Clone)]
103pub struct Message {
104 pub role: String,
105 pub content: String,
106 pub model: Option<String>,
107 pub reasoning: Option<String>,
108 pub tokens: Option<i64>,
109 pub secs: Option<f64>,
110 pub cost: Option<f64>,
114 pub phrase: Option<String>,
116 pub persona: Option<String>,
119 pub created_at: Option<String>,
122}
123
124pub const DEFAULT_SPACE: &str = "default";
126
127const SCHEMA_VERSION: i64 = 2;
134
135const LEGACY_COLUMN_ADDS: &[(&str, &str, &str)] = &[
142 (
143 "messages",
144 "model",
145 "ALTER TABLE messages ADD COLUMN model TEXT",
146 ),
147 (
148 "messages",
149 "reasoning",
150 "ALTER TABLE messages ADD COLUMN reasoning TEXT",
151 ),
152 (
153 "messages",
154 "tokens",
155 "ALTER TABLE messages ADD COLUMN tokens INTEGER",
156 ),
157 (
158 "messages",
159 "secs",
160 "ALTER TABLE messages ADD COLUMN secs REAL",
161 ),
162 (
163 "messages",
164 "cost",
165 "ALTER TABLE messages ADD COLUMN cost REAL",
166 ),
167 (
168 "messages",
169 "phrase",
170 "ALTER TABLE messages ADD COLUMN phrase TEXT",
171 ),
172 (
173 "messages",
174 "persona",
175 "ALTER TABLE messages ADD COLUMN persona TEXT",
176 ),
177 (
178 "model_prefs",
179 "reasoning",
180 "ALTER TABLE model_prefs ADD COLUMN reasoning TEXT",
181 ),
182 (
183 "model_prefs",
184 "updated_at",
185 "ALTER TABLE model_prefs ADD COLUMN updated_at TEXT",
186 ),
187 (
188 "sessions",
189 "slug",
190 "ALTER TABLE sessions ADD COLUMN slug TEXT",
191 ),
192 (
193 "sessions",
194 "space_id",
195 "ALTER TABLE sessions ADD COLUMN space_id TEXT",
196 ),
197 (
198 "sessions",
199 "compact_summary",
200 "ALTER TABLE sessions ADD COLUMN compact_summary TEXT",
201 ),
202 (
203 "sessions",
204 "compact_through",
205 "ALTER TABLE sessions ADD COLUMN compact_through INTEGER NOT NULL DEFAULT 0",
206 ),
207 (
208 "sessions",
209 "web_mode",
210 "ALTER TABLE sessions ADD COLUMN web_mode INTEGER NOT NULL DEFAULT 0",
211 ),
212 (
213 "sessions",
214 "swarm_mode",
215 "ALTER TABLE sessions ADD COLUMN swarm_mode INTEGER NOT NULL DEFAULT 0",
216 ),
217 (
218 "sessions",
219 "kind",
220 "ALTER TABLE sessions ADD COLUMN kind TEXT NOT NULL DEFAULT 'chat'",
221 ),
222 (
223 "sessions",
224 "research_parent_id",
225 "ALTER TABLE sessions ADD COLUMN research_parent_id TEXT",
226 ),
227 (
228 "session_sources",
229 "flag",
230 "ALTER TABLE session_sources ADD COLUMN flag TEXT",
231 ),
232 (
233 "session_sources",
234 "updated_at",
235 "ALTER TABLE session_sources ADD COLUMN updated_at TEXT",
236 ),
237 (
238 "usage_log",
239 "cost_is_provider",
240 "ALTER TABLE usage_log ADD COLUMN cost_is_provider INTEGER",
241 ),
242 (
243 "usage_log",
244 "sync_id",
245 "ALTER TABLE usage_log ADD COLUMN sync_id TEXT",
246 ),
247 (
248 "usage_log",
249 "updated_at",
250 "ALTER TABLE usage_log ADD COLUMN updated_at TEXT",
251 ),
252 (
253 "app_settings",
254 "scope",
255 "ALTER TABLE app_settings ADD COLUMN scope TEXT NOT NULL DEFAULT 'sync'",
256 ),
257 (
258 "app_settings",
259 "updated_at",
260 "ALTER TABLE app_settings ADD COLUMN updated_at TEXT",
261 ),
262 (
263 "spaces",
264 "updated_at",
265 "ALTER TABLE spaces ADD COLUMN updated_at TEXT",
266 ),
267 (
268 "files",
269 "updated_at",
270 "ALTER TABLE files ADD COLUMN updated_at TEXT",
271 ),
272 (
273 "citations",
274 "sync_id",
275 "ALTER TABLE citations ADD COLUMN sync_id TEXT",
276 ),
277 (
278 "watches",
279 "updated_at",
280 "ALTER TABLE watches ADD COLUMN updated_at TEXT",
281 ),
282];
283
284pub fn cache_path_for(db_path: &std::path::Path) -> std::path::PathBuf {
289 db_path
290 .parent()
291 .filter(|p| !p.as_os_str().is_empty())
292 .map_or_else(
293 || std::path::PathBuf::from("cache.db"),
294 |p| p.join("cache.db"),
295 )
296}
297
298pub fn open_attached(db_path: &std::path::Path) -> Result<Connection> {
305 let conn =
306 Connection::open(db_path).with_context(|| format!("opening db {}", db_path.display()))?;
307 let cache = cache_path_for(db_path);
308 let escaped = cache.display().to_string().replace('\'', "''");
309 conn.execute_batch(&format!("ATTACH DATABASE '{escaped}' AS cache"))
310 .with_context(|| format!("attaching cache db {}", cache.display()))?;
311 migrate_cache(&conn, "cache")?;
312 Ok(conn)
313}
314
315pub fn migrate_cache(conn: &Connection, schema: &str) -> Result<()> {
319 conn.execute_batch(&format!(
320 "CREATE TABLE IF NOT EXISTS {schema}.web_cache (
321 url_norm TEXT PRIMARY KEY,
322 url TEXT NOT NULL,
323 title TEXT,
324 text TEXT NOT NULL,
325 fetched_at TEXT NOT NULL
326 );
327 CREATE VIRTUAL TABLE IF NOT EXISTS {schema}.file_chunks USING fts5(
328 file_id UNINDEXED,
329 seq UNINDEXED,
330 location UNINDEXED,
331 text
332 );
333 CREATE TABLE IF NOT EXISTS {schema}.chunk_embeddings (
334 file_id TEXT NOT NULL,
335 seq INTEGER NOT NULL,
336 vec BLOB NOT NULL,
337 PRIMARY KEY (file_id, seq)
338 );
339 CREATE TABLE IF NOT EXISTS {schema}.model_prices (
340 model_id TEXT PRIMARY KEY,
341 backend TEXT NOT NULL,
342 prompt_price REAL NOT NULL,
343 completion_price REAL NOT NULL,
344 cache_read_price REAL,
345 cache_write_price REAL,
346 updated_at TEXT NOT NULL
347 );
348 CREATE TABLE IF NOT EXISTS {schema}.file_index_state (
349 file_id TEXT PRIMARY KEY,
350 mtime INTEGER NOT NULL DEFAULT 0,
351 status TEXT NOT NULL DEFAULT '',
352 updated_at TEXT NOT NULL
353 );",
354 ))?;
355 Ok(())
356}
357
358fn has_column(conn: &Connection, table: &str, column: &str) -> Result<bool> {
362 let mut stmt = conn.prepare(&format!("PRAGMA main.table_info({table})"))?;
363 let rows = stmt.query_map([], |r| r.get::<_, String>(1))?;
364 for row in rows {
365 if row? == column {
366 return Ok(true);
367 }
368 }
369 Ok(false)
370}
371
372pub struct Db {
376 conn: Connection,
377}
378
379impl Db {
380 pub fn open(path: &std::path::Path) -> Result<Self> {
381 let conn = open_attached(path).with_context(|| format!("opening db {}", path.display()))?;
382 let mut db = Self { conn };
383 db.migrate()?;
384 Ok(db)
385 }
386
387 #[cfg(any(test, feature = "test-helpers"))]
388 pub fn open_in_memory() -> Result<Self> {
389 let conn = Connection::open_in_memory()?;
390 conn.execute_batch("ATTACH DATABASE ':memory:' AS cache")?;
391 migrate_cache(&conn, "cache")?;
392 let mut db = Db { conn };
393 db.migrate()?;
394 Ok(db)
395 }
396
397 pub(crate) fn conn(&self) -> &Connection {
400 &self.conn
401 }
402
403 #[cfg(test)]
404 pub fn conn_for_test(&self) -> &Connection {
405 &self.conn
406 }
407
408 fn touch(&self, table: &str, id: &str) -> Result<()> {
413 self.conn.execute(
414 &format!("UPDATE {table} SET updated_at = ?1 WHERE id = ?2"),
415 (Utc::now().to_rfc3339(), id),
416 )?;
417 Ok(())
418 }
419
420 fn tombstone(&self, table: &str, row_id: &str) -> Result<()> {
426 self.conn.execute(
427 "INSERT INTO sync_tombstones (table_name, row_id, deleted_at)
428 VALUES (?1, ?2, ?3)",
429 (table, row_id, Utc::now().to_rfc3339()),
430 )?;
431 Ok(())
432 }
433
434 #[allow(clippy::too_many_lines)]
436 fn migrate(&mut self) -> Result<()> {
437 self.conn.execute_batch(
438 "CREATE TABLE IF NOT EXISTS sessions (
439 id TEXT PRIMARY KEY,
440 title TEXT NOT NULL,
441 model TEXT NOT NULL,
442 created_at TEXT NOT NULL,
443 updated_at TEXT NOT NULL
444 );
445 CREATE TABLE IF NOT EXISTS messages (
446 id TEXT PRIMARY KEY,
447 session_id TEXT NOT NULL REFERENCES sessions(id),
448 role TEXT NOT NULL,
449 content TEXT NOT NULL,
450 created_at TEXT NOT NULL
451 );
452 CREATE INDEX IF NOT EXISTS idx_messages_session
453 ON messages(session_id, created_at);
454 CREATE TABLE IF NOT EXISTS model_prefs (
455 id TEXT PRIMARY KEY,
456 favorite INTEGER NOT NULL DEFAULT 0,
457 last_used TEXT,
458 reasoning TEXT,
459 updated_at TEXT
460 );
461 CREATE TABLE IF NOT EXISTS app_settings (
462 key TEXT PRIMARY KEY,
463 value TEXT NOT NULL,
464 scope TEXT NOT NULL DEFAULT 'sync',
465 updated_at TEXT
466 );
467 CREATE TABLE IF NOT EXISTS spaces (
468 id TEXT PRIMARY KEY,
469 name TEXT NOT NULL UNIQUE,
470 created_at TEXT NOT NULL,
471 updated_at TEXT
472 );
473 CREATE TABLE IF NOT EXISTS files (
474 id TEXT PRIMARY KEY,
475 space_id TEXT NOT NULL,
476 name TEXT NOT NULL,
477 hash TEXT NOT NULL,
478 size INTEGER NOT NULL,
479 created_at TEXT NOT NULL,
480 updated_at TEXT NOT NULL,
481 UNIQUE(space_id, name)
482 );
483 CREATE TABLE IF NOT EXISTS citations (
484 id INTEGER PRIMARY KEY AUTOINCREMENT,
485 sync_id TEXT NOT NULL,
486 space_id TEXT NOT NULL,
487 report_file TEXT NOT NULL,
488 url TEXT NOT NULL,
489 title TEXT
490 );
491 CREATE INDEX IF NOT EXISTS idx_citations_space ON citations(space_id);
492 CREATE TABLE IF NOT EXISTS session_sources (
493 session_id TEXT NOT NULL,
494 url_norm TEXT NOT NULL,
495 flag TEXT,
496 updated_at TEXT,
497 PRIMARY KEY (session_id, url_norm)
498 );
499 CREATE TABLE IF NOT EXISTS watches (
500 id TEXT PRIMARY KEY,
501 space_id TEXT NOT NULL,
502 topic TEXT NOT NULL,
503 interval_hours INTEGER NOT NULL,
504 session_id TEXT NOT NULL,
505 last_run_at TEXT,
506 updated_at TEXT
507 );
508 CREATE TABLE IF NOT EXISTS swarm_personas (
509 session_id TEXT NOT NULL,
510 ord INTEGER NOT NULL,
511 name TEXT NOT NULL,
512 model TEXT NOT NULL,
513 persona TEXT NOT NULL
514 );
515 CREATE INDEX IF NOT EXISTS idx_swarm_personas_session
516 ON swarm_personas(session_id, ord);
517 CREATE TABLE IF NOT EXISTS usage_log (
518 id INTEGER PRIMARY KEY AUTOINCREMENT,
519 sync_id TEXT NOT NULL,
520 created_at TEXT NOT NULL,
521 session_id TEXT,
522 space_id TEXT,
523 backend TEXT NOT NULL,
524 model TEXT NOT NULL,
525 prompt_tokens INTEGER NOT NULL,
526 completion_tokens INTEGER NOT NULL,
527 cache_read_tokens INTEGER NOT NULL DEFAULT 0,
528 cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
529 cost REAL,
530 cost_is_provider INTEGER,
531 updated_at TEXT
532 );
533 CREATE INDEX IF NOT EXISTS idx_usage_log_created ON usage_log(created_at);
534 CREATE INDEX IF NOT EXISTS idx_usage_log_model ON usage_log(model);
535 CREATE TABLE IF NOT EXISTS sync_tombstones (
536 id INTEGER PRIMARY KEY AUTOINCREMENT,
537 table_name TEXT NOT NULL,
538 row_id TEXT NOT NULL,
539 deleted_at TEXT NOT NULL
540 );
541 CREATE INDEX IF NOT EXISTS idx_sync_tombstones_table ON sync_tombstones(table_name, row_id);
542 CREATE TABLE IF NOT EXISTS device_meta (
543 device_id TEXT PRIMARY KEY,
544 created_at TEXT NOT NULL
545 );
546 CREATE TABLE IF NOT EXISTS sync_state (
547 peer_id TEXT NOT NULL,
548 table_name TEXT NOT NULL,
549 pull_cursor TEXT,
550 push_cursor TEXT,
551 last_synced_at TEXT,
552 PRIMARY KEY (peer_id, table_name)
553 );",
554 )?;
555 let version: i64 = self
561 .conn
562 .query_row("PRAGMA user_version", [], |r| r.get(0))?;
563 if version < SCHEMA_VERSION {
564 for (table, column, ddl) in LEGACY_COLUMN_ADDS {
565 if !has_column(&self.conn, table, column)? {
566 self.conn.execute(ddl, []).with_context(|| {
567 format!("migrating column {table}.{column} (user_version {version})")
568 })?;
569 }
570 }
571 let backfill_tx = self.conn.transaction()?;
576 for table in ["citations", "usage_log"] {
577 let ids: Vec<i64> = {
578 let mut stmt = backfill_tx
579 .prepare(&format!("SELECT id FROM {table} WHERE sync_id IS NULL"))?;
580 let rows = stmt.query_map([], |r| r.get(0))?;
581 rows.collect::<rusqlite::Result<Vec<_>>>()?
582 };
583 let mut update = backfill_tx
584 .prepare(&format!("UPDATE {table} SET sync_id = ?1 WHERE id = ?2"))?;
585 for id in ids {
586 update.execute((Uuid::new_v4().to_string(), id))?;
587 }
588 }
589 backfill_tx.commit()?;
590 self.conn.execute_batch(
594 "CREATE UNIQUE INDEX IF NOT EXISTS idx_citations_sync_id ON citations(sync_id);
595 CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_log_sync_id ON usage_log(sync_id);",
596 )?;
597 let now = Utc::now().to_rfc3339();
602 if has_column(&self.conn, "files", "mtime")? {
603 self.conn.execute(
604 "INSERT OR IGNORE INTO cache.file_index_state (file_id, mtime, status, updated_at)
605 SELECT id, mtime, status, ?1 FROM files",
606 [&now],
607 )?;
608 }
609 if has_column(&self.conn, "web_cache", "url_norm")? {
610 self.conn.execute(
611 "INSERT OR IGNORE INTO cache.web_cache (url_norm, url, title, text, fetched_at)
612 SELECT url_norm, url, title, text, fetched_at FROM web_cache",
613 [],
614 )?;
615 }
616 if has_column(&self.conn, "chunk_embeddings", "file_id")? {
617 self.conn.execute(
618 "INSERT OR IGNORE INTO cache.chunk_embeddings (file_id, seq, vec)
619 SELECT file_id, seq, vec FROM chunk_embeddings",
620 [],
621 )?;
622 }
623 if has_column(&self.conn, "file_chunks", "file_id")? {
624 self.conn.execute(
625 "INSERT OR IGNORE INTO cache.file_chunks (file_id, seq, location, text)
626 SELECT file_id, seq, location, text FROM file_chunks",
627 [],
628 )?;
629 }
630 if has_column(&self.conn, "model_prices", "model_id")? {
633 let cols = if has_column(&self.conn, "model_prices", "cache_read_price")? {
634 "model_id, backend, prompt_price, completion_price,\n cache_read_price, cache_write_price, updated_at"
635 } else {
636 "model_id, backend, prompt_price, completion_price, updated_at"
637 };
638 self.conn.execute(
639 &format!(
640 "INSERT OR IGNORE INTO cache.model_prices ({cols}) SELECT {cols} FROM model_prices"
641 ),
642 [],
643 )?;
644 }
645 let old_default: Option<String> = self
652 .conn
653 .query_row(
654 "SELECT id FROM spaces WHERE name = ?1",
655 [DEFAULT_SPACE],
656 |r| r.get(0),
657 )
658 .optional()?;
659 if let Some(old) = old_default.filter(|id| id != DEFAULT_SPACE) {
660 let tx = self.conn.transaction()?;
661 for table in ["sessions", "files", "watches", "usage_log", "citations"] {
662 tx.execute(
663 &format!("UPDATE {table} SET space_id = ?1 WHERE space_id = ?2"),
664 (DEFAULT_SPACE, &old),
665 )?;
666 }
667 tx.execute(
668 "UPDATE sync_tombstones SET row_id = ?1
669 WHERE table_name = 'spaces' AND row_id = ?2",
670 (DEFAULT_SPACE, &old),
671 )?;
672 tx.execute(
673 "UPDATE spaces SET id = ?1 WHERE id = ?2",
674 (DEFAULT_SPACE, &old),
675 )?;
676 tx.commit()?;
677 }
678 self.conn
679 .execute_batch(&format!("PRAGMA user_version = {SCHEMA_VERSION}"))?;
680 }
681 let _ = self
684 .conn
685 .execute_batch("DROP TABLE IF EXISTS message_images;");
686 let now = Utc::now().to_rfc3339();
692 self.conn.execute(
693 "INSERT OR IGNORE INTO spaces (id, name, created_at) VALUES (?1, ?2, ?3)",
694 (DEFAULT_SPACE, DEFAULT_SPACE, &now),
695 )?;
696 let default_id: String = self.conn.query_row(
697 "SELECT id FROM spaces WHERE name = ?1",
698 [DEFAULT_SPACE],
699 |r| r.get(0),
700 )?;
701 self.conn.execute(
702 "UPDATE sessions SET space_id = ?1 WHERE space_id IS NULL",
703 [&default_id],
704 )?;
705 Ok(())
706 }
707
708 pub fn default_space_id(&self) -> Result<String> {
710 Ok(self.conn.query_row(
711 "SELECT id FROM spaces WHERE name = ?1",
712 [DEFAULT_SPACE],
713 |r| r.get(0),
714 )?)
715 }
716
717 pub fn create_space(&self, name: &str) -> Result<Space> {
718 let id = Uuid::new_v4().to_string();
719 let now = Utc::now().to_rfc3339();
720 self.conn.execute(
721 "INSERT INTO spaces (id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?3)",
722 (&id, name, &now),
723 )?;
724 Ok(Space {
725 id,
726 name: name.to_string(),
727 created_at: now,
728 })
729 }
730
731 pub fn list_spaces(&self) -> Result<Vec<Space>> {
733 let mut stmt = self
734 .conn
735 .prepare("SELECT id, name, created_at FROM spaces ORDER BY created_at ASC")?;
736 let rows = stmt.query_map([], |r| {
737 Ok(Space {
738 id: r.get(0)?,
739 name: r.get(1)?,
740 created_at: r.get(2)?,
741 })
742 })?;
743 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
744 }
745
746 pub fn rename_space(&self, id: &str, name: &str) -> Result<()> {
747 self.conn.execute(
748 "UPDATE spaces SET name = ?2, updated_at = ?3 WHERE id = ?1",
749 (id, name, Utc::now().to_rfc3339()),
750 )?;
751 Ok(())
752 }
753
754 pub fn delete_space(&self, id: &str) -> Result<()> {
759 let default_id = self.default_space_id()?;
760 self.conn.execute(
761 "UPDATE sessions SET space_id = ?1, updated_at = ?2 WHERE space_id = ?3",
762 (&default_id, Utc::now().to_rfc3339(), id),
763 )?;
764 self.conn
765 .execute("DELETE FROM spaces WHERE id = ?1", [id])?;
766 self.tombstone("spaces", id)?;
767 Ok(())
768 }
769
770 pub fn count_sessions(&self, space_id: &str) -> Result<u64> {
772 let n: i64 = self.conn.query_row(
773 "SELECT COUNT(*) FROM sessions WHERE space_id = ?1",
774 [space_id],
775 |r| r.get(0),
776 )?;
777 Ok(n as u64)
778 }
779
780 pub fn last_message_preview(&self, session_id: &str) -> Option<String> {
784 let mut stmt = self
785 .conn
786 .prepare(
787 "SELECT content FROM messages WHERE session_id = ?1 \
788 AND role IN ('user','assistant') ORDER BY id DESC LIMIT 1",
789 )
790 .ok()?;
791 let mut rows = stmt
792 .query_map([session_id], |r| r.get::<_, String>(0))
793 .ok()?;
794 rows.next().and_then(Result::ok)
795 }
796
797 pub fn setting_is_local(key: &str) -> bool {
806 matches!(
807 key,
808 "searxng_url"
812 | "langsearch_key"
813 | "search_provider"
814 | "ocr_engine"
815 | "ocr_model"
816 | "local_ocr_model"
817 | "usage_range"
819 | "ui_background"
820 | "last_update_check"
821 | "sync_ssh_peer"
824 )
825 }
826
827 pub fn set_setting(&self, key: &str, value: &str) -> Result<()> {
828 let now = Utc::now().to_rfc3339();
829 let scope = if Self::setting_is_local(key) {
830 "local"
831 } else {
832 "sync"
833 };
834 self.conn.execute(
835 "INSERT INTO app_settings (key, value, scope, updated_at)
836 VALUES (?1, ?2, ?3, ?4)
837 ON CONFLICT(key) DO UPDATE SET value = ?2, scope = ?3, updated_at = ?4",
838 (key, value, scope, &now),
839 )?;
840 Ok(())
841 }
842
843 pub fn load_settings(&self) -> Result<Vec<(String, String)>> {
844 let mut stmt = self.conn.prepare("SELECT key, value FROM app_settings")?;
845 let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
846 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
847 }
848
849 pub fn update_check_due(&self) -> bool {
854 let today = chrono::Local::now().format("%Y-%m-%d").to_string();
855 let last = self
856 .conn
857 .query_row(
858 "SELECT value FROM app_settings WHERE key = 'last_update_check'",
859 [],
860 |r| r.get::<_, String>(0),
861 )
862 .ok();
863 if last.as_deref() == Some(today.as_str()) {
864 return false;
865 }
866 let _ = self.set_setting("last_update_check", &today);
867 true
868 }
869
870 pub fn set_reasoning(&self, model_id: &str, effort: Option<&str>) -> Result<()> {
872 self.conn.execute(
873 "INSERT INTO model_prefs (id, reasoning, updated_at) VALUES (?1, ?2, ?3)
874 ON CONFLICT(id) DO UPDATE SET reasoning = ?2, updated_at = ?3",
875 (model_id, effort, Utc::now().to_rfc3339()),
876 )?;
877 Ok(())
878 }
879
880 pub fn toggle_favorite(&self, model_id: &str) -> Result<bool> {
882 self.conn.execute(
883 "INSERT INTO model_prefs (id, favorite, updated_at) VALUES (?1, 1, ?2)
884 ON CONFLICT(id) DO UPDATE SET favorite = 1 - favorite, updated_at = ?2",
885 (model_id, Utc::now().to_rfc3339()),
886 )?;
887 let fav: i64 = self.conn.query_row(
888 "SELECT favorite FROM model_prefs WHERE id = ?1",
889 [model_id],
890 |r| r.get(0),
891 )?;
892 Ok(fav != 0)
893 }
894
895 pub fn mark_model_used(&self, model_id: &str) -> Result<()> {
897 let now = Utc::now().to_rfc3339();
898 self.conn.execute(
899 "INSERT INTO model_prefs (id, favorite, last_used, updated_at) VALUES (?1, 0, ?2, ?2)
900 ON CONFLICT(id) DO UPDATE SET last_used = ?2, updated_at = ?2",
901 (model_id, &now),
902 )?;
903 Ok(())
904 }
905
906 pub fn load_model_prefs(&self) -> Result<Vec<ModelPref>> {
908 let mut stmt = self
909 .conn
910 .prepare("SELECT id, favorite, last_used, reasoning FROM model_prefs")?;
911 let rows = stmt.query_map([], |r| {
912 Ok(ModelPref {
913 id: r.get(0)?,
914 favorite: r.get::<_, i64>(1)? != 0,
915 last_used: r.get(2)?,
916 reasoning: r.get(3)?,
917 })
918 })?;
919 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
920 }
921
922 pub fn create_session(
923 &self,
924 title: &str,
925 model: &str,
926 space_id: &str,
927 kind: &str,
928 ) -> Result<Session> {
929 let id = Uuid::new_v4().to_string();
930 let now = Utc::now().to_rfc3339();
931 self.conn.execute(
932 "INSERT INTO sessions (id, title, model, space_id, kind, created_at, updated_at)
933 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6)",
934 (&id, title, model, space_id, kind, &now),
935 )?;
936 Ok(Session {
937 id,
938 title: title.to_string(),
939 model: model.to_string(),
940 slug: None,
941 created_at: now,
942 compact_summary: None,
943 compact_through: 0,
944 web_mode: false,
945 swarm_mode: false,
946 kind: kind.to_string(),
947 research_parent_id: None,
948 })
949 }
950
951 pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
954 self.conn
955 .query_row(
956 "SELECT id, title, model, slug, created_at, compact_summary, compact_through, \
957 web_mode, swarm_mode, kind, research_parent_id
958 FROM sessions WHERE id = ?1",
959 [id],
960 |r| {
961 Ok(Session {
962 id: r.get(0)?,
963 title: r.get(1)?,
964 model: r.get(2)?,
965 slug: r.get(3)?,
966 created_at: r.get(4)?,
967 compact_summary: r.get(5)?,
968 compact_through: r.get(6)?,
969 web_mode: r.get::<_, i64>(7)? != 0,
970 swarm_mode: r.get::<_, i64>(8)? != 0,
971 kind: r.get(9)?,
972 research_parent_id: r.get(10)?,
973 })
974 },
975 )
976 .optional()
977 .map_err(Into::into)
978 }
979
980 pub fn list_sessions(&self, space_id: &str) -> Result<Vec<Session>> {
982 let mut stmt = self.conn.prepare(
983 "SELECT id, title, model, slug, created_at, compact_summary, compact_through, \
984 web_mode, swarm_mode, kind, research_parent_id
985 FROM sessions WHERE space_id = ?1 ORDER BY updated_at DESC",
986 )?;
987 let rows = stmt.query_map([space_id], |r| {
988 Ok(Session {
989 id: r.get(0)?,
990 title: r.get(1)?,
991 model: r.get(2)?,
992 slug: r.get(3)?,
993 created_at: r.get(4)?,
994 compact_summary: r.get(5)?,
995 compact_through: r.get(6)?,
996 web_mode: r.get::<_, i64>(7)? != 0,
997 swarm_mode: r.get::<_, i64>(8)? != 0,
998 kind: r.get(9)?,
999 research_parent_id: r.get(10)?,
1000 })
1001 })?;
1002 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1003 }
1004
1005 pub fn latest_session(&self) -> Result<Option<(String, Session)>> {
1008 self.conn
1009 .query_row(
1010 "SELECT space_id, id, title, model, slug, created_at, compact_summary,
1011 compact_through, web_mode, swarm_mode, kind, research_parent_id
1012 FROM sessions
1013 ORDER BY updated_at DESC, rowid DESC
1014 LIMIT 1",
1015 [],
1016 |r| {
1017 Ok((
1018 r.get(0)?,
1019 Session {
1020 id: r.get(1)?,
1021 title: r.get(2)?,
1022 model: r.get(3)?,
1023 slug: r.get(4)?,
1024 created_at: r.get(5)?,
1025 compact_summary: r.get(6)?,
1026 compact_through: r.get(7)?,
1027 web_mode: r.get::<_, i64>(8)? != 0,
1028 swarm_mode: r.get::<_, i64>(9)? != 0,
1029 kind: r.get(10)?,
1030 research_parent_id: r.get(11)?,
1031 },
1032 ))
1033 },
1034 )
1035 .optional()
1036 .map_err(Into::into)
1037 }
1038
1039 pub fn set_compaction(&self, session_id: &str, summary: &str, through: i64) -> Result<()> {
1042 self.conn.execute(
1043 "UPDATE sessions SET compact_summary = ?2, compact_through = ?3, updated_at = ?4
1044 WHERE id = ?1",
1045 (session_id, summary, through, Utc::now().to_rfc3339()),
1046 )?;
1047 Ok(())
1048 }
1049
1050 pub fn set_session_web_mode(&self, session_id: &str, on: bool) -> Result<()> {
1052 self.conn.execute(
1053 "UPDATE sessions SET web_mode = ?2, updated_at = ?3 WHERE id = ?1",
1054 (session_id, i64::from(on), Utc::now().to_rfc3339()),
1055 )?;
1056 Ok(())
1057 }
1058
1059 pub fn set_session_swarm_mode(&self, session_id: &str, on: bool) -> Result<()> {
1061 self.conn.execute(
1062 "UPDATE sessions SET swarm_mode = ?2, updated_at = ?3 WHERE id = ?1",
1063 (session_id, i64::from(on), Utc::now().to_rfc3339()),
1064 )?;
1065 Ok(())
1066 }
1067
1068 pub fn list_swarm_personas(&self, session_id: &str) -> Result<Vec<Persona>> {
1070 let mut stmt = self.conn.prepare(
1071 "SELECT name, model, persona FROM swarm_personas
1072 WHERE session_id = ?1 ORDER BY ord ASC",
1073 )?;
1074 let rows = stmt.query_map([session_id], |r| {
1075 Ok(Persona {
1076 name: r.get(0)?,
1077 model: r.get(1)?,
1078 blurb: r.get(2)?,
1079 })
1080 })?;
1081 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1082 }
1083
1084 pub fn save_swarm_personas(&self, session_id: &str, personas: &[Persona]) -> Result<()> {
1089 let old: Vec<i64> = {
1090 let mut stmt = self
1091 .conn
1092 .prepare("SELECT ord FROM swarm_personas WHERE session_id = ?1")?;
1093 let rows = stmt.query_map([session_id], |r| r.get(0))?;
1094 rows.collect::<rusqlite::Result<Vec<_>>>()?
1095 };
1096 for ord in old {
1097 self.tombstone("swarm_personas", &format!("{session_id}:{ord}"))?;
1098 }
1099 self.conn.execute(
1100 "DELETE FROM swarm_personas WHERE session_id = ?1",
1101 [session_id],
1102 )?;
1103 for (i, p) in personas.iter().enumerate() {
1104 self.conn.execute(
1105 "INSERT INTO swarm_personas (session_id, ord, name, model, persona)
1106 VALUES (?1, ?2, ?3, ?4, ?5)",
1107 (session_id, i as i64, &p.name, &p.model, &p.blurb),
1108 )?;
1109 }
1110 self.touch("sessions", session_id)
1111 }
1112
1113 pub fn set_research_parent(&self, id: &str, parent_id: &str) -> Result<()> {
1117 self.conn.execute(
1118 "UPDATE sessions SET research_parent_id = ?2, updated_at = ?3 WHERE id = ?1",
1119 (id, parent_id, Utc::now().to_rfc3339()),
1120 )?;
1121 Ok(())
1122 }
1123
1124 pub fn set_session_title(&self, id: &str, title: &str, slug: Option<&str>) -> Result<()> {
1126 self.conn.execute(
1127 "UPDATE sessions SET title = ?2, slug = COALESCE(?3, slug), updated_at = ?4
1128 WHERE id = ?1",
1129 (id, title, slug, Utc::now().to_rfc3339()),
1130 )?;
1131 Ok(())
1132 }
1133
1134 pub fn delete_message(&self, id: &str) -> Result<()> {
1138 self.conn
1139 .execute("DELETE FROM messages WHERE id = ?1", [id])?;
1140 self.tombstone("messages", id)?;
1141 Ok(())
1142 }
1143
1144 pub fn delete_session(&self, id: &str) -> Result<()> {
1148 let ids: Vec<String> = {
1149 let mut stmt = self
1150 .conn
1151 .prepare("SELECT id FROM messages WHERE session_id = ?1")?;
1152 let rows = stmt.query_map([id], |r| r.get(0))?;
1153 rows.collect::<rusqlite::Result<Vec<_>>>()?
1154 };
1155 let persona_ids: Vec<String> = {
1156 let mut stmt = self
1157 .conn
1158 .prepare("SELECT ord FROM swarm_personas WHERE session_id = ?1")?;
1159 let rows = stmt.query_map([id], |r| Ok(format!("{id}:{}", r.get::<_, i64>(0)?)))?;
1160 rows.collect::<rusqlite::Result<Vec<_>>>()?
1161 };
1162 for mid in &ids {
1163 self.tombstone("messages", mid)?;
1164 }
1165 for persona_id in &persona_ids {
1166 self.tombstone("swarm_personas", persona_id)?;
1167 }
1168 self.conn
1169 .execute("DELETE FROM messages WHERE session_id = ?1", [id])?;
1170 self.conn
1171 .execute("DELETE FROM swarm_personas WHERE session_id = ?1", [id])?;
1172 self.conn
1173 .execute("DELETE FROM sessions WHERE id = ?1", [id])?;
1174 self.tombstone("sessions", id)?;
1175 Ok(())
1176 }
1177
1178 pub fn load_messages(&self, session_id: &str) -> Result<Vec<Message>> {
1179 let mut stmt = self.conn.prepare(
1180 "SELECT role, content, model, reasoning, tokens, secs, cost, phrase, persona, created_at
1181 FROM messages WHERE session_id = ?1 ORDER BY created_at ASC, rowid ASC",
1182 )?;
1183 let messages = stmt
1184 .query_map([session_id], |r| {
1185 Ok(Message {
1186 role: r.get(0)?,
1187 content: r.get(1)?,
1188 model: r.get(2)?,
1189 reasoning: r.get(3)?,
1190 tokens: r.get(4)?,
1191 secs: r.get(5)?,
1192 cost: r.get(6)?,
1193 phrase: r.get(7)?,
1194 persona: r.get(8)?,
1195 created_at: r.get(9)?,
1196 })
1197 })?
1198 .collect::<rusqlite::Result<Vec<_>>>()?;
1199 Ok(messages)
1200 }
1201
1202 pub fn add_user_message(&self, session_id: &str, content: &str) -> Result<String> {
1204 self.insert_message(
1205 session_id, "user", content, None, None, None, None, None, None,
1206 )
1207 }
1208
1209 pub fn add_gate_reply_message(&self, session_id: &str, content: &str) -> Result<String> {
1215 self.insert_message(
1216 session_id,
1217 "gate_reply",
1218 content,
1219 None,
1220 None,
1221 None,
1222 None,
1223 None,
1224 None,
1225 )
1226 }
1227
1228 pub fn add_tool_call_message(&self, session_id: &str, content: &str) -> Result<String> {
1234 self.insert_message(
1235 session_id,
1236 "tool_call",
1237 content,
1238 None,
1239 None,
1240 None,
1241 None,
1242 None,
1243 None,
1244 )
1245 }
1246
1247 pub fn add_error_message(&self, session_id: &str, content: &str) -> Result<String> {
1250 self.insert_message(
1251 session_id, "error", content, None, None, None, None, None, None,
1252 )
1253 }
1254
1255 pub fn add_research_stage_message(&self, session_id: &str, content: &str) -> Result<String> {
1260 self.insert_message(
1261 session_id,
1262 "research_stage",
1263 content,
1264 None,
1265 None,
1266 None,
1267 None,
1268 None,
1269 None,
1270 )
1271 }
1272
1273 pub fn add_research_plan_message(&self, session_id: &str, content: &str) -> Result<String> {
1276 self.insert_message(
1277 session_id,
1278 "research_plan",
1279 content,
1280 None,
1281 None,
1282 None,
1283 None,
1284 None,
1285 None,
1286 )
1287 }
1288
1289 pub fn add_survey_message(&self, session_id: &str, content: &str) -> Result<String> {
1293 self.insert_message(
1294 session_id, "survey", content, None, None, None, None, None, None,
1295 )
1296 }
1297
1298 pub fn upsert_research_stage_message(
1303 &self,
1304 session_id: &str,
1305 label: &str,
1306 detail: &str,
1307 ) -> Result<()> {
1308 let content = stage_content(label, detail);
1309 let existing: Option<String> = self
1310 .conn
1311 .query_row(
1312 "SELECT id FROM messages WHERE session_id = ?1 AND role = 'research_stage'
1313 AND (content = ?2 OR content LIKE ?3)
1314 ORDER BY created_at DESC LIMIT 1",
1315 (session_id, label, format!("{label}:%")),
1316 |r| r.get(0),
1317 )
1318 .ok();
1319 match existing {
1320 Some(id) => {
1321 let now = Utc::now().to_rfc3339();
1322 self.conn.execute(
1323 "UPDATE messages SET content = ?2, created_at = ?3 WHERE id = ?1",
1324 (&id, &content, &now),
1325 )?;
1326 }
1327 None => {
1328 self.add_research_stage_message(session_id, &content)?;
1329 }
1330 }
1331 Ok(())
1332 }
1333
1334 #[cfg(test)]
1338 pub fn add_session_sources(&self, session_id: &str, url_norms: &[String]) -> Result<()> {
1339 add_session_sources(&self.conn, session_id, url_norms)
1340 }
1341
1342 #[cfg(test)]
1344 pub fn search_session_sources(
1345 &self,
1346 session_id: &str,
1347 query: &str,
1348 ) -> Result<Vec<(String, String)>> {
1349 search_session_sources(&self.conn, session_id, query)
1350 }
1351
1352 pub fn set_source_flag(
1358 &self,
1359 session_id: &str,
1360 url_norm: &str,
1361 flag: Option<&str>,
1362 ) -> Result<()> {
1363 self.conn.execute(
1364 "UPDATE session_sources SET flag = ?3, updated_at = ?4
1365 WHERE session_id = ?1 AND url_norm = ?2",
1366 (session_id, url_norm, flag, Utc::now().to_rfc3339()),
1367 )?;
1368 Ok(())
1369 }
1370
1371 #[allow(clippy::too_many_arguments)]
1377 pub fn add_assistant_message(
1378 &self,
1379 session_id: &str,
1380 content: &str,
1381 model: Option<&str>,
1382 reasoning: Option<&str>,
1383 tokens: Option<i64>,
1384 secs: Option<f64>,
1385 cost: Option<f64>,
1386 phrase: Option<&str>,
1387 ) -> Result<String> {
1388 self.insert_message(
1389 session_id,
1390 "assistant",
1391 content,
1392 model,
1393 reasoning,
1394 tokens,
1395 secs,
1396 cost,
1397 phrase,
1398 )
1399 }
1400
1401 pub fn add_persona_message(
1404 &self,
1405 session_id: &str,
1406 content: &str,
1407 persona_name: &str,
1408 model: &str,
1409 ) -> Result<String> {
1410 let id = self.insert_message(
1411 session_id,
1412 "assistant",
1413 content,
1414 Some(model),
1415 None,
1416 None,
1417 None,
1418 None,
1419 None,
1420 )?;
1421 self.conn.execute(
1422 "UPDATE messages SET persona = ?2 WHERE id = ?1",
1423 (&id, persona_name),
1424 )?;
1425 Ok(id)
1426 }
1427
1428 #[allow(clippy::too_many_arguments)]
1431 pub(crate) fn insert_message(
1432 &self,
1433 session_id: &str,
1434 role: &str,
1435 content: &str,
1436 model: Option<&str>,
1437 reasoning: Option<&str>,
1438 tokens: Option<i64>,
1439 secs: Option<f64>,
1440 cost: Option<f64>,
1441 phrase: Option<&str>,
1442 ) -> Result<String> {
1443 let now = Utc::now().to_rfc3339();
1444 let id = Uuid::new_v4().to_string();
1445 self.conn.execute(
1446 "INSERT INTO messages
1447 (id, session_id, role, content, model, reasoning, tokens, secs, cost, phrase, created_at)
1448 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
1449 (
1450 &id, session_id, role, content, model, reasoning, tokens, secs, cost, phrase, &now,
1451 ),
1452 )?;
1453 self.conn.execute(
1454 "UPDATE sessions SET updated_at = ?2 WHERE id = ?1",
1455 (session_id, &now),
1456 )?;
1457 Ok(id)
1458 }
1459
1460 pub fn message_created_at(&self, session_id: &str, index: usize) -> Result<Option<String>> {
1465 let mut stmt = self.conn.prepare(
1466 "SELECT created_at FROM messages WHERE session_id = ?1
1467 ORDER BY created_at ASC, rowid ASC LIMIT 1 OFFSET ?2",
1468 )?;
1469 Ok(stmt
1470 .query_row((session_id, index as i64), |r| r.get(0))
1471 .optional()?)
1472 }
1473
1474 pub fn add_compaction_message(
1481 &self,
1482 session_id: &str,
1483 content: &str,
1484 at: &str,
1485 ) -> Result<String> {
1486 let id = Uuid::new_v4().to_string();
1487 self.conn.execute(
1488 "INSERT INTO messages
1489 (id, session_id, role, content, model, reasoning, tokens, secs, phrase, created_at)
1490 VALUES (?1, ?2, 'compaction', ?3, NULL, NULL, NULL, NULL, NULL, ?4)",
1491 (&id, session_id, content, at),
1492 )?;
1493 Ok(id)
1494 }
1495
1496 pub fn update_compaction_message(&self, session_id: &str, content: &str) -> Result<usize> {
1501 Ok(self.conn.execute(
1502 "UPDATE messages SET content = ?2
1503 WHERE session_id = ?1 AND role = 'compaction'",
1504 (session_id, content),
1505 )?)
1506 }
1507
1508 pub fn set_session_model(&self, session_id: &str, model: &str) -> Result<()> {
1509 self.conn.execute(
1510 "UPDATE sessions SET model = ?2, updated_at = ?3 WHERE id = ?1",
1511 (session_id, model, Utc::now().to_rfc3339()),
1512 )?;
1513 Ok(())
1514 }
1515
1516 pub fn upsert_file(
1524 &self,
1525 space_id: &str,
1526 name: &str,
1527 hash: &str,
1528 size: i64,
1529 status: &str,
1530 ) -> Result<String> {
1531 let now = Utc::now().to_rfc3339();
1532 if let Ok(existing) = self.conn.query_row(
1533 "SELECT id FROM files WHERE space_id = ?1 AND name = ?2",
1534 (space_id, name),
1535 |r| r.get::<_, String>(0),
1536 ) {
1537 self.conn.execute(
1538 "UPDATE files SET hash = ?2, size = ?3, updated_at = ?4 WHERE id = ?1",
1539 (&existing, hash, size, &now),
1540 )?;
1541 self.conn.execute(
1542 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1543 VALUES (?1, 0, ?2, ?3)
1544 ON CONFLICT(file_id) DO UPDATE SET status = excluded.status,
1545 updated_at = excluded.updated_at",
1546 (&existing, status, &now),
1547 )?;
1548 return Ok(existing);
1549 }
1550 let id = Uuid::new_v4().to_string();
1551 self.conn.execute(
1552 "INSERT INTO files (id, space_id, name, hash, size, created_at, updated_at)
1553 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
1554 (&id, space_id, name, hash, size, &now, &now),
1555 )?;
1556 self.conn.execute(
1557 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1558 VALUES (?1, 0, ?2, ?3)",
1559 (&id, status, &now),
1560 )?;
1561 Ok(id)
1562 }
1563
1564 pub fn list_files(&self, space_id: &str) -> Result<Vec<FileRow>> {
1565 let mut stmt = self.conn.prepare(
1566 "SELECT files.id, files.name, files.hash, files.size,
1567 COALESCE(cache.file_index_state.status, 'not indexed'),
1568 COALESCE(cache.file_index_state.mtime, 0)
1569 FROM files
1570 LEFT JOIN cache.file_index_state
1571 ON cache.file_index_state.file_id = files.id
1572 WHERE files.space_id = ?1 ORDER BY files.name ASC",
1573 )?;
1574 let rows = stmt.query_map([space_id], |r| {
1575 Ok(FileRow {
1576 id: r.get(0)?,
1577 name: r.get(1)?,
1578 hash: r.get(2)?,
1579 size: r.get(3)?,
1580 status: r.get(4)?,
1581 mtime: r.get(5)?,
1582 })
1583 })?;
1584 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1585 }
1586
1587 pub fn file_indexed(&self, file_id: &str) -> Result<bool> {
1592 Ok(self.conn.query_row(
1593 "SELECT EXISTS(SELECT 1 FROM cache.file_index_state WHERE file_id = ?1)",
1594 [file_id],
1595 |r| r.get(0),
1596 )?)
1597 }
1598
1599 pub fn file_has_chunks(&self, file_id: &str) -> Result<bool> {
1603 Ok(self.conn.query_row(
1604 "SELECT EXISTS(SELECT 1 FROM cache.file_chunks WHERE file_id = ?1)",
1605 [file_id],
1606 |r| r.get(0),
1607 )?)
1608 }
1609
1610 pub fn delete_file(&self, file_id: &str) -> Result<()> {
1611 self.conn.execute(
1612 "DELETE FROM cache.file_chunks WHERE file_id = ?1",
1613 [file_id],
1614 )?;
1615 self.conn.execute(
1616 "DELETE FROM cache.chunk_embeddings WHERE file_id = ?1",
1617 [file_id],
1618 )?;
1619 self.conn.execute(
1620 "DELETE FROM cache.file_index_state WHERE file_id = ?1",
1621 [file_id],
1622 )?;
1623 self.conn
1624 .execute("DELETE FROM files WHERE id = ?1", [file_id])?;
1625 self.tombstone("files", file_id)?;
1626 Ok(())
1627 }
1628
1629 pub fn set_file_mtime(&self, file_id: &str, mtime: i64) -> Result<()> {
1633 let now = Utc::now().to_rfc3339();
1634 self.conn.execute(
1635 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1636 VALUES (?1, ?2, '', ?3)
1637 ON CONFLICT(file_id) DO UPDATE SET mtime = excluded.mtime,
1638 updated_at = excluded.updated_at",
1639 (file_id, mtime, &now),
1640 )?;
1641 Ok(())
1642 }
1643
1644 pub fn set_file_status(&self, file_id: &str, status: &str) -> Result<()> {
1648 let now = Utc::now().to_rfc3339();
1649 self.conn.execute(
1650 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1651 VALUES (?1, 0, ?2, ?3)
1652 ON CONFLICT(file_id) DO UPDATE SET status = excluded.status,
1653 updated_at = excluded.updated_at",
1654 (file_id, status, &now),
1655 )?;
1656 Ok(())
1657 }
1658
1659 pub fn rename_file(&self, file_id: &str, new_name: &str) -> Result<()> {
1660 self.conn.execute(
1661 "UPDATE files SET name = ?2, updated_at = ?3 WHERE id = ?1",
1662 (file_id, new_name, Utc::now().to_rfc3339()),
1663 )?;
1664 Ok(())
1665 }
1666
1667 pub fn replace_file_ref_in_messages(
1671 &self,
1672 space_id: &str,
1673 old_name: &str,
1674 new_name: &str,
1675 ) -> Result<()> {
1676 self.conn.execute(
1677 "UPDATE messages SET content = REPLACE(content, ?1, ?2)
1678 WHERE session_id IN (SELECT id FROM sessions WHERE space_id = ?3)",
1679 (old_name, new_name, space_id),
1680 )?;
1681 Ok(())
1682 }
1683
1684 pub fn set_file_chunks(&self, file_id: &str, chunks: &[(String, String)]) -> Result<()> {
1689 self.conn.execute(
1690 "DELETE FROM cache.file_chunks WHERE file_id = ?1",
1691 [file_id],
1692 )?;
1693 self.conn.execute(
1694 "DELETE FROM cache.chunk_embeddings WHERE file_id = ?1",
1695 [file_id],
1696 )?;
1697 for (seq, (location, text)) in chunks.iter().enumerate() {
1698 self.conn.execute(
1699 "INSERT INTO cache.file_chunks (file_id, seq, location, text) VALUES (?1, ?2, ?3, ?4)",
1700 (file_id, seq as i64, location, text),
1701 )?;
1702 }
1703 Ok(())
1704 }
1705
1706 #[cfg(test)]
1709 pub fn raw(&self) -> &Connection {
1710 &self.conn
1711 }
1712
1713 pub fn integrity_check(&self) -> Result<String> {
1717 self.conn
1718 .query_row("PRAGMA integrity_check", [], |r| r.get(0))
1719 .context("running integrity check")
1720 }
1721
1722 pub fn file_chunk_texts(&self, file_id: &str) -> Result<Vec<(i64, String)>> {
1724 let mut stmt = self.conn.prepare(
1725 "SELECT CAST(seq AS INTEGER), text FROM cache.file_chunks
1726 WHERE file_id = ?1 ORDER BY CAST(seq AS INTEGER) ASC",
1727 )?;
1728 let rows = stmt.query_map([file_id], |r| Ok((r.get(0)?, r.get(1)?)))?;
1729 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1730 }
1731
1732 pub fn files_missing_embeddings(&self, space_id: &str) -> Result<Vec<String>> {
1734 files_missing_embeddings(&self.conn, space_id)
1735 }
1736
1737 pub fn add_citations(
1741 &self,
1742 space_id: &str,
1743 report_file: &str,
1744 citations: &[(String, Option<String>)],
1745 ) -> Result<()> {
1746 for (url, title) in citations {
1747 self.conn.execute(
1748 "INSERT INTO citations (sync_id, space_id, report_file, url, title)
1749 VALUES (?1, ?2, ?3, ?4, ?5)",
1750 (
1751 Uuid::new_v4().to_string(),
1752 space_id,
1753 report_file,
1754 url,
1755 title,
1756 ),
1757 )?;
1758 }
1759 Ok(())
1760 }
1761
1762 pub fn search_citations(
1767 &self,
1768 space_id: &str,
1769 query: Option<&str>,
1770 ) -> Result<Vec<(String, String, String)>> {
1771 search_citations(&self.conn, space_id, query)
1772 }
1773
1774 pub fn set_chunk_embeddings(&self, file_id: &str, vecs: &[(i64, Vec<f32>)]) -> Result<()> {
1776 for (seq, v) in vecs {
1777 self.conn.execute(
1778 "INSERT OR REPLACE INTO cache.chunk_embeddings (file_id, seq, vec) VALUES (?1, ?2, ?3)",
1779 (file_id, seq, vec_to_blob(v)),
1780 )?;
1781 }
1782 Ok(())
1783 }
1784
1785 pub fn clear_chunk_embeddings(&self) -> Result<()> {
1789 self.conn
1790 .execute("DELETE FROM cache.chunk_embeddings", [])?;
1791 Ok(())
1792 }
1793
1794 pub fn create_watch(
1795 &self,
1796 space_id: &str,
1797 topic: &str,
1798 interval_hours: i64,
1799 session_id: &str,
1800 ) -> Result<String> {
1801 let id = Uuid::new_v4().to_string();
1802 let now = Utc::now().to_rfc3339();
1803 self.conn.execute(
1804 "INSERT INTO watches (id, space_id, topic, interval_hours, session_id, last_run_at, updated_at)
1805 VALUES (?1, ?2, ?3, ?4, ?5, NULL, ?6)",
1806 (&id, space_id, topic, interval_hours, session_id, &now),
1807 )?;
1808 Ok(id)
1809 }
1810
1811 pub fn list_watches(&self, space_id: &str) -> Result<Vec<Watch>> {
1812 let mut stmt = self.conn.prepare(
1813 "SELECT id, space_id, topic, interval_hours, session_id, last_run_at
1814 FROM watches WHERE space_id = ?1 ORDER BY topic",
1815 )?;
1816 let rows = stmt.query_map([space_id], |r| {
1817 Ok(Watch {
1818 id: r.get(0)?,
1819 space_id: r.get(1)?,
1820 topic: r.get(2)?,
1821 interval_hours: r.get(3)?,
1822 session_id: r.get(4)?,
1823 last_run_at: r.get(5)?,
1824 })
1825 })?;
1826 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1827 }
1828
1829 pub fn list_all_watches(&self) -> Result<Vec<Watch>> {
1832 let mut stmt = self.conn.prepare(
1833 "SELECT id, space_id, topic, interval_hours, session_id, last_run_at FROM watches",
1834 )?;
1835 let rows = stmt.query_map([], |r| {
1836 Ok(Watch {
1837 id: r.get(0)?,
1838 space_id: r.get(1)?,
1839 topic: r.get(2)?,
1840 interval_hours: r.get(3)?,
1841 session_id: r.get(4)?,
1842 last_run_at: r.get(5)?,
1843 })
1844 })?;
1845 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1846 }
1847
1848 pub fn touch_watch(&self, id: &str, now_rfc3339: &str) -> Result<()> {
1849 self.conn.execute(
1850 "UPDATE watches SET last_run_at = ?2, updated_at = ?3 WHERE id = ?1",
1851 (id, now_rfc3339, Utc::now().to_rfc3339()),
1852 )?;
1853 Ok(())
1854 }
1855
1856 pub fn set_watch_session(&self, id: &str, session_id: &str) -> Result<()> {
1860 self.conn.execute(
1861 "UPDATE watches SET session_id = ?2, updated_at = ?3 WHERE id = ?1",
1862 (id, session_id, Utc::now().to_rfc3339()),
1863 )?;
1864 Ok(())
1865 }
1866
1867 pub fn delete_watch(&self, id: &str) -> Result<()> {
1868 self.conn
1869 .execute("DELETE FROM watches WHERE id = ?1", [id])?;
1870 self.tombstone("watches", id)?;
1871 Ok(())
1872 }
1873}
1874
1875pub fn vec_to_blob(v: &[f32]) -> Vec<u8> {
1877 v.iter().flat_map(|f| f.to_le_bytes()).collect()
1878}
1879
1880pub fn blob_to_vec(b: &[u8]) -> Vec<f32> {
1882 b.chunks_exact(4)
1883 .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
1884 .collect()
1885}
1886
1887pub fn search_citations(
1892 conn: &Connection,
1893 space_id: &str,
1894 query: Option<&str>,
1895) -> Result<Vec<(String, String, String)>> {
1896 let mut stmt = conn.prepare(
1897 "SELECT report_file, url, COALESCE(title, '') FROM citations
1898 WHERE space_id = ?1
1899 AND (?2 IS NULL OR url LIKE ?2 OR title LIKE ?2 OR report_file LIKE ?2)
1900 ORDER BY id DESC",
1901 )?;
1902 let pattern = query.map(|q| format!("%{q}%"));
1903 let rows = stmt.query_map((space_id, pattern), |r| {
1904 Ok((r.get(0)?, r.get(1)?, r.get(2)?))
1905 })?;
1906 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1907}
1908
1909pub fn stage_content(label: &str, detail: &str) -> String {
1911 if detail.is_empty() {
1912 label.to_string()
1913 } else {
1914 format!("{label}: {detail}")
1915 }
1916}
1917
1918pub fn add_session_sources(
1921 conn: &Connection,
1922 session_id: &str,
1923 url_norms: &[String],
1924) -> Result<()> {
1925 let now = Utc::now().to_rfc3339();
1926 for u in url_norms {
1927 conn.execute(
1928 "INSERT OR IGNORE INTO session_sources (session_id, url_norm, updated_at)
1929 VALUES (?1, ?2, ?3)",
1930 (session_id, u, &now),
1931 )?;
1932 }
1933 Ok(())
1934}
1935
1936pub fn search_session_sources(
1941 conn: &Connection,
1942 session_id: &str,
1943 query: &str,
1944) -> Result<Vec<(String, String)>> {
1945 let mut stmt = conn.prepare(
1946 "SELECT cache.web_cache.url, cache.web_cache.text FROM session_sources
1947 JOIN cache.web_cache ON cache.web_cache.url_norm = session_sources.url_norm
1948 WHERE session_sources.session_id = ?1",
1949 )?;
1950 let rows = stmt.query_map([session_id], |r| {
1951 Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
1952 })?;
1953 let needle = query.to_lowercase();
1954 Ok(rows
1955 .collect::<rusqlite::Result<Vec<_>>>()?
1956 .into_iter()
1957 .filter(|(_, text)| text.to_lowercase().contains(&needle))
1958 .collect())
1959}
1960
1961pub fn pinned_urls(conn: &Connection, session_id: &str) -> Result<Vec<String>> {
1964 let mut stmt = conn.prepare(
1965 "SELECT url_norm FROM session_sources WHERE session_id = ?1 AND flag = 'pinned'",
1966 )?;
1967 let rows = stmt.query_map([session_id], |r| r.get::<_, String>(0))?;
1968 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1969}
1970
1971pub fn discarded_domains(conn: &Connection, session_id: &str) -> Result<Vec<String>> {
1974 let mut stmt = conn.prepare(
1975 "SELECT url_norm FROM session_sources WHERE session_id = ?1 AND flag = 'discarded'",
1976 )?;
1977 let rows: Vec<String> = stmt
1978 .query_map([session_id], |r| r.get::<_, String>(0))?
1979 .collect::<rusqlite::Result<Vec<_>>>()?;
1980 let mut hosts: Vec<String> = rows
1981 .iter()
1982 .filter_map(|u| {
1983 reqwest::Url::parse(u)
1984 .ok()
1985 .and_then(|p| p.host_str().map(str::to_string))
1986 })
1987 .collect();
1988 hosts.sort();
1989 hosts.dedup();
1990 Ok(hosts)
1991}
1992
1993pub fn is_fresh(fetched_at: &str, now: chrono::DateTime<Utc>) -> bool {
1997 chrono::DateTime::parse_from_rfc3339(fetched_at)
1998 .is_ok_and(|dt| now.signed_duration_since(dt) < chrono::Duration::hours(24))
1999}
2000
2001pub fn cache_get(conn: &Connection, url_norm: &str) -> Result<Option<(String, String, String)>> {
2008 let row = conn.query_row(
2009 "SELECT COALESCE(title, ''), text, fetched_at FROM web_cache WHERE url_norm = ?1",
2010 [url_norm],
2011 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
2012 );
2013 match row {
2014 Ok(v) => Ok(Some(v)),
2015 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
2016 Err(e) => Err(e.into()),
2017 }
2018}
2019
2020pub fn cache_put(
2022 conn: &Connection,
2023 url_norm: &str,
2024 url: &str,
2025 title: Option<&str>,
2026 text: &str,
2027) -> Result<()> {
2028 let now = Utc::now().to_rfc3339();
2029 conn.execute(
2030 "INSERT INTO web_cache (url_norm, url, title, text, fetched_at) VALUES (?1, ?2, ?3, ?4, ?5)
2031 ON CONFLICT(url_norm) DO UPDATE SET url = ?2, title = ?3, text = ?4, fetched_at = ?5",
2032 (url_norm, url, title, text, &now),
2033 )?;
2034 Ok(())
2035}
2036
2037pub fn files_missing_embeddings(conn: &Connection, space_id: &str) -> Result<Vec<String>> {
2044 let mut stmt = conn.prepare(
2045 "SELECT files.id
2046 FROM files
2047 WHERE files.space_id = ?1
2048 AND (
2049 NOT EXISTS (
2050 SELECT 1 FROM cache.file_chunks
2051 WHERE cache.file_chunks.file_id = files.id
2052 )
2053 OR EXISTS (
2054 SELECT 1
2055 FROM cache.file_chunks
2056 LEFT JOIN cache.chunk_embeddings
2057 ON cache.chunk_embeddings.file_id = cache.file_chunks.file_id
2058 AND cache.chunk_embeddings.seq = CAST(cache.file_chunks.seq AS INTEGER)
2059 WHERE cache.file_chunks.file_id = files.id
2060 AND cache.chunk_embeddings.file_id IS NULL
2061 )
2062 )
2063 ORDER BY files.name",
2064 )?;
2065 let rows = stmt.query_map([space_id], |r| r.get(0))?;
2066 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2067}
2068
2069pub fn semantic_chunks(
2074 conn: &Connection,
2075 space_id: &str,
2076 query: &[f32],
2077 limit: usize,
2078) -> Result<Vec<(String, String, String, f32)>> {
2079 let mut stmt = conn.prepare(
2080 "SELECT files.name, cache.file_chunks.location, cache.file_chunks.text,
2081 cache.chunk_embeddings.vec
2082 FROM cache.chunk_embeddings
2083 JOIN files ON files.id = cache.chunk_embeddings.file_id
2084 JOIN cache.file_chunks
2085 ON cache.file_chunks.file_id = cache.chunk_embeddings.file_id
2086 AND CAST(cache.file_chunks.seq AS INTEGER) = cache.chunk_embeddings.seq
2087 WHERE files.space_id = ?1",
2088 )?;
2089 let rows = stmt.query_map([space_id], |r| {
2090 Ok((
2091 r.get::<_, String>(0)?,
2092 r.get::<_, String>(1)?,
2093 r.get::<_, String>(2)?,
2094 r.get::<_, Vec<u8>>(3)?,
2095 ))
2096 })?;
2097 let mut hits: Vec<(String, String, String, f32)> = Vec::new();
2098 for row in rows {
2099 let (name, loc, text, blob) = row?;
2100 let v = blob_to_vec(&blob);
2101 if v.len() != query.len() {
2102 continue;
2103 }
2104 let score = cosine(query, &v);
2105 hits.push((name, loc, text, score));
2106 }
2107 hits.sort_by(|a, b| b.3.partial_cmp(&a.3).unwrap_or(std::cmp::Ordering::Equal));
2108 hits.truncate(limit);
2109 Ok(hits)
2110}
2111
2112fn cosine(a: &[f32], b: &[f32]) -> f32 {
2113 let (mut dot, mut na, mut nb) = (0.0f32, 0.0f32, 0.0f32);
2114 for (x, y) in a.iter().zip(b) {
2115 dot += x * y;
2116 na += x * x;
2117 nb += y * y;
2118 }
2119 let denom = na.sqrt() * nb.sqrt();
2120 if denom == 0.0 { 0.0 } else { dot / denom }
2121}
2122
2123#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2127pub enum UsageRange {
2128 Day,
2129 Week,
2130 Month,
2131 #[default]
2132 All,
2133}
2134
2135impl UsageRange {
2136 pub const CYCLE: [Self; 4] = [Self::Day, Self::Week, Self::Month, Self::All];
2138
2139 pub const fn label(self) -> &'static str {
2141 match self {
2142 Self::Day => "24h",
2143 Self::Week => "7d",
2144 Self::Month => "30d",
2145 Self::All => "all",
2146 }
2147 }
2148
2149 pub const fn title(self) -> &'static str {
2151 match self {
2152 Self::Day => "last 24 hours",
2153 Self::Week => "last 7 days",
2154 Self::Month => "last 30 days",
2155 Self::All => "all time",
2156 }
2157 }
2158
2159 pub const fn key(self) -> &'static str {
2161 match self {
2162 Self::Day => "day",
2163 Self::Week => "week",
2164 Self::Month => "month",
2165 Self::All => "all",
2166 }
2167 }
2168
2169 pub fn from_key(key: &str) -> Self {
2172 Self::CYCLE
2173 .iter()
2174 .copied()
2175 .find(|r| r.key() == key)
2176 .unwrap_or_default()
2177 }
2178
2179 #[must_use]
2181 pub const fn next(self) -> Self {
2182 match self {
2183 Self::Day => Self::Week,
2184 Self::Week => Self::Month,
2185 Self::Month => Self::All,
2186 Self::All => Self::Day,
2187 }
2188 }
2189
2190 #[must_use]
2192 pub const fn prev(self) -> Self {
2193 match self {
2194 Self::Day => Self::All,
2195 Self::Week => Self::Day,
2196 Self::Month => Self::Week,
2197 Self::All => Self::Month,
2198 }
2199 }
2200
2201 pub fn since(self) -> Option<chrono::DateTime<chrono::Utc>> {
2203 use chrono::{Duration, Utc};
2204 match self {
2205 Self::Day => Some(Utc::now() - Duration::hours(24)),
2206 Self::Week => Some(Utc::now() - Duration::days(7)),
2207 Self::Month => Some(Utc::now() - Duration::days(30)),
2208 Self::All => None,
2209 }
2210 }
2211
2212 pub const fn empty_message(self) -> &'static str {
2214 match self {
2215 Self::Day => "no usage in the last 24 hours — ←/→ for a wider window",
2216 Self::Week => "no usage in the last 7 days — ←/→ for a wider window",
2217 Self::Month => "no usage in the last 30 days — ←/→ for a wider window",
2218 Self::All => "no usage logged yet — send a message first",
2219 }
2220 }
2221}
2222
2223pub fn price_name(model: &str) -> &str {
2228 let stripped = ["go:", "openai:", "codex:", "opencode:"]
2229 .iter()
2230 .find_map(|p| model.strip_prefix(p))
2231 .unwrap_or(model);
2232 stripped.rsplit('/').next().unwrap_or(stripped)
2233}
2234
2235fn catalog_price<'a>(
2240 prices: &'a std::collections::HashMap<String, ModelPricing>,
2241 model: &str,
2242) -> Option<&'a ModelPricing> {
2243 if let Some(price) = prices.get(model) {
2244 return Some(price);
2245 }
2246 let name = price_name(model);
2247 if name.is_empty() {
2248 return None;
2249 }
2250 prices
2251 .iter()
2252 .filter(|(id, _)| {
2253 id.strip_suffix(name)
2254 .is_some_and(|rest| rest.ends_with('/'))
2255 })
2256 .min_by_key(|(id, _)| id.len()) .map(|(_, price)| price)
2258}
2259
2260fn catalog_request_cost(
2263 price: ModelPricing,
2264 prompt_tokens: u64,
2265 completion_tokens: u64,
2266 cache_read_tokens: u64,
2267 cache_creation_tokens: u64,
2268) -> f64 {
2269 let reads = cache_read_tokens.min(prompt_tokens);
2270 let writes = cache_creation_tokens.min(prompt_tokens - reads);
2271 let ordinary = prompt_tokens - reads - writes;
2272 let read_price = price.cache_read.unwrap_or(price.prompt);
2273 let write_price = price.cache_write.unwrap_or(price.prompt);
2274 (ordinary as f64 * price.prompt
2275 + reads as f64 * read_price
2276 + writes as f64 * write_price
2277 + completion_tokens as f64 * price.completion)
2278 / 1e6
2279}
2280
2281struct CostBackfillRow {
2282 id: i64,
2283 backend: String,
2284 model: String,
2285 prompt_tokens: u64,
2286 completion_tokens: u64,
2287 cache_read_tokens: u64,
2288 cache_creation_tokens: u64,
2289 old_cost: Option<f64>,
2290 cost_is_provider: Option<bool>,
2292}
2293
2294#[derive(Default)]
2296pub struct UsageTotals {
2297 pub requests: u64,
2298 pub prompt_tokens: u64,
2299 pub completion_tokens: u64,
2300 pub cache_read_tokens: u64,
2301 pub cache_creation_tokens: u64,
2302 pub cost: f64,
2304}
2305
2306#[derive(Default)]
2308pub struct UsageDay {
2309 pub day: String,
2311 pub requests: u64,
2312 pub prompt_tokens: u64,
2313 pub completion_tokens: u64,
2314 pub cache_read_tokens: u64,
2315 pub cost: f64,
2316}
2317
2318#[derive(Default)]
2320pub struct UsageByBackend {
2321 pub backend: String,
2322 pub requests: u64,
2323 pub prompt_tokens: u64,
2324 pub completion_tokens: u64,
2325 pub cache_read_tokens: u64,
2326 pub cost: f64,
2327}
2328
2329#[derive(Default)]
2331pub struct UsageByModel {
2332 pub model: String,
2333 pub requests: u64,
2334 pub prompt_tokens: u64,
2335 pub completion_tokens: u64,
2336 pub cache_read_tokens: u64,
2337 pub cost: f64,
2338}
2339
2340#[derive(Default)]
2342pub struct UsageRow {
2343 pub created_at: String,
2344 pub backend: String,
2345 pub model: String,
2346 pub prompt_tokens: u64,
2347 pub completion_tokens: u64,
2348 pub cache_read_tokens: u64,
2349 pub cost: Option<f64>,
2350}
2351
2352impl Db {
2353 #[allow(clippy::too_many_arguments)]
2356 pub fn log_usage(
2357 &self,
2358 backend: &str,
2359 model: &str,
2360 prompt_tokens: u64,
2361 completion_tokens: u64,
2362 cache_read_tokens: u64,
2363 cache_creation_tokens: u64,
2364 cost: Option<f64>,
2365 cost_is_provider: bool,
2366 session_id: Option<&str>,
2367 space_id: Option<&str>,
2368 ) -> Result<i64> {
2369 self.conn.execute(
2370 "INSERT INTO usage_log (sync_id, created_at, session_id, space_id, backend, model,
2371 prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens,
2372 cost, cost_is_provider, updated_at)
2373 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
2374 (
2375 Uuid::new_v4().to_string(),
2376 Utc::now().to_rfc3339(),
2377 session_id,
2378 space_id,
2379 backend,
2380 model,
2381 prompt_tokens as i64,
2382 completion_tokens as i64,
2383 cache_read_tokens as i64,
2384 cache_creation_tokens as i64,
2385 cost,
2386 i64::from(cost_is_provider),
2387 Utc::now().to_rfc3339(),
2388 ),
2389 )?;
2390 Ok(self.conn.last_insert_rowid())
2391 }
2392
2393 #[allow(clippy::too_many_arguments)]
2398 pub fn update_usage(
2399 &self,
2400 row_id: i64,
2401 prompt_tokens: u64,
2402 completion_tokens: u64,
2403 cache_read_tokens: u64,
2404 cache_creation_tokens: u64,
2405 cost: Option<f64>,
2406 cost_is_provider: bool,
2407 ) -> Result<()> {
2408 self.conn.execute(
2409 "UPDATE usage_log SET prompt_tokens = ?1, completion_tokens = ?2,
2410 cache_read_tokens = ?3, cache_creation_tokens = ?4, cost = ?5,
2411 cost_is_provider = ?6, updated_at = ?7
2412 WHERE id = ?8",
2413 (
2414 prompt_tokens as i64,
2415 completion_tokens as i64,
2416 cache_read_tokens as i64,
2417 cache_creation_tokens as i64,
2418 cost,
2419 i64::from(cost_is_provider),
2420 Utc::now().to_rfc3339(),
2421 row_id,
2422 ),
2423 )?;
2424 Ok(())
2425 }
2426
2427 pub fn request_cost(
2432 &self,
2433 model: &str,
2434 prompt_tokens: u64,
2435 completion_tokens: u64,
2436 cache_read_tokens: u64,
2437 cache_creation_tokens: u64,
2438 ) -> Option<f64> {
2439 self.model_price(model).map(|price| {
2440 catalog_request_cost(
2441 price,
2442 prompt_tokens,
2443 completion_tokens,
2444 cache_read_tokens,
2445 cache_creation_tokens,
2446 )
2447 })
2448 }
2449
2450 pub fn backfill_usage_costs(&mut self) -> Result<usize> {
2462 let max_price: f64 = self.conn.query_row(
2466 "SELECT COALESCE(MAX(prompt_price), 0) FROM cache.model_prices",
2467 [],
2468 |r| r.get(0),
2469 )?;
2470 if max_price > 0.0 && max_price < 0.001 {
2471 self.conn.execute(
2472 "UPDATE cache.model_prices SET
2473 prompt_price = prompt_price * 1e6,
2474 completion_price = completion_price * 1e6,
2475 cache_read_price = cache_read_price * 1e6,
2476 cache_write_price = cache_write_price * 1e6",
2477 [],
2478 )?;
2479 }
2480 let mut prices: std::collections::HashMap<String, ModelPricing> =
2483 std::collections::HashMap::default();
2484 {
2485 let mut stmt = self.conn.prepare(
2486 "SELECT model_id, prompt_price, completion_price,
2487 cache_read_price, cache_write_price
2488 FROM cache.model_prices",
2489 )?;
2490 let rows = stmt.query_map([], |r| {
2491 Ok((
2492 r.get::<_, String>(0)?,
2493 ModelPricing {
2494 prompt: r.get(1)?,
2495 completion: r.get(2)?,
2496 cache_read: r.get(3)?,
2497 cache_write: r.get(4)?,
2498 },
2499 ))
2500 })?;
2501 for row in rows {
2502 let (model, price) = row?;
2503 prices.insert(model, price);
2504 }
2505 }
2506 let rows: Vec<CostBackfillRow> = {
2507 let mut stmt = self.conn.prepare(
2508 "SELECT id, backend, model, prompt_tokens, completion_tokens,
2509 cache_read_tokens, cache_creation_tokens, cost,
2510 cost_is_provider
2511 FROM usage_log",
2512 )?;
2513 let rows = stmt.query_map([], |r| {
2514 Ok(CostBackfillRow {
2515 id: r.get(0)?,
2516 backend: r.get(1)?,
2517 model: r.get(2)?,
2518 prompt_tokens: r.get::<_, i64>(3)? as u64,
2519 completion_tokens: r.get::<_, i64>(4)? as u64,
2520 cache_read_tokens: r.get::<_, i64>(5)? as u64,
2521 cache_creation_tokens: r.get::<_, i64>(6)? as u64,
2522 old_cost: r.get(7)?,
2523 cost_is_provider: r.get::<_, Option<i64>>(8)?.map(|v| v != 0),
2524 })
2525 })?;
2526 rows.collect::<rusqlite::Result<Vec<_>>>()?
2527 };
2528 let tx = self.conn.transaction()?;
2529 {
2530 let mut update = tx.prepare(
2531 "UPDATE usage_log SET cost = ?2, cost_is_provider = 0, updated_at = ?3 WHERE id = ?1",
2532 )?;
2533 for row in &rows {
2534 let legacy_opencode_cost = row.cost_is_provider.is_none()
2538 && row.backend == "OpenCode Go"
2539 && row.old_cost.is_some();
2540 if row.cost_is_provider == Some(true) || legacy_opencode_cost {
2541 continue;
2542 }
2543 let recomputed = catalog_price(&prices, &row.model).map(|price| {
2544 catalog_request_cost(
2545 *price,
2546 row.prompt_tokens,
2547 row.completion_tokens,
2548 row.cache_read_tokens,
2549 row.cache_creation_tokens,
2550 )
2551 });
2552 let write = match (recomputed, row.old_cost) {
2555 (Some(cost), None) => Some(cost),
2556 (Some(cost), Some(old)) if (cost - old).abs() > 1e-12 => Some(cost),
2557 _ => None,
2558 };
2559 if let Some(cost) = write {
2560 update.execute((row.id, cost, Utc::now().to_rfc3339()))?;
2561 }
2562 }
2563 }
2564 tx.commit()?;
2565 Ok(rows.len())
2566 }
2567
2568 pub fn upsert_model_prices(&mut self, prices: &[(String, String, ModelPricing)]) -> Result<()> {
2573 if prices.is_empty() {
2574 return Ok(());
2575 }
2576 let tx = self.conn.transaction()?;
2577 {
2578 let mut stmt = tx.prepare(
2579 "INSERT INTO cache.model_prices (model_id, backend, prompt_price, completion_price,
2580 cache_read_price, cache_write_price, updated_at)
2581 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
2582 ON CONFLICT(model_id) DO UPDATE SET
2583 prompt_price = excluded.prompt_price,
2584 completion_price = excluded.completion_price,
2585 cache_read_price = excluded.cache_read_price,
2586 cache_write_price = excluded.cache_write_price,
2587 updated_at = excluded.updated_at
2588 WHERE cache.model_prices.prompt_price != excluded.prompt_price
2589 OR cache.model_prices.completion_price != excluded.completion_price
2590 OR cache.model_prices.cache_read_price IS NOT excluded.cache_read_price
2591 OR cache.model_prices.cache_write_price IS NOT excluded.cache_write_price",
2592 )?;
2593 let now = Utc::now().to_rfc3339();
2594 for (model, backend, price) in prices {
2595 stmt.execute((
2596 model,
2597 backend,
2598 price.prompt,
2599 price.completion,
2600 price.cache_read,
2601 price.cache_write,
2602 &now,
2603 ))?;
2604 }
2605 }
2606 tx.commit()?;
2607 Ok(())
2608 }
2609
2610 pub fn model_price(&self, model: &str) -> Option<ModelPricing> {
2617 let read_price = |r: &rusqlite::Row| {
2618 Ok(ModelPricing {
2619 prompt: r.get(0)?,
2620 completion: r.get(1)?,
2621 cache_read: r.get(2)?,
2622 cache_write: r.get(3)?,
2623 })
2624 };
2625 if let Ok(price) = self.conn.query_row(
2626 "SELECT prompt_price, completion_price, cache_read_price, cache_write_price
2627 FROM cache.model_prices WHERE model_id = ?1",
2628 [model],
2629 read_price,
2630 ) {
2631 return Some(price);
2632 }
2633 let name = price_name(model);
2634 if name.is_empty() {
2635 return None;
2636 }
2637 self.conn
2640 .query_row(
2641 "SELECT prompt_price, completion_price, cache_read_price, cache_write_price
2642 FROM cache.model_prices
2643 WHERE backend = 'OpenRouter'
2644 AND substr(model_id, -length(?1) - 1) = '/' || ?1
2645 ORDER BY length(model_id) LIMIT 1",
2646 [name],
2647 read_price,
2648 )
2649 .ok()
2650 }
2651
2652 pub fn usage_totals(&self, since: Option<&str>) -> Result<UsageTotals> {
2657 let mut sql = String::from(
2658 "SELECT COUNT(*),
2659 COALESCE(SUM(prompt_tokens), 0),
2660 COALESCE(SUM(completion_tokens), 0),
2661 COALESCE(SUM(cache_read_tokens), 0),
2662 COALESCE(SUM(cache_creation_tokens), 0),
2663 COALESCE(SUM(cost), 0)
2664 FROM usage_log",
2665 );
2666 if since.is_some() {
2667 sql.push_str(" WHERE created_at >= ?1");
2668 }
2669 let map = |r: &rusqlite::Row| {
2670 Ok(UsageTotals {
2671 requests: r.get::<_, i64>(0)? as u64,
2672 prompt_tokens: r.get::<_, i64>(1)? as u64,
2673 completion_tokens: r.get::<_, i64>(2)? as u64,
2674 cache_read_tokens: r.get::<_, i64>(3)? as u64,
2675 cache_creation_tokens: r.get::<_, i64>(4)? as u64,
2676 cost: r.get::<_, f64>(5)?,
2677 })
2678 };
2679 let totals = match since {
2680 Some(s) => self.conn.query_row(&sql, [s], map),
2681 None => self.conn.query_row(&sql, [], map),
2682 }?;
2683 Ok(totals)
2684 }
2685
2686 pub fn usage_by_backend(&self, since: Option<&str>) -> Result<Vec<UsageByBackend>> {
2689 let mut sql = String::from(
2690 "SELECT backend, COUNT(*),
2691 COALESCE(SUM(prompt_tokens), 0),
2692 COALESCE(SUM(completion_tokens), 0),
2693 COALESCE(SUM(cache_read_tokens), 0),
2694 COALESCE(SUM(cost), 0)
2695 FROM usage_log",
2696 );
2697 if since.is_some() {
2698 sql.push_str(" WHERE created_at >= ?1");
2699 }
2700 sql.push_str(" GROUP BY backend ORDER BY COUNT(*) DESC");
2701 let map = |r: &rusqlite::Row| {
2702 Ok(UsageByBackend {
2703 backend: r.get(0)?,
2704 requests: r.get::<_, i64>(1)? as u64,
2705 prompt_tokens: r.get::<_, i64>(2)? as u64,
2706 completion_tokens: r.get::<_, i64>(3)? as u64,
2707 cache_read_tokens: r.get::<_, i64>(4)? as u64,
2708 cost: r.get::<_, f64>(5)?,
2709 })
2710 };
2711 let mut stmt = self.conn.prepare(&sql)?;
2712 let rows = match since {
2713 Some(s) => stmt.query_map([s], map),
2714 None => stmt.query_map([], map),
2715 }?;
2716 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2717 }
2718
2719 pub fn usage_by_model(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageByModel>> {
2722 let mut sql = String::from(
2723 "SELECT model, COUNT(*),
2724 COALESCE(SUM(prompt_tokens), 0),
2725 COALESCE(SUM(completion_tokens), 0),
2726 COALESCE(SUM(cache_read_tokens), 0),
2727 COALESCE(SUM(cost), 0)
2728 FROM usage_log",
2729 );
2730 if since.is_some() {
2731 sql.push_str(" WHERE created_at >= ?1");
2732 sql.push_str(" GROUP BY model ORDER BY COUNT(*) DESC LIMIT ?2");
2733 } else {
2734 sql.push_str(" GROUP BY model ORDER BY COUNT(*) DESC LIMIT ?1");
2735 }
2736 let map = |r: &rusqlite::Row| {
2737 Ok(UsageByModel {
2738 model: r.get(0)?,
2739 requests: r.get::<_, i64>(1)? as u64,
2740 prompt_tokens: r.get::<_, i64>(2)? as u64,
2741 completion_tokens: r.get::<_, i64>(3)? as u64,
2742 cache_read_tokens: r.get::<_, i64>(4)? as u64,
2743 cost: r.get::<_, f64>(5)?,
2744 })
2745 };
2746 let mut stmt = self.conn.prepare(&sql)?;
2747 let rows = match since {
2748 Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
2749 None => stmt.query_map(rusqlite::params![limit as i64], map),
2750 }?;
2751 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2752 }
2753
2754 pub fn usage_recent(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageRow>> {
2757 let mut sql = String::from(
2758 "SELECT created_at, backend, model, prompt_tokens, completion_tokens,
2759 cache_read_tokens, cost
2760 FROM usage_log",
2761 );
2762 if since.is_some() {
2763 sql.push_str(" WHERE created_at >= ?1");
2764 sql.push_str(" ORDER BY id DESC LIMIT ?2");
2765 } else {
2766 sql.push_str(" ORDER BY id DESC LIMIT ?1");
2767 }
2768 let map = |r: &rusqlite::Row| {
2769 Ok(UsageRow {
2770 created_at: r.get(0)?,
2771 backend: r.get(1)?,
2772 model: r.get(2)?,
2773 prompt_tokens: r.get::<_, i64>(3)? as u64,
2774 completion_tokens: r.get::<_, i64>(4)? as u64,
2775 cache_read_tokens: r.get::<_, i64>(5)? as u64,
2776 cost: r.get(6)?,
2777 })
2778 };
2779 let mut stmt = self.conn.prepare(&sql)?;
2780 let rows = match since {
2781 Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
2782 None => stmt.query_map(rusqlite::params![limit as i64], map),
2783 }?;
2784 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2785 }
2786
2787 pub fn usage_by_day(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageDay>> {
2790 let mut sql = String::from(
2791 "SELECT substr(created_at, 1, 10) AS day, COUNT(*),
2792 COALESCE(SUM(prompt_tokens), 0),
2793 COALESCE(SUM(completion_tokens), 0),
2794 COALESCE(SUM(cache_read_tokens), 0),
2795 COALESCE(SUM(cost), 0)
2796 FROM usage_log",
2797 );
2798 if since.is_some() {
2799 sql.push_str(" WHERE created_at >= ?1");
2800 sql.push_str(" GROUP BY day ORDER BY day DESC LIMIT ?2");
2801 } else {
2802 sql.push_str(" GROUP BY day ORDER BY day DESC LIMIT ?1");
2803 }
2804 let map = |r: &rusqlite::Row| {
2805 Ok(UsageDay {
2806 day: r.get(0)?,
2807 requests: r.get::<_, i64>(1)? as u64,
2808 prompt_tokens: r.get::<_, i64>(2)? as u64,
2809 completion_tokens: r.get::<_, i64>(3)? as u64,
2810 cache_read_tokens: r.get::<_, i64>(4)? as u64,
2811 cost: r.get(5)?,
2812 })
2813 };
2814 let mut stmt = self.conn.prepare(&sql)?;
2815 let rows = match since {
2816 Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
2817 None => stmt.query_map(rusqlite::params![limit as i64], map),
2818 }?;
2819 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2820 }
2821}
2822
2823#[allow(dead_code)]
2831#[derive(Debug, Clone, PartialEq, Eq)]
2832pub struct SyncState {
2833 pub peer_id: String,
2834 pub table_name: String,
2835 pub pull_cursor: Option<String>,
2836 pub push_cursor: Option<String>,
2837 pub last_synced_at: Option<String>,
2838}
2839
2840#[allow(dead_code)]
2842impl Db {
2843 pub fn device_id(&self) -> Result<String> {
2847 if let Some(id) = self
2848 .conn
2849 .query_row("SELECT device_id FROM device_meta LIMIT 1", [], |r| {
2850 r.get(0)
2851 })
2852 .optional()?
2853 {
2854 return Ok(id);
2855 }
2856 let id = Uuid::new_v4().to_string();
2857 let now = Utc::now().to_rfc3339();
2858 self.conn.execute(
2859 "INSERT INTO device_meta (device_id, created_at) VALUES (?1, ?2)",
2860 (&id, &now),
2861 )?;
2862 Ok(id)
2863 }
2864
2865 pub fn set_sync_state(
2868 &self,
2869 peer_id: &str,
2870 table_name: &str,
2871 pull_cursor: Option<&str>,
2872 push_cursor: Option<&str>,
2873 ) -> Result<()> {
2874 self.conn.execute(
2875 "INSERT INTO sync_state (peer_id, table_name, pull_cursor, push_cursor, last_synced_at)
2876 VALUES (?1, ?2, ?3, ?4, ?5)
2877 ON CONFLICT(peer_id, table_name) DO UPDATE SET
2878 pull_cursor = COALESCE(?3, pull_cursor),
2879 push_cursor = COALESCE(?4, push_cursor),
2880 last_synced_at = ?5",
2881 (
2882 peer_id,
2883 table_name,
2884 pull_cursor,
2885 push_cursor,
2886 Utc::now().to_rfc3339(),
2887 ),
2888 )?;
2889 Ok(())
2890 }
2891
2892 pub fn load_sync_state(&self) -> Result<Vec<SyncState>> {
2893 let mut stmt = self.conn.prepare(
2894 "SELECT peer_id, table_name, pull_cursor, push_cursor, last_synced_at
2895 FROM sync_state ORDER BY peer_id, table_name",
2896 )?;
2897 let rows = stmt.query_map([], |r| {
2898 Ok(SyncState {
2899 peer_id: r.get(0)?,
2900 table_name: r.get(1)?,
2901 pull_cursor: r.get(2)?,
2902 push_cursor: r.get(3)?,
2903 last_synced_at: r.get(4)?,
2904 })
2905 })?;
2906 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2907 }
2908}
2909
2910pub fn fts_quote(query: &str) -> String {
2914 query
2915 .split_whitespace()
2916 .map(|t| format!("\"{}\"", t.replace('"', "\"\"")))
2917 .collect::<Vec<_>>()
2918 .join(" ")
2919}
2920
2921pub fn search_chunks(
2923 conn: &Connection,
2924 space_id: &str,
2925 query: &str,
2926 limit: usize,
2927) -> Result<Vec<(String, String, String)>> {
2928 let q = fts_quote(query);
2929 if q.is_empty() {
2930 return Ok(Vec::new());
2931 }
2932 let mut stmt = conn.prepare(
2933 "SELECT files.name, cache.file_chunks.location,
2934 snippet(file_chunks, 3, '', '', '…', 24)
2935 FROM cache.file_chunks JOIN files ON files.id = cache.file_chunks.file_id
2936 WHERE file_chunks MATCH ?1 AND files.space_id = ?2
2937 ORDER BY bm25(file_chunks) LIMIT ?3",
2938 )?;
2939 let rows = stmt.query_map((q, space_id, limit as i64), |r| {
2940 Ok((r.get(0)?, r.get(1)?, r.get(2)?))
2941 })?;
2942 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2943}
2944
2945pub fn file_text(conn: &Connection, space_id: &str, name: &str) -> Result<Option<String>> {
2947 let mut stmt = conn.prepare(
2948 "SELECT cache.file_chunks.text
2949 FROM cache.file_chunks JOIN files ON files.id = cache.file_chunks.file_id
2950 WHERE files.space_id = ?1 AND files.name = ?2
2951 AND cache.file_chunks.location != 'file metadata'
2952 ORDER BY CAST(cache.file_chunks.seq AS INTEGER) ASC",
2953 )?;
2954 let rows = stmt.query_map((space_id, name), |r| r.get::<_, String>(0))?;
2955 let parts = rows.collect::<rusqlite::Result<Vec<_>>>()?;
2956 Ok((!parts.is_empty()).then(|| parts.join("\n")))
2957}
2958
2959pub fn count_files(conn: &Connection, space_id: &str) -> Result<u64> {
2960 let n: i64 = conn.query_row(
2961 "SELECT COUNT(*) FROM files WHERE space_id = ?1",
2962 [space_id],
2963 |r| r.get(0),
2964 )?;
2965 Ok(n as u64)
2966}
2967
2968#[cfg(test)]
2969mod tests {
2970 use super::*;
2971
2972 fn price(prompt: f64, completion: f64) -> ModelPricing {
2973 ModelPricing {
2974 prompt,
2975 completion,
2976 cache_read: None,
2977 cache_write: None,
2978 }
2979 }
2980
2981 fn cache_price(
2982 prompt: f64,
2983 completion: f64,
2984 cache_read: f64,
2985 cache_write: f64,
2986 ) -> ModelPricing {
2987 ModelPricing {
2988 prompt,
2989 completion,
2990 cache_read: Some(cache_read),
2991 cache_write: Some(cache_write),
2992 }
2993 }
2994
2995 #[test]
2996 fn is_fresh_true_under_24h_false_over() {
2997 let now = Utc::now();
2998 let recent = (now - chrono::Duration::hours(1)).to_rfc3339();
2999 let stale = (now - chrono::Duration::hours(25)).to_rfc3339();
3000 assert!(is_fresh(&recent, now));
3001 assert!(!is_fresh(&stale, now));
3002 assert!(!is_fresh("not a timestamp", now)); }
3004
3005 #[test]
3006 fn usage_log_round_trips_and_aggregates() {
3007 let mut db = Db::open_in_memory().unwrap();
3008 db.upsert_model_prices(&[(
3009 "anthropic/claude-3.5-sonnet".to_string(),
3010 "OpenRouter".to_string(),
3011 price(3.0, 15.0),
3012 )])
3013 .unwrap();
3014 assert_eq!(
3015 db.model_price("anthropic/claude-3.5-sonnet"),
3016 Some(price(3.0, 15.0))
3017 );
3018 db.upsert_model_prices(&[(
3020 "anthropic/claude-3.5-sonnet".to_string(),
3021 "OpenRouter".to_string(),
3022 cache_price(4.0, 16.0, 0.4, 5.0),
3023 )])
3024 .unwrap();
3025 assert_eq!(
3026 db.model_price("anthropic/claude-3.5-sonnet"),
3027 Some(cache_price(4.0, 16.0, 0.4, 5.0))
3028 );
3029 assert_eq!(db.model_price("unknown/model"), None);
3030
3031 db.log_usage(
3033 "OpenRouter",
3034 "anthropic/claude-3.5-sonnet",
3035 100,
3036 10,
3037 70,
3038 20,
3039 Some(0.00056),
3040 true,
3041 Some("s1"),
3042 Some("space-a"),
3043 )
3044 .unwrap();
3045 db.log_usage(
3046 "Codex",
3047 "gpt-5.1-codex",
3048 50,
3049 5,
3050 0,
3051 0,
3052 None,
3053 false,
3054 None,
3055 None,
3056 )
3057 .unwrap();
3058
3059 let totals = db.usage_totals(None).unwrap();
3060 assert_eq!(totals.requests, 2);
3061 assert_eq!(totals.prompt_tokens, 150);
3062 assert_eq!(totals.completion_tokens, 15);
3063 assert_eq!(totals.cache_read_tokens, 70);
3064 assert_eq!(totals.cache_creation_tokens, 20);
3065 assert!((totals.cost - 0.00056).abs() < 1e-9);
3066
3067 let by_backend = db.usage_by_backend(None).unwrap();
3068 assert_eq!(by_backend.len(), 2);
3069 assert_eq!(by_backend[0].backend, "OpenRouter"); assert_eq!(by_backend[0].requests, 1);
3071
3072 let by_model = db.usage_by_model(5, None).unwrap();
3073 assert_eq!(by_model.len(), 2);
3074 assert!(by_model.iter().any(|m| m.model == "gpt-5.1-codex"));
3075
3076 let recent = db.usage_recent(10, None).unwrap();
3077 assert_eq!(recent.len(), 2);
3078 assert_eq!(recent[0].model, "gpt-5.1-codex"); assert_eq!(recent[0].cost, None);
3080 assert_eq!(recent[1].cache_read_tokens, 70);
3081 }
3082
3083 #[test]
3084 fn usage_queries_filter_by_since_window() {
3085 let db = Db::open_in_memory().unwrap();
3086 let insert = |created: &str| {
3088 db.raw().execute(
3089 "INSERT INTO usage_log (sync_id, created_at, session_id, backend, model,
3090 prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens, cost)
3091 VALUES (?1, ?2, NULL, 'OpenRouter', 'a/model', 100, 10, 0, 0, 0.001)",
3092 (uuid::Uuid::new_v4().to_string(), created),
3093 )
3094 };
3095 insert("2026-01-01T00:00:00+00:00").unwrap();
3096 insert("2026-01-02T00:00:00+00:00").unwrap();
3097 insert("2026-01-03T00:00:00+00:00").unwrap();
3098
3099 let since = Some("2026-01-02T00:00:00+00:00");
3100 let totals = db.usage_totals(since).unwrap();
3101 assert_eq!(totals.requests, 2);
3102 assert_eq!(totals.prompt_tokens, 200);
3103 assert!((totals.cost - 0.002).abs() < 1e-12);
3104 assert_eq!(db.usage_totals(None).unwrap().requests, 3);
3105 assert_eq!(db.usage_by_backend(since).unwrap()[0].requests, 2);
3106 assert_eq!(db.usage_by_model(5, since).unwrap()[0].requests, 2);
3107 let recent = db.usage_recent(10, since).unwrap();
3108 assert_eq!(recent.len(), 2);
3109 assert_eq!(recent[0].created_at, "2026-01-03T00:00:00+00:00");
3110 assert!(
3112 recent
3113 .iter()
3114 .any(|r| r.created_at == "2026-01-02T00:00:00+00:00")
3115 );
3116 }
3117
3118 #[test]
3119 fn usage_range_cycles_and_persists() {
3120 use crate::db::UsageRange;
3121 assert_eq!(UsageRange::Day.next(), UsageRange::Week);
3122 assert_eq!(UsageRange::All.next(), UsageRange::Day);
3123 assert_eq!(UsageRange::Day.prev(), UsageRange::All);
3124 assert_eq!(UsageRange::from_key("month"), UsageRange::Month);
3125 assert_eq!(UsageRange::from_key("bogus"), UsageRange::All);
3126 assert_eq!(UsageRange::Week.key(), "week");
3127 assert_eq!(UsageRange::Day.title(), "last 24 hours");
3128 assert!(UsageRange::All.since().is_none());
3129 assert!(UsageRange::Day.since().is_some());
3130 }
3131
3132 #[test]
3133 fn backfill_usage_costs_recomputes_history_from_catalog() {
3134 let mut db = Db::open_in_memory().unwrap();
3135 db.upsert_model_prices(&[(
3136 "anthropic/claude-3.5-sonnet".to_string(),
3137 "OpenRouter".to_string(),
3138 cache_price(3.0, 15.0, 0.3, 3.75),
3139 )])
3140 .unwrap();
3141 db.log_usage(
3144 "OpenRouter",
3145 "anthropic/claude-3.5-sonnet",
3146 100,
3147 10,
3148 70,
3149 20,
3150 None,
3151 false,
3152 None,
3153 None,
3154 )
3155 .unwrap();
3156 db.log_usage(
3157 "Codex",
3158 "gpt-5.1-codex",
3159 50,
3160 5,
3161 0,
3162 0,
3163 None,
3164 false,
3165 None,
3166 None,
3167 )
3168 .unwrap();
3169
3170 let visited = db.backfill_usage_costs().unwrap();
3171 assert_eq!(visited, 2);
3172 let totals = db.usage_totals(None).unwrap();
3175 assert!((totals.cost - 0.000_276).abs() < 1e-12);
3176 let recent = db.usage_recent(10, None).unwrap();
3177 assert!((recent[1].cost.unwrap() - 0.000_276).abs() < 1e-12); assert_eq!(recent[0].cost, None); db.backfill_usage_costs().unwrap();
3182 assert!((db.usage_totals(None).unwrap().cost - 0.000_276).abs() < 1e-12);
3183 }
3184
3185 #[test]
3186 fn backfill_preserves_provider_reported_cost() {
3187 let mut db = Db::open_in_memory().unwrap();
3188 db.upsert_model_prices(&[(
3189 "anthropic/claude-3.5-sonnet".to_string(),
3190 "OpenRouter".to_string(),
3191 cache_price(3.0, 15.0, 0.3, 3.75),
3192 )])
3193 .unwrap();
3194 db.log_usage(
3195 "OpenRouter",
3196 "anthropic/claude-3.5-sonnet",
3197 100,
3198 10,
3199 70,
3200 20,
3201 Some(0.000_321),
3202 true,
3203 None,
3204 None,
3205 )
3206 .unwrap();
3207
3208 db.backfill_usage_costs().unwrap();
3209
3210 assert_eq!(db.usage_recent(1, None).unwrap()[0].cost, Some(0.000_321));
3211 }
3212
3213 #[test]
3214 fn backfill_usage_costs_heals_per_token_catalog() {
3215 let mut db = Db::open_in_memory().unwrap();
3219 db.upsert_model_prices(&[(
3220 "deepseek/deepseek-v4-flash-0731".to_string(),
3221 "OpenRouter".to_string(),
3222 price(8e-08, 1.8e-07),
3223 )])
3224 .unwrap();
3225 db.log_usage(
3226 "OpenRouter",
3227 "deepseek/deepseek-v4-flash-0731",
3228 122_221,
3229 672,
3230 118_784,
3231 0,
3232 Some(9.89864e-09), false,
3234 None,
3235 None,
3236 )
3237 .unwrap();
3238
3239 db.backfill_usage_costs().unwrap();
3240
3241 assert_eq!(
3242 db.model_price("deepseek/deepseek-v4-flash-0731"),
3243 Some(price(0.08, 0.18))
3244 );
3245 let recent = db.usage_recent(10, None).unwrap();
3247 let cost = recent[0].cost.unwrap();
3248 assert!((cost - 0.009_898_6).abs() < 1e-6, "cost was {cost}");
3249 assert!(cost > 0.009, "cost was {cost}");
3250 }
3251
3252 #[test]
3253 fn request_cost_prices_tokens_against_catalog() {
3254 let mut db = Db::open_in_memory().unwrap();
3255 assert_eq!(db.request_cost("unknown/model", 100, 10, 0, 0), None);
3256 db.upsert_model_prices(&[(
3257 "anthropic/claude-3.5-sonnet".to_string(),
3258 "OpenRouter".to_string(),
3259 cache_price(3.0, 15.0, 0.3, 3.75),
3260 )])
3261 .unwrap();
3262 let cost = db
3263 .request_cost("anthropic/claude-3.5-sonnet", 100, 10, 70, 20)
3264 .unwrap();
3265 assert!((cost - 0.000_276).abs() < 1e-12);
3266 }
3267
3268 #[test]
3269 fn model_price_cross_references_openrouter_catalog_twins() {
3270 let mut db = Db::open_in_memory().unwrap();
3271 db.upsert_model_prices(&[
3272 (
3273 "deepseek/deepseek-v4-flash".to_string(),
3274 "OpenRouter".to_string(),
3275 cache_price(0.08, 0.18, 0.016, 0.08),
3276 ),
3277 (
3278 "openai/gpt-5".to_string(),
3279 "OpenRouter".to_string(),
3280 cache_price(1.25, 10.0, 0.125, 1.25),
3281 ),
3282 ])
3283 .unwrap();
3284 assert_eq!(
3287 db.model_price("deepseek/deepseek-v4-flash"),
3288 Some(cache_price(0.08, 0.18, 0.016, 0.08))
3289 );
3290 assert_eq!(
3291 db.model_price("go:deepseek-v4-flash"),
3292 Some(cache_price(0.08, 0.18, 0.016, 0.08))
3293 );
3294 assert_eq!(
3295 db.model_price("deepseek-v4-flash"),
3296 Some(cache_price(0.08, 0.18, 0.016, 0.08))
3297 );
3298 assert_eq!(
3299 db.model_price("openai:gpt-5"),
3300 Some(cache_price(1.25, 10.0, 0.125, 1.25))
3301 );
3302 assert_eq!(
3303 db.model_price("codex:gpt-5"),
3304 Some(cache_price(1.25, 10.0, 0.125, 1.25))
3305 );
3306 assert_eq!(db.model_price("no-such-model-anywhere"), None);
3308 let cost = db
3310 .request_cost("go:deepseek-v4-flash", 100, 10, 70, 0)
3311 .unwrap();
3312 assert!((cost - 0.000_005_32).abs() < 1e-15);
3313 }
3314
3315 #[test]
3316 fn price_name_strips_backend_prefixes_and_vendors() {
3317 assert_eq!(price_name("go:deepseek-v4-flash"), "deepseek-v4-flash");
3318 assert_eq!(price_name("openai:gpt-5"), "gpt-5");
3319 assert_eq!(price_name("codex:gpt-5.1-codex"), "gpt-5.1-codex");
3320 assert_eq!(price_name("opencode:qwen3.6-plus"), "qwen3.6-plus");
3321 assert_eq!(
3322 price_name("deepseek/deepseek-v4-flash"),
3323 "deepseek-v4-flash"
3324 );
3325 assert_eq!(price_name("gpt-5"), "gpt-5");
3326 }
3327
3328 #[test]
3329 fn backfill_prices_non_openrouter_models_via_catalog_twins() {
3330 let mut db = Db::open_in_memory().unwrap();
3331 db.upsert_model_prices(&[(
3332 "deepseek/deepseek-v4-flash".to_string(),
3333 "OpenRouter".to_string(),
3334 cache_price(0.08, 0.18, 0.016, 0.08),
3335 )])
3336 .unwrap();
3337 db.log_usage(
3340 "OpenCode Go",
3341 "go:deepseek-v4-flash",
3342 100,
3343 10,
3344 0,
3345 0,
3346 None,
3347 false,
3348 None,
3349 None,
3350 )
3351 .unwrap();
3352
3353 db.backfill_usage_costs().unwrap();
3354
3355 let recent = db.usage_recent(10, None).unwrap();
3356 let cost = recent[0].cost.unwrap();
3357 assert!((cost - 0.000_009_8).abs() < 1e-15, "cost was {cost}");
3358 assert!((db.usage_totals(None).unwrap().cost - 0.000_009_8).abs() < 1e-15);
3359 }
3360
3361 #[test]
3362 fn web_cache_roundtrips_and_updates_on_rewrite() {
3363 let db = Db::open_in_memory().unwrap();
3364 assert!(cache_get(db.raw(), "example.com/a").unwrap().is_none());
3365 cache_put(
3366 db.raw(),
3367 "example.com/a",
3368 "https://example.com/a",
3369 Some("Title"),
3370 "body text",
3371 )
3372 .unwrap();
3373 let (title, text, fetched_at) = cache_get(db.raw(), "example.com/a").unwrap().unwrap();
3374 assert_eq!(title, "Title");
3375 assert_eq!(text, "body text");
3376 assert!(!fetched_at.is_empty());
3377
3378 cache_put(
3380 db.raw(),
3381 "example.com/a",
3382 "https://example.com/a",
3383 None,
3384 "new body",
3385 )
3386 .unwrap();
3387 let (title, text, _) = cache_get(db.raw(), "example.com/a").unwrap().unwrap();
3388 assert_eq!(title, "");
3389 assert_eq!(text, "new body");
3390 }
3391
3392 #[test]
3393 fn web_mode_defaults_off_and_toggles() {
3394 let db = Db::open_in_memory().unwrap();
3395 let space = db.default_space_id().unwrap();
3396 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3397 assert!(!s.web_mode);
3398 db.set_session_web_mode(&s.id, true).unwrap();
3399 assert!(db.list_sessions(&space).unwrap()[0].web_mode);
3400 }
3401
3402 #[test]
3403 fn swarm_mode_defaults_off_and_toggles() {
3404 let db = Db::open_in_memory().unwrap();
3405 let space = db.default_space_id().unwrap();
3406 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3407 assert!(!s.swarm_mode);
3408 db.set_session_swarm_mode(&s.id, true).unwrap();
3409 assert!(db.list_sessions(&space).unwrap()[0].swarm_mode);
3410 assert!(db.get_session(&s.id).unwrap().unwrap().swarm_mode);
3411 }
3412
3413 #[test]
3414 fn swarm_personas_roundtrip_and_replace_all_on_save() {
3415 let db = Db::open_in_memory().unwrap();
3416 let space = db.default_space_id().unwrap();
3417 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3418 assert!(db.list_swarm_personas(&s.id).unwrap().is_empty());
3419
3420 let roster = vec![
3421 Persona {
3422 name: "Skeptic".into(),
3423 model: "a/one".into(),
3424 blurb: "pokes holes".into(),
3425 },
3426 Persona {
3427 name: "Advocate".into(),
3428 model: "b/two".into(),
3429 blurb: "user-first".into(),
3430 },
3431 ];
3432 db.save_swarm_personas(&s.id, &roster).unwrap();
3433 let loaded = db.list_swarm_personas(&s.id).unwrap();
3434 assert_eq!(loaded.len(), 2);
3435 assert_eq!(loaded[0].name, "Skeptic");
3436 assert_eq!(loaded[1].name, "Advocate");
3437
3438 db.save_swarm_personas(&s.id, &roster[..1]).unwrap();
3440 assert_eq!(db.list_swarm_personas(&s.id).unwrap().len(), 1);
3441 }
3442
3443 #[test]
3444 fn persona_message_tags_role_assistant_with_persona_and_model() {
3445 let db = Db::open_in_memory().unwrap();
3446 let space = db.default_space_id().unwrap();
3447 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3448 db.add_persona_message(&s.id, "reply text", "Skeptic", "a/one")
3449 .unwrap();
3450 let msgs = db.load_messages(&s.id).unwrap();
3451 assert_eq!(msgs.len(), 1);
3452 assert_eq!(msgs[0].role, "assistant");
3453 assert_eq!(msgs[0].persona.as_deref(), Some("Skeptic"));
3454 assert_eq!(msgs[0].model.as_deref(), Some("a/one"));
3455
3456 db.add_assistant_message(&s.id, "final answer", None, None, None, None, None, None)
3458 .unwrap();
3459 let msgs = db.load_messages(&s.id).unwrap();
3460 assert_eq!(msgs[1].persona, None);
3461 }
3462
3463 #[test]
3464 fn session_sources_link_to_the_web_cache_and_are_keyword_searchable() {
3465 let db = Db::open_in_memory().unwrap();
3466 let space = db.default_space_id().unwrap();
3467 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3468 cache_put(
3469 db.raw(),
3470 "https://example.com/a",
3471 "https://example.com/a",
3472 Some("A"),
3473 "rust borrow checker deep dive",
3474 )
3475 .unwrap();
3476 cache_put(
3477 db.raw(),
3478 "https://example.com/b",
3479 "https://example.com/b",
3480 Some("B"),
3481 "cooking pasta recipes",
3482 )
3483 .unwrap();
3484 db.add_session_sources(
3485 &s.id,
3486 &[
3487 "https://example.com/a".to_string(),
3488 "https://example.com/b".to_string(),
3489 ],
3490 )
3491 .unwrap();
3492
3493 let hits = db.search_session_sources(&s.id, "borrow checker").unwrap();
3494 assert_eq!(hits.len(), 1);
3495 assert!(hits[0].1.contains("borrow checker"));
3496
3497 assert!(
3498 db.search_session_sources(&s.id, "quantum")
3499 .unwrap()
3500 .is_empty()
3501 );
3502 }
3503
3504 #[test]
3505 fn set_source_flag_pins_and_discards_then_clears() {
3506 let db = Db::open_in_memory().unwrap();
3507 let session_id = "sess-1";
3508 add_session_sources(&db.conn, session_id, &["https://a.example/x".to_string()]).unwrap();
3509 db.set_source_flag(session_id, "https://a.example/x", Some("pinned"))
3510 .unwrap();
3511 assert_eq!(
3512 pinned_urls(&db.conn, session_id).unwrap(),
3513 vec!["https://a.example/x".to_string()]
3514 );
3515 assert!(discarded_domains(&db.conn, session_id).unwrap().is_empty());
3516
3517 db.set_source_flag(session_id, "https://a.example/x", Some("discarded"))
3518 .unwrap();
3519 assert!(pinned_urls(&db.conn, session_id).unwrap().is_empty());
3520 assert_eq!(
3521 discarded_domains(&db.conn, session_id).unwrap(),
3522 vec!["a.example".to_string()]
3523 );
3524
3525 db.set_source_flag(session_id, "https://a.example/x", None)
3526 .unwrap();
3527 assert!(discarded_domains(&db.conn, session_id).unwrap().is_empty());
3528 }
3529
3530 #[test]
3531 fn upsert_research_stage_message_replaces_the_same_labels_row() {
3532 let db = Db::open_in_memory().unwrap();
3533 let space = db.default_space_id().unwrap();
3534 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3535 db.upsert_research_stage_message(&s.id, "searching", "round 1, 1/3")
3536 .unwrap();
3537 db.upsert_research_stage_message(&s.id, "searching", "round 1, 2/3")
3538 .unwrap();
3539 db.upsert_research_stage_message(&s.id, "planning", "")
3540 .unwrap();
3541
3542 let msgs = db.load_messages(&s.id).unwrap();
3543 let searching: Vec<_> = msgs
3544 .iter()
3545 .filter(|m| m.content.starts_with("searching:"))
3546 .collect();
3547 assert_eq!(searching.len(), 1, "expected one row, updated in place");
3548 assert!(searching[0].content.contains("2/3"));
3549 assert_eq!(msgs.iter().filter(|m| m.content == "planning").count(), 1);
3550 }
3551
3552 #[test]
3553 fn session_and_message_roundtrip() {
3554 let db = Db::open_in_memory().unwrap();
3555 let space = db.default_space_id().unwrap();
3556 let s = db
3557 .create_session("hello", "openai/gpt-4o", &space, "chat")
3558 .unwrap();
3559
3560 db.add_user_message(&s.id, "hi").unwrap();
3561 db.add_assistant_message(
3562 &s.id,
3563 "hello there",
3564 Some("openai/gpt-4o"),
3565 Some("let me think"),
3566 Some(3),
3567 Some(1.5),
3568 Some(0.0042),
3569 Some("Vibed"),
3570 )
3571 .unwrap();
3572
3573 let msgs = db.load_messages(&s.id).unwrap();
3574 assert_eq!(msgs.len(), 2);
3575 assert_eq!(msgs[0].role, "user");
3576 assert_eq!(msgs[1].content, "hello there");
3577 assert_eq!(msgs[1].model.as_deref(), Some("openai/gpt-4o"));
3578 assert_eq!(msgs[1].reasoning.as_deref(), Some("let me think"));
3579 assert_eq!(msgs[1].tokens, Some(3));
3580 assert_eq!(msgs[1].cost, Some(0.0042));
3581
3582 let sessions = db.list_sessions(&space).unwrap();
3583 assert_eq!(sessions.len(), 1);
3584 assert_eq!(sessions[0].id, s.id);
3585 }
3586
3587 #[test]
3588 fn latest_session_returns_the_newest_session_across_spaces() {
3589 let db = Db::open_in_memory().unwrap();
3590 let default_id = db.default_space_id().unwrap();
3591 let other = db.create_space("other").unwrap();
3592 let older = db
3593 .create_session("older", "a/one", &default_id, "chat")
3594 .unwrap();
3595 let newer = db
3596 .create_session("newer", "b/two", &other.id, "chat")
3597 .unwrap();
3598
3599 db.raw()
3602 .execute(
3603 "UPDATE sessions SET updated_at = ?1 WHERE id = ?2",
3604 ("2026-01-01T00:00:00Z", &older.id),
3605 )
3606 .unwrap();
3607 db.raw()
3608 .execute(
3609 "UPDATE sessions SET updated_at = ?1 WHERE id = ?2",
3610 ("2026-01-02T00:00:00Z", &newer.id),
3611 )
3612 .unwrap();
3613
3614 let (space_id, session) = db.latest_session().unwrap().unwrap();
3615 assert_eq!(space_id, other.id);
3616 assert_eq!(session.id, newer.id);
3617 }
3618
3619 #[test]
3620 fn latest_session_is_empty_without_sessions() {
3621 let db = Db::open_in_memory().unwrap();
3622 assert!(db.latest_session().unwrap().is_none());
3623 }
3624
3625 #[test]
3626 fn markdown_images_in_content_roundtrip() {
3627 let db = Db::open_in_memory().unwrap();
3628 let space = db.default_space_id().unwrap();
3629 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3630 let content = "look at  and ";
3631 db.add_user_message(&s.id, content).unwrap();
3632 let msgs = db.load_messages(&s.id).unwrap();
3633 assert_eq!(msgs.len(), 1);
3634 assert!(msgs[0].content.contains(""));
3635 assert!(msgs[0].content.contains(""));
3636 }
3637
3638 #[test]
3639 fn model_prefs_toggle_and_used() {
3640 let db = Db::open_in_memory().unwrap();
3641 assert!(db.toggle_favorite("a/one").unwrap()); assert!(!db.toggle_favorite("a/one").unwrap()); db.mark_model_used("a/one").unwrap();
3644 db.set_reasoning("a/one", Some("high")).unwrap();
3645
3646 let prefs = db.load_model_prefs().unwrap();
3647 let p = &prefs[0];
3648 assert_eq!(p.id, "a/one");
3649 assert!(!p.favorite);
3650 assert!(p.last_used.is_some());
3651 assert_eq!(p.reasoning.as_deref(), Some("high"));
3652 }
3653
3654 #[test]
3655 fn settings_roundtrip() {
3656 let db = Db::open_in_memory().unwrap();
3657 db.set_setting("temperature", "0.7").unwrap();
3658 db.set_setting("temperature", "0.9").unwrap(); let s = db.load_settings().unwrap();
3660 assert_eq!(s, vec![("temperature".to_string(), "0.9".to_string())]);
3661 }
3662
3663 #[test]
3664 fn set_model_updates_row() {
3665 let db = Db::open_in_memory().unwrap();
3666 let space = db.default_space_id().unwrap();
3667 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3668 db.set_session_model(&s.id, "c/d").unwrap();
3669 assert_eq!(db.list_sessions(&space).unwrap()[0].model, "c/d");
3670 }
3671
3672 #[test]
3673 fn compaction_persists_and_roundtrips() {
3674 let db = Db::open_in_memory().unwrap();
3675 let space = db.default_space_id().unwrap();
3676 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3677 assert_eq!(s.compact_summary, None);
3678 assert_eq!(s.compact_through, 0);
3679
3680 db.set_compaction(&s.id, "digest of earlier turns", 6)
3681 .unwrap();
3682 let reloaded = &db.list_sessions(&space).unwrap()[0];
3683 assert_eq!(
3684 reloaded.compact_summary.as_deref(),
3685 Some("digest of earlier turns")
3686 );
3687 assert_eq!(reloaded.compact_through, 6);
3688 }
3689
3690 #[test]
3691 fn spaces_crud_and_session_reassignment_on_delete() {
3692 let db = Db::open_in_memory().unwrap();
3693 let spaces = db.list_spaces().unwrap();
3694 assert_eq!(spaces.len(), 1);
3695 assert_eq!(spaces[0].name, DEFAULT_SPACE);
3696
3697 let work = db.create_space("work").unwrap();
3698 let s = db.create_session("hi", "a/b", &work.id, "chat").unwrap();
3699 assert_eq!(db.count_sessions(&work.id).unwrap(), 1);
3700
3701 db.rename_space(&work.id, "work-renamed").unwrap();
3702 assert!(
3703 db.list_spaces()
3704 .unwrap()
3705 .iter()
3706 .any(|s| s.name == "work-renamed")
3707 );
3708
3709 db.delete_space(&work.id).unwrap();
3710 assert_eq!(db.list_spaces().unwrap().len(), 1); let default_id = db.default_space_id().unwrap();
3712 let moved = db.list_sessions(&default_id).unwrap();
3713 assert!(moved.iter().any(|c| c.id == s.id)); }
3715
3716 #[test]
3717 fn chunk_embeddings_store_rank_and_invalidate() {
3718 let db = Db::open_in_memory().unwrap();
3719 let space = db.default_space_id().unwrap();
3720 let id = db.upsert_file(&space, "book.pdf", "h1", 10, "ok").unwrap();
3721 db.set_file_chunks(
3722 &id,
3723 &[
3724 ("page 1".into(), "cooking with fire".into()),
3725 ("page 2".into(), "quantum entanglement".into()),
3726 ],
3727 )
3728 .unwrap();
3729
3730 let v = vec![0.25f32, -1.0, 3.5];
3732 assert_eq!(blob_to_vec(&vec_to_blob(&v)), v);
3733
3734 assert_eq!(
3736 files_missing_embeddings(&db.conn, &space).unwrap(),
3737 vec![id.clone()]
3738 );
3739
3740 db.set_chunk_embeddings(&id, &[(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])])
3741 .unwrap();
3742 assert!(
3743 files_missing_embeddings(&db.conn, &space)
3744 .unwrap()
3745 .is_empty()
3746 );
3747
3748 let hits = semantic_chunks(&db.conn, &space, &[0.1, 0.9], 5).unwrap();
3750 assert_eq!(hits[0].1, "page 2");
3751 assert!(hits[0].2.contains("quantum"));
3752 assert!(hits[0].3 > hits[1].3, "scores must be descending");
3753
3754 let hits = semantic_chunks(&db.conn, &space, &[1.0, 0.0, 0.0], 5).unwrap();
3756 assert!(hits.is_empty());
3757
3758 db.set_file_chunks(&id, &[("page 1".into(), "new text".into())])
3760 .unwrap();
3761 assert_eq!(
3762 files_missing_embeddings(&db.conn, &space).unwrap(),
3763 vec![id.clone()]
3764 );
3765
3766 db.set_chunk_embeddings(&id, &[(99, vec![1.0, 0.0])])
3768 .unwrap();
3769 assert_eq!(
3770 files_missing_embeddings(&db.conn, &space).unwrap(),
3771 vec![id.clone()]
3772 );
3773 }
3774
3775 #[test]
3776 fn files_without_chunks_are_embedding_queue_entries() {
3777 let db = Db::open_in_memory().unwrap();
3778 let space = db.default_space_id().unwrap();
3779 let id = db
3780 .upsert_file(&space, "empty.py", "h", 0, "no text (scanned?)")
3781 .unwrap();
3782 assert_eq!(
3783 files_missing_embeddings(&db.conn, &space).unwrap(),
3784 vec![id]
3785 );
3786
3787 let ocr = db.upsert_file(&space, "scan.pdf", "p", 0, "ocr…").unwrap();
3790 assert!(
3791 files_missing_embeddings(&db.conn, &space)
3792 .unwrap()
3793 .contains(&ocr)
3794 );
3795 }
3796
3797 #[test]
3798 fn files_upsert_list_delete_roundtrip() {
3799 let db = Db::open_in_memory().unwrap();
3800 let space = db.default_space_id().unwrap();
3801 let id = db.upsert_file(&space, "notes.md", "h1", 10, "ok").unwrap();
3802 db.set_file_chunks(&id, &[("lines 1-40".into(), "hello fts world".into())])
3803 .unwrap();
3804
3805 let files = db.list_files(&space).unwrap();
3806 assert_eq!(files.len(), 1);
3807 assert_eq!(files[0].name, "notes.md");
3808 assert_eq!(files[0].hash, "h1");
3809 assert_eq!(files[0].status, "ok");
3810
3811 let id2 = db.upsert_file(&space, "notes.md", "h2", 12, "ok").unwrap();
3813 db.set_file_chunks(&id2, &[("lines 1-40".into(), "goodbye".into())])
3814 .unwrap();
3815 let files = db.list_files(&space).unwrap();
3816 assert_eq!(files.len(), 1);
3817 assert_eq!(files[0].hash, "h2");
3818
3819 db.delete_file(&files[0].id).unwrap();
3820 assert!(db.list_files(&space).unwrap().is_empty());
3821 }
3822
3823 #[test]
3824 fn chunk_search_ranks_and_scopes_by_space() {
3825 let db = Db::open_in_memory().unwrap();
3826 let space = db.default_space_id().unwrap();
3827 let other = db.create_space("other").unwrap();
3828 let a = db.upsert_file(&space, "a.md", "h", 1, "ok").unwrap();
3829 let b = db.upsert_file(&other.id, "b.md", "h", 1, "ok").unwrap();
3830 db.set_file_chunks(&a, &[("lines 1-40".into(), "rust borrow checker".into())])
3831 .unwrap();
3832 db.set_file_chunks(&b, &[("lines 1-40".into(), "rust in other space".into())])
3833 .unwrap();
3834
3835 let hits = search_chunks(&db.conn, &space, "rust", 8).unwrap();
3836 assert_eq!(hits.len(), 1); assert_eq!(hits[0].0, "a.md");
3838 assert_eq!(hits[0].1, "lines 1-40");
3839 assert!(hits[0].2.contains("rust"));
3840
3841 assert!(search_chunks(&db.conn, &space, "c++ \"quoted\" -dash", 8).is_ok());
3843 }
3844
3845 #[test]
3846 fn file_text_joins_chunks_in_order() {
3847 let db = Db::open_in_memory().unwrap();
3848 let space = db.default_space_id().unwrap();
3849 let id = db.upsert_file(&space, "doc.txt", "h", 1, "ok").unwrap();
3850 db.set_file_chunks(
3851 &id,
3852 &[
3853 ("lines 1-2".into(), "one\ntwo".into()),
3854 ("lines 3-4".into(), "three\nfour".into()),
3855 ],
3856 )
3857 .unwrap();
3858 let text = file_text(&db.conn, &space, "doc.txt").unwrap().unwrap();
3859 assert_eq!(text, "one\ntwo\nthree\nfour");
3860 assert!(
3861 file_text(&db.conn, &space, "missing.txt")
3862 .unwrap()
3863 .is_none()
3864 );
3865 assert_eq!(count_files(&db.conn, &space).unwrap(), 1);
3866 }
3867
3868 #[test]
3869 fn research_stage_messages_round_trip() {
3870 let db = Db::open_in_memory().unwrap();
3871 let space = db.default_space_id().unwrap();
3872 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3873 db.add_research_stage_message(&s.id, "planning…").unwrap();
3874 let msgs = db.load_messages(&s.id).unwrap();
3875 assert_eq!(msgs.last().unwrap().role, "research_stage");
3876 assert_eq!(msgs.last().unwrap().content, "planning…");
3877 }
3878
3879 #[test]
3880 fn survey_messages_round_trip() {
3881 let db = Db::open_in_memory().unwrap();
3882 let space = db.default_space_id().unwrap();
3883 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3884 db.add_survey_message(&s.id, "For \"topic\":\n 1. Depth or breadth?")
3885 .unwrap();
3886 let msgs = db.load_messages(&s.id).unwrap();
3887 let last = msgs.last().unwrap();
3888 assert_eq!(last.role, "survey");
3889 assert!(last.content.contains("Depth or breadth?"));
3890 }
3891
3892 #[test]
3893 fn gate_reply_round_trip() {
3894 let db = Db::open_in_memory().unwrap();
3895 let space = db.default_space_id().unwrap();
3896 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3897 db.add_gate_reply_message(&s.id, "the second option")
3898 .unwrap();
3899 let msgs = db.load_messages(&s.id).unwrap();
3900 let last = msgs.last().unwrap();
3901 assert_eq!(last.role, "gate_reply");
3902 assert_eq!(last.content, "the second option");
3903 }
3904
3905 #[test]
3906 fn create_list_touch_delete_watch_roundtrip() {
3907 let db = Db::open_in_memory().unwrap();
3908 let id = db
3909 .create_watch("space-1", "rust async runtimes", 24, "sess-1")
3910 .unwrap();
3911 let watches = db.list_watches("space-1").unwrap();
3912 assert_eq!(watches.len(), 1);
3913 assert_eq!(watches[0].topic, "rust async runtimes");
3914 assert_eq!(watches[0].interval_hours, 24);
3915 assert!(watches[0].last_run_at.is_none());
3916
3917 db.touch_watch(&id, "2026-07-07T00:00:00+00:00").unwrap();
3918 let watches = db.list_watches("space-1").unwrap();
3919 assert_eq!(
3920 watches[0].last_run_at.as_deref(),
3921 Some("2026-07-07T00:00:00+00:00")
3922 );
3923
3924 db.delete_watch(&id).unwrap();
3925 assert!(db.list_watches("space-1").unwrap().is_empty());
3926 }
3927
3928 #[test]
3929 #[allow(clippy::similar_names)]
3931 fn list_all_watches_returns_watches_from_all_spaces() {
3932 let db = Db::open_in_memory().unwrap();
3933
3934 let id1 = db.create_watch("space-a", "topic-1", 24, "sess-1").unwrap();
3936 let id2 = db.create_watch("space-b", "topic-2", 48, "sess-2").unwrap();
3937 let id3 = db.create_watch("space-a", "topic-3", 12, "sess-3").unwrap();
3938
3939 let all_watches = db.list_all_watches().unwrap();
3941 assert_eq!(all_watches.len(), 3);
3942 assert!(
3943 all_watches
3944 .iter()
3945 .any(|w| w.id == id1 && w.space_id == "space-a")
3946 );
3947 assert!(
3948 all_watches
3949 .iter()
3950 .any(|w| w.id == id2 && w.space_id == "space-b")
3951 );
3952 assert!(
3953 all_watches
3954 .iter()
3955 .any(|w| w.id == id3 && w.space_id == "space-a")
3956 );
3957
3958 let space_a_watches = db.list_watches("space-a").unwrap();
3961 assert_eq!(space_a_watches.len(), 2);
3962 assert!(space_a_watches.iter().all(|w| w.space_id == "space-a"));
3963
3964 let space_b_watches = db.list_watches("space-b").unwrap();
3965 assert_eq!(space_b_watches.len(), 1);
3966 assert!(space_b_watches.iter().all(|w| w.space_id == "space-b"));
3967 }
3968
3969 fn legacy_db(path: &std::path::Path, now: &str) {
3973 let conn = rusqlite::Connection::open(path).unwrap();
3974 conn.execute_batch(
3975 "CREATE TABLE sessions (id TEXT PRIMARY KEY, title TEXT NOT NULL,
3976 model TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL);
3977 CREATE TABLE messages (id TEXT PRIMARY KEY, session_id TEXT NOT NULL,
3978 role TEXT NOT NULL, content TEXT NOT NULL, created_at TEXT NOT NULL);
3979 CREATE TABLE model_prefs (id TEXT PRIMARY KEY,
3980 favorite INTEGER NOT NULL DEFAULT 0, last_used TEXT);
3981 CREATE TABLE app_settings (key TEXT PRIMARY KEY, value TEXT NOT NULL);
3982 CREATE TABLE spaces (id TEXT PRIMARY KEY, name TEXT NOT NULL UNIQUE,
3983 created_at TEXT NOT NULL);
3984 CREATE TABLE files (id TEXT PRIMARY KEY, space_id TEXT NOT NULL,
3985 name TEXT NOT NULL, hash TEXT NOT NULL, size INTEGER NOT NULL,
3986 status TEXT NOT NULL, created_at TEXT NOT NULL,
3987 UNIQUE(space_id, name));
3988 ALTER TABLE files ADD COLUMN mtime INTEGER NOT NULL DEFAULT 0;
3989 CREATE VIRTUAL TABLE file_chunks USING fts5(
3990 file_id UNINDEXED, seq UNINDEXED, location UNINDEXED, text);
3991 CREATE TABLE chunk_embeddings (file_id TEXT NOT NULL, seq INTEGER NOT NULL,
3992 vec BLOB NOT NULL, PRIMARY KEY (file_id, seq));
3993 CREATE TABLE web_cache (url_norm TEXT PRIMARY KEY, url TEXT NOT NULL,
3994 title TEXT, text TEXT NOT NULL, fetched_at TEXT NOT NULL);
3995 CREATE TABLE citations (id INTEGER PRIMARY KEY AUTOINCREMENT,
3996 space_id TEXT NOT NULL, report_file TEXT NOT NULL,
3997 url TEXT NOT NULL, title TEXT);
3998 CREATE TABLE session_sources (session_id TEXT NOT NULL,
3999 url_norm TEXT NOT NULL, PRIMARY KEY (session_id, url_norm));
4000 CREATE TABLE watches (id TEXT PRIMARY KEY, space_id TEXT NOT NULL,
4001 topic TEXT NOT NULL, interval_hours INTEGER NOT NULL,
4002 session_id TEXT NOT NULL, last_run_at TEXT);
4003 CREATE TABLE swarm_personas (session_id TEXT NOT NULL, ord INTEGER NOT NULL,
4004 name TEXT NOT NULL, model TEXT NOT NULL, persona TEXT NOT NULL);
4005 CREATE TABLE usage_log (id INTEGER PRIMARY KEY AUTOINCREMENT,
4006 created_at TEXT NOT NULL, session_id TEXT, space_id TEXT,
4007 backend TEXT NOT NULL, model TEXT NOT NULL,
4008 prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL,
4009 cache_read_tokens INTEGER NOT NULL DEFAULT 0,
4010 cache_creation_tokens INTEGER NOT NULL DEFAULT 0, cost REAL);
4011 CREATE TABLE model_prices (model_id TEXT PRIMARY KEY, backend TEXT NOT NULL,
4012 prompt_price REAL NOT NULL, completion_price REAL NOT NULL,
4013 cache_read_price REAL, cache_write_price REAL, updated_at TEXT NOT NULL);",
4014 )
4015 .unwrap();
4016 conn.execute(
4017 "INSERT INTO spaces (id, name, created_at) VALUES ('sp', 'default', ?1)",
4018 [now],
4019 )
4020 .unwrap();
4021 conn.execute(
4022 "INSERT INTO files (id, space_id, name, hash, size, status, created_at, mtime)
4023 VALUES ('f1', 'sp', 'a.txt', 'h1', 10, 'ok', ?1, 1234)",
4024 [now],
4025 )
4026 .unwrap();
4027 conn.execute(
4028 "INSERT INTO file_chunks (file_id, seq, location, text)
4029 VALUES ('f1', 0, 'l', 'hello world')",
4030 [],
4031 )
4032 .unwrap();
4033 conn.execute(
4034 "INSERT INTO chunk_embeddings (file_id, seq, vec) VALUES ('f1', 0, ?1)",
4035 [vec_to_blob(&[1.0, 0.0])],
4036 )
4037 .unwrap();
4038 conn.execute(
4039 "INSERT INTO web_cache (url_norm, url, title, text, fetched_at)
4040 VALUES ('https://x.test/', 'https://x.test/', NULL, 'cached body', ?1)",
4041 [now],
4042 )
4043 .unwrap();
4044 conn.execute(
4045 "INSERT INTO model_prices (model_id, backend, prompt_price, completion_price, updated_at)
4046 VALUES ('a/model', 'OpenRouter', 1.0, 2.0, ?1)",
4047 [now],
4048 )
4049 .unwrap();
4050 conn.execute(
4051 "INSERT INTO usage_log (created_at, backend, model, prompt_tokens, completion_tokens)
4052 VALUES (?1, 'OpenRouter', 'a/model', 100, 10)",
4053 [now],
4054 )
4055 .unwrap();
4056 drop(conn);
4057 }
4058
4059 #[test]
4060 #[allow(clippy::float_cmp)]
4062 fn legacy_db_migrates_cache_tables_and_seeds_file_index_state() {
4063 let dir = std::env::temp_dir().join(format!("nexus-migrate-{}", uuid::Uuid::new_v4()));
4064 std::fs::create_dir_all(&dir).unwrap();
4065 let db_path = dir.join("nexus.db");
4066 legacy_db(&db_path, &Utc::now().to_rfc3339());
4067 let mut db = Db::open(&db_path).unwrap();
4068
4069 let v: i64 = db
4071 .raw()
4072 .query_row("PRAGMA user_version", [], |r| r.get(0))
4073 .unwrap();
4074 assert_eq!(v, SCHEMA_VERSION);
4075 assert!(has_column(db.raw(), "files", "mtime").unwrap());
4076
4077 assert_eq!(db.default_space_id().unwrap(), DEFAULT_SPACE);
4080 let files_space: String = db
4081 .raw()
4082 .query_row("SELECT space_id FROM files WHERE id = 'f1'", [], |r| {
4083 r.get(0)
4084 })
4085 .unwrap();
4086 assert_eq!(files_space, DEFAULT_SPACE);
4087
4088 let (mtime, status) = db
4090 .raw()
4091 .query_row(
4092 "SELECT mtime, status FROM cache.file_index_state WHERE file_id = 'f1'",
4093 [],
4094 |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)),
4095 )
4096 .unwrap();
4097 assert_eq!((mtime, status.as_str()), (1234, "ok"));
4098
4099 assert!(cache_path_for(&db_path).is_file());
4102 let space = db.default_space_id().unwrap();
4103 let hits = search_chunks(db.raw(), &space, "hello", 8).unwrap();
4104 assert_eq!(hits.len(), 1);
4105 assert_eq!(
4106 file_text(db.raw(), &space, "a.txt").unwrap().as_deref(),
4107 Some("hello world")
4108 );
4109 let (_, text, _) = cache_get(db.raw(), "https://x.test/").unwrap().unwrap();
4110 assert_eq!(text, "cached body");
4111 let hits = semantic_chunks(db.raw(), &space, &[1.0, 0.0], 4).unwrap();
4112 assert_eq!(hits.len(), 1);
4113 assert_eq!(hits[0].0, "a.txt");
4114
4115 assert_eq!(db.model_price("a/model").unwrap().prompt, 1.0);
4118 assert_eq!(db.backfill_usage_costs().unwrap(), 1);
4119 let sync: String = db
4120 .raw()
4121 .query_row("SELECT sync_id FROM usage_log WHERE id = 1", [], |r| {
4122 r.get(0)
4123 })
4124 .unwrap();
4125 assert!(!sync.is_empty());
4126 }
4127
4128 #[test]
4129 fn fresh_db_creates_sibling_cache_db_with_split_tables() {
4130 let dir = std::env::temp_dir().join(format!("nexus-fresh-{}", uuid::Uuid::new_v4()));
4131 std::fs::create_dir_all(&dir).unwrap();
4132 let db_path = dir.join("nexus.db");
4133 let db = Db::open(&db_path).unwrap();
4134 let space = db.default_space_id().unwrap();
4135 let id = db.upsert_file(&space, "doc.txt", "h", 1, "ok").unwrap();
4136 db.set_file_chunks(&id, &[("l".into(), "needle text".into())])
4137 .unwrap();
4138
4139 assert!(!has_column(db.raw(), "web_cache", "url_norm").unwrap());
4141 assert!(!has_column(db.raw(), "files", "mtime").unwrap());
4142 let cache_path = cache_path_for(&db_path);
4145 assert!(cache_path.is_file());
4146 let conn = rusqlite::Connection::open(&cache_path).unwrap();
4147 assert!(cache_get(&conn, "x").unwrap().is_none());
4148 drop(conn);
4149 assert_eq!(
4150 file_text(db.raw(), &space, "doc.txt").unwrap().as_deref(),
4151 Some("needle text")
4152 );
4153 assert_eq!(db.list_files(&space).unwrap()[0].status, "ok");
4154 }
4155
4156 #[allow(clippy::too_many_lines)]
4162 #[test]
4163 fn every_mutable_table_bumps_updated_at_on_mutation() {
4164 type Mutator<'a> = &'a dyn Fn(&Db) -> Result<()>;
4166 let mut db = Db::open_in_memory().unwrap();
4167 let space = db.default_space_id().unwrap();
4168 let sid = db.create_session("t", "a/b", &space, "chat").unwrap().id;
4169 let read = |table: &str, id: &str| -> String {
4170 std::thread::sleep(std::time::Duration::from_millis(2));
4171 db.raw()
4172 .query_row(
4173 &format!("SELECT updated_at FROM {table} WHERE id = ?1"),
4174 [id],
4175 |r| r.get::<_, String>(0),
4176 )
4177 .unwrap()
4178 };
4179
4180 let mutators: [Mutator<'_>; 6] = [
4182 &|db: &Db| db.set_compaction(&sid, "sum", 3),
4183 &|db: &Db| db.set_session_web_mode(&sid, true),
4184 &|db: &Db| db.set_session_swarm_mode(&sid, true),
4185 &|db: &Db| db.set_session_title(&sid, "new", Some("new-slug")),
4186 &|db: &Db| db.set_session_model(&sid, "m/x"),
4187 &|db: &Db| db.set_research_parent(&sid, "parent"),
4188 ];
4189 for mutate in mutators {
4190 let before = read("sessions", &sid);
4191 mutate(&db).unwrap();
4192 assert!(read("sessions", &sid) > before);
4193 }
4194 let before = read("sessions", &sid);
4196 db.save_swarm_personas(
4197 &sid,
4198 &[Persona {
4199 name: "a".into(),
4200 model: "m".into(),
4201 blurb: "b".into(),
4202 }],
4203 )
4204 .unwrap();
4205 assert!(read("sessions", &sid) > before);
4206
4207 assert!(db.toggle_favorite("a/model").unwrap());
4209 let before = read("model_prefs", "a/model");
4210 db.set_reasoning("a/model", Some("high")).unwrap();
4211 assert!(read("model_prefs", "a/model") > before);
4212 let before = read("model_prefs", "a/model");
4213 db.mark_model_used("a/model").unwrap();
4214 assert!(read("model_prefs", "a/model") > before);
4215
4216 let read_key = |key: &str| -> String {
4218 std::thread::sleep(std::time::Duration::from_millis(2));
4219 db.raw()
4220 .query_row(
4221 "SELECT updated_at FROM app_settings WHERE key = ?1",
4222 [key],
4223 |r| r.get::<_, String>(0),
4224 )
4225 .unwrap()
4226 };
4227 db.set_setting("temperature", "0.5").unwrap();
4228 let before = read_key("temperature");
4229 db.set_setting("temperature", "0.9").unwrap();
4230 assert!(read_key("temperature") > before);
4231
4232 let sp = db.create_space("other").unwrap();
4234 let before = read("spaces", &sp.id);
4235 db.rename_space(&sp.id, "other2").unwrap();
4236 assert!(read("spaces", &sp.id) > before);
4237
4238 let fid = db.upsert_file(&space, "f.txt", "h", 1, "ok").unwrap();
4240 let before = read("files", &fid);
4241 db.rename_file(&fid, "g.txt").unwrap();
4242 assert!(read("files", &fid) > before);
4243
4244 let wid = db.create_watch(&space, "topic", 24, &sid).unwrap();
4246 let before = read("watches", &wid);
4247 db.touch_watch(&wid, "2026-01-01T00:00:00Z").unwrap();
4248 assert!(read("watches", &wid) > before);
4249 let before = read("watches", &wid);
4250 db.set_watch_session(&wid, "other-session").unwrap();
4251 assert!(read("watches", &wid) > before);
4252
4253 let url = "https://x.test/";
4255 add_session_sources(db.raw(), &sid, &[url.to_string()]).unwrap();
4256 let read_src = || -> String {
4257 std::thread::sleep(std::time::Duration::from_millis(2));
4258 db.raw()
4259 .query_row(
4260 "SELECT updated_at FROM session_sources
4261 WHERE session_id = ?1 AND url_norm = ?2",
4262 (&sid, url),
4263 |r| r.get::<_, String>(0),
4264 )
4265 .unwrap()
4266 };
4267 let before = read_src();
4268 db.set_source_flag(&sid, url, Some("pinned")).unwrap();
4269 assert!(read_src() > before);
4270
4271 let row = db
4273 .log_usage("OpenRouter", "a/model", 1, 2, 3, 4, None, false, None, None)
4274 .unwrap();
4275 let read_usage = |db: &Db, row: i64| -> String {
4276 std::thread::sleep(std::time::Duration::from_millis(2));
4277 db.raw()
4278 .query_row(
4279 "SELECT updated_at FROM usage_log WHERE id = ?1",
4280 [&row],
4281 |r| r.get::<_, String>(0),
4282 )
4283 .unwrap()
4284 };
4285 let before = read_usage(&db, row);
4286 db.update_usage(row, 5, 6, 7, 8, Some(0.1), false).unwrap();
4287 assert!(read_usage(&db, row) > before);
4288 db.upsert_model_prices(&[(
4289 "a/model".to_string(),
4290 "OpenRouter".to_string(),
4291 price(1.0, 2.0),
4292 )])
4293 .unwrap();
4294 let before = read_usage(&db, row);
4295 assert_eq!(db.backfill_usage_costs().unwrap(), 1);
4296 assert!(read_usage(&db, row) > before);
4297 }
4298
4299 #[test]
4300 fn delete_paths_write_tombstones_for_sync() {
4301 let db = Db::open_in_memory().unwrap();
4302 let space = db.default_space_id().unwrap();
4303 let sid = db.create_session("t", "a/b", &space, "chat").unwrap().id;
4304 let tombstones = |table: &str| -> Vec<String> {
4305 let mut stmt = db
4306 .raw()
4307 .prepare("SELECT row_id FROM sync_tombstones WHERE table_name = ?1 ORDER BY row_id")
4308 .unwrap();
4309 let rows = stmt.query_map([table], |r| r.get::<_, String>(0)).unwrap();
4310 rows.collect::<rusqlite::Result<Vec<_>>>().unwrap()
4311 };
4312
4313 let mid = db.add_user_message(&sid, "hi").unwrap();
4315 db.delete_message(&mid).unwrap();
4316 assert_eq!(tombstones("messages"), vec![mid.clone()]);
4317
4318 let sid2 = db.create_session("t2", "a/b", &space, "chat").unwrap().id;
4320 let m1 = db.add_user_message(&sid2, "one").unwrap();
4321 let m2 = db.add_user_message(&sid2, "two").unwrap();
4322 db.save_swarm_personas(
4323 &sid2,
4324 &[Persona {
4325 name: "p".into(),
4326 model: "m".into(),
4327 blurb: "b".into(),
4328 }],
4329 )
4330 .unwrap();
4331 db.delete_session(&sid2).unwrap();
4332 let mut expected = vec![mid, m1, m2];
4333 expected.sort();
4334 assert_eq!(tombstones("messages"), expected);
4335 assert_eq!(tombstones("sessions"), vec![sid2.clone()]);
4336 assert_eq!(tombstones("swarm_personas"), vec![format!("{sid2}:0")]);
4337
4338 let fid = db.upsert_file(&space, "f.txt", "h", 1, "ok").unwrap();
4340 db.delete_file(&fid).unwrap();
4341 assert_eq!(tombstones("files"), vec![fid]);
4342
4343 let wid = db.create_watch(&space, "topic", 24, &sid).unwrap();
4345 db.delete_watch(&wid).unwrap();
4346 assert_eq!(tombstones("watches"), vec![wid]);
4347
4348 let sp = db.create_space("doomed").unwrap();
4350 db.delete_space(&sp.id).unwrap();
4351 assert_eq!(tombstones("spaces"), vec![sp.id]);
4352
4353 let persona = |name: &str| Persona {
4355 name: name.into(),
4356 model: "m".into(),
4357 blurb: "b".into(),
4358 };
4359 db.save_swarm_personas(&sid, &[persona("a"), persona("b")])
4360 .unwrap();
4361 db.save_swarm_personas(&sid, &[persona("b")]).unwrap();
4362 let mut expected_persona_tombstones =
4363 vec![format!("{sid2}:0"), format!("{sid}:0"), format!("{sid}:1")];
4364 expected_persona_tombstones.sort();
4365 assert_eq!(tombstones("swarm_personas"), expected_persona_tombstones);
4366 }
4367
4368 #[test]
4373 fn v1_default_space_renumbers_to_deterministic_id() {
4374 let dir = std::env::temp_dir().join(format!("nexus-v1-{}", uuid::Uuid::new_v4()));
4375 std::fs::create_dir_all(&dir).unwrap();
4376 let db_path = dir.join("nexus.db");
4377 {
4378 let db = Db::open(&db_path).unwrap();
4379 let old = db.default_space_id().unwrap();
4380 assert_eq!(old, DEFAULT_SPACE);
4381 db.raw().execute("PRAGMA user_version = 1", []).unwrap();
4384 db.raw()
4385 .execute(
4386 "UPDATE spaces SET id = 'legacy-uuid' WHERE name = ?1",
4387 [DEFAULT_SPACE],
4388 )
4389 .unwrap();
4390 db.raw()
4391 .execute(
4392 "UPDATE sessions SET space_id = 'legacy-uuid' WHERE space_id = ?1",
4393 [&old],
4394 )
4395 .unwrap();
4396 db.raw()
4397 .execute(
4398 "INSERT INTO sync_tombstones (table_name, row_id, deleted_at)
4399 VALUES ('spaces', 'legacy-uuid', '2026-01-01T00:00:00Z')",
4400 [],
4401 )
4402 .unwrap();
4403 let s = db.create_session("t", "m", &old, "chat").unwrap().id;
4404 let _ = s;
4405 drop(db);
4406 }
4407 let db = Db::open(&db_path).unwrap();
4410 assert_eq!(db.default_space_id().unwrap(), DEFAULT_SPACE);
4411 let n: i64 = db
4412 .raw()
4413 .query_row(
4414 "SELECT COUNT(*) FROM sessions WHERE space_id = ?1",
4415 [DEFAULT_SPACE],
4416 |r| r.get(0),
4417 )
4418 .unwrap();
4419 assert_eq!(n, 1, "sessions followed the renumber");
4420 let tomb: i64 = db
4421 .raw()
4422 .query_row(
4423 "SELECT COUNT(*) FROM sync_tombstones WHERE table_name = 'spaces' AND row_id = ?1",
4424 [DEFAULT_SPACE],
4425 |r| r.get(0),
4426 )
4427 .unwrap();
4428 assert_eq!(tomb, 1, "the space tombstone followed too");
4429 let _ = std::fs::remove_dir_all(&dir);
4430 }
4431
4432 #[test]
4433 fn device_id_is_stable_and_sync_state_roundtrips() {
4434 let db = Db::open_in_memory().unwrap();
4435 let a = db.device_id().unwrap();
4436 let b = db.device_id().unwrap();
4437 assert_eq!(a, b);
4438 assert!(!a.is_empty());
4439
4440 db.set_sync_state(
4442 "peer-1",
4443 "sessions",
4444 Some("2026-01-01T00:00:00Z|id1"),
4445 Some("2026-01-02T00:00:00Z|id2"),
4446 )
4447 .unwrap();
4448 db.set_sync_state("peer-1", "sessions", None, Some("2026-01-03T00:00:00Z|id3"))
4449 .unwrap();
4450 db.set_sync_state("peer-1", "messages", Some("c1"), None)
4451 .unwrap();
4452 let states = db.load_sync_state().unwrap();
4453 assert_eq!(states.len(), 2);
4454 let s = states.iter().find(|s| s.table_name == "sessions").unwrap();
4455 assert_eq!(s.pull_cursor.as_deref(), Some("2026-01-01T00:00:00Z|id1"));
4456 assert_eq!(s.push_cursor.as_deref(), Some("2026-01-03T00:00:00Z|id3"));
4457 assert!(s.last_synced_at.is_some());
4458 let m = states.iter().find(|s| s.table_name == "messages").unwrap();
4459 assert_eq!(m.pull_cursor.as_deref(), Some("c1"));
4460 }
4461
4462 #[test]
4463 fn settings_scope_registry_classifies_keys_and_stores_scope() {
4464 for local in [
4466 "searxng_url",
4467 "langsearch_key",
4468 "search_provider",
4469 "ocr_engine",
4470 "ocr_model",
4471 "local_ocr_model",
4472 "usage_range",
4473 "ui_background",
4474 "last_update_check",
4475 ] {
4476 assert!(Db::setting_is_local(local), "{local} should be local");
4477 }
4478 for sync in [
4479 "temperature",
4480 "top_p",
4481 "max_tokens",
4482 "show_stats",
4483 "show_reasoning",
4484 "hide_hints",
4485 "verbosity",
4486 "memory_model",
4487 "embedding_model",
4488 ] {
4489 assert!(!Db::setting_is_local(sync), "{sync} should sync");
4490 }
4491
4492 let db = Db::open_in_memory().unwrap();
4493 db.set_setting("temperature", "0.7").unwrap();
4494 db.set_setting("ocr_engine", "tesseract").unwrap();
4495 let scope = |key: &str| -> String {
4496 db.raw()
4497 .query_row(
4498 "SELECT scope FROM app_settings WHERE key = ?1",
4499 [key],
4500 |r| r.get(0),
4501 )
4502 .unwrap()
4503 };
4504 assert_eq!(scope("temperature"), "sync");
4505 assert_eq!(scope("ocr_engine"), "local");
4506 }
4507
4508 #[test]
4509 fn sync_ids_are_unique_for_citations_and_usage() {
4510 let db = Db::open_in_memory().unwrap();
4511 db.add_citations(
4512 "sp",
4513 "r.md",
4514 &[
4515 ("https://a.test/".to_string(), None),
4516 ("https://b.test/".to_string(), None),
4517 ],
4518 )
4519 .unwrap();
4520 db.log_usage("OpenRouter", "m", 1, 2, 0, 0, None, false, None, None)
4521 .unwrap();
4522 db.log_usage("OpenRouter", "m", 1, 2, 0, 0, None, false, None, None)
4523 .unwrap();
4524 let distinct: i64 = db
4525 .raw()
4526 .query_row(
4527 "SELECT COUNT(DISTINCT sync_id) FROM
4528 (SELECT sync_id FROM citations UNION ALL SELECT sync_id FROM usage_log)",
4529 [],
4530 |r| r.get(0),
4531 )
4532 .unwrap();
4533 let total: i64 = db
4534 .raw()
4535 .query_row(
4536 "SELECT (SELECT COUNT(*) FROM citations) + (SELECT COUNT(*) FROM usage_log)",
4537 [],
4538 |r| r.get(0),
4539 )
4540 .unwrap();
4541 assert_eq!(distinct, total);
4542 }
4543}