From df000da917952024edd14cb624af8d1e1740a2b5 Mon Sep 17 00:00:00 2001 From: Rasmus Rygaard Date: Wed, 4 Feb 2026 06:01:47 -0800 Subject: [PATCH] Add a codex.rate_limits event for websockets (#10324) When communicating over websockets, we can't rely on headers to deliver rate limit information. This PR adds a `codex.rate_limits` event that the server can pass to the client to inform them about rate limit usage. The client parses this data the same way we parse rate limit headers in HTTP mode. This PR also wires up the etag and reasoning headers for websockets --- .../src/endpoint/responses_websocket.rs | 29 ++++++- codex-rs/codex-api/src/rate_limits.rs | 66 ++++++++++++++ codex-rs/codex-api/src/sse/responses.rs | 6 ++ .../core/tests/suite/client_websockets.rs | 86 +++++++++++++++++++ 4 files changed, 183 insertions(+), 4 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index cac686dd0..20c6067b1 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -5,6 +5,7 @@ use crate::common::ResponseStream; use crate::common::ResponsesWsRequest; use crate::error::ApiError; use crate::provider::Provider; +use crate::rate_limits::parse_rate_limit_event; use crate::sse::responses::ResponsesStreamEvent; use crate::sse::responses::process_responses_event; use crate::telemetry::WebsocketTelemetry; @@ -33,6 +34,7 @@ use url::Url; type WsStream = WebSocketStream>; const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; +const X_MODELS_ETAG_HEADER: &str = "x-models-etag"; const X_REASONING_INCLUDED_HEADER: &str = "x-reasoning-included"; pub struct ResponsesWebsocketConnection { @@ -40,6 +42,7 @@ pub struct ResponsesWebsocketConnection { // TODO (pakrym): is this the right place for timeout? idle_timeout: Duration, server_reasoning_included: bool, + models_etag: Option, telemetry: Option>, } @@ -48,12 +51,14 @@ impl ResponsesWebsocketConnection { stream: WsStream, idle_timeout: Duration, server_reasoning_included: bool, + models_etag: Option, telemetry: Option>, ) -> Self { Self { stream: Arc::new(Mutex::new(Some(stream))), idle_timeout, server_reasoning_included, + models_etag, telemetry, } } @@ -71,12 +76,16 @@ impl ResponsesWebsocketConnection { let stream = Arc::clone(&self.stream); let idle_timeout = self.idle_timeout; let server_reasoning_included = self.server_reasoning_included; + let models_etag = self.models_etag.clone(); let telemetry = self.telemetry.clone(); let request_body = serde_json::to_value(&request).map_err(|err| { ApiError::Stream(format!("failed to encode websocket request: {err}")) })?; tokio::spawn(async move { + if let Some(etag) = models_etag { + let _ = tx_event.send(Ok(ResponseEvent::ModelsEtag(etag))).await; + } if server_reasoning_included { let _ = tx_event .send(Ok(ResponseEvent::ServerReasoningIncluded(true))) @@ -136,12 +145,13 @@ impl ResponsesWebsocketClient { headers.extend(extra_headers); add_auth_headers_to_header_map(&self.auth, &mut headers); - let (stream, server_reasoning_included) = - connect_websocket(ws_url, headers, turn_state).await?; + let (stream, server_reasoning_included, models_etag) = + connect_websocket(ws_url, headers, turn_state.clone()).await?; Ok(ResponsesWebsocketConnection::new( stream, self.provider.stream_idle_timeout, server_reasoning_included, + models_etag, telemetry, )) } @@ -151,7 +161,7 @@ async fn connect_websocket( url: Url, headers: HeaderMap, turn_state: Option>>, -) -> Result<(WsStream, bool), ApiError> { +) -> Result<(WsStream, bool, Option), ApiError> { info!("connecting to websocket: {url}"); let mut request = url @@ -177,6 +187,11 @@ async fn connect_websocket( }; let reasoning_included = response.headers().contains_key(X_REASONING_INCLUDED_HEADER); + let models_etag = response + .headers() + .get(X_MODELS_ETAG_HEADER) + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); if let Some(turn_state) = turn_state && let Some(header_value) = response .headers() @@ -185,7 +200,7 @@ async fn connect_websocket( { let _ = turn_state.set(header_value.to_string()); } - Ok((stream, reasoning_included)) + Ok((stream, reasoning_included, models_etag)) } fn map_ws_error(err: WsError, url: &Url) -> ApiError { @@ -273,6 +288,12 @@ async fn run_websocket_response_stream( continue; } }; + if event.kind() == "codex.rate_limits" { + if let Some(snapshot) = parse_rate_limit_event(&text) { + let _ = tx_event.send(Ok(ResponseEvent::RateLimits(snapshot))).await; + } + continue; + } match process_responses_event(event) { Ok(Some(event)) => { let is_completed = matches!(event, ResponseEvent::Completed { .. }); diff --git a/codex-rs/codex-api/src/rate_limits.rs b/codex-rs/codex-api/src/rate_limits.rs index c29aab21f..047c5934b 100644 --- a/codex-rs/codex-api/src/rate_limits.rs +++ b/codex-rs/codex-api/src/rate_limits.rs @@ -1,7 +1,9 @@ +use codex_protocol::account::PlanType; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use http::HeaderMap; +use serde::Deserialize; use std::fmt::Display; #[derive(Debug)] @@ -41,6 +43,70 @@ pub fn parse_rate_limit(headers: &HeaderMap) -> Option { }) } +#[derive(Debug, Deserialize)] +struct RateLimitEventWindow { + used_percent: f64, + window_minutes: Option, + reset_at: Option, +} + +#[derive(Debug, Deserialize)] +struct RateLimitEventDetails { + primary: Option, + secondary: Option, +} + +#[derive(Debug, Deserialize)] +struct RateLimitEventCredits { + has_credits: bool, + unlimited: bool, + balance: Option, +} + +#[derive(Debug, Deserialize)] +struct RateLimitEvent { + #[serde(rename = "type")] + kind: String, + plan_type: Option, + rate_limits: Option, + credits: Option, +} + +pub fn parse_rate_limit_event(payload: &str) -> Option { + let event: RateLimitEvent = serde_json::from_str(payload).ok()?; + if event.kind != "codex.rate_limits" { + return None; + } + let (primary, secondary) = if let Some(details) = event.rate_limits.as_ref() { + ( + map_event_window(details.primary.as_ref()), + map_event_window(details.secondary.as_ref()), + ) + } else { + (None, None) + }; + let credits = event.credits.map(|credits| CreditsSnapshot { + has_credits: credits.has_credits, + unlimited: credits.unlimited, + balance: credits.balance, + }); + Some(RateLimitSnapshot { + primary, + secondary, + credits, + plan_type: event.plan_type, + }) +} + +fn map_event_window(window: Option<&RateLimitEventWindow>) -> Option { + let window = window?; + Some(RateLimitWindow { + used_percent: window.used_percent, + window_minutes: window.window_minutes, + resets_at: window.reset_at, + }) +} + /// Parses the bespoke Codex rate-limit headers into a `RateLimitSnapshot`. pub fn parse_promo_message(headers: &HeaderMap) -> Option { parse_header_str(headers, "x-codex-promo-message") diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 2c9110472..95bda35fb 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -165,6 +165,12 @@ pub struct ResponsesStreamEvent { content_index: Option, } +impl ResponsesStreamEvent { + pub fn kind(&self) -> &str { + &self.kind + } +} + #[derive(Debug)] pub enum ResponsesEventError { Api(ApiError), diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 9f6628344..15005caa4 100644 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -18,6 +18,7 @@ use codex_otel::OtelManager; use codex_otel::metrics::MetricsClient; use codex_otel::metrics::MetricsConfig; use codex_protocol::ThreadId; +use codex_protocol::account::PlanType; use codex_protocol::config_types::ReasoningSummary; use core_test_support::load_default_config_for_test; use core_test_support::responses::WebSocketConnectionConfig; @@ -30,6 +31,7 @@ use core_test_support::skip_if_no_network; use futures::StreamExt; use opentelemetry_sdk::metrics::InMemoryMetricExporter; use pretty_assertions::assert_eq; +use serde_json::json; use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; @@ -203,6 +205,90 @@ async fn responses_websocket_emits_reasoning_included_event() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_emits_rate_limit_events() { + skip_if_no_network!(); + + let rate_limit_event = json!({ + "type": "codex.rate_limits", + "plan_type": "plus", + "rate_limits": { + "allowed": true, + "limit_reached": false, + "primary": { + "used_percent": 42, + "window_minutes": 60, + "reset_at": 1700000000 + }, + "secondary": null + }, + "code_review_rate_limits": null, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "123" + }, + "promo": null + }); + + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![vec![ + rate_limit_event, + ev_response_created("resp-1"), + ev_completed("resp-1"), + ]], + response_headers: vec![ + ("X-Models-Etag".to_string(), "etag-123".to_string()), + ("X-Reasoning-Included".to_string(), "true".to_string()), + ], + }]) + .await; + + let harness = websocket_harness(&server).await; + let mut session = harness.client.new_session(None); + let prompt = prompt_with_input(vec![message_item("hello")]); + + let mut stream = session + .stream(&prompt) + .await + .expect("websocket stream failed"); + + let mut saw_rate_limits = None; + let mut saw_models_etag = None; + let mut saw_reasoning_included = false; + + while let Some(event) = stream.next().await { + match event.expect("event") { + ResponseEvent::RateLimits(snapshot) => { + saw_rate_limits = Some(snapshot); + } + ResponseEvent::ModelsEtag(etag) => { + saw_models_etag = Some(etag); + } + ResponseEvent::ServerReasoningIncluded(true) => { + saw_reasoning_included = true; + } + ResponseEvent::Completed { .. } => break, + _ => {} + } + } + + let rate_limits = saw_rate_limits.expect("missing rate limits"); + let primary = rate_limits.primary.expect("missing primary window"); + assert_eq!(primary.used_percent, 42.0); + assert_eq!(primary.window_minutes, Some(60)); + assert_eq!(primary.resets_at, Some(1_700_000_000)); + assert_eq!(rate_limits.plan_type, Some(PlanType::Plus)); + let credits = rate_limits.credits.expect("missing credits"); + assert!(credits.has_credits); + assert!(!credits.unlimited); + assert_eq!(credits.balance.as_deref(), Some("123")); + assert_eq!(saw_models_etag.as_deref(), Some("etag-123")); + assert!(saw_reasoning_included); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_appends_on_prefix() { skip_if_no_network!();