forge_workspace/
account_loader.rs

1//! Per-account boot-time loading task.
2//!
3//! Probes one account once and settles it to a terminal `LoadingState`
4//! (`Ready` or `Bailed`) in that single round. Spawned at
5//! `Workspace::new`, one tokio task per `[[accounts]]` entry in
6//! forge.toml. The boot wave and the cron spawn path wait until
7//! `AccountStateMap::all_loaded()` returns true (every account in a
8//! terminal state); the gateway's walk only considers accounts whose
9//! terminal state is `Ready`.
10//!
11//! One probe per task, every outcome terminal:
12//! - 200 -> snapshot stored via `set_usage`, `Ready`, task exits.
13//! - Any probe error -> `Bailed` in the same round. Auth failures
14//!   record their status; a 200 whose body maps to nothing
15//!   (`Unmappable`) records nothing; the transient classes
16//!   (`RateLimited` etc.) record status + the server `Retry-After`,
17//!   which also schedules the poller's re-probe.
18//!
19//! Healing is the pollers' job: the 60 s usage poller re-probes
20//! accounts once their `Retry-After` / backoff window passes and flips
21//! `Bailed` -> `Ready`, which makes the account eligible for new
22//! spawns without shifting running sessions. A rate-limited probe does
23//! not mean inference is limited - the session surfaces its own error
24//! if it is.
25
26use std::sync::Weak;
27
28use forge_gateway::ProbeError;
29use forge_primitives::usage::UsageSnapshot;
30use forge_primitives::usage::oauth::OauthUsageError;
31
32use forge_gateway::{AccountKey, AccountPool, LoadingState};
33
34use crate::workspace::Workspace;
35
36/// Apply one probe outcome to the account state map and return the
37/// terminal `LoadingState` it settles to. Every error class is
38/// terminal in this single pass; healing is the usage poller's
39/// schedule.
40fn settle_probe_result(
41    pool: &AccountPool,
42    key: &AccountKey,
43    result: &Result<UsageSnapshot, ProbeError>,
44) -> LoadingState {
45    match result {
46        Ok(snapshot) => {
47            pool.set_usage(key, snapshot.clone());
48            LoadingState::Ready
49        }
50        Err(ProbeError::Unmappable(_)) => {
51            // Response-shape drift is not an account fault: records no
52            // error; any stale record clears with the bail.
53            pool.clear_last_error(key);
54            pool.set_loading(key, LoadingState::Bailed);
55            LoadingState::Bailed
56        }
57        Err(err) => {
58            let failure_class = crate::workspace::classify_oauth_usage_error(err);
59            let retry_after = match err {
60                ProbeError::Fetch(OauthUsageError::RateLimited { retry_after, .. }) => *retry_after,
61                _ => None,
62            };
63            pool.set_last_error(key, failure_class, retry_after);
64            // set_last_error bails only on the auth classes; the
65            // transient ones need the explicit terminal settle.
66            pool.set_loading(key, LoadingState::Bailed);
67            // After the bail, which drops the caches: a 429 is a reading
68            // even though it is not a healthy account.
69            if let Some(snapshot) = err.rate_limited_snapshot() {
70                pool.set_limit_reading(key, snapshot.clone());
71            }
72            LoadingState::Bailed
73        }
74    }
75}
76
77/// Probe the account once and settle its terminal `LoadingState`.
78///
79/// Takes `Weak<Workspace>` so the task auto-exits when the workspace
80/// is dropped during shutdown - same pattern `start_usage_poller`
81/// uses. The lock is only acquired for single mutator calls, never
82/// across an await, so other workspace operations aren't blocked.
83pub async fn run_account_loading(account_key: AccountKey, workspace_weak: Weak<Workspace>) {
84    let Some(workspace) = workspace_weak.upgrade() else {
85        // Workspace dropped during shutdown; exit cleanly.
86        return;
87    };
88    let pool = workspace.account_pool();
89    let Some(provider) = pool.provider(&account_key) else {
90        tracing::warn!(
91            target: "forge_workspace::account_loader",
92            account = %account_key.0,
93            "no provider for the account key; skipping its boot probe",
94        );
95        return;
96    };
97    let account_env = pool.env(&account_key).unwrap_or_default();
98    // The backend owns the probe; the settle helper owns the verdict
99    // against the state machine.
100    let probe_result = crate::provider_probe::probe_via_backend(provider, &account_env).await;
101    match &probe_result {
102        Ok(_) => tracing::info!(
103            target: "forge_workspace::account_loader",
104            account = %account_key.0,
105            "boot probe reached Ready",
106        ),
107        Err(err) => tracing::warn!(
108            target: "forge_workspace::account_loader",
109            account = %account_key.0,
110            error = %err,
111            "boot probe failed; account settled Bailed; the usage poller will re-probe it",
112        ),
113    }
114    settle_probe_result(pool, &account_key, &probe_result);
115    workspace.announce_accounts_changed();
116}
117
118#[cfg(test)]
119#[allow(clippy::expect_used, clippy::unwrap_used)]
120mod tests {
121    use super::*;
122    use forge_gateway::ProbeError;
123    use forge_gateway::UsageFetchStatus;
124    use forge_primitives::usage::oauth::OauthUsageError;
125    use forge_primitives::usage::{UsageSnapshot, UsageSourceKind, UsageWindow};
126    use std::time::Duration;
127
128    // Mirrors the `make_account` helper idiom in `account::tests`.
129    fn states_with_one_account() -> AccountPool {
130        AccountPool::new(&[crate::config::LoadedAccount {
131            display_name: "test".to_owned(),
132            provider: forge_primitives::account::Provider::Anthropic,
133            base_url: None,
134            models: vec!["claude-sonnet-5".to_owned()],
135            model_aliases: std::collections::HashMap::new(),
136            model_slugs: std::collections::HashMap::new(),
137            env: std::collections::HashMap::new(),
138        }])
139    }
140
141    fn snapshot() -> UsageSnapshot {
142        UsageSnapshot {
143            source: UsageSourceKind::Oauth,
144            fetched_at: std::time::SystemTime::UNIX_EPOCH,
145            five_hour: Some(UsageWindow {
146                utilization: 10.0,
147                resets_at: Some(std::time::SystemTime::now() + Duration::from_secs(3600)),
148                reset_description: None,
149                status: None,
150            }),
151            seven_day: Some(UsageWindow {
152                utilization: 20.0,
153                resets_at: Some(std::time::SystemTime::now() + Duration::from_secs(3600)),
154                reset_description: None,
155                status: None,
156            }),
157            seven_day_opus: None,
158            seven_day_sonnet: None,
159            extra_usage: None,
160            spend: None,
161            balance: None,
162        }
163    }
164
165    #[test]
166    fn settle_probe_result_bails_rate_limited_with_status_and_retry_after() {
167        let states = states_with_one_account();
168        let key = AccountKey("test".to_owned());
169        let state = settle_probe_result(
170            &states,
171            &key,
172            &Err(ProbeError::Fetch(OauthUsageError::RateLimited {
173                retry_after: Some(Duration::from_secs(3600)),
174                snapshot: None,
175            })),
176        );
177        assert_eq!(state, LoadingState::Bailed);
178        assert_eq!(
179            states.loading_state(&key),
180            LoadingState::Bailed,
181            "the terminal settle is stored, not just returned",
182        );
183        assert_eq!(
184            states.usage_error("test"),
185            Some(UsageFetchStatus::RateLimited),
186            "the rate-limit class is recorded for the row's reason text",
187        );
188        assert!(
189            states.latest_reading("test").is_none(),
190            "a 429 with no unified headers has no reading to keep",
191        );
192        // The server Retry-After schedules the re-probe (clamped to
193        // the 10-minute ceiling); the local exponential default would
194        // be 30 s, so a ten-minute gap proves the server value was
195        // honoured, not discarded.
196        let gap = states.next_probe_after(&key);
197        assert!(
198            gap.is_some_and(|g| g > Duration::from_secs(599) && g <= Duration::from_secs(601)),
199            "Retry-After retained (clamped) in the probe schedule; got {gap:?}",
200        );
201    }
202
203    /// **A 429's windows survive the boot settle.** The account still bails
204    /// for the walk - a rate-limited probe heals nothing - and the reading
205    /// that says which window is spent is there for a view.
206    #[test]
207    fn settle_probe_result_keeps_a_rate_limited_readings_windows() {
208        let states = states_with_one_account();
209        let key = AccountKey("test".to_owned());
210        let state = settle_probe_result(
211            &states,
212            &key,
213            &Err(ProbeError::Fetch(OauthUsageError::RateLimited {
214                retry_after: Some(Duration::from_secs(3600)),
215                snapshot: Some(Box::new(snapshot())),
216            })),
217        );
218        assert_eq!(state, LoadingState::Bailed, "a 429 does not heal the account");
219        assert_eq!(states.loading_state(&key), LoadingState::Bailed);
220        assert_eq!(states.usage_error("test"), Some(UsageFetchStatus::RateLimited));
221        let reading = states.latest_reading("test").expect("the 429's windows survived the bail");
222        assert!(reading.five_hour.is_some(), "the 5h window was dropped");
223        assert!(reading.seven_day.is_some(), "the 7d window was dropped");
224    }
225
226    #[test]
227    fn settle_probe_result_bails_unauthorized_with_its_status() {
228        let states = states_with_one_account();
229        let key = AccountKey("test".to_owned());
230        let state = settle_probe_result(
231            &states,
232            &key,
233            &Err(ProbeError::Fetch(OauthUsageError::Unauthorized(401))),
234        );
235        assert_eq!(state, LoadingState::Bailed);
236        assert_eq!(
237            states.loading_state(&key),
238            LoadingState::Bailed,
239            "the terminal settle is stored, not just returned",
240        );
241        assert_eq!(states.usage_error("test"), Some(UsageFetchStatus::Unauthorized),);
242    }
243
244    #[test]
245    fn settle_probe_result_bails_unmappable_without_an_error_record() {
246        let states = states_with_one_account();
247        let key = AccountKey("test".to_owned());
248        // A stale record from an earlier probe must not survive the
249        // shape-drift bail - the bail records nothing and clears.
250        states.set_last_error(&key, UsageFetchStatus::RateLimited, None);
251        let state = settle_probe_result(
252            &states,
253            &key,
254            &Err(ProbeError::Unmappable("shape drift".to_owned())),
255        );
256        assert_eq!(state, LoadingState::Bailed);
257        assert_eq!(
258            states.loading_state(&key),
259            LoadingState::Bailed,
260            "the terminal settle is stored, not just returned",
261        );
262        assert!(
263            states.usage_error("test").is_none(),
264            "Unmappable records no error, preserving today's nuance",
265        );
266    }
267
268    #[test]
269    fn settle_probe_result_settles_successful_probe_ready() {
270        let states = states_with_one_account();
271        let key = AccountKey("test".to_owned());
272        let state = settle_probe_result(&states, &key, &Ok(snapshot()));
273        assert_eq!(state, LoadingState::Ready);
274        assert_eq!(
275            states.loading_state(&key),
276            LoadingState::Ready,
277            "the terminal settle is stored, not just returned",
278        );
279        assert!(states.usage("test").is_some(), "the snapshot is cached");
280    }
281}
forge_workspace/account_loader.rs

forge_workspace/
account_loader.rs
