Use PlanType enum when formatting usage-limit CTA (#3495)

- Started using Play type struct
- Added CTA for team/business 
- Refactored a bit to unify the logic

codex-rs/core/src/auth.rs

@@ -17,6 +17,7 @@ use std::time::Duration;
 
 use codex_protocol::mcp_protocol::AuthMode;
 
+use crate::token_data::PlanType;
 use crate::token_data::TokenData;
 use crate::token_data::parse_id_token;
 
@@ -134,9 +135,9 @@ impl CodexAuth {
         self.get_current_token_data().and_then(|t| t.account_id)
     }
 
-    pub fn get_plan_type(&self) -> Option<String> {
+    pub(crate) fn get_plan_type(&self) -> Option<PlanType> {
         self.get_current_token_data()
-            .and_then(|t| t.id_token.chatgpt_plan_type.as_ref().map(|p| p.as_string()))
+            .and_then(|t| t.id_token.chatgpt_plan_type)
     }
 
     fn get_current_auth_json(&self) -> Option<AuthDotJson> {

codex-rs/core/src/client.rs

@@ -41,6 +41,7 @@ use crate::model_provider_info::WireApi;
 use crate::openai_model_info::get_model_info;
 use crate::openai_tools::create_tools_json_for_responses_api;
 use crate::protocol::TokenUsage;
+use crate::token_data::PlanType;
 use crate::util::backoff;
 use codex_protocol::config_types::ReasoningEffort as ReasoningEffortConfig;
 use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
@@ -60,7 +61,7 @@ struct Error {
     message: Option<String>,
 
     // Optional fields available on "usage_limit_reached" and "usage_not_included" errors
-    plan_type: Option<String>,
+    plan_type: Option<PlanType>,
     resets_in_seconds: Option<u64>,
 }
 
@@ -304,7 +305,7 @@ impl ModelClient {
                                 // token.
                                 let plan_type = error
                                     .plan_type
-                                    .or_else(|| auth.and_then(|a| a.get_plan_type()));
+                                    .or_else(|| auth.as_ref().and_then(|a| a.get_plan_type()));
                                 let resets_in_seconds = error.resets_in_seconds;
                                 return Err(CodexErr::UsageLimitReached(UsageLimitReachedError {
                                     plan_type,
@@ -1037,4 +1038,37 @@ mod tests {
         let delay = try_parse_retry_after(&err);
         assert_eq!(delay, Some(Duration::from_secs_f64(1.898)));
     }
+
+    #[test]
+    fn error_response_deserializes_old_schema_known_plan_type_and_serializes_back() {
+        use crate::token_data::KnownPlan;
+        use crate::token_data::PlanType;
+
+        let json = r#"{"error":{"type":"usage_limit_reached","plan_type":"pro","resets_in_seconds":3600}}"#;
+        let resp: ErrorResponse =
+            serde_json::from_str(json).expect("should deserialize old schema");
+
+        assert!(matches!(
+            resp.error.plan_type,
+            Some(PlanType::Known(KnownPlan::Pro))
+        ));
+
+        let plan_json = serde_json::to_string(&resp.error.plan_type).expect("serialize plan_type");
+        assert_eq!(plan_json, "\"pro\"");
+    }
+
+    #[test]
+    fn error_response_deserializes_old_schema_unknown_plan_type_and_serializes_back() {
+        use crate::token_data::PlanType;
+
+        let json =
+            r#"{"error":{"type":"usage_limit_reached","plan_type":"vip","resets_in_seconds":60}}"#;
+        let resp: ErrorResponse =
+            serde_json::from_str(json).expect("should deserialize old schema");
+
+        assert!(matches!(resp.error.plan_type, Some(PlanType::Unknown(ref s)) if s == "vip"));
+
+        let plan_json = serde_json::to_string(&resp.error.plan_type).expect("serialize plan_type");
+        assert_eq!(plan_json, "\"vip\"");
+    }
 }

codex-rs/core/src/error.rs

@@ -1,3 +1,5 @@
+use crate::token_data::KnownPlan;
+use crate::token_data::PlanType;
 use codex_protocol::mcp_protocol::ConversationId;
 use reqwest::StatusCode;
 use serde_json;
@@ -127,38 +129,58 @@ pub enum CodexErr {
 
 #[derive(Debug)]
 pub struct UsageLimitReachedError {
-    pub plan_type: Option<String>,
-    pub resets_in_seconds: Option<u64>,
+    pub(crate) plan_type: Option<PlanType>,
+    pub(crate) resets_in_seconds: Option<u64>,
 }
 
 impl std::fmt::Display for UsageLimitReachedError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        // Base message differs slightly for legacy ChatGPT Plus plan users.
-        if let Some(plan_type) = &self.plan_type
-            && plan_type == "plus"
-        {
-            write!(
-                f,
-                "You've hit your usage limit. Upgrade to Pro (https://openai.com/chatgpt/pricing) or try again"
-            )?;
-            if let Some(secs) = self.resets_in_seconds {
-                let reset_duration = format_reset_duration(secs);
-                write!(f, " in {reset_duration}.")?;
-            } else {
-                write!(f, " later.")?;
+        let message = match self.plan_type.as_ref() {
+            Some(PlanType::Known(KnownPlan::Plus)) => format!(
+                "You've hit your usage limit. Upgrade to Pro (https://openai.com/chatgpt/pricing){}",
+                retry_suffix_after_or(self.resets_in_seconds)
+            ),
+            Some(PlanType::Known(KnownPlan::Team)) | Some(PlanType::Known(KnownPlan::Business)) => {
+                format!(
+                    "You've hit your usage limit. To get more access now, send a request to your admin{}",
+                    retry_suffix_after_or(self.resets_in_seconds)
+                )
             }
-        } else {
-            write!(f, "You've hit your usage limit.")?;
-
-            if let Some(secs) = self.resets_in_seconds {
-                let reset_duration = format_reset_duration(secs);
-                write!(f, " Try again in {reset_duration}.")?;
-            } else {
-                write!(f, " Try again later.")?;
+            Some(PlanType::Known(KnownPlan::Free)) => {
+                "To use Codex with your ChatGPT plan, upgrade to Plus: https://openai.com/chatgpt/pricing."
+                    .to_string()
             }
-        }
+            Some(PlanType::Known(KnownPlan::Pro))
+            | Some(PlanType::Known(KnownPlan::Enterprise))
+            | Some(PlanType::Known(KnownPlan::Edu)) => format!(
+                "You've hit your usage limit.{}",
+                retry_suffix(self.resets_in_seconds)
+            ),
+            Some(PlanType::Unknown(_)) | None => format!(
+                "You've hit your usage limit.{}",
+                retry_suffix(self.resets_in_seconds)
+            ),
+        };
 
-        Ok(())
+        write!(f, "{message}")
+    }
+}
+
+fn retry_suffix(resets_in_seconds: Option<u64>) -> String {
+    if let Some(secs) = resets_in_seconds {
+        let reset_duration = format_reset_duration(secs);
+        format!(" Try again in {reset_duration}.")
+    } else {
+        " Try again later.".to_string()
+    }
+}
+
+fn retry_suffix_after_or(resets_in_seconds: Option<u64>) -> String {
+    if let Some(secs) = resets_in_seconds {
+        let reset_duration = format_reset_duration(secs);
+        format!(" or try again in {reset_duration}.")
+    } else {
+        " or try again later.".to_string()
     }
 }
 
@@ -237,7 +259,7 @@ mod tests {
     #[test]
     fn usage_limit_reached_error_formats_plus_plan() {
         let err = UsageLimitReachedError {
-            plan_type: Some("plus".to_string()),
+            plan_type: Some(PlanType::Known(KnownPlan::Plus)),
             resets_in_seconds: None,
         };
         assert_eq!(
@@ -246,6 +268,18 @@ mod tests {
         );
     }
 
+    #[test]
+    fn usage_limit_reached_error_formats_free_plan() {
+        let err = UsageLimitReachedError {
+            plan_type: Some(PlanType::Known(KnownPlan::Free)),
+            resets_in_seconds: Some(3600),
+        };
+        assert_eq!(
+            err.to_string(),
+            "To use Codex with your ChatGPT plan, upgrade to Plus: https://openai.com/chatgpt/pricing."
+        );
+    }
+
     #[test]
     fn usage_limit_reached_error_formats_default_when_none() {
         let err = UsageLimitReachedError {
@@ -258,10 +292,34 @@ mod tests {
         );
     }
 
+    #[test]
+    fn usage_limit_reached_error_formats_team_plan() {
+        let err = UsageLimitReachedError {
+            plan_type: Some(PlanType::Known(KnownPlan::Team)),
+            resets_in_seconds: Some(3600),
+        };
+        assert_eq!(
+            err.to_string(),
+            "You've hit your usage limit. To get more access now, send a request to your admin or try again in 1 hour."
+        );
+    }
+
+    #[test]
+    fn usage_limit_reached_error_formats_business_plan_without_reset() {
+        let err = UsageLimitReachedError {
+            plan_type: Some(PlanType::Known(KnownPlan::Business)),
+            resets_in_seconds: None,
+        };
+        assert_eq!(
+            err.to_string(),
+            "You've hit your usage limit. To get more access now, send a request to your admin or try again later."
+        );
+    }
+
     #[test]
     fn usage_limit_reached_error_formats_default_for_other_plans() {
         let err = UsageLimitReachedError {
-            plan_type: Some("pro".to_string()),
+            plan_type: Some(PlanType::Known(KnownPlan::Pro)),
             resets_in_seconds: None,
         };
         assert_eq!(
@@ -285,7 +343,7 @@ mod tests {
     #[test]
     fn usage_limit_reached_includes_hours_and_minutes() {
         let err = UsageLimitReachedError {
-            plan_type: Some("plus".to_string()),
+            plan_type: Some(PlanType::Known(KnownPlan::Plus)),
             resets_in_seconds: Some(3 * 3600 + 32 * 60),
         };
         assert_eq!(

codex-rs/core/src/token_data.rs

@@ -47,15 +47,6 @@ pub(crate) enum PlanType {
     Unknown(String),
 }
 
-impl PlanType {
-    pub fn as_string(&self) -> String {
-        match self {
-            Self::Known(known) => format!("{known:?}").to_lowercase(),
-            Self::Unknown(s) => s.clone(),
-        }
-    }
-}
-
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "lowercase")]
 pub(crate) enum KnownPlan {