Richard ZinkandClaude Opus 5.5
1e8418336d
- opencode Go reports plan requests as cost 0 with the usage in allowance;
ModelConfig::cost values them at the model's list prices (charged amounts
and free requests stay as reported). Rows booked at 0 are valued once on
start (idempotent).
- opencode's price list now has per-duration cache prices for some models
({"30m": ..., "12h": ...}): the shortest duration is used, also in tiers.
Saving the provider's models failed with "untagged enum ModelEntry".
- Each answer shows tokens sent/received and the cost of the turn's requests
(usage.cost filled when messages are saved and loaded).
- The chat title keeps a minimum width in a crowded top bar.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1575 lines
69 KiB
Rust
1575 lines
69 KiB
Rust
//! 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 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(())
|
||
}
|
||
|
||
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")?;
|
||
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_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(())
|
||
}
|
||
|
||
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(())
|
||
}
|
||
}
|
||
|
||
// --- 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>>)>> {
|
||
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 id = (SELECT MIN(id) FROM checkpoints WHERE session_id = ?1 AND ts >= ?2 AND path = c.path)",
|
||
)?;
|
||
let rows = stmt
|
||
.query_map(params![session_id, since], |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
|
||
}
|