forge_workspace/store/
state.rs

1//! Machine-local forge state on the redb `account_usage` table: the
2//! per-account usage cache, one row per account. Values are serde-json;
3//! no field schema on disk.
4
5use std::collections::BTreeMap;
6
7use anyhow::Context;
8use redb::{ReadableTable, TableDefinition};
9
10use super::Db;
11use crate::account_cache::CachedAccountUsage;
12
13const ACCOUNT_USAGE: TableDefinition<&str, &[u8]> = TableDefinition::new("account_usage");
14
15/// Every cached per-account usage snapshot, keyed by account display
16/// name. A record that fails to decode is skipped with a warn so one
17/// corrupt blob can't wipe the rest of the cache.
18pub fn account_usage(db: &Db) -> anyhow::Result<BTreeMap<String, CachedAccountUsage>> {
19    let txn = db.database().begin_read()?;
20    let table = match txn.open_table(ACCOUNT_USAGE) {
21        Ok(t) => t,
22        // A fresh database has no table until the first write; an absent
23        // table is an empty cache, not an error.
24        Err(redb::TableError::TableDoesNotExist(_)) => return Ok(BTreeMap::new()),
25        Err(e) => return Err(e.into()),
26    };
27    let mut out = BTreeMap::new();
28    for entry in table.iter()? {
29        let (name, value) = entry?;
30        match serde_json::from_slice::<CachedAccountUsage>(value.value()) {
31            Ok(usage) => {
32                out.insert(name.value().to_owned(), usage);
33            }
34            Err(err) => tracing::warn!(
35                target: "forge_workspace::store::state",
36                account = %name.value(),
37                error = %err,
38                "skipping account-usage record that failed to decode",
39            ),
40        }
41    }
42    Ok(out)
43}
44
45/// Overwrite the usage cache with `entries`: clear every row then insert
46/// the full map in one write txn, mirroring the poller's whole-map write.
47pub fn replace_account_usage(
48    db: &Db,
49    entries: &BTreeMap<String, CachedAccountUsage>,
50) -> anyhow::Result<()> {
51    let txn = db.database().begin_write()?;
52    {
53        let mut table = txn.open_table(ACCOUNT_USAGE)?;
54        let existing: Vec<String> = table
55            .iter()?
56            .map(|entry| entry.map(|(k, _)| k.value().to_owned()))
57            .collect::<Result<_, _>>()?;
58        for key in existing {
59            table.remove(key.as_str())?;
60        }
61        for (name, entry) in entries {
62            let value = serde_json::to_vec(entry).context("serialize account usage")?;
63            table.insert(name.as_str(), value.as_slice())?;
64        }
65    }
66    txn.commit()?;
67    Ok(())
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use forge_primitives::usage::{ApiSpend, UsageSnapshot, UsageSourceKind, UsageWindow};
74    use std::time::{Duration, SystemTime};
75    use tempfile::tempdir;
76
77    fn usage_entry(utilization: f64) -> CachedAccountUsage {
78        CachedAccountUsage {
79            snapshot: UsageSnapshot {
80                source: UsageSourceKind::Oauth,
81                fetched_at: SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000),
82                five_hour: Some(UsageWindow {
83                    utilization,
84                    resets_at: None,
85                    reset_description: None,
86                    status: None,
87                }),
88                seven_day: None,
89                seven_day_opus: None,
90                seven_day_sonnet: None,
91                extra_usage: None,
92                spend: None,
93                balance: None,
94            },
95        }
96    }
97
98    fn usage_map(pairs: &[(&str, f64)]) -> BTreeMap<String, CachedAccountUsage> {
99        pairs.iter().map(|(name, util)| ((*name).to_owned(), usage_entry(*util))).collect()
100    }
101
102    #[test]
103    fn account_usage_replace_mirrors_the_map() {
104        let dir = tempdir().expect("tempdir");
105        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
106
107        replace_account_usage(&db, &usage_map(&[("Gateway", 42.0), ("Stargate", 10.0)]))
108            .expect("first write");
109        let loaded = account_usage(&db).expect("read");
110        assert_eq!(loaded.len(), 2);
111        assert_eq!(
112            loaded
113                .get("Gateway")
114                .and_then(|e| e.snapshot.five_hour.as_ref())
115                .map(|w| w.utilization),
116            Some(42.0),
117        );
118
119        // A second write drops the account no longer present.
120        replace_account_usage(&db, &usage_map(&[("Gateway", 99.0)])).expect("second write");
121        let loaded = account_usage(&db).expect("read again");
122        assert_eq!(loaded.len(), 1, "replace mirrors exactly the new map");
123        assert!(!loaded.contains_key("Stargate"), "the dropped account is gone");
124        assert_eq!(
125            loaded
126                .get("Gateway")
127                .and_then(|e| e.snapshot.five_hour.as_ref())
128                .map(|w| w.utilization),
129            Some(99.0),
130            "the surviving account's snapshot is the new value",
131        );
132    }
133
134    /// A row written before `UsageSnapshot` grew `spend` must still
135    /// decode, because there is no migration step: the cache is read
136    /// straight back at boot. Written as raw JSON in the old shape
137    /// rather than by constructing today's type, so the assertion is
138    /// about what is actually on disk.
139    #[test]
140    fn an_old_shape_account_usage_row_still_decodes() {
141        let dir = tempdir().expect("tempdir");
142        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
143
144        let old_row = br#"{"snapshot":{
145            "source":"Oauth",
146            "fetched_at":{"secs_since_epoch":1735128000,"nanos_since_epoch":0},
147            "five_hour":{"utilization":42.0,"resets_at":null,"reset_description":null},
148            "seven_day":{"utilization":7.0,"resets_at":null,"reset_description":null},
149            "seven_day_opus":null,"seven_day_sonnet":null,
150            "extra_usage":{"monthly_limit":20.0,"used_credits":12.5,
151                           "utilization":62.5,"currency":"USD"}
152        }}"#;
153
154        let txn = db.database().begin_write().expect("begin");
155        {
156            let mut table = txn.open_table(ACCOUNT_USAGE).expect("open table");
157            table.insert("Gateway", old_row.as_slice()).expect("insert old-shape row");
158        }
159        txn.commit().expect("commit");
160
161        let loaded = account_usage(&db).expect("read");
162        let entry = loaded.get("Gateway").expect("the pre-spend row survives the new shape");
163        let snapshot = &entry.snapshot;
164        assert_eq!(snapshot.source, UsageSourceKind::Oauth, "the source it carried is preserved");
165        assert_eq!(
166            snapshot.five_hour.as_ref().map(|w| w.utilization),
167            Some(42.0),
168            "the five-hour window it carried is preserved",
169        );
170        assert_eq!(
171            snapshot.seven_day.as_ref().map(|w| w.utilization),
172            Some(7.0),
173            "the seven-day window it carried is preserved",
174        );
175        assert_eq!(
176            snapshot.extra_usage.as_ref().and_then(|e| e.used_credits),
177            Some(12.5),
178            "the extra-usage block it carried is preserved",
179        );
180        assert!(snapshot.spend.is_none(), "the absent new field decodes to None");
181    }
182
183    /// The sibling case: a row whose `spend` block predates the cap
184    /// fields. The test above never reaches `ApiSpend` at all, so only
185    /// this one fails if a cap field is ever declared bare instead of
186    /// `Option` - which would make every cached API row undecodable.
187    #[test]
188    fn a_spend_row_written_before_the_cap_fields_still_decodes() {
189        let dir = tempdir().expect("tempdir");
190        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
191
192        let old_row = br#"{"snapshot":{
193            "source":"OpenRouterKey",
194            "fetched_at":{"secs_since_epoch":1735128000,"nanos_since_epoch":0},
195            "five_hour":null,"seven_day":null,
196            "seven_day_opus":null,"seven_day_sonnet":null,
197            "extra_usage":null,
198            "spend":{"daily":0.5,"weekly":1.25,"monthly":20.3}
199        }}"#;
200
201        let txn = db.database().begin_write().expect("begin");
202        {
203            let mut table = txn.open_table(ACCOUNT_USAGE).expect("open table");
204            table.insert("Api", old_row.as_slice()).expect("insert pre-cap row");
205        }
206        txn.commit().expect("commit");
207
208        let loaded = account_usage(&db).expect("read");
209        let spend = loaded
210            .get("Api")
211            .expect("the pre-cap row survives the new shape")
212            .snapshot
213            .spend
214            .as_ref()
215            .expect("the spend block it carried is preserved");
216
217        assert_eq!(
218            spend,
219            &ApiSpend {
220                daily: 0.5,
221                weekly: 1.25,
222                monthly: 20.3,
223                limit: None,
224                limit_remaining: None,
225                limit_reset: None,
226                expires_at: None,
227            },
228            "the figures it carried survive and an absent cap stays absent rather than becoming zero",
229        );
230    }
231
232    /// The sibling case for `balance`: a row written before the
233    /// snapshot grew the account credit pool must still decode, because
234    /// redb has no migration step and the cache is read straight back
235    /// at boot.
236    #[test]
237    fn an_account_usage_row_written_before_balance_still_decodes() {
238        let dir = tempdir().expect("tempdir");
239        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
240
241        let old_row = br#"{"snapshot":{
242            "source":"OpenRouterKey",
243            "fetched_at":{"secs_since_epoch":1735128000,"nanos_since_epoch":0},
244            "five_hour":null,"seven_day":null,
245            "seven_day_opus":null,"seven_day_sonnet":null,
246            "extra_usage":null,
247            "spend":{"daily":0.5,"weekly":1.25,"monthly":20.3}
248        }}"#;
249
250        let txn = db.database().begin_write().expect("begin");
251        {
252            let mut table = txn.open_table(ACCOUNT_USAGE).expect("open table");
253            table.insert("Api", old_row.as_slice()).expect("insert pre-balance row");
254        }
255        txn.commit().expect("commit");
256
257        let loaded = account_usage(&db).expect("read");
258        let snapshot = &loaded.get("Api").expect("the pre-balance row survives").snapshot;
259        assert!(snapshot.spend.is_some(), "the spend block it carried is preserved");
260        assert_eq!(
261            snapshot.balance, None,
262            "the absent balance decodes to None rather than failing the row",
263        );
264    }
265
266    #[test]
267    fn corrupt_account_usage_record_is_skipped_not_fatal() {
268        let dir = tempdir().expect("tempdir");
269        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
270
271        replace_account_usage(&db, &usage_map(&[("Gateway", 42.0)])).expect("write good");
272
273        // A blob that isn't a valid entry must not poison the load.
274        let txn = db.database().begin_write().expect("begin");
275        {
276            let mut table = txn.open_table(ACCOUNT_USAGE).expect("open table");
277            table.insert("corrupt", "not usage".as_bytes()).expect("insert corrupt");
278        }
279        txn.commit().expect("commit");
280
281        let loaded = account_usage(&db).expect("read tolerates the corrupt blob");
282        assert_eq!(loaded.len(), 1, "the good record survives a corrupt sibling");
283        assert!(loaded.contains_key("Gateway"));
284    }
285
286    #[test]
287    fn state_survives_db_reopen() {
288        let dir = tempdir().expect("tempdir");
289        let path = dir.path().join("db.redb");
290        {
291            let db = Db::open(&path).expect("open db");
292            replace_account_usage(&db, &usage_map(&[("Gateway", 42.0)])).expect("write usage");
293        }
294        let db = Db::open(&path).expect("reopen db");
295        assert!(
296            account_usage(&db).expect("read").contains_key("Gateway"),
297            "the usage cache survives restart",
298        );
299    }
300}
forge_workspace/store/state.rs

forge_workspace/store/
state.rs
