fix: taking plan type from usage endpoint instead of thru auth token (#7610)

pull plan type from the usage endpoint, persist it in session state /
tui state, and propagate through rate limit snapshots

codex-rs/tui/src/chatwidget.rs

@@ -58,6 +58,7 @@ use codex_core::protocol::WebSearchBeginEvent;
 use codex_core::protocol::WebSearchEndEvent;
 use codex_core::skills::model::SkillMetadata;
 use codex_protocol::ConversationId;
+use codex_protocol::account::PlanType;
 use codex_protocol::approvals::ElicitationRequestEvent;
 use codex_protocol::parse_command::ParsedCommand;
 use codex_protocol::user_input::UserInput;
@@ -282,6 +283,7 @@ pub(crate) struct ChatWidget {
     initial_user_message: Option<UserMessage>,
     token_info: Option<TokenUsageInfo>,
     rate_limit_snapshot: Option<RateLimitSnapshotDisplay>,
+    plan_type: Option<PlanType>,
     rate_limit_warnings: RateLimitWarningState,
     rate_limit_switch_prompt: RateLimitSwitchPromptState,
     rate_limit_poller: Option<JoinHandle<()>>,
@@ -580,6 +582,8 @@ impl ChatWidget {
                     });
             }
 
+            self.plan_type = snapshot.plan_type.or(self.plan_type);
+
             let warnings = self.rate_limit_warnings.take_warnings(
                 snapshot
                     .secondary
@@ -1275,6 +1279,7 @@ impl ChatWidget {
             ),
             token_info: None,
             rate_limit_snapshot: None,
+            plan_type: None,
             rate_limit_warnings: RateLimitWarningState::default(),
             rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
             rate_limit_poller: None,
@@ -1357,6 +1362,7 @@ impl ChatWidget {
             ),
             token_info: None,
             rate_limit_snapshot: None,
+            plan_type: None,
             rate_limit_warnings: RateLimitWarningState::default(),
             rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
             rate_limit_poller: None,
@@ -2001,6 +2007,7 @@ impl ChatWidget {
             context_usage,
             &self.conversation_id,
             self.rate_limit_snapshot.as_ref(),
+            self.plan_type,
             Local::now(),
         ));
     }

codex-rs/tui/src/chatwidget/tests.rs

@@ -47,6 +47,7 @@ use codex_core::protocol::UndoStartedEvent;
 use codex_core::protocol::ViewImageToolCallEvent;
 use codex_core::protocol::WarningEvent;
 use codex_protocol::ConversationId;
+use codex_protocol::account::PlanType;
 use codex_protocol::openai_models::ModelPreset;
 use codex_protocol::openai_models::ReasoningEffortPreset;
 use codex_protocol::parse_command::ParsedCommand;
@@ -90,6 +91,7 @@ fn snapshot(percent: f64) -> RateLimitSnapshot {
         }),
         secondary: None,
         credits: None,
+        plan_type: None,
     }
 }
 
@@ -398,6 +400,7 @@ fn make_chatwidget_manual() -> (
         initial_user_message: None,
         token_info: None,
         rate_limit_snapshot: None,
+        plan_type: None,
         rate_limit_warnings: RateLimitWarningState::default(),
         rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
         rate_limit_poller: None,
@@ -546,6 +549,7 @@ fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
             unlimited: false,
             balance: Some("17.5".to_string()),
         }),
+        plan_type: None,
     }));
     let initial_balance = chat
         .rate_limit_snapshot
@@ -562,6 +566,7 @@ fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
         }),
         secondary: None,
         credits: None,
+        plan_type: None,
     }));
 
     let display = chat
@@ -581,6 +586,59 @@ fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
     );
 }
 
+#[test]
+fn rate_limit_snapshot_updates_and_retains_plan_type() {
+    let (mut chat, _rx, _op_rx) = make_chatwidget_manual();
+
+    chat.on_rate_limit_snapshot(Some(RateLimitSnapshot {
+        primary: Some(RateLimitWindow {
+            used_percent: 10.0,
+            window_minutes: Some(60),
+            resets_at: None,
+        }),
+        secondary: Some(RateLimitWindow {
+            used_percent: 5.0,
+            window_minutes: Some(300),
+            resets_at: None,
+        }),
+        credits: None,
+        plan_type: Some(PlanType::Plus),
+    }));
+    assert_eq!(chat.plan_type, Some(PlanType::Plus));
+
+    chat.on_rate_limit_snapshot(Some(RateLimitSnapshot {
+        primary: Some(RateLimitWindow {
+            used_percent: 25.0,
+            window_minutes: Some(30),
+            resets_at: Some(123),
+        }),
+        secondary: Some(RateLimitWindow {
+            used_percent: 15.0,
+            window_minutes: Some(300),
+            resets_at: Some(234),
+        }),
+        credits: None,
+        plan_type: Some(PlanType::Pro),
+    }));
+    assert_eq!(chat.plan_type, Some(PlanType::Pro));
+
+    chat.on_rate_limit_snapshot(Some(RateLimitSnapshot {
+        primary: Some(RateLimitWindow {
+            used_percent: 30.0,
+            window_minutes: Some(60),
+            resets_at: Some(456),
+        }),
+        secondary: Some(RateLimitWindow {
+            used_percent: 18.0,
+            window_minutes: Some(300),
+            resets_at: Some(567),
+        }),
+        credits: None,
+        plan_type: None,
+    }));
+    assert_eq!(chat.plan_type, Some(PlanType::Pro));
+}
+
 #[test]
 fn rate_limit_switch_prompt_skips_when_on_lower_cost_model() {
     let (mut chat, _, _) = make_chatwidget_manual();

codex-rs/tui/src/status/card.rs

@@ -10,6 +10,7 @@ use codex_core::config::Config;
 use codex_core::protocol::SandboxPolicy;
 use codex_core::protocol::TokenUsage;
 use codex_protocol::ConversationId;
+use codex_protocol::account::PlanType;
 use ratatui::prelude::*;
 use ratatui::style::Stylize;
 use std::collections::BTreeSet;
@@ -65,6 +66,7 @@ struct StatusHistoryCell {
     rate_limits: StatusRateLimitData,
 }
 
+#[allow(clippy::too_many_arguments)]
 pub(crate) fn new_status_output(
     config: &Config,
     auth_manager: &AuthManager,
@@ -72,6 +74,7 @@ pub(crate) fn new_status_output(
     context_usage: Option<&TokenUsage>,
     session_id: &Option<ConversationId>,
     rate_limits: Option<&RateLimitSnapshotDisplay>,
+    plan_type: Option<PlanType>,
     now: DateTime<Local>,
 ) -> CompositeHistoryCell {
     let command = PlainHistoryCell::new(vec!["/status".magenta().into()]);
@@ -82,6 +85,7 @@ pub(crate) fn new_status_output(
         context_usage,
         session_id,
         rate_limits,
+        plan_type,
         now,
     );
 
@@ -89,6 +93,7 @@ pub(crate) fn new_status_output(
 }
 
 impl StatusHistoryCell {
+    #[allow(clippy::too_many_arguments)]
     fn new(
         config: &Config,
         auth_manager: &AuthManager,
@@ -96,6 +101,7 @@ impl StatusHistoryCell {
         context_usage: Option<&TokenUsage>,
         session_id: &Option<ConversationId>,
         rate_limits: Option<&RateLimitSnapshotDisplay>,
+        plan_type: Option<PlanType>,
         now: DateTime<Local>,
     ) -> Self {
         let config_entries = create_config_summary_entries(config);
@@ -111,7 +117,7 @@ impl StatusHistoryCell {
             SandboxPolicy::WorkspaceWrite { .. } => "workspace-write".to_string(),
         };
         let agents_summary = compose_agents_summary(config);
-        let account = compose_account_display(auth_manager);
+        let account = compose_account_display(auth_manager, plan_type);
         let session_id = session_id.as_ref().map(std::string::ToString::to_string);
         let context_window = config.model_context_window.and_then(|window| {
             context_usage.map(|usage| StatusContextWindowData {

codex-rs/tui/src/status/helpers.rs

@@ -6,6 +6,7 @@ use codex_app_server_protocol::AuthMode;
 use codex_core::AuthManager;
 use codex_core::config::Config;
 use codex_core::project_doc::discover_project_doc_paths;
+use codex_protocol::account::PlanType;
 use std::path::Path;
 use unicode_width::UnicodeWidthStr;
 
@@ -83,13 +84,18 @@ pub(crate) fn compose_agents_summary(config: &Config) -> String {
     }
 }
 
-pub(crate) fn compose_account_display(auth_manager: &AuthManager) -> Option<StatusAccountDisplay> {
+pub(crate) fn compose_account_display(
+    auth_manager: &AuthManager,
+    plan: Option<PlanType>,
+) -> Option<StatusAccountDisplay> {
     let auth = auth_manager.auth()?;
 
     match auth.mode {
         AuthMode::ChatGPT => {
             let email = auth.get_account_email();
-            let plan = auth.raw_plan_type().map(|plan| title_case(plan.as_str()));
+            let plan = plan
+                .map(|plan_type| title_case(format!("{plan_type:?}").as_str()))
+                .or_else(|| Some("Unknown".to_string()));
             Some(StatusAccountDisplay::ChatGpt { email, plan })
         }
         AuthMode::ApiKey => Some(StatusAccountDisplay::ApiKey),

codex-rs/tui/src/status/tests.rs

@@ -120,6 +120,7 @@ fn status_snapshot_includes_reasoning_details() {
             resets_at: Some(reset_at_from(&captured_at, 1_200)),
         }),
         credits: None,
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
 
@@ -130,6 +131,7 @@ fn status_snapshot_includes_reasoning_details() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -171,6 +173,7 @@ fn status_snapshot_includes_monthly_limit() {
         }),
         secondary: None,
         credits: None,
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
 
@@ -181,6 +184,7 @@ fn status_snapshot_includes_monthly_limit() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -211,6 +215,7 @@ fn status_snapshot_shows_unlimited_credits() {
             unlimited: true,
             balance: None,
         }),
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
     let composite = new_status_output(
@@ -220,6 +225,7 @@ fn status_snapshot_shows_unlimited_credits() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let rendered = render_lines(&composite.display_lines(120));
@@ -249,6 +255,7 @@ fn status_snapshot_shows_positive_credits() {
             unlimited: false,
             balance: Some("12.5".to_string()),
         }),
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
     let composite = new_status_output(
@@ -258,6 +265,7 @@ fn status_snapshot_shows_positive_credits() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let rendered = render_lines(&composite.display_lines(120));
@@ -287,6 +295,7 @@ fn status_snapshot_hides_zero_credits() {
             unlimited: false,
             balance: Some("0".to_string()),
         }),
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
     let composite = new_status_output(
@@ -296,6 +305,7 @@ fn status_snapshot_hides_zero_credits() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let rendered = render_lines(&composite.display_lines(120));
@@ -323,6 +333,7 @@ fn status_snapshot_hides_when_has_no_credits_flag() {
             unlimited: true,
             balance: None,
         }),
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
     let composite = new_status_output(
@@ -332,6 +343,7 @@ fn status_snapshot_hides_when_has_no_credits_flag() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let rendered = render_lines(&composite.display_lines(120));
@@ -369,6 +381,7 @@ fn status_card_token_usage_excludes_cached_tokens() {
         Some(&usage),
         &None,
         None,
+        None,
         now,
     );
     let rendered = render_lines(&composite.display_lines(120));
@@ -410,6 +423,7 @@ fn status_snapshot_truncates_in_narrow_terminal() {
         }),
         secondary: None,
         credits: None,
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
 
@@ -420,6 +434,7 @@ fn status_snapshot_truncates_in_narrow_terminal() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(70));
@@ -461,6 +476,7 @@ fn status_snapshot_shows_missing_limits_message() {
         Some(&usage),
         &None,
         None,
+        None,
         now,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -509,6 +525,7 @@ fn status_snapshot_includes_credits_and_limits() {
             unlimited: false,
             balance: Some("37.5".to_string()),
         }),
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
 
@@ -519,6 +536,7 @@ fn status_snapshot_includes_credits_and_limits() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -551,6 +569,7 @@ fn status_snapshot_shows_empty_limits_message() {
         primary: None,
         secondary: None,
         credits: None,
+        plan_type: None,
     };
     let captured_at = chrono::Local
         .with_ymd_and_hms(2024, 6, 7, 8, 9, 10)
@@ -565,6 +584,7 @@ fn status_snapshot_shows_empty_limits_message() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         captured_at,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -609,6 +629,7 @@ fn status_snapshot_shows_stale_limits_message() {
             resets_at: Some(reset_at_from(&captured_at, 1_800)),
         }),
         credits: None,
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
     let now = captured_at + ChronoDuration::minutes(20);
@@ -620,6 +641,7 @@ fn status_snapshot_shows_stale_limits_message() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         now,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -668,6 +690,7 @@ fn status_snapshot_cached_limits_hide_credits_without_flag() {
             unlimited: false,
             balance: Some("80".to_string()),
         }),
+        plan_type: None,
     };
     let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
     let now = captured_at + ChronoDuration::minutes(20);
@@ -679,6 +702,7 @@ fn status_snapshot_cached_limits_hide_credits_without_flag() {
         Some(&usage),
         &None,
         Some(&rate_display),
+        None,
         now,
     );
     let mut rendered_lines = render_lines(&composite.display_lines(80));
@@ -725,6 +749,7 @@ fn status_context_window_uses_last_usage() {
         Some(&last_usage),
         &None,
         None,
+        None,
         now,
     );
     let rendered_lines = render_lines(&composite.display_lines(80));