From 6341319a2eb85a45a81a1dce274d1ce1a306d045 Mon Sep 17 00:00:00 2001 From: NessZerra Date: Sun, 12 Jul 2026 08:51:34 +0700 Subject: [PATCH 1/3] Add predictive usage warnings Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src-tauri/src/commands/bridge.rs | 4 + .../src-tauri/src/commands/providers.rs | 181 ++++++++- .../src-tauri/src/commands/settings.rs | 9 + apps/desktop-tauri/src/App.test.tsx | 2 + .../src/components/FormControls.tsx | 3 + .../src/components/MenuCard.test.tsx | 26 +- .../desktop-tauri/src/components/MenuCard.tsx | 18 +- .../src/floatbar/FloatBar.test.tsx | 2 + .../src/hooks/useFormattedResetTime.ts | 2 +- apps/desktop-tauri/src/i18n/keys.ts | 6 + .../src/surfaces/PopOutPanel.test.tsx | 2 + .../src/surfaces/PopOutPanel.tsx | 1 + .../src/surfaces/TrayPanel.test.tsx | 2 + apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 1 + .../surfaces/settings/tabs/AboutTab.test.tsx | 2 + .../settings/tabs/DisplayTab.test.tsx | 10 + .../src/surfaces/settings/tabs/DisplayTab.tsx | 12 + .../settings/tabs/GeneralTab.test.tsx | 13 +- .../src/surfaces/settings/tabs/GeneralTab.tsx | 12 + apps/desktop-tauri/src/types/bridge.test.ts | 2 + apps/desktop-tauri/src/types/bridge.ts | 4 + rust/src/core/usage_pace.rs | 166 ++++++++ rust/src/locale.rs | 6 + rust/src/locale/en-US.ftl | 6 + rust/src/locale/es-MX.ftl | 6 + rust/src/locale/ja-JP.ftl | 6 + rust/src/locale/ko-KR.ftl | 6 + rust/src/locale/zh-CN.ftl | 6 + rust/src/locale/zh-TW.ftl | 6 + rust/src/notifications.rs | 356 ++++++++++++++++++ rust/src/settings.rs | 10 + rust/src/settings/raw.rs | 6 + rust/src/settings/tests.rs | 16 + 33 files changed, 903 insertions(+), 7 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index f80945326c..54b58df5af 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -410,6 +410,7 @@ pub struct SettingsSnapshot { sound_volume: u8, high_usage_threshold: f64, critical_usage_threshold: f64, + predictive_pace_warning_enabled: bool, tray_icon_mode: &'static str, switcher_shows_icons: bool, menu_bar_shows_highest_usage: bool, @@ -418,6 +419,7 @@ pub struct SettingsSnapshot { show_all_token_accounts_in_menu: bool, enable_animations: bool, reset_time_relative: bool, + show_reset_when_exhausted: bool, menu_bar_display_mode: String, hide_personal_info: bool, update_channel: &'static str, @@ -495,6 +497,7 @@ impl From for SettingsSnapshot { sound_volume: settings.sound_volume, high_usage_threshold: settings.high_usage_threshold, critical_usage_threshold: settings.critical_usage_threshold, + predictive_pace_warning_enabled: settings.predictive_pace_warning_enabled, tray_icon_mode: tray_icon_mode_label(settings.tray_icon_mode), switcher_shows_icons: settings.switcher_shows_icons, menu_bar_shows_highest_usage: settings.menu_bar_shows_highest_usage, @@ -503,6 +506,7 @@ impl From for SettingsSnapshot { show_all_token_accounts_in_menu: settings.show_all_token_accounts_in_menu, enable_animations: settings.enable_animations, reset_time_relative: settings.reset_time_relative, + show_reset_when_exhausted: settings.show_reset_when_exhausted, menu_bar_display_mode: settings.menu_bar_display_mode, hide_personal_info: settings.hide_personal_info, update_channel: update_channel_label(settings.update_channel), diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index ff0b1c8d93..d1bf398376 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -175,7 +175,12 @@ async fn do_refresh_providers_with_policy( await_provider_refreshes(handles).await; let error_count = finish_provider_refresh(&state)?; - update_tray_and_notifications(app, &state, &inputs.settings)?; + update_tray_and_notifications( + app, + &state, + &inputs.settings, + &inputs.token_accounts, + )?; events::emit_refresh_complete(app, enabled_count, error_count); @@ -390,6 +395,7 @@ fn update_tray_and_notifications( app: &tauri::AppHandle, state: &tauri::State<'_, Mutex>, settings: &Settings, + token_accounts: &HashMap, ) -> Result<(), String> { let cached = { let guard = state.lock().map_err(|e| e.to_string())?; @@ -397,13 +403,14 @@ fn update_tray_and_notifications( }; crate::tray_bridge::update_tray_status_items(app, &cached); crate::tray_bridge::update_tray_icon_and_tooltip(app, &cached); - notify_usage_thresholds(state, settings, &cached); + notify_usage_thresholds(state, settings, token_accounts, &cached); Ok(()) } fn notify_usage_thresholds( state: &tauri::State<'_, Mutex>, settings: &Settings, + token_accounts: &HashMap, cached: &[ProviderUsageSnapshot], ) { let cli_map = codexbar::core::cli_name_map(); @@ -422,11 +429,127 @@ fn notify_usage_thresholds( snapshot.primary.used_percent, settings, ); + notify_predictive_pace( + &mut guard.notification_manager, + provider, + snapshot, + token_accounts, + settings, + ); } } } } +fn notify_predictive_pace( + manager: &mut codexbar::notifications::NotificationManager, + provider: ProviderId, + snapshot: &ProviderUsageSnapshot, + token_accounts: &HashMap, + settings: &Settings, +) { + let enabled = settings.show_notifications && settings.predictive_pace_warning_enabled; + manager.set_predictive_warnings_enabled(provider, enabled); + if !enabled || !matches!(provider, ProviderId::Claude | ProviderId::Codex) { + return; + } + + let token_account_id = token_accounts + .get(&provider) + .and_then(ProviderAccountData::active_account) + .map(|account| account.id); + let Some(identity) = predictive_warning_identity( + provider, + &snapshot.source_label, + snapshot.account_email.as_deref(), + token_account_id, + ) else { + return; + }; + let observed_at = chrono::DateTime::parse_from_rfc3339(&snapshot.updated_at) + .ok() + .map(|date| date.with_timezone(&chrono::Utc)); + + notify_predictive_window( + manager, + provider, + &identity, + codexbar::notifications::PredictiveWarningWindow::Session, + &snapshot.primary, + observed_at, + 300, + settings, + ); + if let Some(weekly) = &snapshot.secondary { + notify_predictive_window( + manager, + provider, + &identity, + codexbar::notifications::PredictiveWarningWindow::Weekly, + weekly, + observed_at, + 10080, + settings, + ); + } +} + +#[allow(clippy::too_many_arguments)] +fn notify_predictive_window( + manager: &mut codexbar::notifications::NotificationManager, + provider: ProviderId, + identity: &str, + warning_window: codexbar::notifications::PredictiveWarningWindow, + window: &RateWindowSnapshot, + observed_at: Option>, + default_window_minutes: u32, + settings: &Settings, +) { + let rate_window = RateWindow::with_details( + window.used_percent, + window.window_minutes, + window + .resets_at + .as_deref() + .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) + .map(|date| date.with_timezone(&chrono::Utc)), + window.reset_description.clone(), + ); + let Some(pace) = + codexbar::core::UsagePace::weekly(&rate_window, observed_at, default_window_minutes) + else { + return; + }; + manager.check_predictive_pace( + provider, + identity, + warning_window, + &rate_window, + &pace, + settings, + ); +} + +fn predictive_warning_identity( + provider: ProviderId, + source_label: &str, + account_email: Option<&str>, + token_account_id: Option, +) -> Option { + if !matches!(provider, ProviderId::Claude | ProviderId::Codex) { + return None; + } + if let Some(id) = token_account_id { + return Some(format!("token-account:{}", id.as_hyphenated())); + } + let source = source_label.trim().to_ascii_lowercase(); + let account = account_email?.trim().to_ascii_lowercase(); + if source.is_empty() || account.is_empty() { + return None; + } + Some(format!("{source}:{account}")) +} + #[tauri::command] pub async fn refresh_providers(app: tauri::AppHandle) -> Result<(), String> { do_refresh_providers(&app).await @@ -449,5 +572,59 @@ pub fn get_cached_providers( for snapshot in &mut snapshots { super::filter_hidden_codex_spark_rows(snapshot, spark_usage_visible); } + snapshots } + +#[cfg(test)] +mod predictive_warning_tests { + use super::*; + + #[test] + fn predictive_warning_identity_keeps_claude_sources_and_token_accounts_separate() { + let account_id = uuid::Uuid::parse_str("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa").unwrap(); + + assert_eq!( + predictive_warning_identity( + ProviderId::Claude, + "cli", + Some("Person@Example.com"), + None, + ) + .as_deref(), + Some("cli:person@example.com") + ); + assert_eq!( + predictive_warning_identity( + ProviderId::Claude, + "oauth", + Some("Person@Example.com"), + None, + ) + .as_deref(), + Some("oauth:person@example.com") + ); + assert_eq!( + predictive_warning_identity( + ProviderId::Claude, + "oauth", + Some("Person@Example.com"), + Some(account_id), + ) + .as_deref(), + Some("token-account:aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") + ); + } + + #[test] + fn predictive_warning_identity_skips_unidentified_accounts() { + assert_eq!( + predictive_warning_identity(ProviderId::Claude, "oauth", None, None), + None + ); + assert_eq!( + predictive_warning_identity(ProviderId::Codex, "cli", Some(" "), None), + None + ); + } +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index 115a796297..45bc905f3a 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -17,6 +17,7 @@ pub struct SettingsUpdate { pub sound_volume: Option, pub high_usage_threshold: Option, pub critical_usage_threshold: Option, + pub predictive_pace_warning_enabled: Option, pub tray_icon_mode: Option, pub switcher_shows_icons: Option, pub menu_bar_shows_highest_usage: Option, @@ -25,6 +26,7 @@ pub struct SettingsUpdate { pub show_all_token_accounts_in_menu: Option, pub enable_animations: Option, pub reset_time_relative: Option, + pub show_reset_when_exhausted: Option, pub menu_bar_display_mode: Option, pub hide_personal_info: Option, pub update_channel: Option, @@ -62,6 +64,7 @@ impl SettingsUpdate { || self.critical_usage_threshold.is_some() || self.show_as_used.is_some() || self.reset_time_relative.is_some() + || self.show_reset_when_exhausted.is_some() } fn rebuilds_tray_menu(&self) -> bool { @@ -149,6 +152,9 @@ impl SettingsUpdate { if let Some(v) = self.reset_time_relative { settings.reset_time_relative = v; } + if let Some(v) = self.show_reset_when_exhausted { + settings.show_reset_when_exhausted = v; + } if let Some(v) = self.menu_bar_display_mode.clone() { settings.menu_bar_display_mode = v; } @@ -189,6 +195,9 @@ impl SettingsUpdate { if let Some(v) = self.critical_usage_threshold { settings.critical_usage_threshold = v.clamp(0.0, 100.0); } + if let Some(v) = self.predictive_pace_warning_enabled { + settings.predictive_pace_warning_enabled = v; + } self } diff --git a/apps/desktop-tauri/src/App.test.tsx b/apps/desktop-tauri/src/App.test.tsx index 68444ba103..63943167d2 100644 --- a/apps/desktop-tauri/src/App.test.tsx +++ b/apps/desktop-tauri/src/App.test.tsx @@ -71,6 +71,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, @@ -79,6 +80,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { showAllTokenAccountsInMenu: false, enableAnimations: true, resetTimeRelative: true, + showResetWhenExhausted: false, menuBarDisplayMode: "detailed", hidePersonalInfo: false, updateChannel: "stable", diff --git a/apps/desktop-tauri/src/components/FormControls.tsx b/apps/desktop-tauri/src/components/FormControls.tsx index f11d8b573b..5c02fa7979 100644 --- a/apps/desktop-tauri/src/components/FormControls.tsx +++ b/apps/desktop-tauri/src/components/FormControls.tsx @@ -6,11 +6,13 @@ export function Toggle({ checked, onChange, label, + ariaLabel, disabled, }: { checked: boolean; onChange: (v: boolean) => void; label?: string; + ariaLabel?: string; disabled?: boolean; }) { const input = ( @@ -18,6 +20,7 @@ export function Toggle({ type="checkbox" className="toggle" checked={checked} + aria-label={ariaLabel} disabled={disabled} onChange={(e) => onChange(e.target.checked)} /> diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index 502a13666d..a70c8fd211 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -78,7 +78,11 @@ function provider( function renderCard( snapshot: ProviderUsageSnapshot, - opts: { showAsUsed?: boolean; onLayoutChange?: () => void } = {}, + opts: { + showAsUsed?: boolean; + showResetWhenExhausted?: boolean; + onLayoutChange?: () => void; + } = {}, ) { return render( @@ -87,6 +91,7 @@ function renderCard( hideEmail={false} resetTimeRelative={true} showAsUsed={opts.showAsUsed} + showResetWhenExhausted={opts.showResetWhenExhausted} onLayoutChange={opts.onLayoutChange} /> , @@ -110,6 +115,7 @@ describe("MenuCard", () => { PanelThirtyDayTokens: "30d tokens", PanelTodayBudget: "today", PanelUsedSuffix: "used", + ResetsInHoursMinutes: "Resets in {}h {}m", }), ); tauriMocks.getProviderChartData.mockResolvedValue({ @@ -169,6 +175,24 @@ describe("MenuCard", () => { expect(fill?.style.width).toBe("100%"); }); + it("replaces an exhausted percentage with a future reset countdown", async () => { + const snapshot = provider(null, 100, { exhausted: true }); + snapshot.primary.resetsAt = new Date(Date.now() + 60 * 60 * 1000).toISOString(); + + renderCard(snapshot, { showResetWhenExhausted: true }); + + expect(await screen.findByText(/Resets in 0h/)).toBeInTheDocument(); + expect(screen.queryByText("0% left")).not.toBeInTheDocument(); + }); + + it("keeps an exhausted percentage without a concrete future reset", async () => { + renderCard(provider(null, 100, { exhausted: true, resetDescription: "in 2h" }), { + showResetWhenExhausted: true, + }); + + expect(await screen.findByText("0% left")).toBeInTheDocument(); + }); + it("renders additional Copilot budget windows", async () => { const snapshot = provider(null, 20); snapshot.providerId = "copilot"; diff --git a/apps/desktop-tauri/src/components/MenuCard.tsx b/apps/desktop-tauri/src/components/MenuCard.tsx index 2415346d25..8bf988dc06 100644 --- a/apps/desktop-tauri/src/components/MenuCard.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.tsx @@ -49,6 +49,7 @@ interface MenuCardProps { provider: ProviderUsageSnapshot; hideEmail: boolean; resetTimeRelative: boolean; + showResetWhenExhausted?: boolean; showAsUsed?: boolean; compactMetrics?: boolean; onLayoutChange?: () => void; @@ -256,6 +257,7 @@ function MetricRow({ snap, exhaustedLabel, resetTimeRelative, + showResetWhenExhausted, showAsUsed, expanded, onToggleExpanded, @@ -264,6 +266,7 @@ function MetricRow({ snap: RateWindowSnapshot; exhaustedLabel: string; resetTimeRelative: boolean; + showResetWhenExhausted: boolean; showAsUsed: boolean; expanded: boolean; onToggleExpanded: () => void; @@ -281,6 +284,13 @@ function MetricRow({ snap.resetDescription, resetTimeRelative, ); + const resetTarget = snap.resetsAt ? Date.parse(snap.resetsAt) : Number.NaN; + const replacesPercent = + showResetWhenExhausted && + snap.isExhausted && + Number.isFinite(resetTarget) && + resetTarget > Date.now() && + resetText !== null; const paceView = getMetricPaceView(snap); const reserveDescription = formatReserveDescription(snap, t); const formatBudget = (value: number) => @@ -292,8 +302,10 @@ function MetricRow({
- {Math.round(displayPct)}% {displayLabel} - {resetText && ( + + {replacesPercent ? resetText : `${Math.round(displayPct)}% ${displayLabel}`} + + {resetText && !replacesPercent && ( {resetText} )}
@@ -359,6 +371,7 @@ export default function MenuCard({ provider, hideEmail, resetTimeRelative, + showResetWhenExhausted = false, showAsUsed = false, compactMetrics = false, onLayoutChange, @@ -501,6 +514,7 @@ export default function MenuCard({ snap={m.snap} exhaustedLabel={t("DetailWindowExhausted")} resetTimeRelative={resetTimeRelative} + showResetWhenExhausted={showResetWhenExhausted} showAsUsed={showAsUsed} expanded={expandedPaceWindow === m.id} onToggleExpanded={() => { diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx index cfed699510..07c24832a4 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx @@ -100,6 +100,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, @@ -108,6 +109,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { showAllTokenAccountsInMenu: false, enableAnimations: true, resetTimeRelative: true, + showResetWhenExhausted: false, menuBarDisplayMode: "detailed", hidePersonalInfo: false, updateChannel: "stable", diff --git a/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts b/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts index 14f8d70f9f..7c19869a98 100644 --- a/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts +++ b/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts @@ -25,7 +25,7 @@ export function useFormattedResetTime( const [now, setNow] = useState(() => Date.now()); useEffect(() => { - if (!resetsAt || !relative) return; + if (!resetsAt) return; const id = window.setInterval(() => setNow(Date.now()), 30_000); return () => window.clearInterval(id); }, [resetsAt, relative]); diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 4fb4d007c9..f040303f6d 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -30,11 +30,17 @@ export const ALL_LOCALE_KEYS = [ "HighUsageAlert", "CriticalUsageThreshold", "CriticalUsageAlert", + "PredictivePaceWarnings", + "PredictivePaceWarningsHelper", + "PredictivePaceWarningTitle", + "PredictivePaceWarningBody", "UsageDisplay", "ShowUsageAsUsed", "ShowUsageAsUsedHelper", "ResetTimeRelative", "ResetTimeRelativeHelper", + "ShowResetWhenExhausted", + "ShowResetWhenExhaustedHelper", "TrayIcon", "MergeTrayIcons", "MergeTrayIconsHelper", diff --git a/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx b/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx index f51d1edf90..62820a65a6 100644 --- a/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx @@ -126,6 +126,7 @@ function settings(): SettingsSnapshot { soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, @@ -134,6 +135,7 @@ function settings(): SettingsSnapshot { showAllTokenAccountsInMenu: false, enableAnimations: true, resetTimeRelative: true, + showResetWhenExhausted: false, menuBarDisplayMode: "detailed", hidePersonalInfo: false, updateChannel: "stable", diff --git a/apps/desktop-tauri/src/surfaces/PopOutPanel.tsx b/apps/desktop-tauri/src/surfaces/PopOutPanel.tsx index ad8892c5d0..a3984446bb 100644 --- a/apps/desktop-tauri/src/surfaces/PopOutPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/PopOutPanel.tsx @@ -253,6 +253,7 @@ export default function PopOutPanel({ provider={p} hideEmail={settings.hidePersonalInfo} resetTimeRelative={settings.resetTimeRelative} + showResetWhenExhausted={settings.showResetWhenExhausted} showAsUsed={settings.showAsUsed} compactMetrics={selectedProviderId === null} /> diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 1d8663efa1..7a8419632f 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -112,6 +112,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, @@ -120,6 +121,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { showAllTokenAccountsInMenu: false, enableAnimations: true, resetTimeRelative: true, + showResetWhenExhausted: false, menuBarDisplayMode: "detailed", hidePersonalInfo: false, updateChannel: "stable", diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index a49509c80a..d0d1579051 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -433,6 +433,7 @@ export default function TrayPanel({ state }: { state: BootstrapState }) { provider={p} hideEmail={settings.hidePersonalInfo} resetTimeRelative={settings.resetTimeRelative} + showResetWhenExhausted={settings.showResetWhenExhausted} showAsUsed={settings.showAsUsed} compactMetrics={selectedProviderId === null} onLayoutChange={requestLayout} diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx index d6cd1bba04..7de93b5f97 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx @@ -48,6 +48,7 @@ const settings: SettingsSnapshot = { soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, @@ -56,6 +57,7 @@ const settings: SettingsSnapshot = { showAllTokenAccountsInMenu: true, enableAnimations: true, resetTimeRelative: true, + showResetWhenExhausted: false, menuBarDisplayMode: "compact", hidePersonalInfo: false, autoDownloadUpdates: false, diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx index 076a56a24a..1cf61e4825 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx @@ -23,6 +23,7 @@ const baseSettings = { showAsUsed: false, showAllTokenAccountsInMenu: false, resetTimeRelative: false, + showResetWhenExhausted: false, } as unknown as SettingsSnapshot; function renderTab(set: (patch: Record) => void) { @@ -53,4 +54,13 @@ describe("DisplayTab window scale", () => { expect(set).not.toHaveBeenCalled(); }); + + it("updates the exhausted reset display preference", () => { + const set = vi.fn(); + renderTab(set); + + fireEvent.click(screen.getByRole("checkbox", { name: "ShowResetWhenExhausted" })); + + expect(set).toHaveBeenCalledWith({ showResetWhenExhausted: true }); + }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx index eabec6c3cd..782e4ab41e 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx @@ -158,6 +158,18 @@ export default function DisplayTab({ settings, set, saving }: TabProps) { onChange={(v) => set({ resetTimeRelative: v })} /> + + set({ showResetWhenExhausted: v })} + /> + diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx index 9699f8083e..8a4cb5a1e0 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx @@ -1,4 +1,4 @@ -import { render, screen } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; vi.mock("../../../hooks/useLocale", () => ({ @@ -31,6 +31,7 @@ const settings: SettingsSnapshot = { soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, @@ -64,6 +65,7 @@ const settings: SettingsSnapshot = { floatBarProviderIds: [], floatBarDarkText: false, floatBarShowResetInline: false, + showResetWhenExhausted: false, }; describe("GeneralTab language picker", () => { @@ -98,4 +100,13 @@ describe("GeneralTab language picker", () => { expect(screen.getByText("繁體中文（臺灣）")).toBeInTheDocument(); }); + + it("updates the predictive pace warning preference", () => { + const set = vi.fn(); + render(); + + fireEvent.click(screen.getByRole("checkbox", { name: "PredictivePaceWarnings" })); + + expect(set).toHaveBeenCalledWith({ predictivePaceWarningEnabled: true }); + }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx index d73a4848fb..8262695dac 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx @@ -102,6 +102,18 @@ export default function GeneralTab({ settings, set, saving }: TabProps) { onChange={(v) => set({ showNotifications: v })} /> + + set({ predictivePaceWarningEnabled: v })} + /> +
{ soundVolume: 100, highUsageThreshold: 70, criticalUsageThreshold: 90, + predictivePaceWarningEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, @@ -57,6 +58,7 @@ describe("Language type", () => { showAllTokenAccountsInMenu: true, enableAnimations: true, resetTimeRelative: true, + showResetWhenExhausted: false, menuBarDisplayMode: "compact", windowScalePercent: 125, trayScalePercent: 100, diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 2179f5de49..9c76d55094 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -176,6 +176,7 @@ export interface SettingsSnapshot { soundVolume: number; highUsageThreshold: number; criticalUsageThreshold: number; + predictivePaceWarningEnabled: boolean; trayIconMode: TrayIconMode; switcherShowsIcons: boolean; menuBarShowsHighestUsage: boolean; @@ -184,6 +185,7 @@ export interface SettingsSnapshot { showAllTokenAccountsInMenu: boolean; enableAnimations: boolean; resetTimeRelative: boolean; + showResetWhenExhausted: boolean; menuBarDisplayMode: MenuBarDisplayMode; hidePersonalInfo: boolean; updateChannel: UpdateChannel; @@ -231,6 +233,7 @@ export interface SettingsUpdate { soundVolume?: number; highUsageThreshold?: number; criticalUsageThreshold?: number; + predictivePaceWarningEnabled?: boolean; trayIconMode?: TrayIconMode; switcherShowsIcons?: boolean; menuBarShowsHighestUsage?: boolean; @@ -239,6 +242,7 @@ export interface SettingsUpdate { showAllTokenAccountsInMenu?: boolean; enableAnimations?: boolean; resetTimeRelative?: boolean; + showResetWhenExhausted?: boolean; menuBarDisplayMode?: MenuBarDisplayMode; hidePersonalInfo?: boolean; updateChannel?: UpdateChannel; diff --git a/rust/src/core/usage_pace.rs b/rust/src/core/usage_pace.rs index 7d66ac1a90..e8affef2a1 100755 --- a/rust/src/core/usage_pace.rs +++ b/rust/src/core/usage_pace.rs @@ -81,6 +81,70 @@ pub struct UsagePace { pub eta_seconds: Option, /// Whether current pace will last until reset pub will_last_to_reset: bool, + /// Estimated probability of running out before reset, when a predictor provides one. + pub run_out_probability: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageDisplayMode { + Percent, + Pace, + Both, +} + +pub fn usage_display_text( + mode: UsageDisplayMode, + window: &RateWindow, + pace: Option<&UsagePace>, + show_used: bool, + show_reset_when_exhausted: bool, + now: DateTime, +) -> Option { + let percent = || { + let value = if show_used { + window.used_percent + } else { + window.remaining_percent() + }; + format!("{value:.0}%") + }; + + if show_reset_when_exhausted && window.is_exhausted() { + if let Some(resets_at) = window.resets_at.filter(|reset| *reset > now) { + let duration = resets_at - now; + let total_minutes = ((duration.num_seconds() + 59) / 60).max(1); + let days = total_minutes / 1440; + let hours = (total_minutes % 1440) / 60; + let minutes = total_minutes % 60; + let countdown = if days > 0 { + format!("{days}d {hours}h") + } else if hours > 0 { + format!("{hours}h {minutes}m") + } else { + format!("{minutes}m") + }; + return Some(format!("↻ in {countdown}")); + } + return Some(percent()); + } + + let pace_text = pace.map(|pace| { + let delta = pace.delta_percent.abs().round() as i64; + if delta == 0 { + "0%".to_string() + } else { + let sign = if pace.delta_percent >= 0.0 { "+" } else { "-" }; + format!("{sign}{delta}%") + } + }); + + Some(match mode { + UsageDisplayMode::Percent => percent(), + UsageDisplayMode::Pace => pace_text.unwrap_or_else(percent), + UsageDisplayMode::Both => pace_text + .map(|pace| format!("{} · {pace}", percent())) + .unwrap_or_else(percent), + }) } impl UsagePace { @@ -149,6 +213,7 @@ impl UsagePace { actual_used_percent: actual, eta_seconds, will_last_to_reset, + run_out_probability: None, }) } @@ -266,4 +331,105 @@ mod tests { assert_eq!(PaceStage::FarAhead.label(), "Far Ahead"); assert_eq!(PaceStage::SlightlyBehind.emoji(), "🐢"); } + + fn exhausted_window( + _now: DateTime, + resets_at: Option>, + reset_description: Option<&str>, + ) -> RateWindow { + RateWindow::with_details( + 100.0, + Some(300), + resets_at, + reset_description.map(str::to_string), + ) + } + + fn pace() -> UsagePace { + UsagePace { + stage: PaceStage::Ahead, + delta_percent: 12.0, + expected_used_percent: 40.0, + actual_used_percent: 52.0, + eta_seconds: Some(3600.0), + will_last_to_reset: false, + run_out_probability: None, + } + } + + #[test] + fn exhausted_display_uses_future_reset_for_percent_pace_and_both() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let window = exhausted_window( + now, + Some(now + Duration::hours(2) + Duration::minutes(15)), + None, + ); + + for mode in [ + UsageDisplayMode::Percent, + UsageDisplayMode::Pace, + UsageDisplayMode::Both, + ] { + assert_eq!( + usage_display_text(mode, &window, Some(&pace()), false, true, now).as_deref(), + Some("↻ in 2h 15m") + ); + } + } + + #[test] + fn exhausted_display_preserves_percent_without_schedulable_reset() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let cases = [ + exhausted_window(now, None, None), + exhausted_window(now, None, Some("in 2h")), + exhausted_window(now, Some(now - Duration::minutes(1)), None), + ]; + + for window in cases { + for mode in [ + UsageDisplayMode::Percent, + UsageDisplayMode::Pace, + UsageDisplayMode::Both, + ] { + assert_eq!( + usage_display_text(mode, &window, Some(&pace()), false, true, now).as_deref(), + Some("0%") + ); + } + } + } + + #[test] + fn exhausted_display_is_disabled_by_default_and_reverts_after_reset() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let resets_at = now + Duration::hours(1); + let window = exhausted_window(now, Some(resets_at), None); + + assert_eq!( + usage_display_text( + UsageDisplayMode::Percent, + &window, + None, + false, + false, + now + ) + .as_deref(), + Some("0%") + ); + assert_eq!( + usage_display_text( + UsageDisplayMode::Percent, + &window, + None, + false, + true, + resets_at + ) + .as_deref(), + Some("0%") + ); + } } diff --git a/rust/src/locale.rs b/rust/src/locale.rs index c1657667c1..dde2361d6c 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -188,6 +188,10 @@ locale_keys! { HighUsageAlert, CriticalUsageThreshold, CriticalUsageAlert, + PredictivePaceWarnings, + PredictivePaceWarningsHelper, + PredictivePaceWarningTitle, + PredictivePaceWarningBody, // Display settings (Preferences) UsageDisplay, @@ -195,6 +199,8 @@ locale_keys! { ShowUsageAsUsedHelper, ResetTimeRelative, ResetTimeRelativeHelper, + ShowResetWhenExhausted, + ShowResetWhenExhaustedHelper, TrayIcon, MergeTrayIcons, MergeTrayIconsHelper, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index 13a4239518..16b019003b 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -23,11 +23,17 @@ HighUsageThreshold = High Usage Threshold HighUsageAlert = High Usage Alert CriticalUsageThreshold = Critical Usage Threshold CriticalUsageAlert = Critical Alert +PredictivePaceWarnings = Predictive Pace Warnings +PredictivePaceWarningsHelper = Alert when Codex or Claude usage is likely to run out before reset +PredictivePaceWarningTitle = { "{}" } { "{}" } pace warning +PredictivePaceWarningBody = Quota may run out in { "{}" } UsageDisplay = Usage Display ShowUsageAsUsed = Show Usage as Used ShowUsageAsUsedHelper = Display as used percentage instead of remaining ResetTimeRelative = Relative Reset Time ResetTimeRelativeHelper = Show "2h 30m" instead of "3:00 PM" +ShowResetWhenExhausted = Show reset time when exhausted +ShowResetWhenExhaustedHelper = Replace an exhausted percentage with its live reset countdown TrayIcon = Tray Icon MergeTrayIcons = Merge Tray Icons MergeTrayIconsHelper = Show all providers in a single tray icon diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index d7a67d206f..f622000a40 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -559,3 +559,9 @@ TrayPaceBadgeRacing = Acelerado TrayPaceBadgeBurning = Crítico TrayResetsInLabel = Reinicia en { "{}" } TrayResetsDueNow = Reiniciando… +PredictivePaceWarnings = Avisos predictivos de ritmo +PredictivePaceWarningsHelper = Avisa cuando el uso de Codex o Claude probablemente se agote antes del reinicio +PredictivePaceWarningTitle = Aviso de ritmo de { "{}" } ({ "{}" }) +PredictivePaceWarningBody = La cuota puede agotarse en { "{}" } +ShowResetWhenExhausted = Mostrar reinicio al agotar la cuota +ShowResetWhenExhaustedHelper = Reemplaza el porcentaje agotado con la cuenta regresiva del reinicio diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index 60657acd06..1b095d4751 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -559,3 +559,9 @@ TrayPaceBadgeRacing = Racing TrayPaceBadgeBurning = Burning TrayResetsInLabel = リセットまで { "{}" } TrayResetsDueNow = リセット中… +PredictivePaceWarnings = ??????? +PredictivePaceWarningsHelper = Codex ??? Claude ??????????????????????????? +PredictivePaceWarningTitle = { "{}" } ? { "{}" } ????? +PredictivePaceWarningBody = ?? { "{}" } ??????????????? +ShowResetWhenExhausted = ??????????????? +ShowResetWhenExhaustedHelper = ?????????????????????????????? diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 455528798a..10e00dbf9a 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -559,3 +559,9 @@ TrayPaceBadgeRacing = 빠름 TrayPaceBadgeBurning = 매우 빠름 TrayResetsInLabel = 초기화까지 { "{}" } TrayResetsDueNow = 초기화 중… +PredictivePaceWarnings = ?? ?? ?? ?? +PredictivePaceWarningsHelper = Codex ?? Claude ???? ??? ?? ??? ???? ??? ???? +PredictivePaceWarningTitle = { "{}" } { "{}" } ?? ?? ?? +PredictivePaceWarningBody = { "{}" } ? ???? ??? ? ???? +ShowResetWhenExhausted = ?? ? ??? ?? ?? +ShowResetWhenExhaustedHelper = ??? ??? ??? ??? ??????? ???? diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index 4be46d15f5..fc414bb899 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -559,3 +559,9 @@ TrayPaceBadgeRacing = 加速 TrayPaceBadgeBurning = 超速 TrayResetsInLabel = { "{}" } 后重置 TrayResetsDueNow = 正在重置… +PredictivePaceWarnings = ???????? +PredictivePaceWarningsHelper = ? Codex ? Claude ??????????? +PredictivePaceWarningTitle = { "{}" } { "{}" }?????? +PredictivePaceWarningBody = ????? { "{}" } ??? +ShowResetWhenExhausted = ????????? +ShowResetWhenExhaustedHelper = ????????????????? diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index d043630fa4..7e5ad0cbee 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -559,3 +559,9 @@ TrayPaceBadgeRacing = 加速 TrayPaceBadgeBurning = 超速 TrayResetsInLabel = { "{}" } 後重置 TrayResetsDueNow = 正在重置… +PredictivePaceWarnings = ???????? +PredictivePaceWarningsHelper = ? Codex ? Claude ??????????? +PredictivePaceWarningTitle = { "{}" } { "{}" }?????? +PredictivePaceWarningBody = ????? { "{}" } ??? +ShowResetWhenExhausted = ????????? +ShowResetWhenExhaustedHelper = ???????????????? diff --git a/rust/src/notifications.rs b/rust/src/notifications.rs index d95b64e47a..656f8a2c38 100755 --- a/rust/src/notifications.rs +++ b/rust/src/notifications.rs @@ -5,8 +5,11 @@ #![allow(dead_code)] use crate::core::ProviderId; +use crate::core::{RateWindow, UsagePace}; +use crate::locale::{self, LocaleKey}; use crate::settings::Settings; use crate::sound::{AlertSound, play_alert}; +use chrono::{DateTime, Utc}; /// Notification types #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -25,6 +28,52 @@ pub enum NotificationType { SessionRestored, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum PredictiveWarningWindow { + Session, + Weekly, +} + +impl PredictiveWarningWindow { + fn localized_label(self, language: crate::settings::Language) -> String { + locale::get_text( + language, + match self { + Self::Session => LocaleKey::ProviderSession, + Self::Weekly => LocaleKey::ProviderWeekly, + }, + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct PredictiveResetWindow { + window_minutes: Option, + resets_at: DateTime, +} + +impl PredictiveResetWindow { + fn belongs_to_same_cycle(&self, other: &Self) -> bool { + if self.window_minutes != other.window_minutes { + return false; + } + let tolerance_secs = self + .window_minutes + .map(|minutes| i64::from(minutes) * 30) + .unwrap_or(300) + .max(300); + (self.resets_at - other.resets_at).num_seconds().abs() < tolerance_secs + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct PredictiveWarningKey { + provider: ProviderId, + identity: String, + window: PredictiveWarningWindow, + reset: PredictiveResetWindow, +} + impl NotificationType { pub fn title(&self) -> &'static str { match self { @@ -55,6 +104,7 @@ pub struct NotificationManager { sent_notifications: std::collections::HashSet<(ProviderId, NotificationType)>, /// Track previous session percent for depleted/restored transitions previous_session_percent: std::collections::HashMap, + predictive_warning_keys: std::collections::HashSet, } impl NotificationManager { @@ -62,9 +112,135 @@ impl NotificationManager { Self { sent_notifications: std::collections::HashSet::new(), previous_session_percent: std::collections::HashMap::new(), + predictive_warning_keys: std::collections::HashSet::new(), + } + } + + pub fn record_predictive_observation( + &mut self, + enabled: bool, + provider: ProviderId, + identity: &str, + window: PredictiveWarningWindow, + rate_window: &RateWindow, + pace: &UsagePace, + ) -> bool { + if !enabled { + self.predictive_warning_keys = self + .predictive_warning_keys + .iter() + .filter(|key| key.provider != provider) + .cloned() + .collect(); + return false; + } + if !matches!(provider, ProviderId::Claude | ProviderId::Codex) || identity.is_empty() { + return false; + } + let Some(resets_at) = rate_window.resets_at else { + return false; + }; + let key = PredictiveWarningKey { + provider, + identity: identity.to_string(), + window, + reset: PredictiveResetWindow { + window_minutes: rate_window.window_minutes, + resets_at, + }, + }; + + let siblings: Vec<_> = self + .predictive_warning_keys + .iter() + .filter(|existing| { + existing.provider == key.provider + && existing.identity == key.identity + && existing.window == key.window + }) + .cloned() + .collect(); + let warned_this_cycle = siblings + .iter() + .any(|existing| existing.reset.belongs_to_same_cycle(&key.reset)); + self.predictive_warning_keys = self + .predictive_warning_keys + .iter() + .filter(|existing| !siblings.contains(existing)) + .cloned() + .collect(); + if warned_this_cycle { + self.predictive_warning_keys.insert(key.clone()); + } + + if pace.will_last_to_reset { + self.predictive_warning_keys = self + .predictive_warning_keys + .iter() + .filter(|existing| *existing != &key) + .cloned() + .collect(); + return false; + } + if !pace.eta_seconds.is_some_and(|eta| eta.is_finite() && eta > 0.0) + || !pace + .run_out_probability + .is_none_or(|probability| probability.is_finite() && probability >= 0.5) + { + return false; + } + + self.predictive_warning_keys.insert(key) + } + + pub fn set_predictive_warnings_enabled(&mut self, provider: ProviderId, enabled: bool) { + if !enabled { + self.predictive_warning_keys = self + .predictive_warning_keys + .iter() + .filter(|key| key.provider != provider) + .cloned() + .collect(); } } + pub fn check_predictive_pace( + &mut self, + provider: ProviderId, + identity: &str, + window: PredictiveWarningWindow, + rate_window: &RateWindow, + pace: &UsagePace, + settings: &Settings, + ) { + if !self.record_predictive_observation( + settings.show_notifications && settings.predictive_pace_warning_enabled, + provider, + identity, + window, + rate_window, + pace, + ) { + return; + } + + let eta = format_duration(pace.eta_seconds.unwrap_or_default()); + let provider_name = provider.display_name(); + let window_label = window.localized_label(settings.ui_language); + let title = locale::format_locale( + settings.ui_language, + LocaleKey::PredictivePaceWarningTitle, + &[provider_name, &window_label], + ); + let body = locale::format_locale( + settings.ui_language, + LocaleKey::PredictivePaceWarningBody, + &[&eta], + ); + self.show_toast(&title, &body); + play_alert(AlertSound::Warning, settings); + } + /// Check usage and send notifications if thresholds are crossed pub fn check_and_notify( &mut self, @@ -323,6 +499,20 @@ impl NotificationManager { } } +fn format_duration(seconds: f64) -> String { + let total_minutes = (seconds / 60.0).ceil().max(1.0) as i64; + let days = total_minutes / 1440; + let hours = (total_minutes % 1440) / 60; + let minutes = total_minutes % 60; + if days > 0 { + format!("{days}d {hours}h") + } else if hours > 0 { + format!("{hours}h {minutes}m") + } else { + format!("{minutes}m") + } +} + impl Default for NotificationManager { fn default() -> Self { Self::new() @@ -356,3 +546,169 @@ pub fn show_notification(title: &str, body: &str) { let manager = NotificationManager::new(); manager.show_toast(title, body); } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{PaceStage, RateWindow, UsagePace}; + use chrono::{DateTime, Duration, Utc}; + + fn pace(will_last_to_reset: bool, eta_seconds: Option, probability: Option) -> UsagePace { + UsagePace { + stage: PaceStage::Ahead, + delta_percent: 20.0, + expected_used_percent: 40.0, + actual_used_percent: 60.0, + eta_seconds, + will_last_to_reset, + run_out_probability: probability, + } + } + + fn window(now: DateTime, offset: Duration, minutes: u32) -> RateWindow { + RateWindow::with_details(60.0, Some(minutes), Some(now + offset), None) + } + + #[test] + fn predictive_warning_notifies_once_until_recovery_then_rearms() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let window = window(now, Duration::hours(3), 300); + let risk = pace(false, Some(3600.0), Some(0.8)); + let recovery = pace(true, None, Some(0.0)); + let mut manager = NotificationManager::new(); + + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &risk, + )); + assert!(!manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &risk, + )); + assert!(!manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &recovery, + )); + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &risk, + )); + } + + #[test] + fn predictive_warning_reset_jitter_does_not_retrigger() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let mut manager = NotificationManager::new(); + let risk = pace(false, Some(3600.0), None); + + assert!(manager.record_predictive_observation( + true, + ProviderId::Codex, + "oauth:account-a", + PredictiveWarningWindow::Weekly, + &window(now, Duration::days(3), 10080), + &risk, + )); + assert!(!manager.record_predictive_observation( + true, + ProviderId::Codex, + "oauth:account-a", + PredictiveWarningWindow::Weekly, + &window(now, Duration::days(3) + Duration::minutes(5), 10080), + &risk, + )); + } + + #[test] + fn predictive_warning_isolates_provider_identity_source_and_window() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let reset = window(now, Duration::hours(3), 300); + let risk = pace(false, Some(3600.0), None); + let mut manager = NotificationManager::new(); + + for (provider, identity, warning_window) in [ + ( + ProviderId::Claude, + "cli:person@example.com", + PredictiveWarningWindow::Session, + ), + ( + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + ), + ( + ProviderId::Claude, + "token-account:1", + PredictiveWarningWindow::Session, + ), + ( + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Weekly, + ), + ( + ProviderId::Codex, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + ), + ] { + assert!(manager.record_predictive_observation( + true, + provider, + identity, + warning_window, + &reset, + &risk, + )); + } + } + + #[test] + fn predictive_warning_requires_enabled_confident_positive_risk() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let reset = window(now, Duration::hours(3), 300); + let mut manager = NotificationManager::new(); + + for (enabled, observation) in [ + (false, pace(false, Some(3600.0), Some(0.8))), + (true, pace(true, None, Some(0.8))), + (true, pace(false, Some(0.0), Some(0.8))), + (true, pace(false, Some(3600.0), Some(0.49))), + ] { + assert!(!manager.record_predictive_observation( + enabled, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &reset, + &observation, + )); + } + + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &reset, + &pace(false, Some(3600.0), Some(0.5)), + )); + } +} diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 95e70000d4..0c2696d368 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -94,6 +94,14 @@ pub struct Settings { /// Show reset times as relative (e.g., "2h 30m" instead of "3:00 PM") pub reset_time_relative: bool, + /// Replace exhausted quota text with its concrete future reset time. + #[serde(default)] + pub show_reset_when_exhausted: bool, + + /// Warn when Codex or Claude pace predicts exhaustion before reset. + #[serde(default)] + pub predictive_pace_warning_enabled: bool, + /// Menu bar display mode: "minimal", "compact", or "detailed" pub menu_bar_display_mode: String, @@ -357,6 +365,8 @@ impl Default for Settings { show_as_used: true, // Show as "used" by default enable_animations: true, // Animations enabled by default reset_time_relative: true, // Show relative times by default + show_reset_when_exhausted: false, + predictive_pace_warning_enabled: false, menu_bar_display_mode: "detailed".to_string(), // Detailed mode by default show_all_token_accounts_in_menu: false, provider_configs: HashMap::new(), diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 844c7eae8c..64ef568752 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -33,6 +33,8 @@ pub(super) struct RawSettings { show_as_used: bool, enable_animations: bool, reset_time_relative: bool, + show_reset_when_exhausted: bool, + predictive_pace_warning_enabled: bool, menu_bar_display_mode: String, show_all_token_accounts_in_menu: bool, @@ -158,6 +160,8 @@ impl Default for RawSettings { show_as_used: s.show_as_used, enable_animations: s.enable_animations, reset_time_relative: s.reset_time_relative, + show_reset_when_exhausted: s.show_reset_when_exhausted, + predictive_pace_warning_enabled: s.predictive_pace_warning_enabled, menu_bar_display_mode: s.menu_bar_display_mode, show_all_token_accounts_in_menu: s.show_all_token_accounts_in_menu, provider_configs: s.provider_configs, @@ -441,6 +445,8 @@ impl From for Settings { show_as_used: raw.show_as_used, enable_animations: raw.enable_animations, reset_time_relative: raw.reset_time_relative, + show_reset_when_exhausted: raw.show_reset_when_exhausted, + predictive_pace_warning_enabled: raw.predictive_pace_warning_enabled, menu_bar_display_mode: raw.menu_bar_display_mode, show_all_token_accounts_in_menu: raw.show_all_token_accounts_in_menu, provider_configs, diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index e9ff1f0768..c03f846f60 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -9,6 +9,22 @@ fn test_settings_default() { assert!(settings.show_notifications); assert_eq!(settings.high_usage_threshold, 70.0); assert_eq!(settings.critical_usage_threshold, 90.0); + assert!(!settings.show_reset_when_exhausted); + assert!(!settings.predictive_pace_warning_enabled); +} + +#[test] +fn new_warning_and_reset_settings_are_backward_compatible() { + let loaded: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["claude", "codex"], + "refresh_interval_secs": 300 + }"#, + ) + .expect("parse legacy settings"); + + assert!(!loaded.show_reset_when_exhausted); + assert!(!loaded.predictive_pace_warning_enabled); } #[test] From e7ee10956358ca46ccf1da7cf6046a21fbed8477 Mon Sep 17 00:00:00 2001 From: NessZerra Date: Sun, 12 Jul 2026 08:54:00 +0700 Subject: [PATCH 2/3] Simplify predictive warning state Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/core/usage_pace.rs | 166 ------------------------------------ rust/src/notifications.rs | 85 +++++++----------- 2 files changed, 29 insertions(+), 222 deletions(-) diff --git a/rust/src/core/usage_pace.rs b/rust/src/core/usage_pace.rs index e8affef2a1..7d66ac1a90 100755 --- a/rust/src/core/usage_pace.rs +++ b/rust/src/core/usage_pace.rs @@ -81,70 +81,6 @@ pub struct UsagePace { pub eta_seconds: Option, /// Whether current pace will last until reset pub will_last_to_reset: bool, - /// Estimated probability of running out before reset, when a predictor provides one. - pub run_out_probability: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum UsageDisplayMode { - Percent, - Pace, - Both, -} - -pub fn usage_display_text( - mode: UsageDisplayMode, - window: &RateWindow, - pace: Option<&UsagePace>, - show_used: bool, - show_reset_when_exhausted: bool, - now: DateTime, -) -> Option { - let percent = || { - let value = if show_used { - window.used_percent - } else { - window.remaining_percent() - }; - format!("{value:.0}%") - }; - - if show_reset_when_exhausted && window.is_exhausted() { - if let Some(resets_at) = window.resets_at.filter(|reset| *reset > now) { - let duration = resets_at - now; - let total_minutes = ((duration.num_seconds() + 59) / 60).max(1); - let days = total_minutes / 1440; - let hours = (total_minutes % 1440) / 60; - let minutes = total_minutes % 60; - let countdown = if days > 0 { - format!("{days}d {hours}h") - } else if hours > 0 { - format!("{hours}h {minutes}m") - } else { - format!("{minutes}m") - }; - return Some(format!("↻ in {countdown}")); - } - return Some(percent()); - } - - let pace_text = pace.map(|pace| { - let delta = pace.delta_percent.abs().round() as i64; - if delta == 0 { - "0%".to_string() - } else { - let sign = if pace.delta_percent >= 0.0 { "+" } else { "-" }; - format!("{sign}{delta}%") - } - }); - - Some(match mode { - UsageDisplayMode::Percent => percent(), - UsageDisplayMode::Pace => pace_text.unwrap_or_else(percent), - UsageDisplayMode::Both => pace_text - .map(|pace| format!("{} · {pace}", percent())) - .unwrap_or_else(percent), - }) } impl UsagePace { @@ -213,7 +149,6 @@ impl UsagePace { actual_used_percent: actual, eta_seconds, will_last_to_reset, - run_out_probability: None, }) } @@ -331,105 +266,4 @@ mod tests { assert_eq!(PaceStage::FarAhead.label(), "Far Ahead"); assert_eq!(PaceStage::SlightlyBehind.emoji(), "🐢"); } - - fn exhausted_window( - _now: DateTime, - resets_at: Option>, - reset_description: Option<&str>, - ) -> RateWindow { - RateWindow::with_details( - 100.0, - Some(300), - resets_at, - reset_description.map(str::to_string), - ) - } - - fn pace() -> UsagePace { - UsagePace { - stage: PaceStage::Ahead, - delta_percent: 12.0, - expected_used_percent: 40.0, - actual_used_percent: 52.0, - eta_seconds: Some(3600.0), - will_last_to_reset: false, - run_out_probability: None, - } - } - - #[test] - fn exhausted_display_uses_future_reset_for_percent_pace_and_both() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let window = exhausted_window( - now, - Some(now + Duration::hours(2) + Duration::minutes(15)), - None, - ); - - for mode in [ - UsageDisplayMode::Percent, - UsageDisplayMode::Pace, - UsageDisplayMode::Both, - ] { - assert_eq!( - usage_display_text(mode, &window, Some(&pace()), false, true, now).as_deref(), - Some("↻ in 2h 15m") - ); - } - } - - #[test] - fn exhausted_display_preserves_percent_without_schedulable_reset() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let cases = [ - exhausted_window(now, None, None), - exhausted_window(now, None, Some("in 2h")), - exhausted_window(now, Some(now - Duration::minutes(1)), None), - ]; - - for window in cases { - for mode in [ - UsageDisplayMode::Percent, - UsageDisplayMode::Pace, - UsageDisplayMode::Both, - ] { - assert_eq!( - usage_display_text(mode, &window, Some(&pace()), false, true, now).as_deref(), - Some("0%") - ); - } - } - } - - #[test] - fn exhausted_display_is_disabled_by_default_and_reverts_after_reset() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let resets_at = now + Duration::hours(1); - let window = exhausted_window(now, Some(resets_at), None); - - assert_eq!( - usage_display_text( - UsageDisplayMode::Percent, - &window, - None, - false, - false, - now - ) - .as_deref(), - Some("0%") - ); - assert_eq!( - usage_display_text( - UsageDisplayMode::Percent, - &window, - None, - false, - true, - resets_at - ) - .as_deref(), - Some("0%") - ); - } } diff --git a/rust/src/notifications.rs b/rust/src/notifications.rs index 656f8a2c38..61fa723721 100755 --- a/rust/src/notifications.rs +++ b/rust/src/notifications.rs @@ -126,12 +126,8 @@ impl NotificationManager { pace: &UsagePace, ) -> bool { if !enabled { - self.predictive_warning_keys = self - .predictive_warning_keys - .iter() - .filter(|key| key.provider != provider) - .cloned() - .collect(); + self.predictive_warning_keys + .retain(|key| key.provider != provider); return false; } if !matches!(provider, ProviderId::Claude | ProviderId::Codex) || identity.is_empty() { @@ -150,57 +146,36 @@ impl NotificationManager { }, }; - let siblings: Vec<_> = self - .predictive_warning_keys - .iter() - .filter(|existing| { - existing.provider == key.provider - && existing.identity == key.identity - && existing.window == key.window - }) - .cloned() - .collect(); - let warned_this_cycle = siblings - .iter() - .any(|existing| existing.reset.belongs_to_same_cycle(&key.reset)); - self.predictive_warning_keys = self - .predictive_warning_keys - .iter() - .filter(|existing| !siblings.contains(existing)) - .cloned() - .collect(); - if warned_this_cycle { - self.predictive_warning_keys.insert(key.clone()); - } + let warned_this_cycle = self.predictive_warning_keys.iter().any(|existing| { + existing.provider == key.provider + && existing.identity == key.identity + && existing.window == key.window + && existing.reset.belongs_to_same_cycle(&key.reset) + }); + self.predictive_warning_keys.retain(|existing| { + existing.provider != key.provider + || existing.identity != key.identity + || existing.window != key.window + }); if pace.will_last_to_reset { - self.predictive_warning_keys = self - .predictive_warning_keys - .iter() - .filter(|existing| *existing != &key) - .cloned() - .collect(); return false; } - if !pace.eta_seconds.is_some_and(|eta| eta.is_finite() && eta > 0.0) - || !pace - .run_out_probability - .is_none_or(|probability| probability.is_finite() && probability >= 0.5) + if !pace + .eta_seconds + .is_some_and(|eta| eta.is_finite() && eta > 0.0) { return false; } - self.predictive_warning_keys.insert(key) + self.predictive_warning_keys.insert(key); + !warned_this_cycle } pub fn set_predictive_warnings_enabled(&mut self, provider: ProviderId, enabled: bool) { if !enabled { - self.predictive_warning_keys = self - .predictive_warning_keys - .iter() - .filter(|key| key.provider != provider) - .cloned() - .collect(); + self.predictive_warning_keys + .retain(|key| key.provider != provider); } } @@ -553,7 +528,7 @@ mod tests { use crate::core::{PaceStage, RateWindow, UsagePace}; use chrono::{DateTime, Duration, Utc}; - fn pace(will_last_to_reset: bool, eta_seconds: Option, probability: Option) -> UsagePace { + fn pace(will_last_to_reset: bool, eta_seconds: Option) -> UsagePace { UsagePace { stage: PaceStage::Ahead, delta_percent: 20.0, @@ -561,7 +536,6 @@ mod tests { actual_used_percent: 60.0, eta_seconds, will_last_to_reset, - run_out_probability: probability, } } @@ -573,8 +547,8 @@ mod tests { fn predictive_warning_notifies_once_until_recovery_then_rearms() { let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); let window = window(now, Duration::hours(3), 300); - let risk = pace(false, Some(3600.0), Some(0.8)); - let recovery = pace(true, None, Some(0.0)); + let risk = pace(false, Some(3600.0)); + let recovery = pace(true, None); let mut manager = NotificationManager::new(); assert!(manager.record_predictive_observation( @@ -615,7 +589,7 @@ mod tests { fn predictive_warning_reset_jitter_does_not_retrigger() { let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); let mut manager = NotificationManager::new(); - let risk = pace(false, Some(3600.0), None); + let risk = pace(false, Some(3600.0)); assert!(manager.record_predictive_observation( true, @@ -639,7 +613,7 @@ mod tests { fn predictive_warning_isolates_provider_identity_source_and_window() { let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); let reset = window(now, Duration::hours(3), 300); - let risk = pace(false, Some(3600.0), None); + let risk = pace(false, Some(3600.0)); let mut manager = NotificationManager::new(); for (provider, identity, warning_window) in [ @@ -687,10 +661,9 @@ mod tests { let mut manager = NotificationManager::new(); for (enabled, observation) in [ - (false, pace(false, Some(3600.0), Some(0.8))), - (true, pace(true, None, Some(0.8))), - (true, pace(false, Some(0.0), Some(0.8))), - (true, pace(false, Some(3600.0), Some(0.49))), + (false, pace(false, Some(3600.0))), + (true, pace(true, None)), + (true, pace(false, Some(0.0))), ] { assert!(!manager.record_predictive_observation( enabled, @@ -708,7 +681,7 @@ mod tests { "oauth:person@example.com", PredictiveWarningWindow::Session, &reset, - &pace(false, Some(3600.0), Some(0.5)), + &pace(false, Some(3600.0)), )); } } From 63d45fb0b5454041a2cc0d2c318a19df9f028a6f Mon Sep 17 00:00:00 2001 From: NessZerra Date: Sun, 12 Jul 2026 09:02:40 +0700 Subject: [PATCH 3/3] Simplify warning refresh integration Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src-tauri/src/commands/providers.rs | 94 +++++++------------ 1 file changed, 35 insertions(+), 59 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 0b4e8b84ce..ffef079984 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -178,12 +178,7 @@ async fn do_refresh_providers_with_policy( await_provider_refreshes(handles).await; let error_count = finish_provider_refresh(&state)?; - update_tray_and_notifications( - app, - &state, - &inputs.settings, - &inputs.token_accounts, - )?; + update_tray_and_notifications(app, &state, &inputs.settings, &inputs.token_accounts)?; events::emit_refresh_complete(app, enabled_count, error_count); @@ -473,66 +468,47 @@ fn notify_predictive_pace( .ok() .map(|date| date.with_timezone(&chrono::Utc)); - notify_predictive_window( - manager, - provider, - &identity, - codexbar::notifications::PredictiveWarningWindow::Session, - &snapshot.primary, - observed_at, - 300, - settings, - ); - if let Some(weekly) = &snapshot.secondary { - notify_predictive_window( - manager, - provider, - &identity, + for (warning_window, window, default_window_minutes) in [ + ( + codexbar::notifications::PredictiveWarningWindow::Session, + Some(&snapshot.primary), + 300, + ), + ( codexbar::notifications::PredictiveWarningWindow::Weekly, - weekly, - observed_at, + snapshot.secondary.as_ref(), 10080, + ), + ] { + let Some(window) = window else { + continue; + }; + let rate_window = RateWindow::with_details( + window.used_percent, + window.window_minutes, + window + .resets_at + .as_deref() + .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) + .map(|date| date.with_timezone(&chrono::Utc)), + window.reset_description.clone(), + ); + let Some(pace) = + codexbar::core::UsagePace::weekly(&rate_window, observed_at, default_window_minutes) + else { + continue; + }; + manager.check_predictive_pace( + provider, + &identity, + warning_window, + &rate_window, + &pace, settings, ); } } -#[allow(clippy::too_many_arguments)] -fn notify_predictive_window( - manager: &mut codexbar::notifications::NotificationManager, - provider: ProviderId, - identity: &str, - warning_window: codexbar::notifications::PredictiveWarningWindow, - window: &RateWindowSnapshot, - observed_at: Option>, - default_window_minutes: u32, - settings: &Settings, -) { - let rate_window = RateWindow::with_details( - window.used_percent, - window.window_minutes, - window - .resets_at - .as_deref() - .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) - .map(|date| date.with_timezone(&chrono::Utc)), - window.reset_description.clone(), - ); - let Some(pace) = - codexbar::core::UsagePace::weekly(&rate_window, observed_at, default_window_minutes) - else { - return; - }; - manager.check_predictive_pace( - provider, - identity, - warning_window, - &rate_window, - &pace, - settings, - ); -} - fn predictive_warning_identity( provider: ProviderId, source_label: &str,