diff --git a/codex-rs/app-server-protocol/schema/json/EventMsg.json b/codex-rs/app-server-protocol/schema/json/EventMsg.json index 5991401b84f..58526ca8bac 100644 --- a/codex-rs/app-server-protocol/schema/json/EventMsg.json +++ b/codex-rs/app-server-protocol/schema/json/EventMsg.json @@ -3456,6 +3456,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 69c6c2cb9b6..4a9170560f7 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -4392,6 +4392,18 @@ } ] }, + "limitId": { + "type": [ + "string", + "null" + ] + }, + "limitName": { + "type": [ + "string", + "null" + ] + }, "planType": { "anyOf": [ { @@ -4437,6 +4449,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 88fb5ddbdb3..1b21948f6c1 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -6541,6 +6541,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { @@ -11823,7 +11835,22 @@ "$schema": "http://json-schema.org/draft-07/schema#", "properties": { "rateLimits": { - "$ref": "#/definitions/v2/RateLimitSnapshot" + "allOf": [ + { + "$ref": "#/definitions/v2/RateLimitSnapshot" + } + ], + "description": "Backward-compatible single-bucket view; mirrors the historical payload." + }, + "rateLimitsByLimitId": { + "additionalProperties": { + "$ref": "#/definitions/v2/RateLimitSnapshot" + }, + "description": "Multi-bucket view keyed by metered `limit_id` (for example, `codex`).", + "type": [ + "object", + "null" + ] } }, "required": [ @@ -12838,6 +12865,18 @@ } ] }, + "limitId": { + "type": [ + "string", + "null" + ] + }, + "limitName": { + "type": [ + "string", + "null" + ] + }, "planType": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json b/codex-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json index 90949d56f27..676c0011131 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json @@ -3456,6 +3456,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json b/codex-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json index b25d2dbed39..04d0a1a38b2 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json @@ -3456,6 +3456,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json b/codex-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json index fa9e144c48a..0dc049e5978 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json @@ -3456,6 +3456,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json index d168911bdf1..91879645de7 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json @@ -48,6 +48,18 @@ } ] }, + "limitId": { + "type": [ + "string", + "null" + ] + }, + "limitName": { + "type": [ + "string", + "null" + ] + }, "planType": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json b/codex-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json index a7025fbaea1..244156269d2 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json @@ -48,6 +48,18 @@ } ] }, + "limitId": { + "type": [ + "string", + "null" + ] + }, + "limitName": { + "type": [ + "string", + "null" + ] + }, "planType": { "anyOf": [ { @@ -110,7 +122,22 @@ }, "properties": { "rateLimits": { - "$ref": "#/definitions/RateLimitSnapshot" + "allOf": [ + { + "$ref": "#/definitions/RateLimitSnapshot" + } + ], + "description": "Backward-compatible single-bucket view; mirrors the historical payload." + }, + "rateLimitsByLimitId": { + "additionalProperties": { + "$ref": "#/definitions/RateLimitSnapshot" + }, + "description": "Multi-bucket view keyed by metered `limit_id` (for example, `codex`).", + "type": [ + "object", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts b/codex-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts index 9c2dad7f094..8604128b4e4 100644 --- a/codex-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts +++ b/codex-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts @@ -5,4 +5,4 @@ import type { CreditsSnapshot } from "./CreditsSnapshot"; import type { PlanType } from "./PlanType"; import type { RateLimitWindow } from "./RateLimitWindow"; -export type RateLimitSnapshot = { primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, plan_type: PlanType | null, }; +export type RateLimitSnapshot = { limit_id: string | null, limit_name: string | null, primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, plan_type: PlanType | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountRateLimitsResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountRateLimitsResponse.ts index fe970c1d42b..e75e9b8283c 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountRateLimitsResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountRateLimitsResponse.ts @@ -3,4 +3,12 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { RateLimitSnapshot } from "./RateLimitSnapshot"; -export type GetAccountRateLimitsResponse = { rateLimits: RateLimitSnapshot, }; +export type GetAccountRateLimitsResponse = { +/** + * Backward-compatible single-bucket view; mirrors the historical payload. + */ +rateLimits: RateLimitSnapshot, +/** + * Multi-bucket view keyed by metered `limit_id` (for example, `codex`). + */ +rateLimitsByLimitId: { [key in string]?: RateLimitSnapshot } | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts b/codex-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts index f1a33f0b13b..0c2ebe1893f 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts @@ -5,4 +5,4 @@ import type { PlanType } from "../PlanType"; import type { CreditsSnapshot } from "./CreditsSnapshot"; import type { RateLimitWindow } from "./RateLimitWindow"; -export type RateLimitSnapshot = { primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, planType: PlanType | null, }; +export type RateLimitSnapshot = { limitId: string | null, limitName: string | null, primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, planType: PlanType | null, }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 679dc4ab231..bfc4b7a80fa 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -1009,7 +1009,10 @@ pub struct ChatgptAuthTokensRefreshResponse { #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct GetAccountRateLimitsResponse { + /// Backward-compatible single-bucket view; mirrors the historical payload. pub rate_limits: RateLimitSnapshot, + /// Multi-bucket view keyed by metered `limit_id` (for example, `codex`). + pub rate_limits_by_limit_id: Option>, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -3103,6 +3106,8 @@ pub struct AccountRateLimitsUpdatedNotification { #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct RateLimitSnapshot { + pub limit_id: Option, + pub limit_name: Option, pub primary: Option, pub secondary: Option, pub credits: Option, @@ -3112,6 +3117,8 @@ pub struct RateLimitSnapshot { impl From for RateLimitSnapshot { fn from(value: CoreRateLimitSnapshot) -> Self { Self { + limit_id: value.limit_id, + limit_name: value.limit_name, primary: value.primary.map(RateLimitWindow::from), secondary: value.secondary.map(RateLimitWindow::from), credits: value.credits.map(CreditsSnapshot::from), diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 73647b1c08d..d3e397c6f43 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2206,6 +2206,8 @@ mod tests { model_context_window: Some(4096), }; let rate_limits = RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: None, primary: Some(RateLimitWindow { used_percent: 42.5, window_minutes: Some(15), @@ -2258,6 +2260,8 @@ mod tests { OutgoingMessage::AppServerNotification( ServerNotification::AccountRateLimitsUpdated(payload), ) => { + assert_eq!(payload.rate_limits.limit_id.as_deref(), Some("codex")); + assert_eq!(payload.rate_limits.limit_name, None); assert!(payload.rate_limits.primary.is_some()); assert!(payload.rate_limits.credits.is_some()); } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index e8b921980f4..cd8042366c6 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1425,9 +1425,15 @@ impl CodexMessageProcessor { async fn get_account_rate_limits(&self, request_id: RequestId) { match self.fetch_account_rate_limits().await { - Ok(rate_limits) => { + Ok((rate_limits, rate_limits_by_limit_id)) => { let response = GetAccountRateLimitsResponse { rate_limits: rate_limits.into(), + rate_limits_by_limit_id: Some( + rate_limits_by_limit_id + .into_iter() + .map(|(limit_id, snapshot)| (limit_id, snapshot.into())) + .collect(), + ), }; self.outgoing.send_response(request_id, response).await; } @@ -1437,7 +1443,15 @@ impl CodexMessageProcessor { } } - async fn fetch_account_rate_limits(&self) -> Result { + async fn fetch_account_rate_limits( + &self, + ) -> Result< + ( + CoreRateLimitSnapshot, + HashMap, + ), + JSONRPCErrorError, + > { let Some(auth) = self.auth_manager.auth().await else { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, @@ -1461,14 +1475,41 @@ impl CodexMessageProcessor { data: None, })?; - client - .get_rate_limits() + let snapshots = client + .get_rate_limits_many() .await .map_err(|err| JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to fetch codex rate limits: {err}"), data: None, + })?; + if snapshots.is_empty() { + return Err(JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: "failed to fetch codex rate limits: no snapshots returned".to_string(), + data: None, + }); + } + + let rate_limits_by_limit_id: HashMap = snapshots + .iter() + .cloned() + .map(|snapshot| { + let limit_id = snapshot + .limit_id + .clone() + .unwrap_or_else(|| "codex".to_string()); + (limit_id, snapshot) }) + .collect(); + + let primary = snapshots + .iter() + .find(|snapshot| snapshot.limit_id.as_deref() == Some("codex")) + .cloned() + .unwrap_or_else(|| snapshots[0].clone()); + + Ok((primary, rate_limits_by_limit_id)) } async fn get_user_saved_config(&self, request_id: RequestId) { diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index b64bd5bce95..c9fd9ebda84 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -283,6 +283,8 @@ mod tests { let notification = ServerNotification::AccountRateLimitsUpdated(AccountRateLimitsUpdatedNotification { rate_limits: RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: None, primary: Some(RateLimitWindow { used_percent: 25, window_duration_mins: Some(15), @@ -299,7 +301,9 @@ mod tests { json!({ "method": "account/rateLimits/updated", "params": { - "rateLimits": { + "rateLimits": { + "limitId": "codex", + "limitName": null, "primary": { "usedPercent": 25, "windowDurationMins": 15, diff --git a/codex-rs/app-server/tests/suite/v2/rate_limits.rs b/codex-rs/app-server/tests/suite/v2/rate_limits.rs index e4e670310ae..6f66d7bff73 100644 --- a/codex-rs/app-server/tests/suite/v2/rate_limits.rs +++ b/codex-rs/app-server/tests/suite/v2/rate_limits.rs @@ -117,7 +117,23 @@ async fn get_account_rate_limits_returns_snapshot() -> Result<()> { "reset_after_seconds": 43200, "reset_at": secondary_reset_timestamp, } - } + }, + "additional_rate_limits": [ + { + "limit_name": "codex_other", + "metered_feature": "codex_other", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 88, + "limit_window_seconds": 1800, + "reset_after_seconds": 600, + "reset_at": 1735693200 + } + } + } + ] }); Mock::given(method("GET")) @@ -143,6 +159,8 @@ async fn get_account_rate_limits_returns_snapshot() -> Result<()> { let expected = GetAccountRateLimitsResponse { rate_limits: RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: None, primary: Some(RateLimitWindow { used_percent: 42, window_duration_mins: Some(60), @@ -156,6 +174,46 @@ async fn get_account_rate_limits_returns_snapshot() -> Result<()> { credits: None, plan_type: Some(AccountPlanType::Pro), }, + rate_limits_by_limit_id: Some( + [ + ( + "codex".to_string(), + RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: None, + primary: Some(RateLimitWindow { + used_percent: 42, + window_duration_mins: Some(60), + resets_at: Some(primary_reset_timestamp), + }), + secondary: Some(RateLimitWindow { + used_percent: 5, + window_duration_mins: Some(1440), + resets_at: Some(secondary_reset_timestamp), + }), + credits: None, + plan_type: Some(AccountPlanType::Pro), + }, + ), + ( + "codex_other".to_string(), + RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 88, + window_duration_mins: Some(30), + resets_at: Some(1735693200), + }), + secondary: None, + credits: None, + plan_type: Some(AccountPlanType::Pro), + }, + ), + ] + .into_iter() + .collect(), + ), }; assert_eq!(received, expected); diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 6fa36d1ffd1..40bc276f012 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -1,9 +1,7 @@ use crate::types::CodeTaskDetailsResponse; use crate::types::ConfigFileResponse; -use crate::types::CreditStatusDetails; use crate::types::PaginatedListTaskListItem; use crate::types::RateLimitStatusPayload; -use crate::types::RateLimitWindowSnapshot; use crate::types::TurnAttemptsSiblingTurnsResponse; use anyhow::Result; use codex_core::auth::CodexAuth; @@ -160,6 +158,15 @@ impl Client { } pub async fn get_rate_limits(&self) -> Result { + let snapshots = self.get_rate_limits_many().await?; + let preferred = snapshots + .iter() + .find(|snapshot| snapshot.limit_id.as_deref() == Some("codex")) + .cloned(); + Ok(preferred.unwrap_or_else(|| snapshots[0].clone())) + } + + pub async fn get_rate_limits_many(&self) -> Result> { let url = match self.path_style { PathStyle::CodexApi => format!("{}/api/codex/usage", self.base_url), PathStyle::ChatGptApi => format!("{}/wham/usage", self.base_url), @@ -167,7 +174,7 @@ impl Client { let req = self.http.get(&url).headers(self.headers()); let (body, ct) = self.exec_request(req, "GET", &url).await?; let payload: RateLimitStatusPayload = self.decode_json(&url, &ct, &body)?; - Ok(Self::rate_limit_snapshot_from_payload(payload)) + Ok(Self::rate_limit_snapshots_from_payload(payload)) } pub async fn list_tasks( @@ -295,35 +302,59 @@ impl Client { } // rate limit helpers - fn rate_limit_snapshot_from_payload(payload: RateLimitStatusPayload) -> RateLimitSnapshot { - let rate_limit_details = payload - .rate_limit - .and_then(|inner| inner.map(|boxed| *boxed)); + fn rate_limit_snapshots_from_payload( + payload: RateLimitStatusPayload, + ) -> Vec { + let plan_type = Some(Self::map_plan_type(payload.plan_type)); + let mut snapshots = vec![Self::make_rate_limit_snapshot( + Some("codex".to_string()), + None, + payload.rate_limit.flatten().map(|details| *details), + payload.credits.flatten().map(|details| *details), + plan_type, + )]; + if let Some(additional) = payload.additional_rate_limits.flatten() { + snapshots.extend(additional.into_iter().map(|details| { + Self::make_rate_limit_snapshot( + Some(details.metered_feature), + Some(details.limit_name), + details.rate_limit.flatten().map(|rate_limit| *rate_limit), + None, + plan_type, + ) + })); + } + snapshots + } - let (primary, secondary) = if let Some(details) = rate_limit_details { - ( + fn make_rate_limit_snapshot( + limit_id: Option, + limit_name: Option, + rate_limit: Option, + credits: Option, + plan_type: Option, + ) -> RateLimitSnapshot { + let (primary, secondary) = match rate_limit { + Some(details) => ( Self::map_rate_limit_window(details.primary_window), Self::map_rate_limit_window(details.secondary_window), - ) - } else { - (None, None) + ), + None => (None, None), }; - RateLimitSnapshot { + limit_id, + limit_name, primary, secondary, - credits: Self::map_credits(payload.credits), - plan_type: Some(Self::map_plan_type(payload.plan_type)), + credits: Self::map_credits(credits), + plan_type, } } fn map_rate_limit_window( - window: Option>>, + window: Option>>, ) -> Option { - let snapshot = match window { - Some(Some(snapshot)) => *snapshot, - _ => return None, - }; + let snapshot = window.flatten().map(|details| *details)?; let used_percent = f64::from(snapshot.used_percent); let window_minutes = Self::window_minutes_from_seconds(snapshot.limit_window_seconds); @@ -335,16 +366,13 @@ impl Client { }) } - fn map_credits(credits: Option>>) -> Option { - let details = match credits { - Some(Some(details)) => *details, - _ => return None, - }; + fn map_credits(credits: Option) -> Option { + let details = credits?; Some(CreditsSnapshot { has_credits: details.has_credits, unlimited: details.unlimited, - balance: details.balance.and_then(|inner| inner), + balance: details.balance.flatten(), }) } @@ -374,3 +402,142 @@ impl Client { Some((seconds_i64 + 59) / 60) } } + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn usage_payload_maps_primary_and_additional_rate_limits() { + let payload = RateLimitStatusPayload { + plan_type: crate::types::PlanType::Pro, + rate_limit: Some(Some(Box::new(crate::types::RateLimitStatusDetails { + primary_window: Some(Some(Box::new(crate::types::RateLimitWindowSnapshot { + used_percent: 42, + limit_window_seconds: 300, + reset_after_seconds: 0, + reset_at: 123, + }))), + secondary_window: Some(Some(Box::new(crate::types::RateLimitWindowSnapshot { + used_percent: 84, + limit_window_seconds: 3600, + reset_after_seconds: 0, + reset_at: 456, + }))), + ..Default::default() + }))), + additional_rate_limits: Some(Some(vec![crate::types::AdditionalRateLimitDetails { + limit_name: "codex_other".to_string(), + metered_feature: "codex_other".to_string(), + rate_limit: Some(Some(Box::new(crate::types::RateLimitStatusDetails { + primary_window: Some(Some(Box::new(crate::types::RateLimitWindowSnapshot { + used_percent: 70, + limit_window_seconds: 900, + reset_after_seconds: 0, + reset_at: 789, + }))), + secondary_window: None, + ..Default::default() + }))), + }])), + credits: Some(Some(Box::new(crate::types::CreditStatusDetails { + has_credits: true, + unlimited: false, + balance: Some(Some("9.99".to_string())), + ..Default::default() + }))), + }; + + let snapshots = Client::rate_limit_snapshots_from_payload(payload); + assert_eq!(snapshots.len(), 2); + + assert_eq!(snapshots[0].limit_id.as_deref(), Some("codex")); + assert_eq!(snapshots[0].limit_name, None); + assert_eq!( + snapshots[0].primary.as_ref().map(|w| w.used_percent), + Some(42.0) + ); + assert_eq!( + snapshots[0].secondary.as_ref().map(|w| w.used_percent), + Some(84.0) + ); + assert_eq!( + snapshots[0].credits, + Some(CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("9.99".to_string()), + }) + ); + assert_eq!(snapshots[0].plan_type, Some(AccountPlanType::Pro)); + + assert_eq!(snapshots[1].limit_id.as_deref(), Some("codex_other")); + assert_eq!(snapshots[1].limit_name.as_deref(), Some("codex_other")); + assert_eq!( + snapshots[1].primary.as_ref().map(|w| w.used_percent), + Some(70.0) + ); + assert_eq!(snapshots[1].credits, None); + assert_eq!(snapshots[1].plan_type, Some(AccountPlanType::Pro)); + } + + #[test] + fn usage_payload_maps_zero_rate_limit_when_primary_absent() { + let payload = RateLimitStatusPayload { + plan_type: crate::types::PlanType::Plus, + rate_limit: None, + additional_rate_limits: Some(Some(vec![crate::types::AdditionalRateLimitDetails { + limit_name: "codex_other".to_string(), + metered_feature: "codex_other".to_string(), + rate_limit: None, + }])), + credits: None, + }; + + let snapshots = Client::rate_limit_snapshots_from_payload(payload); + assert_eq!(snapshots.len(), 2); + assert_eq!(snapshots[0].limit_id.as_deref(), Some("codex")); + assert_eq!(snapshots[0].limit_name, None); + assert_eq!(snapshots[0].primary, None); + assert_eq!(snapshots[1].limit_id.as_deref(), Some("codex_other")); + assert_eq!(snapshots[1].limit_name.as_deref(), Some("codex_other")); + } + + #[test] + fn preferred_snapshot_selection_matches_get_rate_limits_behavior() { + let snapshots = [ + RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 90.0, + window_minutes: Some(60), + resets_at: Some(1), + }), + secondary: None, + credits: None, + plan_type: Some(AccountPlanType::Pro), + }, + RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: Some("codex".to_string()), + primary: Some(RateLimitWindow { + used_percent: 10.0, + window_minutes: Some(60), + resets_at: Some(2), + }), + secondary: None, + credits: None, + plan_type: Some(AccountPlanType::Pro), + }, + ]; + + let preferred = snapshots + .iter() + .find(|snapshot| snapshot.limit_id.as_deref() == Some("codex")) + .cloned() + .unwrap_or_else(|| snapshots[0].clone()); + assert_eq!(preferred.limit_id.as_deref(), Some("codex")); + } +} diff --git a/codex-rs/backend-client/src/types.rs b/codex-rs/backend-client/src/types.rs index 9deeab79036..2559b2b39ed 100644 --- a/codex-rs/backend-client/src/types.rs +++ b/codex-rs/backend-client/src/types.rs @@ -1,3 +1,4 @@ +pub use codex_backend_openapi_models::models::AdditionalRateLimitDetails; pub use codex_backend_openapi_models::models::ConfigFileResponse; pub use codex_backend_openapi_models::models::CreditStatusDetails; pub use codex_backend_openapi_models::models::PaginatedListTaskListItem; diff --git a/codex-rs/codex-api/src/rate_limits.rs b/codex-rs/codex-api/src/rate_limits.rs index 047c5934bde..909ab06a275 100644 --- a/codex-rs/codex-api/src/rate_limits.rs +++ b/codex-rs/codex-api/src/rate_limits.rs @@ -4,6 +4,7 @@ use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use http::HeaderMap; use serde::Deserialize; +use std::collections::BTreeSet; use std::fmt::Display; #[derive(Debug)] @@ -17,25 +18,77 @@ impl Display for RateLimitError { } } -/// Parses the bespoke Codex rate-limit headers into a `RateLimitSnapshot`. -pub fn parse_rate_limit(headers: &HeaderMap) -> Option { +/// Parses the default Codex rate-limit header family into a `RateLimitSnapshot`. +pub fn parse_default_rate_limit(headers: &HeaderMap) -> Option { + parse_rate_limit_for_limit(headers, None) +} + +/// Parses all known rate-limit header families into update records keyed by limit id. +pub fn parse_all_rate_limits(headers: &HeaderMap) -> Vec { + let mut snapshots = Vec::new(); + if let Some(snapshot) = parse_default_rate_limit(headers) { + snapshots.push(snapshot); + } + + let mut limit_ids: BTreeSet = BTreeSet::new(); + + for name in headers.keys() { + let header_name = name.as_str().to_ascii_lowercase(); + if let Some(limit_id) = header_name_to_limit_id(&header_name) + && limit_id != "codex" + { + limit_ids.insert(limit_id); + } + } + + snapshots.extend(limit_ids.into_iter().filter_map(|limit_id| { + let snapshot = parse_rate_limit_for_limit(headers, Some(limit_id.as_str()))?; + has_rate_limit_data(&snapshot).then_some(snapshot) + })); + + snapshots +} + +/// Parses rate-limit headers for the provided limit id. +/// +/// `limit_id` should match the server-provided metered limit id (e.g. `codex`, +/// `codex_other`). When omitted, this defaults to the legacy `codex` header family. +pub fn parse_rate_limit_for_limit( + headers: &HeaderMap, + limit_id: Option<&str>, +) -> Option { + let normalized_limit = limit_id + .map(str::trim) + .filter(|name| !name.is_empty()) + .unwrap_or("codex") + .to_ascii_lowercase() + .replace('_', "-"); + let prefix = format!("x-{normalized_limit}"); let primary = parse_rate_limit_window( headers, - "x-codex-primary-used-percent", - "x-codex-primary-window-minutes", - "x-codex-primary-reset-at", + &format!("{prefix}-primary-used-percent"), + &format!("{prefix}-primary-window-minutes"), + &format!("{prefix}-primary-reset-at"), ); let secondary = parse_rate_limit_window( headers, - "x-codex-secondary-used-percent", - "x-codex-secondary-window-minutes", - "x-codex-secondary-reset-at", + &format!("{prefix}-secondary-used-percent"), + &format!("{prefix}-secondary-window-minutes"), + &format!("{prefix}-secondary-reset-at"), ); + let normalized_limit_id = normalize_limit_id(normalized_limit); let credits = parse_credits_snapshot(headers); + let limit_name_header = format!("{prefix}-limit-name"); + let parsed_limit_name = parse_header_str(headers, &limit_name_header) + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(std::string::ToString::to_string); Some(RateLimitSnapshot { + limit_id: Some(normalized_limit_id), + limit_name: parsed_limit_name, primary, secondary, credits, @@ -70,6 +123,8 @@ struct RateLimitEvent { plan_type: Option, rate_limits: Option, credits: Option, + metered_limit_name: Option, + limit_name: Option, } pub fn parse_rate_limit_event(payload: &str) -> Option { @@ -90,7 +145,13 @@ pub fn parse_rate_limit_event(payload: &str) -> Option { unlimited: credits.unlimited, balance: credits.balance, }); + let limit_id = event + .metered_limit_name + .or(event.limit_name) + .map(normalize_limit_id); Some(RateLimitSnapshot { + limit_id: Some(limit_id.unwrap_or_else(|| "codex".to_string())), + limit_name: None, primary, secondary, credits, @@ -178,3 +239,128 @@ fn parse_header_bool(headers: &HeaderMap, name: &str) -> Option { fn parse_header_str<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> { headers.get(name)?.to_str().ok() } + +fn has_rate_limit_data(snapshot: &RateLimitSnapshot) -> bool { + snapshot.primary.is_some() || snapshot.secondary.is_some() || snapshot.credits.is_some() +} + +fn header_name_to_limit_id(header_name: &str) -> Option { + let suffix = "-primary-used-percent"; + let prefix = header_name.strip_suffix(suffix)?; + let limit = prefix.strip_prefix("x-")?; + Some(normalize_limit_id(limit.to_string())) +} + +fn normalize_limit_id(name: impl Into) -> String { + name.into().trim().to_ascii_lowercase().replace('-', "_") +} + +#[cfg(test)] +mod tests { + use super::*; + use http::HeaderValue; + use pretty_assertions::assert_eq; + + #[test] + fn parse_rate_limit_for_limit_defaults_to_codex_headers() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-codex-primary-used-percent", + HeaderValue::from_static("12.5"), + ); + headers.insert( + "x-codex-primary-window-minutes", + HeaderValue::from_static("60"), + ); + headers.insert( + "x-codex-primary-reset-at", + HeaderValue::from_static("1704069000"), + ); + + let snapshot = parse_rate_limit_for_limit(&headers, None).expect("snapshot"); + assert_eq!(snapshot.limit_id.as_deref(), Some("codex")); + assert_eq!(snapshot.limit_name, None); + let primary = snapshot.primary.expect("primary"); + assert_eq!(primary.used_percent, 12.5); + assert_eq!(primary.window_minutes, Some(60)); + assert_eq!(primary.resets_at, Some(1704069000)); + } + + #[test] + fn parse_rate_limit_for_limit_reads_secondary_headers() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-codex-secondary-primary-used-percent", + HeaderValue::from_static("80"), + ); + headers.insert( + "x-codex-secondary-primary-window-minutes", + HeaderValue::from_static("1440"), + ); + headers.insert( + "x-codex-secondary-primary-reset-at", + HeaderValue::from_static("1704074400"), + ); + + let snapshot = + parse_rate_limit_for_limit(&headers, Some("codex_secondary")).expect("snapshot"); + assert_eq!(snapshot.limit_id.as_deref(), Some("codex_secondary")); + assert_eq!(snapshot.limit_name, None); + let primary = snapshot.primary.expect("primary"); + assert_eq!(primary.used_percent, 80.0); + assert_eq!(primary.window_minutes, Some(1440)); + assert_eq!(primary.resets_at, Some(1704074400)); + assert_eq!(snapshot.secondary, None); + } + + #[test] + fn parse_rate_limit_for_limit_prefers_limit_name_header() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-codex-bengalfox-primary-used-percent", + HeaderValue::from_static("80"), + ); + headers.insert( + "x-codex-bengalfox-limit-name", + HeaderValue::from_static("gpt-5.2-codex-sonic"), + ); + + let snapshot = + parse_rate_limit_for_limit(&headers, Some("codex_bengalfox")).expect("snapshot"); + assert_eq!(snapshot.limit_id.as_deref(), Some("codex_bengalfox")); + assert_eq!(snapshot.limit_name.as_deref(), Some("gpt-5.2-codex-sonic")); + } + + #[test] + fn parse_all_rate_limits_reads_all_limit_families() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-codex-primary-used-percent", + HeaderValue::from_static("12.5"), + ); + headers.insert( + "x-codex-secondary-primary-used-percent", + HeaderValue::from_static("80"), + ); + + let updates = parse_all_rate_limits(&headers); + assert_eq!(updates.len(), 2); + assert_eq!(updates[0].limit_id.as_deref(), Some("codex")); + assert_eq!(updates[1].limit_id.as_deref(), Some("codex_secondary")); + assert_eq!(updates[0].limit_name, None); + assert_eq!(updates[1].limit_name, None); + } + + #[test] + fn parse_all_rate_limits_includes_default_codex_snapshot() { + let headers = HeaderMap::new(); + + let updates = parse_all_rate_limits(&headers); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].limit_id.as_deref(), Some("codex")); + assert_eq!(updates[0].limit_name, None); + assert_eq!(updates[0].primary, None); + assert_eq!(updates[0].secondary, None); + assert_eq!(updates[0].credits, None); + } +} diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 95bda35fb2c..9e4d3456643 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -1,7 +1,7 @@ use crate::common::ResponseEvent; use crate::common::ResponseStream; use crate::error::ApiError; -use crate::rate_limits::parse_rate_limit; +use crate::rate_limits::parse_all_rate_limits; use crate::telemetry::SseTelemetry; use codex_client::ByteStream; use codex_client::StreamResponse; @@ -54,7 +54,7 @@ pub fn spawn_response_stream( telemetry: Option>, turn_state: Option>>, ) -> ResponseStream { - let rate_limits = parse_rate_limit(&stream_response.headers); + let rate_limit_snapshots = parse_all_rate_limits(&stream_response.headers); let models_etag = stream_response .headers .get("X-Models-Etag") @@ -74,7 +74,7 @@ pub fn spawn_response_stream( } let (tx_event, rx_event) = mpsc::channel::>(1600); tokio::spawn(async move { - if let Some(snapshot) = rate_limits { + for snapshot in rate_limit_snapshots { let _ = tx_event.send(Ok(ResponseEvent::RateLimits(snapshot))).await; } if let Some(etag) = models_etag { diff --git a/codex-rs/codex-backend-openapi-models/src/models/additional_rate_limit_details.rs b/codex-rs/codex-backend-openapi-models/src/models/additional_rate_limit_details.rs new file mode 100644 index 00000000000..d89e6561419 --- /dev/null +++ b/codex-rs/codex-backend-openapi-models/src/models/additional_rate_limit_details.rs @@ -0,0 +1,38 @@ +/* + * codex-backend + * + * codex-backend + * + * The version of the OpenAPI document: 0.0.1 + * + * Generated by: https://openapi-generator.tech + */ + +use crate::models; +use serde::Deserialize; +use serde::Serialize; + +#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] +pub struct AdditionalRateLimitDetails { + #[serde(rename = "limit_name")] + pub limit_name: String, + #[serde(rename = "metered_feature")] + pub metered_feature: String, + #[serde( + rename = "rate_limit", + default, + with = "::serde_with::rust::double_option", + skip_serializing_if = "Option::is_none" + )] + pub rate_limit: Option>>, +} + +impl AdditionalRateLimitDetails { + pub fn new(limit_name: String, metered_feature: String) -> AdditionalRateLimitDetails { + AdditionalRateLimitDetails { + limit_name, + metered_feature, + rate_limit: None, + } + } +} diff --git a/codex-rs/codex-backend-openapi-models/src/models/mod.rs b/codex-rs/codex-backend-openapi-models/src/models/mod.rs index 7072dede5e1..90bc0583342 100644 --- a/codex-rs/codex-backend-openapi-models/src/models/mod.rs +++ b/codex-rs/codex-backend-openapi-models/src/models/mod.rs @@ -27,6 +27,9 @@ pub mod paginated_list_task_list_item_; pub use self::paginated_list_task_list_item_::PaginatedListTaskListItem; // Rate Limits +pub mod additional_rate_limit_details; +pub use self::additional_rate_limit_details::AdditionalRateLimitDetails; + pub mod rate_limit_status_payload; pub use self::rate_limit_status_payload::PlanType; pub use self::rate_limit_status_payload::RateLimitStatusPayload; diff --git a/codex-rs/codex-backend-openapi-models/src/models/rate_limit_status_payload.rs b/codex-rs/codex-backend-openapi-models/src/models/rate_limit_status_payload.rs index daf8e6af39a..3d0492fa834 100644 --- a/codex-rs/codex-backend-openapi-models/src/models/rate_limit_status_payload.rs +++ b/codex-rs/codex-backend-openapi-models/src/models/rate_limit_status_payload.rs @@ -30,6 +30,13 @@ pub struct RateLimitStatusPayload { skip_serializing_if = "Option::is_none" )] pub credits: Option>>, + #[serde( + rename = "additional_rate_limits", + default, + with = "::serde_with::rust::double_option", + skip_serializing_if = "Option::is_none" + )] + pub additional_rate_limits: Option>>, } impl RateLimitStatusPayload { @@ -38,6 +45,7 @@ impl RateLimitStatusPayload { plan_type, rate_limit: None, credits: None, + additional_rate_limits: None, } } } diff --git a/codex-rs/core/src/api_bridge.rs b/codex-rs/core/src/api_bridge.rs index f7aeb570bde..af45182c13d 100644 --- a/codex-rs/core/src/api_bridge.rs +++ b/codex-rs/core/src/api_bridge.rs @@ -4,7 +4,7 @@ use codex_api::AuthProvider as ApiAuthProvider; use codex_api::TransportError; use codex_api::error::ApiError; use codex_api::rate_limits::parse_promo_message; -use codex_api::rate_limits::parse_rate_limit; +use codex_api::rate_limits::parse_rate_limit_for_limit; use http::HeaderMap; use serde::Deserialize; @@ -71,7 +71,10 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr { if let Ok(err) = serde_json::from_str::(&body_text) { if err.error.error_type.as_deref() == Some("usage_limit_reached") { - let rate_limits = headers.as_ref().and_then(parse_rate_limit); + let limit_id = extract_header(headers.as_ref(), ACTIVE_LIMIT_HEADER); + let rate_limits = headers.as_ref().and_then(|map| { + parse_rate_limit_for_limit(map, limit_id.as_deref()) + }); let promo_message = headers.as_ref().and_then(parse_promo_message); let resets_at = err .error @@ -80,8 +83,9 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr { return CodexErr::UsageLimitReached(UsageLimitReachedError { plan_type: err.error.plan_type, resets_at, - rate_limits, + rate_limits: rate_limits.map(Box::new), promo_message, + limit_name: limit_id, }); } else if err.error.error_type.as_deref() == Some("usage_not_included") { return CodexErr::UsageNotIncluded; @@ -117,6 +121,7 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr { const MODEL_CAP_MODEL_HEADER: &str = "x-codex-model-cap-model"; const MODEL_CAP_RESET_AFTER_HEADER: &str = "x-codex-model-cap-reset-after-seconds"; +const ACTIVE_LIMIT_HEADER: &str = "x-codex-active-limit"; const REQUEST_ID_HEADER: &str = "x-request-id"; const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id"; const CF_RAY_HEADER: &str = "cf-ray"; @@ -152,6 +157,33 @@ mod tests { assert_eq!(model_cap.model, "boomslang"); assert_eq!(model_cap.reset_after_seconds, Some(120)); } + + #[test] + fn map_api_error_maps_usage_limit_limit_name_header() { + let mut headers = HeaderMap::new(); + headers.insert( + ACTIVE_LIMIT_HEADER, + http::HeaderValue::from_static("codex_other"), + ); + let body = serde_json::json!({ + "error": { + "type": "usage_limit_reached", + "plan_type": "pro", + } + }) + .to_string(); + let err = map_api_error(ApiError::Transport(TransportError::Http { + status: StatusCode::TOO_MANY_REQUESTS, + url: Some("http://example.com/v1/responses".to_string()), + headers: Some(headers), + body: Some(body), + })); + + let CodexErr::UsageLimitReached(usage_limit) = err else { + panic!("expected CodexErr::UsageLimitReached, got {err:?}"); + }; + assert_eq!(usage_limit.limit_name.as_deref(), Some("codex_other")); + } } fn extract_request_tracking_id(headers: Option<&HeaderMap>) -> Option { diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index cd653793a6d..4a090ab252b 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -4279,7 +4279,7 @@ async fn run_sampling_request( Err(CodexErr::UsageLimitReached(e)) => { let rate_limits = e.rate_limits.clone(); if let Some(rate_limits) = rate_limits { - sess.update_rate_limits(&turn_context, rate_limits).await; + sess.update_rate_limits(&turn_context, *rate_limits).await; } return Err(CodexErr::UsageLimitReached(e)); } @@ -5820,6 +5820,8 @@ mod tests { let mut state = SessionState::new(session_configuration); let initial = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 10.0, window_minutes: Some(15), @@ -5836,6 +5838,8 @@ mod tests { state.set_rate_limits(initial.clone()); let update = RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), primary: Some(RateLimitWindow { used_percent: 40.0, window_minutes: Some(30), @@ -5854,6 +5858,8 @@ mod tests { assert_eq!( state.latest_rate_limits, Some(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), primary: update.primary.clone(), secondary: update.secondary, credits: initial.credits, @@ -5903,6 +5909,8 @@ mod tests { let mut state = SessionState::new(session_configuration); let initial = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 15.0, window_minutes: Some(20), @@ -5923,6 +5931,8 @@ mod tests { state.set_rate_limits(initial.clone()); let update = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 35.0, window_minutes: Some(25), @@ -5937,6 +5947,8 @@ mod tests { assert_eq!( state.latest_rate_limits, Some(RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: None, primary: update.primary, secondary: update.secondary, credits: initial.credits, diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 889335824ec..c284134d97c 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -409,12 +409,23 @@ impl std::fmt::Display for RetryLimitReachedError { pub struct UsageLimitReachedError { pub(crate) plan_type: Option, pub(crate) resets_at: Option>, - pub(crate) rate_limits: Option, + pub(crate) rate_limits: Option>, pub(crate) promo_message: Option, + pub(crate) limit_name: Option, } impl std::fmt::Display for UsageLimitReachedError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(limit_name) = self.limit_name.as_deref() + && !limit_name.eq_ignore_ascii_case("codex") + { + return write!( + f, + "You've hit your usage limit for {limit_name}.{}", + retry_suffix(self.resets_at.as_ref()) + ); + } + if let Some(promo_message) = &self.promo_message { return write!( f, @@ -699,6 +710,8 @@ mod tests { .unwrap() .timestamp(); RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 50.0, window_minutes: Some(60), @@ -728,8 +741,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Plus)), resets_at: None, - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; assert_eq!( err.to_string(), @@ -875,8 +889,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Free)), resets_at: None, - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; assert_eq!( err.to_string(), @@ -889,8 +904,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Go)), resets_at: None, - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; assert_eq!( err.to_string(), @@ -903,8 +919,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_at: None, - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; assert_eq!( err.to_string(), @@ -921,8 +938,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Team)), resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; let expected = format!( "You've hit your usage limit. To get more access now, send a request to your admin or try again at {expected_time}." @@ -936,8 +954,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Business)), resets_at: None, - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; assert_eq!( err.to_string(), @@ -950,8 +969,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Enterprise)), resets_at: None, - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; assert_eq!( err.to_string(), @@ -968,8 +988,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Pro)), resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; let expected = format!( "You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at {expected_time}." @@ -978,6 +999,29 @@ mod tests { }); } + #[test] + fn usage_limit_reached_error_hides_upsell_for_non_codex_limit_name() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::hours(1); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Plus)), + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: Some( + "Visit https://chatgpt.com/codex/settings/usage to purchase more credits" + .to_string(), + ), + limit_name: Some("codex_other".to_string()), + }; + let expected = format!( + "You've hit your usage limit for codex_other. Try again at {expected_time}." + ); + assert_eq!(err.to_string(), expected); + }); + } + #[test] fn usage_limit_reached_includes_minutes_when_available() { let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); @@ -987,8 +1031,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; let expected = format!("You've hit your usage limit. Try again at {expected_time}."); assert_eq!(err.to_string(), expected); @@ -1096,8 +1141,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some(PlanType::Known(KnownPlan::Plus)), resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; let expected = format!( "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at {expected_time}." @@ -1116,8 +1162,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; let expected = format!("You've hit your usage limit. Try again at {expected_time}."); assert_eq!(err.to_string(), expected); @@ -1133,8 +1180,9 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: None, + limit_name: None, }; let expected = format!("You've hit your usage limit. Try again at {expected_time}."); assert_eq!(err.to_string(), expected); @@ -1150,10 +1198,11 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_at: Some(resets_at), - rate_limits: Some(rate_limit_snapshot()), + rate_limits: Some(Box::new(rate_limit_snapshot())), promo_message: Some( "To continue using Codex, start a free trial of today".to_string(), ), + limit_name: None, }; let expected = format!( "You've hit your usage limit. To continue using Codex, start a free trial of today, or try again at {expected_time}." diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index dcbf817d99c..16614f3a6d8 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -170,11 +170,21 @@ impl SessionState { } } -// Sometimes new snapshots don't include credits or plan information. +// Merge partial rate-limit updates: new fields overwrite existing values; +// missing fields retain prior values. If `limit_id` is absent everywhere, +// default it to `"codex"`. fn merge_rate_limit_fields( previous: Option<&RateLimitSnapshot>, mut snapshot: RateLimitSnapshot, ) -> RateLimitSnapshot { + if snapshot.limit_id.is_none() { + snapshot.limit_id = previous + .and_then(|prior| prior.limit_id.clone()) + .or_else(|| Some("codex".to_string())); + } + if snapshot.limit_name.is_none() { + snapshot.limit_name = previous.and_then(|prior| prior.limit_name.clone()); + } if snapshot.credits.is_none() { snapshot.credits = previous.and_then(|prior| prior.credits.clone()); } @@ -188,6 +198,7 @@ fn merge_rate_limit_fields( mod tests { use super::*; use crate::codex::make_session_configuration_for_tests; + use crate::protocol::RateLimitWindow; use pretty_assertions::assert_eq; #[tokio::test] @@ -258,4 +269,126 @@ mod tests { assert_eq!(state.get_mcp_tool_selection(), None); } + + #[tokio::test] + async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + + state.set_rate_limits(RateLimitSnapshot { + limit_id: None, + limit_name: None, + primary: Some(RateLimitWindow { + used_percent: 12.0, + window_minutes: Some(60), + resets_at: Some(100), + }), + secondary: None, + credits: None, + plan_type: None, + }); + + assert_eq!( + state + .latest_rate_limits + .as_ref() + .and_then(|v| v.limit_id.clone()), + Some("codex".to_string()) + ); + } + + #[tokio::test] + async fn set_rate_limits_preserves_previous_limit_id_when_missing() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + + state.set_rate_limits(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 20.0, + window_minutes: Some(60), + resets_at: Some(200), + }), + secondary: None, + credits: None, + plan_type: None, + }); + state.set_rate_limits(RateLimitSnapshot { + limit_id: None, + limit_name: None, + primary: Some(RateLimitWindow { + used_percent: 30.0, + window_minutes: Some(60), + resets_at: Some(300), + }), + secondary: None, + credits: None, + plan_type: None, + }); + + assert_eq!( + state + .latest_rate_limits + .as_ref() + .and_then(|v| v.limit_id.clone()), + Some("codex_other".to_string()) + ); + } + + #[tokio::test] + async fn set_rate_limits_accepts_new_limit_id_bucket() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + + state.set_rate_limits(RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: Some("codex".to_string()), + primary: Some(RateLimitWindow { + used_percent: 10.0, + window_minutes: Some(60), + resets_at: Some(100), + }), + secondary: None, + credits: Some(crate::protocol::CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("50".to_string()), + }), + plan_type: Some(codex_protocol::account::PlanType::Plus), + }); + + state.set_rate_limits(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 30.0, + window_minutes: Some(120), + resets_at: Some(200), + }), + secondary: None, + credits: None, + plan_type: None, + }); + + assert_eq!( + state.latest_rate_limits, + Some(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 30.0, + window_minutes: Some(120), + resets_at: Some(200), + }), + secondary: None, + credits: Some(crate::protocol::CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("50".to_string()), + }), + plan_type: Some(codex_protocol::account::PlanType::Plus), + }) + ); + } } diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index cb5668ec548..62d757e16e8 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1509,6 +1509,8 @@ async fn token_count_includes_rate_limits_snapshot() { json!({ "info": null, "rate_limits": { + "limit_id": "codex", + "limit_name": null, "primary": { "used_percent": 12.5, "window_minutes": 10, @@ -1558,6 +1560,8 @@ async fn token_count_includes_rate_limits_snapshot() { "model_context_window": 258400 }, "rate_limits": { + "limit_id": "codex", + "limit_name": null, "primary": { "used_percent": 12.5, "window_minutes": 10, @@ -1630,6 +1634,8 @@ async fn usage_limit_error_emits_rate_limit_event() -> anyhow::Result<()> { let codex = codex_fixture.codex.clone(); let expected_limits = json!({ + "limit_id": "codex", + "limit_name": null, "primary": { "used_percent": 100.0, "window_minutes": 15, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index d1ebc637772..0f06e9e8c21 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -455,6 +455,8 @@ async fn responses_websocket_usage_limit_error_emits_rate_limit_event() { json!({ "info": null, "rate_limits": { + "limit_id": "codex", + "limit_name": null, "primary": { "used_percent": 100.0, "window_minutes": 15, diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 2368a39ba33..450d7dd1d74 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1249,6 +1249,8 @@ pub struct TokenCountEvent { #[derive(Debug, Clone, PartialEq, Deserialize, Serialize, JsonSchema, TS)] pub struct RateLimitSnapshot { + pub limit_id: Option, + pub limit_name: Option, pub primary: Option, pub secondary: Option, pub credits: Option, diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index e507a39a4dd..bf4ef4dec6d 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -25,6 +25,7 @@ //! the final answer. During streaming we hide the status row to avoid duplicate //! progress indicators; once commentary completes and stream queues drain, we //! re-show it so users still see turn-in-progress state between output bursts. +use std::collections::BTreeMap; use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; @@ -41,6 +42,7 @@ use crate::bottom_pane::StatusLineSetupView; use crate::status::RateLimitWindowDisplay; use crate::status::format_directory_display; use crate::status::format_tokens_compact; +use crate::status::rate_limit_snapshot_display_for_limit; use crate::text_formatting::proper_join; use crate::version::CODEX_CLI_VERSION; use codex_app_server_protocol::ConfigLayerSource; @@ -511,7 +513,7 @@ pub(crate) struct ChatWidget { session_header: SessionHeader, initial_user_message: Option, token_info: Option, - rate_limit_snapshot: Option, + rate_limit_snapshots_by_limit_id: BTreeMap, plan_type: Option, rate_limit_warnings: RateLimitWarningState, rate_limit_switch_prompt: RateLimitSwitchPromptState, @@ -1498,10 +1500,18 @@ impl ChatWidget { pub(crate) fn on_rate_limit_snapshot(&mut self, snapshot: Option) { if let Some(mut snapshot) = snapshot { + let limit_id = snapshot + .limit_id + .clone() + .unwrap_or_else(|| "codex".to_string()); + let limit_label = snapshot + .limit_name + .clone() + .unwrap_or_else(|| limit_id.clone()); if snapshot.credits.is_none() { snapshot.credits = self - .rate_limit_snapshot - .as_ref() + .rate_limit_snapshots_by_limit_id + .get(&limit_id) .and_then(|display| display.credits.as_ref()) .map(|credits| CreditsSnapshot { has_credits: credits.has_credits, @@ -1512,32 +1522,38 @@ impl ChatWidget { self.plan_type = snapshot.plan_type.or(self.plan_type); - let warnings = self.rate_limit_warnings.take_warnings( - snapshot - .secondary - .as_ref() - .map(|window| window.used_percent), - snapshot - .secondary - .as_ref() - .and_then(|window| window.window_minutes), - snapshot.primary.as_ref().map(|window| window.used_percent), - snapshot - .primary - .as_ref() - .and_then(|window| window.window_minutes), - ); + let is_codex_limit = limit_id.eq_ignore_ascii_case("codex"); + let warnings = if is_codex_limit { + self.rate_limit_warnings.take_warnings( + snapshot + .secondary + .as_ref() + .map(|window| window.used_percent), + snapshot + .secondary + .as_ref() + .and_then(|window| window.window_minutes), + snapshot.primary.as_ref().map(|window| window.used_percent), + snapshot + .primary + .as_ref() + .and_then(|window| window.window_minutes), + ) + } else { + vec![] + }; - let high_usage = snapshot - .secondary - .as_ref() - .map(|w| w.used_percent >= RATE_LIMIT_SWITCH_PROMPT_THRESHOLD) - .unwrap_or(false) - || snapshot - .primary + let high_usage = is_codex_limit + && (snapshot + .secondary .as_ref() .map(|w| w.used_percent >= RATE_LIMIT_SWITCH_PROMPT_THRESHOLD) - .unwrap_or(false); + .unwrap_or(false) + || snapshot + .primary + .as_ref() + .map(|w| w.used_percent >= RATE_LIMIT_SWITCH_PROMPT_THRESHOLD) + .unwrap_or(false)); if high_usage && !self.rate_limit_switch_prompt_hidden() @@ -1550,8 +1566,10 @@ impl ChatWidget { self.rate_limit_switch_prompt = RateLimitSwitchPromptState::Pending; } - let display = crate::status::rate_limit_snapshot_display(&snapshot, Local::now()); - self.rate_limit_snapshot = Some(display); + let display = + rate_limit_snapshot_display_for_limit(&snapshot, limit_label, Local::now()); + self.rate_limit_snapshots_by_limit_id + .insert(limit_id, display); if !warnings.is_empty() { for warning in warnings { @@ -1560,7 +1578,7 @@ impl ChatWidget { self.request_redraw(); } } else { - self.rate_limit_snapshot = None; + self.rate_limit_snapshots_by_limit_id.clear(); } self.refresh_status_line(); } @@ -2608,7 +2626,7 @@ impl ChatWidget { session_header: SessionHeader::new(header_model), initial_user_message, token_info: None, - rate_limit_snapshot: None, + rate_limit_snapshots_by_limit_id: BTreeMap::new(), plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -2773,7 +2791,7 @@ impl ChatWidget { session_header: SessionHeader::new(header_model), initial_user_message, token_info: None, - rate_limit_snapshot: None, + rate_limit_snapshots_by_limit_id: BTreeMap::new(), plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -2927,7 +2945,7 @@ impl ChatWidget { session_header: SessionHeader::new(header_model), initial_user_message, token_info: None, - rate_limit_snapshot: None, + rate_limit_snapshots_by_limit_id: BTreeMap::new(), plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -4222,7 +4240,12 @@ impl ChatWidget { .unwrap_or(&default_usage); let collaboration_mode = self.collaboration_mode_label(); let reasoning_effort_override = Some(self.effective_reasoning_effort()); - self.add_to_history(crate::status::new_status_output( + let rate_limit_snapshots: Vec = self + .rate_limit_snapshots_by_limit_id + .values() + .cloned() + .collect(); + self.add_to_history(crate::status::new_status_output_with_rate_limits( &self.config, self.auth_manager.as_ref(), token_info, @@ -4230,7 +4253,7 @@ impl ChatWidget { &self.thread_id, self.thread_name.clone(), self.forked_from, - self.rate_limit_snapshot.as_ref(), + rate_limit_snapshots.as_slice(), self.plan_type, Local::now(), self.model_display_name(), @@ -4376,8 +4399,8 @@ impl ChatWidget { .map(|used| format!("{used}% used")), StatusLineItem::FiveHourLimit => { let window = self - .rate_limit_snapshot - .as_ref() + .rate_limit_snapshots_by_limit_id + .get("codex") .and_then(|s| s.primary.as_ref()); let label = window .and_then(|window| window.window_minutes) @@ -4387,8 +4410,8 @@ impl ChatWidget { } StatusLineItem::WeeklyLimit => { let window = self - .rate_limit_snapshot - .as_ref() + .rate_limit_snapshots_by_limit_id + .get("codex") .and_then(|s| s.secondary.as_ref()); let label = window .and_then(|window| window.window_minutes) @@ -4572,9 +4595,10 @@ impl ChatWidget { loop { if let Some(auth) = auth_manager.auth().await && auth.is_chatgpt_auth() - && let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth).await { - app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); + for snapshot in fetch_rate_limits(base_url.clone(), auth).await { + app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); + } } interval.tick().await; } @@ -7140,18 +7164,18 @@ fn extract_first_bold(s: &str) -> Option { None } -async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option { +async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Vec { match BackendClient::from_auth(base_url, &auth) { - Ok(client) => match client.get_rate_limits().await { - Ok(snapshot) => Some(snapshot), + Ok(client) => match client.get_rate_limits_many().await { + Ok(snapshots) => snapshots, Err(err) => { debug!(error = ?err, "failed to fetch rate limits from /usage"); - None + Vec::new() } }, Err(err) => { debug!(error = ?err, "failed to construct backend client for rate limits"); - None + Vec::new() } } } diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index 1f1af74a115..9a2b1a16c4a 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -95,6 +95,7 @@ use insta::assert_snapshot; use pretty_assertions::assert_eq; #[cfg(target_os = "windows")] use serial_test::serial; +use std::collections::BTreeMap; use std::collections::HashSet; use std::path::PathBuf; use tempfile::NamedTempFile; @@ -124,6 +125,8 @@ fn invalid_value(candidate: impl Into, allowed: impl Into) -> Co fn snapshot(percent: f64) -> RateLimitSnapshot { RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: percent, window_minutes: Some(60), @@ -1061,7 +1064,7 @@ async fn make_chatwidget_manual( session_header: SessionHeader::new(resolved_model.clone()), initial_user_message: None, token_info: None, - rate_limit_snapshot: None, + rate_limit_snapshots_by_limit_id: BTreeMap::new(), plan_type: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), @@ -1277,6 +1280,8 @@ async fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: None, secondary: None, credits: Some(CreditsSnapshot { @@ -1287,13 +1292,15 @@ async fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() { plan_type: None, })); let initial_balance = chat - .rate_limit_snapshot - .as_ref() + .rate_limit_snapshots_by_limit_id + .get("codex") .and_then(|snapshot| snapshot.credits.as_ref()) .and_then(|credits| credits.balance.as_deref()); assert_eq!(initial_balance, Some("17.5")); chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 80.0, window_minutes: Some(60), @@ -1305,8 +1312,8 @@ async fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() { })); let display = chat - .rate_limit_snapshot - .as_ref() + .rate_limit_snapshots_by_limit_id + .get("codex") .expect("rate limits should be cached"); let credits = display .credits @@ -1326,6 +1333,8 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 10.0, window_minutes: Some(60), @@ -1342,6 +1351,8 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() { assert_eq!(chat.plan_type, Some(PlanType::Plus)); chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 25.0, window_minutes: Some(30), @@ -1358,6 +1369,8 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() { assert_eq!(chat.plan_type, Some(PlanType::Pro)); chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 30.0, window_minutes: Some(60), @@ -1374,6 +1387,61 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() { assert_eq!(chat.plan_type, Some(PlanType::Pro)); } +#[tokio::test] +async fn rate_limit_snapshots_keep_separate_entries_per_limit_id() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; + + chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: Some("codex".to_string()), + primary: Some(RateLimitWindow { + used_percent: 20.0, + window_minutes: Some(300), + resets_at: Some(100), + }), + secondary: None, + credits: Some(CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("5.00".to_string()), + }), + plan_type: Some(PlanType::Pro), + })); + + chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 90.0, + window_minutes: Some(60), + resets_at: Some(200), + }), + secondary: None, + credits: None, + plan_type: Some(PlanType::Pro), + })); + + let codex = chat + .rate_limit_snapshots_by_limit_id + .get("codex") + .expect("codex snapshot should exist"); + let other = chat + .rate_limit_snapshots_by_limit_id + .get("codex_other") + .expect("codex_other snapshot should exist"); + + assert_eq!(codex.primary.as_ref().map(|w| w.used_percent), Some(20.0)); + assert_eq!( + codex + .credits + .as_ref() + .and_then(|credits| credits.balance.as_deref()), + Some("5.00") + ); + assert_eq!(other.primary.as_ref().map(|w| w.used_percent), Some(90.0)); + assert!(other.credits.is_none()); +} + #[tokio::test] async fn rate_limit_switch_prompt_skips_when_on_lower_cost_model() { let (mut chat, _, _) = make_chatwidget_manual(Some(NUDGE_MODEL_SLUG)).await; @@ -1388,6 +1456,31 @@ async fn rate_limit_switch_prompt_skips_when_on_lower_cost_model() { )); } +#[tokio::test] +async fn rate_limit_switch_prompt_skips_non_codex_limit() { + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let (mut chat, _, _) = make_chatwidget_manual(Some("gpt-5")).await; + chat.auth_manager = AuthManager::from_auth_for_testing(auth); + + chat.on_rate_limit_snapshot(Some(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + primary: Some(RateLimitWindow { + used_percent: 95.0, + window_minutes: Some(60), + resets_at: None, + }), + secondary: None, + credits: None, + plan_type: None, + })); + + assert!(matches!( + chat.rate_limit_switch_prompt, + RateLimitSwitchPromptState::Idle + )); +} + #[tokio::test] async fn rate_limit_switch_prompt_shows_once_per_session() { let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); diff --git a/codex-rs/tui/src/status/card.rs b/codex-rs/tui/src/status/card.rs index 23005274db0..60f7f937526 100644 --- a/codex-rs/tui/src/status/card.rs +++ b/codex-rs/tui/src/status/card.rs @@ -36,6 +36,7 @@ use super::rate_limits::StatusRateLimitData; use super::rate_limits::StatusRateLimitRow; use super::rate_limits::StatusRateLimitValue; use super::rate_limits::compose_rate_limit_data; +use super::rate_limits::compose_rate_limit_data_many; use super::rate_limits::format_status_limit_summary; use super::rate_limits::render_status_limit_progress_bar; use crate::wrapping::RtOptions; @@ -75,6 +76,7 @@ struct StatusHistoryCell { rate_limits: StatusRateLimitData, } +#[cfg(test)] #[allow(clippy::too_many_arguments)] pub(crate) fn new_status_output( config: &Config, @@ -90,6 +92,40 @@ pub(crate) fn new_status_output( model_name: &str, collaboration_mode: Option<&str>, reasoning_effort_override: Option>, +) -> CompositeHistoryCell { + let snapshots = rate_limits.map(std::slice::from_ref).unwrap_or_default(); + new_status_output_with_rate_limits( + config, + auth_manager, + token_info, + total_usage, + session_id, + thread_name, + forked_from, + snapshots, + plan_type, + now, + model_name, + collaboration_mode, + reasoning_effort_override, + ) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn new_status_output_with_rate_limits( + config: &Config, + auth_manager: &AuthManager, + token_info: Option<&TokenUsageInfo>, + total_usage: &TokenUsage, + session_id: &Option, + thread_name: Option, + forked_from: Option, + rate_limits: &[RateLimitSnapshotDisplay], + plan_type: Option, + now: DateTime, + model_name: &str, + collaboration_mode: Option<&str>, + reasoning_effort_override: Option>, ) -> CompositeHistoryCell { let command = PlainHistoryCell::new(vec!["/status".magenta().into()]); let card = StatusHistoryCell::new( @@ -121,7 +157,7 @@ impl StatusHistoryCell { session_id: &Option, thread_name: Option, forked_from: Option, - rate_limits: Option<&RateLimitSnapshotDisplay>, + rate_limits: &[RateLimitSnapshotDisplay], plan_type: Option, now: DateTime, model_name: &str, @@ -189,7 +225,11 @@ impl StatusHistoryCell { output: total_usage.output_tokens, context_window, }; - let rate_limits = compose_rate_limit_data(rate_limits, now); + let rate_limits = if rate_limits.len() <= 1 { + compose_rate_limit_data(rate_limits.first(), now) + } else { + compose_rate_limit_data_many(rate_limits, now) + }; Self { model_name, diff --git a/codex-rs/tui/src/status/mod.rs b/codex-rs/tui/src/status/mod.rs index 4f46af1cbc0..90e406a52db 100644 --- a/codex-rs/tui/src/status/mod.rs +++ b/codex-rs/tui/src/status/mod.rs @@ -12,12 +12,16 @@ mod format; mod helpers; mod rate_limits; +#[cfg(test)] pub(crate) use card::new_status_output; +pub(crate) use card::new_status_output_with_rate_limits; pub(crate) use helpers::format_directory_display; pub(crate) use helpers::format_tokens_compact; pub(crate) use rate_limits::RateLimitSnapshotDisplay; pub(crate) use rate_limits::RateLimitWindowDisplay; +#[cfg(test)] pub(crate) use rate_limits::rate_limit_snapshot_display; +pub(crate) use rate_limits::rate_limit_snapshot_display_for_limit; #[cfg(test)] mod tests; diff --git a/codex-rs/tui/src/status/rate_limits.rs b/codex-rs/tui/src/status/rate_limits.rs index 830fe158b2e..8d46bc8c99b 100644 --- a/codex-rs/tui/src/status/rate_limits.rs +++ b/codex-rs/tui/src/status/rate_limits.rs @@ -86,6 +86,8 @@ impl RateLimitWindowDisplay { #[derive(Debug, Clone)] pub(crate) struct RateLimitSnapshotDisplay { + /// Canonical limit identifier (for example: `codex` or `codex_other`). + pub limit_name: String, /// Local timestamp representing when this display snapshot was captured. pub captured_at: DateTime, /// Primary usage window (typically short duration). @@ -111,11 +113,21 @@ pub(crate) struct CreditsSnapshotDisplay { /// /// Pass the timestamp from the same observation point as `snapshot`; supplying a significantly /// older or newer `captured_at` can produce misleading reset labels and stale classification. +#[cfg(test)] pub(crate) fn rate_limit_snapshot_display( snapshot: &RateLimitSnapshot, captured_at: DateTime, +) -> RateLimitSnapshotDisplay { + rate_limit_snapshot_display_for_limit(snapshot, "codex".to_string(), captured_at) +} + +pub(crate) fn rate_limit_snapshot_display_for_limit( + snapshot: &RateLimitSnapshot, + limit_name: String, + captured_at: DateTime, ) -> RateLimitSnapshotDisplay { RateLimitSnapshotDisplay { + limit_name, captured_at, primary: snapshot .primary @@ -148,57 +160,120 @@ pub(crate) fn compose_rate_limit_data( now: DateTime, ) -> StatusRateLimitData { match snapshot { - Some(snapshot) => { - let mut rows = Vec::with_capacity(3); + Some(snapshot) => compose_rate_limit_data_many(std::slice::from_ref(snapshot), now), + None => StatusRateLimitData::Missing, + } +} - if let Some(primary) = snapshot.primary.as_ref() { - let label: String = primary +pub(crate) fn compose_rate_limit_data_many( + snapshots: &[RateLimitSnapshotDisplay], + now: DateTime, +) -> StatusRateLimitData { + if snapshots.is_empty() { + return StatusRateLimitData::Missing; + } + + let mut rows = Vec::with_capacity(snapshots.len().saturating_mul(3)); + let mut stale = false; + + for snapshot in snapshots { + stale |= now.signed_duration_since(snapshot.captured_at) + > ChronoDuration::minutes(RATE_LIMIT_STALE_THRESHOLD_MINUTES); + + let limit_bucket_label = snapshot.limit_name.clone(); + let show_limit_prefix = !limit_bucket_label.eq_ignore_ascii_case("codex"); + let primary_label = snapshot + .primary + .as_ref() + .map(|window| { + window .window_minutes .map(get_limits_duration) - .unwrap_or_else(|| "5h".to_string()); - let label = capitalize_first(&label); - rows.push(StatusRateLimitRow { - label: format!("{label} limit"), - value: StatusRateLimitValue::Window { - percent_used: primary.used_percent, - resets_at: primary.resets_at.clone(), - }, - }); - } - - if let Some(secondary) = snapshot.secondary.as_ref() { - let label: String = secondary + .unwrap_or_else(|| "5h".to_string()) + }) + .map(|label| capitalize_first(&label)); + let secondary_label = snapshot + .secondary + .as_ref() + .map(|window| { + window .window_minutes .map(get_limits_duration) - .unwrap_or_else(|| "weekly".to_string()); - let label = capitalize_first(&label); - rows.push(StatusRateLimitRow { - label: format!("{label} limit"), - value: StatusRateLimitValue::Window { - percent_used: secondary.used_percent, - resets_at: secondary.resets_at.clone(), - }, - }); - } - - if let Some(credits) = snapshot.credits.as_ref() - && let Some(row) = credit_status_row(credits) - { - rows.push(row); - } - - let is_stale = now.signed_duration_since(snapshot.captured_at) - > ChronoDuration::minutes(RATE_LIMIT_STALE_THRESHOLD_MINUTES); - - if rows.is_empty() { - StatusRateLimitData::Available(vec![]) - } else if is_stale { - StatusRateLimitData::Stale(rows) + .unwrap_or_else(|| "weekly".to_string()) + }) + .map(|label| capitalize_first(&label)); + let window_count = + usize::from(snapshot.primary.is_some()) + usize::from(snapshot.secondary.is_some()); + let combine_non_codex_single_limit = show_limit_prefix && window_count == 1; + + if show_limit_prefix && !combine_non_codex_single_limit { + rows.push(StatusRateLimitRow { + label: format!("{limit_bucket_label} limit"), + value: StatusRateLimitValue::Text(String::new()), + }); + } + + if let Some(primary) = snapshot.primary.as_ref() { + let label = if combine_non_codex_single_limit { + format!( + "{} {} limit", + limit_bucket_label, + primary_label.clone().unwrap_or_else(|| "5h".to_string()) + ) } else { - StatusRateLimitData::Available(rows) - } + format!( + "{} limit", + primary_label.clone().unwrap_or_else(|| "5h".to_string()) + ) + }; + rows.push(StatusRateLimitRow { + label, + value: StatusRateLimitValue::Window { + percent_used: primary.used_percent, + resets_at: primary.resets_at.clone(), + }, + }); } - None => StatusRateLimitData::Missing, + + if let Some(secondary) = snapshot.secondary.as_ref() { + let label = if combine_non_codex_single_limit { + format!( + "{} {} limit", + limit_bucket_label, + secondary_label + .clone() + .unwrap_or_else(|| "weekly".to_string()) + ) + } else { + format!( + "{} limit", + secondary_label + .clone() + .unwrap_or_else(|| "weekly".to_string()) + ) + }; + rows.push(StatusRateLimitRow { + label, + value: StatusRateLimitValue::Window { + percent_used: secondary.used_percent, + resets_at: secondary.resets_at.clone(), + }, + }); + } + + if let Some(credits) = snapshot.credits.as_ref() + && let Some(row) = credit_status_row(credits) + { + rows.push(row); + } + } + + if rows.is_empty() { + StatusRateLimitData::Available(vec![]) + } else if stale { + StatusRateLimitData::Stale(rows) + } else { + StatusRateLimitData::Available(rows) } } @@ -266,3 +341,100 @@ fn format_credit_balance(raw: &str) -> Option { None } + +#[cfg(test)] +mod tests { + use super::CreditsSnapshotDisplay; + use super::RateLimitSnapshotDisplay; + use super::RateLimitWindowDisplay; + use super::StatusRateLimitData; + use super::compose_rate_limit_data_many; + use chrono::Local; + use pretty_assertions::assert_eq; + + fn window(used_percent: f64) -> RateLimitWindowDisplay { + RateLimitWindowDisplay { + used_percent, + resets_at: Some("soon".to_string()), + window_minutes: Some(300), + } + } + + #[test] + fn non_codex_single_limit_renders_combined_row() { + let now = Local::now(); + let codex = RateLimitSnapshotDisplay { + limit_name: "codex".to_string(), + captured_at: now, + primary: Some(window(10.0)), + secondary: None, + credits: Some(CreditsSnapshotDisplay { + has_credits: true, + unlimited: false, + balance: Some("25".to_string()), + }), + }; + let other = RateLimitSnapshotDisplay { + limit_name: "codex-other".to_string(), + captured_at: now, + primary: Some(window(20.0)), + secondary: None, + credits: Some(CreditsSnapshotDisplay { + has_credits: true, + unlimited: false, + balance: Some("99".to_string()), + }), + }; + + let rows = match compose_rate_limit_data_many(&[codex, other], now) { + StatusRateLimitData::Available(rows) => rows, + other => panic!("unexpected status: {other:?}"), + }; + + let labels: Vec = rows.iter().map(|row| row.label.clone()).collect(); + assert_eq!( + labels, + vec![ + "5h limit".to_string(), + "Credits".to_string(), + "codex-other 5h limit".to_string(), + "Credits".to_string(), + ] + ); + assert_eq!(rows.iter().filter(|row| row.label == "Credits").count(), 2); + } + + #[test] + fn non_codex_multi_limit_keeps_group_row() { + let now = Local::now(); + let other = RateLimitSnapshotDisplay { + limit_name: "codex-other".to_string(), + captured_at: now, + primary: Some(RateLimitWindowDisplay { + used_percent: 20.0, + resets_at: Some("soon".to_string()), + window_minutes: Some(60), + }), + secondary: Some(RateLimitWindowDisplay { + used_percent: 40.0, + resets_at: Some("later".to_string()), + window_minutes: None, + }), + credits: None, + }; + + let rows = match compose_rate_limit_data_many(&[other], now) { + StatusRateLimitData::Available(rows) => rows, + other => panic!("unexpected status: {other:?}"), + }; + let labels: Vec = rows.iter().map(|row| row.label.clone()).collect(); + assert_eq!( + labels, + vec![ + "codex-other limit".to_string(), + "1h limit".to_string(), + "Weekly limit".to_string(), + ] + ); + } +} diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index c8844904abe..bff2fab0488 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -122,6 +122,8 @@ async fn status_snapshot_includes_reasoning_details() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 72.5, window_minutes: Some(300), @@ -242,6 +244,8 @@ async fn status_snapshot_includes_monthly_limit() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 12.0, window_minutes: Some(43_200), @@ -291,6 +295,8 @@ async fn status_snapshot_shows_unlimited_credits() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: None, secondary: None, credits: Some(CreditsSnapshot { @@ -338,6 +344,8 @@ async fn status_snapshot_shows_positive_credits() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: None, secondary: None, credits: Some(CreditsSnapshot { @@ -385,6 +393,8 @@ async fn status_snapshot_hides_zero_credits() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: None, secondary: None, credits: Some(CreditsSnapshot { @@ -430,6 +440,8 @@ async fn status_snapshot_hides_when_has_no_credits_flag() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: None, secondary: None, credits: Some(CreditsSnapshot { @@ -533,6 +545,8 @@ async fn status_snapshot_truncates_in_narrow_terminal() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 72.5, window_minutes: Some(300), @@ -642,6 +656,8 @@ async fn status_snapshot_includes_credits_and_limits() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 45.0, window_minutes: Some(300), @@ -705,6 +721,8 @@ async fn status_snapshot_shows_empty_limits_message() { }; let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: None, secondary: None, credits: None, @@ -764,6 +782,8 @@ async fn status_snapshot_shows_stale_limits_message() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 72.5, window_minutes: Some(300), @@ -828,6 +848,8 @@ async fn status_snapshot_cached_limits_hide_credits_without_flag() { .single() .expect("timestamp"); let snapshot = RateLimitSnapshot { + limit_id: None, + limit_name: None, primary: Some(RateLimitWindow { used_percent: 60.0, window_minutes: Some(300),