From 98e397bf8cab1bf71f0d618e791fcd0b14fde0ec Mon Sep 17 00:00:00 2001 From: zsl Date: Mon, 10 Aug 2026 18:08:47 +0300 Subject: [PATCH] fix: scope session-equivalent forecast history by account --- .../src-tauri/src/commands/bridge.rs | 63 ++++++++- .../src-tauri/src/commands/providers.rs | 54 +++++++- .../src-tauri/src/commands/tests.rs | 28 ++-- rust/src/core/session_equivalent_forecast.rs | 123 +++++++++++++++--- 4 files changed, 231 insertions(+), 37 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 5b58357024..c1c32da352 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -155,6 +155,7 @@ impl ProviderUsageSnapshot { id: ProviderId, metadata: &ProviderMetadata, result: &ProviderFetchResult, + token_account_id: Option, ) -> Self { let usage = &result.usage; @@ -194,8 +195,17 @@ impl ProviderUsageSnapshot { s }); - let session_equivalent_forecast = - session_equivalent_forecast_for(id, &usage.primary, usage.secondary.as_ref()); + // Scope forecast history to the signed-in account so switching accounts on one + // provider does not blend burn samples across plans. Codex publishes no email or + // organization (ADR 0003 ambient/managed lanes), so its discriminator is the + // managed token-account id. + let account_key = forecast_account_key(usage, token_account_id); + let session_equivalent_forecast = session_equivalent_forecast_for( + id, + account_key.as_deref(), + &usage.primary, + usage.secondary.as_ref(), + ); Self { provider_id: id.cli_name().to_string(), @@ -302,8 +312,45 @@ impl ProviderUsageSnapshot { } } +/// Account discriminator that forecast history is scoped to. +/// +/// Deliberately mirrors `quota_notification_account_identity` precedence +/// (token account -> email -> organization) so a single account is never seen as two +/// different identities by the notification and forecast subsystems. Kept as a separate +/// function because that one consumes an already-built `ProviderUsageSnapshot`, while the +/// forecast needs the key *while* the snapshot is being built. +/// +/// `providers::tests::forecast_account_key_matches_notification_identity` pins them +/// together. +pub(super) fn forecast_account_key( + usage: &codexbar::core::UsageSnapshot, + token_account_id: Option, +) -> Option { + if let Some(id) = token_account_id { + return Some(format!("token-account:{}", id.as_hyphenated())); + } + if let Some(email) = usage + .account_email + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + return Some(email.to_ascii_lowercase()); + } + if let Some(org) = usage + .account_organization + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + return Some(format!("org:{}", org.to_ascii_lowercase())); + } + None +} + fn session_equivalent_forecast_for( id: ProviderId, + account_key: Option<&str>, session: &RateWindow, weekly: Option<&RateWindow>, ) -> Option { @@ -313,10 +360,16 @@ fn session_equivalent_forecast_for( let weekly = weekly?; let now = chrono::Utc::now(); let provider_id = id.cli_name(); - codexbar::core::record_provider_windows(provider_id, session, Some(weekly), now); + codexbar::core::record_provider_windows(provider_id, account_key, session, Some(weekly), now); let work_days = Settings::load().weekly_progress_work_days; - let forecast = - codexbar::core::forecast_for_provider(provider_id, session, weekly, now, work_days)?; + let forecast = codexbar::core::forecast_for_provider( + provider_id, + account_key, + session, + weekly, + now, + work_days, + )?; Some(SessionEquivalentForecastSnapshot { estimated_windows_to_exhaust_weekly: forecast.estimated_windows_to_exhaust_weekly, windows_until_reset: forecast.windows_until_reset, diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index cf3111bc1a..8f8e21b71b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -343,12 +343,19 @@ fn spawn_provider_refreshes( &inputs.api_keys, &inputs.token_accounts, ); + // Resolved here rather than inside the fetch: forecast history is keyed by + // account, and the managed-account id is the only discriminator Codex exposes. + let token_account_id = inputs + .token_accounts + .get(&id) + .and_then(ProviderAccountData::active_account) + .map(|account| account.id); handles.push(tokio::spawn(async move { let Ok(_permit) = fetch_permits.acquire_owned().await else { return; }; - refresh_provider(app_handle, id, ctx, generation).await; + refresh_provider(app_handle, id, ctx, generation, token_account_id).await; })); } @@ -373,8 +380,9 @@ async fn refresh_provider( id: ProviderId, ctx: FetchContext, generation: u64, + token_account_id: Option, ) { - let snapshot = fetch_provider_snapshot(id, ctx).await; + let snapshot = fetch_provider_snapshot(id, ctx, token_account_id).await; let state = app.state::>(); let published = if let Ok(mut guard) = state.lock() { @@ -585,7 +593,11 @@ fn is_claude_timeout_failure(error: Option<&str>) -> bool { error.eq_ignore_ascii_case("timeout") || error.to_ascii_lowercase().contains("timed out") } -async fn fetch_provider_snapshot(id: ProviderId, ctx: FetchContext) -> ProviderUsageSnapshot { +async fn fetch_provider_snapshot( + id: ProviderId, + ctx: FetchContext, + token_account_id: Option, +) -> ProviderUsageSnapshot { let provider = instantiate_provider(id); let metadata = provider.metadata().clone(); let started = std::time::Instant::now(); @@ -594,7 +606,9 @@ async fn fetch_provider_snapshot(id: ProviderId, ctx: FetchContext) -> ProviderU match tokio::time::timeout(provider_fetch_timeout(id, &ctx), provider.fetch_usage(&ctx)) .await { - Ok(Ok(result)) => ProviderUsageSnapshot::from_fetch_result(id, &metadata, &result), + Ok(Ok(result)) => { + ProviderUsageSnapshot::from_fetch_result(id, &metadata, &result, token_account_id) + } Ok(Err(e)) => ProviderUsageSnapshot::from_error( id, &metadata, @@ -1006,6 +1020,38 @@ mod predictive_warning_tests { snapshot.plan_name = None; assert_eq!(quota_notification_account_identity(&snapshot, None), ""); } + + /// The forecast scope key and the notification identity must never disagree. + /// If they did, one account would be seen as two identities and its burn history + /// would be split, silently halving the sample count behind every forecast. + #[test] + fn forecast_account_key_matches_notification_identity() { + use crate::commands::bridge::forecast_account_key; + + let token = uuid::Uuid::parse_str("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa").unwrap(); + let mut usage = codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(1.0)); + let mut snapshot = empty_snapshot(); + + for (email, org) in [ + (Some("Person@Example.com"), Some("Acme Org")), + (Some("Person@Example.com"), None), + (None, Some("Acme Org")), + (None, None), + ] { + usage.account_email = email.map(str::to_string); + usage.account_organization = org.map(str::to_string); + snapshot.account_email = usage.account_email.clone(); + snapshot.account_organization = usage.account_organization.clone(); + + for tok in [Some(token), None] { + assert_eq!( + forecast_account_key(&usage, tok).unwrap_or_default(), + quota_notification_account_identity(&snapshot, tok), + "identity drift for email={email:?} org={org:?} token={tok:?}" + ); + } + } + } } #[cfg(test)] diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 2323a7e9ed..1709867ac8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -774,7 +774,8 @@ fn provider_cache_upsert_replaces_existing_provider() { wayfinder_usage: None, source_label: "CLI".to_string(), }; - let mut first = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result); + let mut first = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); let mut second = first.clone(); first.error = Some("old".to_string()); second.error = Some("new".to_string()); @@ -796,10 +797,11 @@ fn provider_cache_prunes_disabled_providers() { wayfinder_usage: None, source_label: "CLI".to_string(), }; - let codex = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result); + let codex = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); let claude_meta = instantiate_provider(ProviderId::Claude).metadata().clone(); let claude = - ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &claude_meta, &result); + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &claude_meta, &result, None); let mut cache = vec![codex, claude]; super::prune_provider_cache_to_enabled(&mut cache, &[ProviderId::Codex]); @@ -826,7 +828,7 @@ fn hiding_codex_spark_rows_preserves_other_extra_usage() { source_label: "CLI".to_string(), }; let mut snapshot = - ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result); + ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); snapshot.extra_rate_windows = vec![ NamedRateWindowSnapshot { id: "codex-spark".to_string(), @@ -855,7 +857,8 @@ fn claude_transient_auth_failure_preserves_first_last_good_snapshot() { wayfinder_usage: None, source_label: "OAuth".to_string(), }; - let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result); + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); let error = ProviderUsageSnapshot::from_error( ProviderId::Claude, &metadata, @@ -883,7 +886,8 @@ fn claude_repeated_auth_failure_surfaces_error() { wayfinder_usage: None, source_label: "OAuth".to_string(), }; - let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result); + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); let first_error = ProviderUsageSnapshot::from_error( ProviderId::Claude, &metadata, @@ -916,7 +920,8 @@ fn claude_cli_parse_failure_keeps_last_good_every_time() { wayfinder_usage: None, source_label: "CLI".to_string(), }; - let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result); + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); let err = ProviderUsageSnapshot::from_error( ProviderId::Claude, &metadata, @@ -949,7 +954,8 @@ fn claude_hard_credentials_missing_does_not_preserve_stale() { wayfinder_usage: None, source_label: "OAuth".to_string(), }; - let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result); + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); let err = ProviderUsageSnapshot::from_error( ProviderId::Claude, &metadata, @@ -1081,7 +1087,8 @@ fn japanese_provider_snapshot_localizes_weekly_label() { source_label: "OAuth".to_string(), }; - let snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result); + let snapshot = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); // Secondary label stays raw; localization happens at render time. assert_eq!(snapshot.secondary_label, Some("Weekly".to_string())); @@ -1109,7 +1116,8 @@ fn japanese_provider_snapshot_localizes_pace_reserve_description() { source_label: "OAuth".to_string(), }; - let snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result); + let snapshot = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); // Reserve data stays raw; localization happens at render time. let secondary = snapshot.secondary.as_ref().expect("secondary window"); diff --git a/rust/src/core/session_equivalent_forecast.rs b/rust/src/core/session_equivalent_forecast.rs index dac48b8ec0..c61e8f3221 100644 --- a/rust/src/core/session_equivalent_forecast.rs +++ b/rust/src/core/session_equivalent_forecast.rs @@ -511,13 +511,37 @@ fn is_workday(date: DateTime, work_days: u8) -> bool { // ── In-process history store ───────────────────────────────────────── +/// Identity that forecast history is scoped to. +/// +/// History must never be shared across accounts. Plan sizes differ, so blending +/// observations from two accounts yields a silently wrong median — see the account +/// isolation regression test below. +/// +/// ponytail: when this is persisted to disk, hash `account_key` rather than writing the +/// raw address; in-process it stays plain because the email is already resident in +/// `UsageSnapshot`. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ForecastScope { + pub provider_id: String, + pub account_key: Option, +} + +impl ForecastScope { + pub fn new(provider_id: &str, account_key: Option<&str>) -> Self { + Self { + provider_id: provider_id.to_string(), + account_key: account_key.map(str::to_string), + } + } +} + /// Append-only ring of session/weekly observations for burn estimation. /// /// ponytail: history is process-local and lost on restart; upgrade path = disk cache /// keyed by provider+account once forecast quality justifies persistence. #[derive(Debug, Default)] pub struct SessionEquivalentHistoryStore { - by_provider: HashMap, + by_scope: HashMap, } #[derive(Debug, Default)] @@ -530,10 +554,10 @@ static HISTORY_STORE: LazyLock> = LazyLock::new(|| Mutex::new(SessionEquivalentHistoryStore::default())); impl SessionEquivalentHistoryStore { - /// Record one observation pair for a provider (session + optional weekly). + /// Record one observation pair for a provider+account (session + optional weekly). pub fn record( &mut self, - provider_id: &str, + scope: &ForecastScope, session: Option, weekly: Option, sample_limit: usize, @@ -541,7 +565,7 @@ impl SessionEquivalentHistoryStore { let limit = sample_limit.max(1); // Keep more raw points than completed groups so grouping still has density. let ring = limit.saturating_mul(8).max(24); - let hist = self.by_provider.entry(provider_id.to_string()).or_default(); + let hist = self.by_scope.entry(scope.clone()).or_default(); if let Some(entry) = session { push_ring(&mut hist.session, entry, ring); } @@ -550,8 +574,8 @@ impl SessionEquivalentHistoryStore { } } - pub fn histories(&self, provider_id: &str) -> Vec { - let Some(hist) = self.by_provider.get(provider_id) else { + pub fn histories(&self, scope: &ForecastScope) -> Vec { + let Some(hist) = self.by_scope.get(scope) else { return Vec::new(); }; let mut out = Vec::new(); @@ -591,8 +615,12 @@ pub fn global_history_store() -> &'static Mutex { } /// Record session/weekly windows from a live provider usage snapshot. +/// +/// `account_key` scopes the history so switching accounts on one provider does not +/// blend burn observations across plans. pub fn record_provider_windows( provider_id: &str, + account_key: Option<&str>, session: &crate::core::RateWindow, weekly: Option<&crate::core::RateWindow>, now: DateTime, @@ -624,7 +652,7 @@ pub fn record_provider_windows( if let Ok(mut guard) = global_history_store().lock() { guard.record( - provider_id, + &ForecastScope::new(provider_id, account_key), Some(session_entry), weekly_entry, SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, @@ -662,7 +690,7 @@ pub fn retain_last_full_session_estimate( // ponytail: in-memory only, lost on restart; upgrade = persist to cache file. #[derive(Debug, Default)] struct LastFullSessionEstimateStore { - by_provider: HashMap, + by_scope: HashMap, } static LAST_FULL_SESSION_ESTIMATE_STORE: LazyLock> = @@ -672,38 +700,40 @@ fn last_full_session_estimate_store() -> &'static Mutex, now: DateTime, ) -> Option { let Ok(mut guard) = last_full_session_estimate_store().lock() else { return fresh.filter(|v| v.is_finite() && *v > 0.0); }; - let previous = guard.by_provider.get(provider_id).copied(); + let previous = guard.by_scope.get(scope).copied(); let retained = retain_last_full_session_estimate(previous, fresh, now); if let Some(entry) = retained { - guard.by_provider.insert(provider_id.to_string(), entry); + guard.by_scope.insert(scope.clone(), entry); Some(entry.estimate) } else { - guard.by_provider.remove(provider_id); + guard.by_scope.remove(scope); None } } -/// Compute forecast for a provider using the in-process history ring. +/// Compute forecast for a provider+account using the in-process history ring. pub fn forecast_for_provider( provider_id: &str, + account_key: Option<&str>, session: &crate::core::RateWindow, weekly: &crate::core::RateWindow, now: DateTime, work_days: Option, ) -> Option { + let scope = ForecastScope::new(provider_id, account_key); let histories = global_history_store() .lock() .ok() - .map(|g| g.histories(provider_id)) + .map(|g| g.histories(&scope)) .unwrap_or_default(); let fresh_burn = SessionEquivalentBurnEstimator::estimate( &histories, @@ -718,7 +748,7 @@ pub fn forecast_for_provider( .as_ref() .map(|b| b.sample_count) .unwrap_or(SessionEquivalentBurnEstimator::MINIMUM_SAMPLE_COUNT); - let median = remember_full_session_estimate(provider_id, fresh_median, now)?; + let median = remember_full_session_estimate(&scope, fresh_median, now)?; let burn = SessionEquivalentBurnEstimate { median_weekly_percent_per_window: median, sample_count, @@ -1020,10 +1050,11 @@ mod tests { #[test] fn history_ring_retains_latest_samples() { let mut store = SessionEquivalentHistoryStore::default(); + let scope = ForecastScope::new("claude", None); let base = ts(1_700_500_000); for i in 0..30 { store.record( - "claude", + &scope, Some(PlanUtilizationHistoryEntry { captured_at: base + Duration::minutes(i), used_percent: i as f64, @@ -1033,12 +1064,68 @@ mod tests { 7, ); } - let h = store.histories("claude"); + let h = store.histories(&scope); assert_eq!(h.len(), 1); assert!(h[0].entries.len() <= 7 * 8); assert!((h[0].entries.last().unwrap().used_percent - 29.0).abs() < 1e-9); } + /// Two accounts on one provider must not share burn samples. + /// + /// Regression: history was keyed by `provider_id` alone, so switching the active + /// account blended both accounts' observations into one ring and produced a median + /// drawn from a mixture of plans. + #[test] + fn history_is_isolated_per_account() { + let mut store = SessionEquivalentHistoryStore::default(); + let base = ts(1_700_500_000); + let alice = ForecastScope::new("codex", Some("alice@example.com")); + let bob = ForecastScope::new("codex", Some("bob@example.com")); + + for i in 0..5 { + store.record( + &alice, + Some(PlanUtilizationHistoryEntry { + captured_at: base + Duration::minutes(i), + used_percent: 10.0, + resets_at: Some(base + Duration::hours(5)), + }), + None, + 7, + ); + } + store.record( + &bob, + Some(PlanUtilizationHistoryEntry { + captured_at: base + Duration::minutes(99), + used_percent: 90.0, + resets_at: Some(base + Duration::hours(5)), + }), + None, + 7, + ); + + let alice_hist = store.histories(&alice); + let bob_hist = store.histories(&bob); + assert_eq!(alice_hist[0].entries.len(), 5); + assert_eq!(bob_hist[0].entries.len(), 1); + assert!( + alice_hist[0] + .entries + .iter() + .all(|e| (e.used_percent - 10.0).abs() < 1e-9), + "bob's 90% sample leaked into alice's history" + ); + assert!((bob_hist[0].entries[0].used_percent - 90.0).abs() < 1e-9); + + // A provider with no account discriminator is its own bucket, not a catch-all. + assert!( + store + .histories(&ForecastScope::new("codex", None)) + .is_empty() + ); + } + #[test] fn make_rejects_bad_windows_and_zero_remaining() { let now = ts(1_700_600_000);