1use super::*;
2
3impl Store {
4 pub fn insert_event(
7 &self,
8 event: &TraceEvent,
9 run_id: Option<i64>,
10 test_id: Option<&str>,
11 ) -> anyhow::Result<()> {
12 let mut conn = self.conn.lock().unwrap();
13 let tx = conn.transaction()?;
14 match event {
15 TraceEvent::EpisodeStart(e) => Self::insert_episode(&tx, e, run_id, test_id)?,
16 TraceEvent::Step(e) => Self::insert_step(&tx, e)?,
17 TraceEvent::ToolCall(e) => Self::insert_tool_call(&tx, e)?,
18 TraceEvent::EpisodeEnd(e) => Self::update_episode_end(&tx, e)?,
19 }
20 tx.commit()?;
21 Ok(())
22 }
23
24 pub fn insert_batch(
25 &self,
26 events: &[TraceEvent],
27 run_id: Option<i64>,
28 test_id: Option<&str>,
29 ) -> anyhow::Result<()> {
30 let mut conn = self.conn.lock().unwrap();
31 let tx = conn.transaction()?;
32 for event in events {
33 match event {
34 TraceEvent::EpisodeStart(e) => Self::insert_episode(&tx, e, run_id, test_id)?,
35 TraceEvent::Step(e) => Self::insert_step(&tx, e)?,
36 TraceEvent::ToolCall(e) => Self::insert_tool_call(&tx, e)?,
37 TraceEvent::EpisodeEnd(e) => Self::update_episode_end(&tx, e)?,
38 }
39 }
40 tx.commit()?;
41 Ok(())
42 }
43
44 pub fn count_rows(&self, table: &str) -> anyhow::Result<i64> {
45 let conn = self.conn.lock().unwrap();
46 if !["episodes", "steps", "tool_calls", "runs", "results"].contains(&table) {
47 anyhow::bail!("Invalid table name for count_rows: {}", table);
48 }
49 let sql = format!("SELECT COUNT(*) FROM {}", table);
50 let n: i64 = conn.query_row(&sql, [], |r| r.get(0))?;
51 Ok(n)
52 }
53
54 pub fn get_latest_episode_graph_by_test_id(
55 &self,
56 test_id: &str,
57 ) -> anyhow::Result<crate::agent_assertions::EpisodeGraph> {
58 let conn = self.conn.lock().unwrap();
59 let mut stmt = conn.prepare(
60 "SELECT id FROM episodes
61 WHERE test_id = ?1
62 ORDER BY timestamp DESC
63 LIMIT 1",
64 )?;
65
66 let episode_id: String = stmt.query_row(params![test_id], |row| row.get(0)).map_err(
67 |e| {
68 anyhow::anyhow!(
69 "E_TRACE_EPISODE_MISSING: No episode found for test_id={} (fallback check) : {}",
70 test_id,
71 e
72 )
73 },
74 )?;
75
76 load_episode_graph_for_episode_id(&conn, &episode_id)
77 }
78
79 fn insert_episode(
80 tx: &rusqlite::Transaction<'_>,
81 e: &EpisodeStart,
82 run_id: Option<i64>,
83 test_id: Option<&str>,
84 ) -> anyhow::Result<()> {
85 let prompt_val = e.input.get("prompt").unwrap_or(&serde_json::Value::Null);
86 let prompt_str = if let Some(s) = prompt_val.as_str() {
87 s.to_string()
88 } else {
89 serde_json::to_string(prompt_val).unwrap_or_default()
90 };
91 let meta = serde_json::to_string(&e.meta).unwrap_or_default();
92
93 let meta_test_id = e.meta.get("test_id").and_then(|v| v.as_str());
94 let effective_test_id = test_id.or(meta_test_id).or(Some(&e.episode_id));
95
96 tx.execute(
97 "INSERT INTO episodes (id, run_id, test_id, timestamp, prompt, meta_json) VALUES (?, ?, ?, ?, ?, ?)
98 ON CONFLICT(id) DO UPDATE SET
99 run_id=COALESCE(excluded.run_id, episodes.run_id),
100 test_id=COALESCE(excluded.test_id, episodes.test_id),
101 timestamp=excluded.timestamp,
102 prompt=excluded.prompt,
103 meta_json=excluded.meta_json",
104 (
105 &e.episode_id,
106 run_id,
107 effective_test_id,
108 e.timestamp,
109 prompt_str,
110 meta,
111 ),
112 )
113 .context("insert episode")?;
114 Ok(())
115 }
116
117 fn insert_step(tx: &rusqlite::Transaction<'_>, e: &StepEntry) -> anyhow::Result<()> {
118 let meta = serde_json::to_string(&e.meta).unwrap_or_default();
119 let trunc = serde_json::to_string(&e.truncations).unwrap_or_default();
120
121 tx.execute(
122 "INSERT INTO steps (id, episode_id, idx, kind, name, content, content_sha256, truncations_json, meta_json)
123 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
124 ON CONFLICT(id) DO UPDATE SET content=excluded.content, meta_json=excluded.meta_json",
125 (
126 &e.step_id,
127 &e.episode_id,
128 e.idx,
129 &e.kind,
130 e.name.as_deref(),
131 e.content.as_deref(),
132 e.content_sha256.as_deref(),
133 trunc,
134 meta,
135 ),
136 )
137 .context("insert step")?;
138 Ok(())
139 }
140
141 fn insert_tool_call(tx: &rusqlite::Transaction<'_>, e: &ToolCallEntry) -> anyhow::Result<()> {
142 let args = serde_json::to_string(&e.args).unwrap_or_default();
143 let result = e
144 .result
145 .as_ref()
146 .map(|r| serde_json::to_string(r).unwrap_or_default());
147 let trunc = serde_json::to_string(&e.truncations).unwrap_or_default();
148
149 let call_idx = e.call_index.unwrap_or(0);
150
151 tx.execute(
152 "INSERT INTO tool_calls (step_id, episode_id, tool_name, call_index, args, args_sha256, result, result_sha256, error, truncations_json)
153 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
154 ON CONFLICT(step_id, call_index) DO NOTHING",
155 (
156 &e.step_id,
157 &e.episode_id,
158 &e.tool_name,
159 call_idx,
160 args,
161 e.args_sha256.as_deref(),
162 result,
163 e.result_sha256.as_deref(),
164 e.error.as_deref(),
165 trunc,
166 ),
167 )
168 .context("insert tool call")?;
169 Ok(())
170 }
171
172 fn update_episode_end(tx: &rusqlite::Transaction<'_>, e: &EpisodeEnd) -> anyhow::Result<()> {
173 tx.execute(
174 "UPDATE episodes SET outcome = ? WHERE id = ?",
175 (e.outcome.as_deref(), &e.episode_id),
176 )
177 .context("update episode outcome")?;
178 Ok(())
179 }
180}