Richard ZinkandClaude Opus 5.5
c850586a81
- Terminals and file managers as tabs of a chat; chat terminals keep running on the server (no idle timeout) and re-attach from any device - "Just chat" in New session and a built-in chat mode with a lean prompt and only web tools; custom modes with kind (coding/conversation), read-only and tool picker - Mode selector as a dropdown - The chat shows running compactions (automatic: tokens and share of the limit) and labels summaries - Docker image builds and serves the terminal client binaries again - Fix: atomic session updates; background title/usage writes no longer undo each other (auto compaction triggered one turn late) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1211 lines
51 KiB
Rust
1211 lines
51 KiB
Rust
//! SQLite persistence for sessions and messages.
|
||
|
||
use crate::types::{Message, Part, 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);
|
||
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"),
|
||
("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 mut snippet = String::new();
|
||
if let Some(m) = mdata.and_then(|d| serde_json::from_str::<Message>(&d).ok()) {
|
||
for p in &m.parts {
|
||
let text = match p {
|
||
Part::Text { text } | Part::Summary { text, .. } | Part::Reasoning { text, .. } => text.as_str(),
|
||
Part::ToolResult { content, .. } => content.as_str(),
|
||
_ => continue,
|
||
};
|
||
if let Some(i) = text.to_lowercase().find(&needle) {
|
||
let mut a = i.saturating_sub(60);
|
||
while !text.is_char_boundary(a) {
|
||
a -= 1;
|
||
}
|
||
let mut b = (i + needle.len() + 100).min(text.len());
|
||
while !text.is_char_boundary(b) {
|
||
b += 1;
|
||
}
|
||
snippet = format!("{}{}{}", if a > 0 { "…" } else { "" }, text[a..b].replace('\n', " "), if b < text.len() { "…" } else { "" });
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
if snippet.is_empty() && !session.title.to_lowercase().contains(&needle) {
|
||
continue; // matched only JSON structure (keys, ids)
|
||
}
|
||
out.push((session, snippet));
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
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<()> {
|
||
let conn = self.conn.lock().unwrap();
|
||
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)
|
||
}
|
||
|
||
/// 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(())
|
||
}
|
||
}
|
||
|
||
/// 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}))
|
||
}
|
||
|
||
/// 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,
|
||
}
|
||
|
||
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)?,
|
||
})
|
||
}
|
||
|
||
const IMAGE_COLS: &str = "id, user_id, session_id, prompt, model, path, width, height, created_at";
|
||
|
||
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)"),
|
||
params![img.id, img.user_id, img.session_id, img.prompt, img.model, img.path, img.width, img.height, img.created_at],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Latest images; `user = None` lists everybody's (admins).
|
||
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(),
|
||
})
|
||
}
|
||
|
||
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 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 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 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_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,
|
||
}
|