From 5ebdc9af1ba4db1ce7d1d34f65bf66dabe45e78e Mon Sep 17 00:00:00 2001 From: zhao-oai Date: Tue, 2 Dec 2025 16:23:24 -0500 Subject: [PATCH] persisting credits if new snapshot does not contain credit info (#7490) in response to incoming changes to responses headers where the header may sometimes not contain credits info (no longer forcing a credit check) --- codex-rs/core/src/codex.rs | 72 ++++++++++++++++++++++++++++ codex-rs/core/src/state/session.rs | 16 ++++++- codex-rs/tui/src/chatwidget.rs | 15 +++++- codex-rs/tui/src/chatwidget/tests.rs | 48 +++++++++++++++++++ 4 files changed, 149 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index ef0a237f7..8b1dfafd2 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2481,7 +2481,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; @@ -2551,6 +2554,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 { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index caebac6b8..8c739c924 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 +} diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 8bb1e78e2..79a184313 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; @@ -558,7 +559,19 @@ 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() { + snapshot.credits = self + .rate_limit_snapshot + .as_ref() + .and_then(|display| display.credits.as_ref()) + .map(|credits| CreditsSnapshot { + has_credits: credits.has_credits, + unlimited: credits.unlimited, + balance: 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 eb65e3a5d..890e8bbe1 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; @@ -521,6 +522,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();