forge_workspace/
account_cache.rs

1//! Persistent forge state: the per-account usage cache.
2//!
3//! It lives in the machine-local redb store ([`crate::store::state`]).
4//! [`load`] and [`store`] are thin wrappers over that tenant;
5//! [`ForgeState`] / [`CachedAccountUsage`] are the in-memory shapes.
6//!
7//! The usage cache solves the cold-boot problem: Anthropic's
8//! `/api/oauth/usage` endpoint rate-limits aggressively on per-IP burst
9//! probes, so the first launch can wait 30 s+ before the warm probe gets
10//! through - during which the launchpad picker has no usage to show and
11//! the walk sees every account as unknown. The cache seeds the in-memory
12//! `AccountStateMap` with the last known values until the 60 s poller
13//! refreshes them.
14//!
15//! Failures are non-fatal: a closed store degrades to "no cache; spawn
16//! paths see empty bars until the poller succeeds."
17
18use forge_primitives::usage::UsageSnapshot;
19use serde::{Deserialize, Serialize};
20
21/// Per-account cache entry. The value type of the redb `account_usage`
22/// table; reachable only in-crate since `account_cache` is a private mod.
23#[derive(Debug, Serialize, Deserialize)]
24pub struct CachedAccountUsage {
25    pub snapshot: UsageSnapshot,
26}
27
28/// In-memory forge state read from the redb store.
29#[derive(Debug)]
30pub(crate) struct ForgeState {
31    /// Account display name to cached snapshot. The redb row shape
32    /// ([`CachedAccountUsage`]) is unwrapped at the store boundary, so
33    /// the in-memory map the account pool consumes is bare snapshots.
34    pub account_usage: std::collections::BTreeMap<String, UsageSnapshot>,
35}
36
37impl ForgeState {
38    pub(crate) fn empty() -> Self {
39        Self { account_usage: std::collections::BTreeMap::new() }
40    }
41}
42
43/// Read the persisted forge state from the store. Each read degrades to
44/// its empty default with a warn rather than failing the boot.
45pub(crate) fn load(db: &crate::store::Db) -> ForgeState {
46    ForgeState {
47        account_usage: crate::store::state::account_usage(db).map_or_else(
48            |error| {
49                tracing::warn!(
50                    target: "forge_workspace::account_cache",
51                    %error,
52                    "reading the account-usage cache from the store failed",
53                );
54                std::collections::BTreeMap::new()
55            },
56            |rows| rows.into_iter().map(|(name, row)| (name, row.snapshot)).collect(),
57        ),
58    }
59}
60
61/// Persist the in-memory account-usage snapshots, replacing the prior
62/// set. Backs the 60 s poller. Non-fatal + logged on failure.
63pub(crate) fn store(
64    db: &crate::store::Db,
65    entries: &std::collections::BTreeMap<String, UsageSnapshot>,
66) {
67    let rows: std::collections::BTreeMap<String, CachedAccountUsage> = entries
68        .iter()
69        .map(|(name, snapshot)| (name.clone(), CachedAccountUsage { snapshot: snapshot.clone() }))
70        .collect();
71    if let Err(error) = crate::store::state::replace_account_usage(db, &rows) {
72        tracing::warn!(
73            target: "forge_workspace::account_cache",
74            %error,
75            "persisting account usage to the store failed",
76        );
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use forge_primitives::usage::{UsageSnapshot, UsageSourceKind, UsageWindow};
84    use std::time::{Duration, SystemTime};
85    use tempfile::tempdir;
86
87    fn cfg() -> tempfile::TempDir {
88        tempdir().expect("cfg tempdir")
89    }
90
91    fn fixture_entry() -> UsageSnapshot {
92        UsageSnapshot {
93            source: UsageSourceKind::Oauth,
94            fetched_at: SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000),
95            five_hour: Some(UsageWindow {
96                utilization: 42.0,
97                resets_at: None,
98                reset_description: None,
99                status: None,
100            }),
101            seven_day: None,
102            seven_day_opus: None,
103            seven_day_sonnet: None,
104            extra_usage: None,
105            spend: None,
106            balance: None,
107        }
108    }
109
110    #[test]
111    fn store_and_load_round_trip_through_redb() {
112        let cfg = cfg();
113        let db = crate::store::Db::open(&cfg.path().join("db.redb")).expect("open db");
114        let mut entries = std::collections::BTreeMap::new();
115        entries.insert("Gateway".to_owned(), fixture_entry());
116        store(&db, &entries);
117
118        let loaded = load(&db);
119        assert!(loaded.account_usage.contains_key("Gateway"), "the usage cache reloads");
120    }
121}
forge_workspace/account_cache.rs

forge_workspace/
account_cache.rs
