Skip to main content

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}