From cb96f4f596dde29db83f11cc8696cebf3d2871eb Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 24 Sep 2025 08:31:08 -0700 Subject: [PATCH] Add Reset in for rate limits (#4111) - Parse the headers - Reorganize the struct because it's getting too long - show the resets at in the tui image --- codex-rs/core/src/client.rs | 45 +++++++++---- codex-rs/core/src/error.rs | 16 +++-- codex-rs/core/tests/suite/client.rs | 63 ++++++++++++++----- codex-rs/protocol/src/protocol.rs | 22 ++++--- codex-rs/tui/src/chatwidget.rs | 75 +++++++++++++--------- codex-rs/tui/src/chatwidget/tests.rs | 12 ++-- codex-rs/tui/src/history_cell.rs | 94 ++++++++++++++++++++++++---- 7 files changed, 235 insertions(+), 92 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 73f36109..0ae057b3 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -43,6 +43,7 @@ use crate::model_provider_info::WireApi; use crate::openai_model_info::get_model_info; use crate::openai_tools::create_tools_json_for_responses_api; use crate::protocol::RateLimitSnapshot; +use crate::protocol::RateLimitWindow; use crate::protocol::TokenUsage; use crate::token_data::PlanType; use crate::util::backoff; @@ -488,19 +489,39 @@ fn attach_item_ids(payload_json: &mut Value, original_items: &[ResponseItem]) { } fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option { - let primary_used_percent = parse_header_f64(headers, "x-codex-primary-used-percent")?; - let secondary_used_percent = parse_header_f64(headers, "x-codex-secondary-used-percent")?; - let primary_to_secondary_ratio_percent = - parse_header_f64(headers, "x-codex-primary-over-secondary-limit-percent")?; - let primary_window_minutes = parse_header_u64(headers, "x-codex-primary-window-minutes")?; - let secondary_window_minutes = parse_header_u64(headers, "x-codex-secondary-window-minutes")?; + let primary = parse_rate_limit_window( + headers, + "x-codex-primary-used-percent", + "x-codex-primary-window-minutes", + "x-codex-primary-reset-after-seconds", + ); - Some(RateLimitSnapshot { - primary_used_percent, - secondary_used_percent, - primary_to_secondary_ratio_percent, - primary_window_minutes, - secondary_window_minutes, + let secondary = parse_rate_limit_window( + headers, + "x-codex-secondary-used-percent", + "x-codex-secondary-window-minutes", + "x-codex-secondary-reset-after-seconds", + ); + + if primary.is_none() && secondary.is_none() { + return None; + } + + Some(RateLimitSnapshot { primary, secondary }) +} + +fn parse_rate_limit_window( + headers: &HeaderMap, + used_percent_header: &str, + window_minutes_header: &str, + resets_header: &str, +) -> Option { + let used_percent = parse_header_f64(headers, used_percent_header)?; + + Some(RateLimitWindow { + used_percent, + window_minutes: parse_header_u64(headers, window_minutes_header), + resets_in_seconds: parse_header_u64(headers, resets_header), }) } diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 9840e167..1fe7e115 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -267,14 +267,20 @@ pub fn get_error_message_ui(e: &CodexErr) -> String { #[cfg(test)] mod tests { use super::*; + use codex_protocol::protocol::RateLimitWindow; fn rate_limit_snapshot() -> RateLimitSnapshot { RateLimitSnapshot { - primary_used_percent: 0.5, - secondary_used_percent: 0.3, - primary_to_secondary_ratio_percent: 0.7, - primary_window_minutes: 60, - secondary_window_minutes: 120, + primary: Some(RateLimitWindow { + used_percent: 50.0, + window_minutes: Some(60), + resets_in_seconds: Some(3600), + }), + secondary: Some(RateLimitWindow { + used_percent: 30.0, + window_minutes: Some(120), + resets_in_seconds: Some(7200), + }), } } diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index f6a73d13..867accd7 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -763,9 +763,10 @@ async fn token_count_includes_rate_limits_snapshot() { .insert_header("content-type", "text/event-stream") .insert_header("x-codex-primary-used-percent", "12.5") .insert_header("x-codex-secondary-used-percent", "40.0") - .insert_header("x-codex-primary-over-secondary-limit-percent", "75.0") .insert_header("x-codex-primary-window-minutes", "10") .insert_header("x-codex-secondary-window-minutes", "60") + .insert_header("x-codex-primary-reset-after-seconds", "1800") + .insert_header("x-codex-secondary-reset-after-seconds", "7200") .set_body_raw(sse_body, "text/event-stream"); Mock::given(method("POST")) @@ -811,11 +812,16 @@ async fn token_count_includes_rate_limits_snapshot() { json!({ "info": null, "rate_limits": { - "primary_used_percent": 12.5, - "secondary_used_percent": 40.0, - "primary_to_secondary_ratio_percent": 75.0, - "primary_window_minutes": 10, - "secondary_window_minutes": 60 + "primary": { + "used_percent": 12.5, + "window_minutes": 10, + "resets_in_seconds": 1800 + }, + "secondary": { + "used_percent": 40.0, + "window_minutes": 60, + "resets_in_seconds": 7200 + } } }) ); @@ -853,11 +859,16 @@ async fn token_count_includes_rate_limits_snapshot() { "model_context_window": 272000 }, "rate_limits": { - "primary_used_percent": 12.5, - "secondary_used_percent": 40.0, - "primary_to_secondary_ratio_percent": 75.0, - "primary_window_minutes": 10, - "secondary_window_minutes": 60 + "primary": { + "used_percent": 12.5, + "window_minutes": 10, + "resets_in_seconds": 1800 + }, + "secondary": { + "used_percent": 40.0, + "window_minutes": 60, + "resets_in_seconds": 7200 + } } }) ); @@ -868,7 +879,20 @@ async fn token_count_includes_rate_limits_snapshot() { let final_snapshot = final_payload .rate_limits .expect("latest rate limit snapshot should be retained"); - assert_eq!(final_snapshot.primary_used_percent, 12.5); + assert_eq!( + final_snapshot + .primary + .as_ref() + .map(|window| window.used_percent), + Some(12.5) + ); + assert_eq!( + final_snapshot + .primary + .as_ref() + .and_then(|window| window.resets_in_seconds), + Some(1800) + ); wait_for_event(&codex, |msg| matches!(msg, EventMsg::TaskComplete(_))).await; } @@ -904,11 +928,16 @@ async fn usage_limit_error_emits_rate_limit_event() -> anyhow::Result<()> { let codex = codex_fixture.codex.clone(); let expected_limits = json!({ - "primary_used_percent": 100.0, - "secondary_used_percent": 87.5, - "primary_to_secondary_ratio_percent": 95.0, - "primary_window_minutes": 15, - "secondary_window_minutes": 60 + "primary": { + "used_percent": 100.0, + "window_minutes": 15, + "resets_in_seconds": null + }, + "secondary": { + "used_percent": 87.5, + "window_minutes": 60, + "resets_in_seconds": null + } }); let submission_id = codex diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 70159bd1..bf0b1925 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -597,16 +597,18 @@ pub struct TokenCountEvent { #[derive(Debug, Clone, Deserialize, Serialize, TS)] pub struct RateLimitSnapshot { - /// Percentage (0-100) of the primary window that has been consumed. - pub primary_used_percent: f64, - /// Percentage (0-100) of the secondary window that has been consumed. - pub secondary_used_percent: f64, - /// Size of the primary window relative to secondary (0-100). - pub primary_to_secondary_ratio_percent: f64, - /// Rolling window duration for the primary limit, in minutes. - pub primary_window_minutes: u64, - /// Rolling window duration for the secondary limit, in minutes. - pub secondary_window_minutes: u64, + pub primary: Option, + pub secondary: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, TS)] +pub struct RateLimitWindow { + /// Percentage (0-100) of the window that has been consumed. + pub used_percent: f64, + /// Rolling window duration, in minutes. + pub window_minutes: Option, + /// Seconds until the window resets. + pub resets_in_seconds: Option, } // Includes prompts, tools and space to call compact. diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 6e4026b4..857030e3 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -77,6 +77,7 @@ use crate::history_cell::CommandOutput; use crate::history_cell::ExecCell; use crate::history_cell::HistoryCell; use crate::history_cell::PatchEventType; +use crate::history_cell::RateLimitSnapshotDisplay; use crate::markdown::append_markdown; use crate::slash_command::SlashCommand; use crate::text_formatting::truncate_text; @@ -94,6 +95,7 @@ use crate::streaming::controller::AppEventHistorySink; use crate::streaming::controller::StreamController; use std::path::Path; +use chrono::Local; use codex_common::approval_presets::ApprovalPreset; use codex_common::approval_presets::builtin_approval_presets; use codex_common::model_presets::ModelPreset; @@ -129,39 +131,46 @@ struct RateLimitWarningState { impl RateLimitWarningState { fn take_warnings( &mut self, - secondary_used_percent: f64, - primary_used_percent: f64, + secondary_used_percent: Option, + primary_used_percent: Option, ) -> Vec { - if secondary_used_percent == 100.0 || primary_used_percent == 100.0 { + let reached_secondary_cap = + matches!(secondary_used_percent, Some(percent) if percent == 100.0); + let reached_primary_cap = matches!(primary_used_percent, Some(percent) if percent == 100.0); + if reached_secondary_cap || reached_primary_cap { return Vec::new(); } let mut warnings = Vec::new(); - let mut highest_secondary: Option = None; - while self.secondary_index < RATE_LIMIT_WARNING_THRESHOLDS.len() - && secondary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[self.secondary_index] - { - highest_secondary = Some(RATE_LIMIT_WARNING_THRESHOLDS[self.secondary_index]); - self.secondary_index += 1; - } - if let Some(threshold) = highest_secondary { - warnings.push(format!( - "Heads up, you've used over {threshold:.0}% of your weekly limit. Run /status for a breakdown." - )); + if let Some(secondary_used_percent) = secondary_used_percent { + let mut highest_secondary: Option = None; + while self.secondary_index < RATE_LIMIT_WARNING_THRESHOLDS.len() + && secondary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[self.secondary_index] + { + highest_secondary = Some(RATE_LIMIT_WARNING_THRESHOLDS[self.secondary_index]); + self.secondary_index += 1; + } + if let Some(threshold) = highest_secondary { + warnings.push(format!( + "Heads up, you've used over {threshold:.0}% of your weekly limit. Run /status for a breakdown." + )); + } } - let mut highest_primary: Option = None; - while self.primary_index < RATE_LIMIT_WARNING_THRESHOLDS.len() - && primary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[self.primary_index] - { - highest_primary = Some(RATE_LIMIT_WARNING_THRESHOLDS[self.primary_index]); - self.primary_index += 1; - } - if let Some(threshold) = highest_primary { - warnings.push(format!( - "Heads up, you've used over {threshold:.0}% of your 5h limit. Run /status for a breakdown." - )); + if let Some(primary_used_percent) = primary_used_percent { + let mut highest_primary: Option = None; + while self.primary_index < RATE_LIMIT_WARNING_THRESHOLDS.len() + && primary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[self.primary_index] + { + highest_primary = Some(RATE_LIMIT_WARNING_THRESHOLDS[self.primary_index]); + self.primary_index += 1; + } + if let Some(threshold) = highest_primary { + warnings.push(format!( + "Heads up, you've used over {threshold:.0}% of your 5h limit. Run /status for a breakdown." + )); + } } warnings @@ -189,7 +198,7 @@ pub(crate) struct ChatWidget { session_header: SessionHeader, initial_user_message: Option, token_info: Option, - rate_limit_snapshot: Option, + rate_limit_snapshot: Option, rate_limit_warnings: RateLimitWarningState, // Stream lifecycle controller stream_controller: Option, @@ -366,16 +375,24 @@ impl ChatWidget { fn on_rate_limit_snapshot(&mut self, snapshot: Option) { if let Some(snapshot) = snapshot { let warnings = self.rate_limit_warnings.take_warnings( - snapshot.secondary_used_percent, - snapshot.primary_used_percent, + snapshot + .secondary + .as_ref() + .map(|window| window.used_percent), + snapshot.primary.as_ref().map(|window| window.used_percent), ); - self.rate_limit_snapshot = Some(snapshot); + + let display = history_cell::rate_limit_snapshot_display(&snapshot, Local::now()); + self.rate_limit_snapshot = Some(display); + if !warnings.is_empty() { for warning in warnings { self.add_to_history(history_cell::new_warning_event(warning)); } self.request_redraw(); } + } else { + self.rate_limit_snapshot = None; } } /// Finalize any active exec as failed and stop/clear running UI state. diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index 2d90b240..fa3235b6 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -384,12 +384,12 @@ fn rate_limit_warnings_emit_thresholds() { let mut state = RateLimitWarningState::default(); let mut warnings: Vec = Vec::new(); - warnings.extend(state.take_warnings(10.0, 55.0)); - warnings.extend(state.take_warnings(55.0, 10.0)); - warnings.extend(state.take_warnings(10.0, 80.0)); - warnings.extend(state.take_warnings(80.0, 10.0)); - warnings.extend(state.take_warnings(10.0, 95.0)); - warnings.extend(state.take_warnings(95.0, 10.0)); + warnings.extend(state.take_warnings(Some(10.0), Some(55.0))); + warnings.extend(state.take_warnings(Some(55.0), Some(10.0))); + warnings.extend(state.take_warnings(Some(10.0), Some(80.0))); + warnings.extend(state.take_warnings(Some(80.0), Some(10.0))); + warnings.extend(state.take_warnings(Some(10.0), Some(95.0))); + warnings.extend(state.take_warnings(Some(95.0), Some(10.0))); assert_eq!( warnings, diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index 7406f98d..5b6c4af3 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -11,6 +11,9 @@ use crate::wrapping::RtOptions; use crate::wrapping::word_wrap_line; use crate::wrapping::word_wrap_lines; use base64::Engine; +use chrono::DateTime; +use chrono::Duration as ChronoDuration; +use chrono::Local; use codex_ansi_escape::ansi_escape_line; use codex_common::create_config_summary_entries; use codex_common::elapsed::format_duration; @@ -25,6 +28,7 @@ use codex_core::project_doc::discover_project_doc_paths; use codex_core::protocol::FileChange; use codex_core::protocol::McpInvocation; use codex_core::protocol::RateLimitSnapshot; +use codex_core::protocol::RateLimitWindow; use codex_core::protocol::SandboxPolicy; use codex_core::protocol::SessionConfiguredEvent; use codex_core::protocol::TokenUsage; @@ -47,6 +51,7 @@ use ratatui::widgets::WidgetRef; use ratatui::widgets::Wrap; use std::any::Any; use std::collections::HashMap; +use std::convert::TryFrom; use std::io::Cursor; use std::path::Path; use std::path::PathBuf; @@ -1078,11 +1083,54 @@ pub(crate) fn new_warning_event(message: String) -> PlainHistoryCell { } } +#[derive(Debug, Clone)] +pub(crate) struct RateLimitWindowDisplay { + pub used_percent: f64, + pub resets_at: Option, +} + +impl RateLimitWindowDisplay { + fn from_window(window: &RateLimitWindow, captured_at: DateTime) -> Self { + let resets_at = window + .resets_in_seconds + .and_then(|seconds| i64::try_from(seconds).ok()) + .and_then(|secs| captured_at.checked_add_signed(ChronoDuration::seconds(secs))) + .map(|dt| dt.format("%b %-d, %Y %-I:%M %p").to_string()); + + Self { + used_percent: window.used_percent, + resets_at, + } + } +} + +#[derive(Debug, Clone)] +pub(crate) struct RateLimitSnapshotDisplay { + pub primary: Option, + pub secondary: Option, +} + +pub(crate) fn rate_limit_snapshot_display( + snapshot: &RateLimitSnapshot, + captured_at: DateTime, +) -> RateLimitSnapshotDisplay { + RateLimitSnapshotDisplay { + primary: snapshot + .primary + .as_ref() + .map(|window| RateLimitWindowDisplay::from_window(window, captured_at)), + secondary: snapshot + .secondary + .as_ref() + .map(|window| RateLimitWindowDisplay::from_window(window, captured_at)), + } +} + pub(crate) fn new_status_output( config: &Config, usage: &TokenUsage, session_id: &Option, - rate_limits: Option<&RateLimitSnapshot>, + rate_limits: Option<&RateLimitSnapshotDisplay>, ) -> PlainHistoryCell { let mut lines: Vec> = Vec::new(); lines.push("/status".magenta().into()); @@ -1611,23 +1659,39 @@ fn format_mcp_invocation<'a>(invocation: McpInvocation) -> Line<'a> { invocation_spans.into() } -fn build_status_limit_lines(snapshot: Option<&RateLimitSnapshot>) -> Vec> { +fn build_status_limit_lines(snapshot: Option<&RateLimitSnapshotDisplay>) -> Vec> { let mut lines: Vec> = vec![vec![padded_emoji("⏱️").into(), "Usage Limits".bold()].into()]; match snapshot { Some(snapshot) => { - let rows = [ - ("5h limit".to_string(), snapshot.primary_used_percent), - ("Weekly limit".to_string(), snapshot.secondary_used_percent), - ]; - let label_width = rows - .iter() - .map(|(label, _)| UnicodeWidthStr::width(label.as_str())) - .max() - .unwrap_or(0); - for (label, percent) in rows { - lines.push(build_status_limit_line(&label, percent, label_width)); + let mut windows: Vec<(&str, &RateLimitWindowDisplay)> = Vec::new(); + if let Some(primary) = snapshot.primary.as_ref() { + windows.push(("5h limit", primary)); + } + if let Some(secondary) = snapshot.secondary.as_ref() { + windows.push(("Weekly limit", secondary)); + } + + if windows.is_empty() { + lines.push(" • No rate limit data available.".into()); + } else { + let label_width = windows + .iter() + .map(|(label, _)| UnicodeWidthStr::width(*label)) + .max() + .unwrap_or(0); + + for (label, window) in windows { + lines.push(build_status_limit_line( + label, + window.used_percent, + label_width, + )); + if let Some(resets_at) = window.resets_at.as_deref() { + lines.push(build_status_reset_line(resets_at)); + } + } } } None => lines.push(" • Send a message to load usage data.".into()), @@ -1651,6 +1715,10 @@ fn build_status_limit_line(label: &str, percent_used: f64, label_width: usize) - Line::from(spans) } +fn build_status_reset_line(resets_at: &str) -> Line<'static> { + vec![" ".into(), format!("Resets at: {resets_at}").dim()].into() +} + fn render_status_limit_progress_bar(percent_used: f64) -> String { let ratio = (percent_used / 100.0).clamp(0.0, 1.0); let filled = (ratio * STATUS_LIMIT_BAR_SEGMENTS as f64).round() as usize;