auth: let AuthManager own external bearer auth

codex-rs/app-server/src/message_processor.rs

@@ -144,11 +144,11 @@ impl ExternalAuthRefresher for ExternalAuthRefreshBridge {
         let response: ChatgptAuthTokensRefreshResponse =
             serde_json::from_value(result).map_err(std::io::Error::other)?;
 
-        Ok(ExternalAuthTokens {
-            access_token: response.access_token,
-            chatgpt_account_id: response.chatgpt_account_id,
-            chatgpt_plan_type: response.chatgpt_plan_type,
-        })
+        Ok(ExternalAuthTokens::chatgpt(
+            response.access_token,
+            response.chatgpt_account_id,
+            response.chatgpt_plan_type,
+        ))
     }
 }
 

codex-rs/login/src/auth/auth_tests.rs

@@ -4,6 +4,7 @@ use crate::auth::storage::get_auth_file;
 use crate::token_data::IdTokenInfo;
 use crate::token_data::KnownPlan as InternalKnownPlan;
 use crate::token_data::PlanType as InternalPlanType;
+use async_trait::async_trait;
 use codex_protocol::account::PlanType as AccountPlanType;
 
 use base64::Engine;
@@ -12,6 +13,7 @@ use pretty_assertions::assert_eq;
 use serde::Serialize;
 use serde_json::json;
 use std::sync::Arc;
+use std::sync::Mutex;
 use tempfile::tempdir;
 
 #[tokio::test]
@@ -252,6 +254,100 @@ fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
     assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None);
 }
 
+#[test]
+fn external_auth_tokens_without_chatgpt_metadata_cannot_seed_chatgpt_auth() {
+    let err = AuthDotJson::from_external_tokens(&ExternalAuthTokens::access_token_only(
+        "test-access-token",
+    ))
+    .expect_err("bearer-only external auth should not seed ChatGPT auth");
+
+    assert_eq!(
+        err.to_string(),
+        "external auth tokens are missing ChatGPT metadata"
+    );
+}
+
+#[tokio::test]
+async fn auth_manager_with_external_bearer_refresher_returns_provider_token_only_for_derived_manager()
+ {
+    let base_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("base-token"));
+    let derived_manager =
+        base_manager.with_external_bearer_refresher(Arc::new(StaticExternalAuthRefresher::new(
+            Some(ExternalAuthTokens::access_token_only("provider-token")),
+            ExternalAuthTokens::access_token_only("refreshed-provider-token"),
+        )));
+
+    assert_eq!(
+        base_manager
+            .auth()
+            .await
+            .and_then(|auth| auth.api_key().map(str::to_string)),
+        Some("base-token".to_string())
+    );
+    assert_eq!(
+        derived_manager
+            .auth()
+            .await
+            .and_then(|auth| auth.api_key().map(str::to_string)),
+        Some("provider-token".to_string())
+    );
+}
+
+#[tokio::test]
+async fn unauthorized_recovery_uses_external_refresh_for_bearer_manager() {
+    let base_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("base-token"));
+    let refresher = Arc::new(StaticExternalAuthRefresher::new(
+        Some(ExternalAuthTokens::access_token_only("provider-token")),
+        ExternalAuthTokens::access_token_only("refreshed-provider-token"),
+    ));
+    let derived_manager = base_manager.with_external_bearer_refresher(refresher.clone());
+    let mut recovery = derived_manager.unauthorized_recovery();
+
+    assert!(recovery.has_next());
+    assert_eq!(recovery.mode_name(), "external");
+    assert_eq!(recovery.step_name(), "external_refresh");
+
+    let result = recovery
+        .next()
+        .await
+        .expect("external refresh should succeed");
+
+    assert_eq!(result.auth_state_changed(), Some(true));
+    assert_eq!(*refresher.refresh_calls.lock().unwrap(), 1);
+}
+
+#[derive(Debug)]
+struct StaticExternalAuthRefresher {
+    resolved: Option<ExternalAuthTokens>,
+    refreshed: ExternalAuthTokens,
+    refresh_calls: Mutex<usize>,
+}
+
+impl StaticExternalAuthRefresher {
+    fn new(resolved: Option<ExternalAuthTokens>, refreshed: ExternalAuthTokens) -> Self {
+        Self {
+            resolved,
+            refreshed,
+            refresh_calls: Mutex::new(0),
+        }
+    }
+}
+
+#[async_trait]
+impl ExternalAuthRefresher for StaticExternalAuthRefresher {
+    async fn resolve(&self) -> std::io::Result<Option<ExternalAuthTokens>> {
+        Ok(self.resolved.clone())
+    }
+
+    async fn refresh(
+        &self,
+        _context: ExternalAuthRefreshContext,
+    ) -> std::io::Result<ExternalAuthTokens> {
+        *self.refresh_calls.lock().unwrap() += 1;
+        Ok(self.refreshed.clone())
+    }
+}
+
 struct AuthFileParams {
     openai_api_key: Option<String>,
     chatgpt_plan_type: Option<String>,

codex-rs/login/src/auth/manager.rs

@@ -90,11 +90,43 @@ pub enum RefreshTokenError {
     Transient(#[from] std::io::Error),
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ExternalAuthChatgptMetadata {
+    pub account_id: String,
+    pub plan_type: Option<String>,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct ExternalAuthTokens {
     pub access_token: String,
-    pub chatgpt_account_id: String,
-    pub chatgpt_plan_type: Option<String>,
+    pub chatgpt_metadata: Option<ExternalAuthChatgptMetadata>,
+}
+
+impl ExternalAuthTokens {
+    pub fn access_token_only(access_token: impl Into<String>) -> Self {
+        Self {
+            access_token: access_token.into(),
+            chatgpt_metadata: None,
+        }
+    }
+
+    pub fn chatgpt(
+        access_token: impl Into<String>,
+        chatgpt_account_id: impl Into<String>,
+        chatgpt_plan_type: Option<String>,
+    ) -> Self {
+        Self {
+            access_token: access_token.into(),
+            chatgpt_metadata: Some(ExternalAuthChatgptMetadata {
+                account_id: chatgpt_account_id.into(),
+                plan_type: chatgpt_plan_type,
+            }),
+        }
+    }
+
+    pub fn chatgpt_metadata(&self) -> Option<&ExternalAuthChatgptMetadata> {
+        self.chatgpt_metadata.as_ref()
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -110,6 +142,10 @@ pub struct ExternalAuthRefreshContext {
 
 #[async_trait]
 pub trait ExternalAuthRefresher: Send + Sync {
+    async fn resolve(&self) -> std::io::Result<Option<ExternalAuthTokens>> {
+        Ok(None)
+    }
+
     async fn refresh(
         &self,
         context: ExternalAuthRefreshContext,
@@ -736,11 +772,16 @@ fn refresh_token_endpoint() -> String {
 
 impl AuthDotJson {
     fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result<Self> {
+        let Some(chatgpt_metadata) = external.chatgpt_metadata() else {
+            return Err(std::io::Error::other(
+                "external auth tokens are missing ChatGPT metadata",
+            ));
+        };
         let mut token_info =
             parse_chatgpt_jwt_claims(&external.access_token).map_err(std::io::Error::other)?;
-        token_info.chatgpt_account_id = Some(external.chatgpt_account_id.clone());
-        token_info.chatgpt_plan_type = external
-            .chatgpt_plan_type
+        token_info.chatgpt_account_id = Some(chatgpt_metadata.account_id.clone());
+        token_info.chatgpt_plan_type = chatgpt_metadata
+            .plan_type
             .as_deref()
             .map(InternalPlanType::from_raw_value)
             .or(token_info.chatgpt_plan_type)
@@ -749,7 +790,7 @@ impl AuthDotJson {
             id_token: token_info,
             access_token: external.access_token.clone(),
             refresh_token: String::new(),
-            account_id: Some(external.chatgpt_account_id.clone()),
+            account_id: Some(chatgpt_metadata.account_id.clone()),
         };
 
         Ok(Self {
@@ -765,11 +806,11 @@ impl AuthDotJson {
         chatgpt_account_id: &str,
         chatgpt_plan_type: Option<&str>,
     ) -> std::io::Result<Self> {
-        let external = ExternalAuthTokens {
-            access_token: access_token.to_string(),
-            chatgpt_account_id: chatgpt_account_id.to_string(),
-            chatgpt_plan_type: chatgpt_plan_type.map(str::to_string),
-        };
+        let external = ExternalAuthTokens::chatgpt(
+            access_token,
+            chatgpt_account_id,
+            chatgpt_plan_type.map(str::to_string),
+        );
         Self::from_external_tokens(&external)
     }
 
@@ -799,8 +840,6 @@ impl AuthDotJson {
 #[derive(Clone)]
 struct CachedAuth {
     auth: Option<CodexAuth>,
-    /// Callback used to refresh external auth by asking the parent app for new tokens.
-    external_refresher: Option<Arc<dyn ExternalAuthRefresher>>,
     /// Permanent refresh failure cached for the current auth snapshot so
     /// later refresh attempts for the same credentials fail fast without network.
     permanent_refresh_failure: Option<AuthScopedRefreshFailure>,
@@ -812,6 +851,27 @@ struct AuthScopedRefreshFailure {
     error: RefreshTokenFailedError,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum ExternalAuthKind {
+    Bearer,
+    Chatgpt,
+}
+
+#[derive(Clone)]
+struct ExternalAuthHandle {
+    kind: ExternalAuthKind,
+    refresher: Arc<dyn ExternalAuthRefresher>,
+}
+
+impl Debug for ExternalAuthHandle {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ExternalAuthHandle")
+            .field("kind", &self.kind)
+            .field("refresher", &"present")
+            .finish()
+    }
+}
+
 impl Debug for CachedAuth {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CachedAuth")
@@ -819,10 +879,6 @@ impl Debug for CachedAuth {
                 "auth_mode",
                 &self.auth.as_ref().map(CodexAuth::api_auth_mode),
             )
-            .field(
-                "external_refresher",
-                &self.external_refresher.as_ref().map(|_| "present"),
-            )
             .field(
                 "permanent_refresh_failure",
                 &self
@@ -866,9 +922,14 @@ 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 external ChatGPT auth tokens (chatgptAuthTokens), UnauthorizedRecovery does not touch disk or refresh
-// tokens locally. Instead it calls the ExternalAuthRefresher (account/chatgptAuthTokens/refresh) to ask the
-// parent app for new tokens, stores them in the ephemeral auth store, and retries once.
+// For external auth sources, UnauthorizedRecovery delegates to the configured
+// ExternalAuthRefresher and retries once.
+//
+// - External ChatGPT auth tokens (`chatgptAuthTokens`) are refreshed by asking
+//   the parent app for new tokens, persisting them in the ephemeral auth
+//   store, and reloading the cached auth snapshot.
+// - External bearer auth sources resolve bearer-only tokens for custom model
+//   providers and refresh them without touching disk.
 pub struct UnauthorizedRecovery {
     manager: Arc<AuthManager>,
     step: UnauthorizedRecoveryStep,
@@ -891,9 +952,10 @@ impl UnauthorizedRecovery {
     fn new(manager: Arc<AuthManager>) -> Self {
         let cached_auth = manager.auth_cached();
         let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id);
-        let mode = if cached_auth
-            .as_ref()
-            .is_some_and(CodexAuth::is_external_chatgpt_tokens)
+        let mode = if manager.external_auth_kind() == Some(ExternalAuthKind::Bearer)
+            || cached_auth
+                .as_ref()
+                .is_some_and(CodexAuth::is_external_chatgpt_tokens)
         {
             UnauthorizedRecoveryMode::External
         } else {
@@ -912,6 +974,10 @@ impl UnauthorizedRecovery {
     }
 
     pub fn has_next(&self) -> bool {
+        if self.manager.external_auth_kind() == Some(ExternalAuthKind::Bearer) {
+            return !matches!(self.step, UnauthorizedRecoveryStep::Done);
+        }
+
         if !self
             .manager
             .auth_cached()
@@ -931,6 +997,16 @@ impl UnauthorizedRecovery {
     }
 
     pub fn unavailable_reason(&self) -> &'static str {
+        if self.manager.external_auth_kind() == Some(ExternalAuthKind::Bearer) {
+            return if matches!(self.step, UnauthorizedRecoveryStep::Done) {
+                "recovery_exhausted"
+            } else if self.manager.has_external_auth_refresher() {
+                "ready"
+            } else {
+                "no_external_refresher"
+            };
+        }
+
         if !self
             .manager
             .auth_cached()
@@ -1039,11 +1115,12 @@ impl UnauthorizedRecovery {
 #[derive(Debug)]
 pub struct AuthManager {
     codex_home: PathBuf,
-    inner: RwLock<CachedAuth>,
+    inner: Arc<RwLock<CachedAuth>>,
     enable_codex_api_key_env: bool,
     auth_credentials_store_mode: AuthCredentialsStoreMode,
-    forced_chatgpt_workspace_id: RwLock<Option<String>>,
-    refresh_lock: AsyncMutex<()>,
+    forced_chatgpt_workspace_id: Arc<RwLock<Option<String>>>,
+    refresh_lock: Arc<AsyncMutex<()>>,
+    external_auth: RwLock<Option<ExternalAuthHandle>>,
 }
 
 impl AuthManager {
@@ -1065,15 +1142,15 @@ impl AuthManager {
         .flatten();
         Self {
             codex_home,
-            inner: RwLock::new(CachedAuth {
+            inner: Arc::new(RwLock::new(CachedAuth {
                 auth: managed_auth,
-                external_refresher: None,
                 permanent_refresh_failure: None,
-            }),
+            })),
             enable_codex_api_key_env,
             auth_credentials_store_mode,
-            forced_chatgpt_workspace_id: RwLock::new(None),
-            refresh_lock: AsyncMutex::new(()),
+            forced_chatgpt_workspace_id: Arc::new(RwLock::new(None)),
+            refresh_lock: Arc::new(AsyncMutex::new(())),
+            external_auth: RwLock::new(None),
         }
     }
 
@@ -1081,17 +1158,17 @@ impl AuthManager {
     pub fn from_auth_for_testing(auth: CodexAuth) -> Arc<Self> {
         let cached = CachedAuth {
             auth: Some(auth),
-            external_refresher: None,
             permanent_refresh_failure: None,
         };
 
         Arc::new(Self {
             codex_home: PathBuf::from("non-existent"),
-            inner: RwLock::new(cached),
+            inner: Arc::new(RwLock::new(cached)),
             enable_codex_api_key_env: false,
             auth_credentials_store_mode: AuthCredentialsStoreMode::File,
-            forced_chatgpt_workspace_id: RwLock::new(None),
-            refresh_lock: AsyncMutex::new(()),
+            forced_chatgpt_workspace_id: Arc::new(RwLock::new(None)),
+            refresh_lock: Arc::new(AsyncMutex::new(())),
+            external_auth: RwLock::new(None),
         })
     }
 
@@ -1099,16 +1176,16 @@ impl AuthManager {
     pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc<Self> {
         let cached = CachedAuth {
             auth: Some(auth),
-            external_refresher: None,
             permanent_refresh_failure: None,
         };
         Arc::new(Self {
             codex_home,
-            inner: RwLock::new(cached),
+            inner: Arc::new(RwLock::new(cached)),
             enable_codex_api_key_env: false,
             auth_credentials_store_mode: AuthCredentialsStoreMode::File,
-            forced_chatgpt_workspace_id: RwLock::new(None),
-            refresh_lock: AsyncMutex::new(()),
+            forced_chatgpt_workspace_id: Arc::new(RwLock::new(None)),
+            refresh_lock: Arc::new(AsyncMutex::new(())),
+            external_auth: RwLock::new(None),
         })
     }
 
@@ -1131,6 +1208,10 @@ impl AuthManager {
     /// For stale managed ChatGPT auth, first performs a guarded reload and then
     /// refreshes only if the on-disk auth is unchanged.
     pub async fn auth(&self) -> Option<CodexAuth> {
+        if let Some(auth) = self.resolve_external_api_key_auth().await {
+            return Some(auth);
+        }
+
         let auth = self.auth_cached()?;
         if Self::is_stale_for_proactive_refresh(&auth)
             && let Err(err) = self.refresh_token().await
@@ -1251,17 +1332,38 @@ impl AuthManager {
     }
 
     pub fn set_external_auth_refresher(&self, refresher: Arc<dyn ExternalAuthRefresher>) {
-        if let Ok(mut guard) = self.inner.write() {
-            guard.external_refresher = Some(refresher);
+        if let Ok(mut guard) = self.external_auth.write() {
+            *guard = Some(ExternalAuthHandle {
+                kind: ExternalAuthKind::Chatgpt,
+                refresher,
+            });
         }
     }
 
     pub fn clear_external_auth_refresher(&self) {
-        if let Ok(mut guard) = self.inner.write() {
-            guard.external_refresher = None;
+        if let Ok(mut guard) = self.external_auth.write() {
+            *guard = None;
         }
     }
 
+    pub fn with_external_bearer_refresher(
+        self: &Arc<Self>,
+        refresher: Arc<dyn ExternalAuthRefresher>,
+    ) -> Arc<Self> {
+        Arc::new(Self {
+            codex_home: self.codex_home.clone(),
+            inner: Arc::clone(&self.inner),
+            enable_codex_api_key_env: self.enable_codex_api_key_env,
+            auth_credentials_store_mode: self.auth_credentials_store_mode,
+            forced_chatgpt_workspace_id: Arc::clone(&self.forced_chatgpt_workspace_id),
+            refresh_lock: Arc::clone(&self.refresh_lock),
+            external_auth: RwLock::new(Some(ExternalAuthHandle {
+                kind: ExternalAuthKind::Bearer,
+                refresher,
+            })),
+        })
+    }
+
     pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option<String>) {
         if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() {
             *guard = workspace_id;
@@ -1276,13 +1378,17 @@ impl AuthManager {
     }
 
     pub fn has_external_auth_refresher(&self) -> bool {
-        self.inner
+        self.external_auth
             .read()
             .ok()
-            .map(|guard| guard.external_refresher.is_some())
+            .map(|guard| guard.is_some())
             .unwrap_or(false)
     }
 
+    pub fn has_external_bearer_refresher(&self) -> bool {
+        self.external_auth_kind() == Some(ExternalAuthKind::Bearer)
+    }
+
     pub fn is_external_auth_active(&self) -> bool {
         self.auth_cached()
             .as_ref()
@@ -1310,6 +1416,35 @@ impl AuthManager {
         UnauthorizedRecovery::new(Arc::clone(self))
     }
 
+    fn external_auth_handle(&self) -> Option<ExternalAuthHandle> {
+        self.external_auth
+            .read()
+            .ok()
+            .and_then(|guard| guard.clone())
+    }
+
+    fn external_auth_kind(&self) -> Option<ExternalAuthKind> {
+        self.external_auth_handle().map(|handle| handle.kind)
+    }
+
+    async fn resolve_external_api_key_auth(&self) -> Option<CodexAuth> {
+        let Some(handle) = self.external_auth_handle() else {
+            return None;
+        };
+        if handle.kind != ExternalAuthKind::Bearer {
+            return None;
+        }
+
+        match handle.refresher.resolve().await {
+            Ok(Some(tokens)) => Some(CodexAuth::from_api_key(&tokens.access_token)),
+            Ok(None) => None,
+            Err(err) => {
+                tracing::error!("Failed to resolve external bearer auth: {err}");
+                None
+            }
+        }
+    }
+
     /// Attempt to refresh the token by first performing a guarded reload. Auth
     /// is reloaded from storage only when the account id matches the currently
     /// cached account id. If the persisted token differs from the cached token, we
@@ -1432,16 +1567,7 @@ impl AuthManager {
         reason: ExternalAuthRefreshReason,
     ) -> Result<(), RefreshTokenError> {
         let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id();
-        let refresher = match self.inner.read() {
-            Ok(guard) => guard.external_refresher.clone(),
-            Err(_) => {
-                return Err(RefreshTokenError::Transient(std::io::Error::other(
-                    "failed to read external auth state",
-                )));
-            }
-        };
-
-        let Some(refresher) = refresher else {
+        let Some(handle) = self.external_auth_handle() else {
             return Err(RefreshTokenError::Transient(std::io::Error::other(
                 "external auth refresher is not configured",
             )));
@@ -1456,14 +1582,22 @@ impl AuthManager {
             previous_account_id,
         };
 
-        let refreshed = refresher.refresh(context).await?;
+        let refreshed = handle.refresher.refresh(context).await?;
+        if handle.kind == ExternalAuthKind::Bearer {
+            return Ok(());
+        }
+        let Some(chatgpt_metadata) = refreshed.chatgpt_metadata() else {
+            return Err(RefreshTokenError::Transient(std::io::Error::other(
+                "external auth refresh did not return ChatGPT metadata",
+            )));
+        };
         if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref()
-            && refreshed.chatgpt_account_id != expected_workspace_id
+            && chatgpt_metadata.account_id != expected_workspace_id
         {
             return Err(RefreshTokenError::Transient(std::io::Error::other(
                 format!(
                     "external auth refresh returned workspace {:?}, expected {expected_workspace_id:?}",
-                    refreshed.chatgpt_account_id,
+                    chatgpt_metadata.account_id,
                 ),
             )));
         }

codex-rs/login/src/lib.rs

@@ -22,6 +22,11 @@ pub use auth::AuthManager;
 pub use auth::CLIENT_ID;
 pub use auth::CODEX_API_KEY_ENV_VAR;
 pub use auth::CodexAuth;
+pub use auth::ExternalAuthChatgptMetadata;
+pub use auth::ExternalAuthRefreshContext;
+pub use auth::ExternalAuthRefreshReason;
+pub use auth::ExternalAuthRefresher;
+pub use auth::ExternalAuthTokens;
 pub use auth::OPENAI_API_KEY_ENV_VAR;
 pub use auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
 pub use auth::RefreshTokenError;