From c7d8c38d2f8b1172b9e72d07eadaf16b80f35b7d Mon Sep 17 00:00:00 2001 From: kevin zhao Date: Mon, 17 Nov 2025 15:04:52 -0800 Subject: [PATCH] plan type reconciliation --- .../app-server/src/codex_message_processor.rs | 1 + codex-rs/backend-client/src/client.rs | 29 +++++++++- codex-rs/backend-client/src/lib.rs | 1 + codex-rs/tui/src/app.rs | 7 ++- codex-rs/tui/src/app_event.rs | 4 +- codex-rs/tui/src/chatwidget.rs | 55 ++++++++++++++++++- 6 files changed, 88 insertions(+), 9 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c694486187..16ae612f2c 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -967,6 +967,7 @@ impl CodexMessageProcessor { client .get_rate_limits() .await + .map(|status| status.snapshot) .map_err(|err| JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to fetch codex rate limits: {err}"), diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 28a51598eb..a1a22d7dd0 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -1,5 +1,6 @@ use crate::types::CodeTaskDetailsResponse; use crate::types::PaginatedListTaskListItem; +use crate::types::PlanType; use crate::types::RateLimitStatusPayload; use crate::types::RateLimitWindowSnapshot; use crate::types::TurnAttemptsSiblingTurnsResponse; @@ -44,6 +45,12 @@ pub struct Client { path_style: PathStyle, } +#[derive(Clone, Debug)] +pub struct RateLimitStatus { + pub plan_type: String, + pub snapshot: RateLimitSnapshot, +} + impl Client { pub fn new(base_url: impl Into) -> Result { let mut base_url = base_url.into(); @@ -155,7 +162,7 @@ impl Client { } } - pub async fn get_rate_limits(&self) -> Result { + pub async fn get_rate_limits(&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), @@ -163,7 +170,9 @@ 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)) + let plan_type = Self::plan_type_to_string(payload.plan_type); + let snapshot = Self::rate_limit_snapshot_from_payload(payload); + Ok(RateLimitStatus { plan_type, snapshot }) } pub async fn list_tasks( @@ -288,6 +297,22 @@ impl Client { } } + fn plan_type_to_string(plan_type: PlanType) -> String { + match plan_type { + PlanType::Free => "free", + PlanType::Go => "go", + PlanType::Plus => "plus", + PlanType::Pro => "pro", + PlanType::Team => "team", + PlanType::Business => "business", + PlanType::Education => "education", + PlanType::Quorum => "quorum", + PlanType::Enterprise => "enterprise", + PlanType::Edu => "edu", + } + .to_string() + } + fn map_rate_limit_window( window: Option>>, ) -> Option { diff --git a/codex-rs/backend-client/src/lib.rs b/codex-rs/backend-client/src/lib.rs index 29fe9f3c6b..3824eefbaa 100644 --- a/codex-rs/backend-client/src/lib.rs +++ b/codex-rs/backend-client/src/lib.rs @@ -2,6 +2,7 @@ mod client; pub mod types; pub use client::Client; +pub use client::RateLimitStatus; pub use types::CodeTaskDetailsResponse; pub use types::CodeTaskDetailsResponseExt; pub use types::PaginatedListTaskListItem; diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 9f996007f8..a1bd66819a 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -496,8 +496,11 @@ impl App { AppEvent::FileSearchResult { query, matches } => { self.chat_widget.apply_file_search_result(query, matches); } - AppEvent::RateLimitSnapshotFetched(snapshot) => { - self.chat_widget.on_rate_limit_snapshot(Some(snapshot)); + AppEvent::RateLimitSnapshotFetched(status) => { + self.chat_widget + .on_rate_limit_snapshot(Some(status.snapshot.clone())); + self.chat_widget + .reconcile_plan_type_from_usage(status.plan_type); } AppEvent::UpdateReasoningEffort(effort) => { self.on_update_reasoning_effort(effort); diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 0c3033c5c9..d00f676d67 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -1,10 +1,10 @@ use std::path::PathBuf; +use codex_backend_client::RateLimitStatus; use codex_common::approval_presets::ApprovalPreset; use codex_common::model_presets::ModelPreset; use codex_core::protocol::ConversationPathResponseEvent; use codex_core::protocol::Event; -use codex_core::protocol::RateLimitSnapshot; use codex_file_search::FileMatch; use crate::bottom_pane::ApprovalRequest; @@ -43,7 +43,7 @@ pub(crate) enum AppEvent { }, /// Result of refreshing rate limits - RateLimitSnapshotFetched(RateLimitSnapshot), + RateLimitSnapshotFetched(RateLimitStatus), /// Result of computing a `/diff` command. DiffResult(String), diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 83fea1e90e..1cfef2d81a 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -6,6 +6,7 @@ use std::time::Duration; use codex_app_server_protocol::AuthMode; use codex_backend_client::Client as BackendClient; +use codex_backend_client::RateLimitStatus; use codex_core::config::Config; use codex_core::config::types::Notifications; use codex_core::git_info::current_branch_name; @@ -261,6 +262,7 @@ pub(crate) struct ChatWidget { rate_limit_warnings: RateLimitWarningState, rate_limit_switch_prompt: RateLimitSwitchPromptState, rate_limit_poller: Option>, + last_plan_refresh_attempt: Option, // Stream lifecycle controller stream_controller: Option, running_commands: HashMap, @@ -553,6 +555,51 @@ impl ChatWidget { self.rate_limit_snapshot = None; } } + + pub(crate) fn reconcile_plan_type_from_usage(&mut self, backend_plan: String) { + let trimmed = backend_plan.trim(); + if trimmed.is_empty() { + return; + } + + let plan_lower = trimmed.to_ascii_lowercase(); + let Some(auth) = self.auth_manager.auth() else { + self.last_plan_refresh_attempt = None; + return; + }; + if auth.mode != AuthMode::ChatGPT { + self.last_plan_refresh_attempt = None; + return; + } + + let current_plan = auth + .raw_plan_type() + .map(|plan| plan.to_ascii_lowercase()); + + if current_plan + .as_deref() + .is_some_and(|plan| plan == plan_lower) + { + self.last_plan_refresh_attempt = None; + return; + } + + if self + .last_plan_refresh_attempt + .as_deref() + .is_some_and(|attempt| attempt == plan_lower) + { + return; + } + + self.last_plan_refresh_attempt = Some(plan_lower); + let auth_manager = self.auth_manager.clone(); + tokio::spawn(async move { + if let Err(err) = auth_manager.refresh_token().await { + debug!(?err, "failed to refresh auth after plan mismatch"); + } + }); + } /// Finalize any active exec as failed and stop/clear running UI state. fn finalize_turn(&mut self) { // Ensure any spinner is replaced by a red ✗ and flushed into history. @@ -1065,6 +1112,7 @@ impl ChatWidget { rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), rate_limit_poller: None, + last_plan_refresh_attempt: None, stream_controller: None, running_commands: HashMap::new(), task_complete_pending: false, @@ -1137,6 +1185,7 @@ impl ChatWidget { rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), rate_limit_poller: None, + last_plan_refresh_attempt: None, stream_controller: None, running_commands: HashMap::new(), task_complete_pending: false, @@ -1776,8 +1825,8 @@ impl ChatWidget { let mut interval = tokio::time::interval(Duration::from_secs(60)); loop { - if let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth.clone()).await { - app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); + if let Some(status) = fetch_rate_limits(base_url.clone(), auth.clone()).await { + app_event_tx.send(AppEvent::RateLimitSnapshotFetched(status)); } interval.tick().await; } @@ -2946,7 +2995,7 @@ 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) -> Option { match BackendClient::from_auth(base_url, &auth).await { Ok(client) => match client.get_rate_limits().await { Ok(snapshot) => Some(snapshot),