From 86d45ed0ddf712131ae2e30f0c8f1d1630654b7c Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 23 Sep 2025 20:43:52 -0700 Subject: [PATCH] reset-in --- codex-rs/core/src/client.rs | 44 ++++++++----- codex-rs/core/src/client_common.rs | 4 +- codex-rs/core/src/codex.rs | 8 +-- codex-rs/core/src/error.rs | 24 ++++--- codex-rs/core/tests/suite/client.rs | 33 ++++++++-- codex-rs/protocol/src/protocol.rs | 26 ++++---- codex-rs/tui/src/chatwidget.rs | 79 ++++++++++++++--------- codex-rs/tui/src/chatwidget/tests.rs | 12 ++-- codex-rs/tui/src/history_cell.rs | 96 ++++++++++++++++++++++++---- 9 files changed, 226 insertions(+), 100 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 73f361091d..ec809ff742 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -42,7 +42,8 @@ use crate::model_provider_info::ModelProviderInfo; 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::RateLimitSnapshotEvent; +use crate::protocol::RateLimitWindow; use crate::protocol::TokenUsage; use crate::token_data::PlanType; use crate::util::backoff; @@ -487,21 +488,34 @@ 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")?; +fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option { + let primary_used_percent = parse_header_f64(headers, "x-codex-primary-used-percent"); + let primary_window_minutes = parse_header_u64(headers, "x-codex-primary-window-minutes"); + let primary_resets_in_seconds = + parse_header_u64(headers, "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 primary = primary_used_percent.map(|used_percent| RateLimitWindow { + used_percent, + window_minutes: primary_window_minutes, + resets_in_seconds: primary_resets_in_seconds, + }); + + let secondary_used_percent = parse_header_f64(headers, "x-codex-secondary-used-percent"); + let secondary_window_minutes = parse_header_u64(headers, "x-codex-secondary-window-minutes"); + let secondary_resets_in_seconds = + parse_header_u64(headers, "x-codex-secondary-reset-after-seconds"); + + let secondary = secondary_used_percent.map(|used_percent| RateLimitWindow { + used_percent, + window_minutes: secondary_window_minutes, + resets_in_seconds: secondary_resets_in_seconds, + }); + + if primary.is_none() && secondary.is_none() { + return None; + } + + Some(RateLimitSnapshotEvent { primary, secondary }) } fn parse_header_f64(headers: &HeaderMap, name: &str) -> Option { diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index b695581deb..b8a6cf4752 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -1,7 +1,7 @@ use crate::error::Result; use crate::model_family::ModelFamily; use crate::openai_tools::OpenAiTool; -use crate::protocol::RateLimitSnapshot; +use crate::protocol::RateLimitSnapshotEvent; use crate::protocol::TokenUsage; use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS; use codex_protocol::config_types::ReasoningEffort as ReasoningEffortConfig; @@ -83,7 +83,7 @@ pub enum ResponseEvent { WebSearchCallBegin { call_id: String, }, - RateLimits(RateLimitSnapshot), + RateLimits(RateLimitSnapshotEvent), } #[derive(Debug, Serialize)] diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a2b0b94934..c3e5ae0457 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -100,7 +100,7 @@ use crate::protocol::ListCustomPromptsResponseEvent; use crate::protocol::Op; use crate::protocol::PatchApplyBeginEvent; use crate::protocol::PatchApplyEndEvent; -use crate::protocol::RateLimitSnapshot; +use crate::protocol::RateLimitSnapshotEvent; use crate::protocol::ReviewDecision; use crate::protocol::ReviewOutputEvent; use crate::protocol::SandboxPolicy; @@ -261,7 +261,7 @@ struct State { pending_input: Vec, history: ConversationHistory, token_info: Option, - latest_rate_limits: Option, + latest_rate_limits: Option, } /// Context for an initialized model agent @@ -757,7 +757,7 @@ impl Session { self.send_token_count_event(sub_id).await; } - async fn update_rate_limits(&self, sub_id: &str, new_rate_limits: RateLimitSnapshot) { + async fn update_rate_limits(&self, sub_id: &str, new_rate_limits: RateLimitSnapshotEvent) { { let mut state = self.state.lock().await; state.latest_rate_limits = Some(new_rate_limits); @@ -2148,7 +2148,7 @@ async fn try_run_turn( ResponseEvent::RateLimits(snapshot) => { // Update internal state with latest rate limits, but defer sending until // token usage is available to avoid duplicate TokenCount events. - sess.update_rate_limits(sub_id, snapshot).await; + sess.update_rate_limits(&sub_id, snapshot).await; } ResponseEvent::Completed { response_id: _, diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 9840e167fb..86d78e0112 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -2,7 +2,7 @@ use crate::exec::ExecToolCallOutput; use crate::token_data::KnownPlan; use crate::token_data::PlanType; use codex_protocol::mcp_protocol::ConversationId; -use codex_protocol::protocol::RateLimitSnapshot; +use codex_protocol::protocol::RateLimitSnapshotEvent; use reqwest::StatusCode; use serde_json; use std::io; @@ -139,7 +139,7 @@ pub enum CodexErr { pub struct UsageLimitReachedError { pub(crate) plan_type: Option, pub(crate) resets_in_seconds: Option, - pub(crate) rate_limits: Option, + pub(crate) rate_limits: Option, } impl std::fmt::Display for UsageLimitReachedError { @@ -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, + fn rate_limit_snapshot() -> RateLimitSnapshotEvent { + RateLimitSnapshotEvent { + 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 f6a73d1362..4b4bc9a4b1 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")) @@ -853,11 +854,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 +874,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; } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 70159bd1d2..7696c6623e 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -592,21 +592,23 @@ impl TokenUsageInfo { #[derive(Debug, Clone, Deserialize, Serialize, TS)] pub struct TokenCountEvent { pub info: Option, - pub rate_limits: Option, + pub rate_limits: Option, } #[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 struct RateLimitSnapshotEvent { + 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 6e4026b42b..c46ec52948 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -30,7 +30,7 @@ use codex_core::protocol::McpToolCallBeginEvent; use codex_core::protocol::McpToolCallEndEvent; use codex_core::protocol::Op; use codex_core::protocol::PatchApplyBeginEvent; -use codex_core::protocol::RateLimitSnapshot; +use codex_core::protocol::RateLimitSnapshotEvent; use codex_core::protocol::ReviewRequest; use codex_core::protocol::StreamErrorEvent; use codex_core::protocol::TaskCompleteEvent; @@ -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, @@ -363,19 +372,27 @@ impl ChatWidget { } } - fn on_rate_limit_snapshot(&mut self, snapshot: Option) { + 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 2d90b240e6..fa3235b625 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 7406f98d7c..430fd91baf 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; @@ -24,7 +27,8 @@ use codex_core::plan_tool::UpdatePlanArgs; 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::RateLimitSnapshotEvent; +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: &RateLimitSnapshotEvent, + 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;