valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Reset in for rate limits (#4111)

Browse Source 
- Parse the headers
- Reorganize the struct because it's getting too long
- show the resets at in the tui

<img width="324" height="79" alt="image"
src="https://github.com/user-attachments/assets/ca15cd48-f112-4556-91ab-1e3a9bc4683d"
/>
This commit is contained in:
Ahmed Ibrahim
2025-09-24 08:31:08 -07:00
committed by GitHub
parent 5b910f1f05
commit cb96f4f596
7 changed files with 235 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
codex-rs/core/src/client.rs
View File

@@ -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<RateLimitSnapshot> {
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<RateLimitWindow> {
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),
})
}

16
codex-rs/core/src/error.rs
View File

@@ -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),
}),
}
}

63
codex-rs/core/tests/suite/client.rs
View File

@@ -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

22
codex-rs/protocol/src/protocol.rs
View File

@@ -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<RateLimitWindow>,
pub secondary: Option<RateLimitWindow>,
}
#[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<u64>,
/// Seconds until the window resets.
pub resets_in_seconds: Option<u64>,
}
// Includes prompts, tools and space to call compact.

75
codex-rs/tui/src/chatwidget.rs
View File

@@ -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<f64>,
primary_used_percent: Option<f64>,
) -> Vec<String> {
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<f64> = 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<f64> = 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<f64> = 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<f64> = 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<UserMessage>,
token_info: Option<TokenUsageInfo>,
rate_limit_snapshot: Option<RateLimitSnapshot>,
rate_limit_snapshot: Option<RateLimitSnapshotDisplay>,
rate_limit_warnings: RateLimitWarningState,
// Stream lifecycle controller
stream_controller: Option<StreamController>,
@@ -366,16 +375,24 @@ impl ChatWidget {
fn on_rate_limit_snapshot(&mut self, snapshot: Option<RateLimitSnapshot>) {
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.

12
codex-rs/tui/src/chatwidget/tests.rs
View File

@@ -384,12 +384,12 @@ fn rate_limit_warnings_emit_thresholds() {
let mut state = RateLimitWarningState::default();
let mut warnings: Vec<String> = 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,

94
codex-rs/tui/src/history_cell.rs
View File

@@ -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<String>,
}
impl RateLimitWindowDisplay {
fn from_window(window: &RateLimitWindow, captured_at: DateTime<Local>) -> 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<RateLimitWindowDisplay>,
pub secondary: Option<RateLimitWindowDisplay>,
}
pub(crate) fn rate_limit_snapshot_display(
snapshot: &RateLimitSnapshot,
captured_at: DateTime<Local>,
) -> 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<ConversationId>,
rate_limits: Option<&RateLimitSnapshot>,
rate_limits: Option<&RateLimitSnapshotDisplay>,
) -> PlainHistoryCell {
let mut lines: Vec<Line<'static>> = 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<Line<'static>> {
fn build_status_limit_lines(snapshot: Option<&RateLimitSnapshotDisplay>) -> Vec<Line<'static>> {
let mut lines: Vec<Line<'static>> =
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;