1use std::io::Write;
2use std::path::PathBuf;
3use std::sync::Mutex;
4
5use alexandria_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);
48CREATE INDEX IF NOT EXISTS traces_session ON traces(session_id);
49CREATE INDEX IF NOT EXISTS traces_ts ON traces(ts_request_ms);
50CREATE INDEX IF NOT EXISTS traces_model ON traces(routed_model);
51
52CREATE TABLE IF NOT EXISTS pricing (
53 model TEXT PRIMARY KEY,
54 input_per_m REAL, cached_input_per_m REAL,
55 cache_creation_per_m REAL, output_per_m REAL
56);
57
58CREATE TABLE IF NOT EXISTS heartbeats (
59 ts_ms INTEGER NOT NULL,
60 provider TEXT NOT NULL,
61 account_id TEXT,
62 ok INTEGER,
63 status INTEGER,
64 latency_ms INTEGER,
65 message TEXT
66);
67CREATE INDEX IF NOT EXISTS heartbeats_ts ON heartbeats(ts_ms);
68"#;
69
70const TRACE_COLS: &str = "id, ts_request_ms, ts_response_ms, harness, client_format, upstream_provider,
71 requested_model, routed_model, status, streamed,
72 input_tokens, cached_input_tokens, cache_creation_tokens, output_tokens, reasoning_tokens,
73 cost_usd, billing_bucket, error, session_id, resp_body_path,
74 upstream_format, req_body_path, upstream_req_body_path, req_headers_json, resp_headers_json,
75 account_id, run_id, tags_json, client_ip, key_fingerprint";
76
77fn trace_row_json(r: &rusqlite::Row) -> rusqlite::Result<Value> {
78 let ts_request_ms = r.get::<_, i64>(1)?;
79 let ts_response_ms = r.get::<_, Option<i64>>(2)?;
80 Ok(json!({
81 "id": r.get::<_, String>(0)?,
82 "ts_request_ms": ts_request_ms,
83 "ts_response_ms": ts_response_ms,
84 "harness": r.get::<_, Option<String>>(3)?,
85 "client_format": r.get::<_, Option<String>>(4)?,
86 "upstream_provider": r.get::<_, Option<String>>(5)?,
87 "requested_model": r.get::<_, Option<String>>(6)?,
88 "routed_model": r.get::<_, Option<String>>(7)?,
89 "status": r.get::<_, Option<i64>>(8)?,
90 "streamed": r.get::<_, Option<i64>>(9)?,
91 "input_tokens": r.get::<_, Option<i64>>(10)?,
92 "cached_input_tokens": r.get::<_, Option<i64>>(11)?,
93 "cache_creation_tokens": r.get::<_, Option<i64>>(12)?,
94 "output_tokens": r.get::<_, Option<i64>>(13)?,
95 "reasoning_tokens": r.get::<_, Option<i64>>(14)?,
96 "cost_usd": r.get::<_, Option<f64>>(15)?,
97 "billing_bucket": r.get::<_, Option<String>>(16)?,
98 "error": r.get::<_, Option<String>>(17)?,
99 "session_id": r.get::<_, Option<String>>(18)?,
100 "resp_body_path": r.get::<_, Option<String>>(19)?,
101 "upstream_format": r.get::<_, Option<String>>(20)?,
102 "req_body_path": r.get::<_, Option<String>>(21)?,
103 "upstream_req_body_path": r.get::<_, Option<String>>(22)?,
104 "req_headers_json": r.get::<_, Option<String>>(23)?,
105 "resp_headers_json": r.get::<_, Option<String>>(24)?,
106 "account_id": r.get::<_, Option<String>>(25)?,
107 "run_id": r.get::<_, Option<String>>(26)?,
108 "tags_json": r.get::<_, Option<String>>(27)?,
109 "client_ip": r.get::<_, Option<String>>(28)?,
110 "key_fingerprint": r.get::<_, Option<String>>(29)?,
111 "latency_ms": ts_response_ms.map(|t| t - ts_request_ms),
112 }))
113}
114
115const DEFAULT_SEARCH_LIMIT: usize = 200;
116const MAX_SEARCH_LIMIT: usize = 5000;
117
118fn effective_limit(limit: usize) -> usize {
119 if limit == 0 {
120 DEFAULT_SEARCH_LIMIT
121 } else {
122 limit.min(MAX_SEARCH_LIMIT)
123 }
124}
125
126fn migrate_traces(conn: &Connection) -> Result<()> {
127 for col in [
128 "run_id TEXT",
129 "tags_json TEXT",
130 "client_ip TEXT",
131 "key_fingerprint TEXT",
132 ] {
133 if let Err(e) = conn.execute_batch(&format!("ALTER TABLE traces ADD COLUMN {col}")) {
134 if !e.to_string().contains("duplicate column name") {
135 return Err(e.into());
136 }
137 }
138 }
139 conn.execute_batch("CREATE INDEX IF NOT EXISTS traces_run ON traces(run_id)")?;
140 Ok(())
141}
142
143#[derive(Debug, Clone)]
144pub struct TraceFilter {
145 pub since_ms: Option<i64>,
146 pub until_ms: Option<i64>,
147 pub run_id: Option<String>,
148 pub session: Option<String>,
149 pub model: Option<String>,
150 pub provider: Option<String>,
151 pub path: Option<String>,
152 pub harness: Option<String>,
153 pub status: Option<i64>,
154 pub errors_only: bool,
155 pub key_fingerprint: Option<String>,
156 pub limit: usize,
157}
158
159impl Default for TraceFilter {
160 fn default() -> Self {
161 Self {
162 since_ms: None,
163 until_ms: None,
164 run_id: None,
165 session: None,
166 model: None,
167 provider: None,
168 path: None,
169 harness: None,
170 status: None,
171 errors_only: false,
172 key_fingerprint: None,
173 limit: DEFAULT_SEARCH_LIMIT,
174 }
175 }
176}
177
178pub struct Store {
179 conn: Mutex<Connection>,
180 pub data_dir: PathBuf,
181}
182
183impl Store {
184 pub fn open(data_dir: PathBuf) -> Result<Self> {
185 std::fs::create_dir_all(&data_dir)?;
186 let db_path = data_dir.join("alexandria.sqlite3");
187 let conn =
188 Connection::open(&db_path).with_context(|| format!("opening sqlite at {db_path:?}"))?;
189 conn.pragma_update(None, "journal_mode", "WAL")?;
190 conn.execute_batch(SCHEMA)?;
191 migrate_traces(&conn)?;
192 seed_pricing(&conn)?;
193 Ok(Self {
194 conn: Mutex::new(conn),
195 data_dir,
196 })
197 }
198
199 pub fn pricing_for(&self, model: &str) -> Option<Pricing> {
200 let conn = self.conn.lock().unwrap();
201 let mut stmt = conn
202 .prepare("SELECT model, input_per_m, cached_input_per_m, cache_creation_per_m, output_per_m FROM pricing")
203 .ok()?;
204 let rows: Vec<(String, Pricing)> = stmt
205 .query_map([], |r| {
206 Ok((
207 r.get::<_, String>(0)?,
208 Pricing {
209 input_per_m: r.get(1)?,
210 cached_input_per_m: r.get(2)?,
211 cache_creation_per_m: r.get(3)?,
212 output_per_m: r.get(4)?,
213 },
214 ))
215 })
216 .ok()?
217 .filter_map(|r| r.ok())
218 .collect();
219 rows.iter()
220 .filter(|(key, _)| model.starts_with(key.as_str()))
221 .max_by_key(|(key, _)| key.len())
222 .map(|(_, p)| p.clone())
223 }
224
225 pub fn pricing_models(&self) -> Vec<String> {
226 let conn = self.conn.lock().unwrap();
227 let Ok(mut stmt) = conn.prepare("SELECT model FROM pricing ORDER BY model") else {
228 return vec![];
229 };
230 stmt.query_map([], |r| r.get::<_, String>(0))
231 .map(|rows| rows.filter_map(|r| r.ok()).collect())
232 .unwrap_or_default()
233 }
234
235 pub fn insert_trace(&self, t: &TraceRecord) -> Result<()> {
236 let conn = self.conn.lock().unwrap();
237 conn.execute(
238 r#"INSERT OR REPLACE INTO traces (
239 id, ts_request_ms, ts_response_ms, session_id, harness, client_format,
240 upstream_provider, upstream_format, requested_model, routed_model,
241 method, path, status, streamed,
242 input_tokens, cached_input_tokens, cache_creation_tokens, output_tokens, reasoning_tokens,
243 cost_usd, billing_bucket,
244 req_body_path, upstream_req_body_path, resp_body_path,
245 req_headers_json, resp_headers_json, error, account_id,
246 run_id, tags_json, client_ip, key_fingerprint
247 ) 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)"#,
248 params![
249 t.id,
250 t.ts_request_ms,
251 t.ts_response_ms,
252 t.session_id,
253 t.harness,
254 t.client_format,
255 t.upstream_provider,
256 t.upstream_format,
257 t.requested_model,
258 t.routed_model,
259 t.method,
260 t.path,
261 t.status,
262 t.streamed.map(|b| b as i64),
263 t.usage.input_tokens,
264 t.usage.cached_input_tokens,
265 t.usage.cache_creation_tokens,
266 t.usage.output_tokens,
267 t.usage.reasoning_tokens,
268 t.cost_usd,
269 t.billing_bucket,
270 t.req_body_path,
271 t.upstream_req_body_path,
272 t.resp_body_path,
273 t.req_headers_json,
274 t.resp_headers_json,
275 t.error,
276 t.account_id,
277 t.run_id,
278 t.tags,
279 t.client_ip,
280 t.key_fingerprint,
281 ],
282 )?;
283 Ok(())
284 }
285
286 pub fn list_traces(
287 &self,
288 limit: usize,
289 session: Option<&str>,
290 model: Option<&str>,
291 ) -> Result<Vec<Value>> {
292 let filter = TraceFilter {
293 session: session.map(String::from),
294 model: model.map(String::from),
295 limit,
296 ..Default::default()
297 };
298 self.search_traces(&filter)
299 }
300
301 pub fn search_traces(&self, f: &TraceFilter) -> Result<Vec<Value>> {
302 let conn = self.conn.lock().unwrap();
303 let mut sql = format!("SELECT {TRACE_COLS} FROM traces WHERE 1=1");
304 let mut args: Vec<String> = vec![];
305 if let Some(since) = f.since_ms {
306 sql.push_str(" AND ts_request_ms >= ?");
307 args.push(since.to_string());
308 }
309 if let Some(until) = f.until_ms {
310 sql.push_str(" AND ts_request_ms <= ?");
311 args.push(until.to_string());
312 }
313 if let Some(r) = &f.run_id {
314 sql.push_str(" AND run_id = ?");
315 args.push(r.clone());
316 }
317 if let Some(s) = &f.session {
318 sql.push_str(" AND session_id = ?");
319 args.push(s.clone());
320 }
321 if let Some(m) = &f.model {
322 sql.push_str(" AND routed_model LIKE ?");
323 args.push(format!("%{m}%"));
324 }
325 if let Some(p) = &f.provider {
326 sql.push_str(" AND upstream_provider = ?");
327 args.push(p.clone());
328 }
329 if let Some(p) = &f.path {
330 sql.push_str(" AND path = ?");
331 args.push(p.clone());
332 }
333 if let Some(h) = &f.harness {
334 sql.push_str(" AND harness LIKE ?");
335 args.push(format!("%{h}%"));
336 }
337 if let Some(s) = f.status {
338 sql.push_str(" AND status = ?");
339 args.push(s.to_string());
340 }
341 if f.errors_only {
342 sql.push_str(" AND error IS NOT NULL");
343 }
344 if let Some(k) = &f.key_fingerprint {
345 sql.push_str(" AND key_fingerprint = ?");
346 args.push(k.clone());
347 }
348 sql.push_str(" ORDER BY ts_request_ms DESC LIMIT ?");
349 args.push(effective_limit(f.limit).to_string());
350 let mut stmt = conn.prepare(&sql)?;
351 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), trace_row_json)?;
352 Ok(rows.filter_map(|r| r.ok()).collect())
353 }
354
355 pub fn sessions(&self, since_ms: Option<i64>, limit: usize) -> Result<Vec<Value>> {
356 let conn = self.conn.lock().unwrap();
357 let mut sql = String::from(
358 "SELECT session_id, MAX(run_id), MIN(ts_request_ms), MAX(ts_request_ms), COUNT(*),
359 GROUP_CONCAT(DISTINCT routed_model), MAX(harness),
360 COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
361 COALESCE(SUM(cost_usd),0.0),
362 COALESCE(SUM(CASE WHEN error IS NOT NULL OR status >= 400 THEN 1 ELSE 0 END),0),
363 (SELECT t2.status FROM traces t2 WHERE t2.session_id = traces.session_id
364 ORDER BY t2.ts_request_ms DESC LIMIT 1),
365 GROUP_CONCAT(tags_json, char(31))
366 FROM traces WHERE session_id IS NOT NULL",
367 );
368 let mut args: Vec<String> = vec![];
369 if let Some(since) = since_ms {
370 sql.push_str(" AND ts_request_ms >= ?");
371 args.push(since.to_string());
372 }
373 sql.push_str(" GROUP BY session_id ORDER BY MAX(ts_request_ms) DESC LIMIT ?");
374 let limit = if limit == 0 {
375 DEFAULT_SEARCH_LIMIT
376 } else {
377 limit.min(1000)
378 };
379 args.push(limit.to_string());
380 let mut stmt = conn.prepare(&sql)?;
381 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), |r| {
382 let models: Vec<String> = r
383 .get::<_, Option<String>>(5)?
384 .map(|s| s.split(',').map(str::to_string).collect())
385 .unwrap_or_default();
386 let mut tags = serde_json::Map::new();
387 if let Some(raw) = r.get::<_, Option<String>>(12)? {
388 for piece in raw.split('\u{1f}') {
389 if let Ok(Value::Object(o)) = serde_json::from_str::<Value>(piece) {
390 tags.extend(o);
391 }
392 }
393 }
394 Ok(json!({
395 "session_id": r.get::<_, String>(0)?,
396 "run_id": r.get::<_, Option<String>>(1)?,
397 "first_ts_ms": r.get::<_, Option<i64>>(2)?,
398 "last_ts_ms": r.get::<_, Option<i64>>(3)?,
399 "trace_count": r.get::<_, i64>(4)?,
400 "models": models,
401 "harness": r.get::<_, Option<String>>(6)?,
402 "total_input_tokens": r.get::<_, i64>(7)?,
403 "total_output_tokens": r.get::<_, i64>(8)?,
404 "total_cost_usd": r.get::<_, f64>(9)?,
405 "errors": r.get::<_, i64>(10)?,
406 "last_status": r.get::<_, Option<i64>>(11)?,
407 "tags": tags,
408 }))
409 })?;
410 Ok(rows.filter_map(|r| r.ok()).collect())
411 }
412
413 pub fn session_traces(&self, session_id: &str, since_ms: Option<i64>) -> Result<Vec<Value>> {
414 let conn = self.conn.lock().unwrap();
415 let mut sql = format!("SELECT {TRACE_COLS} FROM traces WHERE session_id = ?");
416 let mut args = vec![session_id.to_string()];
417 if let Some(since) = since_ms {
418 sql.push_str(" AND ts_request_ms >= ?");
419 args.push(since.to_string());
420 }
421 sql.push_str(" ORDER BY ts_request_ms ASC");
422 let mut stmt = conn.prepare(&sql)?;
423 let rows = stmt.query_map(rusqlite::params_from_iter(args.iter()), trace_row_json)?;
424 Ok(rows.filter_map(|r| r.ok()).collect())
425 }
426
427 pub fn get_trace(&self, id: &str) -> Result<Option<Value>> {
428 let conn = self.conn.lock().unwrap();
429 let row = conn
430 .query_row(
431 &format!("SELECT {TRACE_COLS} FROM traces WHERE id = ?1"),
432 params![id],
433 trace_row_json,
434 )
435 .optional()?;
436 Ok(row)
437 }
438
439 pub fn delete_trace(&self, id: &str) -> Result<Vec<String>> {
440 let conn = self.conn.lock().unwrap();
441 let paths: Option<(Option<String>, Option<String>, Option<String>)> = conn
442 .query_row(
443 "SELECT req_body_path, upstream_req_body_path, resp_body_path
444 FROM traces WHERE id = ?1",
445 params![id],
446 |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
447 )
448 .optional()?;
449 let Some((req, upstream, resp)) = paths else {
450 anyhow::bail!("trace not found: {id}");
451 };
452 conn.execute("DELETE FROM traces WHERE id = ?1", params![id])?;
453 Ok([req, upstream, resp].into_iter().flatten().collect())
454 }
455
456 pub fn run_summary(&self, run_id: &str) -> Result<Value> {
457 let conn = self.conn.lock().unwrap();
458 let (trace_count, first_ts_ms, last_ts_ms, total_input, total_output, total_cost, errors) =
459 conn.query_row(
460 "SELECT COUNT(*), MIN(ts_request_ms), MAX(ts_request_ms),
461 COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
462 COALESCE(SUM(cost_usd),0.0),
463 COALESCE(SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END),0)
464 FROM traces WHERE run_id = ?1",
465 params![run_id],
466 |r| {
467 Ok((
468 r.get::<_, i64>(0)?,
469 r.get::<_, Option<i64>>(1)?,
470 r.get::<_, Option<i64>>(2)?,
471 r.get::<_, i64>(3)?,
472 r.get::<_, i64>(4)?,
473 r.get::<_, f64>(5)?,
474 r.get::<_, i64>(6)?,
475 ))
476 },
477 )?;
478 let mut status_counts = serde_json::Map::new();
479 let mut stmt = conn.prepare(
480 "SELECT status, COUNT(*) FROM traces WHERE run_id = ?1 GROUP BY status",
481 )?;
482 let pairs = stmt.query_map(params![run_id], |r| {
483 Ok((r.get::<_, Option<i64>>(0)?, r.get::<_, i64>(1)?))
484 })?;
485 for pair in pairs.flatten() {
486 let key = pair.0.map(|s| s.to_string()).unwrap_or_else(|| "none".into());
487 status_counts.insert(key, json!(pair.1));
488 }
489 let distinct = |col: &str| -> Result<Vec<String>> {
490 let mut stmt = conn.prepare(&format!(
491 "SELECT DISTINCT {col} FROM traces WHERE run_id = ?1 AND {col} IS NOT NULL ORDER BY {col}"
492 ))?;
493 let vals = stmt
494 .query_map(params![run_id], |r| r.get::<_, String>(0))?
495 .filter_map(|r| r.ok())
496 .collect();
497 Ok(vals)
498 };
499 let models = distinct("routed_model")?;
500 let providers = distinct("upstream_provider")?;
501 let mut tags = serde_json::Map::new();
502 let mut stmt = conn.prepare(
503 "SELECT tags_json FROM traces WHERE run_id = ?1 AND tags_json IS NOT NULL ORDER BY ts_request_ms",
504 )?;
505 let tag_rows = stmt.query_map(params![run_id], |r| r.get::<_, String>(0))?;
506 for raw in tag_rows.flatten() {
507 if let Ok(Value::Object(o)) = serde_json::from_str::<Value>(&raw) {
508 tags.extend(o);
509 }
510 }
511 Ok(json!({
512 "run_id": run_id,
513 "trace_count": trace_count,
514 "first_ts_ms": first_ts_ms,
515 "last_ts_ms": last_ts_ms,
516 "status_counts": status_counts,
517 "models": models,
518 "providers": providers,
519 "total_input_tokens": total_input,
520 "total_output_tokens": total_output,
521 "total_cost_usd": total_cost,
522 "tags": tags,
523 "errors": errors,
524 }))
525 }
526
527 pub fn run_artifacts(&self, run_id: &str) -> Result<Vec<Value>> {
528 let conn = self.conn.lock().unwrap();
529 let mut stmt = conn.prepare(
530 "SELECT id, req_body_path, upstream_req_body_path, resp_body_path
531 FROM traces WHERE run_id = ?1 ORDER BY ts_request_ms",
532 )?;
533 let rows: Vec<(String, Option<String>, Option<String>, Option<String>)> = stmt
534 .query_map(params![run_id], |r| {
535 Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?))
536 })?
537 .filter_map(|r| r.ok())
538 .collect();
539 let mut out = Vec::new();
540 for (trace_id, req, upstream_req, resp) in rows {
541 for (kind, path) in [
542 ("request", req),
543 ("upstream-request", upstream_req),
544 ("response", resp),
545 ] {
546 let Some(path) = path else { continue };
547 let size_bytes = std::fs::metadata(&path).ok().map(|m| m.len());
548 out.push(json!({
549 "trace_id": trace_id,
550 "kind": kind,
551 "path": path,
552 "exists": size_bytes.is_some(),
553 "size_bytes": size_bytes,
554 }));
555 }
556 }
557 Ok(out)
558 }
559
560 pub fn insert_heartbeat(
561 &self,
562 ts_ms: i64,
563 provider: &str,
564 account_id: Option<&str>,
565 ok: bool,
566 status: Option<i64>,
567 latency_ms: i64,
568 message: &str,
569 ) -> Result<()> {
570 let conn = self.conn.lock().unwrap();
571 conn.execute(
572 "INSERT INTO heartbeats (ts_ms, provider, account_id, ok, status, latency_ms, message)
573 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
574 params![ts_ms, provider, account_id, ok as i64, status, latency_ms, message],
575 )?;
576 Ok(())
577 }
578
579 pub fn last_heartbeats(&self) -> Result<Vec<Value>> {
580 let conn = self.conn.lock().unwrap();
581 let mut stmt = conn.prepare(
582 "SELECT h.ts_ms, h.provider, h.account_id, h.ok, h.status, h.latency_ms, h.message
583 FROM heartbeats h
584 JOIN (SELECT provider, MAX(ts_ms) AS ts FROM heartbeats GROUP BY provider) latest
585 ON h.provider = latest.provider AND h.ts_ms = latest.ts",
586 )?;
587 let rows = stmt.query_map([], |r| {
588 Ok(json!({
589 "ts_ms": r.get::<_, i64>(0)?,
590 "provider": r.get::<_, String>(1)?,
591 "account_id": r.get::<_, Option<String>>(2)?,
592 "ok": r.get::<_, i64>(3)? == 1,
593 "status": r.get::<_, Option<i64>>(4)?,
594 "latency_ms": r.get::<_, i64>(5)?,
595 "message": r.get::<_, String>(6)?,
596 }))
597 })?;
598 Ok(rows.filter_map(|r| r.ok()).collect())
599 }
600
601 pub fn latest_provider_headers(&self) -> Result<Vec<(String, i64, String)>> {
602 let conn = self.conn.lock().unwrap();
603 let mut stmt = conn.prepare(
604 "SELECT t.upstream_provider, t.ts_request_ms, t.resp_headers_json
605 FROM traces t
606 JOIN (SELECT upstream_provider p, MAX(ts_request_ms) ts FROM traces
607 WHERE status >= 200 AND status < 300
608 AND resp_headers_json IS NOT NULL AND upstream_provider IS NOT NULL
609 GROUP BY upstream_provider) latest
610 ON t.upstream_provider = latest.p AND t.ts_request_ms = latest.ts
611 WHERE t.resp_headers_json IS NOT NULL",
612 )?;
613 let rows: Vec<(String, i64, String)> = stmt
614 .query_map([], |r| {
615 Ok((
616 r.get::<_, String>(0)?,
617 r.get::<_, i64>(1)?,
618 r.get::<_, String>(2)?,
619 ))
620 })?
621 .filter_map(|r| r.ok())
622 .collect();
623 let mut seen = std::collections::HashMap::new();
624 for row in rows {
625 seen.entry(row.0.clone()).or_insert(row);
626 }
627 Ok(seen.into_values().collect())
628 }
629
630 pub fn analytics(&self, since_ms: i64) -> Result<Value> {
631 let conn = self.conn.lock().unwrap();
632 let mut stmt = conn.prepare(
633 "SELECT routed_model, upstream_provider, billing_bucket, COUNT(*),
634 COALESCE(SUM(input_tokens),0), COALESCE(SUM(cached_input_tokens),0),
635 COALESCE(SUM(output_tokens),0), COALESCE(SUM(cost_usd),0.0),
636 SUM(CASE WHEN status >= 200 AND status < 300 THEN 0 ELSE 1 END),
637 AVG(CASE WHEN ts_response_ms IS NOT NULL THEN ts_response_ms - ts_request_ms END)
638 FROM traces WHERE ts_request_ms >= ?1
639 GROUP BY routed_model, upstream_provider, billing_bucket
640 ORDER BY SUM(cost_usd) DESC",
641 )?;
642 let rows: Vec<Value> = stmt
643 .query_map(params![since_ms], |r| {
644 Ok(json!({
645 "routed_model": r.get::<_, Option<String>>(0)?,
646 "upstream_provider": r.get::<_, Option<String>>(1)?,
647 "billing_bucket": r.get::<_, Option<String>>(2)?,
648 "requests": r.get::<_, i64>(3)?,
649 "input_tokens": r.get::<_, i64>(4)?,
650 "cached_input_tokens": r.get::<_, i64>(5)?,
651 "output_tokens": r.get::<_, i64>(6)?,
652 "cost_usd": r.get::<_, f64>(7)?,
653 "errors": r.get::<_, Option<i64>>(8)?,
654 "avg_latency_ms": r.get::<_, Option<f64>>(9)?,
655 }))
656 })?
657 .filter_map(|r| r.ok())
658 .collect();
659 let (mut requests, mut cost, mut errors) = (0i64, 0f64, 0i64);
660 let mut buckets: std::collections::HashMap<String, f64> = Default::default();
661 for row in &rows {
662 requests += row["requests"].as_i64().unwrap_or(0);
663 cost += row["cost_usd"].as_f64().unwrap_or(0.0);
664 errors += row["errors"].as_i64().unwrap_or(0);
665 let bucket = row["billing_bucket"].as_str().unwrap_or("unknown").to_string();
666 *buckets.entry(bucket).or_default() += row["cost_usd"].as_f64().unwrap_or(0.0);
667 }
668 Ok(json!({
669 "since_ms": since_ms,
670 "totals": {"requests": requests, "cost_usd": cost, "errors": errors, "cost_by_bucket": buckets},
671 "by_model": rows,
672 }))
673 }
674
675 pub fn write_body(&self, trace_id: &str, kind: &str, bytes: &[u8]) -> Result<String> {
676 let date = Utc::now().format("%Y-%m-%d").to_string();
677 let dir = self.data_dir.join("bodies").join(&date);
678 std::fs::create_dir_all(&dir)?;
679 let path = dir.join(format!("{trace_id}.{kind}.gz"));
680 let file = std::fs::File::create(&path)?;
681 let mut enc = GzEncoder::new(file, Compression::default());
682 enc.write_all(bytes)?;
683 enc.finish()?;
684 Ok(path.to_string_lossy().to_string())
685 }
686}
687
688#[cfg(test)]
689mod tests {
690 use super::*;
691
692 fn tmpdir(name: &str) -> PathBuf {
693 let dir = std::env::temp_dir().join(format!(
694 "alexandria-store-test-{name}-{}",
695 std::process::id()
696 ));
697 let _ = std::fs::remove_dir_all(&dir);
698 std::fs::create_dir_all(&dir).unwrap();
699 dir
700 }
701
702 fn trace(id: &str, ts: i64, run: Option<&str>) -> TraceRecord {
703 TraceRecord {
704 id: id.into(),
705 ts_request_ms: ts,
706 ts_response_ms: Some(ts + 250),
707 status: Some(200),
708 routed_model: Some("claude-haiku-4-5".into()),
709 upstream_provider: Some("anthropic".into()),
710 run_id: run.map(String::from),
711 usage: alexandria_core::Usage {
712 input_tokens: Some(10),
713 output_tokens: Some(5),
714 ..Default::default()
715 },
716 cost_usd: Some(0.001),
717 ..Default::default()
718 }
719 }
720
721 #[test]
722 fn run_summary_aggregates() {
723 let store = Store::open(tmpdir("summary")).unwrap();
724 let mut a = trace("a", 1000, Some("run-1"));
725 a.tags = Some(r#"{"suite":"swebench"}"#.into());
726 let mut b = trace("b", 2000, Some("run-1"));
727 b.tags = Some(r#"{"case":"astropy-1"}"#.into());
728 b.status = Some(500);
729 b.error = Some("boom".into());
730 b.routed_model = Some("gpt-5.5".into());
731 b.upstream_provider = Some("openai".into());
732 let c = trace("c", 3000, Some("run-2"));
733 for t in [&a, &b, &c] {
734 store.insert_trace(t).unwrap();
735 }
736 let s = store.run_summary("run-1").unwrap();
737 assert_eq!(s["trace_count"], 2);
738 assert_eq!(s["first_ts_ms"], 1000);
739 assert_eq!(s["last_ts_ms"], 2000);
740 assert_eq!(s["status_counts"]["200"], 1);
741 assert_eq!(s["status_counts"]["500"], 1);
742 assert_eq!(
743 s["models"],
744 json!(["claude-haiku-4-5", "gpt-5.5"])
745 );
746 assert_eq!(s["providers"], json!(["anthropic", "openai"]));
747 assert_eq!(s["total_input_tokens"], 20);
748 assert_eq!(s["total_output_tokens"], 10);
749 assert_eq!(s["tags"]["suite"], "swebench");
750 assert_eq!(s["tags"]["case"], "astropy-1");
751 assert_eq!(s["errors"], 1);
752 let missing = store.run_summary("nope").unwrap();
753 assert_eq!(missing["trace_count"], 0);
754 }
755
756 #[test]
757 fn search_traces_filters() {
758 let store = Store::open(tmpdir("search")).unwrap();
759 let mut a = trace("a", 1000, Some("run-1"));
760 a.key_fingerprint = Some("deadbeefdeadbeef".into());
761 let mut b = trace("b", 2000, Some("run-1"));
762 b.status = Some(429);
763 b.error = Some("rate limited".into());
764 let c = trace("c", 3000, None);
765 for t in [&a, &b, &c] {
766 store.insert_trace(t).unwrap();
767 }
768 let all = store.search_traces(&TraceFilter::default()).unwrap();
769 assert_eq!(all.len(), 3);
770 assert_eq!(all[0]["id"], "c");
771 assert_eq!(all[0]["latency_ms"], 250);
772 let window = store
773 .search_traces(&TraceFilter {
774 since_ms: Some(1500),
775 until_ms: Some(2500),
776 ..Default::default()
777 })
778 .unwrap();
779 assert_eq!(window.len(), 1);
780 assert_eq!(window[0]["id"], "b");
781 let by_run = store
782 .search_traces(&TraceFilter {
783 run_id: Some("run-1".into()),
784 ..Default::default()
785 })
786 .unwrap();
787 assert_eq!(by_run.len(), 2);
788 let by_status = store
789 .search_traces(&TraceFilter {
790 status: Some(429),
791 ..Default::default()
792 })
793 .unwrap();
794 assert_eq!(by_status.len(), 1);
795 let errors = store
796 .search_traces(&TraceFilter {
797 errors_only: true,
798 ..Default::default()
799 })
800 .unwrap();
801 assert_eq!(errors.len(), 1);
802 assert_eq!(errors[0]["id"], "b");
803 let by_key = store
804 .search_traces(&TraceFilter {
805 key_fingerprint: Some("deadbeefdeadbeef".into()),
806 ..Default::default()
807 })
808 .unwrap();
809 assert_eq!(by_key.len(), 1);
810 assert_eq!(by_key[0]["id"], "a");
811 let limited = store
812 .search_traces(&TraceFilter {
813 limit: 1,
814 ..Default::default()
815 })
816 .unwrap();
817 assert_eq!(limited.len(), 1);
818 }
819
820 #[test]
821 fn sessions_aggregate_and_order() {
822 let store = Store::open(tmpdir("sessions")).unwrap();
823 let mut a = trace("a", 1000, Some("run-1"));
824 a.session_id = Some("ses_1".into());
825 a.tags = Some(r#"{"suite":"swebench"}"#.into());
826 a.harness = Some("codex".into());
827 let mut b = trace("b", 2000, None);
828 b.session_id = Some("ses_1".into());
829 b.status = Some(500);
830 b.error = Some("boom".into());
831 b.routed_model = Some("gpt-5.5".into());
832 b.tags = Some(r#"{"case":"x1"}"#.into());
833 let mut c = trace("c", 5000, None);
834 c.session_id = Some("ses_2".into());
835 let d = trace("d", 9000, None);
836 for t in [&a, &b, &c, &d] {
837 store.insert_trace(t).unwrap();
838 }
839 let sessions = store.sessions(None, 0).unwrap();
840 assert_eq!(sessions.len(), 2);
841 assert_eq!(sessions[0]["session_id"], "ses_2");
842 let s1 = &sessions[1];
843 assert_eq!(s1["session_id"], "ses_1");
844 assert_eq!(s1["run_id"], "run-1");
845 assert_eq!(s1["first_ts_ms"], 1000);
846 assert_eq!(s1["last_ts_ms"], 2000);
847 assert_eq!(s1["trace_count"], 2);
848 assert_eq!(s1["harness"], "codex");
849 assert_eq!(s1["total_input_tokens"], 20);
850 assert_eq!(s1["total_output_tokens"], 10);
851 assert_eq!(s1["errors"], 1);
852 assert_eq!(s1["last_status"], 500);
853 assert_eq!(s1["tags"]["suite"], "swebench");
854 assert_eq!(s1["tags"]["case"], "x1");
855 let models: Vec<String> = s1["models"]
856 .as_array()
857 .unwrap()
858 .iter()
859 .map(|m| m.as_str().unwrap().to_string())
860 .collect();
861 assert!(models.contains(&"claude-haiku-4-5".to_string()));
862 assert!(models.contains(&"gpt-5.5".to_string()));
863 let recent = store.sessions(Some(3000), 0).unwrap();
864 assert_eq!(recent.len(), 1);
865 assert_eq!(recent[0]["session_id"], "ses_2");
866 let limited = store.sessions(None, 1).unwrap();
867 assert_eq!(limited.len(), 1);
868 }
869
870 #[test]
871 fn session_traces_ascending() {
872 let store = Store::open(tmpdir("session-traces")).unwrap();
873 for (id, ts) in [("a", 3000i64), ("b", 1000), ("c", 2000)] {
874 let mut t = trace(id, ts, None);
875 t.session_id = Some("ses_1".into());
876 t.upstream_format = Some("anthropic".into());
877 t.req_body_path = Some(format!("/bodies/{id}.request.json.gz"));
878 store.insert_trace(&t).unwrap();
879 }
880 let rows = store.session_traces("ses_1", None).unwrap();
881 assert_eq!(rows.len(), 3);
882 assert_eq!(rows[0]["id"], "b");
883 assert_eq!(rows[2]["id"], "a");
884 assert_eq!(rows[0]["upstream_format"], "anthropic");
885 assert_eq!(rows[0]["client_format"], Value::Null);
886 assert_eq!(rows[0]["req_body_path"], "/bodies/b.request.json.gz");
887 let windowed = store.session_traces("ses_1", Some(1500)).unwrap();
888 assert_eq!(windowed.len(), 2);
889 assert_eq!(windowed[0]["id"], "c");
890 assert!(store.session_traces("nope", None).unwrap().is_empty());
891 }
892
893 #[test]
894 fn get_and_delete_trace() {
895 let store = Store::open(tmpdir("delete")).unwrap();
896 let mut t = trace("a", 1000, None);
897 t.req_body_path = Some(
898 store
899 .write_body("a", "request.json", b"{\"model\":\"x\"}")
900 .unwrap(),
901 );
902 t.resp_body_path = Some("/nonexistent/a.response.body.gz".into());
903 store.insert_trace(&t).unwrap();
904 let row = store.get_trace("a").unwrap().unwrap();
905 assert_eq!(row["id"], "a");
906 assert_eq!(row["resp_body_path"], "/nonexistent/a.response.body.gz");
907 assert!(store.get_trace("missing").unwrap().is_none());
908 let paths = store.delete_trace("a").unwrap();
909 assert_eq!(paths.len(), 2);
910 assert!(store.get_trace("a").unwrap().is_none());
911 assert!(store.delete_trace("a").is_err());
912 }
913
914 #[test]
915 fn limit_defaults_and_caps() {
916 assert_eq!(effective_limit(0), 200);
917 assert_eq!(effective_limit(50), 50);
918 assert_eq!(effective_limit(9000), 5000);
919 }
920
921 #[test]
922 fn run_artifacts_reports_files() {
923 let dir = tmpdir("artifacts");
924 let store = Store::open(dir).unwrap();
925 let mut t = trace("a", 1000, Some("run-1"));
926 t.req_body_path = Some(
927 store
928 .write_body("a", "request.json", b"{\"model\":\"x\"}")
929 .unwrap(),
930 );
931 t.resp_body_path = Some("/nonexistent/a.response.body.gz".into());
932 store.insert_trace(&t).unwrap();
933 let arts = store.run_artifacts("run-1").unwrap();
934 assert_eq!(arts.len(), 2);
935 assert_eq!(arts[0]["kind"], "request");
936 assert_eq!(arts[0]["exists"], true);
937 assert!(arts[0]["size_bytes"].as_u64().unwrap() > 0);
938 assert_eq!(arts[1]["kind"], "response");
939 assert_eq!(arts[1]["exists"], false);
940 assert_eq!(arts[1]["size_bytes"], Value::Null);
941 assert!(store.run_artifacts("nope").unwrap().is_empty());
942 }
943
944 #[test]
945 fn reopen_keeps_working() {
946 let dir = tmpdir("reopen");
947 {
948 let store = Store::open(dir.clone()).unwrap();
949 store.insert_trace(&trace("a", 1000, Some("run-1"))).unwrap();
950 }
951 let store = Store::open(dir).unwrap();
952 store.insert_trace(&trace("b", 2000, Some("run-1"))).unwrap();
953 let s = store.run_summary("run-1").unwrap();
954 assert_eq!(s["trace_count"], 2);
955 }
956
957 #[test]
958 fn migrates_old_schema() {
959 let dir = tmpdir("migrate");
960 let db_path = dir.join("alexandria.sqlite3");
961 {
962 let conn = Connection::open(&db_path).unwrap();
963 conn.execute_batch(
964 "CREATE TABLE traces (
965 id TEXT PRIMARY KEY,
966 ts_request_ms INTEGER NOT NULL,
967 ts_response_ms INTEGER,
968 session_id TEXT, harness TEXT, client_format TEXT,
969 upstream_provider TEXT, upstream_format TEXT,
970 requested_model TEXT, routed_model TEXT,
971 method TEXT, path TEXT, status INTEGER, streamed INTEGER,
972 input_tokens INTEGER, cached_input_tokens INTEGER,
973 cache_creation_tokens INTEGER, output_tokens INTEGER,
974 reasoning_tokens INTEGER, cost_usd REAL, billing_bucket TEXT,
975 req_body_path TEXT, upstream_req_body_path TEXT, resp_body_path TEXT,
976 req_headers_json TEXT, resp_headers_json TEXT,
977 error TEXT, account_id TEXT
978 );
979 INSERT INTO traces (id, ts_request_ms) VALUES ('old', 500);",
980 )
981 .unwrap();
982 }
983 let store = Store::open(dir).unwrap();
984 let mut t = trace("new", 1000, Some("run-1"));
985 t.tags = Some(r#"{"k":"v"}"#.into());
986 t.client_ip = Some("127.0.0.1".into());
987 t.key_fingerprint = Some("deadbeefdeadbeef".into());
988 store.insert_trace(&t).unwrap();
989 let rows = store.search_traces(&TraceFilter::default()).unwrap();
990 assert_eq!(rows.len(), 2);
991 assert_eq!(rows[0]["run_id"], "run-1");
992 assert_eq!(rows[1]["id"], "old");
993 assert_eq!(rows[1]["run_id"], Value::Null);
994 }
995}
996
997fn seed_pricing(conn: &Connection) -> Result<()> {
998 let models: Vec<Value> = serde_json::from_str(include_str!("models.json"))?;
999 for m in models {
1000 conn.execute(
1001 "INSERT OR IGNORE INTO pricing (model, input_per_m, cached_input_per_m, cache_creation_per_m, output_per_m)
1002 VALUES (?1, ?2, ?3, ?4, ?5)",
1003 params![
1004 m["model"].as_str(),
1005 m["input_per_m"].as_f64(),
1006 m["cached_input_per_m"].as_f64(),
1007 m["cache_creation_per_m"].as_f64(),
1008 m["output_per_m"].as_f64(),
1009 ],
1010 )?;
1011 }
1012 Ok(())
1013}