From d9b914738af94e7bd3fb0d1aecd4b645ccfad5b8 Mon Sep 17 00:00:00 2001 From: kevin zhao Date: Tue, 2 Dec 2025 16:20:45 +0000 Subject: [PATCH 1/3] persisting credits if new snapshot does not contain credit info --- codex-rs/core/src/state/session.rs | 16 +++++++++- codex-rs/tui/src/chatwidget.rs | 16 +++++++++- codex-rs/tui/src/chatwidget/tests.rs | 48 ++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index caebac6b86c..48fb3964225 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -62,7 +62,10 @@ impl SessionState { } pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) { - self.latest_rate_limits = Some(snapshot); + self.latest_rate_limits = Some(merge_rate_limit_credits( + self.latest_rate_limits.as_ref(), + snapshot, + )); } pub(crate) fn token_info_and_rate_limits( @@ -79,3 +82,14 @@ impl SessionState { self.history.get_total_token_usage() } } + +// Sometimes new snapshots don't include credits +fn merge_rate_limit_credits( + previous: Option<&RateLimitSnapshot>, + mut snapshot: RateLimitSnapshot, +) -> RateLimitSnapshot { + if snapshot.credits.is_none() { + snapshot.credits = previous.and_then(|prior| prior.credits.clone()); + } + snapshot +} \ No newline at end of file diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index d11d983cf14..445d4aa1ee7 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -20,6 +20,7 @@ use codex_core::protocol::AgentReasoningRawContentDeltaEvent; use codex_core::protocol::AgentReasoningRawContentEvent; use codex_core::protocol::ApplyPatchApprovalRequestEvent; use codex_core::protocol::BackgroundEventEvent; +use codex_core::protocol::CreditsSnapshot; use codex_core::protocol::DeprecationNoticeEvent; use codex_core::protocol::ErrorEvent; use codex_core::protocol::Event; @@ -548,7 +549,20 @@ impl ChatWidget { } pub(crate) fn on_rate_limit_snapshot(&mut self, snapshot: Option) { - if let Some(snapshot) = snapshot { + if let Some(mut snapshot) = snapshot { + if snapshot.credits.is_none() + && let Some(previous_credits) = self + .rate_limit_snapshot + .as_ref() + .and_then(|display| display.credits.as_ref()) + { + snapshot.credits = Some(CreditsSnapshot { + has_credits: previous_credits.has_credits, + unlimited: previous_credits.unlimited, + balance: previous_credits.balance.clone(), + }); + } + let warnings = self.rate_limit_warnings.take_warnings( snapshot .secondary diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index e0a1cc3d3b8..448c4bd79b7 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -18,6 +18,7 @@ use codex_core::protocol::AgentReasoningDeltaEvent; use codex_core::protocol::AgentReasoningEvent; use codex_core::protocol::ApplyPatchApprovalRequestEvent; use codex_core::protocol::BackgroundEventEvent; +use codex_core::protocol::CreditsSnapshot; use codex_core::protocol::Event; use codex_core::protocol::EventMsg; use codex_core::protocol::ExecApprovalRequestEvent; @@ -481,6 +482,53 @@ fn test_rate_limit_warnings_monthly() { ); } +#[test] +fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(); + + chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + primary: None, + secondary: None, + credits: Some(CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("17.5".to_string()), + }), + })); + let initial_balance = chat + .rate_limit_snapshot + .as_ref() + .and_then(|snapshot| snapshot.credits.as_ref()) + .and_then(|credits| credits.balance.as_deref()); + assert_eq!(initial_balance, Some("17.5")); + + chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent: 80.0, + window_minutes: Some(60), + resets_at: Some(123), + }), + secondary: None, + credits: None, + })); + + let display = chat + .rate_limit_snapshot + .as_ref() + .expect("rate limits should be cached"); + let credits = display + .credits + .as_ref() + .expect("credits should persist when headers omit them"); + + assert_eq!(credits.balance.as_deref(), Some("17.5")); + assert!(!credits.unlimited); + assert_eq!( + display.primary.as_ref().map(|window| window.used_percent), + Some(80.0) + ); +} + #[test] fn rate_limit_switch_prompt_skips_when_on_lower_cost_model() { let (mut chat, _, _) = make_chatwidget_manual(); From c3d940d1e234345bc9242f1596a192ce05153f54 Mon Sep 17 00:00:00 2001 From: kevin zhao Date: Tue, 2 Dec 2025 16:31:15 +0000 Subject: [PATCH 2/3] cleanup --- codex-rs/core/src/state/session.rs | 2 +- codex-rs/tui/src/chatwidget.rs | 15 +++++++-------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 48fb3964225..8c739c92435 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -92,4 +92,4 @@ fn merge_rate_limit_credits( snapshot.credits = previous.and_then(|prior| prior.credits.clone()); } snapshot -} \ No newline at end of file +} diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 445d4aa1ee7..3921b9f3cb1 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -550,17 +550,16 @@ impl ChatWidget { pub(crate) fn on_rate_limit_snapshot(&mut self, snapshot: Option) { if let Some(mut snapshot) = snapshot { - if snapshot.credits.is_none() - && let Some(previous_credits) = self + if snapshot.credits.is_none() { + snapshot.credits = self .rate_limit_snapshot .as_ref() .and_then(|display| display.credits.as_ref()) - { - snapshot.credits = Some(CreditsSnapshot { - has_credits: previous_credits.has_credits, - unlimited: previous_credits.unlimited, - balance: previous_credits.balance.clone(), - }); + .map(|credits| CreditsSnapshot { + has_credits: credits.has_credits, + unlimited: credits.unlimited, + balance: credits.balance.clone(), + }); } let warnings = self.rate_limit_warnings.take_warnings( From 06f6f0bba17aad0d8ed999b44a4fbf61c700829b Mon Sep 17 00:00:00 2001 From: kevin zhao Date: Tue, 2 Dec 2025 13:47:47 -0500 Subject: [PATCH 3/3] add teset --- codex-rs/core/src/codex.rs | 72 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a3d21ba8cae..3ae5964db77 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2445,7 +2445,10 @@ mod tests { use crate::tools::format_exec_output_str; use crate::protocol::CompactedItem; + use crate::protocol::CreditsSnapshot; use crate::protocol::InitialHistory; + use crate::protocol::RateLimitSnapshot; + use crate::protocol::RateLimitWindow; use crate::protocol::ResumedHistory; use crate::state::TaskKind; use crate::tasks::SessionTask; @@ -2515,6 +2518,75 @@ mod tests { assert_eq!(expected, actual); } + #[test] + fn set_rate_limits_retains_previous_credits() { + let codex_home = tempfile::tempdir().expect("create temp dir"); + let config = Config::load_from_base_config_with_overrides( + ConfigToml::default(), + ConfigOverrides::default(), + codex_home.path().to_path_buf(), + ) + .expect("load default test config"); + let config = Arc::new(config); + let session_configuration = SessionConfiguration { + provider: config.model_provider.clone(), + model: config.model.clone(), + model_reasoning_effort: config.model_reasoning_effort, + model_reasoning_summary: config.model_reasoning_summary, + developer_instructions: config.developer_instructions.clone(), + user_instructions: config.user_instructions.clone(), + base_instructions: config.base_instructions.clone(), + compact_prompt: config.compact_prompt.clone(), + approval_policy: config.approval_policy, + sandbox_policy: config.sandbox_policy.clone(), + cwd: config.cwd.clone(), + original_config_do_not_use: Arc::clone(&config), + features: Features::default(), + exec_policy: Arc::new(ExecPolicy::empty()), + session_source: SessionSource::Exec, + }; + + let mut state = SessionState::new(session_configuration); + let initial = RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent: 10.0, + window_minutes: Some(15), + resets_at: Some(1_700), + }), + secondary: None, + credits: Some(CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("10.00".to_string()), + }), + }; + state.set_rate_limits(initial.clone()); + + let update = RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent: 40.0, + window_minutes: Some(30), + resets_at: Some(1_800), + }), + secondary: Some(RateLimitWindow { + used_percent: 5.0, + window_minutes: Some(60), + resets_at: Some(1_900), + }), + credits: None, + }; + state.set_rate_limits(update.clone()); + + assert_eq!( + state.latest_rate_limits, + Some(RateLimitSnapshot { + primary: update.primary.clone(), + secondary: update.secondary, + credits: initial.credits, + }) + ); + } + #[test] fn prefers_structured_content_when_present() { let ctr = CallToolResult {