diff --git a/apps/desktop-tauri/src-tauri/src/commands/claude_usage.rs b/apps/desktop-tauri/src-tauri/src/commands/claude_usage.rs index 1ee24e7e24..4cc503237e 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/claude_usage.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/claude_usage.rs @@ -1,8 +1,11 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; +use std::io; +use std::path::PathBuf; use std::sync::{Arc, Mutex}; use codexbar::core::{ProviderError, RateWindow, UsageSnapshot}; use codexbar::providers::claude::accounts::{self, ClaudeAccount}; +use codexbar::{atomic_file, secure_file}; use serde::{Deserialize, Serialize}; use tauri::{Emitter, Manager}; @@ -138,6 +141,74 @@ fn publish_results( true } +const ACCOUNT_USAGE_STORE_VERSION: u32 = 1; +const ACCOUNT_USAGE_STORE_RELATIVE_PATH: &str = "claude/account-usage-v1.json"; + +/// Holds only rendered percentages, reset times and sanitized error text. +/// Credentials stay in the account store and are never written here. +#[derive(Serialize, Deserialize)] +struct ClaudeAccountUsageStore { + version: u32, + #[serde(default)] + accounts: BTreeMap, +} + +fn encode_account_usage( + usage: &HashMap, +) -> serde_json::Result { + serde_json::to_string_pretty(&ClaudeAccountUsageStore { + version: ACCOUNT_USAGE_STORE_VERSION, + accounts: usage + .iter() + .map(|(id, state)| (id.clone(), state.clone())) + .collect(), + }) +} + +/// An unreadable or unsupported store is ignored so it never blocks startup. +fn decode_account_usage(raw: &str) -> Option> { + let store: ClaudeAccountUsageStore = serde_json::from_str(raw).ok()?; + (store.version == ACCOUNT_USAGE_STORE_VERSION).then(|| store.accounts.into_iter().collect()) +} + +fn account_usage_store_path() -> Option { + codexbar::logging::config_root().map(|root| root.join(ACCOUNT_USAGE_STORE_RELATIVE_PATH)) +} + +/// Last-known per-account usage, restored at launch so the tray can show it +/// before the first refresh completes. +pub(crate) fn load_account_usage_cache() -> HashMap { + let Some(path) = account_usage_store_path() else { + return HashMap::new(); + }; + match secure_file::read_string(&path) { + Ok(raw) => decode_account_usage(&raw).unwrap_or_else(|| { + tracing::warn!("ignoring unreadable Claude account usage cache"); + HashMap::new() + }), + Err(error) => { + if error.kind() != io::ErrorKind::NotFound { + tracing::warn!("failed to read Claude account usage cache: {error}"); + } + HashMap::new() + } + } +} + +fn save_account_usage_cache(usage: &HashMap) -> io::Result<()> { + let Some(path) = account_usage_store_path() else { + return Ok(()); + }; + let raw = encode_account_usage(usage).map_err(io::Error::other)?; + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + // One fixed staging name is overwritten on each save, so no temp files accumulate. + let staged = path.with_extension("json.tmp"); + secure_file::write_string(&staged, &raw)?; + atomic_file::replace_staged(&staged, &path) +} + /// Share the provider refresh cadence and permits; each request uses only the /// selected account's OAuth credentials and leaves the active CLI login intact. pub(crate) async fn refresh_claude_account_lanes( @@ -151,37 +222,41 @@ pub(crate) async fn refresh_claude_account_lanes( let Ok(accounts) = super::claude_accounts_list() else { return; }; - let mut workers = Vec::new(); - for account in accounts { - let permits = Arc::clone(&permits); - workers.push(tokio::spawn(async move { - let Ok(_permit) = permits.acquire_owned().await else { - return None; + // One task walks the accounts in turn. Concurrent lanes burst the usage + // endpoint and one account's 429 used to block every other account (#787). + // The task owns each credential renewal and persistence, so a UI that stops + // waiting never cancels a refresh-token rotation half way. + let worker = tokio::spawn(async move { + let mut results = Vec::new(); + for account in accounts { + let Ok(_permit) = permits.acquire().await else { + break; }; - // The worker owns the credential operation until renewal and - // persistence finish. Refresh rotation is never cancelled by a UI. let result = accounts::fetch_usage(&account.id) .await .map(|result| result.usage); - Some((account.id, result)) - })); - } - let mut results = Vec::new(); - for worker in workers { - if let Ok(Some(result)) = worker.await { - results.push(result); + results.push((account.id, result)); } - } + results + }); + let results = worker.await.unwrap_or_default(); let Ok(accounts) = super::claude_accounts_list() else { return; }; let state = app.state::>(); - let published = state - .lock() - .is_ok_and(|mut state| publish_results(&mut state, generation, &accounts, results)); - if published { - let _ = app.emit("claude-accounts-updated", ()); + let usage = state.lock().ok().and_then(|mut state| { + publish_results(&mut state, generation, &accounts, results) + .then(|| state.claude_account_usage.clone()) + }); + let Some(usage) = usage else { + return; + }; + if let Err(error) = save_account_usage_cache(&usage) { + tracing::warn!("failed to persist Claude account usage: {error}"); } + let _ = app.emit("claude-accounts-updated", ()); + let handle = app.clone(); + let _ = app.run_on_main_thread(move || crate::tray_bridge::rebuild_tray_menu(&handle)); } #[cfg(test)] @@ -249,4 +324,50 @@ mod tests { fn missing_session_is_unavailable_instead_of_zero_usage() { assert!(quota_window(&RateWindow::no_active_session()).is_none()); } + + #[test] + fn account_usage_cache_round_trips_per_account_and_rejects_other_versions() { + let mut usage = HashMap::new(); + usage.insert( + "acct-a".to_string(), + ClaudeAccountUsageState { + usage: Some(ClaudeAccountUsage { + five_hour: Some(ClaudeAccountWindow { + used_percent: 42.0, + resets_at: None, + }), + seven_day: None, + updated_at: "2026-10-18T00:00:00+00:00".into(), + }), + usage_error: None, + needs_authentication: false, + }, + ); + usage.insert( + "acct-b".to_string(), + ClaudeAccountUsageState { + usage_error: Some("Sign in again to check this Claude account.".into()), + needs_authentication: true, + ..Default::default() + }, + ); + let raw = encode_account_usage(&usage).unwrap(); + let decoded = decode_account_usage(&raw).expect("current version decodes"); + assert_eq!(decoded.len(), 2); + assert_eq!( + decoded["acct-a"] + .usage + .as_ref() + .unwrap() + .five_hour + .as_ref() + .unwrap() + .used_percent, + 42.0 + ); + assert!(decoded["acct-b"].needs_authentication); + let other_version = raw.replace("\"version\": 1", "\"version\": 2"); + assert!(decode_account_usage(&other_version).is_none()); + assert!(decode_account_usage("not json").is_none()); + } } diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index e945098821..c3c982961b 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -155,6 +155,7 @@ fn main() { let mut initial_state = AppState::new(); initial_state.proof_config = proof_config; + initial_state.claude_account_usage = commands::load_account_usage_cache(); // Validate the complete proof seed before installing any snapshots, so an // invalid multi-provider fixture cannot leave a partial cache behind. if let Some(snapshots) = proof_harness::seed_usage_snapshots_from_env( diff --git a/apps/desktop-tauri/src-tauri/src/tray_accounts.rs b/apps/desktop-tauri/src-tauri/src/tray_accounts.rs index 3cfcf32dbc..41c67d2577 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_accounts.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_accounts.rs @@ -27,7 +27,11 @@ pub(crate) enum AccountMenuAction { SwitchGrokAccount(String), } -pub(crate) fn prepend_account_menus(spec: &mut Vec, settings: &Settings) { +pub(crate) fn prepend_account_menus( + spec: &mut Vec, + settings: &Settings, + claude_usage: &HashMap, +) { let accounts = crate::commands::load_codex_accounts().unwrap_or_default(); let active = codexbar::codex_accounts::CodexAccountManager::new().discover_ambient_account(&accounts); @@ -48,6 +52,7 @@ pub(crate) fn prepend_account_menus(spec: &mut Vec, settings: &Se &claude_accounts, settings.ui_language, settings.hide_personal_info, + claude_usage, ), ); @@ -284,13 +289,56 @@ fn hidden_account_label(lang: Language, ordinal: usize) -> String { format!("{word} {ordinal}") } +/// Session and weekly percentages for one Claude account row, from the last +/// successful check recorded for that account. No recorded check, no percentage. +/// A window whose reset time has passed is skipped: the stored percentage is +/// known to be stale (restored cache, failed checks, or refresh turned off). +fn claude_usage_summary( + lang: Language, + state: Option<&crate::commands::ClaudeAccountUsageState>, + now: chrono::DateTime, +) -> Option { + let usage = state?.usage.as_ref()?; + let parts: Vec = [ + (LocaleKey::ProviderSessionLabel, usage.five_hour.as_ref()), + (LocaleKey::ProviderWeeklyLabel, usage.seven_day.as_ref()), + ] + .into_iter() + .filter_map(|(key, window)| { + window + .filter(|window| !window_has_reset(window, now)) + .map(|window| { + format!( + "{} {:.0}%", + locale::get_text(lang, key), + window.used_percent + ) + }) + }) + .collect(); + (!parts.is_empty()).then(|| parts.join(", ")) +} + +fn window_has_reset( + window: &crate::commands::ClaudeAccountWindow, + now: chrono::DateTime, +) -> bool { + window + .resets_at + .as_deref() + .and_then(|at| chrono::DateTime::parse_from_rfc3339(at).ok()) + .is_some_and(|at| at <= now) +} + fn claude_accounts_menu( accounts: &[codexbar::providers::claude::accounts::ClaudeAccount], lang: Language, hide_personal_info: bool, + usage: &HashMap, ) -> TrayMenuEntry { let text = |key| locale::get_text(lang, key); let ordinals = ordinals_by_sorted_id(accounts.iter().map(|account| account.id.as_str())); + let now = chrono::Utc::now(); let mut children: Vec<_> = accounts .iter() .map(|account| { @@ -302,6 +350,10 @@ fn claude_accounts_menu( |org| format!("{} ({org})", account.email), ) }; + let label = match claude_usage_summary(lang, usage.get(&account.id), now) { + Some(summary) => format!("{label} - {summary}"), + None => label, + }; let mut entry = TrayMenuEntry::check_item( format!("switch_claude_account:{}", account.id), label, @@ -570,7 +622,8 @@ mod tests { is_saved: true, ..current.clone() }; - let menu = claude_accounts_menu(&[current, saved], Language::English, false); + let menu = + claude_accounts_menu(&[current, saved], Language::English, false, &HashMap::new()); assert_eq!(menu.id.as_deref(), Some("claude_accounts")); assert_eq!(menu.children[0].checked, Some(true)); assert!(menu.children[0].disabled); @@ -581,11 +634,95 @@ mod tests { assert!(!menu.children[1].disabled); assert!(menu_contains(&menu.children, "add_claude_account")); assert!(menu_contains( - &claude_accounts_menu(&[], Language::English, false).children, + &claude_accounts_menu(&[], Language::English, false, &HashMap::new()).children, "add_claude_account" )); } + #[test] + fn claude_rows_show_each_accounts_last_recorded_usage() { + use crate::commands::{ClaudeAccountUsage, ClaudeAccountUsageState, ClaudeAccountWindow}; + let accounts = [claude_menu_account("a:org"), claude_menu_account("b:org")]; + let mut usage = HashMap::new(); + usage.insert( + "a:org".to_string(), + ClaudeAccountUsageState { + usage: Some(ClaudeAccountUsage { + five_hour: Some(ClaudeAccountWindow { + used_percent: 42.4, + resets_at: None, + }), + seven_day: Some(ClaudeAccountWindow { + used_percent: 10.0, + resets_at: None, + }), + updated_at: "2026-10-18T00:00:00+00:00".into(), + }), + ..Default::default() + }, + ); + usage.insert( + "b:org".to_string(), + ClaudeAccountUsageState { + usage_error: Some("Claude usage check failed. Will retry automatically.".into()), + ..Default::default() + }, + ); + let menu = claude_accounts_menu(&accounts, Language::English, false, &usage); + let session = locale::get_text(Language::English, LocaleKey::ProviderSessionLabel); + let weekly = locale::get_text(Language::English, LocaleKey::ProviderWeeklyLabel); + assert_eq!( + menu.children[0].label, + format!("a:org@example.com - {session} 42%, {weekly} 10%") + ); + // A row with no successful check recorded shows no percentage. + assert_eq!(menu.children[1].label, "b:org@example.com"); + } + + #[test] + fn claude_usage_summary_skips_windows_past_their_reset() { + use crate::commands::{ClaudeAccountUsage, ClaudeAccountUsageState, ClaudeAccountWindow}; + let now = chrono::DateTime::parse_from_rfc3339("2026-10-18T12:00:00Z") + .unwrap() + .with_timezone(&chrono::Utc); + let state = |five_hour: Option<&str>, seven_day: Option<&str>| ClaudeAccountUsageState { + usage: Some(ClaudeAccountUsage { + five_hour: Some(ClaudeAccountWindow { + used_percent: 42.0, + resets_at: five_hour.map(String::from), + }), + seven_day: Some(ClaudeAccountWindow { + used_percent: 10.0, + resets_at: seven_day.map(String::from), + }), + updated_at: "2026-10-18T00:00:00+00:00".into(), + }), + ..Default::default() + }; + let session = locale::get_text(Language::English, LocaleKey::ProviderSessionLabel); + let weekly = locale::get_text(Language::English, LocaleKey::ProviderWeeklyLabel); + + let expired_session = state(Some("2026-10-18T05:00:00Z"), Some("2026-10-20T00:00:00Z")); + assert_eq!( + claude_usage_summary(Language::English, Some(&expired_session), now), + Some(format!("{weekly} 10%")) + ); + let both_expired = state( + Some("2026-10-18T12:00:00Z"), + Some("2026-10-17T00:00:00+02:00"), + ); + assert_eq!( + claude_usage_summary(Language::English, Some(&both_expired), now), + None + ); + // Missing or unparseable reset times keep the window. + let unknown = state(None, Some("not a date")); + assert_eq!( + claude_usage_summary(Language::English, Some(&unknown), now), + Some(format!("{session} 42%, {weekly} 10%")) + ); + } + #[test] fn grok_menu_checks_current_account_and_routes_saved_accounts() { use codexbar::providers::grok::accounts::GrokAccount; @@ -656,7 +793,12 @@ mod tests { is_active: true, is_saved: true, }; - let claude_menu = claude_accounts_menu(&[claude, claude_other], Language::English, true); + let claude_menu = claude_accounts_menu( + &[claude, claude_other], + Language::English, + true, + &HashMap::new(), + ); let grok_menu = grok_accounts_menu(&[grok], Language::English, true); let labels: Vec<&str> = claude_menu.children[..2] .iter() @@ -680,6 +822,7 @@ mod tests { ], Language::English, true, + &HashMap::new(), ); let by_id: Vec<(&str, &str)> = reordered.children[..3] .iter() diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index cc4cc1d621..2c2b76d3d7 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -124,7 +124,17 @@ fn build_native_tray_menu( }, settings.ui_language, ); - crate::tray_accounts::prepend_account_menus(&mut spec, &settings); + // Callers release the AppState lock before building the tray menu. + let claude_usage = app + .try_state::>() + .and_then(|state| { + state + .lock() + .ok() + .map(|guard| guard.claude_account_usage.clone()) + }) + .unwrap_or_default(); + crate::tray_accounts::prepend_account_menus(&mut spec, &settings, &claude_usage); let entries = spec .iter() .map(|entry| build_native_menu_entry(app, entry)) diff --git a/rust/src/providers/claude/oauth/mod.rs b/rust/src/providers/claude/oauth/mod.rs index 6b60bab73c..99279abed2 100644 --- a/rust/src/providers/claude/oauth/mod.rs +++ b/rust/src/providers/claude/oauth/mod.rs @@ -385,7 +385,9 @@ impl ClaudeOAuthFetcher { { return Err(error); } - self.fetch_with_credentials(credentials).await + let gate_key = + usage_gate::account_key(saved_account_id.as_deref(), &credentials.access_token); + self.fetch_with_credentials(credentials, &gate_key).await } /// Fetch a saved account directly; never switch Claude Code to obtain usage. @@ -454,7 +456,11 @@ impl ClaudeOAuthFetcher { { return Err(account_refresh_error(error)); } - let result = self.fetch_with_credentials(credentials.clone()).await; + // Keyed by the account, so its own 429 block survives a token renewal. + let gate_key = usage_gate::account_key(Some(id), &credentials.access_token); + let result = self + .fetch_with_credentials(credentials.clone(), &gate_key) + .await; // Unknown expiry metadata and early server expiry still get one renewal. if matches!(&result, Err(ProviderError::OAuthExpired(_))) && credentials.refresh_token.is_some() @@ -468,7 +474,7 @@ impl ClaudeOAuthFetcher { if let Some(error) = outcome { return Err(account_refresh_error(error)); } - return self.fetch_with_credentials(refreshed).await; + return self.fetch_with_credentials(refreshed, &gate_key).await; } result } @@ -493,14 +499,16 @@ impl ClaudeOAuthFetcher { rate_limit_tier: None, }; - self.fetch_with_credentials(credentials).await + let gate_key = usage_gate::account_key(None, &credentials.access_token); + self.fetch_with_credentials(credentials, &gate_key).await } async fn fetch_with_credentials( &self, credentials: ClaudeOAuthCredentials, + gate_key: &str, ) -> Result { - let usage_response = self.fetch_usage(&credentials).await?; + let usage_response = self.fetch_usage(&credentials, gate_key).await?; let mut usage = self.build_usage_snapshot(&usage_response, &credentials); let spend_cost = apply_spend_only_extra_usage(&mut usage, &usage_response); // Header email only; a profile failure leaves the usage result intact. @@ -652,6 +660,7 @@ impl ClaudeOAuthFetcher { pub async fn fetch_usage( &self, credentials: &ClaudeOAuthCredentials, + gate_key: &str, ) -> Result { if credentials.is_expired() { return Err(ProviderError::OAuthExpired( @@ -668,8 +677,7 @@ impl ClaudeOAuthFetcher { ))); } - let fingerprint = usage_gate::token_fingerprint(&credentials.access_token); - if let Some(result) = Self::gate_precheck(&USAGE_GATE, usage_gate::now_ms(), &fingerprint) { + if let Some(result) = Self::gate_precheck(&USAGE_GATE, usage_gate::now_ms(), gate_key) { return result; } @@ -724,7 +732,7 @@ impl ClaudeOAuthFetcher { let backoff = USAGE_GATE .lock() .unwrap_or_else(|e| e.into_inner()) - .record_rate_limit(usage_gate::now_ms(), &fingerprint, retry_after); + .record_rate_limit(usage_gate::now_ms(), gate_key, retry_after); return Err(Self::rate_limited_error(backoff)); } @@ -744,7 +752,7 @@ impl ClaudeOAuthFetcher { USAGE_GATE .lock() .unwrap_or_else(|e| e.into_inner()) - .record_success(usage_gate::now_ms(), &fingerprint, body); + .record_success(usage_gate::now_ms(), gate_key, body); Ok(usage) } diff --git a/rust/src/providers/claude/oauth/usage_gate.rs b/rust/src/providers/claude/oauth/usage_gate.rs index 94e5e84fc6..4eda226f73 100644 --- a/rust/src/providers/claude/oauth/usage_gate.rs +++ b/rust/src/providers/claude/oauth/usage_gate.rs @@ -1,24 +1,32 @@ //! Gate in front of the Claude OAuth usage endpoint (nesszer/Win-CodexBar#775). //! -//! * After a 429 no request goes out until +//! * After a 429 no request goes out for that account until //! `max(Retry-After.min(60 min), (5 min * 2^(n-1)).min(30 min))` has passed. -//! `n` counts consecutive 429s and resets only on a success from the sign-in -//! that was blocked, never because the block expired or because another -//! account succeeded. The block itself ignores the token, so a token refresh -//! does not lift it. Concurrent 429s while a block is active do not escalate. +//! `n` counts consecutive 429s of that account and resets only on a success +//! of the same account, never because the block expired or because another +//! account succeeded. A 429 blocks only the account that received it, so one +//! saved account's rate limit never hides another account's usage (#787). +//! Concurrent 429s while a block is active do not escalate. +//! * Accounts are keyed by `account_key`: the hashed saved-account id when it is +//! known, otherwise the token fingerprint. A renewed access token keeps the +//! same account key, so a token refresh does not lift the block. //! * The state lives in one small JSON file under `%LOCALAPPDATA%\CodexBar\`, //! so every CodexBar process of the user (tray app and one-off //! `codexbar usage` runs) shares it. An in-memory copy keeps the gate working //! when the file cannot be read or written. -//! * The last successful response body is kept in the same file and served again -//! for 3 minutes to the same sign-in. The response holds only utilization -//! percentages, reset timestamps and credit amounts (see `OAuthUsageResponse` -//! in `oauth/mod.rs`), no token or account identity. Precedent: the CLI probe -//! screen is persisted across processes the same way. The sign-in is keyed by -//! the first 16 hex characters of the SHA-256 of the access token; the token -//! itself is never written. +//! * The last successful response body of each account is kept in the same file +//! and served again for 3 minutes to that same account. The response holds +//! only utilization percentages, reset timestamps and credit amounts (see +//! `OAuthUsageResponse` in `oauth/mod.rs`), no token or account identity. +//! Precedent: the CLI probe screen is persisted across processes the same way. +//! The token itself is never written; keys are truncated SHA-256 digests. +//! * Blocks and cached bodies are bounded to `MAX_CACHED_SIGN_INS` accounts. +//! * Files written before per-account state still load. A single-slot +//! `last_success` is migrated into the map on the next read. A global block +//! from an older build cannot be attributed to one account, so it is dropped. use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::path::PathBuf; use std::time::Duration; @@ -28,6 +36,9 @@ const MAX_BACKOFF_MS: i64 = 30 * 60 * 1000; /// Cap of an honored `Retry-After`; also the longest wait ever reported. const MAX_RETRY_AFTER_MS: i64 = 60 * 60 * 1000; const MIN_POLL_INTERVAL_MS: i64 = 3 * 60 * 1000; +/// Accounts whose block or last good body is kept at once. Saved accounts are +/// few; the bound only stops rotated-away sign-ins from growing the file forever. +const MAX_CACHED_SIGN_INS: usize = 8; const FILE_NAME: &str = "claude-oauth-usage-gate.json"; /// What the caller should do before hitting the endpoint. @@ -42,33 +53,91 @@ pub(super) enum Decision { #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] struct CachedBody { at_ms: i64, + /// The account key this body belongs to (never a token). fp: String, body: String, } #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] -struct Persisted { - blocked_until_ms: i64, +struct Block { + /// Requests for this account are refused until this epoch millisecond. + until_ms: i64, + /// Consecutive 429s of this account; resets when this account succeeds. consecutive: u32, - /// Fingerprint of the sign-in whose 429 started the current streak. +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +struct Persisted { + /// 429 block and streak per account key, bounded to `MAX_CACHED_SIGN_INS`. #[serde(default)] - blocked_fp: String, + blocks: BTreeMap, + /// Last good body per account key, bounded to `MAX_CACHED_SIGN_INS`. #[serde(default)] + last_successes: BTreeMap, + /// Single-slot cache written before `last_successes`. Read only so old files + /// load; it is migrated into the map and never written back. + #[serde(default, skip_serializing)] last_success: Option, } +impl Persisted { + /// Fold a legacy single slot into the map, then enforce the bounds. + fn normalize(&mut self) { + if let Some(legacy) = self.last_success.take() { + let newer = self + .last_successes + .get(&legacy.fp) + .is_none_or(|kept| kept.at_ms < legacy.at_ms); + if newer { + self.last_successes.insert(legacy.fp.clone(), legacy); + } + } + trim_to_bound(&mut self.last_successes, |cached| cached.at_ms); + trim_to_bound(&mut self.blocks, |block| block.until_ms); + } +} + +/// Evict the entries with the smallest timestamp until the map is within its bound. +fn trim_to_bound(map: &mut BTreeMap, time_ms: impl Fn(&V) -> i64) { + while map.len() > MAX_CACHED_SIGN_INS { + let Some(oldest) = map + .iter() + .min_by_key(|(_, value)| time_ms(value)) + .map(|(key, _)| key.clone()) + else { + break; + }; + map.remove(&oldest); + } +} + pub(super) struct UsageGate { path: Option, state: Persisted, } -/// First 16 hex characters of the token's SHA-256. Never the token. -pub(super) fn token_fingerprint(token: &str) -> String { - let mut hex = crate::core::sha256_hex(token.as_bytes()); +/// First 16 hex characters of the SHA-256 of `text`. +fn digest16(text: &str) -> String { + let mut hex = crate::core::sha256_hex(text.as_bytes()); hex.truncate(16); hex } +/// First 16 hex characters of the token's SHA-256. Never the token. +pub(super) fn token_fingerprint(token: &str) -> String { + digest16(token) +} + +/// Gate key for one Claude account. A known account is keyed by a digest of +/// its id, so renewing its token keeps the same key and its block. A sign-in +/// whose account is not known falls back to its token fingerprint. +pub(super) fn account_key(account_id: Option<&str>, access_token: &str) -> String { + match account_id { + Some(id) => format!("acct-{}", digest16(id)), + None => token_fingerprint(access_token), + } +} + pub(super) fn default_path() -> Option { dirs::data_local_dir().map(|dir| dir.join("CodexBar").join(FILE_NAME)) } @@ -103,7 +172,8 @@ impl UsageGate { let Ok(text) = std::fs::read_to_string(path) else { return; }; - if let Ok(disk) = serde_json::from_str::(&text) { + if let Ok(mut disk) = serde_json::from_str::(&text) { + disk.normalize(); self.state = disk; } } @@ -124,20 +194,20 @@ impl UsageGate { } } - /// Time left on the active block, clamped so a skewed or hand-edited + /// Time left on this account's block, clamped so a skewed or hand-edited /// timestamp cannot lock requests out for longer than any real backoff. - fn remaining(&self, now_ms: i64) -> Option { - let left = self.state.blocked_until_ms.saturating_sub(now_ms); + fn remaining(&self, now_ms: i64, key: &str) -> Option { + let until = self.state.blocks.get(key)?.until_ms; + let left = until.saturating_sub(now_ms); (left > 0).then(|| Duration::from_millis(left.min(MAX_RETRY_AFTER_MS) as u64)) } - pub(super) fn check(&mut self, now_ms: i64, fingerprint: &str) -> Decision { + pub(super) fn check(&mut self, now_ms: i64, key: &str) -> Decision { self.refresh(); - if let Some(remaining) = self.remaining(now_ms) { + if let Some(remaining) = self.remaining(now_ms, key) { return Decision::Blocked(remaining); } - if let Some(last) = &self.state.last_success - && last.fp == fingerprint + if let Some(last) = self.state.last_successes.get(key) && now_ms >= last.at_ms && now_ms - last.at_ms < MIN_POLL_INTERVAL_MS { @@ -146,41 +216,43 @@ impl UsageGate { Decision::Proceed } - /// Record a 429 and return how long requests are blocked. + /// Record a 429 for `key` and return how long that account is blocked. pub(super) fn record_rate_limit( &mut self, now_ms: i64, - fingerprint: &str, + key: &str, retry_after: Duration, ) -> Duration { self.refresh(); - if let Some(remaining) = self.remaining(now_ms) { + if let Some(remaining) = self.remaining(now_ms, key) { // Another request already escalated; do not count this 429 again. return remaining; } - self.state.consecutive = self.state.consecutive.saturating_add(1); - self.state.blocked_fp = fingerprint.to_string(); - let backoff = backoff_ms(retry_after, self.state.consecutive); - self.state.blocked_until_ms = now_ms.saturating_add(backoff); + let block = self.state.blocks.entry(key.to_string()).or_default(); + block.consecutive = block.consecutive.saturating_add(1); + let backoff = backoff_ms(retry_after, block.consecutive); + block.until_ms = now_ms.saturating_add(backoff); + trim_to_bound(&mut self.state.blocks, |block| block.until_ms); self.persist(); Duration::from_millis(backoff as u64) } - /// Remember a successful response body. The 429 streak resets only when - /// the sign-in that was blocked succeeds. The refresh is needed so this - /// write does not overwrite a block another process just recorded. - pub(super) fn record_success(&mut self, now_ms: i64, fingerprint: &str, body: String) { + /// Remember a successful response body for `key`. The account's own 429 + /// streak ends with its own success. The refresh is needed so this write + /// does not overwrite a block another process just recorded or another + /// account's cached body. + pub(super) fn record_success(&mut self, now_ms: i64, key: &str, body: String) { self.refresh(); - if self.state.blocked_fp == fingerprint { - self.state.consecutive = 0; - self.state.blocked_until_ms = 0; - self.state.blocked_fp.clear(); - } - self.state.last_success = Some(CachedBody { - at_ms: now_ms, - fp: fingerprint.to_string(), - body, - }); + self.state.blocks.remove(key); + self.state.last_successes.insert( + key.to_string(), + CachedBody { + at_ms: now_ms, + fp: key.to_string(), + body, + }, + ); + trim_to_bound(&mut self.state.last_successes, |cached| cached.at_ms); self.persist(); } } @@ -244,7 +316,7 @@ mod tests { } #[test] - fn success_of_the_blocked_sign_in_resets_the_count() { + fn success_of_the_blocked_account_resets_the_count() { let dir = tempfile::tempdir().unwrap(); let mut g = gate(&dir); g.record_rate_limit(NOW, "fp", Duration::ZERO); @@ -258,7 +330,7 @@ mod tests { } #[test] - fn success_of_another_sign_in_keeps_the_count() { + fn success_of_another_account_keeps_the_count() { let dir = tempfile::tempdir().unwrap(); let mut g = gate(&dir); g.record_rate_limit(NOW, "fp-a", Duration::ZERO); @@ -271,30 +343,76 @@ mod tests { } #[test] - fn block_is_shared_through_the_file_and_ignores_the_token() { + fn block_is_shared_through_the_file_and_stays_with_its_account() { let dir = tempfile::tempdir().unwrap(); - gate(&dir).record_rate_limit(NOW, "fp", Duration::ZERO); + gate(&dir).record_rate_limit(NOW, "acct-a", Duration::ZERO); let mut other = gate(&dir); - assert_eq!(blocked_mins(other.check(NOW + MIN, "another-token")), 4); - assert_eq!( - other.check(NOW + 5 * MIN + 1, "another-token"), - Decision::Proceed - ); + assert_eq!(blocked_mins(other.check(NOW + MIN, "acct-a")), 4); + assert_eq!(other.check(NOW + 5 * MIN + 1, "acct-a"), Decision::Proceed); // The count survives in the file, so the next 429 doubles. assert_eq!( other - .record_rate_limit(NOW + 6 * MIN, "fp", Duration::ZERO) + .record_rate_limit(NOW + 6 * MIN, "acct-a", Duration::ZERO) .as_secs(), 600 ); } + #[test] + fn a_429_blocks_only_its_own_account() { + let dir = tempfile::tempdir().unwrap(); + let mut g = gate(&dir); + g.record_rate_limit(NOW, "acct-a", Duration::ZERO); + assert_eq!(g.check(NOW + MIN, "acct-b"), Decision::Proceed); + // Account B's own 429 starts B's streak; account A keeps its block. + assert_eq!( + g.record_rate_limit(NOW + MIN, "acct-b", Duration::ZERO) + .as_secs(), + 5 * 60 + ); + assert_eq!(blocked_mins(g.check(NOW + MIN, "acct-a")), 4); + } + + #[test] + fn account_key_is_stable_across_token_renewal_and_hides_the_id() { + let before = account_key(Some("uuid-1:org-1"), "token-old"); + let renewed = account_key(Some("uuid-1:org-1"), "token-new"); + assert_eq!(before, renewed); + assert_ne!(before, account_key(Some("uuid-2:org-1"), "token-old")); + assert!(!before.contains("uuid-1")); + // Without a known account the key is the token's fingerprint. + assert_eq!( + account_key(None, "token-old"), + token_fingerprint("token-old") + ); + assert_ne!( + account_key(None, "token-old"), + account_key(None, "token-new") + ); + } + + #[test] + fn block_survives_token_renewal_of_the_same_account() { + let dir = tempfile::tempdir().unwrap(); + let key = account_key(Some("a:org"), "old-token"); + gate(&dir).record_rate_limit(NOW, &key, Duration::ZERO); + let renewed = account_key(Some("a:org"), "new-token"); + assert_eq!(blocked_mins(gate(&dir).check(NOW + MIN, &renewed)), 4); + } + #[test] fn skewed_block_timestamp_is_clamped() { let dir = tempfile::tempdir().unwrap(); + let mut blocks = BTreeMap::new(); + blocks.insert( + "fp".to_string(), + Block { + until_ms: NOW + 1000 * 60 * MIN, + consecutive: 1, + }, + ); let state = Persisted { - blocked_until_ms: NOW + 1000 * 60 * MIN, - consecutive: 1, + blocks, ..Persisted::default() }; std::fs::write( @@ -305,6 +423,20 @@ mod tests { assert_eq!(blocked_mins(gate(&dir).check(NOW, "fp")), 60); } + #[test] + fn legacy_global_block_is_not_attributed_to_any_account() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join(FILE_NAME), + format!( + r#"{{"blocked_until_ms":{},"consecutive":1,"blocked_fp":"fp"}}"#, + NOW + 10 * MIN + ), + ) + .unwrap(); + assert_eq!(gate(&dir).check(NOW, "fp"), Decision::Proceed); + } + #[test] fn memory_copy_works_without_a_usable_file() { let dir = tempfile::tempdir().unwrap(); @@ -340,7 +472,7 @@ mod tests { } #[test] - fn successful_polls_are_throttled_per_sign_in_across_processes() { + fn successful_polls_are_throttled_per_account_across_processes() { let dir = tempfile::tempdir().unwrap(); gate(&dir).record_success(NOW, "fp-a", "usage-a".into()); let mut other = gate(&dir); @@ -366,4 +498,85 @@ mod tests { assert!(!text.contains("secret-token") && !text.contains(&full)); assert!(text.contains(&fp)); } + + #[test] + fn two_saved_accounts_keep_their_cached_bodies_within_the_ttl() { + let dir = tempfile::tempdir().unwrap(); + let mut g = gate(&dir); + g.record_success(NOW, "fp-a", "usage-a".into()); + g.record_success(NOW + MIN, "fp-b", "usage-b".into()); + let mut other = gate(&dir); + assert_eq!( + other.check(NOW + 2 * MIN, "fp-a"), + Decision::Cached("usage-a".into()) + ); + assert_eq!( + other.check(NOW + 2 * MIN, "fp-b"), + Decision::Cached("usage-b".into()) + ); + // Each slot still expires on its own clock. + assert_eq!(other.check(NOW + 3 * MIN + 1, "fp-a"), Decision::Proceed); + } + + #[test] + fn success_cache_is_bounded_and_evicts_the_oldest_sign_in() { + let dir = tempfile::tempdir().unwrap(); + let mut g = gate(&dir); + for n in 0..=MAX_CACHED_SIGN_INS { + let at = NOW + i64::try_from(n).unwrap() * 1000; + g.record_success(at, &format!("fp-{n}"), format!("usage-{n}")); + } + let text = std::fs::read_to_string(dir.path().join(FILE_NAME)).unwrap(); + let file: Persisted = serde_json::from_str(&text).unwrap(); + assert_eq!(file.last_successes.len(), MAX_CACHED_SIGN_INS); + let now = NOW + 60 * 1000; + assert_eq!(g.check(now, "fp-0"), Decision::Proceed); + assert_eq!( + g.check(now, &format!("fp-{MAX_CACHED_SIGN_INS}")), + Decision::Cached(format!("usage-{MAX_CACHED_SIGN_INS}")) + ); + } + + #[test] + fn legacy_single_slot_file_loads_and_is_served_then_rewritten() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join(FILE_NAME); + std::fs::write( + &path, + format!( + r#"{{"blocked_until_ms":0,"consecutive":0,"last_success":{{"at_ms":{NOW},"fp":"fp-a","body":"usage-a"}}}}"# + ), + ) + .unwrap(); + let mut g = gate(&dir); + assert_eq!( + g.check(NOW + MIN, "fp-a"), + Decision::Cached("usage-a".into()) + ); + assert_eq!(g.check(NOW + MIN, "fp-b"), Decision::Proceed); + g.record_success(NOW + 2 * MIN, "fp-b", "usage-b".into()); + let text = std::fs::read_to_string(&path).unwrap(); + assert!(!text.contains("\"last_success\"")); + let mut reread = gate(&dir); + assert_eq!( + reread.check(NOW + 2 * MIN, "fp-a"), + Decision::Cached("usage-a".into()) + ); + assert_eq!( + reread.check(NOW + 2 * MIN, "fp-b"), + Decision::Cached("usage-b".into()) + ); + } + + #[test] + fn legacy_null_or_missing_slot_and_older_files_load_without_error() { + for body in [ + r#"{"blocked_until_ms":0,"consecutive":0,"last_success":null}"#, + r#"{"blocked_until_ms":0,"consecutive":2}"#, + ] { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join(FILE_NAME), body).unwrap(); + assert_eq!(gate(&dir).check(NOW, "fp-a"), Decision::Proceed); + } + } }