From 230310b4236dc6916d232882c032fb3a06a07a64 Mon Sep 17 00:00:00 2001 From: Colin Young Date: Thu, 2 Oct 2025 20:30:45 -0700 Subject: [PATCH] [VSCE-147] Fetch Initial Rate Limit Usage --- codex-rs/core/src/client.rs | 113 ++++++++++++++++++++++++++++- codex-rs/core/src/codex.rs | 33 ++++++++- codex-rs/core/src/state/session.rs | 10 ++- codex-rs/protocol/src/protocol.rs | 4 + codex-rs/tui/src/status/tests.rs | 8 ++ 5 files changed, 164 insertions(+), 4 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 7cf60f56b9..623cfa9bdd 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -592,7 +592,12 @@ fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option { "x-codex-secondary-reset-after-seconds", ); - Some(RateLimitSnapshot { primary, secondary }) + Some(RateLimitSnapshot { + primary, + secondary, + allowed: None, + limit_reached: None, + }) } fn parse_rate_limit_window( @@ -634,6 +639,112 @@ fn parse_header_str<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> { headers.get(name)?.to_str().ok() } +#[derive(Debug, Deserialize)] +struct UsageStatusPayload { + rate_limit: Option, +} + +#[derive(Debug, Deserialize)] +struct UsageStatusDetails { + allowed: bool, + limit_reached: bool, + #[serde(default)] + primary_window: Option, + #[serde(default)] + secondary_window: Option, +} + +#[derive(Debug, Deserialize)] +struct UsageWindowSnapshot { + used_percent: f64, + limit_window_seconds: u64, + reset_after_seconds: u64, +} + +pub async fn fetch_usage_rate_limits( + auth_manager: Arc, + chatgpt_base_url: String, +) -> Result> { + let Some(auth) = auth_manager.auth() else { + return Ok(None); + }; + + if auth.mode != AuthMode::ChatGPT { + return Ok(None); + } + + let token = match auth.get_token().await { + Ok(token) => token, + Err(err) => { + warn!("failed to load ChatGPT token for usage request: {err}"); + return Ok(None); + } + }; + + let account_id = auth.get_account_id(); + + let base = chatgpt_base_url.trim_end_matches('/'); + let url = format!("{base}/codex/usage"); + let client = create_client(); + + let mut request = client + .get(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json"); + + if let Some(account_id) = account_id { + request = request.header("chatgpt-account-id", account_id); + } + + let response = match request.send().await { + Ok(response) => response, + Err(err) => { + warn!("failed to request usage status: {err:#}"); + return Ok(None); + } + }; + + if !response.status().is_success() { + debug!(status = ?response.status(), "usage status request did not succeed"); + return Ok(None); + } + + let payload: UsageStatusPayload = match response.json().await { + Ok(payload) => payload, + Err(err) => { + warn!("failed to parse usage status response: {err:#}"); + return Ok(None); + } + }; + + Ok(convert_usage_payload(payload)) +} + +fn convert_usage_payload(payload: UsageStatusPayload) -> Option { + let details = payload.rate_limit?; + + Some(RateLimitSnapshot { + primary: details.primary_window.map(convert_usage_window), + secondary: details.secondary_window.map(convert_usage_window), + allowed: Some(details.allowed), + limit_reached: Some(details.limit_reached), + }) +} + +fn convert_usage_window(window: UsageWindowSnapshot) -> RateLimitWindow { + let window_minutes = if window.limit_window_seconds == 0 { + None + } else { + Some(window.limit_window_seconds / 60) + }; + + RateLimitWindow { + used_percent: window.used_percent, + window_minutes, + resets_in_seconds: Some(window.reset_after_seconds), + } +} + async fn process_sse( stream: S, tx_event: mpsc::Sender>, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4baa63ec3e..b16c2633e2 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -49,6 +49,7 @@ use crate::apply_patch::CODEX_APPLY_PATCH_ARG1; use crate::apply_patch::InternalApplyPatchInvocation; use crate::apply_patch::convert_apply_patch_to_protocol; use crate::client::ModelClient; +use crate::client::fetch_usage_rate_limits; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::config::Config; @@ -397,9 +398,33 @@ impl Session { let default_shell_fut = shell::default_user_shell(); let history_meta_fut = crate::message_history::history_metadata(&config); + let usage_auth = auth_manager.clone(); + let usage_base_url = config.chatgpt_base_url.clone(); + let rate_limits_fut = + async move { fetch_usage_rate_limits(usage_auth, usage_base_url).await }; + // Join all independent futures. - let (rollout_recorder, mcp_res, default_shell, (history_log_id, history_entry_count)) = - tokio::join!(rollout_fut, mcp_fut, default_shell_fut, history_meta_fut); + let ( + rollout_recorder, + mcp_res, + default_shell, + (history_log_id, history_entry_count), + initial_rate_limits_result, + ) = tokio::join!( + rollout_fut, + mcp_fut, + default_shell_fut, + history_meta_fut, + rate_limits_fut, + ); + + let initial_rate_limits = match initial_rate_limits_result { + Ok(rate_limits) => rate_limits, + Err(err) => { + warn!("failed to fetch rate limit status: {err:#}"); + None + } + }; let rollout_recorder = rollout_recorder.map_err(|e| { error!("failed to initialize rollout recorder: {e:#}"); @@ -532,6 +557,10 @@ impl Session { sess.send_event(event).await; } + if let Some(snapshot) = initial_rate_limits { + sess.update_rate_limits(INITIAL_SUBMIT_ID, snapshot).await; + } + Ok((sess, turn_context)) } diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index ee0c5fc976..23655b9ac2 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -66,7 +66,15 @@ impl SessionState { ); } - pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) { + pub(crate) fn set_rate_limits(&mut self, mut snapshot: RateLimitSnapshot) { + if let Some(prev) = self.latest_rate_limits.as_ref() { + if snapshot.allowed.is_none() { + snapshot.allowed = prev.allowed; + } + if snapshot.limit_reached.is_none() { + snapshot.limit_reached = prev.limit_reached; + } + } self.latest_rate_limits = Some(snapshot); } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 2c6d3b33f0..45620c09b5 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -602,6 +602,10 @@ pub struct TokenCountEvent { pub struct RateLimitSnapshot { pub primary: Option, pub secondary: Option, + #[serde(default)] + pub allowed: Option, + #[serde(default)] + pub limit_reached: Option, } #[derive(Debug, Clone, Deserialize, Serialize, TS)] diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index 395d952f41..6b851672a1 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -96,6 +96,8 @@ fn status_snapshot_includes_reasoning_details() { window_minutes: Some(10080), resets_in_seconds: Some(1_200), }), + allowed: None, + limit_reached: None, }; let captured_at = chrono::Local .with_ymd_and_hms(2024, 1, 2, 3, 4, 5) @@ -137,6 +139,8 @@ fn status_snapshot_includes_monthly_limit() { resets_in_seconds: Some(86_400), }), secondary: None, + allowed: None, + limit_reached: None, }; let captured_at = chrono::Local .with_ymd_and_hms(2024, 5, 6, 7, 8, 9) @@ -204,6 +208,8 @@ fn status_snapshot_truncates_in_narrow_terminal() { resets_in_seconds: Some(600), }), secondary: None, + allowed: None, + limit_reached: None, }; let captured_at = chrono::Local .with_ymd_and_hms(2024, 1, 2, 3, 4, 5) @@ -267,6 +273,8 @@ fn status_snapshot_shows_empty_limits_message() { let snapshot = RateLimitSnapshot { primary: None, secondary: None, + allowed: None, + limit_reached: None, }; let captured_at = chrono::Local .with_ymd_and_hms(2024, 6, 7, 8, 9, 10)