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, serde::Serialize, serde::Deserialize)]
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 traces for a given session in chronological order.
67    ///
68    /// Use `limit` and `offset` for pagination. Pass `None` for both to
69    /// retrieve all traces.
70    pub fn get_traces(
71        &self,
72        session_id: &str,
73        limit: Option<usize>,
74        offset: Option<usize>,
75    ) -> Result<Vec<Trace>> {
76        let conn = self.conn.lock().unwrap();
77        let lim: i64 = limit.map(|n| n as i64).unwrap_or(i64::MAX);
78        let off: i64 = offset.map(|n| n as i64).unwrap_or(0);
79        let mut stmt = conn.prepare(
80            "SELECT id, session_id, parent_id, trace_type, content, metadata, created_at
81             FROM _adb_traces
82             WHERE session_id = ?1
83             ORDER BY created_at ASC
84             LIMIT ?2 OFFSET ?3",
85        )?;
86        let rows = stmt.query_map(params![session_id, lim, off], parse_trace)?;
87        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
88    }
89
90    /// Return a subtree of traces rooted at `root_id`.
91    ///
92    /// Uses a recursive CTE to follow `parent_id` links. The root trace itself
93    /// is included. Results are ordered by `created_at` ascending.
94    pub fn get_trace_tree(&self, root_id: &str) -> Result<Vec<Trace>> {
95        let conn = self.conn.lock().unwrap();
96        let mut stmt = conn.prepare(
97            "WITH RECURSIVE tree(id) AS (
98                 SELECT id FROM _adb_traces WHERE id = ?1
99                 UNION ALL
100                 SELECT t.id
101                 FROM _adb_traces t
102                 JOIN tree ON t.parent_id = tree.id
103             )
104             SELECT t.id, t.session_id, t.parent_id, t.trace_type,
105                    t.content, t.metadata, t.created_at
106             FROM _adb_traces t
107             JOIN tree ON t.id = tree.id
108             ORDER BY t.created_at ASC",
109        )?;
110        let rows = stmt.query_map(params![root_id], parse_trace)?;
111        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
112    }
113}
114
115// ── Row parsers ──────────────────────────────────────────────────────────────
116
117fn parse_trace(row: &rusqlite::Row) -> rusqlite::Result<Trace> {
118    let meta_str: Option<String> = row.get(5)?;
119    Ok(Trace {
120        id: row.get(0)?,
121        session_id: row.get(1)?,
122        parent_id: row.get(2)?,
123        trace_type: row.get(3)?,
124        content: row.get(4)?,
125        metadata: meta_str
126            .as_deref()
127            .and_then(|s| serde_json::from_str(s).ok()),
128        created_at: row.get(6)?,
129    })
130}
agentdb/traces.rs

agentdb/
traces.rs
