fix(auth): treat prolite as pro-capable

Author: enieuwy

code-rs/core/src/auth.rs

@@ -259,8 +259,9 @@ impl CodexAuth {
     pub fn supports_pro_only_models(&self) -> bool {
         self.uses_codex_backend()
             && self
-                .get_plan_type()
-                .is_some_and(|plan| plan.eq_ignore_ascii_case("pro"))
+                .get_current_token_data()
+                .and_then(|t| t.id_token.chatgpt_plan_type)
+                .is_some_and(|plan| plan.supports_pro_only_models())
     }
 
     fn get_current_auth_json(&self) -> Option<AuthDotJson> {
@@ -430,6 +431,7 @@ pub fn login_with_chatgpt_auth_tokens(
             "free" => PlanType::Known(KnownPlan::Free),
             "plus" => PlanType::Known(KnownPlan::Plus),
             "pro" => PlanType::Known(KnownPlan::Pro),
+            "prolite" => PlanType::Known(KnownPlan::ProLite),
             "team" => PlanType::Known(KnownPlan::Team),
             "business" => PlanType::Known(KnownPlan::Business),
             "enterprise" => PlanType::Known(KnownPlan::Enterprise),
@@ -1120,6 +1122,27 @@ mod tests {
         )
     }
 
+    #[test]
+    fn prolite_account_supports_pro_only_models() {
+        let auth = CodexAuth::from_tokens_with_originator(
+            TokenData {
+                id_token: IdTokenInfo {
+                    email: Some("user@example.com".to_string()),
+                    chatgpt_plan_type: Some(PlanType::Known(KnownPlan::ProLite)),
+                    chatgpt_account_is_fedramp: false,
+                    raw_jwt: "header.payload.signature".to_string(),
+                },
+                access_token: "test-access-token".to_string(),
+                refresh_token: "test-refresh-token".to_string(),
+                account_id: None,
+            },
+            None,
+            "code_cli_rs",
+        );
+
+        assert!(auth.supports_pro_only_models());
+    }
+
     /// Even if the OPENAI_API_KEY is set in auth.json, if the plan is not in
     /// [`TokenData::is_plan_that_should_use_api_key`], it should use
     /// [`AuthMode::ChatGPT`].

code-rs/core/src/token_data.rs

@@ -73,6 +73,11 @@ impl PlanType {
             Self::Unknown(s) => s.clone(),
         }
     }
+
+    pub(crate) fn supports_pro_only_models(&self) -> bool {
+        matches!(self, Self::Known(KnownPlan::Pro | KnownPlan::ProLite))
+            || matches!(self, Self::Unknown(plan) if plan.eq_ignore_ascii_case("prolite"))
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -81,6 +86,7 @@ pub(crate) enum KnownPlan {
     Free,
     Plus,
     Pro,
+    ProLite,
     Team,
     Business,
     Enterprise,
@@ -224,6 +230,39 @@ mod tests {
         assert!(!info.is_fedramp_account());
     }
 
+    #[test]
+    fn prolite_plan_supports_pro_only_models() {
+        #[derive(Serialize)]
+        struct Header {
+            alg: &'static str,
+            typ: &'static str,
+        }
+        let header = Header {
+            alg: "none",
+            typ: "JWT",
+        };
+        let payload = serde_json::json!({
+            "email": "user@example.com",
+            "https://api.openai.com/auth": {
+                "chatgpt_plan_type": "prolite"
+            }
+        });
+
+        fn b64url_no_pad(bytes: &[u8]) -> String {
+            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
+        }
+
+        let header_b64 = b64url_no_pad(&serde_json::to_vec(&header).unwrap());
+        let payload_b64 = b64url_no_pad(&serde_json::to_vec(&payload).unwrap());
+        let signature_b64 = b64url_no_pad(b"sig");
+        let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}");
+
+        let info = parse_id_token(&fake_jwt).expect("should parse");
+        let plan = info.chatgpt_plan_type.expect("plan should parse");
+        assert_eq!(plan.as_string(), "prolite");
+        assert!(plan.supports_pro_only_models());
+    }
+
     #[test]
     fn id_token_info_handles_missing_fields() {
         #[derive(Serialize)]