From c579ae41aee48e7d2a0c9171bb6bd655d61736d0 Mon Sep 17 00:00:00 2001
From: Ahmed Ibrahim <aibrahim@openai.com>
Date: Wed, 20 Aug 2025 14:05:20 -0700
Subject: [PATCH] Fix login for internal employees (#2528)

This PR:
- fixes for internal employee because we currently want to prefer SIWC
for them.
- fixes retrying forever on unauthorized access. we need to break
eventually on max retries.
---
 codex-rs/core/src/client.rs      | 15 ++++++++-------
 codex-rs/login/src/lib.rs        |  2 +-
 codex-rs/login/src/token_data.rs | 17 ++++++++++++++++-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs
index 9552a53d..5534e11f 100644
--- a/codex-rs/core/src/client.rs
+++ b/codex-rs/core/src/client.rs
@@ -248,12 +248,10 @@ impl ModelClient {
                         .and_then(|v| v.to_str().ok())
                         .and_then(|s| s.parse::<u64>().ok());
 
-                    if status == StatusCode::UNAUTHORIZED {
-                        if let Some(a) = auth.as_ref() {
-                            let _ = a.refresh_token().await;
-                        }
-                        // Retry immediately with refreshed credentials.
-                        continue;
+                    if status == StatusCode::UNAUTHORIZED
+                        && let Some(a) = auth.as_ref()
+                    {
+                        let _ = a.refresh_token().await;
                     }
 
                     // The OpenAI Responses endpoint returns structured JSON bodies even for 4xx/5xx
@@ -263,7 +261,10 @@ impl ModelClient {
                     // exact error message (e.g. "Unknown parameter: 'input[0].metadata'"). The body is
                     // small and this branch only runs on error paths so the extra allocation is
                     // negligible.
-                    if !(status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()) {
+                    if !(status == StatusCode::TOO_MANY_REQUESTS
+                        || status == StatusCode::UNAUTHORIZED
+                        || status.is_server_error())
+                    {
                         // Surface the error body to callers. Use `unwrap_or_default` per Clippy.
                         let body = res.text().await.unwrap_or_default();
                         return Err(CodexErr::UnexpectedStatus(status, body));
diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs
index d6fa4aef..8c9a5cf3 100644
--- a/codex-rs/login/src/lib.rs
+++ b/codex-rs/login/src/lib.rs
@@ -242,7 +242,7 @@ fn load_auth(
         // "refreshable" even if we are using the API key for auth?
         match &tokens {
             Some(tokens) => {
-                if tokens.should_use_api_key(preferred_auth_method) {
+                if tokens.should_use_api_key(preferred_auth_method, tokens.is_openai_email()) {
                     return Ok(Some(CodexAuth::from_api_key(api_key)));
                 } else {
                     // Ignore the API key and fall through to ChatGPT auth.
diff --git a/codex-rs/login/src/token_data.rs b/codex-rs/login/src/token_data.rs
index 59e6d972..f6d04f16 100644
--- a/codex-rs/login/src/token_data.rs
+++ b/codex-rs/login/src/token_data.rs
@@ -25,16 +25,31 @@ pub struct TokenData {
 impl TokenData {
     /// Returns true if this is a plan that should use the traditional
     /// "metered" billing via an API key.
-    pub(crate) fn should_use_api_key(&self, preferred_auth_method: AuthMode) -> bool {
+    pub(crate) fn should_use_api_key(
+        &self,
+        preferred_auth_method: AuthMode,
+        is_openai_email: bool,
+    ) -> bool {
         if preferred_auth_method == AuthMode::ApiKey {
             return true;
         }
+        // If the email is an OpenAI email, use AuthMode::ChatGPT unless preferred_auth_method is AuthMode::ApiKey.
+        if is_openai_email {
+            return false;
+        }
 
         self.id_token
             .chatgpt_plan_type
             .as_ref()
             .is_none_or(|plan| plan.is_plan_that_should_use_api_key())
     }
+
+    pub fn is_openai_email(&self) -> bool {
+        self.id_token
+            .email
+            .as_deref()
+            .is_some_and(|email| email.trim().to_ascii_lowercase().ends_with("@openai.com"))
+    }
 }
 
 /// Flat subset of useful claims in id_token from auth.json.