bctx_forge/tracker/
store.rs

1use anyhow::{Context, Result};
2use chrono::Utc;
3use rusqlite::{params, Connection};
4use std::path::PathBuf;
5
6#[derive(Debug, Clone)]
7pub struct ExecutionRecord {
8    pub program: String,
9    pub args: String,
10    pub tokens_before: usize,
11    pub tokens_after: usize,
12    pub exit_code: i32,
13    pub duration_ms: u64,
14    pub timestamp: String,
15}
16
17pub struct ExecutionStore {
18    conn: Connection,
19}
20
21impl ExecutionStore {
22    pub fn open(path: &PathBuf) -> Result<Self> {
23        if let Some(parent) = path.parent() {
24            std::fs::create_dir_all(parent)?;
25        }
26        let conn = Connection::open(path).context("opening execution store")?;
27        conn.execute_batch(
28            "CREATE TABLE IF NOT EXISTS executions (
29                id INTEGER PRIMARY KEY AUTOINCREMENT,
30                program TEXT NOT NULL,
31                args TEXT NOT NULL,
32                tokens_before INTEGER NOT NULL,
33                tokens_after INTEGER NOT NULL,
34                exit_code INTEGER NOT NULL,
35                duration_ms INTEGER NOT NULL,
36                timestamp TEXT NOT NULL
37            );
38            CREATE INDEX IF NOT EXISTS idx_program ON executions(program);
39            CREATE INDEX IF NOT EXISTS idx_timestamp ON executions(timestamp);",
40        )?;
41        Ok(Self { conn })
42    }
43
44    pub fn in_memory() -> Result<Self> {
45        let conn = Connection::open_in_memory()?;
46        conn.execute_batch(
47            "CREATE TABLE IF NOT EXISTS executions (
48                id INTEGER PRIMARY KEY AUTOINCREMENT,
49                program TEXT NOT NULL,
50                args TEXT NOT NULL,
51                tokens_before INTEGER NOT NULL,
52                tokens_after INTEGER NOT NULL,
53                exit_code INTEGER NOT NULL,
54                duration_ms INTEGER NOT NULL,
55                timestamp TEXT NOT NULL
56            );",
57        )?;
58        Ok(Self { conn })
59    }
60
61    pub fn record(&self, rec: &ExecutionRecord) -> Result<()> {
62        self.conn.execute(
63            "INSERT INTO executions (program, args, tokens_before, tokens_after, exit_code, duration_ms, timestamp)
64             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
65            params![
66                rec.program, rec.args, rec.tokens_before, rec.tokens_after,
67                rec.exit_code, rec.duration_ms, rec.timestamp
68            ],
69        )?;
70        Ok(())
71    }
72
73    pub fn total_saved(&self) -> Result<usize> {
74        let val: i64 = self.conn.query_row(
75            "SELECT COALESCE(SUM(tokens_before - tokens_after), 0) FROM executions",
76            [],
77            |row| row.get(0),
78        )?;
79        Ok(val as usize)
80    }
81
82    pub fn summary(&self) -> Result<Vec<(String, usize, usize, f64)>> {
83        let mut stmt = self.conn.prepare(
84            "SELECT program,
85                    SUM(tokens_before) as tb,
86                    SUM(tokens_after) as ta,
87                    AVG(CAST(tokens_before - tokens_after AS REAL) / NULLIF(tokens_before, 0)) * 100.0 as pct
88             FROM executions
89             GROUP BY program
90             ORDER BY tb DESC
91             LIMIT 20",
92        )?;
93        let rows = stmt.query_map([], |row| {
94            Ok((
95                row.get::<_, String>(0)?,
96                row.get::<_, i64>(1)? as usize,
97                row.get::<_, i64>(2)? as usize,
98                row.get::<_, f64>(3).unwrap_or(0.0),
99            ))
100        })?;
101        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
102    }
103
104    pub fn recent(&self, limit: usize) -> Result<Vec<ExecutionRecord>> {
105        let mut stmt = self.conn.prepare(
106            "SELECT program, args, tokens_before, tokens_after, exit_code, duration_ms, timestamp
107             FROM executions ORDER BY id DESC LIMIT ?1",
108        )?;
109        let rows = stmt.query_map(params![limit as i64], |row| {
110            Ok(ExecutionRecord {
111                program: row.get(0)?,
112                args: row.get(1)?,
113                tokens_before: row.get::<_, i64>(2)? as usize,
114                tokens_after: row.get::<_, i64>(3)? as usize,
115                exit_code: row.get(4)?,
116                duration_ms: row.get::<_, i64>(5)? as u64,
117                timestamp: row.get(6)?,
118            })
119        })?;
120        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
121    }
122
123    pub fn new_record(
124        program: &str,
125        args: &[String],
126        tokens_before: usize,
127        tokens_after: usize,
128        exit_code: i32,
129        duration_ms: u64,
130    ) -> ExecutionRecord {
131        ExecutionRecord {
132            program: program.to_string(),
133            args: args.join(" "),
134            tokens_before,
135            tokens_after,
136            exit_code,
137            duration_ms,
138            timestamp: Utc::now().to_rfc3339(),
139        }
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn record_then_read_back_aggregates() {
149        // Guards the local-tracking pipeline that `bctx gain`/dashboard read from:
150        // a recorded execution must be retrievable and summed correctly.
151        let store = ExecutionStore::in_memory().expect("open in-memory store");
152        store
153            .record(&ExecutionStore::new_record(
154                "git",
155                &["log".into()],
156                100,
157                30,
158                0,
159                5,
160            ))
161            .expect("record git");
162        store
163            .record(&ExecutionStore::new_record(
164                "cargo",
165                &["build".into()],
166                50,
167                20,
168                0,
169                9,
170            ))
171            .expect("record cargo");
172
173        assert_eq!(store.recent(10).expect("recent").len(), 2);
174        // (100-30) + (50-20) = 100 tokens saved
175        assert_eq!(store.total_saved().expect("total_saved"), 100);
176        assert_eq!(store.summary().expect("summary").len(), 2);
177    }
178}
bctx_forge/tracker/store.rs

bctx_forge/tracker/
store.rs
