assay_core/storage/
store.rs

1use super::now_rfc3339ish;
2use crate::model::{AttemptRow, EvalConfig, LlmResponse, TestResultRow, TestStatus};
3use crate::trace::schema::{EpisodeEnd, EpisodeStart, StepEntry, ToolCallEntry, TraceEvent};
4use anyhow::Context;
5use rusqlite::{params, Connection};
6use std::path::Path;
7use std::sync::{Arc, Mutex};
8
9#[path = "store_internal/mod.rs"]
10mod store_internal;
11#[path = "store_trace.rs"]
12mod store_trace;
13
14#[derive(Clone)]
15pub struct Store {
16    pub conn: Arc<Mutex<Connection>>,
17}
18
19pub struct StoreStats {
20    pub runs: Option<u64>,
21    pub results: Option<u64>,
22    pub last_run_id: Option<i64>,
23    pub last_run_at: Option<String>,
24    pub version: Option<String>,
25}
26
27impl Store {
28    pub fn open(path: &Path) -> anyhow::Result<Self> {
29        let conn = Connection::open(path).context("failed to open sqlite db")?;
30        conn.execute("PRAGMA foreign_keys = ON", [])?;
31        Ok(Self {
32            conn: Arc::new(Mutex::new(conn)),
33        })
34    }
35
36    pub fn memory() -> anyhow::Result<Self> {
37        // SQLite in-memory DB
38        let conn = Connection::open_in_memory().context("failed to open in-memory sqlite db")?;
39        Ok(Self {
40            conn: Arc::new(Mutex::new(conn)),
41        })
42    }
43
44    pub fn init_schema(&self) -> anyhow::Result<()> {
45        let conn = self.conn.lock().unwrap();
46        conn.execute_batch(crate::storage::schema::DDL)?;
47
48        // v0.3.0 Migrations
49        migrate_v030(&conn)?;
50
51        // Ensure attempts table exists (covered by DDL if creating fresh, but good to be explicit if DDL didn't run on existing DB)
52        // DDL handles IF NOT EXISTS for attempts.
53
54        // Index on fingerprint for speed (CREATE INDEX IF NOT EXISTS is valid sqlite)
55        let _ = conn.execute(
56            "CREATE INDEX IF NOT EXISTS idx_results_fingerprint ON results(fingerprint)",
57            [],
58        );
59
60        Ok(())
61    }
62
63    pub fn fetch_recent_results(
64        &self,
65        suite: &str,
66        limit: u32,
67    ) -> anyhow::Result<Vec<crate::model::TestResultRow>> {
68        let conn = self.conn.lock().unwrap();
69        let mut stmt = conn.prepare(
70            "SELECT
71                r.test_id, r.outcome, r.duration_ms, r.score, r.attempts_json,
72                r.fingerprint, r.skip_reason
73             FROM results r
74             JOIN runs ON r.run_id = runs.id
75             WHERE runs.suite = ?1
76             ORDER BY r.id DESC
77             LIMIT ?2",
78        )?;
79
80        let rows = stmt.query_map(rusqlite::params![suite, limit], row_to_test_result)?;
81
82        let mut results = Vec::new();
83        for r in rows {
84            results.push(r?);
85        }
86        Ok(results)
87    }
88
89    pub fn fetch_results_for_last_n_runs(
90        &self,
91        suite: &str,
92        n: u32,
93    ) -> anyhow::Result<Vec<crate::model::TestResultRow>> {
94        let conn = self.conn.lock().unwrap();
95        let mut stmt = conn.prepare(
96            "SELECT
97                r.test_id, r.outcome, r.duration_ms, r.score, r.attempts_json,
98                r.fingerprint, r.skip_reason
99             FROM results r
100             JOIN runs ON r.run_id = runs.id
101             WHERE runs.id IN (
102                 SELECT id FROM runs WHERE suite = ?1 ORDER BY id DESC LIMIT ?2
103             )
104             ORDER BY r.id DESC",
105        )?;
106
107        let rows = stmt.query_map(rusqlite::params![suite, n], row_to_test_result)?;
108
109        let mut results = Vec::new();
110        for r in rows {
111            results.push(r?);
112        }
113        Ok(results)
114    }
115
116    pub fn get_latest_run_id(&self, suite: &str) -> anyhow::Result<Option<i64>> {
117        let conn = self.conn.lock().unwrap();
118        let mut stmt =
119            conn.prepare("SELECT id FROM runs WHERE suite = ?1 ORDER BY id DESC LIMIT 1")?;
120        let mut rows = stmt.query(params![suite])?;
121        if let Some(row) = rows.next()? {
122            Ok(Some(row.get(0)?))
123        } else {
124            Ok(None)
125        }
126    }
127
128    pub fn fetch_results_for_run(
129        &self,
130        run_id: i64,
131    ) -> anyhow::Result<Vec<crate::model::TestResultRow>> {
132        let conn = self.conn.lock().unwrap();
133        let mut stmt = conn.prepare(
134            "SELECT
135                r.test_id, r.outcome, r.duration_ms, r.score, r.attempts_json,
136                r.fingerprint, r.skip_reason
137             FROM results r
138             WHERE r.run_id = ?1
139             ORDER BY r.test_id ASC",
140        )?;
141
142        let rows = stmt.query_map(params![run_id], row_to_test_result)?;
143
144        let mut results = Vec::new();
145        for r in rows {
146            results.push(r?);
147        }
148        Ok(results)
149    }
150
151    pub fn get_last_passing_by_fingerprint(
152        &self,
153        fingerprint: &str,
154    ) -> anyhow::Result<Option<TestResultRow>> {
155        let conn = self.conn.lock().unwrap();
156        // We want the most recent passing result for this fingerprint.
157        // run_id DESC ensures recency.
158        let mut stmt = conn.prepare(
159            "SELECT r.test_id, r.score, r.duration_ms, r.output_json, r.skip_reason, run.id, run.started_at
160             FROM results r
161             JOIN runs run ON r.run_id = run.id
162             WHERE r.fingerprint = ?1 AND r.outcome = 'pass'
163             ORDER BY r.id DESC LIMIT 1"
164        )?;
165
166        let mut rows = stmt.query(params![fingerprint])?;
167        if let Some(row) = rows.next()? {
168            let status = TestStatus::Pass;
169
170            let skip_reason: Option<String> = row.get(4)?;
171            let run_id: i64 = row.get(5)?;
172            let started_at: String = row.get(6)?;
173
174            let details = serde_json::json!({
175                "skip": {
176                    "reason": skip_reason.clone().unwrap_or_else(|| "fingerprint_match".into()),
177                    "fingerprint": fingerprint,
178                    "previous_run_id": run_id,
179                    "previous_at": started_at,
180                    "origin_run_id": run_id,
181                    "previous_score": row.get::<_, Option<f64>>(1)?
182                }
183            });
184
185            Ok(Some(TestResultRow {
186                test_id: row.get(0)?,
187                status,
188                message: skip_reason.unwrap_or_else(|| "fingerprint_match".to_string()),
189                score: row.get(1)?,
190                duration_ms: row.get(2)?,
191                cached: true,
192                details,
193                fingerprint: Some(fingerprint.to_string()),
194                skip_reason: None,
195                attempts: None,
196                error_policy_applied: None,
197            }))
198        } else {
199            Ok(None)
200        }
201    }
202
203    pub fn insert_run(&self, suite: &str) -> anyhow::Result<i64> {
204        let started_at = now_rfc3339ish();
205        let conn = self.conn.lock().unwrap();
206        insert_run_row(&conn, suite, &started_at, "running", None)
207    }
208
209    pub fn create_run(&self, cfg: &EvalConfig) -> anyhow::Result<i64> {
210        let started_at = now_rfc3339ish();
211        let config_json = serde_json::to_string(cfg)?;
212        let conn = self.conn.lock().unwrap();
213        insert_run_row(
214            &conn,
215            &cfg.suite,
216            &started_at,
217            "running",
218            Some(config_json.as_str()),
219        )
220    }
221
222    pub fn finalize_run(&self, run_id: i64, status: &str) -> anyhow::Result<()> {
223        let conn = self.conn.lock().unwrap();
224        conn.execute(
225            "UPDATE runs SET status=?1 WHERE id=?2",
226            params![status, run_id],
227        )?;
228        Ok(())
229    }
230
231    pub fn insert_result_embedded(
232        &self,
233        run_id: i64,
234        row: &TestResultRow,
235        attempts: &[AttemptRow],
236        output: &LlmResponse,
237    ) -> anyhow::Result<()> {
238        let conn = self.conn.lock().unwrap();
239
240        // 1. Insert into results
241        conn.execute(
242            "INSERT INTO results(run_id, test_id, outcome, score, duration_ms, attempts_json, output_json, fingerprint, skip_reason)
243             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
244            params![
245                run_id,
246                row.test_id,
247                status_to_outcome(&row.status),
248                row.score,
249                row.duration_ms.map(|v| v as i64),
250                serde_json::to_string(attempts)?,
251                serde_json::to_string(output)?,
252                row.fingerprint,
253                row.skip_reason
254            ],
255        )?;
256
257        let result_id = conn.last_insert_rowid();
258
259        // 2. Insert individual attempts
260        let mut stmt = conn.prepare(
261            "INSERT INTO attempts(result_id, attempt_number, outcome, score, duration_ms, output_json, error_message)
262             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)"
263        )?;
264
265        for attempt in attempts {
266            stmt.execute(params![
267                result_id,
268                attempt.attempt_no as i64,
269                status_to_outcome(&attempt.status),
270                0.0, // Score not tracked per attempt yet
271                attempt.duration_ms.map(|v| v as i64),
272                serde_json::to_string(&attempt.details)?,
273                Option::<String>::None
274            ])?;
275        }
276
277        Ok(())
278    }
279
280    // ... existing ...
281
282    // quarantine
283    pub fn quarantine_get_reason(
284        &self,
285        suite: &str,
286        test_id: &str,
287    ) -> anyhow::Result<Option<String>> {
288        let conn = self.conn.lock().unwrap();
289        let mut stmt =
290            conn.prepare("SELECT reason FROM quarantine WHERE suite=?1 AND test_id=?2")?;
291        let mut rows = stmt.query(params![suite, test_id])?;
292        if let Some(row) = rows.next()? {
293            Ok(Some(row.get::<_, Option<String>>(0)?.unwrap_or_default()))
294        } else {
295            Ok(None)
296        }
297    }
298
299    pub fn quarantine_add(&self, suite: &str, test_id: &str, reason: &str) -> anyhow::Result<()> {
300        let conn = self.conn.lock().unwrap();
301        conn.execute(
302            "INSERT INTO quarantine(suite, test_id, reason, added_at)
303             VALUES (?1, ?2, ?3, ?4)
304             ON CONFLICT(suite, test_id) DO UPDATE SET reason=excluded.reason, added_at=excluded.added_at",
305            params![suite, test_id, reason, now_rfc3339ish()],
306        )?;
307        Ok(())
308    }
309
310    pub fn quarantine_remove(&self, suite: &str, test_id: &str) -> anyhow::Result<()> {
311        let conn = self.conn.lock().unwrap();
312        conn.execute(
313            "DELETE FROM quarantine WHERE suite=?1 AND test_id=?2",
314            params![suite, test_id],
315        )?;
316        Ok(())
317    }
318
319    // cache
320    pub fn cache_get(&self, key: &str) -> anyhow::Result<Option<LlmResponse>> {
321        let conn = self.conn.lock().unwrap();
322        let mut stmt = conn.prepare("SELECT response_json FROM cache WHERE key=?1")?;
323        let mut rows = stmt.query(params![key])?;
324        if let Some(row) = rows.next()? {
325            let s: String = row.get(0)?;
326            let mut resp: LlmResponse = serde_json::from_str(&s)?;
327            resp.cached = true;
328            Ok(Some(resp))
329        } else {
330            Ok(None)
331        }
332    }
333
334    pub fn cache_put(&self, key: &str, resp: &LlmResponse) -> anyhow::Result<()> {
335        let conn = self.conn.lock().unwrap();
336        let created_at = now_rfc3339ish();
337        let mut to_store = resp.clone();
338        to_store.cached = false;
339        conn.execute(
340            "INSERT INTO cache(key, response_json, created_at) VALUES (?1, ?2, ?3)
341             ON CONFLICT(key) DO UPDATE SET response_json=excluded.response_json, created_at=excluded.created_at",
342            params![key, serde_json::to_string(&to_store)?, created_at],
343        )?;
344        Ok(())
345    }
346
347    // embeddings
348    pub fn get_embedding(&self, key: &str) -> anyhow::Result<Option<(String, Vec<f32>)>> {
349        let conn = self.conn.lock().unwrap();
350        let mut stmt = conn.prepare("SELECT model, vec FROM embeddings WHERE key = ?1 LIMIT 1")?;
351        let mut rows = stmt.query(params![key])?;
352
353        if let Some(row) = rows.next()? {
354            let model: String = row.get(0)?;
355            let blob: Vec<u8> = row.get(1)?;
356            let vec = crate::embeddings::util::decode_vec_f32(&blob)?;
357            Ok(Some((model, vec)))
358        } else {
359            Ok(None)
360        }
361    }
362
363    pub fn put_embedding(&self, key: &str, model: &str, vec: &[f32]) -> anyhow::Result<()> {
364        let conn = self.conn.lock().unwrap();
365        let blob = crate::embeddings::util::encode_vec_f32(vec);
366        let dims = vec.len() as i64;
367        let created_at = now_rfc3339ish();
368
369        conn.execute(
370            "INSERT OR REPLACE INTO embeddings (key, model, dims, vec, created_at)
371             VALUES (?1, ?2, ?3, ?4, ?5)",
372            params![key, model, dims, blob, created_at],
373        )?;
374        Ok(())
375    }
376    pub fn stats_best_effort(&self) -> anyhow::Result<StoreStats> {
377        let conn = self.conn.lock().unwrap();
378
379        let runs: Option<u64> = conn
380            .query_row("SELECT COUNT(*) FROM runs", [], |r| {
381                r.get::<_, i64>(0).map(|x| x as u64)
382            })
383            .ok();
384        let results: Option<u64> = conn
385            .query_row("SELECT COUNT(*) FROM results", [], |r| {
386                r.get::<_, i64>(0).map(|x| x as u64)
387            })
388            .ok();
389
390        let last: Option<(i64, String)> = conn
391            .query_row(
392                "SELECT id, started_at FROM runs ORDER BY id DESC LIMIT 1",
393                [],
394                |r| Ok((r.get(0)?, r.get(1)?)),
395            )
396            .ok();
397
398        let (last_id, last_started) = if let Some((id, s)) = last {
399            (Some(id), Some(s))
400        } else {
401            (None, None)
402        };
403
404        let v_str: Option<String> = conn
405            .query_row("PRAGMA user_version", [], |r| r.get(0))
406            .ok()
407            .map(|v: i64| v.to_string());
408
409        Ok(StoreStats {
410            runs,
411            results,
412            last_run_id: last_id,
413            last_run_at: last_started,
414            version: v_str,
415        })
416    }
417
418    // --- Assertions Support ---
419
420    pub fn get_episode_graph(
421        &self,
422        run_id: i64,
423        test_id: &str,
424    ) -> anyhow::Result<crate::agent_assertions::EpisodeGraph> {
425        let conn = self.conn.lock().unwrap();
426
427        // 1. Find Episode
428        let mut stmt = conn.prepare("SELECT id FROM episodes WHERE run_id = ? AND test_id = ?")?;
429        let mut rows = stmt.query(params![run_id, test_id])?;
430
431        let mut episode_ids = Vec::new();
432        while let Some(row) = rows.next()? {
433            episode_ids.push(row.get::<_, String>(0)?);
434        }
435
436        if episode_ids.is_empty() {
437            anyhow::bail!(
438                "E_TRACE_EPISODE_MISSING: No episode found for run_id={} test_id={}",
439                run_id,
440                test_id
441            );
442        }
443        if episode_ids.len() > 1 {
444            anyhow::bail!(
445                "E_TRACE_EPISODE_AMBIGUOUS: Multiple episodes ({}) found for run_id={} test_id={}",
446                episode_ids.len(),
447                run_id,
448                test_id
449            );
450        }
451        let episode_id = episode_ids[0].clone();
452
453        load_episode_graph_for_episode_id(&conn, &episode_id)
454    }
455}
456
457fn status_to_outcome(s: &TestStatus) -> &'static str {
458    store_internal::results::status_to_outcome_impl(s)
459}
460
461fn migrate_v030(conn: &Connection) -> anyhow::Result<()> {
462    store_internal::schema::migrate_v030_impl(conn)
463}
464
465fn row_to_test_result(row: &rusqlite::Row<'_>) -> rusqlite::Result<TestResultRow> {
466    store_internal::results::row_to_test_result_impl(row)
467}
468
469fn insert_run_row(
470    conn: &Connection,
471    suite: &str,
472    started_at: &str,
473    status: &str,
474    config_json: Option<&str>,
475) -> anyhow::Result<i64> {
476    store_internal::results::insert_run_row_impl(conn, suite, started_at, status, config_json)
477}
478
479fn load_episode_graph_for_episode_id(
480    conn: &Connection,
481    episode_id: &str,
482) -> anyhow::Result<crate::agent_assertions::EpisodeGraph> {
483    store_internal::episodes::load_episode_graph_for_episode_id_impl(conn, episode_id)
484}
assay_core/storage/store.rs

assay_core/storage/
store.rs
