zink/codecraft
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 40 Wiki Activity
Files
c551171c1833413e49603d4ff2999a0275da7909
codecraft/server/src/store.rs
T

1946 lines
88 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! SQLite persistence for sessions and messages.
use crate::types::{ImageData, Message, Part, Role, Session, now_ms};
use serde::Serialize;
use serde_json::{Value, json};
use anyhow::Result;
use rusqlite::{Connection, OptionalExtension, params};
use std::path::Path;
use std::sync::Mutex;
pub struct Store {
conn: Mutex<Connection>,
}
impl Store {
pub fn open(path: &Path) -> Result<Self> {
let conn = Connection::open(path)?;
conn.execute_batch(
"PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
data TEXT NOT NULL,
parent_id TEXT,
updated_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS messages (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
seq INTEGER NOT NULL,
data TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS messages_session ON messages(session_id, seq);
-- One row per model request. Kept when sessions are deleted (statistics).
CREATE TABLE IF NOT EXISTS usage (
id INTEGER PRIMARY KEY AUTOINCREMENT,
ts INTEGER NOT NULL,
session_id TEXT NOT NULL,
root_session_id TEXT NOT NULL,
provider TEXT NOT NULL,
model TEXT NOT NULL,
agent TEXT NOT NULL,
input_tokens INTEGER NOT NULL,
output_tokens INTEGER NOT NULL,
cache_read_tokens INTEGER NOT NULL,
cache_write_tokens INTEGER NOT NULL,
ttft_ms INTEGER NOT NULL,
duration_ms INTEGER NOT NULL,
cost REAL NOT NULL,
error INTEGER NOT NULL DEFAULT 0
);
CREATE INDEX IF NOT EXISTS usage_ts ON usage(ts);
CREATE INDEX IF NOT EXISTS usage_root ON usage(root_session_id);
CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY,
username TEXT NOT NULL UNIQUE COLLATE NOCASE,
display_name TEXT NOT NULL DEFAULT '',
password_hash TEXT,
role TEXT NOT NULL DEFAULT 'user',
oidc_sub TEXT UNIQUE,
disabled INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS auth_tokens (
token_hash TEXT PRIMARY KEY,
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
created_at INTEGER NOT NULL,
expires_at INTEGER NOT NULL,
user_agent TEXT NOT NULL DEFAULT ''
);
CREATE TABLE IF NOT EXISTS folders (
id TEXT PRIMARY KEY,
path TEXT NOT NULL UNIQUE,
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS folder_access (
folder_id TEXT NOT NULL REFERENCES folders(id) ON DELETE CASCADE,
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
PRIMARY KEY (folder_id, user_id)
);
CREATE TABLE IF NOT EXISTS images (
id TEXT PRIMARY KEY,
user_id TEXT,
session_id TEXT,
prompt TEXT NOT NULL,
model TEXT NOT NULL,
path TEXT NOT NULL,
width INTEGER NOT NULL DEFAULT 0,
height INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS images_user ON images(user_id, created_at);
-- Images of messages, referenced by hash from the message JSON (see blobs.rs).
CREATE TABLE IF NOT EXISTS blobs (
hash TEXT PRIMARY KEY,
media_type TEXT NOT NULL,
data BLOB NOT NULL,
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS queue (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
data TEXT NOT NULL,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS queue_session ON queue(session_id, created_at);
CREATE TABLE IF NOT EXISTS api_tokens (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
name TEXT NOT NULL,
token_hash TEXT NOT NULL UNIQUE,
prefix TEXT NOT NULL,
created_at INTEGER NOT NULL,
last_used INTEGER,
expires_at INTEGER,
scope TEXT NOT NULL DEFAULT 'full'
);
CREATE TABLE IF NOT EXISTS jobs (
id TEXT PRIMARY KEY,
owner_id TEXT NOT NULL,
path TEXT NOT NULL UNIQUE,
enabled INTEGER NOT NULL DEFAULT 1,
last_run INTEGER,
last_session_id TEXT,
last_status TEXT,
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS user_settings (
user_id TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
PRIMARY KEY (user_id, key)
);
CREATE TABLE IF NOT EXISTS push_subs (
endpoint TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
p256dh TEXT NOT NULL,
auth TEXT NOT NULL,
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS audit (
id INTEGER PRIMARY KEY AUTOINCREMENT,
ts INTEGER NOT NULL,
user_id TEXT,
username TEXT NOT NULL DEFAULT '',
action TEXT NOT NULL,
detail TEXT NOT NULL DEFAULT '',
ip TEXT NOT NULL DEFAULT ''
);
CREATE INDEX IF NOT EXISTS audit_ts ON audit(ts);
CREATE TABLE IF NOT EXISTS checkpoints (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
ts INTEGER NOT NULL,
path TEXT NOT NULL,
existed INTEGER NOT NULL,
content BLOB
);
CREATE INDEX IF NOT EXISTS checkpoints_session ON checkpoints(session_id, ts);
CREATE TABLE IF NOT EXISTS project_approvals (
cwd TEXT NOT NULL,
key TEXT NOT NULL,
created_at INTEGER NOT NULL,
PRIMARY KEY (cwd, key)
);
CREATE TABLE IF NOT EXISTS share_links (
token TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
owner_id TEXT,
created_at INTEGER NOT NULL,
expires_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS share_links_session ON share_links(session_id);
CREATE TABLE IF NOT EXISTS factory_issues (
repo_url TEXT NOT NULL,
number INTEGER NOT NULL,
owner_id TEXT NOT NULL,
title TEXT NOT NULL DEFAULT '',
author TEXT NOT NULL DEFAULT '',
state TEXT NOT NULL DEFAULT 'new',
attempts INTEGER NOT NULL DEFAULT 0,
rounds INTEGER NOT NULL DEFAULT 0,
last_comment INTEGER NOT NULL DEFAULT 0,
scenarios TEXT NOT NULL DEFAULT '',
branch TEXT NOT NULL DEFAULT '',
pr_url TEXT NOT NULL DEFAULT '',
note TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
PRIMARY KEY (repo_url, number)
);
CREATE TABLE IF NOT EXISTS factory_runs (
id TEXT PRIMARY KEY,
owner_id TEXT NOT NULL,
repo_url TEXT NOT NULL,
issue INTEGER NOT NULL,
phase TEXT NOT NULL,
round INTEGER NOT NULL DEFAULT 0,
session_id TEXT NOT NULL,
model TEXT NOT NULL DEFAULT '',
started_at INTEGER NOT NULL,
ended_at INTEGER,
status TEXT NOT NULL DEFAULT 'running',
result TEXT NOT NULL DEFAULT ''
);
CREATE INDEX IF NOT EXISTS factory_runs_issue ON factory_runs(repo_url, issue);
CREATE TABLE IF NOT EXISTS project_defaults (
cwd TEXT PRIMARY KEY,
data TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS project_checks (
cwd TEXT PRIMARY KEY,
command TEXT NOT NULL,
enabled INTEGER NOT NULL DEFAULT 1
);
CREATE TABLE IF NOT EXISTS turn_checks (
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
message_id TEXT NOT NULL,
status TEXT NOT NULL,
exit_code INTEGER,
duration_ms INTEGER NOT NULL DEFAULT 0,
output TEXT NOT NULL DEFAULT '',
ts INTEGER NOT NULL,
PRIMARY KEY (session_id, message_id)
);
CREATE TABLE IF NOT EXISTS session_shares (
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
PRIMARY KEY (session_id, user_id)
);",
)?;
// Columns added after the first release.
for (table, column, ddl) in [
("sessions", "owner_id", "ALTER TABLE sessions ADD COLUMN owner_id TEXT"),
("usage", "user_id", "ALTER TABLE usage ADD COLUMN user_id TEXT"),
("users", "monthly_quota", "ALTER TABLE users ADD COLUMN monthly_quota REAL"),
("jobs", "webhook_secret", "ALTER TABLE jobs ADD COLUMN webhook_secret TEXT"),
("jobs", "settings", "ALTER TABLE jobs ADD COLUMN settings TEXT"),
("images", "blob", "ALTER TABLE images ADD COLUMN blob TEXT"),
("usage", "allowance", "ALTER TABLE usage ADD COLUMN allowance REAL"),
] {
let exists: bool = conn
.prepare(&format!("SELECT 1 FROM pragma_table_info('{table}') WHERE name = ?1"))?
.exists([column])?;
if !exists {
conn.execute(ddl, [])?;
}
}
conn.execute_batch(
"CREATE INDEX IF NOT EXISTS sessions_owner ON sessions(owner_id);
CREATE INDEX IF NOT EXISTS usage_user ON usage(user_id);",
)?;
Ok(Self { conn: Mutex::new(conn) })
}
pub fn save_session(&self, s: &Session) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO sessions (id, data, parent_id, updated_at, owner_id) VALUES (?1, ?2, ?3, ?4, ?5)
ON CONFLICT(id) DO UPDATE SET data = excluded.data, updated_at = excluded.updated_at,
owner_id = excluded.owner_id",
params![s.id, serde_json::to_string(s)?, s.parent_id, s.updated_at, s.owner_id],
)?;
Ok(())
}
pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
let conn = self.conn.lock().unwrap();
let data: Option<String> = conn
.query_row("SELECT data FROM sessions WHERE id = ?1", [id], |r| r.get(0))
.optional()?;
Ok(data.map(|d| serde_json::from_str(&d)).transpose()?)
}
/// Top-level sessions (subagent sessions are reachable through their parent).
/// `user = None` lists all sessions (admins); otherwise own and shared ones.
pub fn list_sessions(&self, user: Option<&str>) -> Result<Vec<Session>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT data FROM sessions WHERE parent_id IS NULL AND (?1 IS NULL OR owner_id = ?1
OR id IN (SELECT session_id FROM session_shares WHERE user_id = ?1))
ORDER BY updated_at DESC LIMIT 500",
)?;
let rows = stmt.query_map([user], |r| r.get::<_, String>(0))?;
let mut out = Vec::new();
for row in rows {
out.push(serde_json::from_str(&row?)?);
}
Ok(out)
}
/// Full-text search over titles and message texts of the sessions a user
/// can see (None = all). Returns (session, snippet) pairs, newest first.
pub fn search_sessions(&self, user: Option<&str>, q: &str) -> Result<Vec<(Session, String)>> {
let needle = q.to_lowercase();
let like = format!("%{}%", q.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_"));
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT s.data,
(SELECT m.data FROM messages m WHERE m.session_id = s.id AND m.data LIKE ?2 ESCAPE '\\' ORDER BY m.seq DESC LIMIT 1)
FROM sessions s
WHERE s.parent_id IS NULL
AND (?1 IS NULL OR s.owner_id = ?1 OR s.id IN (SELECT session_id FROM session_shares WHERE user_id = ?1))
AND (s.data LIKE ?2 ESCAPE '\\' OR EXISTS (SELECT 1 FROM messages m WHERE m.session_id = s.id AND m.data LIKE ?2 ESCAPE '\\'))
ORDER BY s.updated_at DESC LIMIT 50",
)?;
let rows = stmt.query_map(rusqlite::params![user, like], |r| Ok((r.get::<_, String>(0)?, r.get::<_, Option<String>>(1)?)))?;
let mut out = Vec::new();
for row in rows {
let (sdata, mdata) = row?;
let session: Session = serde_json::from_str(&sdata)?;
// Snippet from the visible text of the matching message (not JSON keys).
let snippet = mdata.and_then(|d| serde_json::from_str::<Message>(&d).ok()).and_then(|m| snippet_in(&m, &needle)).unwrap_or_default();
if snippet.is_empty() && !session.title.to_lowercase().contains(&needle) {
continue; // matched only JSON structure (keys, ids)
}
out.push((session, snippet));
}
Ok(out)
}
/// A part of a session's messages for the chat view (long chats are not loaded at
/// once). Returns the messages (oldest first) and whether older / newer ones exist.
pub fn messages_window(&self, session_id: &str, w: &MessageWindow) -> Result<(Vec<Message>, bool, bool)> {
let conn = self.conn.lock().unwrap();
let seq_of = |id: &str| -> Result<i64> {
conn.query_row("SELECT seq FROM messages WHERE session_id = ?1 AND id = ?2", params![session_id, id], |r| r.get(0))
.optional()?
.ok_or_else(|| anyhow::anyhow!("message {id} not found"))
};
let (sql, a, limit): (&str, i64, i64) = match w {
// From the last summary (compaction) on, at most `limit`.
MessageWindow::Tail { limit } => {
let from: i64 = conn
.query_row(
"SELECT seq FROM messages WHERE session_id = ?1 AND data LIKE '%\"type\":\"summary\"%' ORDER BY seq DESC LIMIT 1",
[session_id],
|r| r.get(0),
)
.optional()?
.unwrap_or(i64::MIN);
("SELECT data FROM messages WHERE session_id = ?1 AND seq >= ?2 ORDER BY seq DESC LIMIT ?3", from, *limit)
}
// Scrolling up stops at a summary (compaction); the history before it on request.
MessageWindow::Before { id, limit, stop_at_summary } => {
let before = seq_of(id)?;
let summary: Option<i64> = if *stop_at_summary {
conn.query_row(
"SELECT seq FROM messages WHERE session_id = ?1 AND seq < ?2 AND data LIKE '%\"type\":\"summary\"%' ORDER BY seq DESC LIMIT 1",
params![session_id, before],
|r| r.get(0),
)
.optional()?
} else {
None
};
match summary {
Some(from) => {
let mut stmt = conn.prepare("SELECT seq, data FROM messages WHERE session_id = ?1 AND seq >= ?2 AND seq < ?3 ORDER BY seq DESC LIMIT ?4")?;
let mut rows: Vec<(i64, Message)> = Vec::new();
for r in stmt.query_map(params![session_id, from, before, limit], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)))? {
let (seq, data) = r?;
rows.push((seq, serde_json::from_str(&data)?));
}
drop(stmt);
return Self::finish_window(&conn, session_id, rows, false);
}
None => ("SELECT data FROM messages WHERE session_id = ?1 AND seq < ?2 ORDER BY seq DESC LIMIT ?3", before, *limit),
}
}
MessageWindow::After { id, limit } => ("SELECT data FROM messages WHERE session_id = ?1 AND seq > ?2 ORDER BY seq ASC LIMIT ?3", seq_of(id)?, *limit),
// The message and about half the limit before it.
MessageWindow::Around { id, limit } => {
let seq = seq_of(id)?;
let before: i64 = conn.query_row(
"SELECT COALESCE(MIN(seq), ?2) FROM (SELECT seq FROM messages WHERE session_id = ?1 AND seq < ?2 ORDER BY seq DESC LIMIT ?3)",
params![session_id, seq, limit / 2],
|r| r.get(0),
)?;
("SELECT data FROM messages WHERE session_id = ?1 AND seq >= ?2 ORDER BY seq ASC LIMIT ?3", before, *limit)
}
};
let mut rows: Vec<(i64, Message)> = Vec::new();
let mut stmt = conn.prepare(&sql.replace("SELECT data", "SELECT seq, data"))?;
for r in stmt.query_map(params![session_id, a, limit], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)))? {
let (seq, data) = r?;
rows.push((seq, serde_json::from_str(&data)?));
}
drop(stmt);
Self::finish_window(&conn, session_id, rows, matches!(w, MessageWindow::Tail { .. }))
}
/// Sort a window and tell whether older / newer messages exist around it.
fn finish_window(conn: &Connection, session_id: &str, mut rows: Vec<(i64, Message)>, tail: bool) -> Result<(Vec<Message>, bool, bool)> {
rows.sort_by_key(|(seq, _)| *seq);
let exists = |q: &str, seq: i64| -> Result<bool> { Ok(conn.prepare(q)?.exists(params![session_id, seq])?) };
let older = match rows.first() {
Some((first, _)) => exists("SELECT 1 FROM messages WHERE session_id = ?1 AND seq < ?2", *first)?,
None => tail && exists("SELECT 1 FROM messages WHERE session_id = ?1 AND seq > ?2", i64::MIN)?,
};
let newer = match rows.last() {
Some((last, _)) => exists("SELECT 1 FROM messages WHERE session_id = ?1 AND seq > ?2", *last)?,
None => false,
};
Ok((rows.into_iter().map(|(_, m)| m).collect(), older, newer))
}
/// Messages of one session containing `q` (case-insensitive), newest first:
/// (message id, role, created_at, snippet).
pub fn search_messages(&self, session_id: &str, q: &str, limit: usize) -> Result<Vec<(String, Role, i64, String)>> {
let needle = q.to_lowercase();
let like = format!("%{}%", q.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_"));
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT data FROM messages WHERE session_id = ?1 AND data LIKE ?2 ESCAPE '\\' ORDER BY seq DESC")?;
let mut out = Vec::new();
for r in stmt.query_map(params![session_id, like], |r| r.get::<_, String>(0))? {
let Ok(m) = serde_json::from_str::<Message>(&r?) else { continue };
// Only visible text counts (not JSON keys or ids).
if let Some(snippet) = snippet_in(&m, &needle) {
out.push((m.id.clone(), m.role, m.created_at, snippet));
if out.len() >= limit {
break;
}
}
}
Ok(out)
}
/// Ids of the subagent sessions started by a session.
pub fn child_sessions(&self, id: &str) -> Result<Vec<String>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT id FROM sessions WHERE parent_id = ?1")?;
let rows = stmt.query_map([id], |r| r.get(0))?.collect::<rusqlite::Result<Vec<String>>>()?;
Ok(rows)
}
pub fn delete_session(&self, id: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
// Child sessions (subagents) go with their parent.
conn.execute("DELETE FROM sessions WHERE parent_id = ?1", [id])?;
conn.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
Ok(())
}
/// (title, cwd) of all other chats (subagent sessions of this chat excluded).
pub fn other_session_cwds(&self, id: &str) -> Result<Vec<(String, String)>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT COALESCE(json_extract(data, '$.title'), ''), COALESCE(json_extract(data, '$.cwd'), '') FROM sessions WHERE id != ?1 AND COALESCE(parent_id, '') != ?1",
)?;
let rows = stmt.query_map([id], |r| Ok((r.get(0)?, r.get(1)?)))?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
}
pub fn add_message(&self, session_id: &str, m: &Message) -> Result<()> {
// Images go to the blob table (encoded before taking the lock).
let mut m = m.clone();
let blobs = crate::blobs::externalize(&mut m);
let conn = self.conn.lock().unwrap();
for b in &blobs {
conn.execute(
"INSERT OR IGNORE INTO blobs (hash, media_type, data, created_at) VALUES (?1, ?2, ?3, ?4)",
params![b.hash, b.media_type, b.bytes, now_ms()],
)?;
}
conn.execute(
"INSERT INTO messages (id, session_id, seq, data) VALUES (?1, ?2,
(SELECT COALESCE(MAX(seq), 0) + 1 FROM messages WHERE session_id = ?2), ?3)
ON CONFLICT(id) DO UPDATE SET data = excluded.data",
params![m.id, session_id, serde_json::to_string(&m)?],
)?;
conn.execute(
"UPDATE sessions SET updated_at = ?2 WHERE id = ?1",
params![session_id, now_ms()],
)?;
Ok(())
}
pub fn messages(&self, session_id: &str) -> Result<Vec<Message>> {
let conn = self.conn.lock().unwrap();
let mut stmt =
conn.prepare("SELECT data FROM messages WHERE session_id = ?1 ORDER BY seq")?;
let rows = stmt.query_map([session_id], |r| r.get::<_, String>(0))?;
let mut out = Vec::new();
for row in rows {
out.push(serde_json::from_str(&row?)?);
}
Ok(out)
}
/// Messages with the bytes of their images filled in – for model requests and exports.
pub fn messages_for_model(&self, session_id: &str) -> Result<Vec<Message>> {
let mut out = self.messages(session_id)?;
let conn = self.conn.lock().unwrap();
let mut cache: std::collections::HashMap<String, String> = std::collections::HashMap::new();
for m in &mut out {
crate::blobs::for_each_image(m, |img| fill_image(&conn, &mut cache, img));
}
Ok(out)
}
/// Fill in the bytes of one image reference (no-op for inline images).
pub fn resolve_image(&self, img: &ImageData) -> ImageData {
let mut img = img.clone();
let conn = self.conn.lock().unwrap();
fill_image(&conn, &mut std::collections::HashMap::new(), &mut img);
img
}
/// One stored image: (media type, bytes).
pub fn blob(&self, hash: &str) -> Result<Option<(String, Vec<u8>)>> {
let conn = self.conn.lock().unwrap();
Ok(conn
.query_row("SELECT media_type, data FROM blobs WHERE hash = ?1", [hash], |r| Ok((r.get(0)?, r.get(1)?)))
.optional()?)
}
/// Move inline images of older messages to the blob table (once, after the update
/// that introduced it). Does not touch `updated_at` of the sessions.
pub fn migrate_inline_images(&self) -> Result<usize> {
let ids: Vec<String> = {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT id FROM messages WHERE data LIKE '%\"media_type\"%\"data\":\"%'")?;
let rows = stmt.query_map([], |r| r.get(0))?.collect::<rusqlite::Result<Vec<String>>>()?;
rows
};
let mut moved = 0;
for id in ids {
let data: Option<String> = {
let conn = self.conn.lock().unwrap();
conn.query_row("SELECT data FROM messages WHERE id = ?1", [&id], |r| r.get(0)).optional()?
};
let Some(data) = data else { continue };
let Ok(mut m) = serde_json::from_str::<Message>(&data) else { continue };
let blobs = crate::blobs::externalize(&mut m);
if blobs.is_empty() {
continue;
}
let conn = self.conn.lock().unwrap();
for b in &blobs {
conn.execute(
"INSERT OR IGNORE INTO blobs (hash, media_type, data, created_at) VALUES (?1, ?2, ?3, ?4)",
params![b.hash, b.media_type, b.bytes, now_ms()],
)?;
}
conn.execute("UPDATE messages SET data = ?2 WHERE id = ?1", params![id, serde_json::to_string(&m)?])?;
moved += blobs.len();
}
Ok(moved)
}
/// Delete images no message refers to any more (deleted or truncated sessions).
pub fn gc_blobs(&self) -> Result<usize> {
let conn = self.conn.lock().unwrap();
let mut used = std::collections::HashSet::new();
let mut stmt = conn.prepare("SELECT data FROM messages WHERE data LIKE '%\"blob\":\"%'")?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
let data: String = row.get(0)?;
for (i, _) in data.match_indices("\"blob\":\"") {
if let Some(h) = data.get(i + 8..i + 72) {
used.insert(h.to_string());
}
}
}
drop(rows);
drop(stmt);
// Copies of generated images (see ImageRecord::blob).
for h in conn.prepare("SELECT blob FROM images WHERE blob IS NOT NULL")?.query_map([], |r| r.get::<_, String>(0))? {
used.insert(h?);
}
let all: Vec<String> = conn.prepare("SELECT hash FROM blobs")?.query_map([], |r| r.get(0))?.collect::<rusqlite::Result<_>>()?;
let mut removed = 0;
for h in all.iter().filter(|h| !used.contains(*h)) {
removed += conn.execute("DELETE FROM blobs WHERE hash = ?1", [h])?;
}
Ok(removed)
}
/// Give the space of deleted data back to the file system when it is worth it.
pub fn vacuum_if_worthwhile(&self) -> Result<bool> {
let conn = self.conn.lock().unwrap();
let free: i64 = conn.query_row("PRAGMA freelist_count", [], |r| r.get(0))?;
let page: i64 = conn.query_row("PRAGMA page_size", [], |r| r.get(0))?;
if free * page < 32 * 1024 * 1024 {
return Ok(false);
}
conn.execute_batch("VACUUM")?;
Ok(true)
}
/// Remove all messages after (and including) the given one – used for "edit & resend".
pub fn truncate_from(&self, session_id: &str, message_id: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"DELETE FROM messages WHERE session_id = ?1 AND seq >=
(SELECT seq FROM messages WHERE id = ?2 AND session_id = ?1)",
params![session_id, message_id],
)?;
Ok(())
}
}
/// Put the bytes of a stored image into `img.data`.
fn fill_image(conn: &Connection, cache: &mut std::collections::HashMap<String, String>, img: &mut ImageData) {
use base64::Engine;
let Some(hash) = img.blob.clone().filter(|_| img.data.is_empty()) else { return };
if !cache.contains_key(&hash) {
let Ok(bytes) = conn.query_row("SELECT data FROM blobs WHERE hash = ?1", [&hash], |r| r.get::<_, Vec<u8>>(0)) else {
tracing::warn!("image {hash} is missing in the store");
return;
};
cache.insert(hash.clone(), base64::engine::general_purpose::STANDARD.encode(bytes));
}
img.data = cache[&hash].clone();
}
/// One model request, as recorded for statistics.
#[derive(Debug, Clone, Serialize)]
pub struct UsageRecord {
pub ts: i64,
pub session_id: String,
pub root_session_id: String,
pub provider: String,
pub model: String,
pub agent: String,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_read_tokens: u64,
pub cache_write_tokens: u64,
/// Time to first token (prefill + queueing), milliseconds.
pub ttft_ms: u64,
pub duration_ms: u64,
pub cost: f64,
pub error: bool,
/// User to bill when there is no session (CLI requests); else the session owner.
pub user_id: Option<String>,
/// opencode Go: share of the plan's usage limits (USD), as reported by the provider.
pub allowance: Option<f64>,
}
/// Aggregation SQL shared by all statistics queries.
const AGG: &str = "COUNT(*) AS requests,
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0),
COALESCE(SUM(cost),0.0), COALESCE(SUM(ttft_ms),0), COALESCE(SUM(duration_ms),0),
COALESCE(SUM(error),0)";
fn agg_row(r: &rusqlite::Row, offset: usize) -> rusqlite::Result<Value> {
let requests: i64 = r.get(offset)?;
let input: i64 = r.get(offset + 1)?;
let output: i64 = r.get(offset + 2)?;
let cache_read: i64 = r.get(offset + 3)?;
let cache_write: i64 = r.get(offset + 4)?;
let cost: f64 = r.get(offset + 5)?;
let ttft: i64 = r.get(offset + 6)?;
let duration: i64 = r.get(offset + 7)?;
let errors: i64 = r.get(offset + 8)?;
let gen_ms = (duration - ttft).max(0);
let prompt = input + cache_read + cache_write;
Ok(json!({
"requests": requests,
"input_tokens": input,
"output_tokens": output,
"cache_read_tokens": cache_read,
"cache_write_tokens": cache_write,
"cost": cost,
"errors": errors,
"avg_ttft_ms": if requests > 0 { ttft / requests } else { 0 },
// Output speed: tokens generated per second after the first token.
"output_tps": if gen_ms > 0 { output as f64 * 1000.0 / gen_ms as f64 } else { 0.0 },
// Prefill speed: prompt tokens processed per second until the first token.
"prefill_tps": if ttft > 0 { prompt as f64 * 1000.0 / ttft as f64 } else { 0.0 },
"cache_hit_rate": if prompt > 0 { cache_read as f64 / prompt as f64 } else { 0.0 },
"duration_ms": duration,
}))
}
impl Store {
pub fn add_usage(&self, u: &UsageRecord) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO usage (ts, session_id, root_session_id, provider, model, agent,
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
ttft_ms, duration_ms, cost, error, user_id, allowance)
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,
COALESCE(?15, (SELECT owner_id FROM sessions WHERE id = ?3)), ?16)",
params![
u.ts, u.session_id, u.root_session_id, u.provider, u.model, u.agent,
u.input_tokens as i64, u.output_tokens as i64, u.cache_read_tokens as i64,
u.cache_write_tokens as i64, u.ttft_ms as i64, u.duration_ms as i64, u.cost, u.error as i64, u.user_id,
u.allowance
],
)?;
Ok(())
}
/// Statistics for one chat (including its subagents), totals and per model.
pub fn session_stats(&self, root_session_id: &str) -> Result<Value> {
let conn = self.conn.lock().unwrap();
let total = conn.query_row(
&format!("SELECT {AGG} FROM usage WHERE root_session_id = ?1"),
[root_session_id],
|r| agg_row(r, 0),
)?;
let mut stmt = conn.prepare(&format!(
"SELECT provider || '/' || model, {AGG} FROM usage WHERE root_session_id = ?1
GROUP BY 1 ORDER BY 1"
))?;
let by_model = stmt
.query_map([root_session_id], |r| {
let mut v = agg_row(r, 1)?;
v["model"] = json!(r.get::<_, String>(0)?);
Ok(v)
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let mut stmt = conn.prepare(&format!(
"SELECT agent || ' · ' || provider || '/' || model, {AGG} FROM usage WHERE root_session_id = ?1 GROUP BY 1 ORDER BY 1"
))?;
let by_agent = stmt
.query_map([root_session_id], |r| {
let mut v = agg_row(r, 1)?;
v["key"] = json!(r.get::<_, String>(0)?);
Ok(v)
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let (first, last): (Option<i64>, Option<i64>) =
conn.query_row("SELECT MIN(ts), MAX(ts) FROM usage WHERE root_session_id = ?1", [root_session_id], |r| Ok((r.get(0)?, r.get(1)?)))?;
Ok(json!({"total": total, "by_model": by_model, "by_agent": by_agent, "first": first, "last": last}))
}
/// Plan requests booked at 0 (opencode Go reports nothing charged, only `allowance`):
/// (row id, ts, "provider/model", usage).
pub fn plan_requests_without_cost(&self) -> Result<Vec<(i64, i64, String, crate::types::Usage)>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT id, ts, provider || '/' || model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, allowance
FROM usage WHERE cost = 0 AND allowance > 0",
)?;
let rows = stmt
.query_map([], |r| {
Ok((
r.get(0)?,
r.get(1)?,
r.get(2)?,
crate::types::Usage {
input_tokens: r.get::<_, i64>(3)? as u64,
output_tokens: r.get::<_, i64>(4)? as u64,
cache_read_tokens: r.get::<_, i64>(5)? as u64,
cache_write_tokens: r.get::<_, i64>(6)? as u64,
allowance: r.get(7)?,
..Default::default()
},
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn set_usage_cost(&self, row: i64, cost: f64) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE usage SET cost = ?2 WHERE id = ?1", params![row, cost])?;
Ok(())
}
/// Global statistics in [from, to) (ms). `tz_offset_min` shifts day/hour buckets
/// into the viewer's local time.
pub fn global_stats(&self, from: i64, to: i64, tz_offset_min: i64, user: Option<&str>) -> Result<Value> {
let conn = self.conn.lock().unwrap();
let shift = tz_offset_min * 60_000;
let total = conn.query_row(
&format!("SELECT {AGG} FROM usage WHERE ts >= ?1 AND ts < ?2 AND (?3 IS NULL OR user_id = ?3)"),
params![from, to, user],
|r| agg_row(r, 0),
)?;
let grouped_in = |from: i64, key_sql: &str, key_name: &str| -> Result<Vec<Value>> {
let mut stmt = conn.prepare(&format!(
"SELECT {key_sql}, {AGG} FROM usage WHERE ts >= ?1 AND ts < ?2 AND ?3 = ?3
AND (?4 IS NULL OR user_id = ?4) GROUP BY 1 ORDER BY 1"
))?;
let rows = stmt
.query_map(params![from, to, shift, user], |r| {
let mut v = agg_row(r, 1)?;
v[key_name] = match r.get_ref(0)? {
rusqlite::types::ValueRef::Integer(i) => json!(i),
rusqlite::types::ValueRef::Text(t) => json!(String::from_utf8_lossy(t)),
_ => Value::Null,
};
Ok(v)
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
};
let grouped = |key_sql: &str, key_name: &str| grouped_in(from, key_sql, key_name);
let day_sql = "strftime('%Y-%m-%d', (ts + ?3) / 1000, 'unixepoch')";
let by_day = grouped(day_sql, "day")?;
// Long ranges are shown per month.
let by_month = grouped("strftime('%Y-%m', (ts + ?3) / 1000, 'unixepoch')", "month")?;
let first_ts: Option<i64> = conn.query_row("SELECT MIN(ts) FROM usage WHERE (?1 IS NULL OR user_id = ?1)", params![user], |r| r.get(0))?;
let by_model = grouped("provider || '/' || model", "model")?;
let by_provider = grouped("provider", "provider")?;
let by_agent = grouped("agent", "agent")?;
let by_user = grouped(
"COALESCE((SELECT username FROM users WHERE users.id = usage.user_id), '(none)')",
"user",
)?;
// Heatmap: rolling window of the last 14 days (local date + hour),
// independent of the selected range.
let heatmap = grouped_in(
to - 15 * 86_400_000,
"strftime('%Y-%m-%d %H', (ts + ?3) / 1000, 'unixepoch')",
"slot",
)?;
let by_day_model = {
let mut stmt = conn.prepare(&format!(
"SELECT {day_sql}, provider || '/' || model, SUM(cost),
SUM(input_tokens + cache_read_tokens + cache_write_tokens), SUM(output_tokens)
FROM usage WHERE ts >= ?1 AND ts < ?2 AND (?4 IS NULL OR user_id = ?4)
GROUP BY 1, 2 ORDER BY 1"
))?;
stmt.query_map(params![from, to, shift, user], |r| {
Ok(json!({
"day": r.get::<_, String>(0)?,
"model": r.get::<_, String>(1)?,
"cost": r.get::<_, f64>(2)?,
"input_tokens": r.get::<_, i64>(3)?,
"output_tokens": r.get::<_, i64>(4)?,
}))
})?
.collect::<rusqlite::Result<Vec<_>>>()?
};
// Most expensive chats in the range.
let top_sessions = {
let mut stmt = conn.prepare(&format!(
"SELECT u.root_session_id, COALESCE(json_extract(s.data, '$.title'), CASE WHEN u.root_session_id LIKE 'cli:%' THEN 'CLI (local)' ELSE '(deleted)' END), {AGG}
FROM usage u LEFT JOIN sessions s ON s.id = u.root_session_id
WHERE u.ts >= ?1 AND u.ts < ?2 AND (?3 IS NULL OR u.user_id = ?3)
GROUP BY 1 ORDER BY SUM(u.cost) DESC, COUNT(*) DESC LIMIT 15"
))?;
stmt.query_map(params![from, to, user], |r| {
let mut v = agg_row(r, 2)?;
v["session_id"] = json!(r.get::<_, String>(0)?);
v["title"] = json!(r.get::<_, String>(1)?);
Ok(v)
})?
.collect::<rusqlite::Result<Vec<_>>>()?
};
Ok(json!({
"total": total,
"by_day": by_day,
"by_month": by_month,
"first_ts": first_ts,
"by_day_model": by_day_model,
"by_model": by_model,
"by_provider": by_provider,
"by_agent": by_agent,
"by_user": by_user,
"heatmap": heatmap,
"top_sessions": top_sessions,
}))
}
}
// --- users & auth -------------------------------------------------------------
#[derive(Debug, Clone, Serialize)]
pub struct User {
pub id: String,
pub username: String,
pub display_name: String,
pub role: String,
pub has_password: bool,
pub oidc: bool,
pub disabled: bool,
pub created_at: i64,
/// Monthly cost limit for this user (None = unlimited).
pub monthly_quota: Option<f64>,
}
impl User {
pub fn is_admin(&self) -> bool {
self.role == "admin"
}
}
const USER_COLS: &str =
"id, username, display_name, role, password_hash IS NOT NULL, oidc_sub IS NOT NULL, disabled, created_at, monthly_quota";
fn user_row(r: &rusqlite::Row) -> rusqlite::Result<User> {
Ok(User {
id: r.get(0)?,
username: r.get(1)?,
display_name: r.get(2)?,
role: r.get(3)?,
has_password: r.get(4)?,
oidc: r.get(5)?,
disabled: r.get::<_, i64>(6)? != 0,
created_at: r.get(7)?,
monthly_quota: r.get(8)?,
})
}
impl Store {
pub fn user_count(&self) -> Result<i64> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row("SELECT COUNT(*) FROM users", [], |r| r.get(0))?)
}
pub fn create_user(&self, username: &str, display_name: &str, role: &str, password_hash: Option<&str>, oidc_sub: Option<&str>) -> Result<User> {
let id = crate::types::new_id();
{
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO users (id, username, display_name, password_hash, role, oidc_sub, created_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
params![id, username, display_name, password_hash, role, oidc_sub, now_ms()],
)?;
}
self.user(&id)?.ok_or_else(|| anyhow::anyhow!("user vanished"))
}
pub fn user(&self, id: &str) -> Result<Option<User>> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row(&format!("SELECT {USER_COLS} FROM users WHERE id = ?1"), [id], user_row).optional()?)
}
pub fn user_by_name(&self, username: &str) -> Result<Option<(User, Option<String>)>> {
let conn = self.conn.lock().unwrap();
Ok(conn
.query_row(
&format!("SELECT {USER_COLS}, password_hash FROM users WHERE username = ?1"),
[username],
|r| Ok((user_row(r)?, r.get::<_, Option<String>>(9)?)),
)
.optional()?)
}
pub fn user_by_oidc(&self, sub: &str) -> Result<Option<User>> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row(&format!("SELECT {USER_COLS} FROM users WHERE oidc_sub = ?1"), [sub], user_row).optional()?)
}
pub fn users(&self) -> Result<Vec<User>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(&format!("SELECT {USER_COLS} FROM users ORDER BY username"))?;
let rows = stmt.query_map([], user_row)?.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
/// Update selected fields; None leaves a field unchanged.
pub fn update_user(&self, id: &str, display_name: Option<&str>, role: Option<&str>, disabled: Option<bool>, password_hash: Option<&str>, oidc_sub: Option<&str>) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"UPDATE users SET display_name = COALESCE(?2, display_name), role = COALESCE(?3, role),
disabled = COALESCE(?4, disabled), password_hash = COALESCE(?5, password_hash),
oidc_sub = COALESCE(?6, oidc_sub) WHERE id = ?1",
params![id, display_name, role, disabled.map(|d| d as i64), password_hash, oidc_sub],
)?;
if disabled == Some(true) {
conn.execute("DELETE FROM auth_tokens WHERE user_id = ?1", [id])?;
}
Ok(())
}
pub fn delete_user(&self, id: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM users WHERE id = ?1", [id])?;
Ok(())
}
pub fn admin_count(&self) -> Result<i64> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row("SELECT COUNT(*) FROM users WHERE role = 'admin' AND disabled = 0", [], |r| r.get(0))?)
}
pub fn add_token(&self, token_hash: &str, user_id: &str, expires_at: i64, user_agent: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM auth_tokens WHERE expires_at < ?1", [now_ms()])?;
conn.execute(
"INSERT INTO auth_tokens (token_hash, user_id, created_at, expires_at, user_agent) VALUES (?1, ?2, ?3, ?4, ?5)",
params![token_hash, user_id, now_ms(), expires_at, user_agent],
)?;
Ok(())
}
// --- personal API tokens ------------------------------------------------
pub fn api_token_add(&self, user_id: &str, name: &str, token_hash: &str, prefix: &str, expires_at: Option<i64>, scope: &str) -> Result<String> {
let id = crate::types::new_id();
self.conn.lock().unwrap().execute(
"INSERT INTO api_tokens (id, user_id, name, token_hash, prefix, created_at, expires_at, scope) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
params![id, user_id, name, token_hash, prefix, now_ms(), expires_at, scope],
)?;
Ok(id)
}
pub fn api_tokens(&self, user_id: &str) -> Result<Vec<serde_json::Value>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT id, name, prefix, created_at, last_used, expires_at, scope FROM api_tokens WHERE user_id = ?1 ORDER BY created_at")?;
let rows = stmt.query_map([user_id], |r| {
Ok(serde_json::json!({
"id": r.get::<_, String>(0)?, "name": r.get::<_, String>(1)?, "prefix": r.get::<_, String>(2)?,
"created_at": r.get::<_, i64>(3)?, "last_used": r.get::<_, Option<i64>>(4)?, "expires_at": r.get::<_, Option<i64>>(5)?, "scope": r.get::<_, String>(6)?,
}))
})?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
}
pub fn api_token_delete(&self, user_id: &str, id: &str) -> Result<bool> {
Ok(self.conn.lock().unwrap().execute("DELETE FROM api_tokens WHERE id = ?1 AND user_id = ?2", [id, user_id])? > 0)
}
/// Scope ("full" / "read") of a personal API token.
pub fn api_token_scope(&self, token_hash: &str) -> Option<String> {
self.conn.lock().unwrap().query_row("SELECT scope FROM api_tokens WHERE token_hash = ?1", [token_hash], |r| r.get(0)).optional().ok().flatten()
}
/// The (enabled) user behind a valid personal API token.
pub fn user_by_api_token(&self, token_hash: &str) -> Result<Option<User>> {
let conn = self.conn.lock().unwrap();
let user = conn
.query_row(
&format!(
"SELECT {} FROM api_tokens t JOIN users u ON u.id = t.user_id
WHERE t.token_hash = ?1 AND (t.expires_at IS NULL OR t.expires_at > ?2) AND u.disabled = 0",
USER_COLS.split(", ").map(|c| format!("u.{c}")).collect::<Vec<_>>().join(", ")
),
params![token_hash, now_ms()],
user_row,
)
.optional()?;
if user.is_some() {
// At most one write per minute per token.
conn.execute("UPDATE api_tokens SET last_used = ?2 WHERE token_hash = ?1 AND (last_used IS NULL OR last_used < ?2 - 60000)", params![token_hash, now_ms()])?;
}
Ok(user)
}
/// The (enabled) user behind a valid login token.
pub fn user_by_token(&self, token_hash: &str) -> Result<Option<User>> {
let conn = self.conn.lock().unwrap();
Ok(conn
.query_row(
&format!(
"SELECT {} FROM auth_tokens t JOIN users u ON u.id = t.user_id
WHERE t.token_hash = ?1 AND t.expires_at > ?2 AND u.disabled = 0",
USER_COLS.split(", ").map(|c| format!("u.{c}")).collect::<Vec<_>>().join(", ")
),
params![token_hash, now_ms()],
user_row,
)
.optional()?)
}
pub fn delete_token(&self, token_hash: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM auth_tokens WHERE token_hash = ?1", [token_hash])?;
Ok(())
}
pub fn delete_user_tokens(&self, user_id: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM auth_tokens WHERE user_id = ?1", [user_id])?;
Ok(())
}
pub fn shares(&self, session_id: &str) -> Result<Vec<String>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT user_id FROM session_shares WHERE session_id = ?1")?;
let rows = stmt.query_map([session_id], |r| r.get(0))?.collect::<rusqlite::Result<Vec<String>>>()?;
Ok(rows)
}
pub fn set_shares(&self, session_id: &str, users: &[String]) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM session_shares WHERE session_id = ?1", [session_id])?;
for u in users {
conn.execute("INSERT OR IGNORE INTO session_shares (session_id, user_id) VALUES (?1, ?2)", params![session_id, u])?;
}
Ok(())
}
/// Assign sessions without owner (created before user management) to a user.
pub fn adopt_orphans(&self, user_id: &str) -> Result<usize> {
let conn = self.conn.lock().unwrap();
let ids: Vec<(String, String)> = {
let mut stmt = conn.prepare("SELECT id, data FROM sessions WHERE owner_id IS NULL")?;
stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?.collect::<rusqlite::Result<Vec<_>>>()?
};
for (id, data) in &ids {
let mut v: Value = serde_json::from_str(data)?;
v["owner_id"] = json!(user_id);
conn.execute("UPDATE sessions SET owner_id = ?2, data = ?3 WHERE id = ?1", params![id, user_id, v.to_string()])?;
}
conn.execute(
"UPDATE usage SET user_id = ?1 WHERE user_id IS NULL",
[user_id],
)?;
Ok(ids.len())
}
}
// --- shared folders -----------------------------------------------------------
#[derive(Debug, Clone, Serialize)]
pub struct Folder {
pub id: String,
/// Absolute path inside the workspace.
pub path: String,
pub users: Vec<String>,
pub created_at: i64,
}
impl Store {
pub fn folders(&self) -> Result<Vec<Folder>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT f.id, f.path, f.created_at, COALESCE(group_concat(a.user_id), '')
FROM folders f LEFT JOIN folder_access a ON a.folder_id = f.id GROUP BY f.id ORDER BY f.path",
)?;
let rows = stmt
.query_map([], |r| {
let users: String = r.get(3)?;
Ok(Folder {
id: r.get(0)?,
path: r.get(1)?,
created_at: r.get(2)?,
users: users.split(',').filter(|u| !u.is_empty()).map(String::from).collect(),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn add_folder(&self, path: &str) -> Result<String> {
let conn = self.conn.lock().unwrap();
let id = crate::types::new_id();
conn.execute("INSERT INTO folders (id, path, created_at) VALUES (?1, ?2, ?3)", params![id, path, now_ms()])?;
Ok(id)
}
pub fn delete_folder(&self, id: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM folders WHERE id = ?1", [id])?;
Ok(())
}
pub fn set_folder_access(&self, id: &str, users: &[String]) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM folder_access WHERE folder_id = ?1", [id])?;
for u in users {
conn.execute("INSERT OR IGNORE INTO folder_access (folder_id, user_id) VALUES (?1, ?2)", params![id, u])?;
}
Ok(())
}
/// Paths of folders shared with a user.
pub fn folders_for_user(&self, user_id: &str) -> Result<Vec<String>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT f.path FROM folders f JOIN folder_access a ON a.folder_id = f.id WHERE a.user_id = ?1 ORDER BY f.path",
)?;
let rows = stmt.query_map([user_id], |r| r.get(0))?.collect::<rusqlite::Result<Vec<String>>>()?;
Ok(rows)
}
}
// --- generated images ---------------------------------------------------------
#[derive(Debug, Clone, Serialize)]
pub struct ImageRecord {
pub id: String,
pub user_id: Option<String>,
pub session_id: Option<String>,
pub prompt: String,
pub model: String,
pub path: String,
pub width: i64,
pub height: i64,
pub created_at: i64,
/// Copy of the bytes in the image store: the file may be moved or deleted later
/// (agents tidy up folders), the image stays viewable.
#[serde(skip)]
pub blob: Option<String>,
}
fn image_row(r: &rusqlite::Row) -> rusqlite::Result<ImageRecord> {
Ok(ImageRecord {
id: r.get(0)?,
user_id: r.get(1)?,
session_id: r.get(2)?,
prompt: r.get(3)?,
model: r.get(4)?,
path: r.get(5)?,
width: r.get(6)?,
height: r.get(7)?,
created_at: r.get(8)?,
blob: r.get(9)?,
})
}
const IMAGE_COLS: &str = "id, user_id, session_id, prompt, model, path, width, height, created_at, blob";
impl Store {
pub fn add_image(&self, img: &ImageRecord) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
&format!("INSERT INTO images ({IMAGE_COLS}) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10)"),
params![img.id, img.user_id, img.session_id, img.prompt, img.model, img.path, img.width, img.height, img.created_at, img.blob],
)?;
Ok(())
}
/// Store bytes in the image store; returns their hash.
pub fn put_blob(&self, media_type: &str, bytes: &[u8]) -> Result<String> {
let hash = crate::blobs::hash(bytes);
self.conn.lock().unwrap().execute(
"INSERT OR IGNORE INTO blobs (hash, media_type, data, created_at) VALUES (?1, ?2, ?3, ?4)",
params![hash, media_type, bytes, now_ms()],
)?;
Ok(hash)
}
/// Generated images from before the copies in the image store (id, path).
pub fn images_without_copy(&self) -> Result<Vec<(String, String)>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT id, path FROM images WHERE blob IS NULL")?;
let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn set_image_blob(&self, id: &str, hash: &str) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE images SET blob = ?2 WHERE id = ?1", params![id, hash])?;
Ok(())
}
/// A file or folder was moved: generated images in it follow.
pub fn move_image_paths(&self, from: &str, to: &str) -> Result<usize> {
let from = from.trim_end_matches('/');
let to = to.trim_end_matches('/');
Ok(self.conn.lock().unwrap().execute(
"UPDATE images SET path = ?2 || substr(path, length(?1) + 1)
WHERE path = ?1 OR substr(path, 1, length(?1) + 1) = ?1 || '/'",
params![from, to],
)?)
}
/// The file of a generated image was found somewhere else.
pub fn set_image_path(&self, id: &str, path: &str) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE images SET path = ?2 WHERE id = ?1", params![id, path])?;
Ok(())
}
/// Latest images; `user = None` lists everybody's (admins).
/// Generated images of these sessions.
pub fn images_of_sessions(&self, sessions: &[String]) -> Result<Vec<ImageRecord>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(&format!("SELECT {IMAGE_COLS} FROM images WHERE session_id = ?1"))?;
let mut out = Vec::new();
for s in sessions {
out.extend(stmt.query_map([s], image_row)?.collect::<rusqlite::Result<Vec<_>>>()?);
}
Ok(out)
}
pub fn images(&self, user: Option<&str>, limit: i64) -> Result<Vec<ImageRecord>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(&format!(
"SELECT {IMAGE_COLS} FROM images WHERE ?1 IS NULL OR user_id = ?1 ORDER BY created_at DESC LIMIT ?2"
))?;
let rows = stmt.query_map(params![user, limit], image_row)?.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn image(&self, id: &str) -> Result<Option<ImageRecord>> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row(&format!("SELECT {IMAGE_COLS} FROM images WHERE id = ?1"), [id], image_row).optional()?)
}
pub fn delete_image(&self, id: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM images WHERE id = ?1", [id])?;
Ok(())
}
}
impl Store {
/// Cost per provider since `from` (ms).
/// opencode Go: sum of reported allowance of one provider since `from` (ms).
pub fn allowance_since(&self, provider: &str, from: i64) -> Result<f64> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row("SELECT COALESCE(SUM(allowance), 0) FROM usage WHERE provider = ?1 AND ts >= ?2", params![provider, from], |r| r.get(0))?)
}
pub fn cost_by_provider(&self, from: i64) -> Result<std::collections::BTreeMap<String, f64>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT provider, COALESCE(SUM(cost), 0) FROM usage WHERE ts >= ?1 GROUP BY provider")?;
let rows = stmt.query_map([from], |r| Ok((r.get::<_, String>(0)?, r.get::<_, f64>(1)?)))?;
let mut out = std::collections::BTreeMap::new();
for r in rows {
let (p, c) = r?;
out.insert(p, c);
}
Ok(out)
}
}
// --- message queue -------------------------------------------------------------
/// A message waiting to be delivered to a busy session.
#[derive(Debug, Clone, Serialize, serde::Deserialize)]
pub struct QueueItem {
pub id: String,
pub parts: Vec<crate::types::Part>,
/// "user" or "background" (results of background subagents).
pub source: String,
pub created_at: i64,
}
impl Store {
pub fn queue_add(&self, session_id: &str, item: &QueueItem) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO queue (id, session_id, data, created_at) VALUES (?1, ?2, ?3, ?4)",
params![item.id, session_id, serde_json::to_string(item)?, item.created_at],
)?;
Ok(())
}
pub fn queue_list(&self, session_id: &str) -> Result<Vec<QueueItem>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT data FROM queue WHERE session_id = ?1 ORDER BY created_at, rowid")?;
let rows = stmt.query_map([session_id], |r| r.get::<_, String>(0))?;
let mut out = Vec::new();
for r in rows {
out.push(serde_json::from_str(&r?)?);
}
Ok(out)
}
/// Remove one item; false if it was already delivered.
pub fn queue_remove(&self, session_id: &str, id: &str) -> Result<bool> {
let conn = self.conn.lock().unwrap();
Ok(conn.execute("DELETE FROM queue WHERE session_id = ?1 AND id = ?2", params![session_id, id])? > 0)
}
/// Take all items for delivery.
pub fn queue_take(&self, session_id: &str) -> Result<Vec<QueueItem>> {
let items = self.queue_list(session_id)?;
let conn = self.conn.lock().unwrap();
for it in &items {
conn.execute("DELETE FROM queue WHERE id = ?1", [&it.id])?;
}
Ok(items)
}
}
impl Store {
pub fn set_user_quota(&self, id: &str, quota: Option<f64>) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("UPDATE users SET monthly_quota = ?2 WHERE id = ?1", params![id, quota])?;
Ok(())
}
/// Cost of a user's requests since `from` (ms).
pub fn user_cost_since(&self, user_id: &str, from: i64) -> Result<f64> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row("SELECT COALESCE(SUM(cost), 0) FROM usage WHERE user_id = ?1 AND ts >= ?2", params![user_id, from], |r| r.get(0))?)
}
}
// --- audit log & backup ----------------------------------------------------------
impl Store {
// --- metrics ---------------------------------------------------------------
/// All-time usage per provider/model: (provider, model, requests, errors, in, out, cache_read, cache_write, cost, duration_ms).
pub fn usage_totals(&self) -> Result<Vec<(String, String, i64, i64, i64, i64, i64, i64, f64, i64)>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT provider, model, COUNT(*), COALESCE(SUM(error),0), COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0), COALESCE(SUM(cost),0), COALESCE(SUM(duration_ms),0)
FROM usage GROUP BY provider, model",
)?;
let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?, r.get(4)?, r.get(5)?, r.get(6)?, r.get(7)?, r.get(8)?, r.get(9)?)))?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
}
pub fn count(&self, table: &str) -> i64 {
// Table names are fixed strings from the code, never user input.
self.conn.lock().unwrap().query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |r| r.get(0)).unwrap_or(0)
}
// --- scheduled jobs -------------------------------------------------------
fn job_row(r: &rusqlite::Row) -> rusqlite::Result<crate::jobs::Job> {
Ok(crate::jobs::Job {
id: r.get(0)?,
owner_id: r.get(1)?,
path: r.get(2)?,
enabled: r.get::<_, i64>(3)? != 0,
last_run: r.get(4)?,
last_session_id: r.get(5)?,
last_status: r.get(6)?,
created_at: r.get(7)?,
webhook: r.get::<_, Option<String>>(8)?.is_some(),
settings: r.get::<_, Option<String>>(9)?.and_then(|j| serde_json::from_str(&j).ok()),
})
}
pub fn jobs(&self, owner: Option<&str>) -> Result<Vec<crate::jobs::Job>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT id, owner_id, path, enabled, last_run, last_session_id, last_status, created_at, webhook_secret, settings FROM jobs
WHERE ?1 IS NULL OR owner_id = ?1 ORDER BY created_at",
)?;
let rows = stmt.query_map([owner], Self::job_row)?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
}
pub fn job(&self, id: &str) -> Result<Option<crate::jobs::Job>> {
let conn = self.conn.lock().unwrap();
Ok(conn
.query_row("SELECT id, owner_id, path, enabled, last_run, last_session_id, last_status, created_at, webhook_secret, settings FROM jobs WHERE id = ?1", [id], Self::job_row)
.optional()?)
}
/// Register a job file (or return the existing registration of that file).
pub fn job_register(&self, owner: &str, path: &str) -> Result<crate::jobs::Job> {
{
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO jobs (id, owner_id, path, enabled, created_at) VALUES (?1, ?2, ?3, 1, ?4)
ON CONFLICT(path) DO UPDATE SET enabled = 1",
rusqlite::params![crate::types::new_id(), owner, path, now_ms()],
)?;
}
let conn = self.conn.lock().unwrap();
Ok(conn.query_row(
"SELECT id, owner_id, path, enabled, last_run, last_session_id, last_status, created_at, webhook_secret, settings FROM jobs WHERE path = ?1",
[path],
Self::job_row,
)?)
}
pub fn job_started(&self, id: &str, session_id: &str, ts: i64) -> Result<()> {
self.conn.lock().unwrap().execute(
"UPDATE jobs SET last_run = ?2, last_session_id = ?3, last_status = 'running' WHERE id = ?1",
rusqlite::params![id, ts, session_id],
)?;
Ok(())
}
pub fn job_finished(&self, id: &str, session_id: Option<&str>, status: &str) -> Result<()> {
self.conn.lock().unwrap().execute(
"UPDATE jobs SET last_status = ?2, last_session_id = COALESCE(?3, last_session_id) WHERE id = ?1",
rusqlite::params![id, status, session_id],
)?;
Ok(())
}
// --- software factory ---------------------------------------------------------
pub fn factory_issue(&self, repo: &str, number: i64) -> Result<Option<Value>> {
let conn = self.conn.lock().unwrap();
Ok(conn
.query_row("SELECT title, author, state, attempts, rounds, last_comment, scenarios, branch, pr_url, note, created_at, updated_at, owner_id FROM factory_issues WHERE repo_url = ?1 AND number = ?2", params![repo, number], |r| {
Ok(json!({"repo": repo, "number": number, "title": r.get::<_, String>(0)?, "author": r.get::<_, String>(1)?, "state": r.get::<_, String>(2)?, "attempts": r.get::<_, i64>(3)?, "rounds": r.get::<_, i64>(4)?, "last_comment": r.get::<_, i64>(5)?, "scenarios": r.get::<_, String>(6)?, "branch": r.get::<_, String>(7)?, "pr_url": r.get::<_, String>(8)?, "note": r.get::<_, String>(9)?, "created_at": r.get::<_, i64>(10)?, "updated_at": r.get::<_, i64>(11)?, "owner_id": r.get::<_, String>(12)?}))
})
.optional()?)
}
pub fn factory_issue_add(&self, repo: &str, number: i64, owner: &str, title: &str, author: &str) -> Result<()> {
let now = crate::types::now_ms();
self.conn.lock().unwrap().execute(
"INSERT OR IGNORE INTO factory_issues (repo_url, number, owner_id, title, author, created_at, updated_at) VALUES (?1,?2,?3,?4,?5,?6,?6)",
params![repo, number, owner, title, author, now],
)?;
Ok(())
}
/// Update the given columns (whitelisted) of an issue.
pub fn factory_issue_set(&self, repo: &str, number: i64, fields: &[(&str, Value)]) -> Result<()> {
const COLS: [&str; 10] = ["title", "state", "attempts", "rounds", "last_comment", "scenarios", "branch", "pr_url", "note", "author"];
let conn = self.conn.lock().unwrap();
for (k, v) in fields {
if !COLS.contains(k) {
anyhow::bail!("unknown column {k}");
}
let sql = format!("UPDATE factory_issues SET {k} = ?3, updated_at = ?4 WHERE repo_url = ?1 AND number = ?2");
let now = crate::types::now_ms();
match v {
Value::Number(n) => conn.execute(&sql, params![repo, number, n.as_i64().unwrap_or(0), now])?,
other => conn.execute(&sql, params![repo, number, other.as_str().unwrap_or(""), now])?,
};
}
Ok(())
}
pub fn factory_issues(&self, owner: &str) -> Result<Vec<Value>> {
let nums: Vec<(String, i64)> = {
let conn = self.conn.lock().unwrap();
let mut st = conn.prepare("SELECT repo_url, number FROM factory_issues WHERE owner_id = ?1 ORDER BY updated_at DESC")?;
st.query_map([owner], |r| Ok((r.get(0)?, r.get(1)?)))?.filter_map(|r| r.ok()).collect()
};
Ok(nums.into_iter().filter_map(|(r, n)| self.factory_issue(&r, n).ok().flatten()).collect())
}
pub fn factory_run_add(&self, id: &str, owner: &str, repo: &str, issue: i64, phase: &str, round: i64, session: &str, model: &str) -> Result<()> {
self.conn.lock().unwrap().execute(
"INSERT INTO factory_runs (id, owner_id, repo_url, issue, phase, round, session_id, model, started_at) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9)",
params![id, owner, repo, issue, phase, round, session, model, crate::types::now_ms()],
)?;
Ok(())
}
pub fn factory_run_end(&self, id: &str, status: &str, result: &str) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE factory_runs SET ended_at = ?2, status = ?3, result = ?4 WHERE id = ?1", params![id, crate::types::now_ms(), status, result])?;
Ok(())
}
pub fn factory_running(&self, repo: &str) -> Result<i64> {
Ok(self.conn.lock().unwrap().query_row("SELECT COUNT(*) FROM factory_runs WHERE repo_url = ?1 AND status = 'running'", [repo], |r| r.get(0))?)
}
/// Runs with token and cost totals (joined from the usage table).
pub fn factory_runs(&self, owner: &str, repo: Option<&str>, issue: Option<i64>, limit: i64) -> Result<Vec<Value>> {
let conn = self.conn.lock().unwrap();
let mut st = conn.prepare(
"SELECT f.id, f.repo_url, f.issue, f.phase, f.round, f.session_id, f.model, f.started_at, f.ended_at, f.status, f.result,
COALESCE((SELECT SUM(input_tokens + cache_read_tokens + cache_write_tokens) FROM usage WHERE session_id = f.session_id),0),
COALESCE((SELECT SUM(output_tokens) FROM usage WHERE session_id = f.session_id),0),
COALESCE((SELECT SUM(cost) FROM usage WHERE session_id = f.session_id),0.0)
FROM factory_runs f WHERE f.owner_id = ?1 AND (?2 IS NULL OR f.repo_url = ?2) AND (?3 IS NULL OR f.issue = ?3)
ORDER BY f.started_at DESC LIMIT ?4",
)?;
let rows = st.query_map(params![owner, repo, issue, limit], |r| {
Ok(json!({"id": r.get::<_, String>(0)?, "repo": r.get::<_, String>(1)?, "issue": r.get::<_, i64>(2)?, "phase": r.get::<_, String>(3)?, "round": r.get::<_, i64>(4)?, "session_id": r.get::<_, String>(5)?, "model": r.get::<_, String>(6)?, "started_at": r.get::<_, i64>(7)?, "ended_at": r.get::<_, Option<i64>>(8)?, "status": r.get::<_, String>(9)?, "result": r.get::<_, String>(10)?, "input_tokens": r.get::<_, i64>(11)?, "output_tokens": r.get::<_, i64>(12)?, "cost": r.get::<_, f64>(13)?}))
})?;
Ok(rows.filter_map(|r| r.ok()).collect())
}
/// After a restart: runs that were still going are over; their issues go back one step.
pub fn factory_recover(&self) -> Result<usize> {
let conn = self.conn.lock().unwrap();
let n = conn.execute("UPDATE factory_runs SET status = 'interrupted', ended_at = ?1 WHERE status = 'running'", [crate::types::now_ms()])?;
conn.execute("UPDATE factory_issues SET state = 'new' WHERE state = 'review'", [])?;
conn.execute("UPDATE factory_issues SET state = 'approved' WHERE state = 'building'", [])?;
conn.execute("UPDATE factory_issues SET state = 'built' WHERE state = 'qa'", [])?;
Ok(n)
}
pub fn job_touch(&self, id: &str, ts: i64) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE jobs SET last_run = ?2 WHERE id = ?1", rusqlite::params![id, ts])?;
Ok(())
}
pub fn job_set_enabled(&self, id: &str, enabled: bool) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE jobs SET enabled = ?2 WHERE id = ?1", rusqlite::params![id, enabled as i64])?;
Ok(())
}
/// Set (Some) or remove (None) the webhook secret of a job.
pub fn job_set_settings(&self, id: &str, settings: &crate::jobs::JobSettings) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE jobs SET settings = ?2 WHERE id = ?1", rusqlite::params![id, serde_json::to_string(settings)?])?;
Ok(())
}
pub fn job_set_webhook(&self, id: &str, secret: Option<&str>) -> Result<()> {
self.conn.lock().unwrap().execute("UPDATE jobs SET webhook_secret = ?2 WHERE id = ?1", rusqlite::params![id, secret])?;
Ok(())
}
pub fn job_webhook_secret(&self, id: &str) -> Result<Option<String>> {
Ok(self.conn.lock().unwrap().query_row("SELECT webhook_secret FROM jobs WHERE id = ?1", [id], |r| r.get(0)).optional()?.flatten())
}
pub fn job_delete(&self, id: &str) -> Result<()> {
self.conn.lock().unwrap().execute("DELETE FROM jobs WHERE id = ?1", [id])?;
Ok(())
}
// --- per-user settings (JSON) ------------------------------------------
pub fn user_setting<T: serde::de::DeserializeOwned + Default>(&self, user_id: &str, key: &str) -> Result<T> {
let conn = self.conn.lock().unwrap();
let v: Option<String> = conn
.query_row("SELECT value FROM user_settings WHERE user_id = ?1 AND key = ?2", [user_id, key], |r| r.get(0))
.optional()?;
Ok(v.and_then(|v| serde_json::from_str(&v).ok()).unwrap_or_default())
}
pub fn set_user_setting<T: serde::Serialize>(&self, user_id: &str, key: &str, value: &T) -> Result<()> {
self.conn.lock().unwrap().execute(
"INSERT INTO user_settings (user_id, key, value) VALUES (?1, ?2, ?3)
ON CONFLICT(user_id, key) DO UPDATE SET value = ?3",
rusqlite::params![user_id, key, serde_json::to_string(value)?],
)?;
Ok(())
}
// --- web push ---------------------------------------------------------
pub fn push_add(&self, user_id: &str, endpoint: &str, p256dh: &str, auth: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO push_subs (endpoint, user_id, p256dh, auth, created_at) VALUES (?1, ?2, ?3, ?4, ?5)
ON CONFLICT(endpoint) DO UPDATE SET user_id = ?2, p256dh = ?3, auth = ?4",
rusqlite::params![endpoint, user_id, p256dh, auth, now_ms()],
)?;
Ok(())
}
pub fn push_remove(&self, endpoint: &str) -> Result<()> {
self.conn.lock().unwrap().execute("DELETE FROM push_subs WHERE endpoint = ?1", [endpoint])?;
Ok(())
}
pub fn push_subscriptions(&self, user_id: &str) -> Result<Vec<PushSub>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT endpoint, p256dh, auth FROM push_subs WHERE user_id = ?1")?;
let rows = stmt.query_map([user_id], |r| Ok(PushSub { endpoint: r.get(0)?, p256dh: r.get(1)?, auth: r.get(2)? }))?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
}
pub fn audit(&self, user_id: Option<&str>, username: &str, action: &str, detail: &str, ip: &str) {
let conn = self.conn.lock().unwrap();
let _ = conn.execute(
"INSERT INTO audit (ts, user_id, username, action, detail, ip) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
params![now_ms(), user_id, username, action, detail, ip],
);
}
pub fn audit_list(&self, limit: i64, action: Option<&str>) -> Result<Vec<Value>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT ts, username, action, detail, ip FROM audit WHERE ?2 IS NULL OR action LIKE ?2 || '%' ORDER BY id DESC LIMIT ?1",
)?;
let rows = stmt
.query_map(params![limit, action], |r| {
Ok(json!({"ts": r.get::<_, i64>(0)?, "user": r.get::<_, String>(1)?, "action": r.get::<_, String>(2)?, "detail": r.get::<_, String>(3)?, "ip": r.get::<_, String>(4)?}))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
/// Consistent copy of the database (safe while the server runs).
pub fn snapshot(&self, target: &std::path::Path) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("VACUUM INTO ?1", [target.to_string_lossy()])?;
Ok(())
}
}
// --- project overview ------------------------------------------------------------------------
impl Store {
/// Per chat: files the agent changed, the latest check result and the cost.
pub fn project_overview(&self, session_ids: &[String]) -> Result<Value> {
let conn = self.conn.lock().unwrap();
let mut out = serde_json::Map::new();
for id in session_ids {
let files: i64 = conn.query_row("SELECT COUNT(DISTINCT path) FROM checkpoints WHERE session_id = ?1", params![id], |r| r.get(0))?;
let check: Option<String> = conn
.query_row("SELECT status FROM turn_checks WHERE session_id = ?1 ORDER BY ts DESC LIMIT 1", params![id], |r| r.get(0))
.optional()?;
let failed: i64 = conn.query_row("SELECT COUNT(*) FROM turn_checks WHERE session_id = ?1 AND status IN ('failed','timeout')", params![id], |r| r.get(0))?;
let cost: Option<f64> = conn.query_row("SELECT SUM(cost) FROM usage WHERE root_session_id = ?1", params![id], |r| r.get(0))?;
out.insert(id.clone(), json!({"files": files, "check": check, "checks_failed": failed, "cost": cost.unwrap_or(0.0)}));
}
Ok(Value::Object(out))
}
}
// --- weekly overview ----------------------------------------------------------------------
impl Store {
/// (chats with activity, failed checks, checks run) of a user's own top-level chats since `since` (ms).
pub fn weekly_counts(&self, user_id: &str, since: i64) -> Result<(i64, i64, i64)> {
let conn = self.conn.lock().unwrap();
let chats: i64 = conn.query_row("SELECT COUNT(*) FROM sessions WHERE owner_id = ?1 AND parent_id IS NULL AND updated_at >= ?2", params![user_id, since], |r| r.get(0))?;
let (failed, total): (i64, i64) = conn.query_row(
"SELECT COALESCE(SUM(c.status IN ('failed','timeout')),0), COUNT(*) FROM turn_checks c JOIN sessions s ON s.id = c.session_id WHERE s.owner_id = ?1 AND c.ts >= ?2",
params![user_id, since],
|r| Ok((r.get(0)?, r.get(1)?)),
)?;
Ok((chats, failed, total))
}
}
// --- read-only links ----------------------------------------------------------------------
impl Store {
pub fn add_share_link(&self, token: &str, session_id: &str, owner_id: &str, created_at: i64, expires_at: i64) -> Result<()> {
self.conn.lock().unwrap().execute(
"INSERT INTO share_links (token, session_id, owner_id, created_at, expires_at) VALUES (?1, ?2, ?3, ?4, ?5)",
params![token, session_id, owner_id, created_at, expires_at],
)?;
Ok(())
}
/// (token, created_at, expires_at) of the links of a chat, newest first.
pub fn share_links(&self, session_id: &str) -> Result<Vec<(String, i64, i64)>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT token, created_at, expires_at FROM share_links WHERE session_id = ?1 ORDER BY created_at DESC")?;
let rows = stmt.query_map([session_id], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)))?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
}
/// The chat a link points to while it is valid.
pub fn share_link_session(&self, token: &str, now: i64) -> Result<Option<String>> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row("SELECT session_id FROM share_links WHERE token = ?1 AND expires_at > ?2", params![token, now], |r| r.get(0)).optional()?)
}
pub fn delete_share_link(&self, session_id: &str, token: &str) -> Result<()> {
self.conn.lock().unwrap().execute("DELETE FROM share_links WHERE session_id = ?1 AND token = ?2", params![session_id, token])?;
Ok(())
}
/// Links are dropped with their chat and when expired.
pub fn purge_share_links(&self, now: i64) -> Result<()> {
self.conn.lock().unwrap().execute("DELETE FROM share_links WHERE expires_at <= ?1 OR session_id NOT IN (SELECT id FROM sessions)", params![now])?;
Ok(())
}
}
// --- defaults of new chats in a project -------------------------------------------------------
impl Store {
/// {model, agent, permission_mode, color, instructions, folder} new chats of the folder start with.
pub fn project_defaults(&self, cwd: &str) -> Result<Option<Value>> {
let conn = self.conn.lock().unwrap();
let v: Option<String> = conn.query_row("SELECT data FROM project_defaults WHERE cwd = ?1", params![cwd], |r| r.get(0)).optional()?;
Ok(v.and_then(|v| serde_json::from_str(&v).ok()))
}
pub fn set_project_defaults(&self, cwd: &str, data: Option<&Value>) -> Result<()> {
let conn = self.conn.lock().unwrap();
match data {
Some(d) => conn.execute("INSERT INTO project_defaults (cwd, data) VALUES (?1, ?2) ON CONFLICT(cwd) DO UPDATE SET data = ?2", params![cwd, d.to_string()])?,
None => conn.execute("DELETE FROM project_defaults WHERE cwd = ?1", params![cwd])?,
};
Ok(())
}
}
// --- checks after a turn ------------------------------------------------------------------
impl Store {
/// (command, enabled) of a project folder.
pub fn project_check(&self, cwd: &str) -> Result<Option<(String, bool)>> {
let conn = self.conn.lock().unwrap();
Ok(conn.query_row("SELECT command, enabled FROM project_checks WHERE cwd = ?1", params![cwd], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)? != 0))).optional()?)
}
pub fn set_project_check(&self, cwd: &str, command: &str, enabled: bool) -> Result<()> {
let conn = self.conn.lock().unwrap();
if command.trim().is_empty() {
conn.execute("DELETE FROM project_checks WHERE cwd = ?1", params![cwd])?;
} else {
conn.execute(
"INSERT INTO project_checks (cwd, command, enabled) VALUES (?1, ?2, ?3) ON CONFLICT(cwd) DO UPDATE SET command = ?2, enabled = ?3",
params![cwd, command.trim(), enabled as i64],
)?;
}
Ok(())
}
pub fn set_turn_check(&self, session_id: &str, message_id: &str, status: &str, exit_code: Option<i32>, duration_ms: u64, output: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO turn_checks (session_id, message_id, status, exit_code, duration_ms, output, ts) VALUES (?1,?2,?3,?4,?5,?6,?7)
ON CONFLICT(session_id, message_id) DO UPDATE SET status = ?3, exit_code = ?4, duration_ms = ?5, output = ?6, ts = ?7",
params![session_id, message_id, status, exit_code, duration_ms as i64, output, now_ms()],
)?;
Ok(())
}
pub fn turn_checks(&self, session_id: &str) -> Result<Vec<Value>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT message_id, status, exit_code, duration_ms, output FROM turn_checks WHERE session_id = ?1 ORDER BY ts DESC LIMIT 300")?;
let rows = stmt
.query_map(params![session_id], |r| {
Ok(json!({"message_id": r.get::<_, String>(0)?, "status": r.get::<_, String>(1)?, "exit_code": r.get::<_, Option<i32>>(2)?, "duration_ms": r.get::<_, i64>(3)?, "output": r.get::<_, String>(4)?}))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
}
// --- project-wide approvals ---------------------------------------------------------
impl Store {
/// Approvals ("bash:<prefix>" or a tool name) remembered for a project folder.
pub fn project_approvals(&self, cwd: &str) -> Result<Vec<String>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare("SELECT key FROM project_approvals WHERE cwd = ?1 ORDER BY key")?;
let rows = stmt.query_map(params![cwd], |r| r.get::<_, String>(0))?.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn add_project_approval(&self, cwd: &str, key: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("INSERT OR IGNORE INTO project_approvals (cwd, key, created_at) VALUES (?1, ?2, ?3)", params![cwd, key, now_ms()])?;
Ok(())
}
pub fn delete_project_approval(&self, cwd: &str, key: &str) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM project_approvals WHERE cwd = ?1 AND key = ?2", params![cwd, key])?;
Ok(())
}
}
// --- file checkpoints ------------------------------------------------------------
impl Store {
/// Remember a file's content before the agent changes it.
pub fn add_checkpoint(&self, session_id: &str, path: &str, content: Option<&[u8]>) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute(
"INSERT INTO checkpoints (session_id, ts, path, existed, content) VALUES (?1, ?2, ?3, ?4, ?5)",
params![session_id, now_ms(), path, content.is_some() as i64, content],
)?;
Ok(())
}
/// Per path, the oldest snapshot taken at or after `since` (ms): the state to restore.
pub fn checkpoints_since(&self, session_id: &str, since: i64) -> Result<Vec<(String, Option<Vec<u8>>)>> {
self.checkpoints_between(session_id, since, i64::MAX)
}
/// Per file the state before its first change in the time range [since, until).
pub fn checkpoints_between(&self, session_id: &str, since: i64, until: i64) -> Result<Vec<(String, Option<Vec<u8>>)>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT path, existed, content FROM checkpoints c WHERE session_id = ?1 AND ts >= ?2 AND ts < ?3
AND id = (SELECT MIN(id) FROM checkpoints WHERE session_id = ?1 AND ts >= ?2 AND ts < ?3 AND path = c.path)",
)?;
let rows = stmt
.query_map(params![session_id, since, until], |r| {
let existed: i64 = r.get(1)?;
Ok((r.get::<_, String>(0)?, if existed != 0 { r.get::<_, Option<Vec<u8>>>(2)? } else { None }))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
pub fn delete_checkpoints_since(&self, session_id: &str, since: i64) -> Result<()> {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM checkpoints WHERE session_id = ?1 AND ts >= ?2", params![session_id, since])?;
Ok(())
}
pub fn message_created_at(&self, session_id: &str, message_id: &str) -> Result<Option<i64>> {
let conn = self.conn.lock().unwrap();
let data: Option<String> = conn
.query_row("SELECT data FROM messages WHERE session_id = ?1 AND id = ?2", params![session_id, message_id], |r| r.get(0))
.optional()?;
Ok(data.and_then(|d| serde_json::from_str::<Value>(&d).ok()).and_then(|v| v["created_at"].as_i64()))
}
}
pub struct PushSub {
pub endpoint: String,
pub p256dh: String,
pub auth: String,
}
/// Which messages of a session the chat view loads (see Store::messages_window).
pub enum MessageWindow {
Tail { limit: i64 },
Before { id: String, limit: i64, stop_at_summary: bool },
After { id: String, limit: i64 },
Around { id: String, limit: i64 },
}
/// Text around the first occurrence of `needle` (lowercase) in the visible text of a
/// message: its text, reasoning, summaries, tool calls and tool output.
fn snippet_in(m: &Message, needle: &str) -> Option<String> {
for p in &m.parts {
let owned;
let text = match p {
Part::Text { text } | Part::Summary { text, .. } | Part::Reasoning { text, .. } => text.as_str(),
Part::ToolResult { content, .. } => content.as_str(),
Part::ToolCall { input, .. } => {
owned = input.to_string();
owned.as_str()
}
_ => continue,
};
let lower = text.to_lowercase();
// Lowercasing can change byte lengths (rare): then cut the lowercase text.
let (src, i) = match lower.find(needle) {
Some(i) if lower.len() == text.len() => (text, i),
Some(i) => (lower.as_str(), i),
None => continue,
};
let mut a = i.saturating_sub(60);
while !src.is_char_boundary(a) {
a -= 1;
}
let mut b = (i + needle.len() + 100).min(src.len());
while !src.is_char_boundary(b) {
b += 1;
}
return Some(format!("{}{}{}", if a > 0 { "…" } else { "" }, src[a..b].replace('\n', " "), if b < src.len() { "…" } else { "" }));
}
None
}
#[cfg(test)]
mod link_tests {
use super::*;
#[test]
fn share_links_expire_and_can_be_revoked() {
let dir = std::env::temp_dir().join(format!("cc-links-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let store = Store::open(&dir.join("t.db")).unwrap();
store.add_share_link("tok1", "s1", "u", 1_000, 5_000).unwrap();
store.add_share_link("tok2", "s1", "u", 2_000, 9_000).unwrap();
// Valid until the expiry time, then gone.
assert_eq!(store.share_link_session("tok1", 4_999).unwrap().as_deref(), Some("s1"));
assert_eq!(store.share_link_session("tok1", 5_000).unwrap(), None);
assert_eq!(store.share_link_session("nope", 1).unwrap(), None);
// Revoking removes only the named link of that chat.
store.delete_share_link("other", "tok2").unwrap();
assert!(store.share_link_session("tok2", 3_000).unwrap().is_some());
store.delete_share_link("s1", "tok2").unwrap();
assert_eq!(store.share_link_session("tok2", 3_000).unwrap(), None);
std::fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn weekly_counts_of_an_empty_user_are_zero() {
let dir = std::env::temp_dir().join(format!("cc-weekly-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let store = Store::open(&dir.join("t.db")).unwrap();
assert_eq!(store.weekly_counts("nobody", 0).unwrap(), (0, 0, 0));
std::fs::remove_dir_all(&dir).unwrap();
}
}
Reference in new issue View Git Blame Copy Permalink