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 | "last_update_check"
820 | "sync_ssh_peer"
823 )
824 }
825
826 pub fn set_setting(&self, key: &str, value: &str) -> Result<()> {
827 let now = Utc::now().to_rfc3339();
828 let scope = if Self::setting_is_local(key) {
829 "local"
830 } else {
831 "sync"
832 };
833 self.conn.execute(
834 "INSERT INTO app_settings (key, value, scope, updated_at)
835 VALUES (?1, ?2, ?3, ?4)
836 ON CONFLICT(key) DO UPDATE SET value = ?2, scope = ?3, updated_at = ?4",
837 (key, value, scope, &now),
838 )?;
839 Ok(())
840 }
841
842 pub fn load_settings(&self) -> Result<Vec<(String, String)>> {
843 let mut stmt = self.conn.prepare("SELECT key, value FROM app_settings")?;
844 let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
845 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
846 }
847
848 pub fn update_check_due(&self) -> bool {
853 let today = chrono::Local::now().format("%Y-%m-%d").to_string();
854 let last = self
855 .conn
856 .query_row(
857 "SELECT value FROM app_settings WHERE key = 'last_update_check'",
858 [],
859 |r| r.get::<_, String>(0),
860 )
861 .ok();
862 if last.as_deref() == Some(today.as_str()) {
863 return false;
864 }
865 let _ = self.set_setting("last_update_check", &today);
866 true
867 }
868
869 pub fn set_reasoning(&self, model_id: &str, effort: Option<&str>) -> Result<()> {
871 self.conn.execute(
872 "INSERT INTO model_prefs (id, reasoning, updated_at) VALUES (?1, ?2, ?3)
873 ON CONFLICT(id) DO UPDATE SET reasoning = ?2, updated_at = ?3",
874 (model_id, effort, Utc::now().to_rfc3339()),
875 )?;
876 Ok(())
877 }
878
879 pub fn toggle_favorite(&self, model_id: &str) -> Result<bool> {
881 self.conn.execute(
882 "INSERT INTO model_prefs (id, favorite, updated_at) VALUES (?1, 1, ?2)
883 ON CONFLICT(id) DO UPDATE SET favorite = 1 - favorite, updated_at = ?2",
884 (model_id, Utc::now().to_rfc3339()),
885 )?;
886 let fav: i64 = self.conn.query_row(
887 "SELECT favorite FROM model_prefs WHERE id = ?1",
888 [model_id],
889 |r| r.get(0),
890 )?;
891 Ok(fav != 0)
892 }
893
894 pub fn mark_model_used(&self, model_id: &str) -> Result<()> {
896 let now = Utc::now().to_rfc3339();
897 self.conn.execute(
898 "INSERT INTO model_prefs (id, favorite, last_used, updated_at) VALUES (?1, 0, ?2, ?2)
899 ON CONFLICT(id) DO UPDATE SET last_used = ?2, updated_at = ?2",
900 (model_id, &now),
901 )?;
902 Ok(())
903 }
904
905 pub fn load_model_prefs(&self) -> Result<Vec<ModelPref>> {
907 let mut stmt = self
908 .conn
909 .prepare("SELECT id, favorite, last_used, reasoning FROM model_prefs")?;
910 let rows = stmt.query_map([], |r| {
911 Ok(ModelPref {
912 id: r.get(0)?,
913 favorite: r.get::<_, i64>(1)? != 0,
914 last_used: r.get(2)?,
915 reasoning: r.get(3)?,
916 })
917 })?;
918 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
919 }
920
921 pub fn create_session(
922 &self,
923 title: &str,
924 model: &str,
925 space_id: &str,
926 kind: &str,
927 ) -> Result<Session> {
928 let id = Uuid::new_v4().to_string();
929 let now = Utc::now().to_rfc3339();
930 self.conn.execute(
931 "INSERT INTO sessions (id, title, model, space_id, kind, created_at, updated_at)
932 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6)",
933 (&id, title, model, space_id, kind, &now),
934 )?;
935 Ok(Session {
936 id,
937 title: title.to_string(),
938 model: model.to_string(),
939 slug: None,
940 created_at: now,
941 compact_summary: None,
942 compact_through: 0,
943 web_mode: false,
944 swarm_mode: false,
945 kind: kind.to_string(),
946 research_parent_id: None,
947 })
948 }
949
950 pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
953 self.conn
954 .query_row(
955 "SELECT id, title, model, slug, created_at, compact_summary, compact_through, \
956 web_mode, swarm_mode, kind, research_parent_id
957 FROM sessions WHERE id = ?1",
958 [id],
959 |r| {
960 Ok(Session {
961 id: r.get(0)?,
962 title: r.get(1)?,
963 model: r.get(2)?,
964 slug: r.get(3)?,
965 created_at: r.get(4)?,
966 compact_summary: r.get(5)?,
967 compact_through: r.get(6)?,
968 web_mode: r.get::<_, i64>(7)? != 0,
969 swarm_mode: r.get::<_, i64>(8)? != 0,
970 kind: r.get(9)?,
971 research_parent_id: r.get(10)?,
972 })
973 },
974 )
975 .optional()
976 .map_err(Into::into)
977 }
978
979 pub fn list_sessions(&self, space_id: &str) -> Result<Vec<Session>> {
981 let mut stmt = self.conn.prepare(
982 "SELECT id, title, model, slug, created_at, compact_summary, compact_through, \
983 web_mode, swarm_mode, kind, research_parent_id
984 FROM sessions WHERE space_id = ?1 ORDER BY updated_at DESC",
985 )?;
986 let rows = stmt.query_map([space_id], |r| {
987 Ok(Session {
988 id: r.get(0)?,
989 title: r.get(1)?,
990 model: r.get(2)?,
991 slug: r.get(3)?,
992 created_at: r.get(4)?,
993 compact_summary: r.get(5)?,
994 compact_through: r.get(6)?,
995 web_mode: r.get::<_, i64>(7)? != 0,
996 swarm_mode: r.get::<_, i64>(8)? != 0,
997 kind: r.get(9)?,
998 research_parent_id: r.get(10)?,
999 })
1000 })?;
1001 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1002 }
1003
1004 pub fn set_compaction(&self, session_id: &str, summary: &str, through: i64) -> Result<()> {
1007 self.conn.execute(
1008 "UPDATE sessions SET compact_summary = ?2, compact_through = ?3, updated_at = ?4
1009 WHERE id = ?1",
1010 (session_id, summary, through, Utc::now().to_rfc3339()),
1011 )?;
1012 Ok(())
1013 }
1014
1015 pub fn set_session_web_mode(&self, session_id: &str, on: bool) -> Result<()> {
1017 self.conn.execute(
1018 "UPDATE sessions SET web_mode = ?2, updated_at = ?3 WHERE id = ?1",
1019 (session_id, i64::from(on), Utc::now().to_rfc3339()),
1020 )?;
1021 Ok(())
1022 }
1023
1024 pub fn set_session_swarm_mode(&self, session_id: &str, on: bool) -> Result<()> {
1026 self.conn.execute(
1027 "UPDATE sessions SET swarm_mode = ?2, updated_at = ?3 WHERE id = ?1",
1028 (session_id, i64::from(on), Utc::now().to_rfc3339()),
1029 )?;
1030 Ok(())
1031 }
1032
1033 pub fn list_swarm_personas(&self, session_id: &str) -> Result<Vec<Persona>> {
1035 let mut stmt = self.conn.prepare(
1036 "SELECT name, model, persona FROM swarm_personas
1037 WHERE session_id = ?1 ORDER BY ord ASC",
1038 )?;
1039 let rows = stmt.query_map([session_id], |r| {
1040 Ok(Persona {
1041 name: r.get(0)?,
1042 model: r.get(1)?,
1043 blurb: r.get(2)?,
1044 })
1045 })?;
1046 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1047 }
1048
1049 pub fn save_swarm_personas(&self, session_id: &str, personas: &[Persona]) -> Result<()> {
1054 let old: Vec<i64> = {
1055 let mut stmt = self
1056 .conn
1057 .prepare("SELECT ord FROM swarm_personas WHERE session_id = ?1")?;
1058 let rows = stmt.query_map([session_id], |r| r.get(0))?;
1059 rows.collect::<rusqlite::Result<Vec<_>>>()?
1060 };
1061 for ord in old {
1062 self.tombstone("swarm_personas", &format!("{session_id}:{ord}"))?;
1063 }
1064 self.conn.execute(
1065 "DELETE FROM swarm_personas WHERE session_id = ?1",
1066 [session_id],
1067 )?;
1068 for (i, p) in personas.iter().enumerate() {
1069 self.conn.execute(
1070 "INSERT INTO swarm_personas (session_id, ord, name, model, persona)
1071 VALUES (?1, ?2, ?3, ?4, ?5)",
1072 (session_id, i as i64, &p.name, &p.model, &p.blurb),
1073 )?;
1074 }
1075 self.touch("sessions", session_id)
1076 }
1077
1078 pub fn set_research_parent(&self, id: &str, parent_id: &str) -> Result<()> {
1082 self.conn.execute(
1083 "UPDATE sessions SET research_parent_id = ?2, updated_at = ?3 WHERE id = ?1",
1084 (id, parent_id, Utc::now().to_rfc3339()),
1085 )?;
1086 Ok(())
1087 }
1088
1089 pub fn set_session_title(&self, id: &str, title: &str, slug: Option<&str>) -> Result<()> {
1091 self.conn.execute(
1092 "UPDATE sessions SET title = ?2, slug = COALESCE(?3, slug), updated_at = ?4
1093 WHERE id = ?1",
1094 (id, title, slug, Utc::now().to_rfc3339()),
1095 )?;
1096 Ok(())
1097 }
1098
1099 pub fn delete_message(&self, id: &str) -> Result<()> {
1103 self.conn
1104 .execute("DELETE FROM messages WHERE id = ?1", [id])?;
1105 self.tombstone("messages", id)?;
1106 Ok(())
1107 }
1108
1109 pub fn delete_session(&self, id: &str) -> Result<()> {
1113 let ids: Vec<String> = {
1114 let mut stmt = self
1115 .conn
1116 .prepare("SELECT id FROM messages WHERE session_id = ?1")?;
1117 let rows = stmt.query_map([id], |r| r.get(0))?;
1118 rows.collect::<rusqlite::Result<Vec<_>>>()?
1119 };
1120 for mid in &ids {
1121 self.tombstone("messages", mid)?;
1122 }
1123 self.conn
1124 .execute("DELETE FROM messages WHERE session_id = ?1", [id])?;
1125 self.conn
1126 .execute("DELETE FROM sessions WHERE id = ?1", [id])?;
1127 self.tombstone("sessions", id)?;
1128 Ok(())
1129 }
1130
1131 pub fn load_messages(&self, session_id: &str) -> Result<Vec<Message>> {
1132 let mut stmt = self.conn.prepare(
1133 "SELECT role, content, model, reasoning, tokens, secs, cost, phrase, persona, created_at
1134 FROM messages WHERE session_id = ?1 ORDER BY created_at ASC",
1135 )?;
1136 let messages = stmt
1137 .query_map([session_id], |r| {
1138 Ok(Message {
1139 role: r.get(0)?,
1140 content: r.get(1)?,
1141 model: r.get(2)?,
1142 reasoning: r.get(3)?,
1143 tokens: r.get(4)?,
1144 secs: r.get(5)?,
1145 cost: r.get(6)?,
1146 phrase: r.get(7)?,
1147 persona: r.get(8)?,
1148 created_at: r.get(9)?,
1149 })
1150 })?
1151 .collect::<rusqlite::Result<Vec<_>>>()?;
1152 Ok(messages)
1153 }
1154
1155 pub fn add_user_message(&self, session_id: &str, content: &str) -> Result<String> {
1157 self.insert_message(
1158 session_id, "user", content, None, None, None, None, None, None,
1159 )
1160 }
1161
1162 pub fn add_gate_reply_message(&self, session_id: &str, content: &str) -> Result<String> {
1168 self.insert_message(
1169 session_id,
1170 "gate_reply",
1171 content,
1172 None,
1173 None,
1174 None,
1175 None,
1176 None,
1177 None,
1178 )
1179 }
1180
1181 pub fn add_tool_call_message(&self, session_id: &str, content: &str) -> Result<String> {
1184 self.insert_message(
1185 session_id,
1186 "tool_call",
1187 content,
1188 None,
1189 None,
1190 None,
1191 None,
1192 None,
1193 None,
1194 )
1195 }
1196
1197 pub fn add_error_message(&self, session_id: &str, content: &str) -> Result<String> {
1200 self.insert_message(
1201 session_id, "error", content, None, None, None, None, None, None,
1202 )
1203 }
1204
1205 pub fn add_research_stage_message(&self, session_id: &str, content: &str) -> Result<String> {
1210 self.insert_message(
1211 session_id,
1212 "research_stage",
1213 content,
1214 None,
1215 None,
1216 None,
1217 None,
1218 None,
1219 None,
1220 )
1221 }
1222
1223 pub fn add_research_plan_message(&self, session_id: &str, content: &str) -> Result<String> {
1226 self.insert_message(
1227 session_id,
1228 "research_plan",
1229 content,
1230 None,
1231 None,
1232 None,
1233 None,
1234 None,
1235 None,
1236 )
1237 }
1238
1239 pub fn add_survey_message(&self, session_id: &str, content: &str) -> Result<String> {
1243 self.insert_message(
1244 session_id, "survey", content, None, None, None, None, None, None,
1245 )
1246 }
1247
1248 pub fn upsert_research_stage_message(
1253 &self,
1254 session_id: &str,
1255 label: &str,
1256 detail: &str,
1257 ) -> Result<()> {
1258 let content = stage_content(label, detail);
1259 let existing: Option<String> = self
1260 .conn
1261 .query_row(
1262 "SELECT id FROM messages WHERE session_id = ?1 AND role = 'research_stage'
1263 AND (content = ?2 OR content LIKE ?3)
1264 ORDER BY created_at DESC LIMIT 1",
1265 (session_id, label, format!("{label}:%")),
1266 |r| r.get(0),
1267 )
1268 .ok();
1269 match existing {
1270 Some(id) => {
1271 let now = Utc::now().to_rfc3339();
1272 self.conn.execute(
1273 "UPDATE messages SET content = ?2, created_at = ?3 WHERE id = ?1",
1274 (&id, &content, &now),
1275 )?;
1276 }
1277 None => {
1278 self.add_research_stage_message(session_id, &content)?;
1279 }
1280 }
1281 Ok(())
1282 }
1283
1284 #[cfg(test)]
1288 pub fn add_session_sources(&self, session_id: &str, url_norms: &[String]) -> Result<()> {
1289 add_session_sources(&self.conn, session_id, url_norms)
1290 }
1291
1292 #[cfg(test)]
1294 pub fn search_session_sources(
1295 &self,
1296 session_id: &str,
1297 query: &str,
1298 ) -> Result<Vec<(String, String)>> {
1299 search_session_sources(&self.conn, session_id, query)
1300 }
1301
1302 pub fn set_source_flag(
1308 &self,
1309 session_id: &str,
1310 url_norm: &str,
1311 flag: Option<&str>,
1312 ) -> Result<()> {
1313 self.conn.execute(
1314 "UPDATE session_sources SET flag = ?3, updated_at = ?4
1315 WHERE session_id = ?1 AND url_norm = ?2",
1316 (session_id, url_norm, flag, Utc::now().to_rfc3339()),
1317 )?;
1318 Ok(())
1319 }
1320
1321 #[allow(clippy::too_many_arguments)]
1327 pub fn add_assistant_message(
1328 &self,
1329 session_id: &str,
1330 content: &str,
1331 model: Option<&str>,
1332 reasoning: Option<&str>,
1333 tokens: Option<i64>,
1334 secs: Option<f64>,
1335 cost: Option<f64>,
1336 phrase: Option<&str>,
1337 ) -> Result<String> {
1338 self.insert_message(
1339 session_id,
1340 "assistant",
1341 content,
1342 model,
1343 reasoning,
1344 tokens,
1345 secs,
1346 cost,
1347 phrase,
1348 )
1349 }
1350
1351 pub fn add_persona_message(
1354 &self,
1355 session_id: &str,
1356 content: &str,
1357 persona_name: &str,
1358 model: &str,
1359 ) -> Result<String> {
1360 let id = self.insert_message(
1361 session_id,
1362 "assistant",
1363 content,
1364 Some(model),
1365 None,
1366 None,
1367 None,
1368 None,
1369 None,
1370 )?;
1371 self.conn.execute(
1372 "UPDATE messages SET persona = ?2 WHERE id = ?1",
1373 (&id, persona_name),
1374 )?;
1375 Ok(id)
1376 }
1377
1378 #[allow(clippy::too_many_arguments)]
1381 pub(crate) fn insert_message(
1382 &self,
1383 session_id: &str,
1384 role: &str,
1385 content: &str,
1386 model: Option<&str>,
1387 reasoning: Option<&str>,
1388 tokens: Option<i64>,
1389 secs: Option<f64>,
1390 cost: Option<f64>,
1391 phrase: Option<&str>,
1392 ) -> Result<String> {
1393 let now = Utc::now().to_rfc3339();
1394 let id = Uuid::new_v4().to_string();
1395 self.conn.execute(
1396 "INSERT INTO messages
1397 (id, session_id, role, content, model, reasoning, tokens, secs, cost, phrase, created_at)
1398 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
1399 (
1400 &id, session_id, role, content, model, reasoning, tokens, secs, cost, phrase, &now,
1401 ),
1402 )?;
1403 self.conn.execute(
1404 "UPDATE sessions SET updated_at = ?2 WHERE id = ?1",
1405 (session_id, &now),
1406 )?;
1407 Ok(id)
1408 }
1409
1410 pub fn message_created_at(&self, session_id: &str, index: usize) -> Result<Option<String>> {
1415 let mut stmt = self.conn.prepare(
1416 "SELECT created_at FROM messages WHERE session_id = ?1
1417 ORDER BY created_at ASC LIMIT 1 OFFSET ?2",
1418 )?;
1419 Ok(stmt
1420 .query_row((session_id, index as i64), |r| r.get(0))
1421 .optional()?)
1422 }
1423
1424 pub fn add_compaction_message(
1431 &self,
1432 session_id: &str,
1433 content: &str,
1434 at: &str,
1435 ) -> Result<String> {
1436 let id = Uuid::new_v4().to_string();
1437 self.conn.execute(
1438 "INSERT INTO messages
1439 (id, session_id, role, content, model, reasoning, tokens, secs, phrase, created_at)
1440 VALUES (?1, ?2, 'compaction', ?3, NULL, NULL, NULL, NULL, NULL, ?4)",
1441 (&id, session_id, content, at),
1442 )?;
1443 Ok(id)
1444 }
1445
1446 pub fn update_compaction_message(&self, session_id: &str, content: &str) -> Result<usize> {
1451 Ok(self.conn.execute(
1452 "UPDATE messages SET content = ?2
1453 WHERE session_id = ?1 AND role = 'compaction'",
1454 (session_id, content),
1455 )?)
1456 }
1457
1458 pub fn set_session_model(&self, session_id: &str, model: &str) -> Result<()> {
1459 self.conn.execute(
1460 "UPDATE sessions SET model = ?2, updated_at = ?3 WHERE id = ?1",
1461 (session_id, model, Utc::now().to_rfc3339()),
1462 )?;
1463 Ok(())
1464 }
1465
1466 pub fn upsert_file(
1474 &self,
1475 space_id: &str,
1476 name: &str,
1477 hash: &str,
1478 size: i64,
1479 status: &str,
1480 ) -> Result<String> {
1481 let now = Utc::now().to_rfc3339();
1482 if let Ok(existing) = self.conn.query_row(
1483 "SELECT id FROM files WHERE space_id = ?1 AND name = ?2",
1484 (space_id, name),
1485 |r| r.get::<_, String>(0),
1486 ) {
1487 self.conn.execute(
1488 "UPDATE files SET hash = ?2, size = ?3, updated_at = ?4 WHERE id = ?1",
1489 (&existing, hash, size, &now),
1490 )?;
1491 self.conn.execute(
1492 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1493 VALUES (?1, 0, ?2, ?3)
1494 ON CONFLICT(file_id) DO UPDATE SET status = excluded.status,
1495 updated_at = excluded.updated_at",
1496 (&existing, status, &now),
1497 )?;
1498 return Ok(existing);
1499 }
1500 let id = Uuid::new_v4().to_string();
1501 self.conn.execute(
1502 "INSERT INTO files (id, space_id, name, hash, size, created_at, updated_at)
1503 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
1504 (&id, space_id, name, hash, size, &now, &now),
1505 )?;
1506 self.conn.execute(
1507 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1508 VALUES (?1, 0, ?2, ?3)",
1509 (&id, status, &now),
1510 )?;
1511 Ok(id)
1512 }
1513
1514 pub fn list_files(&self, space_id: &str) -> Result<Vec<FileRow>> {
1515 let mut stmt = self.conn.prepare(
1516 "SELECT files.id, files.name, files.hash, files.size,
1517 COALESCE(cache.file_index_state.status, 'not indexed'),
1518 COALESCE(cache.file_index_state.mtime, 0)
1519 FROM files
1520 LEFT JOIN cache.file_index_state
1521 ON cache.file_index_state.file_id = files.id
1522 WHERE files.space_id = ?1 ORDER BY files.name ASC",
1523 )?;
1524 let rows = stmt.query_map([space_id], |r| {
1525 Ok(FileRow {
1526 id: r.get(0)?,
1527 name: r.get(1)?,
1528 hash: r.get(2)?,
1529 size: r.get(3)?,
1530 status: r.get(4)?,
1531 mtime: r.get(5)?,
1532 })
1533 })?;
1534 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1535 }
1536
1537 pub fn file_indexed(&self, file_id: &str) -> Result<bool> {
1542 Ok(self.conn.query_row(
1543 "SELECT EXISTS(SELECT 1 FROM cache.file_index_state WHERE file_id = ?1)",
1544 [file_id],
1545 |r| r.get(0),
1546 )?)
1547 }
1548
1549 pub fn delete_file(&self, file_id: &str) -> Result<()> {
1550 self.conn.execute(
1551 "DELETE FROM cache.file_chunks WHERE file_id = ?1",
1552 [file_id],
1553 )?;
1554 self.conn.execute(
1555 "DELETE FROM cache.chunk_embeddings WHERE file_id = ?1",
1556 [file_id],
1557 )?;
1558 self.conn.execute(
1559 "DELETE FROM cache.file_index_state WHERE file_id = ?1",
1560 [file_id],
1561 )?;
1562 self.conn
1563 .execute("DELETE FROM files WHERE id = ?1", [file_id])?;
1564 self.tombstone("files", file_id)?;
1565 Ok(())
1566 }
1567
1568 pub fn set_file_mtime(&self, file_id: &str, mtime: i64) -> Result<()> {
1572 let now = Utc::now().to_rfc3339();
1573 self.conn.execute(
1574 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1575 VALUES (?1, ?2, '', ?3)
1576 ON CONFLICT(file_id) DO UPDATE SET mtime = excluded.mtime,
1577 updated_at = excluded.updated_at",
1578 (file_id, mtime, &now),
1579 )?;
1580 Ok(())
1581 }
1582
1583 pub fn set_file_status(&self, file_id: &str, status: &str) -> Result<()> {
1587 let now = Utc::now().to_rfc3339();
1588 self.conn.execute(
1589 "INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
1590 VALUES (?1, 0, ?2, ?3)
1591 ON CONFLICT(file_id) DO UPDATE SET status = excluded.status,
1592 updated_at = excluded.updated_at",
1593 (file_id, status, &now),
1594 )?;
1595 Ok(())
1596 }
1597
1598 pub fn rename_file(&self, file_id: &str, new_name: &str) -> Result<()> {
1599 self.conn.execute(
1600 "UPDATE files SET name = ?2, updated_at = ?3 WHERE id = ?1",
1601 (file_id, new_name, Utc::now().to_rfc3339()),
1602 )?;
1603 Ok(())
1604 }
1605
1606 pub fn replace_file_ref_in_messages(
1610 &self,
1611 space_id: &str,
1612 old_name: &str,
1613 new_name: &str,
1614 ) -> Result<()> {
1615 self.conn.execute(
1616 "UPDATE messages SET content = REPLACE(content, ?1, ?2)
1617 WHERE session_id IN (SELECT id FROM sessions WHERE space_id = ?3)",
1618 (old_name, new_name, space_id),
1619 )?;
1620 Ok(())
1621 }
1622
1623 pub fn set_file_chunks(&self, file_id: &str, chunks: &[(String, String)]) -> Result<()> {
1628 self.conn.execute(
1629 "DELETE FROM cache.file_chunks WHERE file_id = ?1",
1630 [file_id],
1631 )?;
1632 self.conn.execute(
1633 "DELETE FROM cache.chunk_embeddings WHERE file_id = ?1",
1634 [file_id],
1635 )?;
1636 for (seq, (location, text)) in chunks.iter().enumerate() {
1637 self.conn.execute(
1638 "INSERT INTO cache.file_chunks (file_id, seq, location, text) VALUES (?1, ?2, ?3, ?4)",
1639 (file_id, seq as i64, location, text),
1640 )?;
1641 }
1642 Ok(())
1643 }
1644
1645 #[cfg(test)]
1648 pub fn raw(&self) -> &Connection {
1649 &self.conn
1650 }
1651
1652 pub fn integrity_check(&self) -> Result<String> {
1656 self.conn
1657 .query_row("PRAGMA integrity_check", [], |r| r.get(0))
1658 .context("running integrity check")
1659 }
1660
1661 pub fn file_chunk_texts(&self, file_id: &str) -> Result<Vec<(i64, String)>> {
1663 let mut stmt = self.conn.prepare(
1664 "SELECT CAST(seq AS INTEGER), text FROM cache.file_chunks
1665 WHERE file_id = ?1 ORDER BY CAST(seq AS INTEGER) ASC",
1666 )?;
1667 let rows = stmt.query_map([file_id], |r| Ok((r.get(0)?, r.get(1)?)))?;
1668 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1669 }
1670
1671 pub fn files_missing_embeddings(&self, space_id: &str) -> Result<Vec<String>> {
1673 files_missing_embeddings(&self.conn, space_id)
1674 }
1675
1676 pub fn add_citations(
1680 &self,
1681 space_id: &str,
1682 report_file: &str,
1683 citations: &[(String, Option<String>)],
1684 ) -> Result<()> {
1685 for (url, title) in citations {
1686 self.conn.execute(
1687 "INSERT INTO citations (sync_id, space_id, report_file, url, title)
1688 VALUES (?1, ?2, ?3, ?4, ?5)",
1689 (
1690 Uuid::new_v4().to_string(),
1691 space_id,
1692 report_file,
1693 url,
1694 title,
1695 ),
1696 )?;
1697 }
1698 Ok(())
1699 }
1700
1701 pub fn search_citations(
1706 &self,
1707 space_id: &str,
1708 query: Option<&str>,
1709 ) -> Result<Vec<(String, String, String)>> {
1710 search_citations(&self.conn, space_id, query)
1711 }
1712
1713 pub fn set_chunk_embeddings(&self, file_id: &str, vecs: &[(i64, Vec<f32>)]) -> Result<()> {
1715 for (seq, v) in vecs {
1716 self.conn.execute(
1717 "INSERT OR REPLACE INTO cache.chunk_embeddings (file_id, seq, vec) VALUES (?1, ?2, ?3)",
1718 (file_id, seq, vec_to_blob(v)),
1719 )?;
1720 }
1721 Ok(())
1722 }
1723
1724 pub fn create_watch(
1725 &self,
1726 space_id: &str,
1727 topic: &str,
1728 interval_hours: i64,
1729 session_id: &str,
1730 ) -> Result<String> {
1731 let id = Uuid::new_v4().to_string();
1732 let now = Utc::now().to_rfc3339();
1733 self.conn.execute(
1734 "INSERT INTO watches (id, space_id, topic, interval_hours, session_id, last_run_at, updated_at)
1735 VALUES (?1, ?2, ?3, ?4, ?5, NULL, ?6)",
1736 (&id, space_id, topic, interval_hours, session_id, &now),
1737 )?;
1738 Ok(id)
1739 }
1740
1741 pub fn list_watches(&self, space_id: &str) -> Result<Vec<Watch>> {
1742 let mut stmt = self.conn.prepare(
1743 "SELECT id, space_id, topic, interval_hours, session_id, last_run_at
1744 FROM watches WHERE space_id = ?1 ORDER BY topic",
1745 )?;
1746 let rows = stmt.query_map([space_id], |r| {
1747 Ok(Watch {
1748 id: r.get(0)?,
1749 space_id: r.get(1)?,
1750 topic: r.get(2)?,
1751 interval_hours: r.get(3)?,
1752 session_id: r.get(4)?,
1753 last_run_at: r.get(5)?,
1754 })
1755 })?;
1756 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1757 }
1758
1759 pub fn list_all_watches(&self) -> Result<Vec<Watch>> {
1762 let mut stmt = self.conn.prepare(
1763 "SELECT id, space_id, topic, interval_hours, session_id, last_run_at FROM watches",
1764 )?;
1765 let rows = stmt.query_map([], |r| {
1766 Ok(Watch {
1767 id: r.get(0)?,
1768 space_id: r.get(1)?,
1769 topic: r.get(2)?,
1770 interval_hours: r.get(3)?,
1771 session_id: r.get(4)?,
1772 last_run_at: r.get(5)?,
1773 })
1774 })?;
1775 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1776 }
1777
1778 pub fn touch_watch(&self, id: &str, now_rfc3339: &str) -> Result<()> {
1779 self.conn.execute(
1780 "UPDATE watches SET last_run_at = ?2, updated_at = ?3 WHERE id = ?1",
1781 (id, now_rfc3339, Utc::now().to_rfc3339()),
1782 )?;
1783 Ok(())
1784 }
1785
1786 pub fn set_watch_session(&self, id: &str, session_id: &str) -> Result<()> {
1790 self.conn.execute(
1791 "UPDATE watches SET session_id = ?2, updated_at = ?3 WHERE id = ?1",
1792 (id, session_id, Utc::now().to_rfc3339()),
1793 )?;
1794 Ok(())
1795 }
1796
1797 pub fn delete_watch(&self, id: &str) -> Result<()> {
1798 self.conn
1799 .execute("DELETE FROM watches WHERE id = ?1", [id])?;
1800 self.tombstone("watches", id)?;
1801 Ok(())
1802 }
1803}
1804
1805pub fn vec_to_blob(v: &[f32]) -> Vec<u8> {
1807 v.iter().flat_map(|f| f.to_le_bytes()).collect()
1808}
1809
1810pub fn blob_to_vec(b: &[u8]) -> Vec<f32> {
1812 b.chunks_exact(4)
1813 .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
1814 .collect()
1815}
1816
1817pub fn search_citations(
1822 conn: &Connection,
1823 space_id: &str,
1824 query: Option<&str>,
1825) -> Result<Vec<(String, String, String)>> {
1826 let mut stmt = conn.prepare(
1827 "SELECT report_file, url, COALESCE(title, '') FROM citations
1828 WHERE space_id = ?1
1829 AND (?2 IS NULL OR url LIKE ?2 OR title LIKE ?2 OR report_file LIKE ?2)
1830 ORDER BY id DESC",
1831 )?;
1832 let pattern = query.map(|q| format!("%{q}%"));
1833 let rows = stmt.query_map((space_id, pattern), |r| {
1834 Ok((r.get(0)?, r.get(1)?, r.get(2)?))
1835 })?;
1836 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1837}
1838
1839pub fn stage_content(label: &str, detail: &str) -> String {
1841 if detail.is_empty() {
1842 label.to_string()
1843 } else {
1844 format!("{label}: {detail}")
1845 }
1846}
1847
1848pub fn add_session_sources(
1851 conn: &Connection,
1852 session_id: &str,
1853 url_norms: &[String],
1854) -> Result<()> {
1855 let now = Utc::now().to_rfc3339();
1856 for u in url_norms {
1857 conn.execute(
1858 "INSERT OR IGNORE INTO session_sources (session_id, url_norm, updated_at)
1859 VALUES (?1, ?2, ?3)",
1860 (session_id, u, &now),
1861 )?;
1862 }
1863 Ok(())
1864}
1865
1866pub fn search_session_sources(
1871 conn: &Connection,
1872 session_id: &str,
1873 query: &str,
1874) -> Result<Vec<(String, String)>> {
1875 let mut stmt = conn.prepare(
1876 "SELECT cache.web_cache.url, cache.web_cache.text FROM session_sources
1877 JOIN cache.web_cache ON cache.web_cache.url_norm = session_sources.url_norm
1878 WHERE session_sources.session_id = ?1",
1879 )?;
1880 let rows = stmt.query_map([session_id], |r| {
1881 Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
1882 })?;
1883 let needle = query.to_lowercase();
1884 Ok(rows
1885 .collect::<rusqlite::Result<Vec<_>>>()?
1886 .into_iter()
1887 .filter(|(_, text)| text.to_lowercase().contains(&needle))
1888 .collect())
1889}
1890
1891pub fn pinned_urls(conn: &Connection, session_id: &str) -> Result<Vec<String>> {
1894 let mut stmt = conn.prepare(
1895 "SELECT url_norm FROM session_sources WHERE session_id = ?1 AND flag = 'pinned'",
1896 )?;
1897 let rows = stmt.query_map([session_id], |r| r.get::<_, String>(0))?;
1898 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1899}
1900
1901pub fn discarded_domains(conn: &Connection, session_id: &str) -> Result<Vec<String>> {
1904 let mut stmt = conn.prepare(
1905 "SELECT url_norm FROM session_sources WHERE session_id = ?1 AND flag = 'discarded'",
1906 )?;
1907 let rows: Vec<String> = stmt
1908 .query_map([session_id], |r| r.get::<_, String>(0))?
1909 .collect::<rusqlite::Result<Vec<_>>>()?;
1910 let mut hosts: Vec<String> = rows
1911 .iter()
1912 .filter_map(|u| {
1913 reqwest::Url::parse(u)
1914 .ok()
1915 .and_then(|p| p.host_str().map(str::to_string))
1916 })
1917 .collect();
1918 hosts.sort();
1919 hosts.dedup();
1920 Ok(hosts)
1921}
1922
1923pub fn is_fresh(fetched_at: &str, now: chrono::DateTime<Utc>) -> bool {
1927 chrono::DateTime::parse_from_rfc3339(fetched_at)
1928 .is_ok_and(|dt| now.signed_duration_since(dt) < chrono::Duration::hours(24))
1929}
1930
1931pub fn cache_get(conn: &Connection, url_norm: &str) -> Result<Option<(String, String, String)>> {
1938 let row = conn.query_row(
1939 "SELECT COALESCE(title, ''), text, fetched_at FROM web_cache WHERE url_norm = ?1",
1940 [url_norm],
1941 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
1942 );
1943 match row {
1944 Ok(v) => Ok(Some(v)),
1945 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
1946 Err(e) => Err(e.into()),
1947 }
1948}
1949
1950pub fn cache_put(
1952 conn: &Connection,
1953 url_norm: &str,
1954 url: &str,
1955 title: Option<&str>,
1956 text: &str,
1957) -> Result<()> {
1958 let now = Utc::now().to_rfc3339();
1959 conn.execute(
1960 "INSERT INTO web_cache (url_norm, url, title, text, fetched_at) VALUES (?1, ?2, ?3, ?4, ?5)
1961 ON CONFLICT(url_norm) DO UPDATE SET url = ?2, title = ?3, text = ?4, fetched_at = ?5",
1962 (url_norm, url, title, text, &now),
1963 )?;
1964 Ok(())
1965}
1966
1967pub fn files_missing_embeddings(conn: &Connection, space_id: &str) -> Result<Vec<String>> {
1971 let mut stmt = conn.prepare(
1972 "SELECT files.id FROM files
1973 WHERE files.space_id = ?1
1974 AND (SELECT COUNT(*) FROM cache.file_chunks WHERE cache.file_chunks.file_id = files.id) >
1975 (SELECT COUNT(*) FROM cache.chunk_embeddings WHERE cache.chunk_embeddings.file_id = files.id)",
1976 )?;
1977 let rows = stmt.query_map([space_id], |r| r.get(0))?;
1978 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
1979}
1980
1981pub fn semantic_chunks(
1986 conn: &Connection,
1987 space_id: &str,
1988 query: &[f32],
1989 limit: usize,
1990) -> Result<Vec<(String, String, String, f32)>> {
1991 let mut stmt = conn.prepare(
1992 "SELECT files.name, cache.file_chunks.location, cache.file_chunks.text,
1993 cache.chunk_embeddings.vec
1994 FROM cache.chunk_embeddings
1995 JOIN files ON files.id = cache.chunk_embeddings.file_id
1996 JOIN cache.file_chunks
1997 ON cache.file_chunks.file_id = cache.chunk_embeddings.file_id
1998 AND CAST(cache.file_chunks.seq AS INTEGER) = cache.chunk_embeddings.seq
1999 WHERE files.space_id = ?1",
2000 )?;
2001 let rows = stmt.query_map([space_id], |r| {
2002 Ok((
2003 r.get::<_, String>(0)?,
2004 r.get::<_, String>(1)?,
2005 r.get::<_, String>(2)?,
2006 r.get::<_, Vec<u8>>(3)?,
2007 ))
2008 })?;
2009 let mut hits: Vec<(String, String, String, f32)> = Vec::new();
2010 for row in rows {
2011 let (name, loc, text, blob) = row?;
2012 let v = blob_to_vec(&blob);
2013 if v.len() != query.len() {
2014 continue;
2015 }
2016 let score = cosine(query, &v);
2017 hits.push((name, loc, text, score));
2018 }
2019 hits.sort_by(|a, b| b.3.partial_cmp(&a.3).unwrap_or(std::cmp::Ordering::Equal));
2020 hits.truncate(limit);
2021 Ok(hits)
2022}
2023
2024fn cosine(a: &[f32], b: &[f32]) -> f32 {
2025 let (mut dot, mut na, mut nb) = (0.0f32, 0.0f32, 0.0f32);
2026 for (x, y) in a.iter().zip(b) {
2027 dot += x * y;
2028 na += x * x;
2029 nb += y * y;
2030 }
2031 let denom = na.sqrt() * nb.sqrt();
2032 if denom == 0.0 { 0.0 } else { dot / denom }
2033}
2034
2035#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2039pub enum UsageRange {
2040 Day,
2041 Week,
2042 Month,
2043 #[default]
2044 All,
2045}
2046
2047impl UsageRange {
2048 pub const CYCLE: [Self; 4] = [Self::Day, Self::Week, Self::Month, Self::All];
2050
2051 pub const fn label(self) -> &'static str {
2053 match self {
2054 Self::Day => "24h",
2055 Self::Week => "7d",
2056 Self::Month => "30d",
2057 Self::All => "all",
2058 }
2059 }
2060
2061 pub const fn title(self) -> &'static str {
2063 match self {
2064 Self::Day => "last 24 hours",
2065 Self::Week => "last 7 days",
2066 Self::Month => "last 30 days",
2067 Self::All => "all time",
2068 }
2069 }
2070
2071 pub const fn key(self) -> &'static str {
2073 match self {
2074 Self::Day => "day",
2075 Self::Week => "week",
2076 Self::Month => "month",
2077 Self::All => "all",
2078 }
2079 }
2080
2081 pub fn from_key(key: &str) -> Self {
2084 Self::CYCLE
2085 .iter()
2086 .copied()
2087 .find(|r| r.key() == key)
2088 .unwrap_or_default()
2089 }
2090
2091 #[must_use]
2093 pub const fn next(self) -> Self {
2094 match self {
2095 Self::Day => Self::Week,
2096 Self::Week => Self::Month,
2097 Self::Month => Self::All,
2098 Self::All => Self::Day,
2099 }
2100 }
2101
2102 #[must_use]
2104 pub const fn prev(self) -> Self {
2105 match self {
2106 Self::Day => Self::All,
2107 Self::Week => Self::Day,
2108 Self::Month => Self::Week,
2109 Self::All => Self::Month,
2110 }
2111 }
2112
2113 pub fn since(self) -> Option<chrono::DateTime<chrono::Utc>> {
2115 use chrono::{Duration, Utc};
2116 match self {
2117 Self::Day => Some(Utc::now() - Duration::hours(24)),
2118 Self::Week => Some(Utc::now() - Duration::days(7)),
2119 Self::Month => Some(Utc::now() - Duration::days(30)),
2120 Self::All => None,
2121 }
2122 }
2123
2124 pub const fn empty_message(self) -> &'static str {
2126 match self {
2127 Self::Day => "no usage in the last 24 hours — ←/→ for a wider window",
2128 Self::Week => "no usage in the last 7 days — ←/→ for a wider window",
2129 Self::Month => "no usage in the last 30 days — ←/→ for a wider window",
2130 Self::All => "no usage logged yet — send a message first",
2131 }
2132 }
2133}
2134
2135pub fn price_name(model: &str) -> &str {
2140 let stripped = ["go:", "openai:", "codex:", "opencode:"]
2141 .iter()
2142 .find_map(|p| model.strip_prefix(p))
2143 .unwrap_or(model);
2144 stripped.rsplit('/').next().unwrap_or(stripped)
2145}
2146
2147fn catalog_price<'a>(
2152 prices: &'a std::collections::HashMap<String, ModelPricing>,
2153 model: &str,
2154) -> Option<&'a ModelPricing> {
2155 if let Some(price) = prices.get(model) {
2156 return Some(price);
2157 }
2158 let name = price_name(model);
2159 if name.is_empty() {
2160 return None;
2161 }
2162 prices
2163 .iter()
2164 .filter(|(id, _)| {
2165 id.strip_suffix(name)
2166 .is_some_and(|rest| rest.ends_with('/'))
2167 })
2168 .min_by_key(|(id, _)| id.len()) .map(|(_, price)| price)
2170}
2171
2172fn catalog_request_cost(
2175 price: ModelPricing,
2176 prompt_tokens: u64,
2177 completion_tokens: u64,
2178 cache_read_tokens: u64,
2179 cache_creation_tokens: u64,
2180) -> f64 {
2181 let reads = cache_read_tokens.min(prompt_tokens);
2182 let writes = cache_creation_tokens.min(prompt_tokens - reads);
2183 let ordinary = prompt_tokens - reads - writes;
2184 let read_price = price.cache_read.unwrap_or(price.prompt);
2185 let write_price = price.cache_write.unwrap_or(price.prompt);
2186 (ordinary as f64 * price.prompt
2187 + reads as f64 * read_price
2188 + writes as f64 * write_price
2189 + completion_tokens as f64 * price.completion)
2190 / 1e6
2191}
2192
2193struct CostBackfillRow {
2194 id: i64,
2195 backend: String,
2196 model: String,
2197 prompt_tokens: u64,
2198 completion_tokens: u64,
2199 cache_read_tokens: u64,
2200 cache_creation_tokens: u64,
2201 old_cost: Option<f64>,
2202 cost_is_provider: Option<bool>,
2204}
2205
2206#[derive(Default)]
2208pub struct UsageTotals {
2209 pub requests: u64,
2210 pub prompt_tokens: u64,
2211 pub completion_tokens: u64,
2212 pub cache_read_tokens: u64,
2213 pub cache_creation_tokens: u64,
2214 pub cost: f64,
2216}
2217
2218#[derive(Default)]
2220pub struct UsageDay {
2221 pub day: String,
2223 pub requests: u64,
2224 pub prompt_tokens: u64,
2225 pub completion_tokens: u64,
2226 pub cache_read_tokens: u64,
2227 pub cost: f64,
2228}
2229
2230#[derive(Default)]
2232pub struct UsageByBackend {
2233 pub backend: String,
2234 pub requests: u64,
2235 pub prompt_tokens: u64,
2236 pub completion_tokens: u64,
2237 pub cache_read_tokens: u64,
2238 pub cost: f64,
2239}
2240
2241#[derive(Default)]
2243pub struct UsageByModel {
2244 pub model: String,
2245 pub requests: u64,
2246 pub prompt_tokens: u64,
2247 pub completion_tokens: u64,
2248 pub cache_read_tokens: u64,
2249 pub cost: f64,
2250}
2251
2252#[derive(Default)]
2254pub struct UsageRow {
2255 pub created_at: String,
2256 pub backend: String,
2257 pub model: String,
2258 pub prompt_tokens: u64,
2259 pub completion_tokens: u64,
2260 pub cache_read_tokens: u64,
2261 pub cost: Option<f64>,
2262}
2263
2264impl Db {
2265 #[allow(clippy::too_many_arguments)]
2268 pub fn log_usage(
2269 &self,
2270 backend: &str,
2271 model: &str,
2272 prompt_tokens: u64,
2273 completion_tokens: u64,
2274 cache_read_tokens: u64,
2275 cache_creation_tokens: u64,
2276 cost: Option<f64>,
2277 cost_is_provider: bool,
2278 session_id: Option<&str>,
2279 space_id: Option<&str>,
2280 ) -> Result<i64> {
2281 self.conn.execute(
2282 "INSERT INTO usage_log (sync_id, created_at, session_id, space_id, backend, model,
2283 prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens,
2284 cost, cost_is_provider, updated_at)
2285 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
2286 (
2287 Uuid::new_v4().to_string(),
2288 Utc::now().to_rfc3339(),
2289 session_id,
2290 space_id,
2291 backend,
2292 model,
2293 prompt_tokens as i64,
2294 completion_tokens as i64,
2295 cache_read_tokens as i64,
2296 cache_creation_tokens as i64,
2297 cost,
2298 i64::from(cost_is_provider),
2299 Utc::now().to_rfc3339(),
2300 ),
2301 )?;
2302 Ok(self.conn.last_insert_rowid())
2303 }
2304
2305 #[allow(clippy::too_many_arguments)]
2310 pub fn update_usage(
2311 &self,
2312 row_id: i64,
2313 prompt_tokens: u64,
2314 completion_tokens: u64,
2315 cache_read_tokens: u64,
2316 cache_creation_tokens: u64,
2317 cost: Option<f64>,
2318 cost_is_provider: bool,
2319 ) -> Result<()> {
2320 self.conn.execute(
2321 "UPDATE usage_log SET prompt_tokens = ?1, completion_tokens = ?2,
2322 cache_read_tokens = ?3, cache_creation_tokens = ?4, cost = ?5,
2323 cost_is_provider = ?6, updated_at = ?7
2324 WHERE id = ?8",
2325 (
2326 prompt_tokens as i64,
2327 completion_tokens as i64,
2328 cache_read_tokens as i64,
2329 cache_creation_tokens as i64,
2330 cost,
2331 i64::from(cost_is_provider),
2332 Utc::now().to_rfc3339(),
2333 row_id,
2334 ),
2335 )?;
2336 Ok(())
2337 }
2338
2339 pub fn request_cost(
2344 &self,
2345 model: &str,
2346 prompt_tokens: u64,
2347 completion_tokens: u64,
2348 cache_read_tokens: u64,
2349 cache_creation_tokens: u64,
2350 ) -> Option<f64> {
2351 self.model_price(model).map(|price| {
2352 catalog_request_cost(
2353 price,
2354 prompt_tokens,
2355 completion_tokens,
2356 cache_read_tokens,
2357 cache_creation_tokens,
2358 )
2359 })
2360 }
2361
2362 pub fn backfill_usage_costs(&mut self) -> Result<usize> {
2374 let max_price: f64 = self.conn.query_row(
2378 "SELECT COALESCE(MAX(prompt_price), 0) FROM cache.model_prices",
2379 [],
2380 |r| r.get(0),
2381 )?;
2382 if max_price > 0.0 && max_price < 0.001 {
2383 self.conn.execute(
2384 "UPDATE cache.model_prices SET
2385 prompt_price = prompt_price * 1e6,
2386 completion_price = completion_price * 1e6,
2387 cache_read_price = cache_read_price * 1e6,
2388 cache_write_price = cache_write_price * 1e6",
2389 [],
2390 )?;
2391 }
2392 let mut prices: std::collections::HashMap<String, ModelPricing> =
2395 std::collections::HashMap::default();
2396 {
2397 let mut stmt = self.conn.prepare(
2398 "SELECT model_id, prompt_price, completion_price,
2399 cache_read_price, cache_write_price
2400 FROM cache.model_prices",
2401 )?;
2402 let rows = stmt.query_map([], |r| {
2403 Ok((
2404 r.get::<_, String>(0)?,
2405 ModelPricing {
2406 prompt: r.get(1)?,
2407 completion: r.get(2)?,
2408 cache_read: r.get(3)?,
2409 cache_write: r.get(4)?,
2410 },
2411 ))
2412 })?;
2413 for row in rows {
2414 let (model, price) = row?;
2415 prices.insert(model, price);
2416 }
2417 }
2418 let rows: Vec<CostBackfillRow> = {
2419 let mut stmt = self.conn.prepare(
2420 "SELECT id, backend, model, prompt_tokens, completion_tokens,
2421 cache_read_tokens, cache_creation_tokens, cost,
2422 cost_is_provider
2423 FROM usage_log",
2424 )?;
2425 let rows = stmt.query_map([], |r| {
2426 Ok(CostBackfillRow {
2427 id: r.get(0)?,
2428 backend: r.get(1)?,
2429 model: r.get(2)?,
2430 prompt_tokens: r.get::<_, i64>(3)? as u64,
2431 completion_tokens: r.get::<_, i64>(4)? as u64,
2432 cache_read_tokens: r.get::<_, i64>(5)? as u64,
2433 cache_creation_tokens: r.get::<_, i64>(6)? as u64,
2434 old_cost: r.get(7)?,
2435 cost_is_provider: r.get::<_, Option<i64>>(8)?.map(|v| v != 0),
2436 })
2437 })?;
2438 rows.collect::<rusqlite::Result<Vec<_>>>()?
2439 };
2440 let tx = self.conn.transaction()?;
2441 {
2442 let mut update = tx.prepare(
2443 "UPDATE usage_log SET cost = ?2, cost_is_provider = 0, updated_at = ?3 WHERE id = ?1",
2444 )?;
2445 for row in &rows {
2446 let legacy_opencode_cost = row.cost_is_provider.is_none()
2450 && row.backend == "OpenCode Go"
2451 && row.old_cost.is_some();
2452 if row.cost_is_provider == Some(true) || legacy_opencode_cost {
2453 continue;
2454 }
2455 let recomputed = catalog_price(&prices, &row.model).map(|price| {
2456 catalog_request_cost(
2457 *price,
2458 row.prompt_tokens,
2459 row.completion_tokens,
2460 row.cache_read_tokens,
2461 row.cache_creation_tokens,
2462 )
2463 });
2464 let write = match (recomputed, row.old_cost) {
2467 (Some(cost), None) => Some(cost),
2468 (Some(cost), Some(old)) if (cost - old).abs() > 1e-12 => Some(cost),
2469 _ => None,
2470 };
2471 if let Some(cost) = write {
2472 update.execute((row.id, cost, Utc::now().to_rfc3339()))?;
2473 }
2474 }
2475 }
2476 tx.commit()?;
2477 Ok(rows.len())
2478 }
2479
2480 pub fn upsert_model_prices(&mut self, prices: &[(String, String, ModelPricing)]) -> Result<()> {
2485 if prices.is_empty() {
2486 return Ok(());
2487 }
2488 let tx = self.conn.transaction()?;
2489 {
2490 let mut stmt = tx.prepare(
2491 "INSERT INTO cache.model_prices (model_id, backend, prompt_price, completion_price,
2492 cache_read_price, cache_write_price, updated_at)
2493 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
2494 ON CONFLICT(model_id) DO UPDATE SET
2495 prompt_price = excluded.prompt_price,
2496 completion_price = excluded.completion_price,
2497 cache_read_price = excluded.cache_read_price,
2498 cache_write_price = excluded.cache_write_price,
2499 updated_at = excluded.updated_at
2500 WHERE cache.model_prices.prompt_price != excluded.prompt_price
2501 OR cache.model_prices.completion_price != excluded.completion_price
2502 OR cache.model_prices.cache_read_price IS NOT excluded.cache_read_price
2503 OR cache.model_prices.cache_write_price IS NOT excluded.cache_write_price",
2504 )?;
2505 let now = Utc::now().to_rfc3339();
2506 for (model, backend, price) in prices {
2507 stmt.execute((
2508 model,
2509 backend,
2510 price.prompt,
2511 price.completion,
2512 price.cache_read,
2513 price.cache_write,
2514 &now,
2515 ))?;
2516 }
2517 }
2518 tx.commit()?;
2519 Ok(())
2520 }
2521
2522 pub fn model_price(&self, model: &str) -> Option<ModelPricing> {
2529 let read_price = |r: &rusqlite::Row| {
2530 Ok(ModelPricing {
2531 prompt: r.get(0)?,
2532 completion: r.get(1)?,
2533 cache_read: r.get(2)?,
2534 cache_write: r.get(3)?,
2535 })
2536 };
2537 if let Ok(price) = self.conn.query_row(
2538 "SELECT prompt_price, completion_price, cache_read_price, cache_write_price
2539 FROM cache.model_prices WHERE model_id = ?1",
2540 [model],
2541 read_price,
2542 ) {
2543 return Some(price);
2544 }
2545 let name = price_name(model);
2546 if name.is_empty() {
2547 return None;
2548 }
2549 self.conn
2552 .query_row(
2553 "SELECT prompt_price, completion_price, cache_read_price, cache_write_price
2554 FROM cache.model_prices
2555 WHERE backend = 'OpenRouter'
2556 AND substr(model_id, -length(?1) - 1) = '/' || ?1
2557 ORDER BY length(model_id) LIMIT 1",
2558 [name],
2559 read_price,
2560 )
2561 .ok()
2562 }
2563
2564 pub fn usage_totals(&self, since: Option<&str>) -> Result<UsageTotals> {
2569 let mut sql = String::from(
2570 "SELECT COUNT(*),
2571 COALESCE(SUM(prompt_tokens), 0),
2572 COALESCE(SUM(completion_tokens), 0),
2573 COALESCE(SUM(cache_read_tokens), 0),
2574 COALESCE(SUM(cache_creation_tokens), 0),
2575 COALESCE(SUM(cost), 0)
2576 FROM usage_log",
2577 );
2578 if since.is_some() {
2579 sql.push_str(" WHERE created_at >= ?1");
2580 }
2581 let map = |r: &rusqlite::Row| {
2582 Ok(UsageTotals {
2583 requests: r.get::<_, i64>(0)? as u64,
2584 prompt_tokens: r.get::<_, i64>(1)? as u64,
2585 completion_tokens: r.get::<_, i64>(2)? as u64,
2586 cache_read_tokens: r.get::<_, i64>(3)? as u64,
2587 cache_creation_tokens: r.get::<_, i64>(4)? as u64,
2588 cost: r.get::<_, f64>(5)?,
2589 })
2590 };
2591 let totals = match since {
2592 Some(s) => self.conn.query_row(&sql, [s], map),
2593 None => self.conn.query_row(&sql, [], map),
2594 }?;
2595 Ok(totals)
2596 }
2597
2598 pub fn usage_by_backend(&self, since: Option<&str>) -> Result<Vec<UsageByBackend>> {
2601 let mut sql = String::from(
2602 "SELECT backend, COUNT(*),
2603 COALESCE(SUM(prompt_tokens), 0),
2604 COALESCE(SUM(completion_tokens), 0),
2605 COALESCE(SUM(cache_read_tokens), 0),
2606 COALESCE(SUM(cost), 0)
2607 FROM usage_log",
2608 );
2609 if since.is_some() {
2610 sql.push_str(" WHERE created_at >= ?1");
2611 }
2612 sql.push_str(" GROUP BY backend ORDER BY COUNT(*) DESC");
2613 let map = |r: &rusqlite::Row| {
2614 Ok(UsageByBackend {
2615 backend: r.get(0)?,
2616 requests: r.get::<_, i64>(1)? as u64,
2617 prompt_tokens: r.get::<_, i64>(2)? as u64,
2618 completion_tokens: r.get::<_, i64>(3)? as u64,
2619 cache_read_tokens: r.get::<_, i64>(4)? as u64,
2620 cost: r.get::<_, f64>(5)?,
2621 })
2622 };
2623 let mut stmt = self.conn.prepare(&sql)?;
2624 let rows = match since {
2625 Some(s) => stmt.query_map([s], map),
2626 None => stmt.query_map([], map),
2627 }?;
2628 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2629 }
2630
2631 pub fn usage_by_model(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageByModel>> {
2634 let mut sql = String::from(
2635 "SELECT model, COUNT(*),
2636 COALESCE(SUM(prompt_tokens), 0),
2637 COALESCE(SUM(completion_tokens), 0),
2638 COALESCE(SUM(cache_read_tokens), 0),
2639 COALESCE(SUM(cost), 0)
2640 FROM usage_log",
2641 );
2642 if since.is_some() {
2643 sql.push_str(" WHERE created_at >= ?1");
2644 sql.push_str(" GROUP BY model ORDER BY COUNT(*) DESC LIMIT ?2");
2645 } else {
2646 sql.push_str(" GROUP BY model ORDER BY COUNT(*) DESC LIMIT ?1");
2647 }
2648 let map = |r: &rusqlite::Row| {
2649 Ok(UsageByModel {
2650 model: r.get(0)?,
2651 requests: r.get::<_, i64>(1)? as u64,
2652 prompt_tokens: r.get::<_, i64>(2)? as u64,
2653 completion_tokens: r.get::<_, i64>(3)? as u64,
2654 cache_read_tokens: r.get::<_, i64>(4)? as u64,
2655 cost: r.get::<_, f64>(5)?,
2656 })
2657 };
2658 let mut stmt = self.conn.prepare(&sql)?;
2659 let rows = match since {
2660 Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
2661 None => stmt.query_map(rusqlite::params![limit as i64], map),
2662 }?;
2663 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2664 }
2665
2666 pub fn usage_recent(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageRow>> {
2669 let mut sql = String::from(
2670 "SELECT created_at, backend, model, prompt_tokens, completion_tokens,
2671 cache_read_tokens, cost
2672 FROM usage_log",
2673 );
2674 if since.is_some() {
2675 sql.push_str(" WHERE created_at >= ?1");
2676 sql.push_str(" ORDER BY id DESC LIMIT ?2");
2677 } else {
2678 sql.push_str(" ORDER BY id DESC LIMIT ?1");
2679 }
2680 let map = |r: &rusqlite::Row| {
2681 Ok(UsageRow {
2682 created_at: r.get(0)?,
2683 backend: r.get(1)?,
2684 model: r.get(2)?,
2685 prompt_tokens: r.get::<_, i64>(3)? as u64,
2686 completion_tokens: r.get::<_, i64>(4)? as u64,
2687 cache_read_tokens: r.get::<_, i64>(5)? as u64,
2688 cost: r.get(6)?,
2689 })
2690 };
2691 let mut stmt = self.conn.prepare(&sql)?;
2692 let rows = match since {
2693 Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
2694 None => stmt.query_map(rusqlite::params![limit as i64], map),
2695 }?;
2696 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2697 }
2698
2699 pub fn usage_by_day(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageDay>> {
2702 let mut sql = String::from(
2703 "SELECT substr(created_at, 1, 10) AS day, COUNT(*),
2704 COALESCE(SUM(prompt_tokens), 0),
2705 COALESCE(SUM(completion_tokens), 0),
2706 COALESCE(SUM(cache_read_tokens), 0),
2707 COALESCE(SUM(cost), 0)
2708 FROM usage_log",
2709 );
2710 if since.is_some() {
2711 sql.push_str(" WHERE created_at >= ?1");
2712 sql.push_str(" GROUP BY day ORDER BY day DESC LIMIT ?2");
2713 } else {
2714 sql.push_str(" GROUP BY day ORDER BY day DESC LIMIT ?1");
2715 }
2716 let map = |r: &rusqlite::Row| {
2717 Ok(UsageDay {
2718 day: r.get(0)?,
2719 requests: r.get::<_, i64>(1)? as u64,
2720 prompt_tokens: r.get::<_, i64>(2)? as u64,
2721 completion_tokens: r.get::<_, i64>(3)? as u64,
2722 cache_read_tokens: r.get::<_, i64>(4)? as u64,
2723 cost: r.get(5)?,
2724 })
2725 };
2726 let mut stmt = self.conn.prepare(&sql)?;
2727 let rows = match since {
2728 Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
2729 None => stmt.query_map(rusqlite::params![limit as i64], map),
2730 }?;
2731 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2732 }
2733}
2734
2735#[allow(dead_code)]
2743#[derive(Debug, Clone, PartialEq, Eq)]
2744pub struct SyncState {
2745 pub peer_id: String,
2746 pub table_name: String,
2747 pub pull_cursor: Option<String>,
2748 pub push_cursor: Option<String>,
2749 pub last_synced_at: Option<String>,
2750}
2751
2752#[allow(dead_code)]
2754impl Db {
2755 pub fn device_id(&self) -> Result<String> {
2759 if let Some(id) = self
2760 .conn
2761 .query_row("SELECT device_id FROM device_meta LIMIT 1", [], |r| {
2762 r.get(0)
2763 })
2764 .optional()?
2765 {
2766 return Ok(id);
2767 }
2768 let id = Uuid::new_v4().to_string();
2769 let now = Utc::now().to_rfc3339();
2770 self.conn.execute(
2771 "INSERT INTO device_meta (device_id, created_at) VALUES (?1, ?2)",
2772 (&id, &now),
2773 )?;
2774 Ok(id)
2775 }
2776
2777 pub fn set_sync_state(
2780 &self,
2781 peer_id: &str,
2782 table_name: &str,
2783 pull_cursor: Option<&str>,
2784 push_cursor: Option<&str>,
2785 ) -> Result<()> {
2786 self.conn.execute(
2787 "INSERT INTO sync_state (peer_id, table_name, pull_cursor, push_cursor, last_synced_at)
2788 VALUES (?1, ?2, ?3, ?4, ?5)
2789 ON CONFLICT(peer_id, table_name) DO UPDATE SET
2790 pull_cursor = COALESCE(?3, pull_cursor),
2791 push_cursor = COALESCE(?4, push_cursor),
2792 last_synced_at = ?5",
2793 (
2794 peer_id,
2795 table_name,
2796 pull_cursor,
2797 push_cursor,
2798 Utc::now().to_rfc3339(),
2799 ),
2800 )?;
2801 Ok(())
2802 }
2803
2804 pub fn load_sync_state(&self) -> Result<Vec<SyncState>> {
2805 let mut stmt = self.conn.prepare(
2806 "SELECT peer_id, table_name, pull_cursor, push_cursor, last_synced_at
2807 FROM sync_state ORDER BY peer_id, table_name",
2808 )?;
2809 let rows = stmt.query_map([], |r| {
2810 Ok(SyncState {
2811 peer_id: r.get(0)?,
2812 table_name: r.get(1)?,
2813 pull_cursor: r.get(2)?,
2814 push_cursor: r.get(3)?,
2815 last_synced_at: r.get(4)?,
2816 })
2817 })?;
2818 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2819 }
2820}
2821
2822pub fn fts_quote(query: &str) -> String {
2826 query
2827 .split_whitespace()
2828 .map(|t| format!("\"{}\"", t.replace('"', "\"\"")))
2829 .collect::<Vec<_>>()
2830 .join(" ")
2831}
2832
2833pub fn search_chunks(
2835 conn: &Connection,
2836 space_id: &str,
2837 query: &str,
2838 limit: usize,
2839) -> Result<Vec<(String, String, String)>> {
2840 let q = fts_quote(query);
2841 if q.is_empty() {
2842 return Ok(Vec::new());
2843 }
2844 let mut stmt = conn.prepare(
2845 "SELECT files.name, cache.file_chunks.location,
2846 snippet(file_chunks, 3, '', '', '…', 24)
2847 FROM cache.file_chunks JOIN files ON files.id = cache.file_chunks.file_id
2848 WHERE file_chunks MATCH ?1 AND files.space_id = ?2
2849 ORDER BY bm25(file_chunks) LIMIT ?3",
2850 )?;
2851 let rows = stmt.query_map((q, space_id, limit as i64), |r| {
2852 Ok((r.get(0)?, r.get(1)?, r.get(2)?))
2853 })?;
2854 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
2855}
2856
2857pub fn file_text(conn: &Connection, space_id: &str, name: &str) -> Result<Option<String>> {
2859 let mut stmt = conn.prepare(
2860 "SELECT cache.file_chunks.text
2861 FROM cache.file_chunks JOIN files ON files.id = cache.file_chunks.file_id
2862 WHERE files.space_id = ?1 AND files.name = ?2
2863 ORDER BY CAST(cache.file_chunks.seq AS INTEGER) ASC",
2864 )?;
2865 let rows = stmt.query_map((space_id, name), |r| r.get::<_, String>(0))?;
2866 let parts = rows.collect::<rusqlite::Result<Vec<_>>>()?;
2867 Ok((!parts.is_empty()).then(|| parts.join("\n")))
2868}
2869
2870pub fn count_files(conn: &Connection, space_id: &str) -> Result<u64> {
2871 let n: i64 = conn.query_row(
2872 "SELECT COUNT(*) FROM files WHERE space_id = ?1",
2873 [space_id],
2874 |r| r.get(0),
2875 )?;
2876 Ok(n as u64)
2877}
2878
2879#[cfg(test)]
2880mod tests {
2881 use super::*;
2882
2883 fn price(prompt: f64, completion: f64) -> ModelPricing {
2884 ModelPricing {
2885 prompt,
2886 completion,
2887 cache_read: None,
2888 cache_write: None,
2889 }
2890 }
2891
2892 fn cache_price(
2893 prompt: f64,
2894 completion: f64,
2895 cache_read: f64,
2896 cache_write: f64,
2897 ) -> ModelPricing {
2898 ModelPricing {
2899 prompt,
2900 completion,
2901 cache_read: Some(cache_read),
2902 cache_write: Some(cache_write),
2903 }
2904 }
2905
2906 #[test]
2907 fn is_fresh_true_under_24h_false_over() {
2908 let now = Utc::now();
2909 let recent = (now - chrono::Duration::hours(1)).to_rfc3339();
2910 let stale = (now - chrono::Duration::hours(25)).to_rfc3339();
2911 assert!(is_fresh(&recent, now));
2912 assert!(!is_fresh(&stale, now));
2913 assert!(!is_fresh("not a timestamp", now)); }
2915
2916 #[test]
2917 fn usage_log_round_trips_and_aggregates() {
2918 let mut db = Db::open_in_memory().unwrap();
2919 db.upsert_model_prices(&[(
2920 "anthropic/claude-3.5-sonnet".to_string(),
2921 "OpenRouter".to_string(),
2922 price(3.0, 15.0),
2923 )])
2924 .unwrap();
2925 assert_eq!(
2926 db.model_price("anthropic/claude-3.5-sonnet"),
2927 Some(price(3.0, 15.0))
2928 );
2929 db.upsert_model_prices(&[(
2931 "anthropic/claude-3.5-sonnet".to_string(),
2932 "OpenRouter".to_string(),
2933 cache_price(4.0, 16.0, 0.4, 5.0),
2934 )])
2935 .unwrap();
2936 assert_eq!(
2937 db.model_price("anthropic/claude-3.5-sonnet"),
2938 Some(cache_price(4.0, 16.0, 0.4, 5.0))
2939 );
2940 assert_eq!(db.model_price("unknown/model"), None);
2941
2942 db.log_usage(
2944 "OpenRouter",
2945 "anthropic/claude-3.5-sonnet",
2946 100,
2947 10,
2948 70,
2949 20,
2950 Some(0.00056),
2951 true,
2952 Some("s1"),
2953 Some("space-a"),
2954 )
2955 .unwrap();
2956 db.log_usage(
2957 "Codex",
2958 "gpt-5.1-codex",
2959 50,
2960 5,
2961 0,
2962 0,
2963 None,
2964 false,
2965 None,
2966 None,
2967 )
2968 .unwrap();
2969
2970 let totals = db.usage_totals(None).unwrap();
2971 assert_eq!(totals.requests, 2);
2972 assert_eq!(totals.prompt_tokens, 150);
2973 assert_eq!(totals.completion_tokens, 15);
2974 assert_eq!(totals.cache_read_tokens, 70);
2975 assert_eq!(totals.cache_creation_tokens, 20);
2976 assert!((totals.cost - 0.00056).abs() < 1e-9);
2977
2978 let by_backend = db.usage_by_backend(None).unwrap();
2979 assert_eq!(by_backend.len(), 2);
2980 assert_eq!(by_backend[0].backend, "OpenRouter"); assert_eq!(by_backend[0].requests, 1);
2982
2983 let by_model = db.usage_by_model(5, None).unwrap();
2984 assert_eq!(by_model.len(), 2);
2985 assert!(by_model.iter().any(|m| m.model == "gpt-5.1-codex"));
2986
2987 let recent = db.usage_recent(10, None).unwrap();
2988 assert_eq!(recent.len(), 2);
2989 assert_eq!(recent[0].model, "gpt-5.1-codex"); assert_eq!(recent[0].cost, None);
2991 assert_eq!(recent[1].cache_read_tokens, 70);
2992 }
2993
2994 #[test]
2995 fn usage_queries_filter_by_since_window() {
2996 let db = Db::open_in_memory().unwrap();
2997 let insert = |created: &str| {
2999 db.raw().execute(
3000 "INSERT INTO usage_log (sync_id, created_at, session_id, backend, model,
3001 prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens, cost)
3002 VALUES (?1, ?2, NULL, 'OpenRouter', 'a/model', 100, 10, 0, 0, 0.001)",
3003 (uuid::Uuid::new_v4().to_string(), created),
3004 )
3005 };
3006 insert("2026-01-01T00:00:00+00:00").unwrap();
3007 insert("2026-01-02T00:00:00+00:00").unwrap();
3008 insert("2026-01-03T00:00:00+00:00").unwrap();
3009
3010 let since = Some("2026-01-02T00:00:00+00:00");
3011 let totals = db.usage_totals(since).unwrap();
3012 assert_eq!(totals.requests, 2);
3013 assert_eq!(totals.prompt_tokens, 200);
3014 assert!((totals.cost - 0.002).abs() < 1e-12);
3015 assert_eq!(db.usage_totals(None).unwrap().requests, 3);
3016 assert_eq!(db.usage_by_backend(since).unwrap()[0].requests, 2);
3017 assert_eq!(db.usage_by_model(5, since).unwrap()[0].requests, 2);
3018 let recent = db.usage_recent(10, since).unwrap();
3019 assert_eq!(recent.len(), 2);
3020 assert_eq!(recent[0].created_at, "2026-01-03T00:00:00+00:00");
3021 assert!(
3023 recent
3024 .iter()
3025 .any(|r| r.created_at == "2026-01-02T00:00:00+00:00")
3026 );
3027 }
3028
3029 #[test]
3030 fn usage_range_cycles_and_persists() {
3031 use crate::db::UsageRange;
3032 assert_eq!(UsageRange::Day.next(), UsageRange::Week);
3033 assert_eq!(UsageRange::All.next(), UsageRange::Day);
3034 assert_eq!(UsageRange::Day.prev(), UsageRange::All);
3035 assert_eq!(UsageRange::from_key("month"), UsageRange::Month);
3036 assert_eq!(UsageRange::from_key("bogus"), UsageRange::All);
3037 assert_eq!(UsageRange::Week.key(), "week");
3038 assert_eq!(UsageRange::Day.title(), "last 24 hours");
3039 assert!(UsageRange::All.since().is_none());
3040 assert!(UsageRange::Day.since().is_some());
3041 }
3042
3043 #[test]
3044 fn backfill_usage_costs_recomputes_history_from_catalog() {
3045 let mut db = Db::open_in_memory().unwrap();
3046 db.upsert_model_prices(&[(
3047 "anthropic/claude-3.5-sonnet".to_string(),
3048 "OpenRouter".to_string(),
3049 cache_price(3.0, 15.0, 0.3, 3.75),
3050 )])
3051 .unwrap();
3052 db.log_usage(
3055 "OpenRouter",
3056 "anthropic/claude-3.5-sonnet",
3057 100,
3058 10,
3059 70,
3060 20,
3061 None,
3062 false,
3063 None,
3064 None,
3065 )
3066 .unwrap();
3067 db.log_usage(
3068 "Codex",
3069 "gpt-5.1-codex",
3070 50,
3071 5,
3072 0,
3073 0,
3074 None,
3075 false,
3076 None,
3077 None,
3078 )
3079 .unwrap();
3080
3081 let visited = db.backfill_usage_costs().unwrap();
3082 assert_eq!(visited, 2);
3083 let totals = db.usage_totals(None).unwrap();
3086 assert!((totals.cost - 0.000_276).abs() < 1e-12);
3087 let recent = db.usage_recent(10, None).unwrap();
3088 assert!((recent[1].cost.unwrap() - 0.000_276).abs() < 1e-12); assert_eq!(recent[0].cost, None); db.backfill_usage_costs().unwrap();
3093 assert!((db.usage_totals(None).unwrap().cost - 0.000_276).abs() < 1e-12);
3094 }
3095
3096 #[test]
3097 fn backfill_preserves_provider_reported_cost() {
3098 let mut db = Db::open_in_memory().unwrap();
3099 db.upsert_model_prices(&[(
3100 "anthropic/claude-3.5-sonnet".to_string(),
3101 "OpenRouter".to_string(),
3102 cache_price(3.0, 15.0, 0.3, 3.75),
3103 )])
3104 .unwrap();
3105 db.log_usage(
3106 "OpenRouter",
3107 "anthropic/claude-3.5-sonnet",
3108 100,
3109 10,
3110 70,
3111 20,
3112 Some(0.000_321),
3113 true,
3114 None,
3115 None,
3116 )
3117 .unwrap();
3118
3119 db.backfill_usage_costs().unwrap();
3120
3121 assert_eq!(db.usage_recent(1, None).unwrap()[0].cost, Some(0.000_321));
3122 }
3123
3124 #[test]
3125 fn backfill_usage_costs_heals_per_token_catalog() {
3126 let mut db = Db::open_in_memory().unwrap();
3130 db.upsert_model_prices(&[(
3131 "deepseek/deepseek-v4-flash-0731".to_string(),
3132 "OpenRouter".to_string(),
3133 price(8e-08, 1.8e-07),
3134 )])
3135 .unwrap();
3136 db.log_usage(
3137 "OpenRouter",
3138 "deepseek/deepseek-v4-flash-0731",
3139 122_221,
3140 672,
3141 118_784,
3142 0,
3143 Some(9.89864e-09), false,
3145 None,
3146 None,
3147 )
3148 .unwrap();
3149
3150 db.backfill_usage_costs().unwrap();
3151
3152 assert_eq!(
3153 db.model_price("deepseek/deepseek-v4-flash-0731"),
3154 Some(price(0.08, 0.18))
3155 );
3156 let recent = db.usage_recent(10, None).unwrap();
3158 let cost = recent[0].cost.unwrap();
3159 assert!((cost - 0.009_898_6).abs() < 1e-6, "cost was {cost}");
3160 assert!(cost > 0.009, "cost was {cost}");
3161 }
3162
3163 #[test]
3164 fn request_cost_prices_tokens_against_catalog() {
3165 let mut db = Db::open_in_memory().unwrap();
3166 assert_eq!(db.request_cost("unknown/model", 100, 10, 0, 0), None);
3167 db.upsert_model_prices(&[(
3168 "anthropic/claude-3.5-sonnet".to_string(),
3169 "OpenRouter".to_string(),
3170 cache_price(3.0, 15.0, 0.3, 3.75),
3171 )])
3172 .unwrap();
3173 let cost = db
3174 .request_cost("anthropic/claude-3.5-sonnet", 100, 10, 70, 20)
3175 .unwrap();
3176 assert!((cost - 0.000_276).abs() < 1e-12);
3177 }
3178
3179 #[test]
3180 fn model_price_cross_references_openrouter_catalog_twins() {
3181 let mut db = Db::open_in_memory().unwrap();
3182 db.upsert_model_prices(&[
3183 (
3184 "deepseek/deepseek-v4-flash".to_string(),
3185 "OpenRouter".to_string(),
3186 cache_price(0.08, 0.18, 0.016, 0.08),
3187 ),
3188 (
3189 "openai/gpt-5".to_string(),
3190 "OpenRouter".to_string(),
3191 cache_price(1.25, 10.0, 0.125, 1.25),
3192 ),
3193 ])
3194 .unwrap();
3195 assert_eq!(
3198 db.model_price("deepseek/deepseek-v4-flash"),
3199 Some(cache_price(0.08, 0.18, 0.016, 0.08))
3200 );
3201 assert_eq!(
3202 db.model_price("go:deepseek-v4-flash"),
3203 Some(cache_price(0.08, 0.18, 0.016, 0.08))
3204 );
3205 assert_eq!(
3206 db.model_price("deepseek-v4-flash"),
3207 Some(cache_price(0.08, 0.18, 0.016, 0.08))
3208 );
3209 assert_eq!(
3210 db.model_price("openai:gpt-5"),
3211 Some(cache_price(1.25, 10.0, 0.125, 1.25))
3212 );
3213 assert_eq!(
3214 db.model_price("codex:gpt-5"),
3215 Some(cache_price(1.25, 10.0, 0.125, 1.25))
3216 );
3217 assert_eq!(db.model_price("no-such-model-anywhere"), None);
3219 let cost = db
3221 .request_cost("go:deepseek-v4-flash", 100, 10, 70, 0)
3222 .unwrap();
3223 assert!((cost - 0.000_005_32).abs() < 1e-15);
3224 }
3225
3226 #[test]
3227 fn price_name_strips_backend_prefixes_and_vendors() {
3228 assert_eq!(price_name("go:deepseek-v4-flash"), "deepseek-v4-flash");
3229 assert_eq!(price_name("openai:gpt-5"), "gpt-5");
3230 assert_eq!(price_name("codex:gpt-5.1-codex"), "gpt-5.1-codex");
3231 assert_eq!(price_name("opencode:qwen3.6-plus"), "qwen3.6-plus");
3232 assert_eq!(
3233 price_name("deepseek/deepseek-v4-flash"),
3234 "deepseek-v4-flash"
3235 );
3236 assert_eq!(price_name("gpt-5"), "gpt-5");
3237 }
3238
3239 #[test]
3240 fn backfill_prices_non_openrouter_models_via_catalog_twins() {
3241 let mut db = Db::open_in_memory().unwrap();
3242 db.upsert_model_prices(&[(
3243 "deepseek/deepseek-v4-flash".to_string(),
3244 "OpenRouter".to_string(),
3245 cache_price(0.08, 0.18, 0.016, 0.08),
3246 )])
3247 .unwrap();
3248 db.log_usage(
3251 "OpenCode Go",
3252 "go:deepseek-v4-flash",
3253 100,
3254 10,
3255 0,
3256 0,
3257 None,
3258 false,
3259 None,
3260 None,
3261 )
3262 .unwrap();
3263
3264 db.backfill_usage_costs().unwrap();
3265
3266 let recent = db.usage_recent(10, None).unwrap();
3267 let cost = recent[0].cost.unwrap();
3268 assert!((cost - 0.000_009_8).abs() < 1e-15, "cost was {cost}");
3269 assert!((db.usage_totals(None).unwrap().cost - 0.000_009_8).abs() < 1e-15);
3270 }
3271
3272 #[test]
3273 fn web_cache_roundtrips_and_updates_on_rewrite() {
3274 let db = Db::open_in_memory().unwrap();
3275 assert!(cache_get(db.raw(), "example.com/a").unwrap().is_none());
3276 cache_put(
3277 db.raw(),
3278 "example.com/a",
3279 "https://example.com/a",
3280 Some("Title"),
3281 "body text",
3282 )
3283 .unwrap();
3284 let (title, text, fetched_at) = cache_get(db.raw(), "example.com/a").unwrap().unwrap();
3285 assert_eq!(title, "Title");
3286 assert_eq!(text, "body text");
3287 assert!(!fetched_at.is_empty());
3288
3289 cache_put(
3291 db.raw(),
3292 "example.com/a",
3293 "https://example.com/a",
3294 None,
3295 "new body",
3296 )
3297 .unwrap();
3298 let (title, text, _) = cache_get(db.raw(), "example.com/a").unwrap().unwrap();
3299 assert_eq!(title, "");
3300 assert_eq!(text, "new body");
3301 }
3302
3303 #[test]
3304 fn web_mode_defaults_off_and_toggles() {
3305 let db = Db::open_in_memory().unwrap();
3306 let space = db.default_space_id().unwrap();
3307 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3308 assert!(!s.web_mode);
3309 db.set_session_web_mode(&s.id, true).unwrap();
3310 assert!(db.list_sessions(&space).unwrap()[0].web_mode);
3311 }
3312
3313 #[test]
3314 fn swarm_mode_defaults_off_and_toggles() {
3315 let db = Db::open_in_memory().unwrap();
3316 let space = db.default_space_id().unwrap();
3317 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3318 assert!(!s.swarm_mode);
3319 db.set_session_swarm_mode(&s.id, true).unwrap();
3320 assert!(db.list_sessions(&space).unwrap()[0].swarm_mode);
3321 assert!(db.get_session(&s.id).unwrap().unwrap().swarm_mode);
3322 }
3323
3324 #[test]
3325 fn swarm_personas_roundtrip_and_replace_all_on_save() {
3326 let db = Db::open_in_memory().unwrap();
3327 let space = db.default_space_id().unwrap();
3328 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3329 assert!(db.list_swarm_personas(&s.id).unwrap().is_empty());
3330
3331 let roster = vec![
3332 Persona {
3333 name: "Skeptic".into(),
3334 model: "a/one".into(),
3335 blurb: "pokes holes".into(),
3336 },
3337 Persona {
3338 name: "Advocate".into(),
3339 model: "b/two".into(),
3340 blurb: "user-first".into(),
3341 },
3342 ];
3343 db.save_swarm_personas(&s.id, &roster).unwrap();
3344 let loaded = db.list_swarm_personas(&s.id).unwrap();
3345 assert_eq!(loaded.len(), 2);
3346 assert_eq!(loaded[0].name, "Skeptic");
3347 assert_eq!(loaded[1].name, "Advocate");
3348
3349 db.save_swarm_personas(&s.id, &roster[..1]).unwrap();
3351 assert_eq!(db.list_swarm_personas(&s.id).unwrap().len(), 1);
3352 }
3353
3354 #[test]
3355 fn persona_message_tags_role_assistant_with_persona_and_model() {
3356 let db = Db::open_in_memory().unwrap();
3357 let space = db.default_space_id().unwrap();
3358 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3359 db.add_persona_message(&s.id, "reply text", "Skeptic", "a/one")
3360 .unwrap();
3361 let msgs = db.load_messages(&s.id).unwrap();
3362 assert_eq!(msgs.len(), 1);
3363 assert_eq!(msgs[0].role, "assistant");
3364 assert_eq!(msgs[0].persona.as_deref(), Some("Skeptic"));
3365 assert_eq!(msgs[0].model.as_deref(), Some("a/one"));
3366
3367 db.add_assistant_message(&s.id, "final answer", None, None, None, None, None, None)
3369 .unwrap();
3370 let msgs = db.load_messages(&s.id).unwrap();
3371 assert_eq!(msgs[1].persona, None);
3372 }
3373
3374 #[test]
3375 fn session_sources_link_to_the_web_cache_and_are_keyword_searchable() {
3376 let db = Db::open_in_memory().unwrap();
3377 let space = db.default_space_id().unwrap();
3378 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3379 cache_put(
3380 db.raw(),
3381 "https://example.com/a",
3382 "https://example.com/a",
3383 Some("A"),
3384 "rust borrow checker deep dive",
3385 )
3386 .unwrap();
3387 cache_put(
3388 db.raw(),
3389 "https://example.com/b",
3390 "https://example.com/b",
3391 Some("B"),
3392 "cooking pasta recipes",
3393 )
3394 .unwrap();
3395 db.add_session_sources(
3396 &s.id,
3397 &[
3398 "https://example.com/a".to_string(),
3399 "https://example.com/b".to_string(),
3400 ],
3401 )
3402 .unwrap();
3403
3404 let hits = db.search_session_sources(&s.id, "borrow checker").unwrap();
3405 assert_eq!(hits.len(), 1);
3406 assert!(hits[0].1.contains("borrow checker"));
3407
3408 assert!(
3409 db.search_session_sources(&s.id, "quantum")
3410 .unwrap()
3411 .is_empty()
3412 );
3413 }
3414
3415 #[test]
3416 fn set_source_flag_pins_and_discards_then_clears() {
3417 let db = Db::open_in_memory().unwrap();
3418 let session_id = "sess-1";
3419 add_session_sources(&db.conn, session_id, &["https://a.example/x".to_string()]).unwrap();
3420 db.set_source_flag(session_id, "https://a.example/x", Some("pinned"))
3421 .unwrap();
3422 assert_eq!(
3423 pinned_urls(&db.conn, session_id).unwrap(),
3424 vec!["https://a.example/x".to_string()]
3425 );
3426 assert!(discarded_domains(&db.conn, session_id).unwrap().is_empty());
3427
3428 db.set_source_flag(session_id, "https://a.example/x", Some("discarded"))
3429 .unwrap();
3430 assert!(pinned_urls(&db.conn, session_id).unwrap().is_empty());
3431 assert_eq!(
3432 discarded_domains(&db.conn, session_id).unwrap(),
3433 vec!["a.example".to_string()]
3434 );
3435
3436 db.set_source_flag(session_id, "https://a.example/x", None)
3437 .unwrap();
3438 assert!(discarded_domains(&db.conn, session_id).unwrap().is_empty());
3439 }
3440
3441 #[test]
3442 fn upsert_research_stage_message_replaces_the_same_labels_row() {
3443 let db = Db::open_in_memory().unwrap();
3444 let space = db.default_space_id().unwrap();
3445 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3446 db.upsert_research_stage_message(&s.id, "searching", "round 1, 1/3")
3447 .unwrap();
3448 db.upsert_research_stage_message(&s.id, "searching", "round 1, 2/3")
3449 .unwrap();
3450 db.upsert_research_stage_message(&s.id, "planning", "")
3451 .unwrap();
3452
3453 let msgs = db.load_messages(&s.id).unwrap();
3454 let searching: Vec<_> = msgs
3455 .iter()
3456 .filter(|m| m.content.starts_with("searching:"))
3457 .collect();
3458 assert_eq!(searching.len(), 1, "expected one row, updated in place");
3459 assert!(searching[0].content.contains("2/3"));
3460 assert_eq!(msgs.iter().filter(|m| m.content == "planning").count(), 1);
3461 }
3462
3463 #[test]
3464 fn session_and_message_roundtrip() {
3465 let db = Db::open_in_memory().unwrap();
3466 let space = db.default_space_id().unwrap();
3467 let s = db
3468 .create_session("hello", "openai/gpt-4o", &space, "chat")
3469 .unwrap();
3470
3471 db.add_user_message(&s.id, "hi").unwrap();
3472 db.add_assistant_message(
3473 &s.id,
3474 "hello there",
3475 Some("openai/gpt-4o"),
3476 Some("let me think"),
3477 Some(3),
3478 Some(1.5),
3479 Some(0.0042),
3480 Some("Vibed"),
3481 )
3482 .unwrap();
3483
3484 let msgs = db.load_messages(&s.id).unwrap();
3485 assert_eq!(msgs.len(), 2);
3486 assert_eq!(msgs[0].role, "user");
3487 assert_eq!(msgs[1].content, "hello there");
3488 assert_eq!(msgs[1].model.as_deref(), Some("openai/gpt-4o"));
3489 assert_eq!(msgs[1].reasoning.as_deref(), Some("let me think"));
3490 assert_eq!(msgs[1].tokens, Some(3));
3491 assert_eq!(msgs[1].cost, Some(0.0042));
3492
3493 let sessions = db.list_sessions(&space).unwrap();
3494 assert_eq!(sessions.len(), 1);
3495 assert_eq!(sessions[0].id, s.id);
3496 }
3497
3498 #[test]
3499 fn markdown_images_in_content_roundtrip() {
3500 let db = Db::open_in_memory().unwrap();
3501 let space = db.default_space_id().unwrap();
3502 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3503 let content = "look at  and ";
3504 db.add_user_message(&s.id, content).unwrap();
3505 let msgs = db.load_messages(&s.id).unwrap();
3506 assert_eq!(msgs.len(), 1);
3507 assert!(msgs[0].content.contains(""));
3508 assert!(msgs[0].content.contains(""));
3509 }
3510
3511 #[test]
3512 fn model_prefs_toggle_and_used() {
3513 let db = Db::open_in_memory().unwrap();
3514 assert!(db.toggle_favorite("a/one").unwrap()); assert!(!db.toggle_favorite("a/one").unwrap()); db.mark_model_used("a/one").unwrap();
3517 db.set_reasoning("a/one", Some("high")).unwrap();
3518
3519 let prefs = db.load_model_prefs().unwrap();
3520 let p = &prefs[0];
3521 assert_eq!(p.id, "a/one");
3522 assert!(!p.favorite);
3523 assert!(p.last_used.is_some());
3524 assert_eq!(p.reasoning.as_deref(), Some("high"));
3525 }
3526
3527 #[test]
3528 fn settings_roundtrip() {
3529 let db = Db::open_in_memory().unwrap();
3530 db.set_setting("temperature", "0.7").unwrap();
3531 db.set_setting("temperature", "0.9").unwrap(); let s = db.load_settings().unwrap();
3533 assert_eq!(s, vec![("temperature".to_string(), "0.9".to_string())]);
3534 }
3535
3536 #[test]
3537 fn set_model_updates_row() {
3538 let db = Db::open_in_memory().unwrap();
3539 let space = db.default_space_id().unwrap();
3540 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3541 db.set_session_model(&s.id, "c/d").unwrap();
3542 assert_eq!(db.list_sessions(&space).unwrap()[0].model, "c/d");
3543 }
3544
3545 #[test]
3546 fn compaction_persists_and_roundtrips() {
3547 let db = Db::open_in_memory().unwrap();
3548 let space = db.default_space_id().unwrap();
3549 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3550 assert_eq!(s.compact_summary, None);
3551 assert_eq!(s.compact_through, 0);
3552
3553 db.set_compaction(&s.id, "digest of earlier turns", 6)
3554 .unwrap();
3555 let reloaded = &db.list_sessions(&space).unwrap()[0];
3556 assert_eq!(
3557 reloaded.compact_summary.as_deref(),
3558 Some("digest of earlier turns")
3559 );
3560 assert_eq!(reloaded.compact_through, 6);
3561 }
3562
3563 #[test]
3564 fn spaces_crud_and_session_reassignment_on_delete() {
3565 let db = Db::open_in_memory().unwrap();
3566 let spaces = db.list_spaces().unwrap();
3567 assert_eq!(spaces.len(), 1);
3568 assert_eq!(spaces[0].name, DEFAULT_SPACE);
3569
3570 let work = db.create_space("work").unwrap();
3571 let s = db.create_session("hi", "a/b", &work.id, "chat").unwrap();
3572 assert_eq!(db.count_sessions(&work.id).unwrap(), 1);
3573
3574 db.rename_space(&work.id, "work-renamed").unwrap();
3575 assert!(
3576 db.list_spaces()
3577 .unwrap()
3578 .iter()
3579 .any(|s| s.name == "work-renamed")
3580 );
3581
3582 db.delete_space(&work.id).unwrap();
3583 assert_eq!(db.list_spaces().unwrap().len(), 1); let default_id = db.default_space_id().unwrap();
3585 let moved = db.list_sessions(&default_id).unwrap();
3586 assert!(moved.iter().any(|c| c.id == s.id)); }
3588
3589 #[test]
3590 fn chunk_embeddings_store_rank_and_invalidate() {
3591 let db = Db::open_in_memory().unwrap();
3592 let space = db.default_space_id().unwrap();
3593 let id = db.upsert_file(&space, "book.pdf", "h1", 10, "ok").unwrap();
3594 db.set_file_chunks(
3595 &id,
3596 &[
3597 ("page 1".into(), "cooking with fire".into()),
3598 ("page 2".into(), "quantum entanglement".into()),
3599 ],
3600 )
3601 .unwrap();
3602
3603 let v = vec![0.25f32, -1.0, 3.5];
3605 assert_eq!(blob_to_vec(&vec_to_blob(&v)), v);
3606
3607 assert_eq!(
3609 files_missing_embeddings(&db.conn, &space).unwrap(),
3610 vec![id.clone()]
3611 );
3612
3613 db.set_chunk_embeddings(&id, &[(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])])
3614 .unwrap();
3615 assert!(
3616 files_missing_embeddings(&db.conn, &space)
3617 .unwrap()
3618 .is_empty()
3619 );
3620
3621 let hits = semantic_chunks(&db.conn, &space, &[0.1, 0.9], 5).unwrap();
3623 assert_eq!(hits[0].1, "page 2");
3624 assert!(hits[0].2.contains("quantum"));
3625 assert!(hits[0].3 > hits[1].3, "scores must be descending");
3626
3627 let hits = semantic_chunks(&db.conn, &space, &[1.0, 0.0, 0.0], 5).unwrap();
3629 assert!(hits.is_empty());
3630
3631 db.set_file_chunks(&id, &[("page 1".into(), "new text".into())])
3633 .unwrap();
3634 assert_eq!(
3635 files_missing_embeddings(&db.conn, &space).unwrap(),
3636 vec![id.clone()]
3637 );
3638 }
3639
3640 #[test]
3641 fn files_upsert_list_delete_roundtrip() {
3642 let db = Db::open_in_memory().unwrap();
3643 let space = db.default_space_id().unwrap();
3644 let id = db.upsert_file(&space, "notes.md", "h1", 10, "ok").unwrap();
3645 db.set_file_chunks(&id, &[("lines 1-40".into(), "hello fts world".into())])
3646 .unwrap();
3647
3648 let files = db.list_files(&space).unwrap();
3649 assert_eq!(files.len(), 1);
3650 assert_eq!(files[0].name, "notes.md");
3651 assert_eq!(files[0].hash, "h1");
3652 assert_eq!(files[0].status, "ok");
3653
3654 let id2 = db.upsert_file(&space, "notes.md", "h2", 12, "ok").unwrap();
3656 db.set_file_chunks(&id2, &[("lines 1-40".into(), "goodbye".into())])
3657 .unwrap();
3658 let files = db.list_files(&space).unwrap();
3659 assert_eq!(files.len(), 1);
3660 assert_eq!(files[0].hash, "h2");
3661
3662 db.delete_file(&files[0].id).unwrap();
3663 assert!(db.list_files(&space).unwrap().is_empty());
3664 }
3665
3666 #[test]
3667 fn chunk_search_ranks_and_scopes_by_space() {
3668 let db = Db::open_in_memory().unwrap();
3669 let space = db.default_space_id().unwrap();
3670 let other = db.create_space("other").unwrap();
3671 let a = db.upsert_file(&space, "a.md", "h", 1, "ok").unwrap();
3672 let b = db.upsert_file(&other.id, "b.md", "h", 1, "ok").unwrap();
3673 db.set_file_chunks(&a, &[("lines 1-40".into(), "rust borrow checker".into())])
3674 .unwrap();
3675 db.set_file_chunks(&b, &[("lines 1-40".into(), "rust in other space".into())])
3676 .unwrap();
3677
3678 let hits = search_chunks(&db.conn, &space, "rust", 8).unwrap();
3679 assert_eq!(hits.len(), 1); assert_eq!(hits[0].0, "a.md");
3681 assert_eq!(hits[0].1, "lines 1-40");
3682 assert!(hits[0].2.contains("rust"));
3683
3684 assert!(search_chunks(&db.conn, &space, "c++ \"quoted\" -dash", 8).is_ok());
3686 }
3687
3688 #[test]
3689 fn file_text_joins_chunks_in_order() {
3690 let db = Db::open_in_memory().unwrap();
3691 let space = db.default_space_id().unwrap();
3692 let id = db.upsert_file(&space, "doc.txt", "h", 1, "ok").unwrap();
3693 db.set_file_chunks(
3694 &id,
3695 &[
3696 ("lines 1-2".into(), "one\ntwo".into()),
3697 ("lines 3-4".into(), "three\nfour".into()),
3698 ],
3699 )
3700 .unwrap();
3701 let text = file_text(&db.conn, &space, "doc.txt").unwrap().unwrap();
3702 assert_eq!(text, "one\ntwo\nthree\nfour");
3703 assert!(
3704 file_text(&db.conn, &space, "missing.txt")
3705 .unwrap()
3706 .is_none()
3707 );
3708 assert_eq!(count_files(&db.conn, &space).unwrap(), 1);
3709 }
3710
3711 #[test]
3712 fn research_stage_messages_round_trip() {
3713 let db = Db::open_in_memory().unwrap();
3714 let space = db.default_space_id().unwrap();
3715 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3716 db.add_research_stage_message(&s.id, "planning…").unwrap();
3717 let msgs = db.load_messages(&s.id).unwrap();
3718 assert_eq!(msgs.last().unwrap().role, "research_stage");
3719 assert_eq!(msgs.last().unwrap().content, "planning…");
3720 }
3721
3722 #[test]
3723 fn survey_messages_round_trip() {
3724 let db = Db::open_in_memory().unwrap();
3725 let space = db.default_space_id().unwrap();
3726 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3727 db.add_survey_message(&s.id, "For \"topic\":\n 1. Depth or breadth?")
3728 .unwrap();
3729 let msgs = db.load_messages(&s.id).unwrap();
3730 let last = msgs.last().unwrap();
3731 assert_eq!(last.role, "survey");
3732 assert!(last.content.contains("Depth or breadth?"));
3733 }
3734
3735 #[test]
3736 fn gate_reply_round_trip() {
3737 let db = Db::open_in_memory().unwrap();
3738 let space = db.default_space_id().unwrap();
3739 let s = db.create_session("t", "a/b", &space, "chat").unwrap();
3740 db.add_gate_reply_message(&s.id, "the second option")
3741 .unwrap();
3742 let msgs = db.load_messages(&s.id).unwrap();
3743 let last = msgs.last().unwrap();
3744 assert_eq!(last.role, "gate_reply");
3745 assert_eq!(last.content, "the second option");
3746 }
3747
3748 #[test]
3749 fn create_list_touch_delete_watch_roundtrip() {
3750 let db = Db::open_in_memory().unwrap();
3751 let id = db
3752 .create_watch("space-1", "rust async runtimes", 24, "sess-1")
3753 .unwrap();
3754 let watches = db.list_watches("space-1").unwrap();
3755 assert_eq!(watches.len(), 1);
3756 assert_eq!(watches[0].topic, "rust async runtimes");
3757 assert_eq!(watches[0].interval_hours, 24);
3758 assert!(watches[0].last_run_at.is_none());
3759
3760 db.touch_watch(&id, "2026-07-07T00:00:00+00:00").unwrap();
3761 let watches = db.list_watches("space-1").unwrap();
3762 assert_eq!(
3763 watches[0].last_run_at.as_deref(),
3764 Some("2026-07-07T00:00:00+00:00")
3765 );
3766
3767 db.delete_watch(&id).unwrap();
3768 assert!(db.list_watches("space-1").unwrap().is_empty());
3769 }
3770
3771 #[test]
3772 #[allow(clippy::similar_names)]
3774 fn list_all_watches_returns_watches_from_all_spaces() {
3775 let db = Db::open_in_memory().unwrap();
3776
3777 let id1 = db.create_watch("space-a", "topic-1", 24, "sess-1").unwrap();
3779 let id2 = db.create_watch("space-b", "topic-2", 48, "sess-2").unwrap();
3780 let id3 = db.create_watch("space-a", "topic-3", 12, "sess-3").unwrap();
3781
3782 let all_watches = db.list_all_watches().unwrap();
3784 assert_eq!(all_watches.len(), 3);
3785 assert!(
3786 all_watches
3787 .iter()
3788 .any(|w| w.id == id1 && w.space_id == "space-a")
3789 );
3790 assert!(
3791 all_watches
3792 .iter()
3793 .any(|w| w.id == id2 && w.space_id == "space-b")
3794 );
3795 assert!(
3796 all_watches
3797 .iter()
3798 .any(|w| w.id == id3 && w.space_id == "space-a")
3799 );
3800
3801 let space_a_watches = db.list_watches("space-a").unwrap();
3804 assert_eq!(space_a_watches.len(), 2);
3805 assert!(space_a_watches.iter().all(|w| w.space_id == "space-a"));
3806
3807 let space_b_watches = db.list_watches("space-b").unwrap();
3808 assert_eq!(space_b_watches.len(), 1);
3809 assert!(space_b_watches.iter().all(|w| w.space_id == "space-b"));
3810 }
3811
3812 fn legacy_db(path: &std::path::Path, now: &str) {
3816 let conn = rusqlite::Connection::open(path).unwrap();
3817 conn.execute_batch(
3818 "CREATE TABLE sessions (id TEXT PRIMARY KEY, title TEXT NOT NULL,
3819 model TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL);
3820 CREATE TABLE messages (id TEXT PRIMARY KEY, session_id TEXT NOT NULL,
3821 role TEXT NOT NULL, content TEXT NOT NULL, created_at TEXT NOT NULL);
3822 CREATE TABLE model_prefs (id TEXT PRIMARY KEY,
3823 favorite INTEGER NOT NULL DEFAULT 0, last_used TEXT);
3824 CREATE TABLE app_settings (key TEXT PRIMARY KEY, value TEXT NOT NULL);
3825 CREATE TABLE spaces (id TEXT PRIMARY KEY, name TEXT NOT NULL UNIQUE,
3826 created_at TEXT NOT NULL);
3827 CREATE TABLE files (id TEXT PRIMARY KEY, space_id TEXT NOT NULL,
3828 name TEXT NOT NULL, hash TEXT NOT NULL, size INTEGER NOT NULL,
3829 status TEXT NOT NULL, created_at TEXT NOT NULL,
3830 UNIQUE(space_id, name));
3831 ALTER TABLE files ADD COLUMN mtime INTEGER NOT NULL DEFAULT 0;
3832 CREATE VIRTUAL TABLE file_chunks USING fts5(
3833 file_id UNINDEXED, seq UNINDEXED, location UNINDEXED, text);
3834 CREATE TABLE chunk_embeddings (file_id TEXT NOT NULL, seq INTEGER NOT NULL,
3835 vec BLOB NOT NULL, PRIMARY KEY (file_id, seq));
3836 CREATE TABLE web_cache (url_norm TEXT PRIMARY KEY, url TEXT NOT NULL,
3837 title TEXT, text TEXT NOT NULL, fetched_at TEXT NOT NULL);
3838 CREATE TABLE citations (id INTEGER PRIMARY KEY AUTOINCREMENT,
3839 space_id TEXT NOT NULL, report_file TEXT NOT NULL,
3840 url TEXT NOT NULL, title TEXT);
3841 CREATE TABLE session_sources (session_id TEXT NOT NULL,
3842 url_norm TEXT NOT NULL, PRIMARY KEY (session_id, url_norm));
3843 CREATE TABLE watches (id TEXT PRIMARY KEY, space_id TEXT NOT NULL,
3844 topic TEXT NOT NULL, interval_hours INTEGER NOT NULL,
3845 session_id TEXT NOT NULL, last_run_at TEXT);
3846 CREATE TABLE swarm_personas (session_id TEXT NOT NULL, ord INTEGER NOT NULL,
3847 name TEXT NOT NULL, model TEXT NOT NULL, persona TEXT NOT NULL);
3848 CREATE TABLE usage_log (id INTEGER PRIMARY KEY AUTOINCREMENT,
3849 created_at TEXT NOT NULL, session_id TEXT, space_id TEXT,
3850 backend TEXT NOT NULL, model TEXT NOT NULL,
3851 prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL,
3852 cache_read_tokens INTEGER NOT NULL DEFAULT 0,
3853 cache_creation_tokens INTEGER NOT NULL DEFAULT 0, cost REAL);
3854 CREATE TABLE model_prices (model_id TEXT PRIMARY KEY, backend TEXT NOT NULL,
3855 prompt_price REAL NOT NULL, completion_price REAL NOT NULL,
3856 cache_read_price REAL, cache_write_price REAL, updated_at TEXT NOT NULL);",
3857 )
3858 .unwrap();
3859 conn.execute(
3860 "INSERT INTO spaces (id, name, created_at) VALUES ('sp', 'default', ?1)",
3861 [now],
3862 )
3863 .unwrap();
3864 conn.execute(
3865 "INSERT INTO files (id, space_id, name, hash, size, status, created_at, mtime)
3866 VALUES ('f1', 'sp', 'a.txt', 'h1', 10, 'ok', ?1, 1234)",
3867 [now],
3868 )
3869 .unwrap();
3870 conn.execute(
3871 "INSERT INTO file_chunks (file_id, seq, location, text)
3872 VALUES ('f1', 0, 'l', 'hello world')",
3873 [],
3874 )
3875 .unwrap();
3876 conn.execute(
3877 "INSERT INTO chunk_embeddings (file_id, seq, vec) VALUES ('f1', 0, ?1)",
3878 [vec_to_blob(&[1.0, 0.0])],
3879 )
3880 .unwrap();
3881 conn.execute(
3882 "INSERT INTO web_cache (url_norm, url, title, text, fetched_at)
3883 VALUES ('https://x.test/', 'https://x.test/', NULL, 'cached body', ?1)",
3884 [now],
3885 )
3886 .unwrap();
3887 conn.execute(
3888 "INSERT INTO model_prices (model_id, backend, prompt_price, completion_price, updated_at)
3889 VALUES ('a/model', 'OpenRouter', 1.0, 2.0, ?1)",
3890 [now],
3891 )
3892 .unwrap();
3893 conn.execute(
3894 "INSERT INTO usage_log (created_at, backend, model, prompt_tokens, completion_tokens)
3895 VALUES (?1, 'OpenRouter', 'a/model', 100, 10)",
3896 [now],
3897 )
3898 .unwrap();
3899 drop(conn);
3900 }
3901
3902 #[test]
3903 #[allow(clippy::float_cmp)]
3905 fn legacy_db_migrates_cache_tables_and_seeds_file_index_state() {
3906 let dir = std::env::temp_dir().join(format!("nexus-migrate-{}", uuid::Uuid::new_v4()));
3907 std::fs::create_dir_all(&dir).unwrap();
3908 let db_path = dir.join("nexus.db");
3909 legacy_db(&db_path, &Utc::now().to_rfc3339());
3910 let mut db = Db::open(&db_path).unwrap();
3911
3912 let v: i64 = db
3914 .raw()
3915 .query_row("PRAGMA user_version", [], |r| r.get(0))
3916 .unwrap();
3917 assert_eq!(v, SCHEMA_VERSION);
3918 assert!(has_column(db.raw(), "files", "mtime").unwrap());
3919
3920 assert_eq!(db.default_space_id().unwrap(), DEFAULT_SPACE);
3923 let files_space: String = db
3924 .raw()
3925 .query_row("SELECT space_id FROM files WHERE id = 'f1'", [], |r| {
3926 r.get(0)
3927 })
3928 .unwrap();
3929 assert_eq!(files_space, DEFAULT_SPACE);
3930
3931 let (mtime, status) = db
3933 .raw()
3934 .query_row(
3935 "SELECT mtime, status FROM cache.file_index_state WHERE file_id = 'f1'",
3936 [],
3937 |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)),
3938 )
3939 .unwrap();
3940 assert_eq!((mtime, status.as_str()), (1234, "ok"));
3941
3942 assert!(cache_path_for(&db_path).is_file());
3945 let space = db.default_space_id().unwrap();
3946 let hits = search_chunks(db.raw(), &space, "hello", 8).unwrap();
3947 assert_eq!(hits.len(), 1);
3948 assert_eq!(
3949 file_text(db.raw(), &space, "a.txt").unwrap().as_deref(),
3950 Some("hello world")
3951 );
3952 let (_, text, _) = cache_get(db.raw(), "https://x.test/").unwrap().unwrap();
3953 assert_eq!(text, "cached body");
3954 let hits = semantic_chunks(db.raw(), &space, &[1.0, 0.0], 4).unwrap();
3955 assert_eq!(hits.len(), 1);
3956 assert_eq!(hits[0].0, "a.txt");
3957
3958 assert_eq!(db.model_price("a/model").unwrap().prompt, 1.0);
3961 assert_eq!(db.backfill_usage_costs().unwrap(), 1);
3962 let sync: String = db
3963 .raw()
3964 .query_row("SELECT sync_id FROM usage_log WHERE id = 1", [], |r| {
3965 r.get(0)
3966 })
3967 .unwrap();
3968 assert!(!sync.is_empty());
3969 }
3970
3971 #[test]
3972 fn fresh_db_creates_sibling_cache_db_with_split_tables() {
3973 let dir = std::env::temp_dir().join(format!("nexus-fresh-{}", uuid::Uuid::new_v4()));
3974 std::fs::create_dir_all(&dir).unwrap();
3975 let db_path = dir.join("nexus.db");
3976 let db = Db::open(&db_path).unwrap();
3977 let space = db.default_space_id().unwrap();
3978 let id = db.upsert_file(&space, "doc.txt", "h", 1, "ok").unwrap();
3979 db.set_file_chunks(&id, &[("l".into(), "needle text".into())])
3980 .unwrap();
3981
3982 assert!(!has_column(db.raw(), "web_cache", "url_norm").unwrap());
3984 assert!(!has_column(db.raw(), "files", "mtime").unwrap());
3985 let cache_path = cache_path_for(&db_path);
3988 assert!(cache_path.is_file());
3989 let conn = rusqlite::Connection::open(&cache_path).unwrap();
3990 assert!(cache_get(&conn, "x").unwrap().is_none());
3991 drop(conn);
3992 assert_eq!(
3993 file_text(db.raw(), &space, "doc.txt").unwrap().as_deref(),
3994 Some("needle text")
3995 );
3996 assert_eq!(db.list_files(&space).unwrap()[0].status, "ok");
3997 }
3998
3999 #[allow(clippy::too_many_lines)]
4005 #[test]
4006 fn every_mutable_table_bumps_updated_at_on_mutation() {
4007 type Mutator<'a> = &'a dyn Fn(&Db) -> Result<()>;
4009 let mut db = Db::open_in_memory().unwrap();
4010 let space = db.default_space_id().unwrap();
4011 let sid = db.create_session("t", "a/b", &space, "chat").unwrap().id;
4012 let read = |table: &str, id: &str| -> String {
4013 std::thread::sleep(std::time::Duration::from_millis(2));
4014 db.raw()
4015 .query_row(
4016 &format!("SELECT updated_at FROM {table} WHERE id = ?1"),
4017 [id],
4018 |r| r.get::<_, String>(0),
4019 )
4020 .unwrap()
4021 };
4022
4023 let mutators: [Mutator<'_>; 6] = [
4025 &|db: &Db| db.set_compaction(&sid, "sum", 3),
4026 &|db: &Db| db.set_session_web_mode(&sid, true),
4027 &|db: &Db| db.set_session_swarm_mode(&sid, true),
4028 &|db: &Db| db.set_session_title(&sid, "new", Some("new-slug")),
4029 &|db: &Db| db.set_session_model(&sid, "m/x"),
4030 &|db: &Db| db.set_research_parent(&sid, "parent"),
4031 ];
4032 for mutate in mutators {
4033 let before = read("sessions", &sid);
4034 mutate(&db).unwrap();
4035 assert!(read("sessions", &sid) > before);
4036 }
4037 let before = read("sessions", &sid);
4039 db.save_swarm_personas(
4040 &sid,
4041 &[Persona {
4042 name: "a".into(),
4043 model: "m".into(),
4044 blurb: "b".into(),
4045 }],
4046 )
4047 .unwrap();
4048 assert!(read("sessions", &sid) > before);
4049
4050 assert!(db.toggle_favorite("a/model").unwrap());
4052 let before = read("model_prefs", "a/model");
4053 db.set_reasoning("a/model", Some("high")).unwrap();
4054 assert!(read("model_prefs", "a/model") > before);
4055 let before = read("model_prefs", "a/model");
4056 db.mark_model_used("a/model").unwrap();
4057 assert!(read("model_prefs", "a/model") > before);
4058
4059 let read_key = |key: &str| -> String {
4061 std::thread::sleep(std::time::Duration::from_millis(2));
4062 db.raw()
4063 .query_row(
4064 "SELECT updated_at FROM app_settings WHERE key = ?1",
4065 [key],
4066 |r| r.get::<_, String>(0),
4067 )
4068 .unwrap()
4069 };
4070 db.set_setting("temperature", "0.5").unwrap();
4071 let before = read_key("temperature");
4072 db.set_setting("temperature", "0.9").unwrap();
4073 assert!(read_key("temperature") > before);
4074
4075 let sp = db.create_space("other").unwrap();
4077 let before = read("spaces", &sp.id);
4078 db.rename_space(&sp.id, "other2").unwrap();
4079 assert!(read("spaces", &sp.id) > before);
4080
4081 let fid = db.upsert_file(&space, "f.txt", "h", 1, "ok").unwrap();
4083 let before = read("files", &fid);
4084 db.rename_file(&fid, "g.txt").unwrap();
4085 assert!(read("files", &fid) > before);
4086
4087 let wid = db.create_watch(&space, "topic", 24, &sid).unwrap();
4089 let before = read("watches", &wid);
4090 db.touch_watch(&wid, "2026-01-01T00:00:00Z").unwrap();
4091 assert!(read("watches", &wid) > before);
4092 let before = read("watches", &wid);
4093 db.set_watch_session(&wid, "other-session").unwrap();
4094 assert!(read("watches", &wid) > before);
4095
4096 let url = "https://x.test/";
4098 add_session_sources(db.raw(), &sid, &[url.to_string()]).unwrap();
4099 let read_src = || -> String {
4100 std::thread::sleep(std::time::Duration::from_millis(2));
4101 db.raw()
4102 .query_row(
4103 "SELECT updated_at FROM session_sources
4104 WHERE session_id = ?1 AND url_norm = ?2",
4105 (&sid, url),
4106 |r| r.get::<_, String>(0),
4107 )
4108 .unwrap()
4109 };
4110 let before = read_src();
4111 db.set_source_flag(&sid, url, Some("pinned")).unwrap();
4112 assert!(read_src() > before);
4113
4114 let row = db
4116 .log_usage("OpenRouter", "a/model", 1, 2, 3, 4, None, false, None, None)
4117 .unwrap();
4118 let read_usage = |db: &Db, row: i64| -> String {
4119 std::thread::sleep(std::time::Duration::from_millis(2));
4120 db.raw()
4121 .query_row(
4122 "SELECT updated_at FROM usage_log WHERE id = ?1",
4123 [&row],
4124 |r| r.get::<_, String>(0),
4125 )
4126 .unwrap()
4127 };
4128 let before = read_usage(&db, row);
4129 db.update_usage(row, 5, 6, 7, 8, Some(0.1), false).unwrap();
4130 assert!(read_usage(&db, row) > before);
4131 db.upsert_model_prices(&[(
4132 "a/model".to_string(),
4133 "OpenRouter".to_string(),
4134 price(1.0, 2.0),
4135 )])
4136 .unwrap();
4137 let before = read_usage(&db, row);
4138 assert_eq!(db.backfill_usage_costs().unwrap(), 1);
4139 assert!(read_usage(&db, row) > before);
4140 }
4141
4142 #[test]
4143 fn delete_paths_write_tombstones_for_sync() {
4144 let db = Db::open_in_memory().unwrap();
4145 let space = db.default_space_id().unwrap();
4146 let sid = db.create_session("t", "a/b", &space, "chat").unwrap().id;
4147 let tombstones = |table: &str| -> Vec<String> {
4148 let mut stmt = db
4149 .raw()
4150 .prepare("SELECT row_id FROM sync_tombstones WHERE table_name = ?1 ORDER BY row_id")
4151 .unwrap();
4152 let rows = stmt.query_map([table], |r| r.get::<_, String>(0)).unwrap();
4153 rows.collect::<rusqlite::Result<Vec<_>>>().unwrap()
4154 };
4155
4156 let mid = db.add_user_message(&sid, "hi").unwrap();
4158 db.delete_message(&mid).unwrap();
4159 assert_eq!(tombstones("messages"), vec![mid.clone()]);
4160
4161 let sid2 = db.create_session("t2", "a/b", &space, "chat").unwrap().id;
4163 let m1 = db.add_user_message(&sid2, "one").unwrap();
4164 let m2 = db.add_user_message(&sid2, "two").unwrap();
4165 db.delete_session(&sid2).unwrap();
4166 let mut expected = vec![mid, m1, m2];
4167 expected.sort();
4168 assert_eq!(tombstones("messages"), expected);
4169 assert_eq!(tombstones("sessions"), vec![sid2]);
4170
4171 let fid = db.upsert_file(&space, "f.txt", "h", 1, "ok").unwrap();
4173 db.delete_file(&fid).unwrap();
4174 assert_eq!(tombstones("files"), vec![fid]);
4175
4176 let wid = db.create_watch(&space, "topic", 24, &sid).unwrap();
4178 db.delete_watch(&wid).unwrap();
4179 assert_eq!(tombstones("watches"), vec![wid]);
4180
4181 let sp = db.create_space("doomed").unwrap();
4183 db.delete_space(&sp.id).unwrap();
4184 assert_eq!(tombstones("spaces"), vec![sp.id]);
4185
4186 let persona = |name: &str| Persona {
4188 name: name.into(),
4189 model: "m".into(),
4190 blurb: "b".into(),
4191 };
4192 db.save_swarm_personas(&sid, &[persona("a"), persona("b")])
4193 .unwrap();
4194 db.save_swarm_personas(&sid, &[persona("b")]).unwrap();
4195 assert_eq!(
4196 tombstones("swarm_personas"),
4197 vec![format!("{sid}:0"), format!("{sid}:1")]
4198 );
4199 }
4200
4201 #[test]
4206 fn v1_default_space_renumbers_to_deterministic_id() {
4207 let dir = std::env::temp_dir().join(format!("nexus-v1-{}", uuid::Uuid::new_v4()));
4208 std::fs::create_dir_all(&dir).unwrap();
4209 let db_path = dir.join("nexus.db");
4210 {
4211 let db = Db::open(&db_path).unwrap();
4212 let old = db.default_space_id().unwrap();
4213 assert_eq!(old, DEFAULT_SPACE);
4214 db.raw().execute("PRAGMA user_version = 1", []).unwrap();
4217 db.raw()
4218 .execute(
4219 "UPDATE spaces SET id = 'legacy-uuid' WHERE name = ?1",
4220 [DEFAULT_SPACE],
4221 )
4222 .unwrap();
4223 db.raw()
4224 .execute(
4225 "UPDATE sessions SET space_id = 'legacy-uuid' WHERE space_id = ?1",
4226 [&old],
4227 )
4228 .unwrap();
4229 db.raw()
4230 .execute(
4231 "INSERT INTO sync_tombstones (table_name, row_id, deleted_at)
4232 VALUES ('spaces', 'legacy-uuid', '2026-01-01T00:00:00Z')",
4233 [],
4234 )
4235 .unwrap();
4236 let s = db.create_session("t", "m", &old, "chat").unwrap().id;
4237 let _ = s;
4238 drop(db);
4239 }
4240 let db = Db::open(&db_path).unwrap();
4243 assert_eq!(db.default_space_id().unwrap(), DEFAULT_SPACE);
4244 let n: i64 = db
4245 .raw()
4246 .query_row(
4247 "SELECT COUNT(*) FROM sessions WHERE space_id = ?1",
4248 [DEFAULT_SPACE],
4249 |r| r.get(0),
4250 )
4251 .unwrap();
4252 assert_eq!(n, 1, "sessions followed the renumber");
4253 let tomb: i64 = db
4254 .raw()
4255 .query_row(
4256 "SELECT COUNT(*) FROM sync_tombstones WHERE table_name = 'spaces' AND row_id = ?1",
4257 [DEFAULT_SPACE],
4258 |r| r.get(0),
4259 )
4260 .unwrap();
4261 assert_eq!(tomb, 1, "the space tombstone followed too");
4262 let _ = std::fs::remove_dir_all(&dir);
4263 }
4264
4265 #[test]
4266 fn device_id_is_stable_and_sync_state_roundtrips() {
4267 let db = Db::open_in_memory().unwrap();
4268 let a = db.device_id().unwrap();
4269 let b = db.device_id().unwrap();
4270 assert_eq!(a, b);
4271 assert!(!a.is_empty());
4272
4273 db.set_sync_state(
4275 "peer-1",
4276 "sessions",
4277 Some("2026-01-01T00:00:00Z|id1"),
4278 Some("2026-01-02T00:00:00Z|id2"),
4279 )
4280 .unwrap();
4281 db.set_sync_state("peer-1", "sessions", None, Some("2026-01-03T00:00:00Z|id3"))
4282 .unwrap();
4283 db.set_sync_state("peer-1", "messages", Some("c1"), None)
4284 .unwrap();
4285 let states = db.load_sync_state().unwrap();
4286 assert_eq!(states.len(), 2);
4287 let s = states.iter().find(|s| s.table_name == "sessions").unwrap();
4288 assert_eq!(s.pull_cursor.as_deref(), Some("2026-01-01T00:00:00Z|id1"));
4289 assert_eq!(s.push_cursor.as_deref(), Some("2026-01-03T00:00:00Z|id3"));
4290 assert!(s.last_synced_at.is_some());
4291 let m = states.iter().find(|s| s.table_name == "messages").unwrap();
4292 assert_eq!(m.pull_cursor.as_deref(), Some("c1"));
4293 }
4294
4295 #[test]
4296 fn settings_scope_registry_classifies_keys_and_stores_scope() {
4297 for local in [
4299 "searxng_url",
4300 "langsearch_key",
4301 "search_provider",
4302 "ocr_engine",
4303 "ocr_model",
4304 "local_ocr_model",
4305 "usage_range",
4306 "last_update_check",
4307 ] {
4308 assert!(Db::setting_is_local(local), "{local} should be local");
4309 }
4310 for sync in [
4311 "temperature",
4312 "top_p",
4313 "max_tokens",
4314 "show_stats",
4315 "show_reasoning",
4316 "hide_hints",
4317 "verbosity",
4318 "memory_model",
4319 "embedding_model",
4320 ] {
4321 assert!(!Db::setting_is_local(sync), "{sync} should sync");
4322 }
4323
4324 let db = Db::open_in_memory().unwrap();
4325 db.set_setting("temperature", "0.7").unwrap();
4326 db.set_setting("ocr_engine", "tesseract").unwrap();
4327 let scope = |key: &str| -> String {
4328 db.raw()
4329 .query_row(
4330 "SELECT scope FROM app_settings WHERE key = ?1",
4331 [key],
4332 |r| r.get(0),
4333 )
4334 .unwrap()
4335 };
4336 assert_eq!(scope("temperature"), "sync");
4337 assert_eq!(scope("ocr_engine"), "local");
4338 }
4339
4340 #[test]
4341 fn sync_ids_are_unique_for_citations_and_usage() {
4342 let db = Db::open_in_memory().unwrap();
4343 db.add_citations(
4344 "sp",
4345 "r.md",
4346 &[
4347 ("https://a.test/".to_string(), None),
4348 ("https://b.test/".to_string(), None),
4349 ],
4350 )
4351 .unwrap();
4352 db.log_usage("OpenRouter", "m", 1, 2, 0, 0, None, false, None, None)
4353 .unwrap();
4354 db.log_usage("OpenRouter", "m", 1, 2, 0, 0, None, false, None, None)
4355 .unwrap();
4356 let distinct: i64 = db
4357 .raw()
4358 .query_row(
4359 "SELECT COUNT(DISTINCT sync_id) FROM
4360 (SELECT sync_id FROM citations UNION ALL SELECT sync_id FROM usage_log)",
4361 [],
4362 |r| r.get(0),
4363 )
4364 .unwrap();
4365 let total: i64 = db
4366 .raw()
4367 .query_row(
4368 "SELECT (SELECT COUNT(*) FROM citations) + (SELECT COUNT(*) FROM usage_log)",
4369 [],
4370 |r| r.get(0),
4371 )
4372 .unwrap();
4373 assert_eq!(distinct, total);
4374 }
4375}