forge_workspace/store/
token_usage.rs

1//! Per-file usage-summary cache on the redb `token_usage` table.
2//!
3//! Keyed by canonical session-file path; the value is that file's
4//! deduped `FileUsageSummary`. The cache is incremental: the scan
5//! reuses a file whose mtime and size still match the cached entry
6//! rather than re-parsing a transcript that can reach 100 MB.
7
8use anyhow::Context;
9use forge_agent::env::token_usage::FileUsageSummary;
10use redb::TableDefinition;
11
12use super::Db;
13
14const TOKEN_USAGE: TableDefinition<&str, &[u8]> = TableDefinition::new("token_usage");
15
16/// The cached summary for `path`, or `None` when absent (or the table
17/// doesn't exist yet). A decode failure is surfaced so a corrupt row is
18/// diagnosable rather than silently treated as a cache miss.
19pub fn load(db: &Db, path: &str) -> anyhow::Result<Option<FileUsageSummary>> {
20    let txn = db.database().begin_read()?;
21    let table = match txn.open_table(TOKEN_USAGE) {
22        Ok(t) => t,
23        Err(redb::TableError::TableDoesNotExist(_)) => return Ok(None),
24        Err(e) => return Err(e.into()),
25    };
26    match table.get(path)? {
27        Some(value) => Ok(Some(decode(value.value(), path)?)),
28        None => Ok(None),
29    }
30}
31
32/// Store `summary` for `path`, overwriting any prior entry.
33pub fn store(db: &Db, path: &str, summary: &FileUsageSummary) -> anyhow::Result<()> {
34    let value = serde_json::to_vec(summary).context("serialize usage summary")?;
35    let txn = db.database().begin_write()?;
36    {
37        let mut table = txn.open_table(TOKEN_USAGE)?;
38        table.insert(path, value.as_slice())?;
39    }
40    txn.commit()?;
41    Ok(())
42}
43
44/// Decode a stored blob, tagging a failure with the owning path.
45fn decode(bytes: &[u8], path: &str) -> anyhow::Result<FileUsageSummary> {
46    serde_json::from_slice(bytes).with_context(|| format!("decode usage summary for {path}"))
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use forge_agent::env::token_usage::TokenCounts;
53    use std::collections::BTreeMap;
54    use std::time::{Duration, SystemTime};
55    use tempfile::tempdir;
56
57    fn summary(project: &str, output: u64) -> FileUsageSummary {
58        let mut by_model_day = BTreeMap::new();
59        let mut days = BTreeMap::new();
60        days.insert("2026-07-08".to_owned(), TokenCounts { output, ..TokenCounts::default() });
61        by_model_day.insert("m".to_owned(), days);
62        FileUsageSummary {
63            mtime: SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000),
64            size: 42,
65            folded_project: project.to_owned(),
66            project_resolved: true,
67            by_model_day,
68        }
69    }
70
71    fn open_db() -> (tempfile::TempDir, Db) {
72        let dir = tempdir().expect("tempdir");
73        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
74        (dir, db)
75    }
76
77    #[test]
78    fn load_is_none_on_miss() {
79        let (_dir, db) = open_db();
80        assert!(load(&db, "/nope.jsonl").expect("load").is_none());
81    }
82
83    #[test]
84    fn store_then_load_round_trips() {
85        let (_dir, db) = open_db();
86        let entry = summary("forge", 7);
87        store(&db, "/a.jsonl", &entry).expect("store");
88        assert_eq!(load(&db, "/a.jsonl").expect("load"), Some(entry));
89    }
90
91    #[test]
92    fn a_row_predating_project_resolved_decodes_as_unresolved() {
93        let (_dir, db) = open_db();
94        // A row exactly as the previous version wrote it: the same blob
95        // minus the field that did not exist yet.
96        let mut old = serde_json::to_value(summary("auto", 7)).expect("serialize");
97        old.as_object_mut()
98            .expect("a summary serializes to an object")
99            .remove("project_resolved")
100            .expect("project_resolved is the field being dropped");
101        let bytes = serde_json::to_vec(&old).expect("serialize old shape");
102        let txn = db.database().begin_write().expect("begin write");
103        {
104            let mut table = txn.open_table(TOKEN_USAGE).expect("open table");
105            table.insert("/a.jsonl", bytes.as_slice()).expect("insert");
106        }
107        txn.commit().expect("commit");
108
109        let loaded = load(&db, "/a.jsonl").expect("load").expect("an old row still decodes");
110        assert_eq!(loaded.by_model_day["m"]["2026-07-08"].output, 7, "its counts survive");
111        assert!(
112            !loaded.project_resolved,
113            "an old row reads as unresolved, so its label is re-derived rather than frozen",
114        );
115    }
116
117    #[test]
118    fn store_overwrites_prior_entry() {
119        let (_dir, db) = open_db();
120        store(&db, "/a.jsonl", &summary("forge", 1)).expect("first");
121        store(&db, "/a.jsonl", &summary("forge", 9)).expect("second");
122        let loaded = load(&db, "/a.jsonl").expect("load").expect("present");
123        assert_eq!(loaded.by_model_day["m"]["2026-07-08"].output, 9);
124    }
125}
forge_workspace/store/token_usage.rs

forge_workspace/store/
token_usage.rs
