agentdb/
traces.rs

1use crate::error::{AgentDbError, Result};
2use crate::schema::now_ms;
3use rusqlite::params;
4use rusqlite::Connection;
5use serde_json::Value;
6use std::sync::{Arc, Mutex};
7use uuid::Uuid;
8
9/// A single reasoning trace entry.
10#[derive(Debug, Clone)]
11pub struct Trace {
12    /// Unique identifier for this trace entry.
13    pub id: String,
14    /// Optional session identifier that groups related traces.
15    pub session_id: Option<String>,
16    /// ID of the parent trace (for tree structures), if any.
17    pub parent_id: Option<String>,
18    /// Semantic type of the trace (e.g. `"thought"`, `"tool_call"`, `"observation"`).
19    pub trace_type: String,
20    /// Text content of the trace entry.
21    pub content: String,
22    /// Arbitrary JSON payload attached to this trace.
23    pub metadata: Option<Value>,
24    /// Unix-millisecond timestamp when the trace was recorded.
25    pub created_at: i64,
26}
27
28/// Stores and retrieves reasoning traces.
29pub struct TraceStore {
30    conn: Arc<Mutex<Connection>>,
31}
32
33impl TraceStore {
34    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
35        Self { conn }
36    }
37
38    /// Record a new trace entry. Returns the generated trace ID.
39    ///
40    /// - `session_id`: optional grouping key (e.g. a request or agent run ID).
41    /// - `parent_id`: optional ID of a parent trace for tree structures.
42    /// - `trace_type`: semantic label (e.g. `"thought"`, `"tool_call"`).
43    /// - `content`: the text body of the trace.
44    /// - `metadata`: optional JSON payload.
45    pub fn add_trace(
46        &self,
47        session_id: Option<&str>,
48        parent_id: Option<&str>,
49        trace_type: &str,
50        content: &str,
51        metadata: Option<Value>,
52    ) -> Result<String> {
53        let trace_id = Uuid::new_v4().to_string();
54        let meta_str = metadata.as_ref().map(|m| m.to_string());
55        let now = now_ms();
56        let conn = self.conn.lock().unwrap();
57        conn.execute(
58            "INSERT INTO _adb_traces
59                 (id, session_id, parent_id, trace_type, content, metadata, created_at)
60             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
61            params![trace_id, session_id, parent_id, trace_type, content, meta_str, now],
62        )?;
63        Ok(trace_id)
64    }
65
66    /// Return all traces for a given session in chronological order.
67    pub fn get_traces(&self, session_id: &str) -> Result<Vec<Trace>> {
68        let conn = self.conn.lock().unwrap();
69        let mut stmt = conn.prepare(
70            "SELECT id, session_id, parent_id, trace_type, content, metadata, created_at
71             FROM _adb_traces
72             WHERE session_id = ?1
73             ORDER BY created_at ASC",
74        )?;
75        let rows = stmt.query_map(params![session_id], parse_trace)?;
76        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
77    }
78
79    /// Return a subtree of traces rooted at `root_id`.
80    ///
81    /// Uses a recursive CTE to follow `parent_id` links. The root trace itself
82    /// is included. Results are ordered by `created_at` ascending.
83    pub fn get_trace_tree(&self, root_id: &str) -> Result<Vec<Trace>> {
84        let conn = self.conn.lock().unwrap();
85        let mut stmt = conn.prepare(
86            "WITH RECURSIVE tree(id) AS (
87                 SELECT id FROM _adb_traces WHERE id = ?1
88                 UNION ALL
89                 SELECT t.id
90                 FROM _adb_traces t
91                 JOIN tree ON t.parent_id = tree.id
92             )
93             SELECT t.id, t.session_id, t.parent_id, t.trace_type,
94                    t.content, t.metadata, t.created_at
95             FROM _adb_traces t
96             JOIN tree ON t.id = tree.id
97             ORDER BY t.created_at ASC",
98        )?;
99        let rows = stmt.query_map(params![root_id], parse_trace)?;
100        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
101    }
102}
103
104// ── Row parsers ──────────────────────────────────────────────────────────────
105
106fn parse_trace(row: &rusqlite::Row) -> rusqlite::Result<Trace> {
107    let meta_str: Option<String> = row.get(5)?;
108    Ok(Trace {
109        id: row.get(0)?,
110        session_id: row.get(1)?,
111        parent_id: row.get(2)?,
112        trace_type: row.get(3)?,
113        content: row.get(4)?,
114        metadata: meta_str
115            .as_deref()
116            .and_then(|s| serde_json::from_str(s).ok()),
117        created_at: row.get(6)?,
118    })
119}
agentdb/traces.rs

agentdb/
traces.rs
