diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index fe882c5d..09ad13b4 100644 --- a/codex-rs/core/src/auth.rs +++ b/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 { + pub(crate) fn get_plan_type(&self) -> Option { 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 { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1268be65..8c194518 100644 --- a/codex-rs/core/src/client.rs +++ b/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, // Optional fields available on "usage_limit_reached" and "usage_not_included" errors - plan_type: Option, + plan_type: Option, resets_in_seconds: Option, } @@ -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\""); + } } diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 36b7abe6..b1cda024 100644 --- a/codex-rs/core/src/error.rs +++ b/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, - pub resets_in_seconds: Option, + pub(crate) plan_type: Option, + pub(crate) resets_in_seconds: Option, } 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) -> 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) -> 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!( diff --git a/codex-rs/core/src/token_data.rs b/codex-rs/core/src/token_data.rs index 2c4f859c..7185a54f 100644 --- a/codex-rs/core/src/token_data.rs +++ b/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 {