1use std::path::PathBuf;
21use std::sync::Arc;
22
23use anyhow::{Context, Result};
24use rusqlite::{params, Connection};
25use serde::{Deserialize, Serialize};
26use tokio::sync::Mutex;
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct Span {
31 pub id: String,
33 pub conversation_id: String,
35 pub kind: String,
37 pub name: String,
39 pub args_hash: Option<String>,
42 pub started_at: i64,
44 pub ended_at: Option<i64>,
46 pub error: Option<String>,
48 pub session_id: Option<String>,
51 pub seq: i64,
53}
54
55#[derive(Clone)]
57pub struct TraceStore {
58 conn: Arc<Mutex<Connection>>,
59}
60
61fn now_millis() -> i64 {
62 chrono::Utc::now().timestamp_millis()
63}
64
65pub fn hash_args(value: &serde_json::Value) -> String {
67 use std::fmt::Write;
68 let bytes = serde_json::to_vec(value).unwrap_or_default();
71 let digest = sha2_digest(&bytes);
72 let mut out = String::with_capacity(64);
73 for b in digest {
74 let _ = write!(out, "{b:02x}");
75 }
76 out
77}
78
79fn sha2_digest(data: &[u8]) -> [u8; 32] {
83 use sha2::Digest;
84 let mut h = sha2::Sha256::new();
85 h.update(data);
86 h.finalize().into()
87}
88
89impl TraceStore {
90 pub fn open(path: PathBuf) -> Result<Self> {
92 if let Some(parent) = path.parent() {
93 std::fs::create_dir_all(parent)
94 .with_context(|| format!("creating trace db dir {}", parent.display()))?;
95 }
96 let conn = Connection::open(&path)
97 .with_context(|| format!("opening trace db {}", path.display()))?;
98 Self::init_schema(&conn)?;
99 Ok(Self {
100 conn: Arc::new(Mutex::new(conn)),
101 })
102 }
103
104 pub fn open_in_memory() -> Result<Self> {
106 let conn = Connection::open_in_memory().context("opening in-memory trace db")?;
107 Self::init_schema(&conn)?;
108 Ok(Self {
109 conn: Arc::new(Mutex::new(conn)),
110 })
111 }
112
113 fn init_schema(conn: &Connection) -> Result<()> {
114 conn.execute_batch(
115 "PRAGMA journal_mode = WAL;
116 CREATE TABLE IF NOT EXISTS spans (
117 seq INTEGER PRIMARY KEY AUTOINCREMENT,
118 id TEXT NOT NULL UNIQUE,
119 conversation_id TEXT NOT NULL,
120 kind TEXT NOT NULL,
121 name TEXT NOT NULL,
122 args_hash TEXT,
123 started_at INTEGER NOT NULL,
124 ended_at INTEGER,
125 error TEXT,
126 session_id TEXT
127 );
128 CREATE INDEX IF NOT EXISTS idx_spans_conversation
129 ON spans(conversation_id, seq);",
130 )
131 .context("initializing trace schema")?;
132
133 let existing: std::collections::HashSet<String> = {
135 let mut stmt = conn.prepare("PRAGMA table_info(spans)")?;
136 let names = stmt.query_map([], |row| row.get::<_, String>(1))?;
137 names.filter_map(|r| r.ok()).collect()
138 };
139 if !existing.contains("session_id") {
140 conn.execute_batch("ALTER TABLE spans ADD COLUMN session_id TEXT")
141 .context("adding session_id column")?;
142 }
143
144 Ok(())
145 }
146
147 pub async fn open_span(
152 &self,
153 conversation_id: &str,
154 kind: &str,
155 name: &str,
156 args_hash: Option<&str>,
157 session_id: Option<&str>,
158 ) -> Result<String> {
159 let span_id = uuid::Uuid::new_v4().to_string();
160 let now = now_millis();
161 let conn = self.conn.lock().await;
162 conn.execute(
163 "INSERT INTO spans (id, conversation_id, kind, name, args_hash, started_at, session_id)
164 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
165 params![
166 span_id,
167 conversation_id,
168 kind,
169 name,
170 args_hash,
171 now,
172 session_id
173 ],
174 )
175 .context("inserting span")?;
176 Ok(span_id)
177 }
178
179 pub async fn close_span(&self, span_id: &str, error: Option<&str>) -> Result<()> {
181 let now = now_millis();
182 let conn = self.conn.lock().await;
183 conn.execute(
184 "UPDATE spans SET ended_at = ?1, error = ?2 WHERE id = ?3",
185 params![now, error, span_id],
186 )
187 .context("closing span")?;
188 Ok(())
189 }
190
191 pub async fn get_spans(&self, conversation_id: &str) -> Result<Vec<Span>> {
193 let conn = self.conn.lock().await;
194 let mut stmt = conn.prepare(
195 "SELECT seq, id, conversation_id, kind, name, args_hash,
196 started_at, ended_at, error, session_id
197 FROM spans
198 WHERE conversation_id = ?1
199 ORDER BY seq ASC",
200 )?;
201 let rows = stmt.query_map(params![conversation_id], |row| {
202 Ok(Span {
203 seq: row.get(0)?,
204 id: row.get(1)?,
205 conversation_id: row.get(2)?,
206 kind: row.get(3)?,
207 name: row.get(4)?,
208 args_hash: row.get(5)?,
209 started_at: row.get(6)?,
210 ended_at: row.get(7)?,
211 error: row.get(8)?,
212 session_id: row.get(9)?,
213 })
214 })?;
215 rows.collect::<std::result::Result<Vec<_>, _>>()
216 .context("reading spans")
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[tokio::test]
225 async fn write_and_read_back_tool_call_span() {
226 let store = TraceStore::open_in_memory().unwrap();
227 let conv_id = "test-conv-1";
228
229 let input = serde_json::json!({ "path": "/tmp/foo.txt" });
231 let ah = hash_args(&input);
232 let span_id = store
233 .open_span(conv_id, "tool-call", "read_file", Some(&ah), None)
234 .await
235 .unwrap();
236
237 store.close_span(&span_id, None).await.unwrap();
239
240 let spans = store.get_spans(conv_id).await.unwrap();
242 assert_eq!(spans.len(), 1);
243 let s = &spans[0];
244 assert_eq!(s.conversation_id, conv_id);
245 assert_eq!(s.kind, "tool-call");
246 assert_eq!(s.name, "read_file");
247 assert!(s.args_hash.is_some());
248 assert!(s.ended_at.is_some());
249 assert!(s.error.is_none());
250 }
251
252 #[tokio::test]
253 async fn error_span_records_message() {
254 let store = TraceStore::open_in_memory().unwrap();
255 let span_id = store
256 .open_span("conv-2", "tool-call", "bash", None, None)
257 .await
258 .unwrap();
259 store
260 .close_span(&span_id, Some("permission denied"))
261 .await
262 .unwrap();
263 let spans = store.get_spans("conv-2").await.unwrap();
264 assert_eq!(spans[0].error.as_deref(), Some("permission denied"));
265 }
266
267 #[tokio::test]
268 async fn multiple_spans_ordered_by_seq() {
269 let store = TraceStore::open_in_memory().unwrap();
270 let conv = "conv-order";
271 for name in ["alpha", "beta", "gamma"] {
272 let id = store
273 .open_span(conv, "tool-call", name, None, None)
274 .await
275 .unwrap();
276 store.close_span(&id, None).await.unwrap();
277 }
278 let spans = store.get_spans(conv).await.unwrap();
279 assert_eq!(spans.len(), 3);
280 assert!(spans[0].seq < spans[1].seq);
281 assert!(spans[1].seq < spans[2].seq);
282 assert_eq!(spans[0].name, "alpha");
283 assert_eq!(spans[2].name, "gamma");
284 }
285}