#![allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::cast_precision_loss,
clippy::cast_sign_loss
)]
use anyhow::{Context, Result};
use chrono::Utc;
use rusqlite::{Connection, OptionalExtension};
use uuid::Uuid;
use crate::provider::ModelPricing;
#[derive(Debug, Clone)]
pub struct ModelPref {
pub id: String,
pub favorite: bool,
pub last_used: Option<String>,
pub reasoning: Option<String>,
}
#[derive(Debug, Clone)]
pub struct Space {
pub id: String,
pub name: String,
pub created_at: String,
}
#[derive(Debug, Clone)]
pub struct Session {
pub id: String,
pub title: String,
pub model: String,
pub slug: Option<String>,
pub created_at: String,
pub compact_summary: Option<String>,
pub compact_through: i64,
pub web_mode: bool,
pub swarm_mode: bool,
pub kind: String,
pub research_parent_id: Option<String>,
}
#[derive(Debug, Clone)]
pub struct Persona {
pub name: String,
pub model: String,
pub blurb: String,
}
#[derive(Debug, Clone)]
pub struct Watch {
pub id: String,
pub space_id: String,
pub topic: String,
pub interval_hours: i64,
pub session_id: String,
pub last_run_at: Option<String>,
}
#[derive(Debug, Clone)]
pub struct FileRow {
pub id: String,
pub name: String,
pub hash: String,
pub size: i64,
pub status: String,
pub mtime: i64,
}
#[derive(Debug, Clone)]
pub struct Message {
pub role: String,
pub content: String,
pub model: Option<String>,
pub reasoning: Option<String>,
pub tokens: Option<i64>,
pub secs: Option<f64>,
pub cost: Option<f64>,
pub phrase: Option<String>,
pub persona: Option<String>,
pub created_at: Option<String>,
}
pub const DEFAULT_SPACE: &str = "default";
const SCHEMA_VERSION: i64 = 2;
const LEGACY_COLUMN_ADDS: &[(&str, &str, &str)] = &[
(
"messages",
"model",
"ALTER TABLE messages ADD COLUMN model TEXT",
),
(
"messages",
"reasoning",
"ALTER TABLE messages ADD COLUMN reasoning TEXT",
),
(
"messages",
"tokens",
"ALTER TABLE messages ADD COLUMN tokens INTEGER",
),
(
"messages",
"secs",
"ALTER TABLE messages ADD COLUMN secs REAL",
),
(
"messages",
"cost",
"ALTER TABLE messages ADD COLUMN cost REAL",
),
(
"messages",
"phrase",
"ALTER TABLE messages ADD COLUMN phrase TEXT",
),
(
"messages",
"persona",
"ALTER TABLE messages ADD COLUMN persona TEXT",
),
(
"model_prefs",
"reasoning",
"ALTER TABLE model_prefs ADD COLUMN reasoning TEXT",
),
(
"model_prefs",
"updated_at",
"ALTER TABLE model_prefs ADD COLUMN updated_at TEXT",
),
(
"sessions",
"slug",
"ALTER TABLE sessions ADD COLUMN slug TEXT",
),
(
"sessions",
"space_id",
"ALTER TABLE sessions ADD COLUMN space_id TEXT",
),
(
"sessions",
"compact_summary",
"ALTER TABLE sessions ADD COLUMN compact_summary TEXT",
),
(
"sessions",
"compact_through",
"ALTER TABLE sessions ADD COLUMN compact_through INTEGER NOT NULL DEFAULT 0",
),
(
"sessions",
"web_mode",
"ALTER TABLE sessions ADD COLUMN web_mode INTEGER NOT NULL DEFAULT 0",
),
(
"sessions",
"swarm_mode",
"ALTER TABLE sessions ADD COLUMN swarm_mode INTEGER NOT NULL DEFAULT 0",
),
(
"sessions",
"kind",
"ALTER TABLE sessions ADD COLUMN kind TEXT NOT NULL DEFAULT 'chat'",
),
(
"sessions",
"research_parent_id",
"ALTER TABLE sessions ADD COLUMN research_parent_id TEXT",
),
(
"session_sources",
"flag",
"ALTER TABLE session_sources ADD COLUMN flag TEXT",
),
(
"session_sources",
"updated_at",
"ALTER TABLE session_sources ADD COLUMN updated_at TEXT",
),
(
"usage_log",
"cost_is_provider",
"ALTER TABLE usage_log ADD COLUMN cost_is_provider INTEGER",
),
(
"usage_log",
"sync_id",
"ALTER TABLE usage_log ADD COLUMN sync_id TEXT",
),
(
"usage_log",
"updated_at",
"ALTER TABLE usage_log ADD COLUMN updated_at TEXT",
),
(
"app_settings",
"scope",
"ALTER TABLE app_settings ADD COLUMN scope TEXT NOT NULL DEFAULT 'sync'",
),
(
"app_settings",
"updated_at",
"ALTER TABLE app_settings ADD COLUMN updated_at TEXT",
),
(
"spaces",
"updated_at",
"ALTER TABLE spaces ADD COLUMN updated_at TEXT",
),
(
"files",
"updated_at",
"ALTER TABLE files ADD COLUMN updated_at TEXT",
),
(
"citations",
"sync_id",
"ALTER TABLE citations ADD COLUMN sync_id TEXT",
),
(
"watches",
"updated_at",
"ALTER TABLE watches ADD COLUMN updated_at TEXT",
),
];
pub fn cache_path_for(db_path: &std::path::Path) -> std::path::PathBuf {
db_path
.parent()
.filter(|p| !p.as_os_str().is_empty())
.map_or_else(
|| std::path::PathBuf::from("cache.db"),
|p| p.join("cache.db"),
)
}
pub fn open_attached(db_path: &std::path::Path) -> Result<Connection> {
let conn =
Connection::open(db_path).with_context(|| format!("opening db {}", db_path.display()))?;
let cache = cache_path_for(db_path);
let escaped = cache.display().to_string().replace('\'', "''");
conn.execute_batch(&format!("ATTACH DATABASE '{escaped}' AS cache"))
.with_context(|| format!("attaching cache db {}", cache.display()))?;
migrate_cache(&conn, "cache")?;
Ok(conn)
}
pub fn migrate_cache(conn: &Connection, schema: &str) -> Result<()> {
conn.execute_batch(&format!(
"CREATE TABLE IF NOT EXISTS {schema}.web_cache (
url_norm TEXT PRIMARY KEY,
url TEXT NOT NULL,
title TEXT,
text TEXT NOT NULL,
fetched_at TEXT NOT NULL
);
CREATE VIRTUAL TABLE IF NOT EXISTS {schema}.file_chunks USING fts5(
file_id UNINDEXED,
seq UNINDEXED,
location UNINDEXED,
text
);
CREATE TABLE IF NOT EXISTS {schema}.chunk_embeddings (
file_id TEXT NOT NULL,
seq INTEGER NOT NULL,
vec BLOB NOT NULL,
PRIMARY KEY (file_id, seq)
);
CREATE TABLE IF NOT EXISTS {schema}.model_prices (
model_id TEXT PRIMARY KEY,
backend TEXT NOT NULL,
prompt_price REAL NOT NULL,
completion_price REAL NOT NULL,
cache_read_price REAL,
cache_write_price REAL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS {schema}.file_index_state (
file_id TEXT PRIMARY KEY,
mtime INTEGER NOT NULL DEFAULT 0,
status TEXT NOT NULL DEFAULT '',
updated_at TEXT NOT NULL
);",
))?;
Ok(())
}
fn has_column(conn: &Connection, table: &str, column: &str) -> Result<bool> {
let mut stmt = conn.prepare(&format!("PRAGMA main.table_info({table})"))?;
let rows = stmt.query_map([], |r| r.get::<_, String>(1))?;
for row in rows {
if row? == column {
return Ok(true);
}
}
Ok(false)
}
pub struct Db {
conn: Connection,
}
impl Db {
pub fn open(path: &std::path::Path) -> Result<Self> {
let conn = open_attached(path).with_context(|| format!("opening db {}", path.display()))?;
let mut db = Self { conn };
db.migrate()?;
Ok(db)
}
#[cfg(any(test, feature = "test-helpers"))]
pub fn open_in_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
conn.execute_batch("ATTACH DATABASE ':memory:' AS cache")?;
migrate_cache(&conn, "cache")?;
let mut db = Db { conn };
db.migrate()?;
Ok(db)
}
pub(crate) fn conn(&self) -> &Connection {
&self.conn
}
#[cfg(test)]
pub fn conn_for_test(&self) -> &Connection {
&self.conn
}
fn touch(&self, table: &str, id: &str) -> Result<()> {
self.conn.execute(
&format!("UPDATE {table} SET updated_at = ?1 WHERE id = ?2"),
(Utc::now().to_rfc3339(), id),
)?;
Ok(())
}
fn tombstone(&self, table: &str, row_id: &str) -> Result<()> {
self.conn.execute(
"INSERT INTO sync_tombstones (table_name, row_id, deleted_at)
VALUES (?1, ?2, ?3)",
(table, row_id, Utc::now().to_rfc3339()),
)?;
Ok(())
}
#[allow(clippy::too_many_lines)]
fn migrate(&mut self) -> Result<()> {
self.conn.execute_batch(
"CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
title TEXT NOT NULL,
model TEXT NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS messages (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL REFERENCES sessions(id),
role TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_messages_session
ON messages(session_id, created_at);
CREATE TABLE IF NOT EXISTS model_prefs (
id TEXT PRIMARY KEY,
favorite INTEGER NOT NULL DEFAULT 0,
last_used TEXT,
reasoning TEXT,
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS app_settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL,
scope TEXT NOT NULL DEFAULT 'sync',
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS spaces (
id TEXT PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
created_at TEXT NOT NULL,
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS files (
id TEXT PRIMARY KEY,
space_id TEXT NOT NULL,
name TEXT NOT NULL,
hash TEXT NOT NULL,
size INTEGER NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
UNIQUE(space_id, name)
);
CREATE TABLE IF NOT EXISTS citations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
sync_id TEXT NOT NULL,
space_id TEXT NOT NULL,
report_file TEXT NOT NULL,
url TEXT NOT NULL,
title TEXT
);
CREATE INDEX IF NOT EXISTS idx_citations_space ON citations(space_id);
CREATE TABLE IF NOT EXISTS session_sources (
session_id TEXT NOT NULL,
url_norm TEXT NOT NULL,
flag TEXT,
updated_at TEXT,
PRIMARY KEY (session_id, url_norm)
);
CREATE TABLE IF NOT EXISTS watches (
id TEXT PRIMARY KEY,
space_id TEXT NOT NULL,
topic TEXT NOT NULL,
interval_hours INTEGER NOT NULL,
session_id TEXT NOT NULL,
last_run_at TEXT,
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS swarm_personas (
session_id TEXT NOT NULL,
ord INTEGER NOT NULL,
name TEXT NOT NULL,
model TEXT NOT NULL,
persona TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_swarm_personas_session
ON swarm_personas(session_id, ord);
CREATE TABLE IF NOT EXISTS usage_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
sync_id TEXT NOT NULL,
created_at TEXT NOT NULL,
session_id TEXT,
space_id TEXT,
backend TEXT NOT NULL,
model TEXT NOT NULL,
prompt_tokens INTEGER NOT NULL,
completion_tokens INTEGER NOT NULL,
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
cost REAL,
cost_is_provider INTEGER,
updated_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_usage_log_created ON usage_log(created_at);
CREATE INDEX IF NOT EXISTS idx_usage_log_model ON usage_log(model);
CREATE TABLE IF NOT EXISTS sync_tombstones (
id INTEGER PRIMARY KEY AUTOINCREMENT,
table_name TEXT NOT NULL,
row_id TEXT NOT NULL,
deleted_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_sync_tombstones_table ON sync_tombstones(table_name, row_id);
CREATE TABLE IF NOT EXISTS device_meta (
device_id TEXT PRIMARY KEY,
created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS sync_state (
peer_id TEXT NOT NULL,
table_name TEXT NOT NULL,
pull_cursor TEXT,
push_cursor TEXT,
last_synced_at TEXT,
PRIMARY KEY (peer_id, table_name)
);",
)?;
let version: i64 = self
.conn
.query_row("PRAGMA user_version", [], |r| r.get(0))?;
if version < SCHEMA_VERSION {
for (table, column, ddl) in LEGACY_COLUMN_ADDS {
if !has_column(&self.conn, table, column)? {
self.conn.execute(ddl, []).with_context(|| {
format!("migrating column {table}.{column} (user_version {version})")
})?;
}
}
let backfill_tx = self.conn.transaction()?;
for table in ["citations", "usage_log"] {
let ids: Vec<i64> = {
let mut stmt = backfill_tx
.prepare(&format!("SELECT id FROM {table} WHERE sync_id IS NULL"))?;
let rows = stmt.query_map([], |r| r.get(0))?;
rows.collect::<rusqlite::Result<Vec<_>>>()?
};
let mut update = backfill_tx
.prepare(&format!("UPDATE {table} SET sync_id = ?1 WHERE id = ?2"))?;
for id in ids {
update.execute((Uuid::new_v4().to_string(), id))?;
}
}
backfill_tx.commit()?;
self.conn.execute_batch(
"CREATE UNIQUE INDEX IF NOT EXISTS idx_citations_sync_id ON citations(sync_id);
CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_log_sync_id ON usage_log(sync_id);",
)?;
let now = Utc::now().to_rfc3339();
if has_column(&self.conn, "files", "mtime")? {
self.conn.execute(
"INSERT OR IGNORE INTO cache.file_index_state (file_id, mtime, status, updated_at)
SELECT id, mtime, status, ?1 FROM files",
[&now],
)?;
}
if has_column(&self.conn, "web_cache", "url_norm")? {
self.conn.execute(
"INSERT OR IGNORE INTO cache.web_cache (url_norm, url, title, text, fetched_at)
SELECT url_norm, url, title, text, fetched_at FROM web_cache",
[],
)?;
}
if has_column(&self.conn, "chunk_embeddings", "file_id")? {
self.conn.execute(
"INSERT OR IGNORE INTO cache.chunk_embeddings (file_id, seq, vec)
SELECT file_id, seq, vec FROM chunk_embeddings",
[],
)?;
}
if has_column(&self.conn, "file_chunks", "file_id")? {
self.conn.execute(
"INSERT OR IGNORE INTO cache.file_chunks (file_id, seq, location, text)
SELECT file_id, seq, location, text FROM file_chunks",
[],
)?;
}
if has_column(&self.conn, "model_prices", "model_id")? {
let cols = if has_column(&self.conn, "model_prices", "cache_read_price")? {
"model_id, backend, prompt_price, completion_price,\n cache_read_price, cache_write_price, updated_at"
} else {
"model_id, backend, prompt_price, completion_price, updated_at"
};
self.conn.execute(
&format!(
"INSERT OR IGNORE INTO cache.model_prices ({cols}) SELECT {cols} FROM model_prices"
),
[],
)?;
}
let old_default: Option<String> = self
.conn
.query_row(
"SELECT id FROM spaces WHERE name = ?1",
[DEFAULT_SPACE],
|r| r.get(0),
)
.optional()?;
if let Some(old) = old_default.filter(|id| id != DEFAULT_SPACE) {
let tx = self.conn.transaction()?;
for table in ["sessions", "files", "watches", "usage_log", "citations"] {
tx.execute(
&format!("UPDATE {table} SET space_id = ?1 WHERE space_id = ?2"),
(DEFAULT_SPACE, &old),
)?;
}
tx.execute(
"UPDATE sync_tombstones SET row_id = ?1
WHERE table_name = 'spaces' AND row_id = ?2",
(DEFAULT_SPACE, &old),
)?;
tx.execute(
"UPDATE spaces SET id = ?1 WHERE id = ?2",
(DEFAULT_SPACE, &old),
)?;
tx.commit()?;
}
self.conn
.execute_batch(&format!("PRAGMA user_version = {SCHEMA_VERSION}"))?;
}
let _ = self
.conn
.execute_batch("DROP TABLE IF EXISTS message_images;");
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT OR IGNORE INTO spaces (id, name, created_at) VALUES (?1, ?2, ?3)",
(DEFAULT_SPACE, DEFAULT_SPACE, &now),
)?;
let default_id: String = self.conn.query_row(
"SELECT id FROM spaces WHERE name = ?1",
[DEFAULT_SPACE],
|r| r.get(0),
)?;
self.conn.execute(
"UPDATE sessions SET space_id = ?1 WHERE space_id IS NULL",
[&default_id],
)?;
Ok(())
}
pub fn default_space_id(&self) -> Result<String> {
Ok(self.conn.query_row(
"SELECT id FROM spaces WHERE name = ?1",
[DEFAULT_SPACE],
|r| r.get(0),
)?)
}
pub fn create_space(&self, name: &str) -> Result<Space> {
let id = Uuid::new_v4().to_string();
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO spaces (id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?3)",
(&id, name, &now),
)?;
Ok(Space {
id,
name: name.to_string(),
created_at: now,
})
}
pub fn list_spaces(&self) -> Result<Vec<Space>> {
let mut stmt = self
.conn
.prepare("SELECT id, name, created_at FROM spaces ORDER BY created_at ASC")?;
let rows = stmt.query_map([], |r| {
Ok(Space {
id: r.get(0)?,
name: r.get(1)?,
created_at: r.get(2)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn rename_space(&self, id: &str, name: &str) -> Result<()> {
self.conn.execute(
"UPDATE spaces SET name = ?2, updated_at = ?3 WHERE id = ?1",
(id, name, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn delete_space(&self, id: &str) -> Result<()> {
let default_id = self.default_space_id()?;
self.conn.execute(
"UPDATE sessions SET space_id = ?1, updated_at = ?2 WHERE space_id = ?3",
(&default_id, Utc::now().to_rfc3339(), id),
)?;
self.conn
.execute("DELETE FROM spaces WHERE id = ?1", [id])?;
self.tombstone("spaces", id)?;
Ok(())
}
pub fn count_sessions(&self, space_id: &str) -> Result<u64> {
let n: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM sessions WHERE space_id = ?1",
[space_id],
|r| r.get(0),
)?;
Ok(n as u64)
}
pub fn last_message_preview(&self, session_id: &str) -> Option<String> {
let mut stmt = self
.conn
.prepare(
"SELECT content FROM messages WHERE session_id = ?1 \
AND role IN ('user','assistant') ORDER BY id DESC LIMIT 1",
)
.ok()?;
let mut rows = stmt
.query_map([session_id], |r| r.get::<_, String>(0))
.ok()?;
rows.next().and_then(Result::ok)
}
pub fn setting_is_local(key: &str) -> bool {
matches!(
key,
"searxng_url"
| "langsearch_key"
| "search_provider"
| "ocr_engine"
| "ocr_model"
| "local_ocr_model"
| "usage_range"
| "ui_background"
| "last_update_check"
| "sync_ssh_peer"
)
}
pub fn set_setting(&self, key: &str, value: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
let scope = if Self::setting_is_local(key) {
"local"
} else {
"sync"
};
self.conn.execute(
"INSERT INTO app_settings (key, value, scope, updated_at)
VALUES (?1, ?2, ?3, ?4)
ON CONFLICT(key) DO UPDATE SET value = ?2, scope = ?3, updated_at = ?4",
(key, value, scope, &now),
)?;
Ok(())
}
pub fn load_settings(&self) -> Result<Vec<(String, String)>> {
let mut stmt = self.conn.prepare("SELECT key, value FROM app_settings")?;
let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn update_check_due(&self) -> bool {
let today = chrono::Local::now().format("%Y-%m-%d").to_string();
let last = self
.conn
.query_row(
"SELECT value FROM app_settings WHERE key = 'last_update_check'",
[],
|r| r.get::<_, String>(0),
)
.ok();
if last.as_deref() == Some(today.as_str()) {
return false;
}
let _ = self.set_setting("last_update_check", &today);
true
}
pub fn set_reasoning(&self, model_id: &str, effort: Option<&str>) -> Result<()> {
self.conn.execute(
"INSERT INTO model_prefs (id, reasoning, updated_at) VALUES (?1, ?2, ?3)
ON CONFLICT(id) DO UPDATE SET reasoning = ?2, updated_at = ?3",
(model_id, effort, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn toggle_favorite(&self, model_id: &str) -> Result<bool> {
self.conn.execute(
"INSERT INTO model_prefs (id, favorite, updated_at) VALUES (?1, 1, ?2)
ON CONFLICT(id) DO UPDATE SET favorite = 1 - favorite, updated_at = ?2",
(model_id, Utc::now().to_rfc3339()),
)?;
let fav: i64 = self.conn.query_row(
"SELECT favorite FROM model_prefs WHERE id = ?1",
[model_id],
|r| r.get(0),
)?;
Ok(fav != 0)
}
pub fn mark_model_used(&self, model_id: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO model_prefs (id, favorite, last_used, updated_at) VALUES (?1, 0, ?2, ?2)
ON CONFLICT(id) DO UPDATE SET last_used = ?2, updated_at = ?2",
(model_id, &now),
)?;
Ok(())
}
pub fn load_model_prefs(&self) -> Result<Vec<ModelPref>> {
let mut stmt = self
.conn
.prepare("SELECT id, favorite, last_used, reasoning FROM model_prefs")?;
let rows = stmt.query_map([], |r| {
Ok(ModelPref {
id: r.get(0)?,
favorite: r.get::<_, i64>(1)? != 0,
last_used: r.get(2)?,
reasoning: r.get(3)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn create_session(
&self,
title: &str,
model: &str,
space_id: &str,
kind: &str,
) -> Result<Session> {
let id = Uuid::new_v4().to_string();
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO sessions (id, title, model, space_id, kind, created_at, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6)",
(&id, title, model, space_id, kind, &now),
)?;
Ok(Session {
id,
title: title.to_string(),
model: model.to_string(),
slug: None,
created_at: now,
compact_summary: None,
compact_through: 0,
web_mode: false,
swarm_mode: false,
kind: kind.to_string(),
research_parent_id: None,
})
}
pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
self.conn
.query_row(
"SELECT id, title, model, slug, created_at, compact_summary, compact_through, \
web_mode, swarm_mode, kind, research_parent_id
FROM sessions WHERE id = ?1",
[id],
|r| {
Ok(Session {
id: r.get(0)?,
title: r.get(1)?,
model: r.get(2)?,
slug: r.get(3)?,
created_at: r.get(4)?,
compact_summary: r.get(5)?,
compact_through: r.get(6)?,
web_mode: r.get::<_, i64>(7)? != 0,
swarm_mode: r.get::<_, i64>(8)? != 0,
kind: r.get(9)?,
research_parent_id: r.get(10)?,
})
},
)
.optional()
.map_err(Into::into)
}
pub fn list_sessions(&self, space_id: &str) -> Result<Vec<Session>> {
let mut stmt = self.conn.prepare(
"SELECT id, title, model, slug, created_at, compact_summary, compact_through, \
web_mode, swarm_mode, kind, research_parent_id
FROM sessions WHERE space_id = ?1 ORDER BY updated_at DESC",
)?;
let rows = stmt.query_map([space_id], |r| {
Ok(Session {
id: r.get(0)?,
title: r.get(1)?,
model: r.get(2)?,
slug: r.get(3)?,
created_at: r.get(4)?,
compact_summary: r.get(5)?,
compact_through: r.get(6)?,
web_mode: r.get::<_, i64>(7)? != 0,
swarm_mode: r.get::<_, i64>(8)? != 0,
kind: r.get(9)?,
research_parent_id: r.get(10)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn set_compaction(&self, session_id: &str, summary: &str, through: i64) -> Result<()> {
self.conn.execute(
"UPDATE sessions SET compact_summary = ?2, compact_through = ?3, updated_at = ?4
WHERE id = ?1",
(session_id, summary, through, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn set_session_web_mode(&self, session_id: &str, on: bool) -> Result<()> {
self.conn.execute(
"UPDATE sessions SET web_mode = ?2, updated_at = ?3 WHERE id = ?1",
(session_id, i64::from(on), Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn set_session_swarm_mode(&self, session_id: &str, on: bool) -> Result<()> {
self.conn.execute(
"UPDATE sessions SET swarm_mode = ?2, updated_at = ?3 WHERE id = ?1",
(session_id, i64::from(on), Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn list_swarm_personas(&self, session_id: &str) -> Result<Vec<Persona>> {
let mut stmt = self.conn.prepare(
"SELECT name, model, persona FROM swarm_personas
WHERE session_id = ?1 ORDER BY ord ASC",
)?;
let rows = stmt.query_map([session_id], |r| {
Ok(Persona {
name: r.get(0)?,
model: r.get(1)?,
blurb: r.get(2)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn save_swarm_personas(&self, session_id: &str, personas: &[Persona]) -> Result<()> {
let old: Vec<i64> = {
let mut stmt = self
.conn
.prepare("SELECT ord FROM swarm_personas WHERE session_id = ?1")?;
let rows = stmt.query_map([session_id], |r| r.get(0))?;
rows.collect::<rusqlite::Result<Vec<_>>>()?
};
for ord in old {
self.tombstone("swarm_personas", &format!("{session_id}:{ord}"))?;
}
self.conn.execute(
"DELETE FROM swarm_personas WHERE session_id = ?1",
[session_id],
)?;
for (i, p) in personas.iter().enumerate() {
self.conn.execute(
"INSERT INTO swarm_personas (session_id, ord, name, model, persona)
VALUES (?1, ?2, ?3, ?4, ?5)",
(session_id, i as i64, &p.name, &p.model, &p.blurb),
)?;
}
self.touch("sessions", session_id)
}
pub fn set_research_parent(&self, id: &str, parent_id: &str) -> Result<()> {
self.conn.execute(
"UPDATE sessions SET research_parent_id = ?2, updated_at = ?3 WHERE id = ?1",
(id, parent_id, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn set_session_title(&self, id: &str, title: &str, slug: Option<&str>) -> Result<()> {
self.conn.execute(
"UPDATE sessions SET title = ?2, slug = COALESCE(?3, slug), updated_at = ?4
WHERE id = ?1",
(id, title, slug, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn delete_message(&self, id: &str) -> Result<()> {
self.conn
.execute("DELETE FROM messages WHERE id = ?1", [id])?;
self.tombstone("messages", id)?;
Ok(())
}
pub fn delete_session(&self, id: &str) -> Result<()> {
let ids: Vec<String> = {
let mut stmt = self
.conn
.prepare("SELECT id FROM messages WHERE session_id = ?1")?;
let rows = stmt.query_map([id], |r| r.get(0))?;
rows.collect::<rusqlite::Result<Vec<_>>>()?
};
for mid in &ids {
self.tombstone("messages", mid)?;
}
self.conn
.execute("DELETE FROM messages WHERE session_id = ?1", [id])?;
self.conn
.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
self.tombstone("sessions", id)?;
Ok(())
}
pub fn load_messages(&self, session_id: &str) -> Result<Vec<Message>> {
let mut stmt = self.conn.prepare(
"SELECT role, content, model, reasoning, tokens, secs, cost, phrase, persona, created_at
FROM messages WHERE session_id = ?1 ORDER BY created_at ASC",
)?;
let messages = stmt
.query_map([session_id], |r| {
Ok(Message {
role: r.get(0)?,
content: r.get(1)?,
model: r.get(2)?,
reasoning: r.get(3)?,
tokens: r.get(4)?,
secs: r.get(5)?,
cost: r.get(6)?,
phrase: r.get(7)?,
persona: r.get(8)?,
created_at: r.get(9)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(messages)
}
pub fn add_user_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id, "user", content, None, None, None, None, None, None,
)
}
pub fn add_gate_reply_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id,
"gate_reply",
content,
None,
None,
None,
None,
None,
None,
)
}
pub fn add_tool_call_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id,
"tool_call",
content,
None,
None,
None,
None,
None,
None,
)
}
pub fn add_error_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id, "error", content, None, None, None, None, None, None,
)
}
pub fn add_research_stage_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id,
"research_stage",
content,
None,
None,
None,
None,
None,
None,
)
}
pub fn add_research_plan_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id,
"research_plan",
content,
None,
None,
None,
None,
None,
None,
)
}
pub fn add_survey_message(&self, session_id: &str, content: &str) -> Result<String> {
self.insert_message(
session_id, "survey", content, None, None, None, None, None, None,
)
}
pub fn upsert_research_stage_message(
&self,
session_id: &str,
label: &str,
detail: &str,
) -> Result<()> {
let content = stage_content(label, detail);
let existing: Option<String> = self
.conn
.query_row(
"SELECT id FROM messages WHERE session_id = ?1 AND role = 'research_stage'
AND (content = ?2 OR content LIKE ?3)
ORDER BY created_at DESC LIMIT 1",
(session_id, label, format!("{label}:%")),
|r| r.get(0),
)
.ok();
match existing {
Some(id) => {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE messages SET content = ?2, created_at = ?3 WHERE id = ?1",
(&id, &content, &now),
)?;
}
None => {
self.add_research_stage_message(session_id, &content)?;
}
}
Ok(())
}
#[cfg(test)]
pub fn add_session_sources(&self, session_id: &str, url_norms: &[String]) -> Result<()> {
add_session_sources(&self.conn, session_id, url_norms)
}
#[cfg(test)]
pub fn search_session_sources(
&self,
session_id: &str,
query: &str,
) -> Result<Vec<(String, String)>> {
search_session_sources(&self.conn, session_id, query)
}
pub fn set_source_flag(
&self,
session_id: &str,
url_norm: &str,
flag: Option<&str>,
) -> Result<()> {
self.conn.execute(
"UPDATE session_sources SET flag = ?3, updated_at = ?4
WHERE session_id = ?1 AND url_norm = ?2",
(session_id, url_norm, flag, Utc::now().to_rfc3339()),
)?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub fn add_assistant_message(
&self,
session_id: &str,
content: &str,
model: Option<&str>,
reasoning: Option<&str>,
tokens: Option<i64>,
secs: Option<f64>,
cost: Option<f64>,
phrase: Option<&str>,
) -> Result<String> {
self.insert_message(
session_id,
"assistant",
content,
model,
reasoning,
tokens,
secs,
cost,
phrase,
)
}
pub fn add_persona_message(
&self,
session_id: &str,
content: &str,
persona_name: &str,
model: &str,
) -> Result<String> {
let id = self.insert_message(
session_id,
"assistant",
content,
Some(model),
None,
None,
None,
None,
None,
)?;
self.conn.execute(
"UPDATE messages SET persona = ?2 WHERE id = ?1",
(&id, persona_name),
)?;
Ok(id)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn insert_message(
&self,
session_id: &str,
role: &str,
content: &str,
model: Option<&str>,
reasoning: Option<&str>,
tokens: Option<i64>,
secs: Option<f64>,
cost: Option<f64>,
phrase: Option<&str>,
) -> Result<String> {
let now = Utc::now().to_rfc3339();
let id = Uuid::new_v4().to_string();
self.conn.execute(
"INSERT INTO messages
(id, session_id, role, content, model, reasoning, tokens, secs, cost, phrase, created_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
(
&id, session_id, role, content, model, reasoning, tokens, secs, cost, phrase, &now,
),
)?;
self.conn.execute(
"UPDATE sessions SET updated_at = ?2 WHERE id = ?1",
(session_id, &now),
)?;
Ok(id)
}
pub fn message_created_at(&self, session_id: &str, index: usize) -> Result<Option<String>> {
let mut stmt = self.conn.prepare(
"SELECT created_at FROM messages WHERE session_id = ?1
ORDER BY created_at ASC LIMIT 1 OFFSET ?2",
)?;
Ok(stmt
.query_row((session_id, index as i64), |r| r.get(0))
.optional()?)
}
pub fn add_compaction_message(
&self,
session_id: &str,
content: &str,
at: &str,
) -> Result<String> {
let id = Uuid::new_v4().to_string();
self.conn.execute(
"INSERT INTO messages
(id, session_id, role, content, model, reasoning, tokens, secs, phrase, created_at)
VALUES (?1, ?2, 'compaction', ?3, NULL, NULL, NULL, NULL, NULL, ?4)",
(&id, session_id, content, at),
)?;
Ok(id)
}
pub fn update_compaction_message(&self, session_id: &str, content: &str) -> Result<usize> {
Ok(self.conn.execute(
"UPDATE messages SET content = ?2
WHERE session_id = ?1 AND role = 'compaction'",
(session_id, content),
)?)
}
pub fn set_session_model(&self, session_id: &str, model: &str) -> Result<()> {
self.conn.execute(
"UPDATE sessions SET model = ?2, updated_at = ?3 WHERE id = ?1",
(session_id, model, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn upsert_file(
&self,
space_id: &str,
name: &str,
hash: &str,
size: i64,
status: &str,
) -> Result<String> {
let now = Utc::now().to_rfc3339();
if let Ok(existing) = self.conn.query_row(
"SELECT id FROM files WHERE space_id = ?1 AND name = ?2",
(space_id, name),
|r| r.get::<_, String>(0),
) {
self.conn.execute(
"UPDATE files SET hash = ?2, size = ?3, updated_at = ?4 WHERE id = ?1",
(&existing, hash, size, &now),
)?;
self.conn.execute(
"INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
VALUES (?1, 0, ?2, ?3)
ON CONFLICT(file_id) DO UPDATE SET status = excluded.status,
updated_at = excluded.updated_at",
(&existing, status, &now),
)?;
return Ok(existing);
}
let id = Uuid::new_v4().to_string();
self.conn.execute(
"INSERT INTO files (id, space_id, name, hash, size, created_at, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
(&id, space_id, name, hash, size, &now, &now),
)?;
self.conn.execute(
"INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
VALUES (?1, 0, ?2, ?3)",
(&id, status, &now),
)?;
Ok(id)
}
pub fn list_files(&self, space_id: &str) -> Result<Vec<FileRow>> {
let mut stmt = self.conn.prepare(
"SELECT files.id, files.name, files.hash, files.size,
COALESCE(cache.file_index_state.status, 'not indexed'),
COALESCE(cache.file_index_state.mtime, 0)
FROM files
LEFT JOIN cache.file_index_state
ON cache.file_index_state.file_id = files.id
WHERE files.space_id = ?1 ORDER BY files.name ASC",
)?;
let rows = stmt.query_map([space_id], |r| {
Ok(FileRow {
id: r.get(0)?,
name: r.get(1)?,
hash: r.get(2)?,
size: r.get(3)?,
status: r.get(4)?,
mtime: r.get(5)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn file_indexed(&self, file_id: &str) -> Result<bool> {
Ok(self.conn.query_row(
"SELECT EXISTS(SELECT 1 FROM cache.file_index_state WHERE file_id = ?1)",
[file_id],
|r| r.get(0),
)?)
}
pub fn delete_file(&self, file_id: &str) -> Result<()> {
self.conn.execute(
"DELETE FROM cache.file_chunks WHERE file_id = ?1",
[file_id],
)?;
self.conn.execute(
"DELETE FROM cache.chunk_embeddings WHERE file_id = ?1",
[file_id],
)?;
self.conn.execute(
"DELETE FROM cache.file_index_state WHERE file_id = ?1",
[file_id],
)?;
self.conn
.execute("DELETE FROM files WHERE id = ?1", [file_id])?;
self.tombstone("files", file_id)?;
Ok(())
}
pub fn set_file_mtime(&self, file_id: &str, mtime: i64) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
VALUES (?1, ?2, '', ?3)
ON CONFLICT(file_id) DO UPDATE SET mtime = excluded.mtime,
updated_at = excluded.updated_at",
(file_id, mtime, &now),
)?;
Ok(())
}
pub fn set_file_status(&self, file_id: &str, status: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO cache.file_index_state (file_id, mtime, status, updated_at)
VALUES (?1, 0, ?2, ?3)
ON CONFLICT(file_id) DO UPDATE SET status = excluded.status,
updated_at = excluded.updated_at",
(file_id, status, &now),
)?;
Ok(())
}
pub fn rename_file(&self, file_id: &str, new_name: &str) -> Result<()> {
self.conn.execute(
"UPDATE files SET name = ?2, updated_at = ?3 WHERE id = ?1",
(file_id, new_name, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn replace_file_ref_in_messages(
&self,
space_id: &str,
old_name: &str,
new_name: &str,
) -> Result<()> {
self.conn.execute(
"UPDATE messages SET content = REPLACE(content, ?1, ?2)
WHERE session_id IN (SELECT id FROM sessions WHERE space_id = ?3)",
(old_name, new_name, space_id),
)?;
Ok(())
}
pub fn set_file_chunks(&self, file_id: &str, chunks: &[(String, String)]) -> Result<()> {
self.conn.execute(
"DELETE FROM cache.file_chunks WHERE file_id = ?1",
[file_id],
)?;
self.conn.execute(
"DELETE FROM cache.chunk_embeddings WHERE file_id = ?1",
[file_id],
)?;
for (seq, (location, text)) in chunks.iter().enumerate() {
self.conn.execute(
"INSERT INTO cache.file_chunks (file_id, seq, location, text) VALUES (?1, ?2, ?3, ?4)",
(file_id, seq as i64, location, text),
)?;
}
Ok(())
}
#[cfg(test)]
pub fn raw(&self) -> &Connection {
&self.conn
}
pub fn integrity_check(&self) -> Result<String> {
self.conn
.query_row("PRAGMA integrity_check", [], |r| r.get(0))
.context("running integrity check")
}
pub fn file_chunk_texts(&self, file_id: &str) -> Result<Vec<(i64, String)>> {
let mut stmt = self.conn.prepare(
"SELECT CAST(seq AS INTEGER), text FROM cache.file_chunks
WHERE file_id = ?1 ORDER BY CAST(seq AS INTEGER) ASC",
)?;
let rows = stmt.query_map([file_id], |r| Ok((r.get(0)?, r.get(1)?)))?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn files_missing_embeddings(&self, space_id: &str) -> Result<Vec<String>> {
files_missing_embeddings(&self.conn, space_id)
}
pub fn add_citations(
&self,
space_id: &str,
report_file: &str,
citations: &[(String, Option<String>)],
) -> Result<()> {
for (url, title) in citations {
self.conn.execute(
"INSERT INTO citations (sync_id, space_id, report_file, url, title)
VALUES (?1, ?2, ?3, ?4, ?5)",
(
Uuid::new_v4().to_string(),
space_id,
report_file,
url,
title,
),
)?;
}
Ok(())
}
pub fn search_citations(
&self,
space_id: &str,
query: Option<&str>,
) -> Result<Vec<(String, String, String)>> {
search_citations(&self.conn, space_id, query)
}
pub fn set_chunk_embeddings(&self, file_id: &str, vecs: &[(i64, Vec<f32>)]) -> Result<()> {
for (seq, v) in vecs {
self.conn.execute(
"INSERT OR REPLACE INTO cache.chunk_embeddings (file_id, seq, vec) VALUES (?1, ?2, ?3)",
(file_id, seq, vec_to_blob(v)),
)?;
}
Ok(())
}
pub fn create_watch(
&self,
space_id: &str,
topic: &str,
interval_hours: i64,
session_id: &str,
) -> Result<String> {
let id = Uuid::new_v4().to_string();
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO watches (id, space_id, topic, interval_hours, session_id, last_run_at, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, NULL, ?6)",
(&id, space_id, topic, interval_hours, session_id, &now),
)?;
Ok(id)
}
pub fn list_watches(&self, space_id: &str) -> Result<Vec<Watch>> {
let mut stmt = self.conn.prepare(
"SELECT id, space_id, topic, interval_hours, session_id, last_run_at
FROM watches WHERE space_id = ?1 ORDER BY topic",
)?;
let rows = stmt.query_map([space_id], |r| {
Ok(Watch {
id: r.get(0)?,
space_id: r.get(1)?,
topic: r.get(2)?,
interval_hours: r.get(3)?,
session_id: r.get(4)?,
last_run_at: r.get(5)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn list_all_watches(&self) -> Result<Vec<Watch>> {
let mut stmt = self.conn.prepare(
"SELECT id, space_id, topic, interval_hours, session_id, last_run_at FROM watches",
)?;
let rows = stmt.query_map([], |r| {
Ok(Watch {
id: r.get(0)?,
space_id: r.get(1)?,
topic: r.get(2)?,
interval_hours: r.get(3)?,
session_id: r.get(4)?,
last_run_at: r.get(5)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn touch_watch(&self, id: &str, now_rfc3339: &str) -> Result<()> {
self.conn.execute(
"UPDATE watches SET last_run_at = ?2, updated_at = ?3 WHERE id = ?1",
(id, now_rfc3339, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn set_watch_session(&self, id: &str, session_id: &str) -> Result<()> {
self.conn.execute(
"UPDATE watches SET session_id = ?2, updated_at = ?3 WHERE id = ?1",
(id, session_id, Utc::now().to_rfc3339()),
)?;
Ok(())
}
pub fn delete_watch(&self, id: &str) -> Result<()> {
self.conn
.execute("DELETE FROM watches WHERE id = ?1", [id])?;
self.tombstone("watches", id)?;
Ok(())
}
}
pub fn vec_to_blob(v: &[f32]) -> Vec<u8> {
v.iter().flat_map(|f| f.to_le_bytes()).collect()
}
pub fn blob_to_vec(b: &[u8]) -> Vec<f32> {
b.chunks_exact(4)
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
.collect()
}
pub fn search_citations(
conn: &Connection,
space_id: &str,
query: Option<&str>,
) -> Result<Vec<(String, String, String)>> {
let mut stmt = conn.prepare(
"SELECT report_file, url, COALESCE(title, '') FROM citations
WHERE space_id = ?1
AND (?2 IS NULL OR url LIKE ?2 OR title LIKE ?2 OR report_file LIKE ?2)
ORDER BY id DESC",
)?;
let pattern = query.map(|q| format!("%{q}%"));
let rows = stmt.query_map((space_id, pattern), |r| {
Ok((r.get(0)?, r.get(1)?, r.get(2)?))
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn stage_content(label: &str, detail: &str) -> String {
if detail.is_empty() {
label.to_string()
} else {
format!("{label}: {detail}")
}
}
pub fn add_session_sources(
conn: &Connection,
session_id: &str,
url_norms: &[String],
) -> Result<()> {
let now = Utc::now().to_rfc3339();
for u in url_norms {
conn.execute(
"INSERT OR IGNORE INTO session_sources (session_id, url_norm, updated_at)
VALUES (?1, ?2, ?3)",
(session_id, u, &now),
)?;
}
Ok(())
}
pub fn search_session_sources(
conn: &Connection,
session_id: &str,
query: &str,
) -> Result<Vec<(String, String)>> {
let mut stmt = conn.prepare(
"SELECT cache.web_cache.url, cache.web_cache.text FROM session_sources
JOIN cache.web_cache ON cache.web_cache.url_norm = session_sources.url_norm
WHERE session_sources.session_id = ?1",
)?;
let rows = stmt.query_map([session_id], |r| {
Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
})?;
let needle = query.to_lowercase();
Ok(rows
.collect::<rusqlite::Result<Vec<_>>>()?
.into_iter()
.filter(|(_, text)| text.to_lowercase().contains(&needle))
.collect())
}
pub fn pinned_urls(conn: &Connection, session_id: &str) -> Result<Vec<String>> {
let mut stmt = conn.prepare(
"SELECT url_norm FROM session_sources WHERE session_id = ?1 AND flag = 'pinned'",
)?;
let rows = stmt.query_map([session_id], |r| r.get::<_, String>(0))?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn discarded_domains(conn: &Connection, session_id: &str) -> Result<Vec<String>> {
let mut stmt = conn.prepare(
"SELECT url_norm FROM session_sources WHERE session_id = ?1 AND flag = 'discarded'",
)?;
let rows: Vec<String> = stmt
.query_map([session_id], |r| r.get::<_, String>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
let mut hosts: Vec<String> = rows
.iter()
.filter_map(|u| {
reqwest::Url::parse(u)
.ok()
.and_then(|p| p.host_str().map(str::to_string))
})
.collect();
hosts.sort();
hosts.dedup();
Ok(hosts)
}
pub fn is_fresh(fetched_at: &str, now: chrono::DateTime<Utc>) -> bool {
chrono::DateTime::parse_from_rfc3339(fetched_at)
.is_ok_and(|dt| now.signed_duration_since(dt) < chrono::Duration::hours(24))
}
pub fn cache_get(conn: &Connection, url_norm: &str) -> Result<Option<(String, String, String)>> {
let row = conn.query_row(
"SELECT COALESCE(title, ''), text, fetched_at FROM web_cache WHERE url_norm = ?1",
[url_norm],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
);
match row {
Ok(v) => Ok(Some(v)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(e) => Err(e.into()),
}
}
pub fn cache_put(
conn: &Connection,
url_norm: &str,
url: &str,
title: Option<&str>,
text: &str,
) -> Result<()> {
let now = Utc::now().to_rfc3339();
conn.execute(
"INSERT INTO web_cache (url_norm, url, title, text, fetched_at) VALUES (?1, ?2, ?3, ?4, ?5)
ON CONFLICT(url_norm) DO UPDATE SET url = ?2, title = ?3, text = ?4, fetched_at = ?5",
(url_norm, url, title, text, &now),
)?;
Ok(())
}
pub fn files_missing_embeddings(conn: &Connection, space_id: &str) -> Result<Vec<String>> {
let mut stmt = conn.prepare(
"SELECT files.id FROM files
WHERE files.space_id = ?1
AND (SELECT COUNT(*) FROM cache.file_chunks WHERE cache.file_chunks.file_id = files.id) >
(SELECT COUNT(*) FROM cache.chunk_embeddings WHERE cache.chunk_embeddings.file_id = files.id)",
)?;
let rows = stmt.query_map([space_id], |r| r.get(0))?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn semantic_chunks(
conn: &Connection,
space_id: &str,
query: &[f32],
limit: usize,
) -> Result<Vec<(String, String, String, f32)>> {
let mut stmt = conn.prepare(
"SELECT files.name, cache.file_chunks.location, cache.file_chunks.text,
cache.chunk_embeddings.vec
FROM cache.chunk_embeddings
JOIN files ON files.id = cache.chunk_embeddings.file_id
JOIN cache.file_chunks
ON cache.file_chunks.file_id = cache.chunk_embeddings.file_id
AND CAST(cache.file_chunks.seq AS INTEGER) = cache.chunk_embeddings.seq
WHERE files.space_id = ?1",
)?;
let rows = stmt.query_map([space_id], |r| {
Ok((
r.get::<_, String>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
r.get::<_, Vec<u8>>(3)?,
))
})?;
let mut hits: Vec<(String, String, String, f32)> = Vec::new();
for row in rows {
let (name, loc, text, blob) = row?;
let v = blob_to_vec(&blob);
if v.len() != query.len() {
continue;
}
let score = cosine(query, &v);
hits.push((name, loc, text, score));
}
hits.sort_by(|a, b| b.3.partial_cmp(&a.3).unwrap_or(std::cmp::Ordering::Equal));
hits.truncate(limit);
Ok(hits)
}
fn cosine(a: &[f32], b: &[f32]) -> f32 {
let (mut dot, mut na, mut nb) = (0.0f32, 0.0f32, 0.0f32);
for (x, y) in a.iter().zip(b) {
dot += x * y;
na += x * x;
nb += y * y;
}
let denom = na.sqrt() * nb.sqrt();
if denom == 0.0 { 0.0 } else { dot / denom }
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum UsageRange {
Day,
Week,
Month,
#[default]
All,
}
impl UsageRange {
pub const CYCLE: [Self; 4] = [Self::Day, Self::Week, Self::Month, Self::All];
pub const fn label(self) -> &'static str {
match self {
Self::Day => "24h",
Self::Week => "7d",
Self::Month => "30d",
Self::All => "all",
}
}
pub const fn title(self) -> &'static str {
match self {
Self::Day => "last 24 hours",
Self::Week => "last 7 days",
Self::Month => "last 30 days",
Self::All => "all time",
}
}
pub const fn key(self) -> &'static str {
match self {
Self::Day => "day",
Self::Week => "week",
Self::Month => "month",
Self::All => "all",
}
}
pub fn from_key(key: &str) -> Self {
Self::CYCLE
.iter()
.copied()
.find(|r| r.key() == key)
.unwrap_or_default()
}
#[must_use]
pub const fn next(self) -> Self {
match self {
Self::Day => Self::Week,
Self::Week => Self::Month,
Self::Month => Self::All,
Self::All => Self::Day,
}
}
#[must_use]
pub const fn prev(self) -> Self {
match self {
Self::Day => Self::All,
Self::Week => Self::Day,
Self::Month => Self::Week,
Self::All => Self::Month,
}
}
pub fn since(self) -> Option<chrono::DateTime<chrono::Utc>> {
use chrono::{Duration, Utc};
match self {
Self::Day => Some(Utc::now() - Duration::hours(24)),
Self::Week => Some(Utc::now() - Duration::days(7)),
Self::Month => Some(Utc::now() - Duration::days(30)),
Self::All => None,
}
}
pub const fn empty_message(self) -> &'static str {
match self {
Self::Day => "no usage in the last 24 hours — ←/→ for a wider window",
Self::Week => "no usage in the last 7 days — ←/→ for a wider window",
Self::Month => "no usage in the last 30 days — ←/→ for a wider window",
Self::All => "no usage logged yet — send a message first",
}
}
}
pub fn price_name(model: &str) -> &str {
let stripped = ["go:", "openai:", "codex:", "opencode:"]
.iter()
.find_map(|p| model.strip_prefix(p))
.unwrap_or(model);
stripped.rsplit('/').next().unwrap_or(stripped)
}
fn catalog_price<'a>(
prices: &'a std::collections::HashMap<String, ModelPricing>,
model: &str,
) -> Option<&'a ModelPricing> {
if let Some(price) = prices.get(model) {
return Some(price);
}
let name = price_name(model);
if name.is_empty() {
return None;
}
prices
.iter()
.filter(|(id, _)| {
id.strip_suffix(name)
.is_some_and(|rest| rest.ends_with('/'))
})
.min_by_key(|(id, _)| id.len()) .map(|(_, price)| price)
}
fn catalog_request_cost(
price: ModelPricing,
prompt_tokens: u64,
completion_tokens: u64,
cache_read_tokens: u64,
cache_creation_tokens: u64,
) -> f64 {
let reads = cache_read_tokens.min(prompt_tokens);
let writes = cache_creation_tokens.min(prompt_tokens - reads);
let ordinary = prompt_tokens - reads - writes;
let read_price = price.cache_read.unwrap_or(price.prompt);
let write_price = price.cache_write.unwrap_or(price.prompt);
(ordinary as f64 * price.prompt
+ reads as f64 * read_price
+ writes as f64 * write_price
+ completion_tokens as f64 * price.completion)
/ 1e6
}
struct CostBackfillRow {
id: i64,
backend: String,
model: String,
prompt_tokens: u64,
completion_tokens: u64,
cache_read_tokens: u64,
cache_creation_tokens: u64,
old_cost: Option<f64>,
cost_is_provider: Option<bool>,
}
#[derive(Default)]
pub struct UsageTotals {
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub cache_read_tokens: u64,
pub cache_creation_tokens: u64,
pub cost: f64,
}
#[derive(Default)]
pub struct UsageDay {
pub day: String,
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub cache_read_tokens: u64,
pub cost: f64,
}
#[derive(Default)]
pub struct UsageByBackend {
pub backend: String,
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub cache_read_tokens: u64,
pub cost: f64,
}
#[derive(Default)]
pub struct UsageByModel {
pub model: String,
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub cache_read_tokens: u64,
pub cost: f64,
}
#[derive(Default)]
pub struct UsageRow {
pub created_at: String,
pub backend: String,
pub model: String,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub cache_read_tokens: u64,
pub cost: Option<f64>,
}
impl Db {
#[allow(clippy::too_many_arguments)]
pub fn log_usage(
&self,
backend: &str,
model: &str,
prompt_tokens: u64,
completion_tokens: u64,
cache_read_tokens: u64,
cache_creation_tokens: u64,
cost: Option<f64>,
cost_is_provider: bool,
session_id: Option<&str>,
space_id: Option<&str>,
) -> Result<i64> {
self.conn.execute(
"INSERT INTO usage_log (sync_id, created_at, session_id, space_id, backend, model,
prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens,
cost, cost_is_provider, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
(
Uuid::new_v4().to_string(),
Utc::now().to_rfc3339(),
session_id,
space_id,
backend,
model,
prompt_tokens as i64,
completion_tokens as i64,
cache_read_tokens as i64,
cache_creation_tokens as i64,
cost,
i64::from(cost_is_provider),
Utc::now().to_rfc3339(),
),
)?;
Ok(self.conn.last_insert_rowid())
}
#[allow(clippy::too_many_arguments)]
pub fn update_usage(
&self,
row_id: i64,
prompt_tokens: u64,
completion_tokens: u64,
cache_read_tokens: u64,
cache_creation_tokens: u64,
cost: Option<f64>,
cost_is_provider: bool,
) -> Result<()> {
self.conn.execute(
"UPDATE usage_log SET prompt_tokens = ?1, completion_tokens = ?2,
cache_read_tokens = ?3, cache_creation_tokens = ?4, cost = ?5,
cost_is_provider = ?6, updated_at = ?7
WHERE id = ?8",
(
prompt_tokens as i64,
completion_tokens as i64,
cache_read_tokens as i64,
cache_creation_tokens as i64,
cost,
i64::from(cost_is_provider),
Utc::now().to_rfc3339(),
row_id,
),
)?;
Ok(())
}
pub fn request_cost(
&self,
model: &str,
prompt_tokens: u64,
completion_tokens: u64,
cache_read_tokens: u64,
cache_creation_tokens: u64,
) -> Option<f64> {
self.model_price(model).map(|price| {
catalog_request_cost(
price,
prompt_tokens,
completion_tokens,
cache_read_tokens,
cache_creation_tokens,
)
})
}
pub fn backfill_usage_costs(&mut self) -> Result<usize> {
let max_price: f64 = self.conn.query_row(
"SELECT COALESCE(MAX(prompt_price), 0) FROM cache.model_prices",
[],
|r| r.get(0),
)?;
if max_price > 0.0 && max_price < 0.001 {
self.conn.execute(
"UPDATE cache.model_prices SET
prompt_price = prompt_price * 1e6,
completion_price = completion_price * 1e6,
cache_read_price = cache_read_price * 1e6,
cache_write_price = cache_write_price * 1e6",
[],
)?;
}
let mut prices: std::collections::HashMap<String, ModelPricing> =
std::collections::HashMap::default();
{
let mut stmt = self.conn.prepare(
"SELECT model_id, prompt_price, completion_price,
cache_read_price, cache_write_price
FROM cache.model_prices",
)?;
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, String>(0)?,
ModelPricing {
prompt: r.get(1)?,
completion: r.get(2)?,
cache_read: r.get(3)?,
cache_write: r.get(4)?,
},
))
})?;
for row in rows {
let (model, price) = row?;
prices.insert(model, price);
}
}
let rows: Vec<CostBackfillRow> = {
let mut stmt = self.conn.prepare(
"SELECT id, backend, model, prompt_tokens, completion_tokens,
cache_read_tokens, cache_creation_tokens, cost,
cost_is_provider
FROM usage_log",
)?;
let rows = stmt.query_map([], |r| {
Ok(CostBackfillRow {
id: r.get(0)?,
backend: r.get(1)?,
model: r.get(2)?,
prompt_tokens: r.get::<_, i64>(3)? as u64,
completion_tokens: r.get::<_, i64>(4)? as u64,
cache_read_tokens: r.get::<_, i64>(5)? as u64,
cache_creation_tokens: r.get::<_, i64>(6)? as u64,
old_cost: r.get(7)?,
cost_is_provider: r.get::<_, Option<i64>>(8)?.map(|v| v != 0),
})
})?;
rows.collect::<rusqlite::Result<Vec<_>>>()?
};
let tx = self.conn.transaction()?;
{
let mut update = tx.prepare(
"UPDATE usage_log SET cost = ?2, cost_is_provider = 0, updated_at = ?3 WHERE id = ?1",
)?;
for row in &rows {
let legacy_opencode_cost = row.cost_is_provider.is_none()
&& row.backend == "OpenCode Go"
&& row.old_cost.is_some();
if row.cost_is_provider == Some(true) || legacy_opencode_cost {
continue;
}
let recomputed = catalog_price(&prices, &row.model).map(|price| {
catalog_request_cost(
*price,
row.prompt_tokens,
row.completion_tokens,
row.cache_read_tokens,
row.cache_creation_tokens,
)
});
let write = match (recomputed, row.old_cost) {
(Some(cost), None) => Some(cost),
(Some(cost), Some(old)) if (cost - old).abs() > 1e-12 => Some(cost),
_ => None,
};
if let Some(cost) = write {
update.execute((row.id, cost, Utc::now().to_rfc3339()))?;
}
}
}
tx.commit()?;
Ok(rows.len())
}
pub fn upsert_model_prices(&mut self, prices: &[(String, String, ModelPricing)]) -> Result<()> {
if prices.is_empty() {
return Ok(());
}
let tx = self.conn.transaction()?;
{
let mut stmt = tx.prepare(
"INSERT INTO cache.model_prices (model_id, backend, prompt_price, completion_price,
cache_read_price, cache_write_price, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
ON CONFLICT(model_id) DO UPDATE SET
prompt_price = excluded.prompt_price,
completion_price = excluded.completion_price,
cache_read_price = excluded.cache_read_price,
cache_write_price = excluded.cache_write_price,
updated_at = excluded.updated_at
WHERE cache.model_prices.prompt_price != excluded.prompt_price
OR cache.model_prices.completion_price != excluded.completion_price
OR cache.model_prices.cache_read_price IS NOT excluded.cache_read_price
OR cache.model_prices.cache_write_price IS NOT excluded.cache_write_price",
)?;
let now = Utc::now().to_rfc3339();
for (model, backend, price) in prices {
stmt.execute((
model,
backend,
price.prompt,
price.completion,
price.cache_read,
price.cache_write,
&now,
))?;
}
}
tx.commit()?;
Ok(())
}
pub fn model_price(&self, model: &str) -> Option<ModelPricing> {
let read_price = |r: &rusqlite::Row| {
Ok(ModelPricing {
prompt: r.get(0)?,
completion: r.get(1)?,
cache_read: r.get(2)?,
cache_write: r.get(3)?,
})
};
if let Ok(price) = self.conn.query_row(
"SELECT prompt_price, completion_price, cache_read_price, cache_write_price
FROM cache.model_prices WHERE model_id = ?1",
[model],
read_price,
) {
return Some(price);
}
let name = price_name(model);
if name.is_empty() {
return None;
}
self.conn
.query_row(
"SELECT prompt_price, completion_price, cache_read_price, cache_write_price
FROM cache.model_prices
WHERE backend = 'OpenRouter'
AND substr(model_id, -length(?1) - 1) = '/' || ?1
ORDER BY length(model_id) LIMIT 1",
[name],
read_price,
)
.ok()
}
pub fn usage_totals(&self, since: Option<&str>) -> Result<UsageTotals> {
let mut sql = String::from(
"SELECT COUNT(*),
COALESCE(SUM(prompt_tokens), 0),
COALESCE(SUM(completion_tokens), 0),
COALESCE(SUM(cache_read_tokens), 0),
COALESCE(SUM(cache_creation_tokens), 0),
COALESCE(SUM(cost), 0)
FROM usage_log",
);
if since.is_some() {
sql.push_str(" WHERE created_at >= ?1");
}
let map = |r: &rusqlite::Row| {
Ok(UsageTotals {
requests: r.get::<_, i64>(0)? as u64,
prompt_tokens: r.get::<_, i64>(1)? as u64,
completion_tokens: r.get::<_, i64>(2)? as u64,
cache_read_tokens: r.get::<_, i64>(3)? as u64,
cache_creation_tokens: r.get::<_, i64>(4)? as u64,
cost: r.get::<_, f64>(5)?,
})
};
let totals = match since {
Some(s) => self.conn.query_row(&sql, [s], map),
None => self.conn.query_row(&sql, [], map),
}?;
Ok(totals)
}
pub fn usage_by_backend(&self, since: Option<&str>) -> Result<Vec<UsageByBackend>> {
let mut sql = String::from(
"SELECT backend, COUNT(*),
COALESCE(SUM(prompt_tokens), 0),
COALESCE(SUM(completion_tokens), 0),
COALESCE(SUM(cache_read_tokens), 0),
COALESCE(SUM(cost), 0)
FROM usage_log",
);
if since.is_some() {
sql.push_str(" WHERE created_at >= ?1");
}
sql.push_str(" GROUP BY backend ORDER BY COUNT(*) DESC");
let map = |r: &rusqlite::Row| {
Ok(UsageByBackend {
backend: r.get(0)?,
requests: r.get::<_, i64>(1)? as u64,
prompt_tokens: r.get::<_, i64>(2)? as u64,
completion_tokens: r.get::<_, i64>(3)? as u64,
cache_read_tokens: r.get::<_, i64>(4)? as u64,
cost: r.get::<_, f64>(5)?,
})
};
let mut stmt = self.conn.prepare(&sql)?;
let rows = match since {
Some(s) => stmt.query_map([s], map),
None => stmt.query_map([], map),
}?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn usage_by_model(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageByModel>> {
let mut sql = String::from(
"SELECT model, COUNT(*),
COALESCE(SUM(prompt_tokens), 0),
COALESCE(SUM(completion_tokens), 0),
COALESCE(SUM(cache_read_tokens), 0),
COALESCE(SUM(cost), 0)
FROM usage_log",
);
if since.is_some() {
sql.push_str(" WHERE created_at >= ?1");
sql.push_str(" GROUP BY model ORDER BY COUNT(*) DESC LIMIT ?2");
} else {
sql.push_str(" GROUP BY model ORDER BY COUNT(*) DESC LIMIT ?1");
}
let map = |r: &rusqlite::Row| {
Ok(UsageByModel {
model: r.get(0)?,
requests: r.get::<_, i64>(1)? as u64,
prompt_tokens: r.get::<_, i64>(2)? as u64,
completion_tokens: r.get::<_, i64>(3)? as u64,
cache_read_tokens: r.get::<_, i64>(4)? as u64,
cost: r.get::<_, f64>(5)?,
})
};
let mut stmt = self.conn.prepare(&sql)?;
let rows = match since {
Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
None => stmt.query_map(rusqlite::params![limit as i64], map),
}?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn usage_recent(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageRow>> {
let mut sql = String::from(
"SELECT created_at, backend, model, prompt_tokens, completion_tokens,
cache_read_tokens, cost
FROM usage_log",
);
if since.is_some() {
sql.push_str(" WHERE created_at >= ?1");
sql.push_str(" ORDER BY id DESC LIMIT ?2");
} else {
sql.push_str(" ORDER BY id DESC LIMIT ?1");
}
let map = |r: &rusqlite::Row| {
Ok(UsageRow {
created_at: r.get(0)?,
backend: r.get(1)?,
model: r.get(2)?,
prompt_tokens: r.get::<_, i64>(3)? as u64,
completion_tokens: r.get::<_, i64>(4)? as u64,
cache_read_tokens: r.get::<_, i64>(5)? as u64,
cost: r.get(6)?,
})
};
let mut stmt = self.conn.prepare(&sql)?;
let rows = match since {
Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
None => stmt.query_map(rusqlite::params![limit as i64], map),
}?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn usage_by_day(&self, limit: u64, since: Option<&str>) -> Result<Vec<UsageDay>> {
let mut sql = String::from(
"SELECT substr(created_at, 1, 10) AS day, COUNT(*),
COALESCE(SUM(prompt_tokens), 0),
COALESCE(SUM(completion_tokens), 0),
COALESCE(SUM(cache_read_tokens), 0),
COALESCE(SUM(cost), 0)
FROM usage_log",
);
if since.is_some() {
sql.push_str(" WHERE created_at >= ?1");
sql.push_str(" GROUP BY day ORDER BY day DESC LIMIT ?2");
} else {
sql.push_str(" GROUP BY day ORDER BY day DESC LIMIT ?1");
}
let map = |r: &rusqlite::Row| {
Ok(UsageDay {
day: r.get(0)?,
requests: r.get::<_, i64>(1)? as u64,
prompt_tokens: r.get::<_, i64>(2)? as u64,
completion_tokens: r.get::<_, i64>(3)? as u64,
cache_read_tokens: r.get::<_, i64>(4)? as u64,
cost: r.get(5)?,
})
};
let mut stmt = self.conn.prepare(&sql)?;
let rows = match since {
Some(s) => stmt.query_map(rusqlite::params![s, limit as i64], map),
None => stmt.query_map(rusqlite::params![limit as i64], map),
}?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
}
#[allow(dead_code)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SyncState {
pub peer_id: String,
pub table_name: String,
pub pull_cursor: Option<String>,
pub push_cursor: Option<String>,
pub last_synced_at: Option<String>,
}
#[allow(dead_code)]
impl Db {
pub fn device_id(&self) -> Result<String> {
if let Some(id) = self
.conn
.query_row("SELECT device_id FROM device_meta LIMIT 1", [], |r| {
r.get(0)
})
.optional()?
{
return Ok(id);
}
let id = Uuid::new_v4().to_string();
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO device_meta (device_id, created_at) VALUES (?1, ?2)",
(&id, &now),
)?;
Ok(id)
}
pub fn set_sync_state(
&self,
peer_id: &str,
table_name: &str,
pull_cursor: Option<&str>,
push_cursor: Option<&str>,
) -> Result<()> {
self.conn.execute(
"INSERT INTO sync_state (peer_id, table_name, pull_cursor, push_cursor, last_synced_at)
VALUES (?1, ?2, ?3, ?4, ?5)
ON CONFLICT(peer_id, table_name) DO UPDATE SET
pull_cursor = COALESCE(?3, pull_cursor),
push_cursor = COALESCE(?4, push_cursor),
last_synced_at = ?5",
(
peer_id,
table_name,
pull_cursor,
push_cursor,
Utc::now().to_rfc3339(),
),
)?;
Ok(())
}
pub fn load_sync_state(&self) -> Result<Vec<SyncState>> {
let mut stmt = self.conn.prepare(
"SELECT peer_id, table_name, pull_cursor, push_cursor, last_synced_at
FROM sync_state ORDER BY peer_id, table_name",
)?;
let rows = stmt.query_map([], |r| {
Ok(SyncState {
peer_id: r.get(0)?,
table_name: r.get(1)?,
pull_cursor: r.get(2)?,
push_cursor: r.get(3)?,
last_synced_at: r.get(4)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
}
pub fn fts_quote(query: &str) -> String {
query
.split_whitespace()
.map(|t| format!("\"{}\"", t.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(" ")
}
pub fn search_chunks(
conn: &Connection,
space_id: &str,
query: &str,
limit: usize,
) -> Result<Vec<(String, String, String)>> {
let q = fts_quote(query);
if q.is_empty() {
return Ok(Vec::new());
}
let mut stmt = conn.prepare(
"SELECT files.name, cache.file_chunks.location,
snippet(file_chunks, 3, '', '', '…', 24)
FROM cache.file_chunks JOIN files ON files.id = cache.file_chunks.file_id
WHERE file_chunks MATCH ?1 AND files.space_id = ?2
ORDER BY bm25(file_chunks) LIMIT ?3",
)?;
let rows = stmt.query_map((q, space_id, limit as i64), |r| {
Ok((r.get(0)?, r.get(1)?, r.get(2)?))
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn file_text(conn: &Connection, space_id: &str, name: &str) -> Result<Option<String>> {
let mut stmt = conn.prepare(
"SELECT cache.file_chunks.text
FROM cache.file_chunks JOIN files ON files.id = cache.file_chunks.file_id
WHERE files.space_id = ?1 AND files.name = ?2
ORDER BY CAST(cache.file_chunks.seq AS INTEGER) ASC",
)?;
let rows = stmt.query_map((space_id, name), |r| r.get::<_, String>(0))?;
let parts = rows.collect::<rusqlite::Result<Vec<_>>>()?;
Ok((!parts.is_empty()).then(|| parts.join("\n")))
}
pub fn count_files(conn: &Connection, space_id: &str) -> Result<u64> {
let n: i64 = conn.query_row(
"SELECT COUNT(*) FROM files WHERE space_id = ?1",
[space_id],
|r| r.get(0),
)?;
Ok(n as u64)
}
#[cfg(test)]
mod tests {
use super::*;
fn price(prompt: f64, completion: f64) -> ModelPricing {
ModelPricing {
prompt,
completion,
cache_read: None,
cache_write: None,
}
}
fn cache_price(
prompt: f64,
completion: f64,
cache_read: f64,
cache_write: f64,
) -> ModelPricing {
ModelPricing {
prompt,
completion,
cache_read: Some(cache_read),
cache_write: Some(cache_write),
}
}
#[test]
fn is_fresh_true_under_24h_false_over() {
let now = Utc::now();
let recent = (now - chrono::Duration::hours(1)).to_rfc3339();
let stale = (now - chrono::Duration::hours(25)).to_rfc3339();
assert!(is_fresh(&recent, now));
assert!(!is_fresh(&stale, now));
assert!(!is_fresh("not a timestamp", now)); }
#[test]
fn usage_log_round_trips_and_aggregates() {
let mut db = Db::open_in_memory().unwrap();
db.upsert_model_prices(&[(
"anthropic/claude-3.5-sonnet".to_string(),
"OpenRouter".to_string(),
price(3.0, 15.0),
)])
.unwrap();
assert_eq!(
db.model_price("anthropic/claude-3.5-sonnet"),
Some(price(3.0, 15.0))
);
db.upsert_model_prices(&[(
"anthropic/claude-3.5-sonnet".to_string(),
"OpenRouter".to_string(),
cache_price(4.0, 16.0, 0.4, 5.0),
)])
.unwrap();
assert_eq!(
db.model_price("anthropic/claude-3.5-sonnet"),
Some(cache_price(4.0, 16.0, 0.4, 5.0))
);
assert_eq!(db.model_price("unknown/model"), None);
db.log_usage(
"OpenRouter",
"anthropic/claude-3.5-sonnet",
100,
10,
70,
20,
Some(0.00056),
true,
Some("s1"),
Some("space-a"),
)
.unwrap();
db.log_usage(
"Codex",
"gpt-5.1-codex",
50,
5,
0,
0,
None,
false,
None,
None,
)
.unwrap();
let totals = db.usage_totals(None).unwrap();
assert_eq!(totals.requests, 2);
assert_eq!(totals.prompt_tokens, 150);
assert_eq!(totals.completion_tokens, 15);
assert_eq!(totals.cache_read_tokens, 70);
assert_eq!(totals.cache_creation_tokens, 20);
assert!((totals.cost - 0.00056).abs() < 1e-9);
let by_backend = db.usage_by_backend(None).unwrap();
assert_eq!(by_backend.len(), 2);
assert_eq!(by_backend[0].backend, "OpenRouter"); assert_eq!(by_backend[0].requests, 1);
let by_model = db.usage_by_model(5, None).unwrap();
assert_eq!(by_model.len(), 2);
assert!(by_model.iter().any(|m| m.model == "gpt-5.1-codex"));
let recent = db.usage_recent(10, None).unwrap();
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].model, "gpt-5.1-codex"); assert_eq!(recent[0].cost, None);
assert_eq!(recent[1].cache_read_tokens, 70);
}
#[test]
fn usage_queries_filter_by_since_window() {
let db = Db::open_in_memory().unwrap();
let insert = |created: &str| {
db.raw().execute(
"INSERT INTO usage_log (sync_id, created_at, session_id, backend, model,
prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens, cost)
VALUES (?1, ?2, NULL, 'OpenRouter', 'a/model', 100, 10, 0, 0, 0.001)",
(uuid::Uuid::new_v4().to_string(), created),
)
};
insert("2026-01-01T00:00:00+00:00").unwrap();
insert("2026-01-02T00:00:00+00:00").unwrap();
insert("2026-01-03T00:00:00+00:00").unwrap();
let since = Some("2026-01-02T00:00:00+00:00");
let totals = db.usage_totals(since).unwrap();
assert_eq!(totals.requests, 2);
assert_eq!(totals.prompt_tokens, 200);
assert!((totals.cost - 0.002).abs() < 1e-12);
assert_eq!(db.usage_totals(None).unwrap().requests, 3);
assert_eq!(db.usage_by_backend(since).unwrap()[0].requests, 2);
assert_eq!(db.usage_by_model(5, since).unwrap()[0].requests, 2);
let recent = db.usage_recent(10, since).unwrap();
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].created_at, "2026-01-03T00:00:00+00:00");
assert!(
recent
.iter()
.any(|r| r.created_at == "2026-01-02T00:00:00+00:00")
);
}
#[test]
fn usage_range_cycles_and_persists() {
use crate::db::UsageRange;
assert_eq!(UsageRange::Day.next(), UsageRange::Week);
assert_eq!(UsageRange::All.next(), UsageRange::Day);
assert_eq!(UsageRange::Day.prev(), UsageRange::All);
assert_eq!(UsageRange::from_key("month"), UsageRange::Month);
assert_eq!(UsageRange::from_key("bogus"), UsageRange::All);
assert_eq!(UsageRange::Week.key(), "week");
assert_eq!(UsageRange::Day.title(), "last 24 hours");
assert!(UsageRange::All.since().is_none());
assert!(UsageRange::Day.since().is_some());
}
#[test]
fn backfill_usage_costs_recomputes_history_from_catalog() {
let mut db = Db::open_in_memory().unwrap();
db.upsert_model_prices(&[(
"anthropic/claude-3.5-sonnet".to_string(),
"OpenRouter".to_string(),
cache_price(3.0, 15.0, 0.3, 3.75),
)])
.unwrap();
db.log_usage(
"OpenRouter",
"anthropic/claude-3.5-sonnet",
100,
10,
70,
20,
None,
false,
None,
None,
)
.unwrap();
db.log_usage(
"Codex",
"gpt-5.1-codex",
50,
5,
0,
0,
None,
false,
None,
None,
)
.unwrap();
let visited = db.backfill_usage_costs().unwrap();
assert_eq!(visited, 2);
let totals = db.usage_totals(None).unwrap();
assert!((totals.cost - 0.000_276).abs() < 1e-12);
let recent = db.usage_recent(10, None).unwrap();
assert!((recent[1].cost.unwrap() - 0.000_276).abs() < 1e-12); assert_eq!(recent[0].cost, None);
db.backfill_usage_costs().unwrap();
assert!((db.usage_totals(None).unwrap().cost - 0.000_276).abs() < 1e-12);
}
#[test]
fn backfill_preserves_provider_reported_cost() {
let mut db = Db::open_in_memory().unwrap();
db.upsert_model_prices(&[(
"anthropic/claude-3.5-sonnet".to_string(),
"OpenRouter".to_string(),
cache_price(3.0, 15.0, 0.3, 3.75),
)])
.unwrap();
db.log_usage(
"OpenRouter",
"anthropic/claude-3.5-sonnet",
100,
10,
70,
20,
Some(0.000_321),
true,
None,
None,
)
.unwrap();
db.backfill_usage_costs().unwrap();
assert_eq!(db.usage_recent(1, None).unwrap()[0].cost, Some(0.000_321));
}
#[test]
fn backfill_usage_costs_heals_per_token_catalog() {
let mut db = Db::open_in_memory().unwrap();
db.upsert_model_prices(&[(
"deepseek/deepseek-v4-flash-0731".to_string(),
"OpenRouter".to_string(),
price(8e-08, 1.8e-07),
)])
.unwrap();
db.log_usage(
"OpenRouter",
"deepseek/deepseek-v4-flash-0731",
122_221,
672,
118_784,
0,
Some(9.89864e-09), false,
None,
None,
)
.unwrap();
db.backfill_usage_costs().unwrap();
assert_eq!(
db.model_price("deepseek/deepseek-v4-flash-0731"),
Some(price(0.08, 0.18))
);
let recent = db.usage_recent(10, None).unwrap();
let cost = recent[0].cost.unwrap();
assert!((cost - 0.009_898_6).abs() < 1e-6, "cost was {cost}");
assert!(cost > 0.009, "cost was {cost}");
}
#[test]
fn request_cost_prices_tokens_against_catalog() {
let mut db = Db::open_in_memory().unwrap();
assert_eq!(db.request_cost("unknown/model", 100, 10, 0, 0), None);
db.upsert_model_prices(&[(
"anthropic/claude-3.5-sonnet".to_string(),
"OpenRouter".to_string(),
cache_price(3.0, 15.0, 0.3, 3.75),
)])
.unwrap();
let cost = db
.request_cost("anthropic/claude-3.5-sonnet", 100, 10, 70, 20)
.unwrap();
assert!((cost - 0.000_276).abs() < 1e-12);
}
#[test]
fn model_price_cross_references_openrouter_catalog_twins() {
let mut db = Db::open_in_memory().unwrap();
db.upsert_model_prices(&[
(
"deepseek/deepseek-v4-flash".to_string(),
"OpenRouter".to_string(),
cache_price(0.08, 0.18, 0.016, 0.08),
),
(
"openai/gpt-5".to_string(),
"OpenRouter".to_string(),
cache_price(1.25, 10.0, 0.125, 1.25),
),
])
.unwrap();
assert_eq!(
db.model_price("deepseek/deepseek-v4-flash"),
Some(cache_price(0.08, 0.18, 0.016, 0.08))
);
assert_eq!(
db.model_price("go:deepseek-v4-flash"),
Some(cache_price(0.08, 0.18, 0.016, 0.08))
);
assert_eq!(
db.model_price("deepseek-v4-flash"),
Some(cache_price(0.08, 0.18, 0.016, 0.08))
);
assert_eq!(
db.model_price("openai:gpt-5"),
Some(cache_price(1.25, 10.0, 0.125, 1.25))
);
assert_eq!(
db.model_price("codex:gpt-5"),
Some(cache_price(1.25, 10.0, 0.125, 1.25))
);
assert_eq!(db.model_price("no-such-model-anywhere"), None);
let cost = db
.request_cost("go:deepseek-v4-flash", 100, 10, 70, 0)
.unwrap();
assert!((cost - 0.000_005_32).abs() < 1e-15);
}
#[test]
fn price_name_strips_backend_prefixes_and_vendors() {
assert_eq!(price_name("go:deepseek-v4-flash"), "deepseek-v4-flash");
assert_eq!(price_name("openai:gpt-5"), "gpt-5");
assert_eq!(price_name("codex:gpt-5.1-codex"), "gpt-5.1-codex");
assert_eq!(price_name("opencode:qwen3.6-plus"), "qwen3.6-plus");
assert_eq!(
price_name("deepseek/deepseek-v4-flash"),
"deepseek-v4-flash"
);
assert_eq!(price_name("gpt-5"), "gpt-5");
}
#[test]
fn backfill_prices_non_openrouter_models_via_catalog_twins() {
let mut db = Db::open_in_memory().unwrap();
db.upsert_model_prices(&[(
"deepseek/deepseek-v4-flash".to_string(),
"OpenRouter".to_string(),
cache_price(0.08, 0.18, 0.016, 0.08),
)])
.unwrap();
db.log_usage(
"OpenCode Go",
"go:deepseek-v4-flash",
100,
10,
0,
0,
None,
false,
None,
None,
)
.unwrap();
db.backfill_usage_costs().unwrap();
let recent = db.usage_recent(10, None).unwrap();
let cost = recent[0].cost.unwrap();
assert!((cost - 0.000_009_8).abs() < 1e-15, "cost was {cost}");
assert!((db.usage_totals(None).unwrap().cost - 0.000_009_8).abs() < 1e-15);
}
#[test]
fn web_cache_roundtrips_and_updates_on_rewrite() {
let db = Db::open_in_memory().unwrap();
assert!(cache_get(db.raw(), "example.com/a").unwrap().is_none());
cache_put(
db.raw(),
"example.com/a",
"https://example.com/a",
Some("Title"),
"body text",
)
.unwrap();
let (title, text, fetched_at) = cache_get(db.raw(), "example.com/a").unwrap().unwrap();
assert_eq!(title, "Title");
assert_eq!(text, "body text");
assert!(!fetched_at.is_empty());
cache_put(
db.raw(),
"example.com/a",
"https://example.com/a",
None,
"new body",
)
.unwrap();
let (title, text, _) = cache_get(db.raw(), "example.com/a").unwrap().unwrap();
assert_eq!(title, "");
assert_eq!(text, "new body");
}
#[test]
fn web_mode_defaults_off_and_toggles() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
assert!(!s.web_mode);
db.set_session_web_mode(&s.id, true).unwrap();
assert!(db.list_sessions(&space).unwrap()[0].web_mode);
}
#[test]
fn swarm_mode_defaults_off_and_toggles() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
assert!(!s.swarm_mode);
db.set_session_swarm_mode(&s.id, true).unwrap();
assert!(db.list_sessions(&space).unwrap()[0].swarm_mode);
assert!(db.get_session(&s.id).unwrap().unwrap().swarm_mode);
}
#[test]
fn swarm_personas_roundtrip_and_replace_all_on_save() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
assert!(db.list_swarm_personas(&s.id).unwrap().is_empty());
let roster = vec![
Persona {
name: "Skeptic".into(),
model: "a/one".into(),
blurb: "pokes holes".into(),
},
Persona {
name: "Advocate".into(),
model: "b/two".into(),
blurb: "user-first".into(),
},
];
db.save_swarm_personas(&s.id, &roster).unwrap();
let loaded = db.list_swarm_personas(&s.id).unwrap();
assert_eq!(loaded.len(), 2);
assert_eq!(loaded[0].name, "Skeptic");
assert_eq!(loaded[1].name, "Advocate");
db.save_swarm_personas(&s.id, &roster[..1]).unwrap();
assert_eq!(db.list_swarm_personas(&s.id).unwrap().len(), 1);
}
#[test]
fn persona_message_tags_role_assistant_with_persona_and_model() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
db.add_persona_message(&s.id, "reply text", "Skeptic", "a/one")
.unwrap();
let msgs = db.load_messages(&s.id).unwrap();
assert_eq!(msgs.len(), 1);
assert_eq!(msgs[0].role, "assistant");
assert_eq!(msgs[0].persona.as_deref(), Some("Skeptic"));
assert_eq!(msgs[0].model.as_deref(), Some("a/one"));
db.add_assistant_message(&s.id, "final answer", None, None, None, None, None, None)
.unwrap();
let msgs = db.load_messages(&s.id).unwrap();
assert_eq!(msgs[1].persona, None);
}
#[test]
fn session_sources_link_to_the_web_cache_and_are_keyword_searchable() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
cache_put(
db.raw(),
"https://example.com/a",
"https://example.com/a",
Some("A"),
"rust borrow checker deep dive",
)
.unwrap();
cache_put(
db.raw(),
"https://example.com/b",
"https://example.com/b",
Some("B"),
"cooking pasta recipes",
)
.unwrap();
db.add_session_sources(
&s.id,
&[
"https://example.com/a".to_string(),
"https://example.com/b".to_string(),
],
)
.unwrap();
let hits = db.search_session_sources(&s.id, "borrow checker").unwrap();
assert_eq!(hits.len(), 1);
assert!(hits[0].1.contains("borrow checker"));
assert!(
db.search_session_sources(&s.id, "quantum")
.unwrap()
.is_empty()
);
}
#[test]
fn set_source_flag_pins_and_discards_then_clears() {
let db = Db::open_in_memory().unwrap();
let session_id = "sess-1";
add_session_sources(&db.conn, session_id, &["https://a.example/x".to_string()]).unwrap();
db.set_source_flag(session_id, "https://a.example/x", Some("pinned"))
.unwrap();
assert_eq!(
pinned_urls(&db.conn, session_id).unwrap(),
vec!["https://a.example/x".to_string()]
);
assert!(discarded_domains(&db.conn, session_id).unwrap().is_empty());
db.set_source_flag(session_id, "https://a.example/x", Some("discarded"))
.unwrap();
assert!(pinned_urls(&db.conn, session_id).unwrap().is_empty());
assert_eq!(
discarded_domains(&db.conn, session_id).unwrap(),
vec!["a.example".to_string()]
);
db.set_source_flag(session_id, "https://a.example/x", None)
.unwrap();
assert!(discarded_domains(&db.conn, session_id).unwrap().is_empty());
}
#[test]
fn upsert_research_stage_message_replaces_the_same_labels_row() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
db.upsert_research_stage_message(&s.id, "searching", "round 1, 1/3")
.unwrap();
db.upsert_research_stage_message(&s.id, "searching", "round 1, 2/3")
.unwrap();
db.upsert_research_stage_message(&s.id, "planning", "")
.unwrap();
let msgs = db.load_messages(&s.id).unwrap();
let searching: Vec<_> = msgs
.iter()
.filter(|m| m.content.starts_with("searching:"))
.collect();
assert_eq!(searching.len(), 1, "expected one row, updated in place");
assert!(searching[0].content.contains("2/3"));
assert_eq!(msgs.iter().filter(|m| m.content == "planning").count(), 1);
}
#[test]
fn session_and_message_roundtrip() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db
.create_session("hello", "openai/gpt-4o", &space, "chat")
.unwrap();
db.add_user_message(&s.id, "hi").unwrap();
db.add_assistant_message(
&s.id,
"hello there",
Some("openai/gpt-4o"),
Some("let me think"),
Some(3),
Some(1.5),
Some(0.0042),
Some("Vibed"),
)
.unwrap();
let msgs = db.load_messages(&s.id).unwrap();
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[0].role, "user");
assert_eq!(msgs[1].content, "hello there");
assert_eq!(msgs[1].model.as_deref(), Some("openai/gpt-4o"));
assert_eq!(msgs[1].reasoning.as_deref(), Some("let me think"));
assert_eq!(msgs[1].tokens, Some(3));
assert_eq!(msgs[1].cost, Some(0.0042));
let sessions = db.list_sessions(&space).unwrap();
assert_eq!(sessions.len(), 1);
assert_eq!(sessions[0].id, s.id);
}
#[test]
fn markdown_images_in_content_roundtrip() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
let content = "look at  and ";
db.add_user_message(&s.id, content).unwrap();
let msgs = db.load_messages(&s.id).unwrap();
assert_eq!(msgs.len(), 1);
assert!(msgs[0].content.contains(""));
assert!(msgs[0].content.contains(""));
}
#[test]
fn model_prefs_toggle_and_used() {
let db = Db::open_in_memory().unwrap();
assert!(db.toggle_favorite("a/one").unwrap()); assert!(!db.toggle_favorite("a/one").unwrap()); db.mark_model_used("a/one").unwrap();
db.set_reasoning("a/one", Some("high")).unwrap();
let prefs = db.load_model_prefs().unwrap();
let p = &prefs[0];
assert_eq!(p.id, "a/one");
assert!(!p.favorite);
assert!(p.last_used.is_some());
assert_eq!(p.reasoning.as_deref(), Some("high"));
}
#[test]
fn settings_roundtrip() {
let db = Db::open_in_memory().unwrap();
db.set_setting("temperature", "0.7").unwrap();
db.set_setting("temperature", "0.9").unwrap(); let s = db.load_settings().unwrap();
assert_eq!(s, vec![("temperature".to_string(), "0.9".to_string())]);
}
#[test]
fn set_model_updates_row() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
db.set_session_model(&s.id, "c/d").unwrap();
assert_eq!(db.list_sessions(&space).unwrap()[0].model, "c/d");
}
#[test]
fn compaction_persists_and_roundtrips() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
assert_eq!(s.compact_summary, None);
assert_eq!(s.compact_through, 0);
db.set_compaction(&s.id, "digest of earlier turns", 6)
.unwrap();
let reloaded = &db.list_sessions(&space).unwrap()[0];
assert_eq!(
reloaded.compact_summary.as_deref(),
Some("digest of earlier turns")
);
assert_eq!(reloaded.compact_through, 6);
}
#[test]
fn spaces_crud_and_session_reassignment_on_delete() {
let db = Db::open_in_memory().unwrap();
let spaces = db.list_spaces().unwrap();
assert_eq!(spaces.len(), 1);
assert_eq!(spaces[0].name, DEFAULT_SPACE);
let work = db.create_space("work").unwrap();
let s = db.create_session("hi", "a/b", &work.id, "chat").unwrap();
assert_eq!(db.count_sessions(&work.id).unwrap(), 1);
db.rename_space(&work.id, "work-renamed").unwrap();
assert!(
db.list_spaces()
.unwrap()
.iter()
.any(|s| s.name == "work-renamed")
);
db.delete_space(&work.id).unwrap();
assert_eq!(db.list_spaces().unwrap().len(), 1); let default_id = db.default_space_id().unwrap();
let moved = db.list_sessions(&default_id).unwrap();
assert!(moved.iter().any(|c| c.id == s.id)); }
#[test]
fn chunk_embeddings_store_rank_and_invalidate() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let id = db.upsert_file(&space, "book.pdf", "h1", 10, "ok").unwrap();
db.set_file_chunks(
&id,
&[
("page 1".into(), "cooking with fire".into()),
("page 2".into(), "quantum entanglement".into()),
],
)
.unwrap();
let v = vec![0.25f32, -1.0, 3.5];
assert_eq!(blob_to_vec(&vec_to_blob(&v)), v);
assert_eq!(
files_missing_embeddings(&db.conn, &space).unwrap(),
vec![id.clone()]
);
db.set_chunk_embeddings(&id, &[(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])])
.unwrap();
assert!(
files_missing_embeddings(&db.conn, &space)
.unwrap()
.is_empty()
);
let hits = semantic_chunks(&db.conn, &space, &[0.1, 0.9], 5).unwrap();
assert_eq!(hits[0].1, "page 2");
assert!(hits[0].2.contains("quantum"));
assert!(hits[0].3 > hits[1].3, "scores must be descending");
let hits = semantic_chunks(&db.conn, &space, &[1.0, 0.0, 0.0], 5).unwrap();
assert!(hits.is_empty());
db.set_file_chunks(&id, &[("page 1".into(), "new text".into())])
.unwrap();
assert_eq!(
files_missing_embeddings(&db.conn, &space).unwrap(),
vec![id.clone()]
);
}
#[test]
fn files_upsert_list_delete_roundtrip() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let id = db.upsert_file(&space, "notes.md", "h1", 10, "ok").unwrap();
db.set_file_chunks(&id, &[("lines 1-40".into(), "hello fts world".into())])
.unwrap();
let files = db.list_files(&space).unwrap();
assert_eq!(files.len(), 1);
assert_eq!(files[0].name, "notes.md");
assert_eq!(files[0].hash, "h1");
assert_eq!(files[0].status, "ok");
let id2 = db.upsert_file(&space, "notes.md", "h2", 12, "ok").unwrap();
db.set_file_chunks(&id2, &[("lines 1-40".into(), "goodbye".into())])
.unwrap();
let files = db.list_files(&space).unwrap();
assert_eq!(files.len(), 1);
assert_eq!(files[0].hash, "h2");
db.delete_file(&files[0].id).unwrap();
assert!(db.list_files(&space).unwrap().is_empty());
}
#[test]
fn chunk_search_ranks_and_scopes_by_space() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let other = db.create_space("other").unwrap();
let a = db.upsert_file(&space, "a.md", "h", 1, "ok").unwrap();
let b = db.upsert_file(&other.id, "b.md", "h", 1, "ok").unwrap();
db.set_file_chunks(&a, &[("lines 1-40".into(), "rust borrow checker".into())])
.unwrap();
db.set_file_chunks(&b, &[("lines 1-40".into(), "rust in other space".into())])
.unwrap();
let hits = search_chunks(&db.conn, &space, "rust", 8).unwrap();
assert_eq!(hits.len(), 1); assert_eq!(hits[0].0, "a.md");
assert_eq!(hits[0].1, "lines 1-40");
assert!(hits[0].2.contains("rust"));
assert!(search_chunks(&db.conn, &space, "c++ \"quoted\" -dash", 8).is_ok());
}
#[test]
fn file_text_joins_chunks_in_order() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let id = db.upsert_file(&space, "doc.txt", "h", 1, "ok").unwrap();
db.set_file_chunks(
&id,
&[
("lines 1-2".into(), "one\ntwo".into()),
("lines 3-4".into(), "three\nfour".into()),
],
)
.unwrap();
let text = file_text(&db.conn, &space, "doc.txt").unwrap().unwrap();
assert_eq!(text, "one\ntwo\nthree\nfour");
assert!(
file_text(&db.conn, &space, "missing.txt")
.unwrap()
.is_none()
);
assert_eq!(count_files(&db.conn, &space).unwrap(), 1);
}
#[test]
fn research_stage_messages_round_trip() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
db.add_research_stage_message(&s.id, "planning…").unwrap();
let msgs = db.load_messages(&s.id).unwrap();
assert_eq!(msgs.last().unwrap().role, "research_stage");
assert_eq!(msgs.last().unwrap().content, "planning…");
}
#[test]
fn survey_messages_round_trip() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
db.add_survey_message(&s.id, "For \"topic\":\n 1. Depth or breadth?")
.unwrap();
let msgs = db.load_messages(&s.id).unwrap();
let last = msgs.last().unwrap();
assert_eq!(last.role, "survey");
assert!(last.content.contains("Depth or breadth?"));
}
#[test]
fn gate_reply_round_trip() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let s = db.create_session("t", "a/b", &space, "chat").unwrap();
db.add_gate_reply_message(&s.id, "the second option")
.unwrap();
let msgs = db.load_messages(&s.id).unwrap();
let last = msgs.last().unwrap();
assert_eq!(last.role, "gate_reply");
assert_eq!(last.content, "the second option");
}
#[test]
fn create_list_touch_delete_watch_roundtrip() {
let db = Db::open_in_memory().unwrap();
let id = db
.create_watch("space-1", "rust async runtimes", 24, "sess-1")
.unwrap();
let watches = db.list_watches("space-1").unwrap();
assert_eq!(watches.len(), 1);
assert_eq!(watches[0].topic, "rust async runtimes");
assert_eq!(watches[0].interval_hours, 24);
assert!(watches[0].last_run_at.is_none());
db.touch_watch(&id, "2026-07-07T00:00:00+00:00").unwrap();
let watches = db.list_watches("space-1").unwrap();
assert_eq!(
watches[0].last_run_at.as_deref(),
Some("2026-07-07T00:00:00+00:00")
);
db.delete_watch(&id).unwrap();
assert!(db.list_watches("space-1").unwrap().is_empty());
}
#[test]
#[allow(clippy::similar_names)]
fn list_all_watches_returns_watches_from_all_spaces() {
let db = Db::open_in_memory().unwrap();
let id1 = db.create_watch("space-a", "topic-1", 24, "sess-1").unwrap();
let id2 = db.create_watch("space-b", "topic-2", 48, "sess-2").unwrap();
let id3 = db.create_watch("space-a", "topic-3", 12, "sess-3").unwrap();
let all_watches = db.list_all_watches().unwrap();
assert_eq!(all_watches.len(), 3);
assert!(
all_watches
.iter()
.any(|w| w.id == id1 && w.space_id == "space-a")
);
assert!(
all_watches
.iter()
.any(|w| w.id == id2 && w.space_id == "space-b")
);
assert!(
all_watches
.iter()
.any(|w| w.id == id3 && w.space_id == "space-a")
);
let space_a_watches = db.list_watches("space-a").unwrap();
assert_eq!(space_a_watches.len(), 2);
assert!(space_a_watches.iter().all(|w| w.space_id == "space-a"));
let space_b_watches = db.list_watches("space-b").unwrap();
assert_eq!(space_b_watches.len(), 1);
assert!(space_b_watches.iter().all(|w| w.space_id == "space-b"));
}
fn legacy_db(path: &std::path::Path, now: &str) {
let conn = rusqlite::Connection::open(path).unwrap();
conn.execute_batch(
"CREATE TABLE sessions (id TEXT PRIMARY KEY, title TEXT NOT NULL,
model TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL);
CREATE TABLE messages (id TEXT PRIMARY KEY, session_id TEXT NOT NULL,
role TEXT NOT NULL, content TEXT NOT NULL, created_at TEXT NOT NULL);
CREATE TABLE model_prefs (id TEXT PRIMARY KEY,
favorite INTEGER NOT NULL DEFAULT 0, last_used TEXT);
CREATE TABLE app_settings (key TEXT PRIMARY KEY, value TEXT NOT NULL);
CREATE TABLE spaces (id TEXT PRIMARY KEY, name TEXT NOT NULL UNIQUE,
created_at TEXT NOT NULL);
CREATE TABLE files (id TEXT PRIMARY KEY, space_id TEXT NOT NULL,
name TEXT NOT NULL, hash TEXT NOT NULL, size INTEGER NOT NULL,
status TEXT NOT NULL, created_at TEXT NOT NULL,
UNIQUE(space_id, name));
ALTER TABLE files ADD COLUMN mtime INTEGER NOT NULL DEFAULT 0;
CREATE VIRTUAL TABLE file_chunks USING fts5(
file_id UNINDEXED, seq UNINDEXED, location UNINDEXED, text);
CREATE TABLE chunk_embeddings (file_id TEXT NOT NULL, seq INTEGER NOT NULL,
vec BLOB NOT NULL, PRIMARY KEY (file_id, seq));
CREATE TABLE web_cache (url_norm TEXT PRIMARY KEY, url TEXT NOT NULL,
title TEXT, text TEXT NOT NULL, fetched_at TEXT NOT NULL);
CREATE TABLE citations (id INTEGER PRIMARY KEY AUTOINCREMENT,
space_id TEXT NOT NULL, report_file TEXT NOT NULL,
url TEXT NOT NULL, title TEXT);
CREATE TABLE session_sources (session_id TEXT NOT NULL,
url_norm TEXT NOT NULL, PRIMARY KEY (session_id, url_norm));
CREATE TABLE watches (id TEXT PRIMARY KEY, space_id TEXT NOT NULL,
topic TEXT NOT NULL, interval_hours INTEGER NOT NULL,
session_id TEXT NOT NULL, last_run_at TEXT);
CREATE TABLE swarm_personas (session_id TEXT NOT NULL, ord INTEGER NOT NULL,
name TEXT NOT NULL, model TEXT NOT NULL, persona TEXT NOT NULL);
CREATE TABLE usage_log (id INTEGER PRIMARY KEY AUTOINCREMENT,
created_at TEXT NOT NULL, session_id TEXT, space_id TEXT,
backend TEXT NOT NULL, model TEXT NOT NULL,
prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL,
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_tokens INTEGER NOT NULL DEFAULT 0, cost REAL);
CREATE TABLE model_prices (model_id TEXT PRIMARY KEY, backend TEXT NOT NULL,
prompt_price REAL NOT NULL, completion_price REAL NOT NULL,
cache_read_price REAL, cache_write_price REAL, updated_at TEXT NOT NULL);",
)
.unwrap();
conn.execute(
"INSERT INTO spaces (id, name, created_at) VALUES ('sp', 'default', ?1)",
[now],
)
.unwrap();
conn.execute(
"INSERT INTO files (id, space_id, name, hash, size, status, created_at, mtime)
VALUES ('f1', 'sp', 'a.txt', 'h1', 10, 'ok', ?1, 1234)",
[now],
)
.unwrap();
conn.execute(
"INSERT INTO file_chunks (file_id, seq, location, text)
VALUES ('f1', 0, 'l', 'hello world')",
[],
)
.unwrap();
conn.execute(
"INSERT INTO chunk_embeddings (file_id, seq, vec) VALUES ('f1', 0, ?1)",
[vec_to_blob(&[1.0, 0.0])],
)
.unwrap();
conn.execute(
"INSERT INTO web_cache (url_norm, url, title, text, fetched_at)
VALUES ('https://x.test/', 'https://x.test/', NULL, 'cached body', ?1)",
[now],
)
.unwrap();
conn.execute(
"INSERT INTO model_prices (model_id, backend, prompt_price, completion_price, updated_at)
VALUES ('a/model', 'OpenRouter', 1.0, 2.0, ?1)",
[now],
)
.unwrap();
conn.execute(
"INSERT INTO usage_log (created_at, backend, model, prompt_tokens, completion_tokens)
VALUES (?1, 'OpenRouter', 'a/model', 100, 10)",
[now],
)
.unwrap();
drop(conn);
}
#[test]
#[allow(clippy::float_cmp)]
fn legacy_db_migrates_cache_tables_and_seeds_file_index_state() {
let dir = std::env::temp_dir().join(format!("nexus-migrate-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).unwrap();
let db_path = dir.join("nexus.db");
legacy_db(&db_path, &Utc::now().to_rfc3339());
let mut db = Db::open(&db_path).unwrap();
let v: i64 = db
.raw()
.query_row("PRAGMA user_version", [], |r| r.get(0))
.unwrap();
assert_eq!(v, SCHEMA_VERSION);
assert!(has_column(db.raw(), "files", "mtime").unwrap());
assert_eq!(db.default_space_id().unwrap(), DEFAULT_SPACE);
let files_space: String = db
.raw()
.query_row("SELECT space_id FROM files WHERE id = 'f1'", [], |r| {
r.get(0)
})
.unwrap();
assert_eq!(files_space, DEFAULT_SPACE);
let (mtime, status) = db
.raw()
.query_row(
"SELECT mtime, status FROM cache.file_index_state WHERE file_id = 'f1'",
[],
|r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)),
)
.unwrap();
assert_eq!((mtime, status.as_str()), (1234, "ok"));
assert!(cache_path_for(&db_path).is_file());
let space = db.default_space_id().unwrap();
let hits = search_chunks(db.raw(), &space, "hello", 8).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(
file_text(db.raw(), &space, "a.txt").unwrap().as_deref(),
Some("hello world")
);
let (_, text, _) = cache_get(db.raw(), "https://x.test/").unwrap().unwrap();
assert_eq!(text, "cached body");
let hits = semantic_chunks(db.raw(), &space, &[1.0, 0.0], 4).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].0, "a.txt");
assert_eq!(db.model_price("a/model").unwrap().prompt, 1.0);
assert_eq!(db.backfill_usage_costs().unwrap(), 1);
let sync: String = db
.raw()
.query_row("SELECT sync_id FROM usage_log WHERE id = 1", [], |r| {
r.get(0)
})
.unwrap();
assert!(!sync.is_empty());
}
#[test]
fn fresh_db_creates_sibling_cache_db_with_split_tables() {
let dir = std::env::temp_dir().join(format!("nexus-fresh-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).unwrap();
let db_path = dir.join("nexus.db");
let db = Db::open(&db_path).unwrap();
let space = db.default_space_id().unwrap();
let id = db.upsert_file(&space, "doc.txt", "h", 1, "ok").unwrap();
db.set_file_chunks(&id, &[("l".into(), "needle text".into())])
.unwrap();
assert!(!has_column(db.raw(), "web_cache", "url_norm").unwrap());
assert!(!has_column(db.raw(), "files", "mtime").unwrap());
let cache_path = cache_path_for(&db_path);
assert!(cache_path.is_file());
let conn = rusqlite::Connection::open(&cache_path).unwrap();
assert!(cache_get(&conn, "x").unwrap().is_none());
drop(conn);
assert_eq!(
file_text(db.raw(), &space, "doc.txt").unwrap().as_deref(),
Some("needle text")
);
assert_eq!(db.list_files(&space).unwrap()[0].status, "ok");
}
#[allow(clippy::too_many_lines)]
#[test]
fn every_mutable_table_bumps_updated_at_on_mutation() {
type Mutator<'a> = &'a dyn Fn(&Db) -> Result<()>;
let mut db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let sid = db.create_session("t", "a/b", &space, "chat").unwrap().id;
let read = |table: &str, id: &str| -> String {
std::thread::sleep(std::time::Duration::from_millis(2));
db.raw()
.query_row(
&format!("SELECT updated_at FROM {table} WHERE id = ?1"),
[id],
|r| r.get::<_, String>(0),
)
.unwrap()
};
let mutators: [Mutator<'_>; 6] = [
&|db: &Db| db.set_compaction(&sid, "sum", 3),
&|db: &Db| db.set_session_web_mode(&sid, true),
&|db: &Db| db.set_session_swarm_mode(&sid, true),
&|db: &Db| db.set_session_title(&sid, "new", Some("new-slug")),
&|db: &Db| db.set_session_model(&sid, "m/x"),
&|db: &Db| db.set_research_parent(&sid, "parent"),
];
for mutate in mutators {
let before = read("sessions", &sid);
mutate(&db).unwrap();
assert!(read("sessions", &sid) > before);
}
let before = read("sessions", &sid);
db.save_swarm_personas(
&sid,
&[Persona {
name: "a".into(),
model: "m".into(),
blurb: "b".into(),
}],
)
.unwrap();
assert!(read("sessions", &sid) > before);
assert!(db.toggle_favorite("a/model").unwrap());
let before = read("model_prefs", "a/model");
db.set_reasoning("a/model", Some("high")).unwrap();
assert!(read("model_prefs", "a/model") > before);
let before = read("model_prefs", "a/model");
db.mark_model_used("a/model").unwrap();
assert!(read("model_prefs", "a/model") > before);
let read_key = |key: &str| -> String {
std::thread::sleep(std::time::Duration::from_millis(2));
db.raw()
.query_row(
"SELECT updated_at FROM app_settings WHERE key = ?1",
[key],
|r| r.get::<_, String>(0),
)
.unwrap()
};
db.set_setting("temperature", "0.5").unwrap();
let before = read_key("temperature");
db.set_setting("temperature", "0.9").unwrap();
assert!(read_key("temperature") > before);
let sp = db.create_space("other").unwrap();
let before = read("spaces", &sp.id);
db.rename_space(&sp.id, "other2").unwrap();
assert!(read("spaces", &sp.id) > before);
let fid = db.upsert_file(&space, "f.txt", "h", 1, "ok").unwrap();
let before = read("files", &fid);
db.rename_file(&fid, "g.txt").unwrap();
assert!(read("files", &fid) > before);
let wid = db.create_watch(&space, "topic", 24, &sid).unwrap();
let before = read("watches", &wid);
db.touch_watch(&wid, "2026-01-01T00:00:00Z").unwrap();
assert!(read("watches", &wid) > before);
let before = read("watches", &wid);
db.set_watch_session(&wid, "other-session").unwrap();
assert!(read("watches", &wid) > before);
let url = "https://x.test/";
add_session_sources(db.raw(), &sid, &[url.to_string()]).unwrap();
let read_src = || -> String {
std::thread::sleep(std::time::Duration::from_millis(2));
db.raw()
.query_row(
"SELECT updated_at FROM session_sources
WHERE session_id = ?1 AND url_norm = ?2",
(&sid, url),
|r| r.get::<_, String>(0),
)
.unwrap()
};
let before = read_src();
db.set_source_flag(&sid, url, Some("pinned")).unwrap();
assert!(read_src() > before);
let row = db
.log_usage("OpenRouter", "a/model", 1, 2, 3, 4, None, false, None, None)
.unwrap();
let read_usage = |db: &Db, row: i64| -> String {
std::thread::sleep(std::time::Duration::from_millis(2));
db.raw()
.query_row(
"SELECT updated_at FROM usage_log WHERE id = ?1",
[&row],
|r| r.get::<_, String>(0),
)
.unwrap()
};
let before = read_usage(&db, row);
db.update_usage(row, 5, 6, 7, 8, Some(0.1), false).unwrap();
assert!(read_usage(&db, row) > before);
db.upsert_model_prices(&[(
"a/model".to_string(),
"OpenRouter".to_string(),
price(1.0, 2.0),
)])
.unwrap();
let before = read_usage(&db, row);
assert_eq!(db.backfill_usage_costs().unwrap(), 1);
assert!(read_usage(&db, row) > before);
}
#[test]
fn delete_paths_write_tombstones_for_sync() {
let db = Db::open_in_memory().unwrap();
let space = db.default_space_id().unwrap();
let sid = db.create_session("t", "a/b", &space, "chat").unwrap().id;
let tombstones = |table: &str| -> Vec<String> {
let mut stmt = db
.raw()
.prepare("SELECT row_id FROM sync_tombstones WHERE table_name = ?1 ORDER BY row_id")
.unwrap();
let rows = stmt.query_map([table], |r| r.get::<_, String>(0)).unwrap();
rows.collect::<rusqlite::Result<Vec<_>>>().unwrap()
};
let mid = db.add_user_message(&sid, "hi").unwrap();
db.delete_message(&mid).unwrap();
assert_eq!(tombstones("messages"), vec![mid.clone()]);
let sid2 = db.create_session("t2", "a/b", &space, "chat").unwrap().id;
let m1 = db.add_user_message(&sid2, "one").unwrap();
let m2 = db.add_user_message(&sid2, "two").unwrap();
db.delete_session(&sid2).unwrap();
let mut expected = vec![mid, m1, m2];
expected.sort();
assert_eq!(tombstones("messages"), expected);
assert_eq!(tombstones("sessions"), vec![sid2]);
let fid = db.upsert_file(&space, "f.txt", "h", 1, "ok").unwrap();
db.delete_file(&fid).unwrap();
assert_eq!(tombstones("files"), vec![fid]);
let wid = db.create_watch(&space, "topic", 24, &sid).unwrap();
db.delete_watch(&wid).unwrap();
assert_eq!(tombstones("watches"), vec![wid]);
let sp = db.create_space("doomed").unwrap();
db.delete_space(&sp.id).unwrap();
assert_eq!(tombstones("spaces"), vec![sp.id]);
let persona = |name: &str| Persona {
name: name.into(),
model: "m".into(),
blurb: "b".into(),
};
db.save_swarm_personas(&sid, &[persona("a"), persona("b")])
.unwrap();
db.save_swarm_personas(&sid, &[persona("b")]).unwrap();
assert_eq!(
tombstones("swarm_personas"),
vec![format!("{sid}:0"), format!("{sid}:1")]
);
}
#[test]
fn v1_default_space_renumbers_to_deterministic_id() {
let dir = std::env::temp_dir().join(format!("nexus-v1-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).unwrap();
let db_path = dir.join("nexus.db");
{
let db = Db::open(&db_path).unwrap();
let old = db.default_space_id().unwrap();
assert_eq!(old, DEFAULT_SPACE);
db.raw().execute("PRAGMA user_version = 1", []).unwrap();
db.raw()
.execute(
"UPDATE spaces SET id = 'legacy-uuid' WHERE name = ?1",
[DEFAULT_SPACE],
)
.unwrap();
db.raw()
.execute(
"UPDATE sessions SET space_id = 'legacy-uuid' WHERE space_id = ?1",
[&old],
)
.unwrap();
db.raw()
.execute(
"INSERT INTO sync_tombstones (table_name, row_id, deleted_at)
VALUES ('spaces', 'legacy-uuid', '2026-01-01T00:00:00Z')",
[],
)
.unwrap();
let s = db.create_session("t", "m", &old, "chat").unwrap().id;
let _ = s;
drop(db);
}
let db = Db::open(&db_path).unwrap();
assert_eq!(db.default_space_id().unwrap(), DEFAULT_SPACE);
let n: i64 = db
.raw()
.query_row(
"SELECT COUNT(*) FROM sessions WHERE space_id = ?1",
[DEFAULT_SPACE],
|r| r.get(0),
)
.unwrap();
assert_eq!(n, 1, "sessions followed the renumber");
let tomb: i64 = db
.raw()
.query_row(
"SELECT COUNT(*) FROM sync_tombstones WHERE table_name = 'spaces' AND row_id = ?1",
[DEFAULT_SPACE],
|r| r.get(0),
)
.unwrap();
assert_eq!(tomb, 1, "the space tombstone followed too");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn device_id_is_stable_and_sync_state_roundtrips() {
let db = Db::open_in_memory().unwrap();
let a = db.device_id().unwrap();
let b = db.device_id().unwrap();
assert_eq!(a, b);
assert!(!a.is_empty());
db.set_sync_state(
"peer-1",
"sessions",
Some("2026-01-01T00:00:00Z|id1"),
Some("2026-01-02T00:00:00Z|id2"),
)
.unwrap();
db.set_sync_state("peer-1", "sessions", None, Some("2026-01-03T00:00:00Z|id3"))
.unwrap();
db.set_sync_state("peer-1", "messages", Some("c1"), None)
.unwrap();
let states = db.load_sync_state().unwrap();
assert_eq!(states.len(), 2);
let s = states.iter().find(|s| s.table_name == "sessions").unwrap();
assert_eq!(s.pull_cursor.as_deref(), Some("2026-01-01T00:00:00Z|id1"));
assert_eq!(s.push_cursor.as_deref(), Some("2026-01-03T00:00:00Z|id3"));
assert!(s.last_synced_at.is_some());
let m = states.iter().find(|s| s.table_name == "messages").unwrap();
assert_eq!(m.pull_cursor.as_deref(), Some("c1"));
}
#[test]
fn settings_scope_registry_classifies_keys_and_stores_scope() {
for local in [
"searxng_url",
"langsearch_key",
"search_provider",
"ocr_engine",
"ocr_model",
"local_ocr_model",
"usage_range",
"ui_background",
"last_update_check",
] {
assert!(Db::setting_is_local(local), "{local} should be local");
}
for sync in [
"temperature",
"top_p",
"max_tokens",
"show_stats",
"show_reasoning",
"hide_hints",
"verbosity",
"memory_model",
"embedding_model",
] {
assert!(!Db::setting_is_local(sync), "{sync} should sync");
}
let db = Db::open_in_memory().unwrap();
db.set_setting("temperature", "0.7").unwrap();
db.set_setting("ocr_engine", "tesseract").unwrap();
let scope = |key: &str| -> String {
db.raw()
.query_row(
"SELECT scope FROM app_settings WHERE key = ?1",
[key],
|r| r.get(0),
)
.unwrap()
};
assert_eq!(scope("temperature"), "sync");
assert_eq!(scope("ocr_engine"), "local");
}
#[test]
fn sync_ids_are_unique_for_citations_and_usage() {
let db = Db::open_in_memory().unwrap();
db.add_citations(
"sp",
"r.md",
&[
("https://a.test/".to_string(), None),
("https://b.test/".to_string(), None),
],
)
.unwrap();
db.log_usage("OpenRouter", "m", 1, 2, 0, 0, None, false, None, None)
.unwrap();
db.log_usage("OpenRouter", "m", 1, 2, 0, 0, None, false, None, None)
.unwrap();
let distinct: i64 = db
.raw()
.query_row(
"SELECT COUNT(DISTINCT sync_id) FROM
(SELECT sync_id FROM citations UNION ALL SELECT sync_id FROM usage_log)",
[],
|r| r.get(0),
)
.unwrap();
let total: i64 = db
.raw()
.query_row(
"SELECT (SELECT COUNT(*) FROM citations) + (SELECT COUNT(*) FROM usage_log)",
[],
|r| r.get(0),
)
.unwrap();
assert_eq!(distinct, total);
}
}