forge_gateway/

account.rs

//! Account state, health and probe scheduling.
//!
//! The per-account state map every selection reads: it holds each
//! account's declared models, its `LoadingState`, its latest usage
//! snapshot and its last probe outcome. The walk that turns that into
//! a picked account lives in `selection.rs`; this module owns the
//! state and the classification the walk and the accounts read both
//! use.
//!
//! **Two-tier classification:**
//!
//! 1. **Usable** - usage is unknown, OR usage shows under 100% on
//!    both windows, OR last probe failed transiently (network /
//!    other HTTP). A session can try this account; the spawned
//!    `claude` either succeeds or surfaces its own rate-limit error
//!    we can react to.
//! 2. **Unusable** - usage shows 100% on at least one window, OR
//!    `/api/oauth/usage` returned 429, OR credentials are expired /
//!    unauthorized. Known to be either at the cap or unable to
//!    authenticate.
//!
//! Saturation and probe failures are classification only - they feed
//! the walk's preference order and the accounts read.
//!
//! Rate-limited accounts are visible in a view's usage bars so the
//! user can wait for the window to reset.

use forge_primitives::usage::{UsageSnapshot, UsageWindow};

use forge_primitives::account::LoadedAccount;

/// Newtype wrapping the account's `display_name`.
#[derive(Clone, Debug, Hash, Eq, PartialEq)]
pub struct AccountKey(pub String);

/// Classification of the latest usage-poll attempt outcome for an
/// account. Feeds `unusable_reason`'s probe-blocked branch and the
/// backoff schedule.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageFetchStatus {
    /// Anthropic returned HTTP 429 - too many concurrent polls
    /// against the OAuth `/api/oauth/usage` endpoint (typical when
    /// multiple forge instances poll from the same machine).
    RateLimited,
    /// The account's credential is dead or absent (expired, or no
    /// token at all). Repair is minting the setup token on the account
    /// block plus a restart.
    Expired,
    /// API returned 401/403 - token rejected (may be revoked).
    Unauthorized,
    /// Network failure reaching the API (DNS, TLS, timeout, …).
    NetworkFailed,
    /// Decode / unknown-HTTP-status fallthrough, kept distinct so an
    /// unmapped cause names itself instead of wearing a known bucket.
    /// Like `NetworkFailed` it does not block the account.
    Other,
}

/// Why an account is held out of the usable tier.
/// Which reason can apply is keyed on the account's billing kind: only
/// a window-billed account can saturate, so on a spend row an unusable
/// account can only be probe-blocked. A snapshot whose source is not
/// the account backend's is ignored here, so a cache surviving a
/// `forge.toml` provider edit cannot saturate the row it landed on.
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Unusable {
    /// A plan window is at or over its cap and has not reset yet.
    Saturated,
    /// The last probe failed before it could read usage: throttled,
    /// rejected, or expired credentials.
    ProbeBlocked,
    /// The boot-time loading task ended in `Bailed` on an auth
    /// failure.
    Bailed,
}

/// Boot-time loading state for an account. The boot wave and the cron
/// spawn path wait for every account in the map to resolve to `Ready`
/// or `Bailed`; the selection walk skips `Loading` and keeps `Bailed`
/// as its last resort. A bailed account's `usage` is `None` by
/// construction (the loader clears it on the transition).
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LoadingState {
    /// First-pass probe in progress.
    Loading,
    /// Probe returned 200; account is usable and may be assigned to
    /// sessions.
    Ready,
    /// The probe ended in an auth failure the account cannot recover
    /// from in place. The user fixes the credential and lets the 60 s
    /// usage poller flip the account back to `Ready`.
    Bailed,
}

#[derive(Clone)]
pub struct AccountState {
    /// Declared backend from `[[accounts]] provider`. Decides which
    /// endpoint the usage probe hits and which credential it carries.
    pub provider: forge_primitives::account::Provider,
    /// The canonical model names the account serves. Selection and
    /// the route's model gate match through [`account_serves`].
    pub models: Vec<String>,
    /// Canonical name -> the other names a request may arrive under for
    /// that model. Keys are always members of [`Self::models`].
    pub model_aliases: std::collections::HashMap<String, Vec<String>>,
    /// Canonical name -> upstream slug, only where the spellings
    /// differ. The forward splices the slug into the request body.
    pub model_slugs: std::collections::HashMap<String, String>,
    /// Per-account environment: provider-behaviour extras plus the
    /// derived credential and base-url variable stamps. Stamped onto
    /// the account's `claude` subprocess at spawn, and read by the
    /// usage probe for the `ANTHROPIC_BASE_URL` and `ANTHROPIC_AUTH_TOKEN`
    /// a base-url provider authenticates with.
    pub env: std::collections::HashMap<String, String>,
    /// Latest usage snapshot fetched by the workspace's background
    /// poller. `None` until the first successful fetch. Drives the
    /// selection walk's order and the cooldown's reset; a view reads
    /// [`Self::limit_reading`] with it, in that order.
    pub usage: Option<UsageSnapshot>,
    /// The windows a rate-limited probe carried, kept for the view beside
    /// `usage` rather than in it: they never heal the account, and
    /// selection and cooldowns must keep reading the successful probe.
    pub limit_reading: Option<UsageSnapshot>,
    /// Latest poll-attempt outcome when the fetch failed. Cleared
    /// (set to `None`) on the next successful fetch. `unusable_reason`
    /// reads this to classify a probe-blocked account, and the write
    /// schedules the next probe's backoff (`next_probe_at`).
    pub last_error: Option<UsageFetchStatus>,
    /// Earliest wall-clock time we're allowed to probe this account
    /// again. The poller skips accounts whose `next_probe_at` is in
    /// the future. Used for per-account exponential backoff when
    /// Anthropic's `/api/oauth/usage` endpoint keeps returning 429
    /// (without backoff every poll cycle re-trips the per-IP rate
    /// limit and we never accumulate usage data). `None` means
    /// "probe on the next tick" (default).
    pub next_probe_at: Option<std::time::Instant>,
    /// Consecutive probe failures since the last success. Drives
    /// the backoff schedule: each consecutive failure doubles the
    /// next-probe delay (capped at a sensible ceiling). Reset to
    /// 0 on success.
    pub consecutive_failures: u32,
    /// One-shot arming flag for the `has_just_cleared_cap_window`
    /// scheduler hook. Without it, a stale-reset snapshot paired with
    /// sustained probe failure would re-trigger the override every
    /// poll cycle, defeating the exponential backoff schedule -
    /// forge would hammer Anthropic's `/api/oauth/usage` endpoint once
    /// per minute through a multi-hour outage instead of backing off.
    /// `set_usage` arms it (a fresh snapshot earns one override
    /// attempt once its `resets_at` passes); `disarm_override` clears
    /// it after the scheduler fires an override. False until the
    /// first successful probe lands a snapshot. Distinct from
    /// `LoadingState` - rate-limit concern, not auth-recovery.
    pub override_armed: bool,
    /// Where this account sits in the boot-time loading state
    /// machine. The boot wave and the cron spawn path gate on this
    /// across every account; the selection walk only considers
    /// accounts whose state is `Ready` (or a saturated `Ready`).
    pub loading: LoadingState,
    /// A config reload removed this account from the file. It is not in the
    /// selection walk any more, but a session already bound to it keeps
    /// working until the session ends.
    pub retired: bool,
}

/// Hand-written because the env map carries the account's credential
/// stamps, which the derived output would print.
impl std::fmt::Debug for AccountState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let env: std::collections::BTreeMap<&str, &str> =
            self.env.keys().map(|key| (key.as_str(), "[redacted]")).collect();
        f.debug_struct("AccountState")
            .field("provider", &self.provider)
            .field("models", &self.models)
            .field("model_aliases", &self.model_aliases)
            .field("model_slugs", &self.model_slugs)
            .field("env", &env)
            .field("usage", &self.usage)
            .field("limit_reading", &self.limit_reading)
            .field("last_error", &self.last_error)
            .field("next_probe_at", &self.next_probe_at)
            .field("consecutive_failures", &self.consecutive_failures)
            .field("override_armed", &self.override_armed)
            .field("loading", &self.loading)
            .field("retired", &self.retired)
            .finish()
    }
}

/// `true` when an account declaring `models` serves `model`: it
/// declares the name, or declares it as an alias of a model it does.
/// The one home for the rule: the selection walk, the route's binding
/// gate, and the config load gate all match a request's model here.
pub fn account_serves<S: std::hash::BuildHasher>(
    models: &[String],
    model_aliases: &std::collections::HashMap<String, Vec<String>, S>,
    model: &str,
) -> bool {
    models.iter().any(|m| m == model)
        || model_aliases.values().flatten().any(|alias| alias == model)
}

pub struct AccountStateMap {
    pub ordered_keys: Vec<AccountKey>, // forge.toml definition order
    pub by_key: std::collections::HashMap<AccountKey, AccountState>,
}

/// Hand-written so a field added later cannot print through it until it
/// is deliberately listed; the states it holds redact themselves.
impl std::fmt::Debug for AccountStateMap {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AccountStateMap")
            .field("ordered_keys", &self.ordered_keys)
            .field("by_key", &self.by_key)
            .finish()
    }
}

impl AccountStateMap {
    /// Empty map for the `testing` feature's `Workspace::testing_stub`.
    /// Production code paths reach this map only through the gateway's
    /// selection walk, which a test fixture should never exercise.
    #[cfg(any(test, feature = "testing"))]
    pub fn empty_for_test() -> Self {
        Self { ordered_keys: Vec::new(), by_key: std::collections::HashMap::new() }
    }

    pub fn new(accounts: &[LoadedAccount]) -> Self {
        let mut ordered_keys = Vec::with_capacity(accounts.len());
        let mut by_key = std::collections::HashMap::with_capacity(accounts.len());
        for account in accounts {
            let key = AccountKey(account.display_name.clone());
            ordered_keys.push(key.clone());
            by_key.insert(
                key,
                AccountState {
                    provider: account.provider,
                    models: account.models.clone(),
                    model_aliases: account.model_aliases.clone(),
                    model_slugs: account.model_slugs.clone(),
                    env: account.env.clone(),
                    usage: None,
                    limit_reading: None,
                    last_error: None,
                    next_probe_at: None,
                    consecutive_failures: 0,
                    override_armed: false,
                    loading: LoadingState::Loading,
                    retired: false,
                },
            );
        }
        Self { ordered_keys, by_key }
    }

    /// Seed the in-memory map from a previously-persisted cache.
    /// Each known account that appears in `cached` gets its `usage`
    /// populated; unknown cache entries (account removed from
    /// forge.toml) are ignored. Does NOT clear `last_error` or
    /// `next_probe_at` - the cache is purely seed data; the live
    /// poller still drives backoff. Used by `Workspace::new` so
    /// usage reads have data on cold boot.
    pub fn seed_from_cache(&mut self, cached: &std::collections::BTreeMap<String, UsageSnapshot>) {
        for (name, snapshot) in cached {
            let key = AccountKey(name.clone());
            if let Some(state) = self.by_key.get_mut(&key) {
                state.usage = Some(snapshot.clone());
            }
        }
    }

    /// Snapshot the per-account `usage` for writing to the on-disk
    /// cache. Accounts with no live snapshot are omitted so the cache
    /// file doesn't grow placeholders.
    pub fn snapshots_for_cache(&self) -> std::collections::BTreeMap<String, UsageSnapshot> {
        let mut out = std::collections::BTreeMap::new();
        for (key, state) in &self.by_key {
            if let Some(snapshot) = state.usage.clone() {
                out.insert(key.0.clone(), snapshot);
            }
        }
        out
    }

    /// The account's env map: global `[env]` extended by its own
    /// `[accounts.env]`. Consumed by the spawn path (stamped onto the
    /// child) and the usage poller / loader (base url, bearer, setup
    /// token). `None` for unknown keys.
    pub fn env(&self, key: &AccountKey) -> Option<&std::collections::HashMap<String, String>> {
        self.by_key.get(key).map(|s| &s.env)
    }

    /// Declared [`forge_primitives::account::Provider`] for `key`, the input the probe plan and the
    /// preflight repair copy both branch on. `None` for unknown keys.
    pub fn provider(&self, key: &AccountKey) -> Option<forge_primitives::account::Provider> {
        self.by_key.get(key).map(|s| s.provider)
    }

    /// How the account proves who it is, which is the only thing that
    /// changes what an auth-repair hint tells the user to do. Base-url
    /// wins: such an account authenticates with the flat token beside
    /// its base_url. Every other account's credential is its setup
    /// token, so both its token-mode and its token-less shapes repair
    /// the same way. `None` for unknown keys.
    pub fn auth(&self, key: &AccountKey) -> Option<forge_primitives::account::AccountAuth> {
        if self.provider(key)?.uses_base_url() {
            return Some(forge_primitives::account::AccountAuth::BaseUrl);
        }
        Some(forge_primitives::account::AccountAuth::Token)
    }

    /// Replace the cached usage snapshot for `key` and clear any
    /// stale `last_error`. Resets the consecutive-failure counter
    /// and clears the next-probe gate so this account returns to
    /// the default cadence. Also transitions `loading` to `Ready` -
    /// a successful probe is the terminal-success state for the
    /// boot-time loading task. Silent no-op when `key` isn't
    /// registered (defensive - invariant says every poller key was
    /// inserted in `new()`).
    pub fn set_usage(&mut self, key: &AccountKey, snapshot: UsageSnapshot) {
        if let Some(state) = self.by_key.get_mut(key) {
            state.usage = Some(snapshot);
            state.limit_reading = None;
            state.last_error = None;
            state.next_probe_at = None;
            state.consecutive_failures = 0;
            state.loading = LoadingState::Ready;
            // Arm the scheduler-hook override: this fresh snapshot
            // earns ONE probe attempt once its `resets_at` passes.
            // Without re-arming, the hook stays disarmed forever after
            // its first fire and a stale cached window never gets a
            // fresh probe even on a healthy account.
            state.override_armed = true;
        }
    }

    /// Drive a `LoadingState` transition for `key`, used by the
    /// boot-time loading task. Setting `Bailed` clears the cached
    /// `usage` so no stale bar is drawn (replaces the PR #238
    /// 3-strike counter; bailed accounts have no live snapshot by
    /// construction).
    pub fn set_loading(&mut self, key: &AccountKey, loading: LoadingState) {
        if let Some(state) = self.by_key.get_mut(key) {
            state.loading = loading;
            if loading == LoadingState::Bailed {
                state.usage = None;
                state.limit_reading = None;
            }
        }
    }

    /// Keep the windows a failed probe still carried. It never touches
    /// `loading` or `last_error`: a rate-limited probe does not heal the
    /// account, and nothing selection or a cooldown reads may move.
    pub fn set_limit_reading(&mut self, key: &AccountKey, snapshot: UsageSnapshot) {
        if let Some(state) = self.by_key.get_mut(key) {
            state.limit_reading = Some(snapshot);
        }
    }

    /// The newest reading the account gave, for a view's bars: a
    /// rate-limited response's windows when one landed after the last
    /// successful probe, else that probe's snapshot. Absent while both are.
    pub fn latest_reading(&self, key: &AccountKey) -> Option<&UsageSnapshot> {
        let state = self.by_key.get(key)?;
        state.limit_reading.as_ref().or(state.usage.as_ref())
    }

    /// Disarm the scheduler hook for `key` after firing an override
    /// probe attempt. One-shot semantics: the hook fires once per
    /// arming (i.e. once per fresh snapshot). Subsequent stale-reset
    /// state must wait for the next successful probe (which re-arms
    /// via `set_usage`) or until the existing backoff timer
    /// (`next_probe_at`) elapses.
    pub fn disarm_override(&mut self, key: &AccountKey) {
        if let Some(state) = self.by_key.get_mut(key) {
            state.override_armed = false;
        }
    }

    /// Record the latest poll-attempt failure for `key` and schedule
    /// the next probe.
    ///
    /// `retry_after` (when `Some`) is the server-provided hold-down
    /// duration - typically Anthropic's `Retry-After` header on 429.
    /// We honour it verbatim because the server knows when its
    /// per-account bucket will reset; guessing with our own backoff
    /// either over- or under-shoots and keeps the limit hot. When
    /// `None` (network failures, unknown HTTP status, server didn't
    /// send Retry-After), fall back to a local exponential schedule.
    ///
    /// The cached `usage` snapshot is preserved so the panel keeps
    /// showing the last known-good bars (a fresh 429 doesn't blank
    /// them out).
    pub fn set_last_error(
        &mut self,
        key: &AccountKey,
        status: UsageFetchStatus,
        retry_after: Option<std::time::Duration>,
    ) {
        if let Some(state) = self.by_key.get_mut(key) {
            state.last_error = Some(status);
            state.consecutive_failures = state.consecutive_failures.saturating_add(1);
            // Auth-recovery cache-clear: an Unauthorized or Expired
            // probe response means the account's credential is dead.
            // Transition `loading` to `Bailed` and drop the
            // cached `usage` so no stale bar is drawn and the account
            // reads bailed. The account recovers via the
            // 60 s usage poller, which reads the env the config now
            // holds. Other statuses leave
            // `loading` alone (a transient `RateLimited` or
            // `NetworkFailed` is not auth-related; the cache stays
            // and the account remains Ready for the walk).
            // Replaces the PR #238 `consecutive_unauthorized` 3-strike
            // counter.
            if matches!(status, UsageFetchStatus::Unauthorized | UsageFetchStatus::Expired) {
                let prev = state.loading;
                state.loading = LoadingState::Bailed;
                state.usage = None;
                state.limit_reading = None;
                // Surface the transition so an operator triaging
                // "account suddenly stopped working" can see WHEN the
                // bail happened + which probe-class triggered it.
                // Without this log a Ready -> Bailed flip is silent
                // until a poller runs again.
                if prev != LoadingState::Bailed {
                    tracing::warn!(
                        target: "forge_workspace::account",
                        event_name = "account_bailed",
                        account = %key.0,
                        prev_state = ?prev,
                        status = ?status,
                        "account transitioned to Bailed via probe failure; the pollers will retry it",
                    );
                }
            }
            // Anthropic returns `Retry-After: 0` on the /api/oauth/usage
            // 429 path, which is "no specific hint" rather than "you can
            // retry now". Trusting it literally schedules next_probe_at
            // at the same instant and the next round re-trips the same
            // rate limit - that's the burst-probe behaviour the warm
            // loop kept hitting. Treat any sub-second hint as missing
            // and fall through to the exponential default
            // (starts at 30 s), so the leaky-bucket actually gets
            // a chance to refill.
            let delay = match retry_after {
                Some(d) if d >= std::time::Duration::from_secs(1) => d.min(BACKOFF_CAP),
                _ => backoff_delay(state.consecutive_failures),
            };
            state.next_probe_at = Some(std::time::Instant::now() + delay);
        }
    }

    /// `true` when the poller may probe `key` now. `false` when the
    /// account is in an active backoff window from a recent failure.
    /// Cold-cache accounts (no last_error) always return `true`.
    pub fn should_probe_now(&self, key: &AccountKey) -> bool {
        let Some(state) = self.by_key.get(key) else { return true };
        state.next_probe_at.is_none_or(|t| t <= std::time::Instant::now())
    }

    /// `true` when every account in the map has reached a terminal
    /// `LoadingState` (`Ready` or `Bailed`). The boot wave and the cron
    /// spawn path gate on this. Empty maps return `true` (vacuous; no
    /// accounts means no work to wait on - relevant in the testing-stub
    /// path). A retired account counts as settled: the edit removed it, so
    /// nothing will ever load it.
    pub fn all_loaded(&self) -> bool {
        self.by_key
            .values()
            .all(|s| s.retired || matches!(s.loading, LoadingState::Ready | LoadingState::Bailed))
    }

    /// True when `key`'s loaded usage snapshot shows it at-or-beyond
    /// the plan cap - the same saturation signal `unusable_reason`
    /// classifies on. A Ready-but-saturated account logs in
    /// fine but trips the rate limit on its next request, so the
    /// selection walk prefers other accounts when one is available.
    pub fn is_saturated(&self, key: &AccountKey) -> bool {
        self.by_key.get(key).and_then(|s| s.usage.as_ref()).is_some_and(is_rate_limited)
    }

    /// Why `key` is out of the usable tier, `None` when pickable: what
    /// the walk reads to demote an account. `Bailed`
    /// outranks the probe classes: it is the loading task's terminal
    /// verdict, and a bailed account has no cached windows to saturate
    /// with. A snapshot whose source is not the account backend's is
    /// treated as absent, the same refusal `budget()` applies - the
    /// redb cache survives a provider edit and is re-seeded at boot,
    /// and without this a stale windowed snapshot would classify a
    /// spend row as saturated.
    pub fn unusable_reason(&self, key: &AccountKey) -> Option<Unusable> {
        if self.loading_state(key) == LoadingState::Bailed {
            return Some(Unusable::Bailed);
        }
        let usage = self
            .provider(key)
            .and_then(crate::backend)
            .and_then(|backend| self.usage(key).filter(|s| s.source == backend.source()));
        unusable_reason(usage, self.usage_error(key))
    }

    /// Snapshot the current `LoadingState` for `key`. Returns
    /// `LoadingState::Loading` by default for unknown keys
    /// (defensive - a caller may hold a key that hasn't been
    /// registered in the map).
    pub fn loading_state(&self, key: &AccountKey) -> LoadingState {
        self.by_key.get(key).map_or(LoadingState::Loading, |s| s.loading)
    }

    /// Combined scheduler signal: probe `key` on the current cycle if
    /// the normal backoff gate (`should_probe_now`) is open OR the
    /// one-shot reset-clear override hook (`has_just_cleared_cap_window`)
    /// is firing. Both sides of the OR carry through to the workspace
    /// poller, which calls this method to decide which accounts to
    /// probe each tick. Centralising the OR here keeps a future
    /// refactor from accidentally dropping one half of the signal -
    /// the unit test pins the truth table that the scheduler relies
    /// on, regardless of whether the call site spells out both
    /// predicates inline.
    pub fn scheduler_should_probe(&self, key: &AccountKey) -> bool {
        self.should_probe_now(key) || self.has_just_cleared_cap_window(key)
    }

    /// `true` when at least one cached window shows the account just
    /// transitioned out of its cap (utilization at-or-above 100%,
    /// `resets_at` now in the past) AND the override hook is armed.
    /// The scheduler ORs this with `should_probe_now` so a fresh probe
    /// lands on the next poll cycle after the reset moment, instead of
    /// waiting through the remainder of an active backoff window.
    /// Without the hook, an account 429'd with a multi-hour
    /// `Retry-After` keeps painting the stale "100%" bar for hours
    /// past the actual reset because the next probe is gated until the
    /// backoff timer elapses.
    ///
    /// The `override_armed` gate enforces one-shot semantics: each
    /// successful probe arms the hook exactly once, and the scheduler
    /// disarms it after firing. A persistently failing probe series
    /// won't keep tripping the override every cycle.
    pub fn has_just_cleared_cap_window(&self, key: &AccountKey) -> bool {
        let Some(state) = self.by_key.get(key) else {
            return false;
        };
        if !state.override_armed {
            return false;
        }
        let Some(usage) = state.usage.as_ref() else {
            return false;
        };
        let now = std::time::SystemTime::now();
        let windows = [
            usage.five_hour.as_ref(),
            usage.seven_day.as_ref(),
            usage.seven_day_opus.as_ref(),
            usage.seven_day_sonnet.as_ref(),
        ];
        windows
            .into_iter()
            .flatten()
            .any(|w| w.utilization >= 100.0 && w.resets_at.is_some_and(|when| when <= now))
    }

    /// Look up the cached usage snapshot for `key`. `None` when the
    /// poller hasn't yet succeeded for this account.
    pub fn usage(&self, key: &AccountKey) -> Option<&UsageSnapshot> {
        self.by_key.get(key).and_then(|s| s.usage.as_ref())
    }

    /// Look up the latest poll-attempt failure for `key`. `None`
    /// when the most recent attempt succeeded (or no attempt has
    /// been made yet).
    pub fn usage_error(&self, key: &AccountKey) -> Option<UsageFetchStatus> {
        self.by_key.get(key).and_then(|s| s.last_error)
    }

    /// Drop the recorded failure for `key`, so a bail that never
    /// recorded its own cause falls back to the unrecorded default
    /// rather than wearing a stale class.
    pub fn clear_last_error(&mut self, key: &AccountKey) {
        if let Some(state) = self.by_key.get_mut(key) {
            state.last_error = None;
        }
    }
}

/// The classification behind [`AccountStateMap::unusable_reason`] for a
/// key already resolved to its inputs. `None` means usable: usage
/// unknown, or under 100% on every window, or the last probe failed
/// transiently (network / unknown HTTP - we just don't know yet).
/// `Some` means at the cap, or unable to authenticate.
///
/// Saturation outranks a failed probe: the windows came from a
/// successful read, and on an at-cap account the failed probe usually
/// IS the cap throttling the endpoint.
fn unusable_reason(
    usage: Option<&UsageSnapshot>,
    last_error: Option<UsageFetchStatus>,
) -> Option<Unusable> {
    if usage.is_some_and(is_rate_limited) {
        return Some(Unusable::Saturated);
    }
    if matches!(
        last_error,
        Some(
            UsageFetchStatus::RateLimited
                | UsageFetchStatus::Expired
                | UsageFetchStatus::Unauthorized
        )
    ) {
        return Some(Unusable::ProbeBlocked);
    }
    None
}

/// True when ANY window in the snapshot is currently at-or-beyond
/// the plan cap AND its scheduled reset has not yet passed. Such an
/// account will immediately trip the API rate limit on the next
/// request, so it's excluded from the "available" tier and only used
/// as a fallback when every pinned account is rate-limited.
///
/// Each window's predicate (`UsageWindow::is_currently_limited`) gates
/// on `resets_at` so a stale cached snapshot from before the window
/// reset doesn't keep the account permanently classified as
/// rate-limited. The pair "utilization >= 100% AND resets_at > now"
/// is what makes the classification self-clearing across the reset
/// boundary - earlier shapes that checked utilization alone needed
/// a separate cache-invalidation pathway, which has churned through
/// several abandoned designs.
fn is_rate_limited(snapshot: &UsageSnapshot) -> bool {
    let windows = [
        snapshot.five_hour.as_ref(),
        snapshot.seven_day.as_ref(),
        snapshot.seven_day_opus.as_ref(),
        snapshot.seven_day_sonnet.as_ref(),
    ];
    windows.into_iter().flatten().any(UsageWindow::is_currently_limited)
}

/// Ceiling shared by the exponential schedule and the server
/// `Retry-After` clamp. The hold-down this caps is the sole re-probe
/// path for a settled account, so a bogus header must not hold one
/// down arbitrarily long.
const BACKOFF_CAP: std::time::Duration = std::time::Duration::from_secs(600);

/// Per-account exponential backoff schedule for usage-probe
/// failures. Doubles each consecutive failure, capped at 10
/// minutes. The 60 s poll loop ticks ~once a minute; without
/// backoff, a transient per-IP /usage 429 would persist across
/// every cycle for as long as Anthropic kept rate-limiting,
/// preventing any account from accumulating fresh usage data.
///
/// | consecutive | delay        |
/// |-------------|--------------|
/// | 1           | 30 s         |
/// | 2           | 1 min        |
/// | 3           | 2 min        |
/// | 4           | 4 min        |
/// | 5           | 8 min        |
/// | 6+          | 10 min (cap) |
fn backoff_delay(consecutive_failures: u32) -> std::time::Duration {
    use std::time::Duration;
    let exp = consecutive_failures.min(20); // shift saturates at 20; well under u64::MAX
    let seconds = 30_u64.saturating_mul(1_u64 << exp.saturating_sub(1));
    Duration::from_secs(seconds).min(BACKOFF_CAP)
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
    use super::*;
    use forge_primitives::usage::{UsageSnapshot, UsageSourceKind, UsageWindow};
    use std::time::SystemTime;

    fn make_account(name: &str) -> LoadedAccount {
        LoadedAccount {
            display_name: name.to_owned(),
            provider: forge_primitives::account::Provider::Anthropic,
            base_url: None,
            models: vec!["claude-sonnet-5".to_owned()],
            model_aliases: std::collections::HashMap::new(),
            model_slugs: std::collections::HashMap::new(),
            env: std::collections::HashMap::new(),
        }
    }

    /// `AccountStateMap::new` copies `LoadedAccount.env` verbatim - the
    /// live credential included - so neither the state nor the map may
    /// print a value from it.
    #[test]
    fn a_credential_never_reaches_the_account_state_debug_output() {
        let mut account = make_account("Stargate");
        account.env.insert("CLAUDE_CODE_OAUTH_TOKEN".to_owned(), "gateway-supersecret".to_owned());
        let map = AccountStateMap::new(&[account]);
        let state = map.by_key.values().next().expect("the state is there");
        let rendered = format!("{state:?}");
        assert!(
            !rendered.contains("gateway-supersecret"),
            "the credential stamp leaked: {rendered}",
        );
        assert!(rendered.contains("Anthropic"), "the provider stays readable: {rendered}");
        let rendered = format!("{map:?}");
        assert!(
            !rendered.contains("gateway-supersecret"),
            "the credential stamp leaked through the map: {rendered}",
        );
        assert!(rendered.contains("Stargate"), "the account key stays readable: {rendered}");
    }

    /// Default snapshot fixture used by the account-selection tests. The
    /// `resets_at` value matters: under the resets_at-driven predicate,
    /// a window at 100% utilization with `resets_at = None` is NOT
    /// classified as rate-limited (None means "no future reset to clear
    /// this", which would strand the account forever). Real probes
    /// always emit a future `resets_at` alongside the percentage, so
    /// the fixture mirrors that by defaulting to a reset 1 hour out.
    /// The few tests that need an EXPIRED reset (the staleness case)
    /// build their snapshot inline rather than through this helper.
    fn snapshot(five_hour: Option<f64>, seven_day: Option<f64>) -> UsageSnapshot {
        let future = SystemTime::now() + std::time::Duration::from_secs(3600);
        UsageSnapshot {
            source: UsageSourceKind::Oauth,
            fetched_at: SystemTime::UNIX_EPOCH,
            five_hour: five_hour.map(|util| UsageWindow {
                utilization: util,
                resets_at: Some(future),
                reset_description: None,
                status: None,
            }),
            seven_day: seven_day.map(|util| UsageWindow {
                utilization: util,
                resets_at: Some(future),
                reset_description: None,
                status: None,
            }),
            seven_day_opus: None,
            seven_day_sonnet: None,
            extra_usage: None,
            spend: None,
            balance: None,
        }
    }

    #[test]
    fn account_serves_matches_declared_names_and_their_aliases() {
        let models = vec!["claude-opus-5[1m]".to_owned()];
        let aliases = std::collections::HashMap::from([(
            "claude-opus-5[1m]".to_owned(),
            vec!["claude-opus-5".to_owned()],
        )]);
        assert!(
            account_serves(&models, &aliases, "claude-opus-5[1m]"),
            "a declared name is served",
        );
        assert!(account_serves(&models, &aliases, "claude-opus-5"), "a declared alias is served");
        assert!(
            !account_serves(&models, &aliases, "claude-sonnet-5"),
            "a name that is neither declared nor an alias is not served",
        );
        assert!(
            !account_serves(&models, &std::collections::HashMap::new(), "claude-opus-5"),
            "with no aliases declared the account serves only its declared names",
        );
    }

    #[test]
    fn an_aliased_name_reaches_the_route_gate_through_the_pool() {
        let mut granite = make_account("Granite");
        granite.models = vec!["claude-opus-5[1m]".to_owned()];
        granite
            .model_aliases
            .insert("claude-opus-5[1m]".to_owned(), vec!["claude-opus-5".to_owned()]);
        let pool = crate::AccountPool::new(&[granite]);
        let key = AccountKey("Granite".to_owned());
        assert!(
            pool.declares(&key, "claude-opus-5"),
            "the route gate reads the alias the config declared",
        );
        assert!(
            !pool.declares(&key, "claude-sonnet-5"),
            "the route gate still refuses a model the account does not have",
        );
    }

    #[test]
    fn env_accessor_returns_account_env() {
        let mut codex = make_account("Codex");
        codex.env.insert("ANTHROPIC_BASE_URL".to_owned(), "http://localhost:18765".to_owned());
        let map = AccountStateMap::new(&[codex, make_account("Plain")]);
        assert_eq!(
            map.env(&AccountKey("Codex".to_owned()))
                .and_then(|e| e.get("ANTHROPIC_BASE_URL"))
                .map(String::as_str),
            Some("http://localhost:18765"),
        );
        assert!(
            map.env(&AccountKey("Plain".to_owned()))
                .is_some_and(std::collections::HashMap::is_empty),
            "account with no env -> empty map",
        );
        assert!(map.env(&AccountKey("Unknown".to_owned())).is_none(), "unknown key -> None");
    }

    /// The repair-route table every auth hint branches on. Base-url
    /// wins: an account whose provider owns a base url re-keys its env
    /// token, whatever else its env carries. A token-less Anthropic
    /// account classifies Token - the repair it needs is the same
    /// setup-token copy.
    #[test]
    fn auth_classifies_the_repair_route_per_account() {
        let mut token = make_account("Token");
        token.env.insert("CLAUDE_CODE_OAUTH_TOKEN".to_owned(), "setup-token".to_owned());
        let mut base_url = make_account("Base");
        base_url.provider = forge_primitives::account::Provider::Openrouter;
        base_url.env.insert("CLAUDE_CODE_OAUTH_TOKEN".to_owned(), "setup-token".to_owned());
        let map = AccountStateMap::new(&[token, base_url, make_account("Tokenless")]);
        assert_eq!(
            map.auth(&AccountKey("Token".to_owned())),
            Some(forge_primitives::account::AccountAuth::Token),
        );
        assert_eq!(
            map.auth(&AccountKey("Base".to_owned())),
            Some(forge_primitives::account::AccountAuth::BaseUrl),
            "a base-url provider classifies BaseUrl even beside a setup token",
        );
        assert_eq!(
            map.auth(&AccountKey("Tokenless".to_owned())),
            Some(forge_primitives::account::AccountAuth::Token),
            "a token-less anthropic account repairs as a token account",
        );
        assert_eq!(map.auth(&AccountKey("Unknown".to_owned())), None, "unknown key -> None");
    }

    #[test]
    fn unusable_reason_names_why_the_account_is_out() {
        let mut map = AccountStateMap::new(&[
            make_account("saturated"),
            make_account("probe-blocked"),
            make_account("bailed"),
            make_account("both"),
        ]);
        map.set_usage(&AccountKey("saturated".to_owned()), snapshot(Some(100.0), None));
        // RateLimited leaves loading alone, so this isolates the probe
        // class; an Unauthorized probe would flip loading to Bailed and
        // read as Bailed here instead.
        map.set_last_error(
            &AccountKey("probe-blocked".to_owned()),
            UsageFetchStatus::RateLimited,
            None,
        );
        map.set_loading(&AccountKey("bailed".to_owned()), LoadingState::Bailed);
        // Saturation outranks a failed probe: the windows came from a
        // successful read.
        map.set_usage(&AccountKey("both".to_owned()), snapshot(Some(100.0), None));
        map.set_last_error(&AccountKey("both".to_owned()), UsageFetchStatus::RateLimited, None);

        assert_eq!(map.unusable_reason(&key("saturated")), Some(Unusable::Saturated));
        assert_eq!(map.unusable_reason(&key("probe-blocked")), Some(Unusable::ProbeBlocked));
        assert_eq!(map.unusable_reason(&key("bailed")), Some(Unusable::Bailed));
        assert_eq!(map.unusable_reason(&key("both")), Some(Unusable::Saturated));
    }

    /// The redb cache survives a `forge.toml` provider edit, so a
    /// spend-billed account can boot holding a windowed snapshot under
    /// another backend's source. That snapshot must not classify
    /// Saturated: spend rows have no window to hit, and the tag would
    /// assert a spent allowance the account does not have.
    #[test]
    fn a_source_mismatched_snapshot_cannot_saturate_a_spend_row() {
        let mut router = make_account("Router");
        router.provider = forge_primitives::account::Provider::Openrouter;
        let mut map = AccountStateMap::new(&[router]);
        // The snapshot helper stamps Oauth; the Openrouter backend's
        // source is OpenRouterKey, so the cached windows are refused.
        map.set_usage(&key("Router"), snapshot(Some(100.0), Some(100.0)));
        assert_eq!(
            map.unusable_reason(&key("Router")),
            None,
            "stale cache alone must not tag the spend row at all",
        );

        // With the probe also failing, the row wears the probe class -
        // never the saturated one.
        map.set_last_error(&key("Router"), UsageFetchStatus::RateLimited, None);
        assert_eq!(
            map.unusable_reason(&key("Router")),
            Some(Unusable::ProbeBlocked),
            "the stale snapshot falls through to the probe class",
        );
    }

    #[test]
    fn backoff_delay_doubles_then_caps_at_10_minutes() {
        use std::time::Duration;
        assert_eq!(backoff_delay(1), Duration::from_secs(30));
        assert_eq!(backoff_delay(2), Duration::from_secs(60));
        assert_eq!(backoff_delay(3), Duration::from_secs(120));
        assert_eq!(backoff_delay(4), Duration::from_secs(240));
        assert_eq!(backoff_delay(5), Duration::from_secs(480));
        assert_eq!(backoff_delay(6), Duration::from_secs(600), "capped at 10 min");
        assert_eq!(backoff_delay(20), Duration::from_secs(600), "still capped");
    }

    #[test]
    fn should_probe_now_true_for_cold_cache() {
        let map = AccountStateMap::new(&[make_account("Gateway")]);
        assert!(map.should_probe_now(&AccountKey("Gateway".to_owned())));
    }

    #[test]
    fn set_last_error_schedules_next_probe_in_future() {
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let key = AccountKey("Gateway".to_owned());
        map.set_last_error(&key, UsageFetchStatus::RateLimited, None);
        assert!(!map.should_probe_now(&key), "first failure puts account in backoff");
    }

    #[test]
    fn retry_after_overrides_exponential_backoff() {
        // The server-provided retry_after must win over the local
        // exponential schedule - 30 s for a first failure would
        // under-shoot the actual reset and re-trip the limit.
        use std::time::Duration;
        let mut map = AccountStateMap::new(&[make_account("Gateway1")]);
        let key = AccountKey("Gateway1".to_owned());
        let t0 = std::time::Instant::now();
        map.set_last_error(&key, UsageFetchStatus::RateLimited, Some(Duration::from_secs(480)));
        let next = map.by_key.get(&key).and_then(|s| s.next_probe_at).expect("scheduled");
        let gap = next.saturating_duration_since(t0);
        assert!(
            gap >= Duration::from_secs(479) && gap <= Duration::from_secs(481),
            "next_probe_at ≈ now + 480s; got {gap:?}",
        );
    }

    #[test]
    fn retry_after_is_clamped_to_the_backoff_ceiling() {
        // The hold-down is the sole re-probe path for a settled
        // account, so a bogus (or merely very large) header must not
        // hold one down past the schedule's own 10-minute ceiling.
        use std::time::Duration;
        let mut map = AccountStateMap::new(&[make_account("Gateway1")]);
        let key = AccountKey("Gateway1".to_owned());
        let t0 = std::time::Instant::now();
        map.set_last_error(&key, UsageFetchStatus::RateLimited, Some(Duration::from_secs(100_000)));
        let next = map.by_key.get(&key).and_then(|s| s.next_probe_at).expect("scheduled");
        let gap = next.saturating_duration_since(t0);
        assert!(
            gap >= Duration::from_secs(599) && gap <= Duration::from_secs(601),
            "next_probe_at clamps to ≈ 600s; got {gap:?}",
        );
    }

    #[test]
    fn set_usage_clears_backoff_so_account_probes_again() {
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let key = AccountKey("Gateway".to_owned());
        map.set_last_error(&key, UsageFetchStatus::RateLimited, None);
        assert!(!map.should_probe_now(&key));
        map.set_usage(&key, snapshot(Some(10.0), Some(10.0)));
        assert!(map.should_probe_now(&key), "successful probe clears the backoff gate");
    }

    #[test]
    fn seven_day_util_takes_max_across_windows() {
        // seven_day = 30, seven_day_opus = 80 → binding = max = 80
        let mut s = snapshot(Some(20.0), Some(30.0));
        s.seven_day_opus = Some(UsageWindow {
            utilization: 80.0,
            resets_at: None,
            reset_description: None,
            status: None,
        });
        assert_eq!(s.seven_day_util(), Some(80.0), "the most-used 7d window is the binding one");
    }

    /// An empty fold produced 0.0 with nothing that looked like a
    /// default, so a 200 carrying only the session window reported the
    /// account as 0% used across seven days.
    #[test]
    fn seven_day_util_is_none_when_every_window_is_absent() {
        let mut s = snapshot(Some(20.0), None);
        s.seven_day_opus = None;
        s.seven_day_sonnet = None;
        assert_eq!(s.seven_day_util(), None, "no 7d window means no 7d reading");
        assert_eq!(s.five_hour_util(), Some(20.0), "the window that is present still reports");
    }

    #[test]
    fn is_rate_limited_fires_on_either_window() {
        assert!(is_rate_limited(&snapshot(Some(100.0), Some(0.0))));
        assert!(is_rate_limited(&snapshot(Some(0.0), Some(100.0))));
        assert!(is_rate_limited(&snapshot(Some(100.0), Some(100.0))));
        assert!(!is_rate_limited(&snapshot(Some(99.9), Some(99.9))));
    }

    /// Build a snapshot with explicit `resets_at` values per window. The
    /// public `snapshot()` helper hides this for the common test case
    /// (resets in the future); this variant lets the
    /// staleness/transition tests drive the resets_at clock.
    fn snapshot_with_resets(
        five_hour: Option<(f64, Option<SystemTime>)>,
        seven_day: Option<(f64, Option<SystemTime>)>,
    ) -> UsageSnapshot {
        UsageSnapshot {
            source: UsageSourceKind::Oauth,
            fetched_at: SystemTime::UNIX_EPOCH,
            five_hour: five_hour.map(|(util, resets_at)| UsageWindow {
                utilization: util,
                resets_at,
                reset_description: None,
                status: None,
            }),
            seven_day: seven_day.map(|(util, resets_at)| UsageWindow {
                utilization: util,
                resets_at,
                reset_description: None,
                status: None,
            }),
            seven_day_opus: None,
            seven_day_sonnet: None,
            extra_usage: None,
            spend: None,
            balance: None,
        }
    }

    #[test]
    fn is_rate_limited_false_when_cached_window_reset_already_passed() {
        // Stale cached snapshot: probe reported 100% an hour ago but
        // the reset moment has come and gone. Predicate must return
        // false so the rotation re-considers the account WITHOUT
        // waiting for a fresh probe to overwrite the cache.
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let snap = snapshot_with_resets(Some((100.0, Some(past))), Some((50.0, Some(past))));
        assert!(!is_rate_limited(&snap), "stale 100% reading must not strand the account");
    }

    #[test]
    fn is_rate_limited_true_when_only_one_window_is_capped_and_still_in_window() {
        // 5h at 100% + still in window; 7d below cap. Either window
        // can keep the account out of the pool while it's in its
        // hold-down period.
        let future = SystemTime::now() + std::time::Duration::from_secs(60);
        let snap = snapshot_with_resets(Some((100.0, Some(future))), Some((30.0, Some(future))));
        assert!(is_rate_limited(&snap));
    }

    /// Build a snapshot exercising the opus-specific 7-day window.
    /// Max-plan accounts hit the opus cap independently of the
    /// shared 7-day envelope, and `is_rate_limited` must walk every
    /// window in the four-tuple - a future short-circuit on
    /// `seven_day` would silently break Max-plan classification.
    fn opus_only_capped_snapshot(util: f64, resets_at: Option<SystemTime>) -> UsageSnapshot {
        UsageSnapshot {
            source: UsageSourceKind::Oauth,
            fetched_at: SystemTime::UNIX_EPOCH,
            five_hour: Some(UsageWindow {
                utilization: 30.0,
                resets_at: Some(SystemTime::now() + std::time::Duration::from_secs(3600)),
                reset_description: None,
                status: None,
            }),
            seven_day: Some(UsageWindow {
                utilization: 30.0,
                resets_at: Some(SystemTime::now() + std::time::Duration::from_secs(3600)),
                reset_description: None,
                status: None,
            }),
            seven_day_opus: Some(UsageWindow {
                utilization: util,
                resets_at,
                reset_description: None,
                status: None,
            }),
            seven_day_sonnet: None,
            extra_usage: None,
            spend: None,
            balance: None,
        }
    }

    /// Sister of `opus_only_capped_snapshot` for the sonnet 7-day
    /// window. Same rationale: a sonnet-only cap must bubble through
    /// `is_rate_limited` even when no other window is at the limit.
    fn sonnet_only_capped_snapshot(util: f64, resets_at: Option<SystemTime>) -> UsageSnapshot {
        UsageSnapshot {
            source: UsageSourceKind::Oauth,
            fetched_at: SystemTime::UNIX_EPOCH,
            five_hour: Some(UsageWindow {
                utilization: 30.0,
                resets_at: Some(SystemTime::now() + std::time::Duration::from_secs(3600)),
                reset_description: None,
                status: None,
            }),
            seven_day: Some(UsageWindow {
                utilization: 30.0,
                resets_at: Some(SystemTime::now() + std::time::Duration::from_secs(3600)),
                reset_description: None,
                status: None,
            }),
            seven_day_opus: None,
            seven_day_sonnet: Some(UsageWindow {
                utilization: util,
                resets_at,
                reset_description: None,
                status: None,
            }),
            extra_usage: None,
            spend: None,
            balance: None,
        }
    }

    #[test]
    fn is_rate_limited_fires_on_opus_only_capped_window() {
        // Max-plan reality: shared 5h + 7d windows healthy, opus-7d
        // at cap. The selection walk must classify as rate-limited or
        // the account keeps getting selected and trips API 429s.
        let future = SystemTime::now() + std::time::Duration::from_secs(60);
        let snap = opus_only_capped_snapshot(100.0, Some(future));
        assert!(is_rate_limited(&snap), "opus 7d cap (alone) must be detected");
    }

    #[test]
    fn is_rate_limited_fires_on_sonnet_only_capped_window() {
        // Symmetry with opus: sonnet-7d at cap must also bubble up.
        let future = SystemTime::now() + std::time::Duration::from_secs(60);
        let snap = sonnet_only_capped_snapshot(100.0, Some(future));
        assert!(is_rate_limited(&snap), "sonnet 7d cap (alone) must be detected");
    }

    #[test]
    fn is_rate_limited_false_when_opus_capped_but_reset_passed() {
        // Stale opus-window snapshot - resets_at has come and gone.
        // Must classify as NOT rate-limited so the selection walk
        // reconsiders the account; the same self-clearing applies
        // per-window.
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let snap = opus_only_capped_snapshot(100.0, Some(past));
        assert!(!is_rate_limited(&snap), "stale opus window must self-clear like any other");
    }

    #[test]
    fn is_rate_limited_false_when_window_at_cap_but_no_reset_scheduled() {
        // None means "no scheduled reset" - we cannot prove the limit
        // is still in effect. Treating it as limited would forever
        // strand the account on a single bad cached reading.
        let snap = snapshot_with_resets(Some((100.0, None)), Some((100.0, None)));
        assert!(!is_rate_limited(&snap), "None resets_at must not classify as limited");
    }

    #[test]
    fn has_just_cleared_cap_window_true_when_cached_window_reset_passed() {
        // Cache snapshot has util 100% with resets_at in the past:
        // the scheduler hook must fire so a fresh probe overwrites
        // the stale bar.
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let stale = snapshot_with_resets(Some((100.0, Some(past))), Some((30.0, Some(past))));
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, stale);
        assert!(map.has_just_cleared_cap_window(&k));
    }

    #[test]
    fn has_just_cleared_cap_window_false_when_window_still_in_cap() {
        // Cap window still active: hook must not fire (the existing
        // `should_probe_now` gate covers normal cadence).
        let future = SystemTime::now() + std::time::Duration::from_secs(60);
        let live = snapshot_with_resets(Some((100.0, Some(future))), Some((30.0, Some(future))));
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, live);
        assert!(!map.has_just_cleared_cap_window(&k));
    }

    #[test]
    fn has_just_cleared_cap_window_false_when_no_snapshot_cached() {
        // Cold cache - nothing to compare against. Hook returns false;
        // the normal cold-cache `should_probe_now == true` already
        // schedules the first probe.
        let map = AccountStateMap::new(&[make_account("Gateway")]);
        assert!(!map.has_just_cleared_cap_window(&AccountKey("Gateway".to_owned())));
    }

    #[test]
    fn has_just_cleared_cap_window_false_when_below_cap() {
        // Below-cap snapshot with a past resets_at: NOT a freshly
        // cleared cap, just a stale-but-irrelevant timestamp. The
        // hook only cares about the at-cap transition.
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let snap = snapshot_with_resets(Some((50.0, Some(past))), Some((30.0, Some(past))));
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, snap);
        assert!(!map.has_just_cleared_cap_window(&k));
    }

    #[test]
    fn override_disarms_after_firing_and_does_not_refire_until_rearmed() {
        // Setup mirrors the production sustained-failure shape: a
        // stale cached snapshot (util=100, resets_at in the past)
        // plus an active backoff timer from a recent 429. Without
        // the disarm gate, every poll cycle would re-fire the hook
        // and hammer Anthropic through the entire backoff window.
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let stale = snapshot_with_resets(Some((100.0, Some(past))), None);
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, stale);
        // Push next_probe_at into the future so should_probe_now is
        // false: this is the "in backoff" state the scheduler hook is
        // meant to override.
        map.set_last_error(
            &k,
            UsageFetchStatus::RateLimited,
            Some(std::time::Duration::from_secs(3000)),
        );
        assert!(!map.should_probe_now(&k), "backoff active");
        assert!(map.has_just_cleared_cap_window(&k), "hook fires while armed");

        // Scheduler picks up the account via the hook + fires its
        // override probe; disarm runs.
        map.disarm_override(&k);
        assert!(!map.has_just_cleared_cap_window(&k), "disarmed; no second override this cycle");
        assert!(!map.should_probe_now(&k), "backoff still active");
        // Combined scheduler signal: neither predicate fires - the
        // account waits for the backoff timer to elapse.
        assert!(
            !map.should_probe_now(&k) && !map.has_just_cleared_cap_window(&k),
            "scheduler must respect backoff once the override has been consumed",
        );
    }

    #[test]
    fn fresh_set_usage_rearms_override_for_next_reset_boundary() {
        // After the override fires + disarms once, a NEW successful
        // probe (set_usage) must re-arm the hook so future stale-reset
        // boundaries can still trigger their one-shot override.
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let stale = snapshot_with_resets(Some((100.0, Some(past))), None);
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, stale.clone());
        map.disarm_override(&k);
        assert!(!map.has_just_cleared_cap_window(&k), "disarmed");

        // A fresh probe lands - even if the snapshot is STILL
        // stale-reset, the hook re-arms because set_usage represents
        // a brand-new probe attempt's success.
        map.set_usage(&k, stale);
        assert!(map.has_just_cleared_cap_window(&k), "fresh set_usage re-arms the hook");
    }

    #[test]
    fn override_not_armed_on_construction_so_hook_skips_cold_accounts() {
        // Without an arming step, a brand-new AccountStateMap entry
        // must NOT fire the hook even if some hypothetical stale
        // snapshot were planted manually. Cold accounts are picked up
        // via cold-cache `should_probe_now == true`, not the override.
        let map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        assert!(!map.has_just_cleared_cap_window(&k), "no arm = no override");
    }

    // -------------------------------------------------------------
    // Scheduler OR truth table - pins the combined signal the
    // workspace poller relies on. A future refactor that swaps the
    // OR for AND or drops one half loses the behavior with CI green
    // unless this table covers it.
    // -------------------------------------------------------------

    #[test]
    fn scheduler_should_probe_fires_via_should_probe_now_alone() {
        // Cold-cache account: no failures yet, no snapshot. The normal
        // `should_probe_now` returns true (next_probe_at is None);
        // the hook returns false (no snapshot to override against).
        // OR result: true.
        let map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        assert!(map.should_probe_now(&k));
        assert!(!map.has_just_cleared_cap_window(&k));
        assert!(map.scheduler_should_probe(&k), "OR true via should_probe_now");
    }

    #[test]
    fn scheduler_should_probe_fires_via_hook_alone_during_active_backoff() {
        // The signature override scenario from the planner: a stale
        // cached snapshot (util=100, resets_at in the past) paired
        // with an active backoff timer from a recent 429. The hook
        // must be the deciding factor:
        // - should_probe_now == false (backoff active)
        // - has_just_cleared_cap_window == true (armed + stale reset)
        // - scheduler_should_probe == true (OR fires)
        let past = SystemTime::now() - std::time::Duration::from_secs(60);
        let stale = snapshot_with_resets(Some((100.0, Some(past))), None);
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, stale);
        map.set_last_error(
            &k,
            UsageFetchStatus::RateLimited,
            Some(std::time::Duration::from_secs(3000)),
        );

        assert!(!map.should_probe_now(&k), "backoff active");
        assert!(map.has_just_cleared_cap_window(&k), "armed + stale reset");
        assert!(map.scheduler_should_probe(&k), "OR fires when hook alone is true");
    }

    #[test]
    fn scheduler_should_probe_skips_when_both_predicates_false() {
        // Backoff active AND the snapshot has a future resets_at (no
        // stale state). Both sides of the OR return false; the
        // account is skipped this cycle and waits for the backoff
        // timer to elapse.
        let future = SystemTime::now() + std::time::Duration::from_secs(60);
        let live = snapshot_with_resets(Some((100.0, Some(future))), None);
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = AccountKey("Gateway".to_owned());
        map.set_usage(&k, live);
        map.set_last_error(
            &k,
            UsageFetchStatus::RateLimited,
            Some(std::time::Duration::from_secs(3000)),
        );

        assert!(!map.should_probe_now(&k));
        assert!(!map.has_just_cleared_cap_window(&k));
        assert!(!map.scheduler_should_probe(&k), "OR false when both sides are false");
    }

    // ---------------------------------------------------------------
    // #246: LoadingState gates boot + feeds the selection walk.
    //
    // Replaces the PR #238 `consecutive_unauthorized` 3-strike counter:
    // a single 401 now transitions to `LoadingState::Bailed` (clearing
    // the cached usage). The 60 s usage poller absorbs transient
    // failures by re-probing. The user-visible effect is the same -
    // a bailed account reads as bailed instead of showing a stale bar -
    // just the storage shape moved.
    // ---------------------------------------------------------------

    fn key(name: &str) -> AccountKey {
        AccountKey(name.to_owned())
    }

    #[test]
    fn loading_state_initial_is_loading() {
        // Brand-new account state has no probe result yet; it stays
        // Loading until the loading task lands a terminal verdict (Ready
        // or Bailed).
        let map = AccountStateMap::new(&[make_account("Gateway")]);
        assert_eq!(map.loading_state(&key("Gateway")), LoadingState::Loading);
    }

    #[test]
    fn loading_state_returns_loading_for_unknown_key() {
        // Defensive: a caller may hold a key that hasn't been registered
        // yet; treat unknown keys as Loading rather than panicking.
        let map = AccountStateMap::new(&[make_account("Gateway")]);
        assert_eq!(map.loading_state(&key("NotRegistered")), LoadingState::Loading);
    }

    #[test]
    fn set_usage_transitions_loading_to_ready() {
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        assert_eq!(map.loading_state(&k), LoadingState::Loading);
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));
        assert_eq!(map.loading_state(&k), LoadingState::Ready);
    }

    #[test]
    fn set_last_error_unauthorized_transitions_to_bailed_and_clears_usage() {
        // Subsumes the PR #238 three-strike test. Single Unauthorized
        // (not three) now flips loading to Bailed and clears the
        // cached usage, so no stale bar is drawn and the account reads
        // bailed. The 60 s usage poller re-probes the account.
        let mut map = AccountStateMap::new(&[make_account("Personal")]);
        let k = key("Personal");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));
        assert!(map.usage(&k).is_some(), "cache primed");
        assert_eq!(map.loading_state(&k), LoadingState::Ready);

        map.set_last_error(&k, UsageFetchStatus::Unauthorized, None);

        assert_eq!(map.loading_state(&k), LoadingState::Bailed);
        assert!(map.usage(&k).is_none(), "Bailed clears usage by construction");
    }

    #[test]
    fn set_last_error_expired_transitions_to_bailed_and_clears_usage() {
        // Expired is in the same auth-recovery family as Unauthorized;
        // both drive the same Bailed transition.
        let mut map = AccountStateMap::new(&[make_account("Personal")]);
        let k = key("Personal");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));

        map.set_last_error(&k, UsageFetchStatus::Expired, None);

        assert_eq!(map.loading_state(&k), LoadingState::Bailed);
        assert!(map.usage(&k).is_none());
    }

    /// **A rate-limited response's windows are a reading, not an absence:**
    /// a view states them while the account has no successful probe behind
    /// it. The reading is kept beside the successful snapshot, so the
    /// walk's saturation input and the cooldown's reset stay on the last
    /// probe verdict.
    #[test]
    fn a_limit_reading_is_kept_beside_the_successful_snapshot() {
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));
        map.set_last_error(&k, UsageFetchStatus::RateLimited, None);

        map.set_limit_reading(&k, snapshot(Some(100.0), Some(42.0)));

        let reading = map.latest_reading(&k).expect("the 429's own windows");
        assert_eq!(reading.five_hour.as_ref().map(|w| w.utilization), Some(100.0));
        assert_eq!(
            map.usage(&k).and_then(|s| s.five_hour.as_ref().map(|w| w.utilization)),
            Some(30.0),
            "the walk's saturation input stays on the last successful probe",
        );
        assert!(!map.is_saturated(&k), "a 429 reading must not demote the account in the walk");
    }

    /// A successful probe supersedes the reading, and an auth bail drops
    /// it with the cache: a view states the newest thing the account said,
    /// and a dead credential has said nothing since.
    #[test]
    fn a_success_or_an_auth_bail_clears_the_limit_reading() {
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_limit_reading(&k, snapshot(Some(100.0), None));

        map.set_usage(&k, snapshot(Some(20.0), Some(10.0)));
        assert_eq!(
            map.latest_reading(&k).and_then(|s| s.five_hour.as_ref().map(|w| w.utilization)),
            Some(20.0),
            "a successful probe is newer than the reading it replaces",
        );

        map.set_limit_reading(&k, snapshot(Some(100.0), None));
        map.set_last_error(&k, UsageFetchStatus::Unauthorized, None);
        assert!(map.latest_reading(&k).is_none(), "auth bail drops the reading with the cache");
    }

    #[test]
    fn set_last_error_rate_limited_leaves_loading_unchanged() {
        // RateLimited is a probe-side throttle, not an auth failure;
        // loading state must NOT transition to Bailed. The cached
        // usage is also preserved so a view keeps showing the last
        // known-good bars while the probe backs off.
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));
        assert_eq!(map.loading_state(&k), LoadingState::Ready);

        map.set_last_error(&k, UsageFetchStatus::RateLimited, None);

        assert_eq!(map.loading_state(&k), LoadingState::Ready, "RateLimited preserves Ready");
        assert!(map.usage(&k).is_some(), "cache preserved under RateLimited");
    }

    #[test]
    fn set_last_error_network_failed_leaves_loading_unchanged() {
        // Network errors are transient and unknown - the loading state
        // stays at whatever it was, and the cache is preserved.
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));

        map.set_last_error(&k, UsageFetchStatus::NetworkFailed, None);

        assert_eq!(map.loading_state(&k), LoadingState::Ready);
        assert!(map.usage(&k).is_some());
    }

    #[test]
    fn set_loading_bailed_clears_usage() {
        // Driven by the loading task to clear stale snapshots when
        // transitioning to Bailed; mirrors the side-effect that
        // set_last_error performs for the auth-error case.
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));

        map.set_loading(&k, LoadingState::Bailed);

        assert_eq!(map.loading_state(&k), LoadingState::Bailed);
        assert!(map.usage(&k).is_none(), "explicit set_loading(Bailed) also clears usage");
    }

    #[test]
    fn set_loading_to_loading_does_not_clear_usage() {
        // The transition itself shouldn't wipe a cache that might've
        // been re-primed since the bail. (In practice the cache is
        // already None on a Bailed account, but the contract should
        // be explicit.)
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));

        map.set_loading(&k, LoadingState::Loading);

        assert_eq!(map.loading_state(&k), LoadingState::Loading);
        assert!(map.usage(&k).is_some(), "transition to Loading preserves any cached snapshot");
    }

    #[test]
    fn bailed_account_recovers_to_ready_via_set_usage() {
        // Recovery flow: account got Bailed by a probe failure; the
        // 60 s usage poller's re-probe lands a fresh snapshot;
        // set_usage transitions back to Ready, re-priming the cache.
        let mut map = AccountStateMap::new(&[make_account("Personal")]);
        let k = key("Personal");
        map.set_usage(&k, snapshot(Some(30.0), Some(40.0)));
        map.set_last_error(&k, UsageFetchStatus::Unauthorized, None);
        assert_eq!(map.loading_state(&k), LoadingState::Bailed);

        map.set_usage(&k, snapshot(Some(50.0), Some(60.0)));

        assert_eq!(map.loading_state(&k), LoadingState::Ready);
        assert!(map.usage(&k).is_some());
    }

    #[test]
    fn all_loaded_true_for_empty_map() {
        let map = AccountStateMap::empty_for_test();
        assert!(map.all_loaded(), "empty map is vacuously all-loaded");
    }

    #[test]
    fn all_loaded_false_when_any_loading() {
        let map = AccountStateMap::new(&[make_account("Gateway"), make_account("Personal")]);
        assert!(!map.all_loaded(), "fresh accounts start in Loading; gate must stay closed");
    }

    #[test]
    fn all_loaded_true_when_mix_of_ready_and_bailed() {
        let mut map = AccountStateMap::new(&[make_account("Gateway"), make_account("Personal")]);
        map.set_usage(&key("Gateway"), snapshot(Some(10.0), Some(20.0)));
        map.set_last_error(&key("Personal"), UsageFetchStatus::Unauthorized, None);
        assert_eq!(map.loading_state(&key("Gateway")), LoadingState::Ready);
        assert_eq!(map.loading_state(&key("Personal")), LoadingState::Bailed);
        assert!(map.all_loaded(), "Ready + Bailed are both terminal states");
    }

    /// Terminality does not gate re-probing: a Bailed account re-enters
    /// the poller's probe set the moment its hold-down expires, which
    /// is the single-pass loader's whole healing path.
    #[test]
    fn a_bailed_account_probes_again_once_its_hold_down_expires() {
        let mut map = AccountStateMap::new(&[make_account("Gateway")]);
        let k = key("Gateway");
        map.set_last_error(
            &k,
            UsageFetchStatus::RateLimited,
            Some(std::time::Duration::from_secs(3600)),
        );
        map.set_loading(&k, LoadingState::Bailed);
        assert!(!map.scheduler_should_probe(&k), "an active hold-down gates the re-probe");
        map.by_key.get_mut(&k).unwrap().next_probe_at = Some(
            std::time::Instant::now()
                .checked_sub(std::time::Duration::from_secs(1))
                .expect("test clock is not at the Instant floor"),
        );
        assert!(
            map.scheduler_should_probe(&k),
            "an expired hold-down re-opens the probe even though the account is terminal",
        );
    }
}