From fdd0cd1de974bebe26c8c31d955eea628c8ef96a Mon Sep 17 00:00:00 2001 From: xl-openai Date: Tue, 10 Feb 2026 20:09:31 -0800 Subject: [PATCH 1/2] feat: support multiple rate limits (#11260) Added multi-limit support end-to-end by carrying limit_name in rate-limit snapshots and handling multiple buckets instead of only codex. Extended /usage client parsing to consume additional_rate_limits Updated TUI /status and in-memory state to store/render per-limit snapshots Extended app-server rate-limit read response: kept rate_limits and added rate_limits_by_name. Adjusted usage-limit error messaging for non-default codex limit buckets --- .../schema/json/EventMsg.json | 12 + .../schema/json/ServerNotification.json | 24 ++ .../codex_app_server_protocol.schemas.json | 41 ++- .../json/v1/ForkConversationResponse.json | 12 + .../json/v1/ResumeConversationResponse.json | 12 + .../v1/SessionConfiguredNotification.json | 12 + .../AccountRateLimitsUpdatedNotification.json | 12 + .../json/v2/GetAccountRateLimitsResponse.json | 29 +- .../schema/typescript/RateLimitSnapshot.ts | 2 +- .../v2/GetAccountRateLimitsResponse.ts | 10 +- .../schema/typescript/v2/RateLimitSnapshot.ts | 2 +- .../app-server-protocol/src/protocol/v2.rs | 7 + .../app-server/src/bespoke_event_handling.rs | 4 + .../app-server/src/codex_message_processor.rs | 49 +++- codex-rs/app-server/src/outgoing_message.rs | 6 +- .../app-server/tests/suite/v2/rate_limits.rs | 60 +++- codex-rs/backend-client/src/client.rs | 219 ++++++++++++-- codex-rs/backend-client/src/types.rs | 1 + codex-rs/codex-api/src/rate_limits.rs | 202 ++++++++++++- codex-rs/codex-api/src/sse/responses.rs | 6 +- .../models/additional_rate_limit_details.rs | 38 +++ .../src/models/mod.rs | 3 + .../src/models/rate_limit_status_payload.rs | 8 + codex-rs/core/src/api_bridge.rs | 38 ++- codex-rs/core/src/codex.rs | 14 +- codex-rs/core/src/error.rs | 77 ++++- codex-rs/core/src/state/session.rs | 135 ++++++++- codex-rs/core/tests/suite/client.rs | 6 + .../core/tests/suite/client_websockets.rs | 2 + codex-rs/protocol/src/protocol.rs | 2 + codex-rs/tui/src/chatwidget.rs | 114 +++++--- codex-rs/tui/src/chatwidget/tests.rs | 103 ++++++- codex-rs/tui/src/status/card.rs | 44 ++- codex-rs/tui/src/status/mod.rs | 4 + codex-rs/tui/src/status/rate_limits.rs | 272 ++++++++++++++---- codex-rs/tui/src/status/tests.rs | 22 ++ 36 files changed, 1435 insertions(+), 169 deletions(-) create mode 100644 codex-rs/codex-backend-openapi-models/src/models/additional_rate_limit_details.rs diff --git a/codex-rs/app-server-protocol/schema/json/EventMsg.json b/codex-rs/app-server-protocol/schema/json/EventMsg.json index 7151f0d6b6..32ae8b715c 100644 --- a/codex-rs/app-server-protocol/schema/json/EventMsg.json +++ b/codex-rs/app-server-protocol/schema/json/EventMsg.json @@ -3470,6 +3470,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 39fba5b5ac..71bed38d7e 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -4406,6 +4406,18 @@ } ] }, + "limitId": { + "type": [ + "string", + "null" + ] + }, + "limitName": { + "type": [ + "string", + "null" + ] + }, "planType": { "anyOf": [ { @@ -4451,6 +4463,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 633c485e2a..dd63418f44 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 @@ -6555,6 +6555,18 @@ } ] }, + "limit_id": { + "type": [ + "string", + "null" + ] + }, + "limit_name": { + "type": [ + "string", + "null" + ] + }, "plan_type": { "anyOf": [ { @@ -11837,7 +11849,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": [ @@ -12852,6 +12879,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 027caee1ca..54b0219931 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json @@ -3470,6 +3470,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 e6f3b145a2..aaf063dacc 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json @@ -3470,6 +3470,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 1b717eefc4..07ab5710b9 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json @@ -3470,6 +3470,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 d168911bdf..91879645de 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 a7025fbaea..244156269d 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 9c2dad7f09..8604128b4e 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 fe970c1d42..e75e9b8283 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 f1a33f0b13..0c2ebe1893 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 679dc4ab23..bfc4b7a80f 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 4f0c526f7d..7d9b15db7b 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 61507d79d8..c44b5e83d5 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 b64bd5bce9..c9fd9ebda8 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 e4e670310a..6f66d7bff7 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 6fa36d1ffd..40bc276f01 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 9deeab7903..2559b2b39e 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 047c5934bd..909ab06a27 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 95bda35fb2..9e4d345664 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 0000000000..d89e656141 --- /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 7072dede5e..90bc058334 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 daf8e6af39..3d0492fa83 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 f7aeb570bd..af45182c13 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 b64eedbe4a..5b88d38d52 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -4280,7 +4280,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)); } @@ -5823,6 +5823,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), @@ -5839,6 +5841,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), @@ -5857,6 +5861,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, @@ -5906,6 +5912,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), @@ -5926,6 +5934,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), @@ -5940,6 +5950,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 889335824e..c284134d97 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 dcbf817d99..16614f3a6d 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 cb5668ec54..62d757e16e 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 18527ce713..e5ef5c17a4 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -486,6 +486,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 a56f55ba9a..cbf51fc651 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1251,6 +1251,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 18c4b9b162..29ab15cf0c 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(), @@ -4228,7 +4246,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, @@ -4236,7 +4259,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(), @@ -4382,8 +4405,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) @@ -4393,8 +4416,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) @@ -4578,9 +4601,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; } @@ -7146,18 +7170,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 6fdefb36c4..51cc2b61ea 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -96,6 +96,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; @@ -125,6 +126,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), @@ -1064,7 +1067,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(), @@ -1280,6 +1283,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 { @@ -1290,13 +1295,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), @@ -1308,8 +1315,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 @@ -1329,6 +1336,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), @@ -1345,6 +1354,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), @@ -1361,6 +1372,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), @@ -1377,6 +1390,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; @@ -1391,6 +1459,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 23005274db..60f7f93752 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 4f46af1cbc..90e406a52d 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 830fe158b2..8d46bc8c99 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,60 +160,123 @@ pub(crate) fn compose_rate_limit_data( now: DateTime, ) -> StatusRateLimitData { match snapshot { - Some(snapshot) => { - let mut rows = Vec::with_capacity(3); - - if let Some(primary) = snapshot.primary.as_ref() { - let label: String = primary - .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 - .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) - } else { - StatusRateLimitData::Available(rows) - } - } + Some(snapshot) => compose_rate_limit_data_many(std::slice::from_ref(snapshot), now), None => StatusRateLimitData::Missing, } } +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()) + }) + .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()) + }) + .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 { + 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(), + }, + }); + } + + 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) + } +} + /// Renders a fixed-width progress bar from remaining percentage. /// /// This function expects a remaining value in the `0..=100` range and clamps out-of-range input. @@ -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 c8844904ab..bff2fab048 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), From 22f98d9fb473e248d80799c84d9f07e484ecef57 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 10 Feb 2026 21:34:49 -0800 Subject: [PATCH 2/2] tui: show non-file layer content in /debug-config ## Why The debug output listed non-file-backed layers such as session flags and MDM managed config, but it did not show their values. That made it difficult to explain unexpected effective settings because users could not inspect those layers on disk. ## What - Extend layer rendering to include details for non-file-backed layers. - For session flags, render an indented bulleted list of flattened key = value pairs. - For MDM-backed managed config layers, render an indented mdm value block. - Keep file-backed layer output unchanged. - Add focused tests for session-flags and legacy MDM rendering. ## Validation - just fmt - cargo test -p codex-tui --- codex-rs/core/src/config_loader/layer_io.rs | 23 +++- codex-rs/core/src/config_loader/macos.rs | 20 ++- codex-rs/core/src/config_loader/mod.rs | 12 +- codex-rs/core/src/config_loader/state.rs | 18 +++ codex-rs/core/src/config_loader/tests.rs | 31 +++-- codex-rs/tui/src/debug_config.rs | 135 ++++++++++++++++++++ 6 files changed, 221 insertions(+), 18 deletions(-) diff --git a/codex-rs/core/src/config_loader/layer_io.rs b/codex-rs/core/src/config_loader/layer_io.rs index f6f7d2e37d..a97f285549 100644 --- a/codex-rs/core/src/config_loader/layer_io.rs +++ b/codex-rs/core/src/config_loader/layer_io.rs @@ -2,6 +2,8 @@ use super::LoaderOverrides; use super::diagnostics::config_error_from_toml; use super::diagnostics::io_error_from_config_error; #[cfg(target_os = "macos")] +use super::macos::ManagedAdminConfigLayer; +#[cfg(target_os = "macos")] use super::macos::load_managed_admin_config_layer; use codex_utils_absolute_path::AbsolutePathBuf; use std::io; @@ -19,12 +21,18 @@ pub(super) struct MangedConfigFromFile { pub file: AbsolutePathBuf, } +#[derive(Debug, Clone)] +pub(super) struct ManagedConfigFromMdm { + pub managed_config: TomlValue, + pub raw_toml: String, +} + #[derive(Debug, Clone)] pub(super) struct LoadedConfigLayers { /// If present, data read from a file such as `/etc/codex/managed_config.toml`. pub managed_config: Option, /// If present, data read from managed preferences (macOS only). - pub managed_config_from_mdm: Option, + pub managed_config_from_mdm: Option, } pub(super) async fn load_config_layers_internal( @@ -57,7 +65,9 @@ pub(super) async fn load_config_layers_internal( #[cfg(target_os = "macos")] let managed_preferences = - load_managed_admin_config_layer(managed_preferences_base64.as_deref()).await?; + load_managed_admin_config_layer(managed_preferences_base64.as_deref()) + .await? + .map(map_managed_admin_layer); #[cfg(not(target_os = "macos"))] let managed_preferences = None; @@ -68,6 +78,15 @@ pub(super) async fn load_config_layers_internal( }) } +#[cfg(target_os = "macos")] +fn map_managed_admin_layer(layer: ManagedAdminConfigLayer) -> ManagedConfigFromMdm { + let ManagedAdminConfigLayer { config, raw_toml } = layer; + ManagedConfigFromMdm { + managed_config: config, + raw_toml, + } +} + pub(super) async fn read_config_from_path( path: impl AsRef, log_missing_as_info: bool, diff --git a/codex-rs/core/src/config_loader/macos.rs b/codex-rs/core/src/config_loader/macos.rs index e3bbcc07e2..b474e81eaa 100644 --- a/codex-rs/core/src/config_loader/macos.rs +++ b/codex-rs/core/src/config_loader/macos.rs @@ -15,6 +15,12 @@ const MANAGED_PREFERENCES_APPLICATION_ID: &str = "com.openai.codex"; const MANAGED_PREFERENCES_CONFIG_KEY: &str = "config_toml_base64"; const MANAGED_PREFERENCES_REQUIREMENTS_KEY: &str = "requirements_toml_base64"; +#[derive(Debug, Clone)] +pub(super) struct ManagedAdminConfigLayer { + pub config: TomlValue, + pub raw_toml: String, +} + pub(super) fn managed_preferences_requirements_source() -> RequirementSource { RequirementSource::MdmManagedPreferences { domain: MANAGED_PREFERENCES_APPLICATION_ID.to_string(), @@ -24,7 +30,7 @@ pub(super) fn managed_preferences_requirements_source() -> RequirementSource { pub(crate) async fn load_managed_admin_config_layer( override_base64: Option<&str>, -) -> io::Result> { +) -> io::Result> { if let Some(encoded) = override_base64 { let trimmed = encoded.trim(); return if trimmed.is_empty() { @@ -47,7 +53,7 @@ pub(crate) async fn load_managed_admin_config_layer( } } -fn load_managed_admin_config() -> io::Result> { +fn load_managed_admin_config() -> io::Result> { load_managed_preference(MANAGED_PREFERENCES_CONFIG_KEY)? .as_deref() .map(str::trim) @@ -122,9 +128,13 @@ fn load_managed_preference(key_name: &str) -> io::Result> { Ok(Some(value)) } -fn parse_managed_config_base64(encoded: &str) -> io::Result { - match toml::from_str::(&decode_managed_preferences_base64(encoded)?) { - Ok(TomlValue::Table(parsed)) => Ok(TomlValue::Table(parsed)), +fn parse_managed_config_base64(encoded: &str) -> io::Result { + let raw_toml = decode_managed_preferences_base64(encoded)?; + match toml::from_str::(&raw_toml) { + Ok(TomlValue::Table(parsed)) => Ok(ManagedAdminConfigLayer { + config: TomlValue::Table(parsed), + raw_toml, + }), Ok(other) => { tracing::error!("Managed config TOML must have a table at the root, found {other:?}",); Err(io::Error::new( diff --git a/codex-rs/core/src/config_loader/mod.rs b/codex-rs/core/src/config_loader/mod.rs index a38baf7d47..68973718fe 100644 --- a/codex-rs/core/src/config_loader/mod.rs +++ b/codex-rs/core/src/config_loader/mod.rs @@ -278,9 +278,10 @@ pub async fn load_config_layers_state( )); } if let Some(config) = managed_config_from_mdm { - layers.push(ConfigLayerEntry::new( + layers.push(ConfigLayerEntry::new_with_raw_toml( ConfigLayerSource::LegacyManagedConfigTomlFromMdm, - config, + config.managed_config, + config.raw_toml, )); } @@ -485,7 +486,12 @@ async fn load_requirements_from_legacy_scheme( } = loaded_config_layers; for (source, config) in managed_config_from_mdm - .map(|config| (RequirementSource::LegacyManagedConfigTomlFromMdm, config)) + .map(|config| { + ( + RequirementSource::LegacyManagedConfigTomlFromMdm, + config.managed_config, + ) + }) .into_iter() .chain(managed_config.map(|c| { ( diff --git a/codex-rs/core/src/config_loader/state.rs b/codex-rs/core/src/config_loader/state.rs index 847b19d7e5..30afecc7ca 100644 --- a/codex-rs/core/src/config_loader/state.rs +++ b/codex-rs/core/src/config_loader/state.rs @@ -27,6 +27,7 @@ pub struct LoaderOverrides { pub struct ConfigLayerEntry { pub name: ConfigLayerSource, pub config: TomlValue, + pub raw_toml: Option, pub version: String, pub disabled_reason: Option, } @@ -37,6 +38,18 @@ impl ConfigLayerEntry { Self { name, config, + raw_toml: None, + version, + disabled_reason: None, + } + } + + pub fn new_with_raw_toml(name: ConfigLayerSource, config: TomlValue, raw_toml: String) -> Self { + let version = version_for_toml(&config); + Self { + name, + config, + raw_toml: Some(raw_toml), version, disabled_reason: None, } @@ -51,6 +64,7 @@ impl ConfigLayerEntry { Self { name, config, + raw_toml: None, version, disabled_reason: Some(disabled_reason.into()), } @@ -60,6 +74,10 @@ impl ConfigLayerEntry { self.disabled_reason.is_some() } + pub fn raw_toml(&self) -> Option<&str> { + self.raw_toml.as_deref() + } + pub fn metadata(&self) -> ConfigLayerMetadata { ConfigLayerMetadata { name: self.name.clone(), diff --git a/codex-rs/core/src/config_loader/tests.rs b/codex-rs/core/src/config_loader/tests.rs index 0a6f67d85b..79c657ad96 100644 --- a/codex-rs/core/src/config_loader/tests.rs +++ b/codex-rs/core/src/config_loader/tests.rs @@ -264,6 +264,7 @@ async fn returns_empty_when_all_layers_missing() { .expect("resolve user config.toml path") }, config: TomlValue::Table(toml::map::Map::new()), + raw_toml: None, version: version_for_toml(&TomlValue::Table(toml::map::Map::new())), disabled_reason: None, }, @@ -325,18 +326,17 @@ flag = true "#, ) .expect("write managed config"); + let raw_managed_preferences = r#" +# managed profile +[nested] +value = "managed" +flag = false +"#; let overrides = LoaderOverrides { managed_config_path: Some(managed_path), managed_preferences_base64: Some( - base64::prelude::BASE64_STANDARD.encode( - r#" -[nested] -value = "managed" -flag = false -"# - .as_bytes(), - ), + base64::prelude::BASE64_STANDARD.encode(raw_managed_preferences.as_bytes()), ), macos_managed_config_requirements_base64: None, }; @@ -361,6 +361,19 @@ flag = false Some(&TomlValue::String("managed".to_string())) ); assert_eq!(nested.get("flag"), Some(&TomlValue::Boolean(false))); + let mdm_layer = state + .layers_high_to_low() + .into_iter() + .find(|layer| { + matches!( + layer.name, + super::ConfigLayerSource::LegacyManagedConfigTomlFromMdm + ) + }) + .expect("mdm layer"); + let raw = mdm_layer.raw_toml().expect("preserved mdm toml"); + assert!(raw.contains("# managed profile")); + assert!(raw.contains("value = \"managed\"")); } #[cfg(target_os = "macos")] @@ -862,6 +875,7 @@ async fn project_layer_is_added_when_dot_codex_exists_without_config_toml() -> s dot_codex_folder: AbsolutePathBuf::from_absolute_path(project_root.join(".codex"))?, }, config: TomlValue::Table(toml::map::Map::new()), + raw_toml: None, version: version_for_toml(&TomlValue::Table(toml::map::Map::new())), disabled_reason: None, }], @@ -955,6 +969,7 @@ async fn codex_home_within_project_tree_is_not_double_loaded() -> std::io::Resul dot_codex_folder: AbsolutePathBuf::from_absolute_path(&nested_dot_codex)?, }, config: child_config.clone(), + raw_toml: None, version: version_for_toml(&child_config), disabled_reason: None, }], diff --git a/codex-rs/tui/src/debug_config.rs b/codex-rs/tui/src/debug_config.rs index 681d9c3a9a..2b04670ed8 100644 --- a/codex-rs/tui/src/debug_config.rs +++ b/codex-rs/tui/src/debug_config.rs @@ -1,6 +1,7 @@ use crate::history_cell::PlainHistoryCell; use codex_app_server_protocol::ConfigLayerSource; use codex_core::config::Config; +use codex_core::config_loader::ConfigLayerEntry; use codex_core::config_loader::ConfigLayerStack; use codex_core::config_loader::ConfigLayerStackOrdering; use codex_core::config_loader::NetworkConstraints; @@ -11,6 +12,7 @@ use codex_core::config_loader::WebSearchModeRequirement; use codex_core::protocol::SessionNetworkProxyRuntime; use ratatui::style::Stylize; use ratatui::text::Line; +use toml::Value as TomlValue; pub(crate) fn new_debug_config_output( config: &Config, @@ -71,6 +73,7 @@ fn render_debug_config_lines(stack: &ConfigLayerStack) -> Vec> { "enabled" }; lines.push(format!(" {}. {source} ({status})", index + 1).into()); + lines.extend(render_non_file_layer_details(layer)); if let Some(reason) = &layer.disabled_reason { lines.push(format!(" reason: {reason}").dim().into()); } @@ -169,6 +172,80 @@ fn render_debug_config_lines(stack: &ConfigLayerStack) -> Vec> { lines } +fn render_non_file_layer_details(layer: &ConfigLayerEntry) -> Vec> { + match &layer.name { + ConfigLayerSource::SessionFlags => render_session_flag_details(&layer.config), + ConfigLayerSource::Mdm { .. } | ConfigLayerSource::LegacyManagedConfigTomlFromMdm => { + render_mdm_layer_details(layer) + } + ConfigLayerSource::System { .. } + | ConfigLayerSource::User { .. } + | ConfigLayerSource::Project { .. } + | ConfigLayerSource::LegacyManagedConfigTomlFromFile { .. } => Vec::new(), + } +} + +fn render_session_flag_details(config: &TomlValue) -> Vec> { + let mut pairs = Vec::new(); + flatten_toml_key_values(config, None, &mut pairs); + + if pairs.is_empty() { + return vec![" - ".dim().into()]; + } + + pairs + .into_iter() + .map(|(key, value)| format!(" - {key} = {value}").into()) + .collect() +} + +fn render_mdm_layer_details(layer: &ConfigLayerEntry) -> Vec> { + let value = layer + .raw_toml() + .map(ToString::to_string) + .unwrap_or_else(|| format_toml_value(&layer.config)); + if value.is_empty() { + return vec![" mdm value: ".dim().into()]; + } + + if value.contains('\n') { + let mut lines = vec![" mdm value:".into()]; + lines.extend(value.lines().map(|line| format!(" {line}").into())); + lines + } else { + vec![format!(" mdm value: {value}").into()] + } +} + +fn flatten_toml_key_values( + value: &TomlValue, + prefix: Option<&str>, + out: &mut Vec<(String, String)>, +) { + match value { + TomlValue::Table(table) => { + let mut entries = table.iter().collect::>(); + entries.sort_by_key(|(key, _)| key.as_str()); + for (key, child) in entries { + let next_prefix = if let Some(prefix) = prefix { + format!("{prefix}.{key}") + } else { + key.to_string() + }; + flatten_toml_key_values(child, Some(&next_prefix), out); + } + } + _ => { + let key = prefix.unwrap_or("").to_string(); + out.push((key, format_toml_value(value))); + } + } +} + +fn format_toml_value(value: &TomlValue) -> String { + value.to_string() +} + fn requirement_line( name: &str, value: String, @@ -494,6 +571,64 @@ mod tests { assert!(!rendered.contains(" - rules:")); } + #[test] + fn debug_config_output_lists_session_flag_key_value_pairs() { + let session_flags = toml::from_str::( + r#" +model = "gpt-5" +[sandbox_workspace_write] +network_access = true +writable_roots = ["/tmp"] +"#, + ) + .expect("session flags"); + + let stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::SessionFlags, + session_flags, + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let rendered = render_to_text(&render_debug_config_lines(&stack)); + assert!(rendered.contains("session-flags (enabled)")); + assert!(rendered.contains(" - model = \"gpt-5\"")); + assert!(rendered.contains(" - sandbox_workspace_write.network_access = true")); + assert!(rendered.contains("sandbox_workspace_write.writable_roots")); + assert!(rendered.contains("/tmp")); + } + + #[test] + fn debug_config_output_shows_legacy_mdm_layer_value() { + let raw_mdm_toml = r#" +# managed by MDM +model = "managed_model" +approval_policy = "never" +"#; + let mdm_value = toml::from_str::(raw_mdm_toml).expect("mdm value"); + + let stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new_with_raw_toml( + ConfigLayerSource::LegacyManagedConfigTomlFromMdm, + mdm_value, + raw_mdm_toml.to_string(), + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let rendered = render_to_text(&render_debug_config_lines(&stack)); + assert!(rendered.contains("legacy managed_config.toml (mdm) (enabled)")); + assert!(rendered.contains("mdm value:")); + assert!(rendered.contains("# managed by MDM")); + assert!(rendered.contains("model = \"managed_model\"")); + assert!(rendered.contains("approval_policy = \"never\"")); + } + #[test] fn debug_config_output_normalizes_empty_web_search_mode_list() { let mut requirements = ConfigRequirements::default();