From f3ba12b33ede316b5ebe40e22bdec447826086b9 Mon Sep 17 00:00:00 2001
From: Edward Frazer <efrazer@openai.com>
Date: Wed, 6 May 2026 14:31:20 -0700
Subject: [PATCH] fix: retry agent identity auth after guarded reload

---
 codex-rs/login/src/auth/auth_tests.rs | 48 +++++++++++++++++++++++++++
 codex-rs/login/src/auth/manager.rs    | 23 ++++++++++---
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs
index 80ec9d07a4..cee35d46ac 100644
--- a/codex-rs/login/src/auth/auth_tests.rs
+++ b/codex-rs/login/src/auth/auth_tests.rs
@@ -464,6 +464,54 @@ async fn unauthorized_recovery_uses_external_refresh_for_bearer_manager() {
     assert_eq!(refreshed_token.as_deref(), Some("refreshed-provider-token"));
 }
 
+#[tokio::test]
+#[serial(codex_auth_env)]
+async fn unauthorized_recovery_reloads_agent_identity_once() {
+    let codex_home = tempdir().unwrap();
+    let record = agent_identity_record("account-123");
+    let agent_identity =
+        signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity");
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .and(path("/backend-api/wham/agent-identities/jwks"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
+        .expect(2)
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "task_id": "task-123",
+        })))
+        .expect(2)
+        .mount(&server)
+        .await;
+    let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
+    let chatgpt_base_url = format!("{}/backend-api", server.uri());
+    let _authapi_guard =
+        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
+    let manager = AuthManager::shared(
+        codex_home.path().to_path_buf(),
+        /*enable_codex_api_key_env*/ false,
+        AuthCredentialsStoreMode::File,
+        Some(chatgpt_base_url),
+    )
+    .await;
+
+    let mut recovery = manager.unauthorized_recovery();
+
+    assert!(recovery.has_next());
+    assert_eq!(recovery.mode_name(), "managed");
+    assert_eq!(recovery.step_name(), "reload");
+
+    let result = recovery.next().await.expect("reload should succeed");
+
+    assert_eq!(result.auth_state_changed(), Some(false));
+    assert!(!recovery.has_next());
+    assert_eq!(recovery.step_name(), "done");
+    server.verify().await;
+}
+
 struct ProviderAuthScript {
     tempdir: TempDir,
     command: String,
diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs
index 29897db7be..e526668a8e 100644
--- a/codex-rs/login/src/auth/manager.rs
+++ b/codex-rs/login/src/auth/manager.rs
@@ -1056,6 +1056,8 @@ enum UnauthorizedRecoveryMode {
 // 2. Attempt to refresh the token using OAuth token refresh flow.
 // If after both steps the server still responds with 401 we let the error bubble to the user.
 //
+// For agent identity based authentication, we retry once after the guarded reload step.
+//
 // For external auth sources, UnauthorizedRecovery retries once.
 //
 // - External ChatGPT auth tokens (`chatgptAuthTokens`) are refreshed by asking
@@ -1116,7 +1118,7 @@ impl UnauthorizedRecovery {
             .manager
             .auth_cached()
             .as_ref()
-            .is_some_and(CodexAuth::is_chatgpt_auth)
+            .is_some_and(CodexAuth::uses_codex_backend)
         {
             return false;
         }
@@ -1141,7 +1143,7 @@ impl UnauthorizedRecovery {
             .manager
             .auth_cached()
             .as_ref()
-            .is_some_and(CodexAuth::is_chatgpt_auth)
+            .is_some_and(CodexAuth::uses_codex_backend)
         {
             return "not_chatgpt_auth";
         }
@@ -1189,13 +1191,13 @@ impl UnauthorizedRecovery {
                     .await
                 {
                     ReloadOutcome::ReloadedChanged => {
-                        self.step = UnauthorizedRecoveryStep::RefreshToken;
+                        self.step = self.next_step_after_reload();
                         return Ok(UnauthorizedRecoveryStepResult {
                             auth_state_changed: Some(true),
                         });
                     }
                     ReloadOutcome::ReloadedNoChange => {
-                        self.step = UnauthorizedRecoveryStep::RefreshToken;
+                        self.step = self.next_step_after_reload();
                         return Ok(UnauthorizedRecoveryStepResult {
                             auth_state_changed: Some(false),
                         });
@@ -1231,6 +1233,19 @@ impl UnauthorizedRecovery {
             auth_state_changed: None,
         })
     }
+
+    fn next_step_after_reload(&self) -> UnauthorizedRecoveryStep {
+        if self
+            .manager
+            .auth_cached()
+            .as_ref()
+            .is_some_and(|auth| matches!(auth, CodexAuth::AgentIdentity(_)))
+        {
+            UnauthorizedRecoveryStep::Done
+        } else {
+            UnauthorizedRecoveryStep::RefreshToken
+        }
+    }
 }
 
 /// Central manager providing a single source of truth for auth.json derived