From fed9068b2a2ac1cc17eadc3899b5395d91fb60fc Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:05:55 +0700 Subject: [PATCH] Port upstream 0.66.0: preserve Claude threshold warnings across identity gaps --- .../src-tauri/src/commands/providers.rs | 53 +++- .../src/commands/warning_identity.rs | 56 ++++ rust/src/notifications.rs | 20 +- rust/src/notifications/identity_gaps.rs | 264 ++++++++++++++++ rust/src/notifications/identity_gaps/tests.rs | 283 ++++++++++++++++++ 5 files changed, 672 insertions(+), 4 deletions(-) create mode 100644 rust/src/notifications/identity_gaps.rs create mode 100644 rust/src/notifications/identity_gaps/tests.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..d79aa4a646 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -6,6 +6,7 @@ use super::warning_identity::WarningIdentity; use super::*; use chrono::{Local, Utc}; use codexbar::core::HookUsageWindow; +use codexbar::notifications::WarningScope; use serde::Serialize; use std::sync::Arc; @@ -976,12 +977,24 @@ fn notify_usage_thresholds( snapshot.account_organization.as_deref(), token_account_id, ); + // Hooks keep their own per-source baselines (edge-triggered, first sample + // never fires), so they stay on the source key; only toast dedupe below + // bridges account-identity gaps. let account = warning_identity.threshold_key(); + let scope = warning_identity.gap_scope(); // Skip all session consumers for synthetic/no-session // placeholders (e.g. Claude OAuth five_hour: null). + let session_account = resolve_toast_account( + &mut guard.notification_manager, + provider, + &scope, + "session", + &snapshot.primary, + settings, + ); if guard.notification_manager.check_session_lane( provider, - &account, + &session_account, snapshot.primary.used_percent, snapshot.primary.is_informational, settings, @@ -997,9 +1010,17 @@ fn notify_usage_thresholds( if let Some(weekly) = &snapshot.secondary && !weekly.is_informational { + let weekly_account = resolve_toast_account( + &mut guard.notification_manager, + provider, + &scope, + "weekly", + weekly, + settings, + ); guard.notification_manager.check_and_notify( provider, - &account, + &weekly_account, "weekly", weekly.used_percent, settings, @@ -1024,6 +1045,34 @@ fn notify_usage_thresholds( } } +/// Account key a toast lane is deduped under (see `NotificationManager::resolve_warning_account`). +/// Informational placeholders are not observed, matching `check_session_lane`. +fn resolve_toast_account( + manager: &mut codexbar::notifications::NotificationManager, + provider: ProviderId, + scope: &WarningScope, + window: &str, + lane: &RateWindowSnapshot, + settings: &Settings, +) -> String { + if lane.is_informational { + return scope.key().to_string(); + } + let resets_at = lane + .resets_at + .as_deref() + .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) + .map(|date| date.with_timezone(&chrono::Utc)); + manager.resolve_warning_account( + provider, + scope, + window, + lane.used_percent, + resets_at, + settings, + ) +} + fn dispatch_quota_hooks( settings: &Settings, provider: ProviderId, diff --git a/apps/desktop-tauri/src-tauri/src/commands/warning_identity.rs b/apps/desktop-tauri/src-tauri/src/commands/warning_identity.rs index 7a9bd92b88..1d9e8d7efe 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/warning_identity.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/warning_identity.rs @@ -1,4 +1,5 @@ use codexbar::core::ProviderId; +use codexbar::notifications::WarningScope; #[derive(Debug, Clone, PartialEq, Eq)] enum WarningSourceLane { @@ -106,6 +107,21 @@ impl WarningIdentity { } } + /// How this identity takes part in threshold-warning continuity across + /// identity gaps. Only resolved emails and unresolved Claude CLI/OAuth + /// samples are bridged; token accounts and organizations stay independent. + pub(super) fn gap_scope(&self) -> WarningScope { + let key = self.threshold_key(); + match (&self.account, self.unresolved_key()) { + (WarningAccountState::Unresolved, _) => WarningScope::Unresolved(key), + (WarningAccountState::Email(_), Some(unresolved)) => WarningScope::Resolved { + account: key, + unresolved, + }, + _ => WarningScope::Independent(key), + } + } + pub(super) fn predictive_key(&self) -> Option { match &self.account { WarningAccountState::Token(id) => Some(format!("token-account:{}", id.as_hyphenated())), @@ -223,4 +239,44 @@ mod tests { Some("cli:person@example.com") ); } + + #[test] + fn gap_scope_bridges_only_resolved_email_and_unresolved_claude_sources() { + let unresolved = WarningIdentity::new(ProviderId::Claude, "cli", None, None, None); + assert_eq!( + unresolved.gap_scope(), + WarningScope::Unresolved("claude:cli:unknown".to_string()) + ); + let resolved = WarningIdentity::new( + ProviderId::Claude, + "cli", + Some("Person@Example.com"), + None, + None, + ); + assert_eq!( + resolved.gap_scope(), + WarningScope::Resolved { + account: "person@example.com".to_string(), + unresolved: "claude:cli:unknown".to_string(), + } + ); + let token = WarningIdentity::new( + ProviderId::Claude, + "cli", + Some("person@example.com"), + None, + Some(uuid::Uuid::nil()), + ); + assert!(matches!(token.gap_scope(), WarningScope::Independent(_))); + let org = WarningIdentity::new(ProviderId::Claude, "cli", None, Some("Acme"), None); + assert_eq!( + org.gap_scope(), + WarningScope::Independent("org:acme".to_string()) + ); + let web = WarningIdentity::new(ProviderId::Claude, "web", Some("a@b.c"), None, None); + assert!(matches!(web.gap_scope(), WarningScope::Independent(_))); + let codex = WarningIdentity::new(ProviderId::Codex, "cli", Some("a@b.c"), None, None); + assert!(matches!(codex.gap_scope(), WarningScope::Independent(_))); + } } diff --git a/rust/src/notifications.rs b/rust/src/notifications.rs index 94a109bf40..b7b4be743c 100755 --- a/rust/src/notifications.rs +++ b/rust/src/notifications.rs @@ -14,6 +14,10 @@ use crate::settings::Settings; use crate::sound::{NotificationSoundEvent, play_alert}; use chrono::{DateTime, Utc}; +mod identity_gaps; + +pub use identity_gaps::WarningScope; + /// Notification types #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum NotificationType { @@ -132,6 +136,10 @@ pub struct NotificationManager { previous_session_percent: std::collections::HashMap, predictive_warning_keys: std::collections::HashSet, deepseek_pricing_period: Option, + identity_gaps: identity_gaps::IdentityGapState, + /// Toasts that would have been shown; tests assert on this instead of popping real toasts. + #[cfg(test)] + toasts: std::cell::RefCell>, } impl NotificationManager { @@ -141,6 +149,9 @@ impl NotificationManager { previous_session_percent: std::collections::HashMap::new(), predictive_warning_keys: std::collections::HashSet::new(), deepseek_pricing_period: None, + identity_gaps: identity_gaps::IdentityGapState::default(), + #[cfg(test)] + toasts: std::cell::RefCell::new(Vec::new()), } } @@ -514,7 +525,12 @@ impl NotificationManager { Self::play_notification_sound(NotificationSoundEvent::StatusIssue, settings); } - #[cfg(target_os = "windows")] + #[cfg(test)] + fn show_toast(&self, title: &str, body: &str) { + self.toasts.borrow_mut().push(format!("{title}: {body}")); + } + + #[cfg(all(target_os = "windows", not(test)))] fn show_toast(&self, title: &str, body: &str) { use std::os::windows::process::CommandExt; use std::process::Command; @@ -578,7 +594,7 @@ impl NotificationManager { } } - #[cfg(not(target_os = "windows"))] + #[cfg(all(not(target_os = "windows"), not(test)))] fn show_toast(&self, title: &str, body: &str) { use std::process::Command; diff --git a/rust/src/notifications/identity_gaps.rs b/rust/src/notifications/identity_gaps.rs new file mode 100644 index 0000000000..1eaf4af154 --- /dev/null +++ b/rust/src/notifications/identity_gaps.rs @@ -0,0 +1,264 @@ +//! Threshold-warning continuity across account-identity gaps. +//! +//! Some sources (Claude CLI / OAuth) alternate between samples that resolve an +//! account and samples that cannot. Keyed independently, each identity would +//! fire the same threshold toast once. This module keeps one warning episode +//! across those gaps (upstream CodexBar 0.66.0, `UsageStore+QuotaWarnings`): +//! +//! - A resolved sample moves any fallback ("unresolved") history for the same +//! lane onto the account and marks the account as `shared_with_unresolved`. +//! - An unresolved sample with no fallback history yet follows the most recently +//! known account when the window `resets_at` is unchanged and usage has not +//! dropped; otherwise it keeps the already-fired thresholds if the account +//! history has previously joined the fallback, and starts an independent +//! episode if not. +//! +//! The manager stays provider-agnostic: the caller decides which identities are +//! resolved, unresolved or independent via [`WarningScope`]. + +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; + +use super::{NotificationManager, ThresholdKey}; +use crate::core::ProviderId; +use crate::settings::Settings; + +/// How a sample's account key relates to identity-gap continuity. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WarningScope { + /// No gap semantics: the key is used as-is and never merged. + Independent(String), + /// A resolved account. `unresolved` is the fallback key whose history this + /// account absorbs (the same source lane's unresolved scope). + Resolved { account: String, unresolved: String }, + /// The source could not resolve an account; carries its fallback key. + Unresolved(String), +} + +impl WarningScope { + pub fn key(&self) -> &str { + match self { + Self::Independent(key) | Self::Unresolved(key) => key, + Self::Resolved { account, .. } => account, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct LaneKey { + provider: ProviderId, + account: String, + window: String, +} + +impl LaneKey { + fn new(provider: ProviderId, account: &str, window: &str) -> Self { + Self { + provider, + account: account.to_string(), + window: window.to_string(), + } + } +} + +/// Last observation for one (provider, account, window) lane. +#[derive(Debug, Clone)] +struct LaneEpisode { + resets_at: Option>, + used_percent: f64, + /// Monotonic observation order; the newer episode wins when merging. + sequence: u64, + shared_with_unresolved: bool, +} + +#[derive(Debug, Default)] +pub(super) struct IdentityGapState { + last_known: HashMap, + lanes: HashMap, + sequence: u64, +} + +impl NotificationManager { + /// Resolve which account key a lane sample should be deduped under. + /// + /// Call once per non-informational lane sample, before `check_and_notify` / + /// `check_session_lane`, and pass the returned key to them. + pub fn resolve_warning_account( + &mut self, + provider: ProviderId, + scope: &WarningScope, + window: &str, + used_percent: f64, + resets_at: Option>, + settings: &Settings, + ) -> String { + if matches!(scope, WarningScope::Independent(_)) || !settings.show_notifications { + return scope.key().to_string(); + } + let key = match scope { + WarningScope::Resolved { + account, + unresolved, + } => { + self.merge_unresolved_into_account(provider, unresolved, account); + self.identity_gaps + .last_known + .insert(provider, account.clone()); + account.clone() + } + WarningScope::Unresolved(unresolved) => { + self.unresolved_lane_account(provider, unresolved, window, used_percent, resets_at) + } + WarningScope::Independent(key) => key.clone(), + }; + self.record_lane(provider, &key, window, used_percent, resets_at); + key + } + + fn unresolved_lane_account( + &mut self, + provider: ProviderId, + unresolved: &str, + window: &str, + used_percent: f64, + resets_at: Option>, + ) -> String { + let own_lane = LaneKey::new(provider, unresolved, window); + if self.identity_gaps.lanes.contains_key(&own_lane) { + return unresolved.to_string(); + } + let Some(account) = self.identity_gaps.last_known.get(&provider).cloned() else { + return unresolved.to_string(); + }; + let Some(prior) = self + .identity_gaps + .lanes + .get(&LaneKey::new(provider, &account, window)) + .cloned() + else { + return unresolved.to_string(); + }; + let continues_episode = resets_at.is_some() + && resets_at == prior.resets_at + && used_percent >= prior.used_percent; + if continues_episode { + return account; + } + if prior.shared_with_unresolved { + // A discontinuity may change the key but cannot discard reconciled thresholds. + self.copy_lane_state(provider, &account, unresolved, window); + self.identity_gaps.lanes.insert(own_lane, prior); + } + unresolved.to_string() + } + + fn record_lane( + &mut self, + provider: ProviderId, + account: &str, + window: &str, + used_percent: f64, + resets_at: Option>, + ) { + let gaps = &mut self.identity_gaps; + gaps.sequence += 1; + let lane = gaps + .lanes + .entry(LaneKey::new(provider, account, window)) + .or_insert_with(|| LaneEpisode { + resets_at, + used_percent, + sequence: 0, + shared_with_unresolved: false, + }); + lane.resets_at = resets_at; + lane.used_percent = used_percent; + lane.sequence = gaps.sequence; + } + + /// Move every lane's fired-threshold state from the fallback key onto the + /// account, keeping whichever episode was observed most recently. + fn merge_unresolved_into_account( + &mut self, + provider: ProviderId, + unresolved: &str, + account: &str, + ) { + let windows: Vec = self + .identity_gaps + .lanes + .keys() + .filter(|lane| lane.provider == provider && lane.account == unresolved) + .map(|lane| lane.window.clone()) + .collect(); + for window in windows { + let Some(prior) = self + .identity_gaps + .lanes + .remove(&LaneKey::new(provider, unresolved, &window)) + else { + continue; + }; + let account_lane = LaneKey::new(provider, account, &window); + let account_sequence = self + .identity_gaps + .lanes + .get(&account_lane) + .map(|lane| lane.sequence); + if account_sequence.is_none_or(|sequence| prior.sequence >= sequence) { + self.copy_lane_state(provider, unresolved, account, &window); + self.identity_gaps.lanes.insert(account_lane.clone(), prior); + } + if let Some(lane) = self.identity_gaps.lanes.get_mut(&account_lane) { + lane.shared_with_unresolved = true; + } + self.drop_lane_state(provider, unresolved, &window); + } + } + + /// Replace `to`'s fired thresholds (and session-transition baseline) with `from`'s. + fn copy_lane_state(&mut self, provider: ProviderId, from: &str, to: &str, window: &str) { + let fired: Vec = self + .sent_notifications + .iter() + .filter(|(p, a, w, _)| *p == provider && a == from && w == window) + .cloned() + .collect(); + self.sent_notifications + .retain(|(p, a, w, _)| *p != provider || a != to || w != window); + self.sent_notifications.extend( + fired + .into_iter() + .map(|(p, _, w, t)| (p, to.to_string(), w, t)), + ); + if window == "session" { + match self + .previous_session_percent + .get(&(provider, from.to_string())) + .copied() + { + Some(percent) => { + self.previous_session_percent + .insert((provider, to.to_string()), percent); + } + None => { + self.previous_session_percent + .remove(&(provider, to.to_string())); + } + } + } + } + + fn drop_lane_state(&mut self, provider: ProviderId, account: &str, window: &str) { + self.sent_notifications + .retain(|(p, a, w, _)| *p != provider || a != account || w != window); + if window == "session" { + self.previous_session_percent + .remove(&(provider, account.to_string())); + } + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/notifications/identity_gaps/tests.rs b/rust/src/notifications/identity_gaps/tests.rs new file mode 100644 index 0000000000..4ffd1ac9bb --- /dev/null +++ b/rust/src/notifications/identity_gaps/tests.rs @@ -0,0 +1,283 @@ +//! Sequences mirror upstream `ClaudeCredentialQuotaWarningTests` (`checkIdentitySamples`). +//! Upstream thresholds are remaining-percent (50 / 20); locally high = 50% used and +//! critical = 80% used, so `remaining` values map to `100 - remaining` used. + +use super::*; +use crate::notifications::NotificationType; +use chrono::TimeZone; + +const CLI_UNKNOWN: &str = "claude:cli:unknown"; +const ACCOUNT_A: &str = "account-a@example.com"; +const ACCOUNT_B: &str = "account-b@example.com"; + +#[derive(Clone, Copy)] +enum Lane { + Session, + Weekly, +} + +impl Lane { + fn window(self) -> &'static str { + match self { + Lane::Session => "session", + Lane::Weekly => "weekly", + } + } +} + +struct Sample { + account: Option<&'static str>, + remaining: f64, + reset_offset_secs: Option, +} + +fn sample(account: Option<&'static str>, remaining: f64, reset: Option) -> Sample { + Sample { + account, + remaining, + reset_offset_secs: reset, + } +} + +fn settings() -> Settings { + Settings { + high_usage_threshold: 50.0, + critical_usage_threshold: 80.0, + sound_enabled: false, + ..Settings::default() + } +} + +fn scope_for(account: Option<&str>) -> WarningScope { + match account { + Some(account) => WarningScope::Resolved { + account: account.to_string(), + unresolved: CLI_UNKNOWN.to_string(), + }, + None => WarningScope::Unresolved(CLI_UNKNOWN.to_string()), + } +} + +/// Feed samples the way the shell does and return the toast thresholds fired +/// (`50` for the high toast, `20` for the critical toast). +fn run_samples(samples: &[Sample], lane: Lane) -> Vec { + let settings = settings(); + let mut manager = NotificationManager::new(); + let base = Utc.timestamp_opt(1_900_000_000, 0).unwrap(); + for sample in samples { + let used = 100.0 - sample.remaining; + let resets_at = sample + .reset_offset_secs + .map(|secs| base + chrono::Duration::seconds(secs)); + let account = manager.resolve_warning_account( + ProviderId::Claude, + &scope_for(sample.account), + lane.window(), + used, + resets_at, + &settings, + ); + manager.check_and_notify(ProviderId::Claude, &account, lane.window(), used, &settings); + } + let toasts = manager.toasts.borrow(); + toasts + .iter() + .map(|toast| { + if toast.starts_with(NotificationType::CriticalUsage.title()) { + 20 + } else if toast.starts_with(NotificationType::HighUsage.title()) { + 50 + } else { + panic!("unexpected toast {toast}") + } + }) + .collect() +} + +fn alternating(remaining: &[f64], reset: Option) -> Vec { + remaining + .iter() + .enumerate() + .map(|(index, value)| sample((index % 2 == 0).then_some(ACCOUNT_A), *value, reset)) + .collect() +} + +#[test] +fn repeated_identity_gaps_preserve_threshold_history() { + for lane in [Lane::Session, Lane::Weekly] { + for has_reset in [true, false] { + let samples = alternating( + &[49.0, 48.0, 47.0, 46.0, 45.0, 44.0], + has_reset.then_some(3600), + ); + let expected: &[u32] = if has_reset { &[50] } else { &[50, 50] }; + assert_eq!(run_samples(&samples, lane), expected, "reset={has_reset}"); + } + } +} + +#[test] +fn later_thresholds_do_not_repeat_across_either_identity_key() { + for has_reset in [true, false] { + for crossing_is_known in [true, false] { + let remaining: &[f64] = if crossing_is_known { + &[49.0, 48.0, 19.0, 18.0, 17.0, 16.0] + } else { + &[49.0, 48.0, 47.0, 19.0, 18.0, 17.0] + }; + let expected: &[u32] = if has_reset { &[50, 20] } else { &[50, 50, 20] }; + let samples = alternating(remaining, has_reset.then_some(3600)); + assert_eq!( + run_samples(&samples, Lane::Session), + expected, + "reset={has_reset} crossing_is_known={crossing_is_known}" + ); + } + } +} + +#[test] +fn discontinuous_identity_gaps_start_one_independent_episode() { + for discontinuity in ["reset", "increase", "missing"] { + let reset = match discontinuity { + "missing" => None, + "reset" => Some(7200), + _ => Some(3600), + }; + let unresolved = |remaining: f64| { + sample( + None, + if discontinuity == "increase" { + 49.0 + } else { + remaining + }, + reset, + ) + }; + let samples = [ + sample(Some(ACCOUNT_A), 40.0, Some(3600)), + unresolved(39.0), + sample(Some(ACCOUNT_A), 38.0, Some(3600)), + unresolved(37.0), + sample(Some(ACCOUNT_A), 36.0, Some(3600)), + unresolved(35.0), + ]; + assert_eq!( + run_samples(&samples, Lane::Session), + [50, 50], + "{discontinuity}" + ); + } +} + +#[test] +fn identity_gaps_follow_the_most_recent_account_without_merging_known_accounts() { + let samples = [ + sample(Some(ACCOUNT_A), 49.0, Some(3600)), + sample(Some(ACCOUNT_B), 48.0, Some(3600)), + sample(None, 47.0, Some(3600)), + sample(Some(ACCOUNT_A), 46.0, Some(3600)), + sample(None, 45.0, Some(3600)), + sample(Some(ACCOUNT_B), 44.0, Some(3600)), + sample(None, 43.0, Some(3600)), + ]; + assert_eq!(run_samples(&samples, Lane::Session), [50, 50]); +} + +#[test] +fn initial_unresolved_history_survives_repeated_resolution_and_recovery() { + for has_reset in [true, false] { + let reset = has_reset.then_some(3600); + let samples = [ + sample(None, 49.0, reset), + sample(Some(ACCOUNT_A), 48.0, reset), + sample(None, 47.0, reset), + sample(Some(ACCOUNT_A), 19.0, reset), + sample(None, 18.0, reset), + sample(Some(ACCOUNT_A), 60.0, reset), + sample(None, 49.0, reset), + sample(Some(ACCOUNT_A), 48.0, reset), + sample(None, 47.0, reset), + ]; + assert_eq!( + run_samples(&samples, Lane::Session), + [50, 20, 50], + "reset={has_reset}" + ); + } +} + +#[test] +fn independent_scopes_are_never_merged() { + let settings = settings(); + let mut manager = NotificationManager::new(); + let scope = WarningScope::Independent("token-account:1".to_string()); + assert_eq!( + manager.resolve_warning_account( + ProviderId::Claude, + &scope, + "session", + 60.0, + None, + &settings + ), + "token-account:1" + ); + assert!(manager.identity_gaps.last_known.is_empty()); + assert!(manager.identity_gaps.lanes.is_empty()); +} + +#[test] +fn disabled_notifications_leave_scope_untouched() { + let settings = Settings { + show_notifications: false, + ..settings() + }; + let mut manager = NotificationManager::new(); + let key = manager.resolve_warning_account( + ProviderId::Claude, + &scope_for(None), + "session", + 60.0, + None, + &settings, + ); + assert_eq!(key, CLI_UNKNOWN); + assert!(manager.identity_gaps.lanes.is_empty()); +} + +#[test] +fn session_depleted_state_follows_the_merged_account() { + let settings = settings(); + let mut manager = NotificationManager::new(); + let base = Utc.timestamp_opt(1_900_000_000, 0).unwrap(); + let reset = Some(base + chrono::Duration::hours(1)); + + let unresolved = manager.resolve_warning_account( + ProviderId::Claude, + &scope_for(None), + "session", + 100.0, + reset, + &settings, + ); + manager.check_session_lane(ProviderId::Claude, &unresolved, 100.0, false, &settings); + let account = manager.resolve_warning_account( + ProviderId::Claude, + &scope_for(Some(ACCOUNT_A)), + "session", + 100.0, + reset, + &settings, + ); + manager.check_session_lane(ProviderId::Claude, &account, 100.0, false, &settings); + + let depleted = manager + .toasts + .borrow() + .iter() + .filter(|toast| toast.starts_with(NotificationType::SessionDepleted.title())) + .count(); + assert_eq!(depleted, 1, "depleted toast must not repeat after merge"); +}