1use std::io::Write;
2use std::path::PathBuf;
3use std::sync::Mutex;
4
5use alex_core::{Pricing, TraceRecord};
6use anyhow::{Context, Result};
7use chrono::Utc;
8use flate2::write::GzEncoder;
9use flate2::Compression;
10use rusqlite::{params, Connection, OptionalExtension};
11use serde_json::{json, Value};
12
13const SCHEMA: &str = r#"
14CREATE TABLE IF NOT EXISTS traces (
15 id TEXT PRIMARY KEY,
16 ts_request_ms INTEGER NOT NULL,
17 ts_response_ms INTEGER,
18 session_id TEXT,
19 harness TEXT,
20 client_format TEXT,
21 upstream_provider TEXT,
22 upstream_format TEXT,
23 requested_model TEXT,
24 routed_model TEXT,
25 method TEXT,
26 path TEXT,
27 status INTEGER,
28 streamed INTEGER,
29 input_tokens INTEGER,
30 cached_input_tokens INTEGER,
31 cache_creation_tokens INTEGER,
32 output_tokens INTEGER,
33 reasoning_tokens INTEGER,
34 cost_usd REAL,
35 billing_bucket TEXT,
36 req_body_path TEXT,
37 upstream_req_body_path TEXT,
38 resp_body_path TEXT,
39 req_headers_json TEXT,
40 resp_headers_json TEXT,
41 error TEXT,
42 account_id TEXT,
43 run_id TEXT,
44 tags_json TEXT,
45 client_ip TEXT,
46 key_fingerprint TEXT,
47 reasoning_effort TEXT,
48 thinking_budget INTEGER
49);
50CREATE INDEX IF NOT EXISTS traces_session ON traces(session_id);
51CREATE INDEX IF NOT EXISTS traces_ts ON traces(ts_request_ms);
52CREATE INDEX IF NOT EXISTS traces_model ON traces(routed_model);
53
54CREATE TABLE IF NOT EXISTS pricing (
55 model TEXT PRIMARY KEY,
56 input_per_m REAL, cached_input_per_m REAL,
57 cache_creation_per_m REAL, output_per_m REAL
58);
59
60CREATE TABLE IF NOT EXISTS heartbeats (
61 ts_ms INTEGER NOT NULL,
62 provider TEXT NOT NULL,
63 account_id TEXT,
64 ok INTEGER,
65 status INTEGER,
66 latency_ms INTEGER,
67 message TEXT
68);
69CREATE INDEX IF NOT EXISTS heartbeats_ts ON heartbeats(ts_ms);
70
71CREATE TABLE IF NOT EXISTS run_keys (
72 id TEXT PRIMARY KEY,
73 key_hash TEXT UNIQUE NOT NULL,
74 kind TEXT NOT NULL DEFAULT 'run',
75 run_id TEXT,
76 tags_json TEXT,
77 label TEXT,
78 created_ms INTEGER NOT NULL,
79 expires_ms INTEGER,
80 revoked INTEGER DEFAULT 0,
81 use_count INTEGER DEFAULT 0,
82 last_used_ms INTEGER
83);
84"#;
85
86const RUN_KEY_COLS: &str =
87 "id, key_hash, run_id, tags_json, label, created_ms, expires_ms, revoked, use_count, last_used_ms, kind";
88
89fn run_key_row_json(r: &rusqlite::Row) -> rusqlite::Result<Value> {
90 let key_hash = r.get::<_, String>(1)?;
91 let tags = r
92 .get::<_, Option<String>>(3)?
93 .and_then(|s| serde_json::from_str::<Value>(&s).ok())
94 .filter(|v| v.is_object())
95 .unwrap_or_else(|| json!({}));
96 Ok(json!({
97 "id": r.get::<_, String>(0)?,
98 "key_fingerprint": key_hash.chars().take(16).collect::<String>(),
99 "run_id": r.get::<_, Option<String>>(2)?,
100 "tags": tags,
101 "label": r.get::<_, Option<String>>(4)?,
102 "created_ms": r.get::<_, i64>(5)?,
103 "expires_ms": r.get::<_, Option<i64>>(6)?,
104 "revoked": r.get::<_, i64>(7)? != 0,
105 "use_count": r.get::<_, i64>(8)?,
106 "last_used_ms": r.get::<_, Option<i64>>(9)?,
107 "kind": r.get::<_, String>(10)?,
108 }))
109}
110
111const TRACE_COLS: &str = "id, ts_request_ms, ts_response_ms, harness, client_format, upstream_provider,
112 requested_model, routed_model, status, streamed,
113 input_tokens, cached_input_tokens, cache_creation_tokens, output_tokens, reasoning_tokens,
114 cost_usd, billing_bucket, error, session_id, resp_body_path,
115 upstream_format, req_body_path, upstream_req_body_path, req_headers_json, resp_headers_json,
116 account_id, run_id, tags_json, client_ip, key_fingerprint, reasoning_effort, thinking_budget";
117
118fn trace_row_json(r: &rusqlite::Row) -> rusqlite::Result<Value> {
119 let ts_request_ms = r.get::<_, i64>(1)?;
120 let ts_response_ms = r.get::<_, Option<i64>>(2)?;
121 Ok(json!({
122 "id": r.get::<_, String>(0)?,
123 "ts_request_ms": ts_request_ms,
124 "ts_response_ms": ts_response_ms,
125 "harness": r.get::<_, Option<String>>(3)?,
126 "client_format": r.get::<_, Option<String>>(4)?,
127 "upstream_provider": r.get::<_, Option<String>>(5)?,
128 "requested_model": r.get::<_, Option<String>>(6)?,
129 "routed_model": r.get::<_, Option<String>>(7)?,
130 "status": r.get::<_, Option<i64>>(8)?,
131 "streamed": r.get::<_, Option<i64>>(9)?,
132 "input_tokens": r.get::<_, Option<i64>>(10)?,
133 "cached_input_tokens": r.get::<_, Option<i64>>(11)?,
134 "cache_creation_tokens": r.get::<_, Option<i64>>(12)?,
135 "output_tokens": r.get::<_, Option<i64>>(13)?,
136 "reasoning_tokens": r.get::<_, Option<i64>>(14)?,
137 "cost_usd": r.get::<_, Option<f64>>(15)?,
138 "billing_bucket": r.get::<_, Option<String>>(16)?,
139 "error": r.get::<_, Option<String>>(17)?,
140 "session_id": r.get::<_, Option<String>>(18)?,
141 "resp_body_path": r.get::<_, Option<String>>(19)?,
142 "upstream_format": r.get::<_, Option<String>>(20)?,
143 "req_body_path": r.get::<_, Option<String>>(21)?,
144 "upstream_req_body_path": r.get::<_, Option<String>>(22)?,
145 "req_headers_json": r.get::<_, Option<String>>(23)?,
146 "resp_headers_json": r.get::<_, Option<String>>(24)?,
147 "account_id": r.get::<_, Option<String>>(25)?,
148 "run_id": r.get::<_, Option<String>>(26)?,
149 "tags_json": r.get::<_, Option<String>>(27)?,
150 "client_ip": r.get::<_, Option<String>>(28)?,
151 "key_fingerprint": r.get::<_, Option<String>>(29)?,
152 "reasoning_effort": r.get::<_, Option<String>>(30)?,
153 "thinking_budget": r.get::<_, Option<i64>>(31)?,
154 "latency_ms": ts_response_ms.map(|t| t - ts_request_ms),
155 }))
156}
157
158const DEFAULT_SEARCH_LIMIT: usize = 200;
159const MAX_SEARCH_LIMIT: usize = 5000;
160
161fn effective_limit(limit: usize) -> usize {
162 if limit == 0 {
163 DEFAULT_SEARCH_LIMIT
164 } else {
165 limit.min(MAX_SEARCH_LIMIT)
166 }
167}
168
169fn migrate_traces(conn: &Connection) -> Result<()> {
170 for col in [
171 "run_id TEXT",
172 "tags_json TEXT",
173 "client_ip TEXT",
174 "key_fingerprint TEXT",
175 "reasoning_effort TEXT",
176 "thinking_budget INTEGER",
177 ] {
178 if let Err(e) = conn.execute_batch(&format!("ALTER TABLE traces ADD COLUMN {col}")) {
179 if !e.to_string().contains("duplicate column name") {
180 return Err(e.into());
181 }
182 }
183 }
184 conn.execute_batch("CREATE INDEX IF NOT EXISTS traces_run ON traces(run_id)")?;
185 Ok(())
186}
187
188fn migrate_run_keys(conn: &Connection) -> Result<()> {
189 if let Err(e) = conn.execute_batch("ALTER TABLE run_keys ADD COLUMN kind TEXT NOT NULL DEFAULT 'run'") {
190 if !e.to_string().contains("duplicate column name") {
191 return Err(e.into());
192 }
193 }
194 Ok(())
195}
196
197#[derive(Debug, Clone)]
198pub struct TraceFilter {
199 pub since_ms: Option<i64>,
200 pub until_ms: Option<i64>,
201 pub run_id: Option<String>,
202 pub session: Option<String>,
203 pub model: Option<String>,
204 pub provider: Option<String>,
205 pub path: Option<String>,
206 pub harness: Option<String>,
207 pub status: Option<i64>,
208 pub errors_only: bool,
209 pub key_fingerprint: Option<String>,
210 pub reasoning_effort: Option<String>,
211 pub limit: usize,
212}
213
214impl Default for TraceFilter {
215 fn default() -> Self {
216 Self {
217 since_ms: None,
218 until_ms: None,
219 run_id: None,
220 session: None,
221 model: None,
222 provider: None,
223 path: None,
224 harness: None,
225 status: None,
226 errors_only: false,
227 key_fingerprint: None,
228 reasoning_effort: None,
229 limit: DEFAULT_SEARCH_LIMIT,
230 }
231 }
232}
233
234#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize)]
235pub struct PruneReport {
236 pub bodies_deleted: u64,
237 pub bytes_freed: u64,
238 pub rows_affected: u64,
239 pub rows_deleted: u64,
240 pub dirs_removed: u64,
241}
242
243fn date_dir_name(name: &str) -> bool {
244 name.len() == 10
245 && name.bytes().enumerate().all(|(i, b)| match i {
246 4 | 7 => b == b'-',
247 _ => b.is_ascii_digit(),
248 })
249}
250
251pub struct Store {
252 conn: Mutex<Connection>,
253 pub data_dir: PathBuf,
254}
255
256impl Store {
257 pub fn open(data_dir: PathBuf) -> Result<Self> {
258 std::fs::create_dir_all(&data_dir)?;
259 let db_path = data_dir.join("alexandria.sqlite3");
260 let conn =
261 Connection::open(&db_path).with_context(|| format!("opening sqlite at {db_path:?}"))?;
262 conn.pragma_update(None, "journal_mode", "WAL")?;
263 conn.execute_batch(SCHEMA)?;
264 migrate_traces(&conn)?;
265 migrate_run_keys(&conn)?;
266 seed_pricing(&conn)?;
267 Ok(Self {
268 conn: Mutex::new(conn),
269 data_dir,
270 })
271 }
272
273 pub fn pricing_for(&self, model: &str) -> Option<Pricing> {
274 let conn = self.conn.lock().unwrap();
275 let mut stmt = conn
276 .prepare("SELECT model, input_per_m, cached_input_per_m, cache_creation_per_m, output_per_m FROM pricing")
277 .ok()?;
278 let rows: Vec<(String, Pricing)> = stmt
279 .query_map([], |r| {
280 Ok((
281 r.get::<_, String>(0)?,
282 Pricing {
283 input_per_m: r.get(1)?,
284 cached_input_per_m: r.get(2)?,
285 cache_creation_per_m: r.get(3)?,
286 output_per_m: r.get(4)?,
287 },
288 ))
289 })
290 .ok()?
291 .filter_map(|r| r.ok())
292 .collect();
293 rows.iter()
294 .filter(|(key, _)| model.starts_with(key.as_str()))
295 .max_by_key(|(key, _)| key.len())
296 .map(|(_, p)| p.clone())
297 }
298
299 pub fn pricing_models(&self) -> Vec<String> {
300 let conn = self.conn.lock().unwrap();
301 let Ok(mut stmt) = conn.prepare("SELECT model FROM pricing ORDER BY model") else {
302 return vec![];
303 };
304 stmt.query_map([], |r| r.get::<_, String>(0))
305 .map(|rows| rows.filter_map(|r| r.ok()).collect())
306 .unwrap_or_default()
307 }
308
309 pub fn insert_trace(&self, t: &TraceRecord) -> Result<()> {
310 let conn = self.conn.lock().unwrap();
311 conn.execute(
312 r#"INSERT OR REPLACE INTO traces (
313 id, ts_request_ms, ts_response_ms, session_id, harness, client_format,
314 upstream_provider, upstream_format, requested_model, routed_model,
315 method, path, status, streamed,
316 input_tokens, cached_input_tokens, cache_creation_tokens, output_tokens, reasoning_tokens,
317 cost_usd, billing_bucket,
318 req_body_path, upstream_req_body_path, resp_body_path,
319 req_headers_json, resp_headers_json, error, account_id,
320 run_id, tags_json, client_ip, key_fingerprint, reasoning_effort, thinking_budget
321 ) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16,?17,?18,?19,?20,?21,?22,?23,?24,?25,?26,?27,?28,?29,?30,?31,?32,?33,?34)"#,
322 params![
323 t.id,
324 t.ts_request_ms,
325 t.ts_response_ms,
326 t.session_id,
327 t.harness,
328 t.client_format,
329 t.upstream_provider,
330 t.upstream_format,
331 t.requested_model,
332 t.routed_model,
333 t.method,
334 t.path,
335 t.status,
336 t.streamed.map(|b| b as i64),
337 t.usage.input_tokens,
338 t.usage.cached_input_tokens,
339 t.usage.cache_creation_tokens,
340 t.usage.output_tokens,
341 t.usage.reasoning_tokens,
342 t.cost_usd,
343 t.billing_bucket,
344 t.req_body_path,
345 t.upstream_req_body_path,
346 t.resp_body_path,
347 t.req_headers_json,
348 t.resp_headers_json,
349 t.error,
350 t.account_id,
351 t.run_id,
352 t.tags,
353 t.client_ip,
354 t.key_fingerprint,
355 t.reasoning_effort,
356 t.thinking_budget,
357 ],
358 )?;
359 Ok(())
360 }
361
362 pub fn list_traces(
363 &self,
364 limit: usize,
365 session: Option<&str>,
366 model: Option<&str>,
367 ) -> Result<Vec<Value>> {
368 let filter = TraceFilter {
369 session: session.map(String::from),
370 model: model.map(String::from),
371 limit,
372 ..Default::default()
373 };
374 self.search_traces(&filter)
375 }
376
377 pub fn search_traces(&self, f: &TraceFilter) -> Result<Vec<Value>> {
378 let conn = self.conn.lock().unwrap();
379 let mut sql = format!("SELECT {TRACE_COLS} FROM traces WHERE 1=1");
380 let mut args: Vec<String> = vec![];
381 if let Some(since) = f.since_ms {
382 sql.push_str(" AND ts_request_ms >= ?");
383 args.push(since.to_string());
384 }
385 if let Some(until) = f.until_ms {
386 sql.push_str(" AND ts_request_ms <= ?");
387 args.push(until.to_string());
388 }
389 if let Some(r) = &f.run_id {
390 sql.push_str(" AND run_id = ?");
391 args.push(r.clone());
392 }
393 if let Some(s) = &f.session {
394 sql.push_str(" AND session_id = ?");
395 args.push(s.clone());
396 }
397 if let Some(m) = &f.model {
398 sql.push_str(" AND routed_model LIKE ?");
399 args.push(format!("%{m}%"));
400 }
401 if let Some(p) = &f.provider {
402 sql.push_str(" AND upstream_provider = ?");
403 args.push(p.clone());
404 }
405 if let Some(p) = &f.path {
406 sql.push_str(" AND path = ?");
407 args.push(p.clone());
408 }
409 if let Some(h) = &f.harness {
410 sql.push_str(" AND harness LIKE ?");
411 args.push(format!("%{h}%"));
412 }
413 if let Some(s) = f.status {
414 sql.push_str(" AND status = ?");
415 args.push(s.to_string());
416 }
417 if f.errors_only {
418 sql.push_str(" AND error IS NOT NULL");
419 }
420 if let Some(k) = &f.key_fingerprint {
421 sql.push_str(" AND key_fingerprint = ?");
422 args.push(k.clone());
423 }
424 if let Some(e) = &f.reasoning_effort {
425 sql.push_str(" AND reasoning_effort = ?");
426 args.push(e.clone());
427 }
428 sql.push_str(" ORDER BY ts_request_ms DESC LIMIT ?");
429 args.push(effective_limit(f.limit).to_string());
430 let mut stmt = conn.prepare(&sql)?;
431 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), trace_row_json)?;
432 Ok(rows.filter_map(|r| r.ok()).collect())
433 }
434
435 pub fn sessions(&self, since_ms: Option<i64>, limit: usize) -> Result<Vec<Value>> {
436 let conn = self.conn.lock().unwrap();
437 let mut sql = String::from(
438 "SELECT session_id, MAX(run_id), MIN(ts_request_ms), MAX(ts_request_ms), COUNT(*),
439 GROUP_CONCAT(DISTINCT routed_model), MAX(harness),
440 COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
441 COALESCE(SUM(cost_usd),0.0),
442 COALESCE(SUM(CASE WHEN error IS NOT NULL OR status >= 400 THEN 1 ELSE 0 END),0),
443 (SELECT t2.status FROM traces t2 WHERE t2.session_id = traces.session_id
444 ORDER BY t2.ts_request_ms DESC LIMIT 1),
445 GROUP_CONCAT(tags_json, char(31)),
446 GROUP_CONCAT(DISTINCT reasoning_effort)
447 FROM traces WHERE session_id IS NOT NULL",
448 );
449 let mut args: Vec<String> = vec![];
450 if let Some(since) = since_ms {
451 sql.push_str(" AND ts_request_ms >= ?");
452 args.push(since.to_string());
453 }
454 sql.push_str(" GROUP BY session_id ORDER BY MAX(ts_request_ms) DESC LIMIT ?");
455 let limit = if limit == 0 {
456 DEFAULT_SEARCH_LIMIT
457 } else {
458 limit.min(1000)
459 };
460 args.push(limit.to_string());
461 let mut stmt = conn.prepare(&sql)?;
462 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), |r| {
463 let models: Vec<String> = r
464 .get::<_, Option<String>>(5)?
465 .map(|s| s.split(',').map(str::to_string).collect())
466 .unwrap_or_default();
467 let mut tags = serde_json::Map::new();
468 if let Some(raw) = r.get::<_, Option<String>>(12)? {
469 for piece in raw.split('\u{1f}') {
470 if let Ok(Value::Object(o)) = serde_json::from_str::<Value>(piece) {
471 tags.extend(o);
472 }
473 }
474 }
475 let efforts: Vec<String> = r
476 .get::<_, Option<String>>(13)?
477 .map(|s| s.split(',').map(str::to_string).collect())
478 .unwrap_or_default();
479 Ok(json!({
480 "session_id": r.get::<_, String>(0)?,
481 "run_id": r.get::<_, Option<String>>(1)?,
482 "first_ts_ms": r.get::<_, Option<i64>>(2)?,
483 "last_ts_ms": r.get::<_, Option<i64>>(3)?,
484 "trace_count": r.get::<_, i64>(4)?,
485 "models": models,
486 "harness": r.get::<_, Option<String>>(6)?,
487 "total_input_tokens": r.get::<_, i64>(7)?,
488 "total_output_tokens": r.get::<_, i64>(8)?,
489 "total_cost_usd": r.get::<_, f64>(9)?,
490 "errors": r.get::<_, i64>(10)?,
491 "last_status": r.get::<_, Option<i64>>(11)?,
492 "tags": tags,
493 "efforts": efforts,
494 }))
495 })?;
496 Ok(rows.filter_map(|r| r.ok()).collect())
497 }
498
499 pub fn session_traces(&self, session_id: &str, since_ms: Option<i64>) -> Result<Vec<Value>> {
500 let conn = self.conn.lock().unwrap();
501 let mut sql = format!("SELECT {TRACE_COLS} FROM traces WHERE session_id = ?");
502 let mut args = vec![session_id.to_string()];
503 if let Some(since) = since_ms {
504 sql.push_str(" AND ts_request_ms >= ?");
505 args.push(since.to_string());
506 }
507 sql.push_str(" ORDER BY ts_request_ms ASC");
508 let mut stmt = conn.prepare(&sql)?;
509 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), trace_row_json)?;
510 Ok(rows.filter_map(|r| r.ok()).collect())
511 }
512
513 pub fn get_trace(&self, id: &str) -> Result<Option<Value>> {
514 let conn = self.conn.lock().unwrap();
515 let row = conn
516 .query_row(
517 &format!("SELECT {TRACE_COLS} FROM traces WHERE id = ?1"),
518 params![id],
519 trace_row_json,
520 )
521 .optional()?;
522 Ok(row)
523 }
524
525 pub fn delete_trace(&self, id: &str) -> Result<Vec<String>> {
526 let conn = self.conn.lock().unwrap();
527 let paths: Option<(Option<String>, Option<String>, Option<String>)> = conn
528 .query_row(
529 "SELECT req_body_path, upstream_req_body_path, resp_body_path
530 FROM traces WHERE id = ?1",
531 params![id],
532 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
533 )
534 .optional()?;
535 let Some((req, upstream, resp)) = paths else {
536 anyhow::bail!("trace not found: {id}");
537 };
538 conn.execute("DELETE FROM traces WHERE id = ?1", params![id])?;
539 Ok([req, upstream, resp].into_iter().flatten().collect())
540 }
541
542 pub fn run_summary(&self, run_id: &str) -> Result<Value> {
543 let conn = self.conn.lock().unwrap();
544 #[allow(clippy::type_complexity)]
545 let (
546 trace_count,
547 first_ts_ms,
548 last_ts_ms,
549 last_response_ms,
550 pending,
551 total_input,
552 total_output,
553 total_cost,
554 errors,
555 ) = conn.query_row(
556 "SELECT COUNT(*), MIN(ts_request_ms), MAX(ts_request_ms), MAX(ts_response_ms),
557 COALESCE(SUM(CASE WHEN ts_response_ms IS NULL THEN 1 ELSE 0 END),0),
558 COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
559 COALESCE(SUM(cost_usd),0.0),
560 COALESCE(SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END),0)
561 FROM traces WHERE run_id = ?1",
562 params![run_id],
563 |r| {
564 Ok((
565 r.get::<_, i64>(0)?,
566 r.get::<_, Option<i64>>(1)?,
567 r.get::<_, Option<i64>>(2)?,
568 r.get::<_, Option<i64>>(3)?,
569 r.get::<_, i64>(4)?,
570 r.get::<_, i64>(5)?,
571 r.get::<_, i64>(6)?,
572 r.get::<_, f64>(7)?,
573 r.get::<_, i64>(8)?,
574 ))
575 },
576 )?;
577 let mut status_counts = serde_json::Map::new();
578 let mut stmt = conn.prepare(
579 "SELECT status, COUNT(*) FROM traces WHERE run_id = ?1 GROUP BY status",
580 )?;
581 let pairs = stmt.query_map(params![run_id], |r| {
582 Ok((r.get::<_, Option<i64>>(0)?, r.get::<_, i64>(1)?))
583 })?;
584 for pair in pairs.flatten() {
585 let key = pair.0.map(|s| s.to_string()).unwrap_or_else(|| "none".into());
586 status_counts.insert(key, json!(pair.1));
587 }
588 let distinct = |col: &str| -> Result<Vec<String>> {
589 let mut stmt = conn.prepare(&format!(
590 "SELECT DISTINCT {col} FROM traces WHERE run_id = ?1 AND {col} IS NOT NULL ORDER BY {col}"
591 ))?;
592 let vals = stmt
593 .query_map(params![run_id], |r| r.get::<_, String>(0))?
594 .filter_map(|r| r.ok())
595 .collect();
596 Ok(vals)
597 };
598 let models = distinct("routed_model")?;
599 let providers = distinct("upstream_provider")?;
600 let mut tags = serde_json::Map::new();
601 let mut stmt = conn.prepare(
602 "SELECT tags_json FROM traces WHERE run_id = ?1 AND tags_json IS NOT NULL ORDER BY ts_request_ms",
603 )?;
604 let tag_rows = stmt.query_map(params![run_id], |r| r.get::<_, String>(0))?;
605 for raw in tag_rows.flatten() {
606 if let Ok(Value::Object(o)) = serde_json::from_str::<Value>(&raw) {
607 tags.extend(o);
608 }
609 }
610 Ok(json!({
611 "run_id": run_id,
612 "trace_count": trace_count,
613 "first_ts_ms": first_ts_ms,
614 "last_ts_ms": last_ts_ms,
615 "last_request_ms": last_ts_ms,
616 "last_response_ms": last_response_ms,
617 "last_activity_ms": last_response_ms.max(last_ts_ms),
618 "pending": pending,
619 "status_counts": status_counts,
620 "models": models,
621 "providers": providers,
622 "total_input_tokens": total_input,
623 "total_output_tokens": total_output,
624 "total_cost_usd": total_cost,
625 "tags": tags,
626 "errors": errors,
627 }))
628 }
629
630 pub fn run_artifacts(&self, run_id: &str) -> Result<Vec<Value>> {
631 let conn = self.conn.lock().unwrap();
632 let mut stmt = conn.prepare(
633 "SELECT id, req_body_path, upstream_req_body_path, resp_body_path
634 FROM traces WHERE run_id = ?1 ORDER BY ts_request_ms",
635 )?;
636 let rows: Vec<(String, Option<String>, Option<String>, Option<String>)> = stmt
637 .query_map(params![run_id], |r| {
638 Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?))
639 })?
640 .filter_map(|r| r.ok())
641 .collect();
642 let mut out = Vec::new();
643 for (trace_id, req, upstream_req, resp) in rows {
644 for (kind, path) in [
645 ("request", req),
646 ("upstream-request", upstream_req),
647 ("response", resp),
648 ] {
649 let Some(path) = path else { continue };
650 let size_bytes = std::fs::metadata(&path).ok().map(|m| m.len());
651 out.push(json!({
652 "trace_id": trace_id,
653 "kind": kind,
654 "path": path,
655 "exists": size_bytes.is_some(),
656 "size_bytes": size_bytes,
657 }));
658 }
659 }
660 Ok(out)
661 }
662
663 pub fn insert_heartbeat(
664 &self,
665 ts_ms: i64,
666 provider: &str,
667 account_id: Option<&str>,
668 ok: bool,
669 status: Option<i64>,
670 latency_ms: i64,
671 message: &str,
672 ) -> Result<()> {
673 let conn = self.conn.lock().unwrap();
674 conn.execute(
675 "INSERT INTO heartbeats (ts_ms, provider, account_id, ok, status, latency_ms, message)
676 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
677 params![ts_ms, provider, account_id, ok as i64, status, latency_ms, message],
678 )?;
679 Ok(())
680 }
681
682 pub fn last_heartbeats(&self) -> Result<Vec<Value>> {
683 let conn = self.conn.lock().unwrap();
684 let mut stmt = conn.prepare(
685 "SELECT h.ts_ms, h.provider, h.account_id, h.ok, h.status, h.latency_ms, h.message
686 FROM heartbeats h
687 JOIN (SELECT provider, MAX(ts_ms) AS ts FROM heartbeats GROUP BY provider) latest
688 ON h.provider = latest.provider AND h.ts_ms = latest.ts",
689 )?;
690 let rows = stmt.query_map([], |r| {
691 Ok(json!({
692 "ts_ms": r.get::<_, i64>(0)?,
693 "provider": r.get::<_, String>(1)?,
694 "account_id": r.get::<_, Option<String>>(2)?,
695 "ok": r.get::<_, i64>(3)? == 1,
696 "status": r.get::<_, Option<i64>>(4)?,
697 "latency_ms": r.get::<_, i64>(5)?,
698 "message": r.get::<_, String>(6)?,
699 }))
700 })?;
701 Ok(rows.filter_map(|r| r.ok()).collect())
702 }
703
704 pub fn latest_provider_headers(&self) -> Result<Vec<(String, i64, String)>> {
705 let conn = self.conn.lock().unwrap();
706 let mut stmt = conn.prepare(
707 "SELECT t.upstream_provider, t.ts_request_ms, t.resp_headers_json
708 FROM traces t
709 JOIN (SELECT upstream_provider p, MAX(ts_request_ms) ts FROM traces
710 WHERE status >= 200 AND status < 300
711 AND resp_headers_json IS NOT NULL AND upstream_provider IS NOT NULL
712 GROUP BY upstream_provider) latest
713 ON t.upstream_provider = latest.p AND t.ts_request_ms = latest.ts
714 WHERE t.resp_headers_json IS NOT NULL",
715 )?;
716 let rows: Vec<(String, i64, String)> = stmt
717 .query_map([], |r| {
718 Ok((
719 r.get::<_, String>(0)?,
720 r.get::<_, i64>(1)?,
721 r.get::<_, String>(2)?,
722 ))
723 })?
724 .filter_map(|r| r.ok())
725 .collect();
726 let mut seen = std::collections::HashMap::new();
727 for row in rows {
728 seen.entry(row.0.clone()).or_insert(row);
729 }
730 Ok(seen.into_values().collect())
731 }
732
733 pub fn analytics(&self, since_ms: i64) -> Result<Value> {
734 let conn = self.conn.lock().unwrap();
735 let mut stmt = conn.prepare(
736 "SELECT routed_model, upstream_provider, billing_bucket, COUNT(*),
737 COALESCE(SUM(input_tokens),0), COALESCE(SUM(cached_input_tokens),0),
738 COALESCE(SUM(output_tokens),0), COALESCE(SUM(cost_usd),0.0),
739 SUM(CASE WHEN status >= 200 AND status < 300 THEN 0 ELSE 1 END),
740 AVG(CASE WHEN ts_response_ms IS NOT NULL THEN ts_response_ms - ts_request_ms END)
741 FROM traces WHERE ts_request_ms >= ?1
742 GROUP BY routed_model, upstream_provider, billing_bucket
743 ORDER BY SUM(cost_usd) DESC",
744 )?;
745 let rows: Vec<Value> = stmt
746 .query_map(params![since_ms], |r| {
747 Ok(json!({
748 "routed_model": r.get::<_, Option<String>>(0)?,
749 "upstream_provider": r.get::<_, Option<String>>(1)?,
750 "billing_bucket": r.get::<_, Option<String>>(2)?,
751 "requests": r.get::<_, i64>(3)?,
752 "input_tokens": r.get::<_, i64>(4)?,
753 "cached_input_tokens": r.get::<_, i64>(5)?,
754 "output_tokens": r.get::<_, i64>(6)?,
755 "cost_usd": r.get::<_, f64>(7)?,
756 "errors": r.get::<_, Option<i64>>(8)?,
757 "avg_latency_ms": r.get::<_, Option<f64>>(9)?,
758 }))
759 })?
760 .filter_map(|r| r.ok())
761 .collect();
762 let (mut requests, mut cost, mut errors) = (0i64, 0f64, 0i64);
763 let mut buckets: std::collections::HashMap<String, f64> = Default::default();
764 for row in &rows {
765 requests += row["requests"].as_i64().unwrap_or(0);
766 cost += row["cost_usd"].as_f64().unwrap_or(0.0);
767 errors += row["errors"].as_i64().unwrap_or(0);
768 let bucket = row["billing_bucket"].as_str().unwrap_or("unknown").to_string();
769 *buckets.entry(bucket).or_default() += row["cost_usd"].as_f64().unwrap_or(0.0);
770 }
771 Ok(json!({
772 "since_ms": since_ms,
773 "totals": {"requests": requests, "cost_usd": cost, "errors": errors, "cost_by_bucket": buckets},
774 "by_model": rows,
775 }))
776 }
777
778 #[allow(clippy::too_many_arguments)]
779 pub fn insert_run_key(
780 &self,
781 id: &str,
782 key_hash: &str,
783 kind: &str,
784 run_id: Option<&str>,
785 tags_json: Option<&str>,
786 label: Option<&str>,
787 created_ms: i64,
788 expires_ms: Option<i64>,
789 ) -> Result<()> {
790 let conn = self.conn.lock().unwrap();
791 conn.execute(
792 "INSERT INTO run_keys (id, key_hash, kind, run_id, tags_json, label, created_ms, expires_ms)
793 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
794 params![id, key_hash, kind, run_id, tags_json, label, created_ms, expires_ms],
795 )?;
796 Ok(())
797 }
798
799 pub fn lookup_run_key(&self, key_hash: &str, now_ms: i64) -> Result<Option<Value>> {
800 let conn = self.conn.lock().unwrap();
801 let row = conn
802 .query_row(
803 &format!(
804 "SELECT {RUN_KEY_COLS} FROM run_keys
805 WHERE key_hash = ?1 AND revoked = 0
806 AND (expires_ms IS NULL OR expires_ms > ?2)"
807 ),
808 params![key_hash, now_ms],
809 run_key_row_json,
810 )
811 .optional()?;
812 Ok(row)
813 }
814
815 pub fn touch_run_key(&self, key_hash: &str, now_ms: i64) -> Result<()> {
816 let conn = self.conn.lock().unwrap();
817 conn.execute(
818 "UPDATE run_keys SET use_count = use_count + 1, last_used_ms = ?2 WHERE key_hash = ?1",
819 params![key_hash, now_ms],
820 )?;
821 Ok(())
822 }
823
824 pub fn list_run_keys(&self, include_inactive: bool) -> Result<Vec<Value>> {
825 let conn = self.conn.lock().unwrap();
826 let mut sql = format!("SELECT {RUN_KEY_COLS} FROM run_keys");
827 let mut args: Vec<String> = vec![];
828 if !include_inactive {
829 sql.push_str(" WHERE revoked = 0 AND (expires_ms IS NULL OR expires_ms > ?1)");
830 args.push(Utc::now().timestamp_millis().to_string());
831 }
832 sql.push_str(" ORDER BY created_ms DESC");
833 let mut stmt = conn.prepare(&sql)?;
834 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), run_key_row_json)?;
835 Ok(rows.filter_map(|r| r.ok()).collect())
836 }
837
838 pub fn revoke_run_key(&self, id_or_prefix: &str) -> Result<bool> {
839 let conn = self.conn.lock().unwrap();
840 let changed = conn.execute(
841 "UPDATE run_keys SET revoked = 1 WHERE id = ?1 OR id LIKE ?1 || '%'",
842 params![id_or_prefix],
843 )?;
844 Ok(changed > 0)
845 }
846
847 pub fn prune(&self, older_than_ms: i64, bodies_only: bool, dry_run: bool) -> Result<PruneReport> {
848 let mut report = PruneReport::default();
849 let rows: Vec<(String, Option<String>, Option<String>, Option<String>)> = {
850 let conn = self.conn.lock().unwrap();
851 let mut stmt = conn.prepare(
852 "SELECT id, req_body_path, upstream_req_body_path, resp_body_path
853 FROM traces WHERE ts_request_ms < ?1
854 AND (req_body_path IS NOT NULL OR upstream_req_body_path IS NOT NULL
855 OR resp_body_path IS NOT NULL OR req_headers_json IS NOT NULL
856 OR resp_headers_json IS NOT NULL)",
857 )?;
858 let rows = stmt
859 .query_map(params![older_than_ms], |r| {
860 Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?))
861 })?
862 .filter_map(|r| r.ok())
863 .collect();
864 rows
865 };
866 for (id, req, upstream, resp) in rows {
867 for path in [req, upstream, resp].into_iter().flatten() {
868 let Ok(meta) = std::fs::metadata(&path) else {
869 continue;
870 };
871 report.bytes_freed += meta.len();
872 report.bodies_deleted += 1;
873 if !dry_run {
874 std::fs::remove_file(&path)
875 .with_context(|| format!("deleting body file {path}"))?;
876 }
877 }
878 if !dry_run {
879 let conn = self.conn.lock().unwrap();
880 conn.execute(
881 "UPDATE traces SET req_body_path = NULL, upstream_req_body_path = NULL,
882 resp_body_path = NULL, req_headers_json = NULL, resp_headers_json = NULL
883 WHERE id = ?1",
884 params![id],
885 )?;
886 }
887 report.rows_affected += 1;
888 }
889 if !bodies_only {
890 let conn = self.conn.lock().unwrap();
891 if dry_run {
892 report.rows_deleted = conn.query_row(
893 "SELECT COUNT(*) FROM traces WHERE ts_request_ms < ?1",
894 params![older_than_ms],
895 |r| r.get::<_, i64>(0),
896 )? as u64;
897 } else {
898 report.rows_deleted = conn.execute(
899 "DELETE FROM traces WHERE ts_request_ms < ?1",
900 params![older_than_ms],
901 )? as u64;
902 if report.rows_deleted > 0 {
903 if let Err(e) = conn.execute_batch("VACUUM") {
904 tracing::warn!("prune: VACUUM skipped: {e}");
905 }
906 }
907 }
908 }
909 if !dry_run {
910 let bodies = self.data_dir.join("bodies");
911 if let Ok(entries) = std::fs::read_dir(&bodies) {
912 for entry in entries.flatten() {
913 let path = entry.path();
914 let name = entry.file_name().to_string_lossy().to_string();
915 if !path.is_dir() || !date_dir_name(&name) {
916 continue;
917 }
918 let empty = std::fs::read_dir(&path)
919 .map(|mut d| d.next().is_none())
920 .unwrap_or(false);
921 if empty && std::fs::remove_dir(&path).is_ok() {
922 report.dirs_removed += 1;
923 }
924 }
925 }
926 }
927 Ok(report)
928 }
929
930 pub fn disk_usage(&self) -> Result<Value> {
931 let mut sqlite_bytes = 0u64;
932 for suffix in ["", "-wal", "-shm"] {
933 let path = self.data_dir.join(format!("alexandria.sqlite3{suffix}"));
934 if let Ok(m) = std::fs::metadata(&path) {
935 sqlite_bytes += m.len();
936 }
937 }
938 let mut bodies_bytes = 0u64;
939 let mut days: Vec<Value> = vec![];
940 if let Ok(entries) = std::fs::read_dir(self.data_dir.join("bodies")) {
941 for entry in entries.flatten() {
942 let name = entry.file_name().to_string_lossy().to_string();
943 if !entry.path().is_dir() || !date_dir_name(&name) {
944 continue;
945 }
946 let (mut files, mut bytes) = (0u64, 0u64);
947 if let Ok(inner) = std::fs::read_dir(entry.path()) {
948 for f in inner.flatten() {
949 if let Ok(m) = f.metadata() {
950 if m.is_file() {
951 files += 1;
952 bytes += m.len();
953 }
954 }
955 }
956 }
957 bodies_bytes += bytes;
958 days.push(json!({"date": name, "files": files, "bytes": bytes}));
959 }
960 }
961 days.sort_by(|a, b| b["date"].as_str().cmp(&a["date"].as_str()));
962 let conn = self.conn.lock().unwrap();
963 let (trace_rows, oldest_ts_ms, newest_ts_ms) = conn.query_row(
964 "SELECT COUNT(*), MIN(ts_request_ms), MAX(ts_request_ms) FROM traces",
965 [],
966 |r| {
967 Ok((
968 r.get::<_, i64>(0)?,
969 r.get::<_, Option<i64>>(1)?,
970 r.get::<_, Option<i64>>(2)?,
971 ))
972 },
973 )?;
974 Ok(json!({
975 "sqlite_bytes": sqlite_bytes,
976 "bodies_bytes": bodies_bytes,
977 "days": days,
978 "trace_rows": trace_rows,
979 "oldest_ts_ms": oldest_ts_ms,
980 "newest_ts_ms": newest_ts_ms,
981 }))
982 }
983
984 pub fn write_body(&self, trace_id: &str, kind: &str, bytes: &[u8]) -> Result<String> {
985 let date = Utc::now().format("%Y-%m-%d").to_string();
986 self.write_body_dated(&date, trace_id, kind, bytes)
987 }
988
989 fn write_body_dated(&self, date: &str, trace_id: &str, kind: &str, bytes: &[u8]) -> Result<String> {
990 let dir = self.data_dir.join("bodies").join(date);
991 std::fs::create_dir_all(&dir)?;
992 let path = dir.join(format!("{trace_id}.{kind}.gz"));
993 let file = std::fs::File::create(&path)?;
994 let mut enc = GzEncoder::new(file, Compression::default());
995 enc.write_all(bytes)?;
996 enc.finish()?;
997 Ok(path.to_string_lossy().to_string())
998 }
999}
1000
1001#[cfg(test)]
1002mod tests {
1003 use super::*;
1004
1005 fn tmpdir(name: &str) -> PathBuf {
1006 let dir = std::env::temp_dir().join(format!(
1007 "alexandria-store-test-{name}-{}",
1008 std::process::id()
1009 ));
1010 let _ = std::fs::remove_dir_all(&dir);
1011 std::fs::create_dir_all(&dir).unwrap();
1012 dir
1013 }
1014
1015 fn trace(id: &str, ts: i64, run: Option<&str>) -> TraceRecord {
1016 TraceRecord {
1017 id: id.into(),
1018 ts_request_ms: ts,
1019 ts_response_ms: Some(ts + 250),
1020 status: Some(200),
1021 routed_model: Some("claude-haiku-4-5".into()),
1022 upstream_provider: Some("anthropic".into()),
1023 run_id: run.map(String::from),
1024 usage: alex_core::Usage {
1025 input_tokens: Some(10),
1026 output_tokens: Some(5),
1027 ..Default::default()
1028 },
1029 cost_usd: Some(0.001),
1030 ..Default::default()
1031 }
1032 }
1033
1034 #[test]
1035 fn run_summary_aggregates() {
1036 let store = Store::open(tmpdir("summary")).unwrap();
1037 let mut a = trace("a", 1000, Some("run-1"));
1038 a.tags = Some(r#"{"suite":"swebench"}"#.into());
1039 let mut b = trace("b", 2000, Some("run-1"));
1040 b.tags = Some(r#"{"case":"astropy-1"}"#.into());
1041 b.status = Some(500);
1042 b.error = Some("boom".into());
1043 b.routed_model = Some("gpt-5.5".into());
1044 b.upstream_provider = Some("openai".into());
1045 let c = trace("c", 3000, Some("run-2"));
1046 for t in [&a, &b, &c] {
1047 store.insert_trace(t).unwrap();
1048 }
1049 let s = store.run_summary("run-1").unwrap();
1050 assert_eq!(s["trace_count"], 2);
1051 assert_eq!(s["first_ts_ms"], 1000);
1052 assert_eq!(s["last_ts_ms"], 2000);
1053 assert_eq!(s["status_counts"]["200"], 1);
1054 assert_eq!(s["status_counts"]["500"], 1);
1055 assert_eq!(
1056 s["models"],
1057 json!(["claude-haiku-4-5", "gpt-5.5"])
1058 );
1059 assert_eq!(s["providers"], json!(["anthropic", "openai"]));
1060 assert_eq!(s["total_input_tokens"], 20);
1061 assert_eq!(s["total_output_tokens"], 10);
1062 assert_eq!(s["tags"]["suite"], "swebench");
1063 assert_eq!(s["tags"]["case"], "astropy-1");
1064 assert_eq!(s["errors"], 1);
1065 let missing = store.run_summary("nope").unwrap();
1066 assert_eq!(missing["trace_count"], 0);
1067 }
1068
1069 #[test]
1070 fn search_traces_filters() {
1071 let store = Store::open(tmpdir("search")).unwrap();
1072 let mut a = trace("a", 1000, Some("run-1"));
1073 a.key_fingerprint = Some("deadbeefdeadbeef".into());
1074 a.reasoning_effort = Some("high".into());
1075 a.thinking_budget = Some(16_384);
1076 let mut b = trace("b", 2000, Some("run-1"));
1077 b.status = Some(429);
1078 b.error = Some("rate limited".into());
1079 let c = trace("c", 3000, None);
1080 for t in [&a, &b, &c] {
1081 store.insert_trace(t).unwrap();
1082 }
1083 let all = store.search_traces(&TraceFilter::default()).unwrap();
1084 assert_eq!(all.len(), 3);
1085 assert_eq!(all[0]["id"], "c");
1086 assert_eq!(all[0]["latency_ms"], 250);
1087 assert_eq!(all[2]["reasoning_effort"], "high");
1088 assert_eq!(all[2]["thinking_budget"], 16_384);
1089 let window = store
1090 .search_traces(&TraceFilter {
1091 since_ms: Some(1500),
1092 until_ms: Some(2500),
1093 ..Default::default()
1094 })
1095 .unwrap();
1096 assert_eq!(window.len(), 1);
1097 assert_eq!(window[0]["id"], "b");
1098 let by_run = store
1099 .search_traces(&TraceFilter {
1100 run_id: Some("run-1".into()),
1101 ..Default::default()
1102 })
1103 .unwrap();
1104 assert_eq!(by_run.len(), 2);
1105 let by_status = store
1106 .search_traces(&TraceFilter {
1107 status: Some(429),
1108 ..Default::default()
1109 })
1110 .unwrap();
1111 assert_eq!(by_status.len(), 1);
1112 let errors = store
1113 .search_traces(&TraceFilter {
1114 errors_only: true,
1115 ..Default::default()
1116 })
1117 .unwrap();
1118 assert_eq!(errors.len(), 1);
1119 assert_eq!(errors[0]["id"], "b");
1120 let by_key = store
1121 .search_traces(&TraceFilter {
1122 key_fingerprint: Some("deadbeefdeadbeef".into()),
1123 ..Default::default()
1124 })
1125 .unwrap();
1126 assert_eq!(by_key.len(), 1);
1127 assert_eq!(by_key[0]["id"], "a");
1128 let by_effort = store
1129 .search_traces(&TraceFilter {
1130 reasoning_effort: Some("high".into()),
1131 ..Default::default()
1132 })
1133 .unwrap();
1134 assert_eq!(by_effort.len(), 1);
1135 assert_eq!(by_effort[0]["id"], "a");
1136 let limited = store
1137 .search_traces(&TraceFilter {
1138 limit: 1,
1139 ..Default::default()
1140 })
1141 .unwrap();
1142 assert_eq!(limited.len(), 1);
1143 }
1144
1145 #[test]
1146 fn sessions_aggregate_and_order() {
1147 let store = Store::open(tmpdir("sessions")).unwrap();
1148 let mut a = trace("a", 1000, Some("run-1"));
1149 a.session_id = Some("ses_1".into());
1150 a.tags = Some(r#"{"suite":"swebench"}"#.into());
1151 a.harness = Some("codex".into());
1152 a.reasoning_effort = Some("high".into());
1153 let mut b = trace("b", 2000, None);
1154 b.session_id = Some("ses_1".into());
1155 b.status = Some(500);
1156 b.error = Some("boom".into());
1157 b.routed_model = Some("gpt-5.5".into());
1158 b.tags = Some(r#"{"case":"x1"}"#.into());
1159 b.reasoning_effort = Some("minimal".into());
1160 let mut c = trace("c", 5000, None);
1161 c.session_id = Some("ses_2".into());
1162 let d = trace("d", 9000, None);
1163 for t in [&a, &b, &c, &d] {
1164 store.insert_trace(t).unwrap();
1165 }
1166 let sessions = store.sessions(None, 0).unwrap();
1167 assert_eq!(sessions.len(), 2);
1168 assert_eq!(sessions[0]["session_id"], "ses_2");
1169 let s1 = &sessions[1];
1170 assert_eq!(s1["session_id"], "ses_1");
1171 assert_eq!(s1["run_id"], "run-1");
1172 assert_eq!(s1["first_ts_ms"], 1000);
1173 assert_eq!(s1["last_ts_ms"], 2000);
1174 assert_eq!(s1["trace_count"], 2);
1175 assert_eq!(s1["harness"], "codex");
1176 assert_eq!(s1["total_input_tokens"], 20);
1177 assert_eq!(s1["total_output_tokens"], 10);
1178 assert_eq!(s1["errors"], 1);
1179 assert_eq!(s1["last_status"], 500);
1180 assert_eq!(s1["tags"]["suite"], "swebench");
1181 assert_eq!(s1["tags"]["case"], "x1");
1182 let efforts: Vec<String> = s1["efforts"]
1183 .as_array()
1184 .unwrap()
1185 .iter()
1186 .map(|m| m.as_str().unwrap().to_string())
1187 .collect();
1188 assert!(efforts.contains(&"high".to_string()));
1189 assert!(efforts.contains(&"minimal".to_string()));
1190 let models: Vec<String> = s1["models"]
1191 .as_array()
1192 .unwrap()
1193 .iter()
1194 .map(|m| m.as_str().unwrap().to_string())
1195 .collect();
1196 assert!(models.contains(&"claude-haiku-4-5".to_string()));
1197 assert!(models.contains(&"gpt-5.5".to_string()));
1198 let recent = store.sessions(Some(3000), 0).unwrap();
1199 assert_eq!(recent.len(), 1);
1200 assert_eq!(recent[0]["session_id"], "ses_2");
1201 let limited = store.sessions(None, 1).unwrap();
1202 assert_eq!(limited.len(), 1);
1203 }
1204
1205 #[test]
1206 fn session_traces_ascending() {
1207 let store = Store::open(tmpdir("session-traces")).unwrap();
1208 for (id, ts) in [("a", 3000i64), ("b", 1000), ("c", 2000)] {
1209 let mut t = trace(id, ts, None);
1210 t.session_id = Some("ses_1".into());
1211 t.upstream_format = Some("anthropic".into());
1212 t.req_body_path = Some(format!("/bodies/{id}.request.json.gz"));
1213 store.insert_trace(&t).unwrap();
1214 }
1215 let rows = store.session_traces("ses_1", None).unwrap();
1216 assert_eq!(rows.len(), 3);
1217 assert_eq!(rows[0]["id"], "b");
1218 assert_eq!(rows[2]["id"], "a");
1219 assert_eq!(rows[0]["upstream_format"], "anthropic");
1220 assert_eq!(rows[0]["client_format"], Value::Null);
1221 assert_eq!(rows[0]["req_body_path"], "/bodies/b.request.json.gz");
1222 let windowed = store.session_traces("ses_1", Some(1500)).unwrap();
1223 assert_eq!(windowed.len(), 2);
1224 assert_eq!(windowed[0]["id"], "c");
1225 assert!(store.session_traces("nope", None).unwrap().is_empty());
1226 }
1227
1228 #[test]
1229 fn get_and_delete_trace() {
1230 let store = Store::open(tmpdir("delete")).unwrap();
1231 let mut t = trace("a", 1000, None);
1232 t.req_body_path = Some(
1233 store
1234 .write_body("a", "request.json", b"{\"model\":\"x\"}")
1235 .unwrap(),
1236 );
1237 t.resp_body_path = Some("/nonexistent/a.response.body.gz".into());
1238 store.insert_trace(&t).unwrap();
1239 let row = store.get_trace("a").unwrap().unwrap();
1240 assert_eq!(row["id"], "a");
1241 assert_eq!(row["resp_body_path"], "/nonexistent/a.response.body.gz");
1242 assert!(store.get_trace("missing").unwrap().is_none());
1243 let paths = store.delete_trace("a").unwrap();
1244 assert_eq!(paths.len(), 2);
1245 assert!(store.get_trace("a").unwrap().is_none());
1246 assert!(store.delete_trace("a").is_err());
1247 }
1248
1249 fn seed_dated_trace(store: &Store, id: &str, ts: i64, date: &str) {
1250 let mut t = trace(id, ts, None);
1251 t.req_body_path = Some(
1252 store
1253 .write_body_dated(date, id, "request.json", b"{\"model\":\"x\"}")
1254 .unwrap(),
1255 );
1256 t.resp_body_path = Some(
1257 store
1258 .write_body_dated(date, id, "response.body", b"response bytes here")
1259 .unwrap(),
1260 );
1261 t.req_headers_json = Some(r#"{"authorization":"[redacted]"}"#.into());
1262 t.resp_headers_json = Some(r#"{"content-type":"application/json"}"#.into());
1263 store.insert_trace(&t).unwrap();
1264 }
1265
1266 #[test]
1267 fn date_dir_name_shape() {
1268 assert!(date_dir_name("2024-01-31"));
1269 assert!(!date_dir_name("2024-1-31"));
1270 assert!(!date_dir_name("20240131"));
1271 assert!(!date_dir_name("accounts"));
1272 assert!(!date_dir_name("2024-01-31x"));
1273 }
1274
1275 #[test]
1276 fn prune_bodies_dry_run_then_real() {
1277 let store = Store::open(tmpdir("prune-bodies")).unwrap();
1278 seed_dated_trace(&store, "a", 1000, "2024-01-01");
1279 seed_dated_trace(&store, "b", 2000, "2024-01-02");
1280 seed_dated_trace(&store, "c", 10_000, "2024-01-03");
1281 let scratch = store.data_dir.join("bodies").join("scratch");
1282 std::fs::create_dir_all(&scratch).unwrap();
1283 let dry = store.prune(5000, true, true).unwrap();
1284 assert_eq!(dry.bodies_deleted, 4);
1285 assert!(dry.bytes_freed > 0);
1286 assert_eq!(dry.rows_affected, 2);
1287 assert_eq!(dry.rows_deleted, 0);
1288 assert_eq!(dry.dirs_removed, 0);
1289 let before = store.get_trace("a").unwrap().unwrap();
1290 assert!(before["req_body_path"].is_string());
1291 assert!(store.data_dir.join("bodies").join("2024-01-01").exists());
1292 let real = store.prune(5000, true, false).unwrap();
1293 assert_eq!(real.bodies_deleted, dry.bodies_deleted);
1294 assert_eq!(real.bytes_freed, dry.bytes_freed);
1295 assert_eq!(real.rows_affected, dry.rows_affected);
1296 assert_eq!(real.rows_deleted, 0);
1297 assert_eq!(real.dirs_removed, 2);
1298 assert!(!store.data_dir.join("bodies").join("2024-01-01").exists());
1299 assert!(!store.data_dir.join("bodies").join("2024-01-02").exists());
1300 assert!(store.data_dir.join("bodies").join("2024-01-03").exists());
1301 assert!(scratch.exists());
1302 assert_eq!(store.search_traces(&TraceFilter::default()).unwrap().len(), 3);
1303 let pruned = store.get_trace("a").unwrap().unwrap();
1304 for key in [
1305 "req_body_path",
1306 "upstream_req_body_path",
1307 "resp_body_path",
1308 "req_headers_json",
1309 "resp_headers_json",
1310 ] {
1311 assert_eq!(pruned[key], Value::Null, "{key} not nulled");
1312 }
1313 let kept = store.get_trace("c").unwrap().unwrap();
1314 assert!(kept["req_body_path"].is_string());
1315 assert!(kept["req_headers_json"].is_string());
1316 assert_eq!(store.prune(5000, true, false).unwrap(), PruneReport::default());
1317 }
1318
1319 #[test]
1320 fn prune_rows_deletes_and_du_reflects() {
1321 let store = Store::open(tmpdir("prune-rows")).unwrap();
1322 seed_dated_trace(&store, "a", 1000, "2024-01-01");
1323 seed_dated_trace(&store, "b", 2000, "2024-01-02");
1324 seed_dated_trace(&store, "c", 10_000, "2024-01-03");
1325 let before = store.disk_usage().unwrap();
1326 assert_eq!(before["trace_rows"], 3);
1327 assert_eq!(before["oldest_ts_ms"], 1000);
1328 assert_eq!(before["newest_ts_ms"], 10_000);
1329 assert_eq!(before["days"].as_array().unwrap().len(), 3);
1330 assert_eq!(before["days"][0]["date"], "2024-01-03");
1331 assert_eq!(before["days"][0]["files"], 2);
1332 assert!(before["bodies_bytes"].as_u64().unwrap() > 0);
1333 assert!(before["sqlite_bytes"].as_u64().unwrap() > 0);
1334 let dry = store.prune(5000, false, true).unwrap();
1335 assert_eq!(dry.rows_deleted, 2);
1336 assert_eq!(dry.rows_affected, 2);
1337 assert_eq!(store.disk_usage().unwrap()["trace_rows"], 3);
1338 let real = store.prune(5000, false, false).unwrap();
1339 assert_eq!(real.bodies_deleted, 4);
1340 assert_eq!(real.rows_affected, 2);
1341 assert_eq!(real.rows_deleted, 2);
1342 assert_eq!(real.dirs_removed, 2);
1343 assert_eq!(real.bytes_freed, dry.bytes_freed);
1344 let after = store.disk_usage().unwrap();
1345 assert_eq!(after["trace_rows"], 1);
1346 assert_eq!(after["oldest_ts_ms"], 10_000);
1347 assert_eq!(after["newest_ts_ms"], 10_000);
1348 assert_eq!(after["days"].as_array().unwrap().len(), 1);
1349 assert_eq!(after["days"][0]["date"], "2024-01-03");
1350 assert_eq!(
1351 after["bodies_bytes"].as_u64().unwrap(),
1352 before["bodies_bytes"].as_u64().unwrap() - real.bytes_freed
1353 );
1354 assert!(store.get_trace("a").unwrap().is_none());
1355 assert!(store.get_trace("c").unwrap().is_some());
1356 }
1357
1358 #[test]
1359 fn limit_defaults_and_caps() {
1360 assert_eq!(effective_limit(0), 200);
1361 assert_eq!(effective_limit(50), 50);
1362 assert_eq!(effective_limit(9000), 5000);
1363 }
1364
1365 #[test]
1366 fn run_artifacts_reports_files() {
1367 let dir = tmpdir("artifacts");
1368 let store = Store::open(dir).unwrap();
1369 let mut t = trace("a", 1000, Some("run-1"));
1370 t.req_body_path = Some(
1371 store
1372 .write_body("a", "request.json", b"{\"model\":\"x\"}")
1373 .unwrap(),
1374 );
1375 t.resp_body_path = Some("/nonexistent/a.response.body.gz".into());
1376 store.insert_trace(&t).unwrap();
1377 let arts = store.run_artifacts("run-1").unwrap();
1378 assert_eq!(arts.len(), 2);
1379 assert_eq!(arts[0]["kind"], "request");
1380 assert_eq!(arts[0]["exists"], true);
1381 assert!(arts[0]["size_bytes"].as_u64().unwrap() > 0);
1382 assert_eq!(arts[1]["kind"], "response");
1383 assert_eq!(arts[1]["exists"], false);
1384 assert_eq!(arts[1]["size_bytes"], Value::Null);
1385 assert!(store.run_artifacts("nope").unwrap().is_empty());
1386 }
1387
1388 #[test]
1389 fn reopen_keeps_working() {
1390 let dir = tmpdir("reopen");
1391 {
1392 let store = Store::open(dir.clone()).unwrap();
1393 store.insert_trace(&trace("a", 1000, Some("run-1"))).unwrap();
1394 }
1395 let store = Store::open(dir).unwrap();
1396 store.insert_trace(&trace("b", 2000, Some("run-1"))).unwrap();
1397 let s = store.run_summary("run-1").unwrap();
1398 assert_eq!(s["trace_count"], 2);
1399 }
1400
1401 #[test]
1402 fn run_keys_lifecycle() {
1403 let store = Store::open(tmpdir("run-keys")).unwrap();
1404 store
1405 .insert_run_key(
1406 "rk-aaaa1111",
1407 "aaaa1111bbbb2222cccc",
1408 "run",
1409 Some("run-1"),
1410 Some(r#"{"task":"demo"}"#),
1411 Some("demo job"),
1412 1000,
1413 Some(10_000),
1414 )
1415 .unwrap();
1416 store
1417 .insert_run_key("rk-dddd4444", "dddd4444eeee5555ffff", "run", None, None, None, 2000, None)
1418 .unwrap();
1419 let k = store
1420 .lookup_run_key("aaaa1111bbbb2222cccc", 5000)
1421 .unwrap()
1422 .unwrap();
1423 assert_eq!(k["id"], "rk-aaaa1111");
1424 assert_eq!(k["run_id"], "run-1");
1425 assert_eq!(k["tags"]["task"], "demo");
1426 assert_eq!(k["label"], "demo job");
1427 assert_eq!(k["key_fingerprint"], "aaaa1111bbbb2222");
1428 assert_eq!(k["revoked"], false);
1429 assert!(store
1430 .lookup_run_key("aaaa1111bbbb2222cccc", 10_000)
1431 .unwrap()
1432 .is_none());
1433 assert!(store
1434 .lookup_run_key("dddd4444eeee5555ffff", 10_000)
1435 .unwrap()
1436 .is_some());
1437 assert!(store.lookup_run_key("unknown-hash", 5000).unwrap().is_none());
1438 store.touch_run_key("aaaa1111bbbb2222cccc", 3000).unwrap();
1439 store.touch_run_key("aaaa1111bbbb2222cccc", 4000).unwrap();
1440 let all = store.list_run_keys(true).unwrap();
1441 assert_eq!(all.len(), 2);
1442 let touched = all.iter().find(|r| r["id"] == "rk-aaaa1111").unwrap();
1443 assert_eq!(touched["use_count"], 2);
1444 assert_eq!(touched["last_used_ms"], 4000);
1445 assert_eq!(touched["tags"], json!({"task": "demo"}));
1446 assert_eq!(touched["kind"], "run");
1447 assert!(store.revoke_run_key("rk-aaaa").unwrap());
1448 assert!(!store.revoke_run_key("rk-zzzz").unwrap());
1449 assert!(store
1450 .lookup_run_key("aaaa1111bbbb2222cccc", 5000)
1451 .unwrap()
1452 .is_none());
1453 let active = store.list_run_keys(false).unwrap();
1454 assert_eq!(active.len(), 1);
1455 assert_eq!(active[0]["id"], "rk-dddd4444");
1456 assert!(store
1457 .insert_run_key("rk-ffff6666", "aaaa1111bbbb2222cccc", "run", None, None, None, 1, None)
1458 .is_err());
1459 }
1460
1461 #[test]
1462 fn harness_run_keys_never_expire_and_revoke() {
1463 let store = Store::open(tmpdir("run-keys-harness")).unwrap();
1464 store
1465 .insert_run_key(
1466 "rk-harness1",
1467 "hhhh1111bbbb2222cccc",
1468 "harness",
1469 None,
1470 Some(r#"{"harness":"pi"}"#),
1471 Some("pi"),
1472 1000,
1473 None,
1474 )
1475 .unwrap();
1476 let k = store
1477 .lookup_run_key("hhhh1111bbbb2222cccc", i64::MAX)
1478 .unwrap()
1479 .unwrap();
1480 assert_eq!(k["kind"], "harness");
1481 assert_eq!(k["label"], "pi");
1482 assert_eq!(k["expires_ms"], Value::Null);
1483 assert_eq!(k["tags"], json!({"harness": "pi"}));
1484 let active = store.list_run_keys(false).unwrap();
1485 assert_eq!(active.len(), 1);
1486 assert_eq!(active[0]["id"], "rk-harness1");
1487 assert!(store.revoke_run_key("rk-harness1").unwrap());
1488 assert!(store
1489 .lookup_run_key("hhhh1111bbbb2222cccc", i64::MAX)
1490 .unwrap()
1491 .is_none());
1492 assert!(store.list_run_keys(false).unwrap().is_empty());
1493 }
1494
1495 #[test]
1496 fn run_keys_table_added_to_existing_db() {
1497 let dir = tmpdir("run-keys-migrate");
1498 let db_path = dir.join("alexandria.sqlite3");
1499 {
1500 let conn = Connection::open(&db_path).unwrap();
1501 conn.execute_batch(
1502 "CREATE TABLE traces (id TEXT PRIMARY KEY, ts_request_ms INTEGER NOT NULL,
1503 session_id TEXT, routed_model TEXT);",
1504 )
1505 .unwrap();
1506 }
1507 let store = Store::open(dir).unwrap();
1508 store
1509 .insert_run_key("rk-11112222", "hash-x", "run", None, None, None, 1000, None)
1510 .unwrap();
1511 assert_eq!(store.list_run_keys(true).unwrap().len(), 1);
1512 }
1513
1514 #[test]
1515 fn migrates_old_schema() {
1516 let dir = tmpdir("migrate");
1517 let db_path = dir.join("alexandria.sqlite3");
1518 {
1519 let conn = Connection::open(&db_path).unwrap();
1520 conn.execute_batch(
1521 "CREATE TABLE traces (
1522 id TEXT PRIMARY KEY,
1523 ts_request_ms INTEGER NOT NULL,
1524 ts_response_ms INTEGER,
1525 session_id TEXT, harness TEXT, client_format TEXT,
1526 upstream_provider TEXT, upstream_format TEXT,
1527 requested_model TEXT, routed_model TEXT,
1528 method TEXT, path TEXT, status INTEGER, streamed INTEGER,
1529 input_tokens INTEGER, cached_input_tokens INTEGER,
1530 cache_creation_tokens INTEGER, output_tokens INTEGER,
1531 reasoning_tokens INTEGER, cost_usd REAL, billing_bucket TEXT,
1532 req_body_path TEXT, upstream_req_body_path TEXT, resp_body_path TEXT,
1533 req_headers_json TEXT, resp_headers_json TEXT,
1534 error TEXT, account_id TEXT
1535 );
1536 INSERT INTO traces (id, ts_request_ms) VALUES ('old', 500);",
1537 )
1538 .unwrap();
1539 }
1540 let store = Store::open(dir).unwrap();
1541 let mut t = trace("new", 1000, Some("run-1"));
1542 t.tags = Some(r#"{"k":"v"}"#.into());
1543 t.client_ip = Some("127.0.0.1".into());
1544 t.key_fingerprint = Some("deadbeefdeadbeef".into());
1545 store.insert_trace(&t).unwrap();
1546 let rows = store.search_traces(&TraceFilter::default()).unwrap();
1547 assert_eq!(rows.len(), 2);
1548 assert_eq!(rows[0]["run_id"], "run-1");
1549 assert_eq!(rows[1]["id"], "old");
1550 assert_eq!(rows[1]["run_id"], Value::Null);
1551 assert_eq!(rows[1]["reasoning_effort"], Value::Null);
1552 assert_eq!(rows[1]["thinking_budget"], Value::Null);
1553 }
1554}
1555
1556fn seed_pricing(conn: &Connection) -> Result<()> {
1557 let models: Vec<Value> = serde_json::from_str(include_str!("models.json"))?;
1558 for m in models {
1559 conn.execute(
1560 "INSERT OR IGNORE INTO pricing (model, input_per_m, cached_input_per_m, cache_creation_per_m, output_per_m)
1561 VALUES (?1, ?2, ?3, ?4, ?5)",
1562 params![
1563 m["model"].as_str(),
1564 m["input_per_m"].as_f64(),
1565 m["cached_input_per_m"].as_f64(),
1566 m["cache_creation_per_m"].as_f64(),
1567 m["output_per_m"].as_f64(),
1568 ],
1569 )?;
1570 }
1571 Ok(())
1572}