From 74e1df42260cc82abb2fc74d8295a65e6ab5364e Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 08:38:40 -0700 Subject: [PATCH 01/10] refactor: unify external auth resolution --- codex-rs/app-server/src/message_processor.rs | 5 - codex-rs/app-server/src/request_processors.rs | 1 - .../request_processors/account_processor.rs | 8 +- codex-rs/login/src/auth/auth_tests.rs | 40 +++++++ codex-rs/login/src/auth/external_bearer.rs | 5 - codex-rs/login/src/auth/manager.rs | 100 +++++++++++------- codex-rs/models-manager/src/manager.rs | 3 + codex-rs/models-manager/src/manager_tests.rs | 8 -- 8 files changed, 112 insertions(+), 58 deletions(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 29ace7d99fc..f9b88f52067 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -79,7 +79,6 @@ use codex_login::auth::ExternalAuth; use codex_login::auth::ExternalAuthRefreshContext; use codex_login::auth::ExternalAuthRefreshReason; use codex_protocol::ThreadId; -use codex_protocol::auth::AuthMode as LoginAuthMode; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::StateDbHandle; @@ -168,10 +167,6 @@ impl ExternalAuthRefreshBridge { } impl ExternalAuth for ExternalAuthRefreshBridge { - fn auth_mode(&self) -> LoginAuthMode { - LoginAuthMode::Chatgpt - } - fn refresh( &self, context: ExternalAuthRefreshContext, diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index f03351468fa..7f5ee69a09d 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -382,7 +382,6 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::ServerOptions as LoginServerOptions; use codex_login::ShutdownHandle; -use codex_login::auth::login_with_chatgpt_auth_tokens; use codex_login::complete_device_code_login; use codex_login::login_with_api_key; use codex_login::oauth_client_id; diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 9592563bece..97e7a18d356 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -617,14 +617,16 @@ impl AccountRequestProcessor { ))); } - login_with_chatgpt_auth_tokens( - &self.config.codex_home, + let auth = CodexAuth::from_external_chatgpt_tokens( &access_token, &chatgpt_account_id, chatgpt_plan_type.as_deref(), ) .map_err(|err| internal_error(format!("failed to set external auth: {err}")))?; - self.auth_manager.reload().await; + self.auth_manager + .install_external_auth(auth) + .await + .map_err(|err| internal_error(format!("failed to set external auth: {err}")))?; self.config_manager.replace_cloud_config_bundle_loader( self.auth_manager.clone(), self.config.chatgpt_base_url.clone(), diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 77d8c4fdfb2..cdc680c7511 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1015,6 +1015,46 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None); } +struct RefreshOnlyExternalChatgptAuth; + +impl ExternalAuth for RefreshOnlyExternalChatgptAuth { + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Err(std::io::Error::other("refresh should not be called")) }) + } +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn external_chatgpt_auth_resolves_from_the_manager_cache() { + let codex_home = tempdir().unwrap(); + let params = AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), + }; + let access_token = fake_jwt_for_auth_file_params(¶ms).unwrap(); + let auth = + CodexAuth::from_external_chatgpt_tokens(&access_token, WORKSPACE_ID_ALLOWED, Some("pro")) + .unwrap(); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await; + manager.set_external_auth(Arc::new(RefreshOnlyExternalChatgptAuth)); + manager.install_external_auth(auth).await.unwrap(); + + let auth = manager.auth().await.expect("external auth should resolve"); + + assert_eq!(auth.api_auth_mode(), AuthMode::ChatgptAuthTokens); + assert_eq!(auth.get_account_id().as_deref(), Some(WORKSPACE_ID_ALLOWED)); +} + #[tokio::test] async fn external_bearer_only_auth_manager_uses_cached_provider_token() { let script = ProviderAuthScript::new(&["provider-token", "next-token"]).unwrap(); diff --git a/codex-rs/login/src/auth/external_bearer.rs b/codex-rs/login/src/auth/external_bearer.rs index 1dbd2c3eeef..7e8a47e22fa 100644 --- a/codex-rs/login/src/auth/external_bearer.rs +++ b/codex-rs/login/src/auth/external_bearer.rs @@ -2,7 +2,6 @@ use super::manager::CodexAuth; use super::manager::ExternalAuth; use super::manager::ExternalAuthFuture; use super::manager::ExternalAuthRefreshContext; -use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ModelProviderAuthInfo; use std::fmt; use std::io; @@ -67,10 +66,6 @@ impl BearerTokenRefresher { } impl ExternalAuth for BearerTokenRefresher { - fn auth_mode(&self) -> AuthMode { - AuthMode::ApiKey - } - fn resolve(&self) -> ExternalAuthFuture<'_, Option> { Box::pin(BearerTokenRefresher::resolve(self)) } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 06d125422c4..ac8d20870d3 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -214,11 +214,8 @@ pub struct ExternalAuthRefreshContext { /// Implementations may either resolve auth eagerly via `resolve()` or provide refreshed /// auth on demand via `refresh()`. pub trait ExternalAuth: Send + Sync { - /// Indicates which top-level auth mode this external provider supplies. - fn auth_mode(&self) -> AuthMode; - - /// Returns cached or immediately available auth, if this provider can resolve it synchronously - /// from the caller's perspective. + /// Returns cached or immediately available auth, if this provider owns the current value. + /// Returning `None` lets the manager use externally installed ChatGPT auth from its cache. fn resolve(&self) -> ExternalAuthFuture<'_, Option> { Box::pin(async { Ok(None) }) } @@ -1763,6 +1760,7 @@ pub struct AuthManager { agent_identity_lock: Semaphore, agent_identity_bootstrap_cooldown: Mutex, external_auth: RwLock>>, + resolved_external_auth: RwLock>, auth_route_config: Option, } @@ -1865,6 +1863,7 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), + resolved_external_auth: RwLock::new(None), auth_route_config, } } @@ -1891,6 +1890,7 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), + resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } @@ -1916,6 +1916,7 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), + resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } @@ -1949,6 +1950,7 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), + resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } @@ -1974,13 +1976,15 @@ impl AuthManager { external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), + resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } /// Current cached auth (clone) without attempting a refresh. pub fn auth_cached(&self) -> Option { - self.inner.read().ok().and_then(|c| c.auth.clone()) + self.resolved_external_auth() + .or_else(|| self.inner.read().ok().and_then(|c| c.auth.clone())) } /// Subscribes to cached auth changes that can affect request recovery. @@ -2003,7 +2007,7 @@ impl AuthManager { /// a guarded reload and then refreshes only if the on-disk auth is unchanged. #[instrument(level = "trace", skip_all)] pub async fn auth(&self) -> Option { - if let Some(auth) = self.resolve_external_api_key_auth().await { + if let Some(auth) = self.resolve_external_auth().await { return Some(auth); } @@ -2214,12 +2218,14 @@ impl AuthManager { if let Ok(mut guard) = self.external_auth.write() { *guard = Some(external_auth); } + self.set_resolved_external_auth(None); } pub fn clear_external_auth(&self) { if let Ok(mut guard) = self.external_auth.write() { *guard = None; } + self.set_resolved_external_auth(None); } pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option>) { @@ -2303,35 +2309,45 @@ impl AuthManager { .and_then(|guard| guard.as_ref().cloned()) } - fn external_auth_mode(&self) -> Option { - self.external_auth() - .as_ref() - .map(|external_auth| external_auth.auth_mode()) + fn resolved_external_auth(&self) -> Option { + self.resolved_external_auth + .read() + .ok() + .and_then(|guard| guard.clone()) } - fn has_external_api_key_auth(&self) -> bool { - self.external_auth_mode() == Some(AuthMode::ApiKey) + fn set_resolved_external_auth(&self, auth: Option) { + if let Ok(mut guard) = self.resolved_external_auth.write() { + *guard = auth; + } } - async fn resolve_external_api_key_auth(&self) -> Option { - if !self.has_external_api_key_auth() { - return None; - } + fn has_external_api_key_auth(&self) -> bool { + self.has_external_auth() + && self + .resolved_external_auth() + .as_ref() + .is_some_and(CodexAuth::is_api_key_auth) + } + async fn resolve_external_auth(&self) -> Option { let external_auth = self.external_auth()?; match external_auth.resolve().await { - Ok(Some(auth)) if auth.api_auth_mode() == AuthMode::ApiKey => Some(auth), Ok(Some(auth)) => { - tracing::error!( - resolved_mode = ?auth.api_auth_mode(), - "ignoring non-API-key auth from external API key provider" - ); + if let Err(err) = self.install_external_auth(auth.clone()).await { + tracing::error!("Failed to install resolved external auth: {err}"); + return None; + } + Some(auth) + } + Ok(None) => { + self.set_resolved_external_auth(None); None } - Ok(None) => None, Err(err) => { - tracing::error!("Failed to resolve external API key auth: {err}"); + self.set_resolved_external_auth(None); + tracing::error!("Failed to resolve external auth: {err}"); None } } @@ -2465,17 +2481,11 @@ impl AuthManager { /// Returns the precise kind of credentials backing the current authentication. pub fn get_api_auth_mode(&self) -> Option { - if self.has_external_api_key_auth() { - return Some(AuthMode::ApiKey); - } self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) } /// Returns the effective backend auth mode for the current authentication. pub fn auth_mode(&self) -> Option { - if self.has_external_api_key_auth() { - return Some(AuthMode::ApiKey); - } self.auth_cached().as_ref().map(CodexAuth::auth_mode) } @@ -2517,7 +2527,6 @@ impl AuthManager { "external auth is not configured", ))); }; - let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); let previous_account_id = self .auth_cached() .as_ref() @@ -2531,15 +2540,33 @@ impl AuthManager { .refresh(context) .await .map_err(RefreshTokenError::Transient)?; - if external_auth.auth_mode() == AuthMode::ApiKey { + self.install_external_auth(refreshed).await + } + + /// Validates and installs auth returned by an [`ExternalAuth`] provider. + pub async fn install_external_auth(&self, auth: CodexAuth) -> Result<(), RefreshTokenError> { + if auth.api_auth_mode() == AuthMode::ApiKey { + self.set_resolved_external_auth(Some(auth)); return Ok(()); } - if refreshed.api_auth_mode() != AuthMode::ChatgptAuthTokens { + if !auth.is_external_chatgpt_tokens() { return Err(RefreshTokenError::Transient(std::io::Error::other( - "external auth refresh did not return ChatGPT auth tokens", + format!( + "external auth returned unsupported auth mode {:?}", + auth.api_auth_mode() + ), ))); } - let account_id = refreshed.get_account_id().ok_or_else(|| { + if self + .inner + .read() + .ok() + .is_some_and(|cached| Self::auths_equal_for_refresh(cached.auth.as_ref(), Some(&auth))) + { + return Ok(()); + } + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); + let account_id = auth.get_account_id().ok_or_else(|| { RefreshTokenError::Transient(std::io::Error::other( "external ChatGPT auth tokens are missing an account id", )) @@ -2553,7 +2580,7 @@ impl AuthManager { ), ))); } - let auth_dot_json = refreshed.get_current_auth_json().ok_or_else(|| { + let auth_dot_json = auth.get_current_auth_json().ok_or_else(|| { RefreshTokenError::Transient(std::io::Error::other( "external ChatGPT auth tokens are missing auth state", )) @@ -2566,6 +2593,7 @@ impl AuthManager { ) .map_err(RefreshTokenError::Transient)?; self.reload().await; + self.set_resolved_external_auth(None); Ok(()) } diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 67b939051b6..9ab09469cc4 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -86,6 +86,9 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { ) -> ModelsManagerFuture<'_, Vec> { Box::pin( async move { + if let Some(auth_manager) = self.auth_manager() { + let _ = auth_manager.auth().await; + } let catalog = self.raw_model_catalog(refresh_strategy).await; self.build_available_models(catalog.models) } diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 2e97dfac431..fc8c100a4f6 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -117,10 +117,6 @@ impl TestModelsEndpoint { struct TestExternalApiKeyAuth; impl ExternalAuth for TestExternalApiKeyAuth { - fn auth_mode(&self) -> AuthMode { - AuthMode::ApiKey - } - fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, Option> { Box::pin(async { Ok(Some(CodexAuth::from_api_key("test-external-api-key"))) }) } @@ -137,10 +133,6 @@ impl ExternalAuth for TestExternalApiKeyAuth { struct TestUnresolvedExternalApiKeyAuth; impl ExternalAuth for TestUnresolvedExternalApiKeyAuth { - fn auth_mode(&self) -> AuthMode { - AuthMode::ApiKey - } - fn refresh( &self, _context: ExternalAuthRefreshContext, From 1b6c68b21fc2e56e61a22916261bf3550f5f3398 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 08:51:29 -0700 Subject: [PATCH 02/10] refactor: make external auth own credentials --- codex-rs/app-server/src/external_auth.rs | 98 +++++++++++++++++++ codex-rs/app-server/src/lib.rs | 1 + codex-rs/app-server/src/message_processor.rs | 80 --------------- .../request_processors/account_processor.rs | 6 +- codex-rs/login/src/auth/auth_tests.rs | 16 ++- codex-rs/login/src/auth/external_bearer.rs | 10 +- codex-rs/login/src/auth/manager.rs | 39 ++++---- codex-rs/models-manager/src/manager_tests.rs | 23 ++++- 8 files changed, 157 insertions(+), 116 deletions(-) create mode 100644 codex-rs/app-server/src/external_auth.rs diff --git a/codex-rs/app-server/src/external_auth.rs b/codex-rs/app-server/src/external_auth.rs new file mode 100644 index 00000000000..48d3499659b --- /dev/null +++ b/codex-rs/app-server/src/external_auth.rs @@ -0,0 +1,98 @@ +use std::sync::Arc; +use std::sync::RwLock; + +use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; +use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; +use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse; +use codex_app_server_protocol::ServerRequestPayload; +use codex_login::CodexAuth; +use codex_login::ExternalAuthFuture; +use codex_login::auth::ExternalAuth; +use codex_login::auth::ExternalAuthRefreshContext; +use codex_login::auth::ExternalAuthRefreshReason; +use tokio::time::Duration; +use tokio::time::timeout; + +use crate::outgoing_message::OutgoingMessageSender; + +const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Clone)] +pub(crate) struct ExternalAuthBridge { + outgoing: Arc, + auth: Arc>, +} + +impl ExternalAuthBridge { + pub(crate) fn new(outgoing: Arc, auth: CodexAuth) -> Self { + Self { + outgoing, + auth: Arc::new(RwLock::new(auth)), + } + } + + fn current_auth(&self) -> std::io::Result { + self.auth + .read() + .map(|auth| auth.clone()) + .map_err(|_| std::io::Error::other("external auth lock is poisoned")) + } + + async fn refresh(&self, context: ExternalAuthRefreshContext) -> std::io::Result { + let reason = match context.reason { + ExternalAuthRefreshReason::Unauthorized => ChatgptAuthTokensRefreshReason::Unauthorized, + }; + let params = ChatgptAuthTokensRefreshParams { + reason, + previous_account_id: context.previous_account_id, + }; + + let (request_id, rx) = self + .outgoing + .send_request(ServerRequestPayload::ChatgptAuthTokensRefresh(params)) + .await; + let result = match timeout(EXTERNAL_AUTH_REFRESH_TIMEOUT, rx).await { + Ok(result) => { + let result = result.map_err(|err| { + std::io::Error::other(format!("auth refresh request canceled: {err}")) + })?; + result.map_err(|err| { + std::io::Error::other(format!( + "auth refresh request failed: code={} message={}", + err.code, err.message + )) + })? + } + Err(_) => { + let _canceled = self.outgoing.cancel_request(&request_id).await; + return Err(std::io::Error::other(format!( + "auth refresh request timed out after {}s", + EXTERNAL_AUTH_REFRESH_TIMEOUT.as_secs() + ))); + } + }; + + let response: ChatgptAuthTokensRefreshResponse = + serde_json::from_value(result).map_err(std::io::Error::other)?; + let auth = CodexAuth::from_external_chatgpt_tokens( + response.access_token.as_str(), + response.chatgpt_account_id.as_str(), + response.chatgpt_plan_type.as_deref(), + )?; + *self + .auth + .write() + .map_err(|_| std::io::Error::other("external auth lock is poisoned"))? = auth.clone(); + Ok(auth) + } +} + +impl ExternalAuth for ExternalAuthBridge { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { self.current_auth() }) + } + + fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(ExternalAuthBridge::refresh(self, context)) + } +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index fa6b6cd811f..91da0c1645c 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -96,6 +96,7 @@ mod current_time; mod dynamic_tools; mod error_code; mod extensions; +mod external_auth; mod filters; mod fs_watch; mod fuzzy_file_search; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index f9b88f52067..c7f2df9a6a2 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -50,9 +50,6 @@ use crate::transport::AppServerTransport; use crate::transport::RemoteControlHandle; use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; -use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; -use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; -use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse; use codex_app_server_protocol::ClientNotification; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ClientResponsePayload; @@ -63,7 +60,6 @@ use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::ServerRequestPayload; use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; @@ -74,10 +70,6 @@ use codex_feedback::CodexFeedback; use codex_goal_extension::GoalService; use codex_home::CodexHomeUserInstructionsProvider; use codex_login::AuthManager; -use codex_login::CodexAuth; -use codex_login::auth::ExternalAuth; -use codex_login::auth::ExternalAuthRefreshContext; -use codex_login::auth::ExternalAuthRefreshReason; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; @@ -88,13 +80,11 @@ use tokio::sync::Semaphore; use tokio::sync::broadcast; use tokio::sync::watch; use tokio::time::Duration; -use tokio::time::timeout; use tokio_util::sync::CancellationToken; use tracing::Instrument; use crate::models_refresh_worker::ModelsRefreshWorker; -const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); const CONNECTION_RPC_DRAIN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30); fn deserialize_client_request( @@ -108,73 +98,6 @@ fn deserialize_client_request( }) } -#[derive(Clone)] -struct ExternalAuthRefreshBridge { - outgoing: Arc, -} - -impl ExternalAuthRefreshBridge { - fn map_reason(reason: ExternalAuthRefreshReason) -> ChatgptAuthTokensRefreshReason { - match reason { - ExternalAuthRefreshReason::Unauthorized => ChatgptAuthTokensRefreshReason::Unauthorized, - } - } - - async fn refresh(&self, context: ExternalAuthRefreshContext) -> std::io::Result { - let params = ChatgptAuthTokensRefreshParams { - reason: Self::map_reason(context.reason), - previous_account_id: context.previous_account_id, - }; - - let (request_id, rx) = self - .outgoing - .send_request(ServerRequestPayload::ChatgptAuthTokensRefresh(params)) - .await; - - let result = match timeout(EXTERNAL_AUTH_REFRESH_TIMEOUT, rx).await { - Ok(result) => { - // Two failure scenarios: - // 1) `oneshot::Receiver` failed (sender dropped) => request canceled/channel closed. - // 2) client answered with JSON-RPC error payload => propagate code/message. - let result = result.map_err(|err| { - std::io::Error::other(format!("auth refresh request canceled: {err}")) - })?; - result.map_err(|err| { - std::io::Error::other(format!( - "auth refresh request failed: code={} message={}", - err.code, err.message - )) - })? - } - Err(_) => { - let _canceled = self.outgoing.cancel_request(&request_id).await; - return Err(std::io::Error::other(format!( - "auth refresh request timed out after {}s", - EXTERNAL_AUTH_REFRESH_TIMEOUT.as_secs() - ))); - } - }; - - let response: ChatgptAuthTokensRefreshResponse = - serde_json::from_value(result).map_err(std::io::Error::other)?; - - CodexAuth::from_external_chatgpt_tokens( - response.access_token.as_str(), - response.chatgpt_account_id.as_str(), - response.chatgpt_plan_type.as_deref(), - ) - } -} - -impl ExternalAuth for ExternalAuthRefreshBridge { - fn refresh( - &self, - context: ExternalAuthRefreshContext, - ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { - Box::pin(ExternalAuthRefreshBridge::refresh(self, context)) - } -} - pub(crate) struct MessageProcessor { outgoing: Arc, models_refresh_worker: ModelsRefreshWorker, @@ -319,9 +242,6 @@ impl MessageProcessor { remote_control_handle, plugin_startup_tasks, } = args; - auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { - outgoing: outgoing.clone(), - })); let thread_state_manager = ThreadStateManager::new(); // The thread store is intentionally process-scoped. Config reloads can // affect per-thread behavior, but they must not move newly started, diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 97e7a18d356..5eaf6e8d33f 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,5 +1,6 @@ use super::*; use crate::auth_mode::auth_mode_to_api; +use crate::external_auth::ExternalAuthBridge; use chrono::DateTime; mod rate_limit_resets; @@ -624,7 +625,10 @@ impl AccountRequestProcessor { ) .map_err(|err| internal_error(format!("failed to set external auth: {err}")))?; self.auth_manager - .install_external_auth(auth) + .set_external_auth(Arc::new(ExternalAuthBridge::new( + Arc::clone(&self.outgoing), + auth, + ))) .await .map_err(|err| internal_error(format!("failed to set external auth: {err}")))?; self.config_manager.replace_cloud_config_bundle_loader( diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index cdc680c7511..380a17abcf1 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1015,9 +1015,13 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None); } -struct RefreshOnlyExternalChatgptAuth; +struct TestExternalChatgptAuth(CodexAuth); + +impl ExternalAuth for TestExternalChatgptAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.clone()) }) + } -impl ExternalAuth for RefreshOnlyExternalChatgptAuth { fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { Box::pin(async { Err(std::io::Error::other("refresh should not be called")) }) } @@ -1025,7 +1029,7 @@ impl ExternalAuth for RefreshOnlyExternalChatgptAuth { #[tokio::test] #[serial(codex_auth_env)] -async fn external_chatgpt_auth_resolves_from_the_manager_cache() { +async fn external_chatgpt_auth_resolves_from_its_provider() { let codex_home = tempdir().unwrap(); let params = AuthFileParams { openai_api_key: None, @@ -1046,8 +1050,10 @@ async fn external_chatgpt_auth_resolves_from_the_manager_cache() { /*auth_route_config*/ None, ) .await; - manager.set_external_auth(Arc::new(RefreshOnlyExternalChatgptAuth)); - manager.install_external_auth(auth).await.unwrap(); + manager + .set_external_auth(Arc::new(TestExternalChatgptAuth(auth))) + .await + .unwrap(); let auth = manager.auth().await.expect("external auth should resolve"); diff --git a/codex-rs/login/src/auth/external_bearer.rs b/codex-rs/login/src/auth/external_bearer.rs index 7e8a47e22fa..0a276543eab 100644 --- a/codex-rs/login/src/auth/external_bearer.rs +++ b/codex-rs/login/src/auth/external_bearer.rs @@ -29,7 +29,7 @@ impl BearerTokenRefresher { clippy::await_holding_invalid_type, reason = "external bearer cache misses intentionally hold cached_token across the provider command to avoid duplicate refreshes" )] - async fn resolve(&self) -> io::Result> { + async fn resolve(&self) -> io::Result { let access_token = { let mut cached = self.state.cached_token.lock().await; if let Some(cached_token) = cached.as_ref() { @@ -38,9 +38,7 @@ impl BearerTokenRefresher { None => true, }; if should_use_cached_token { - return Ok(Some(CodexAuth::from_api_key( - cached_token.access_token.as_str(), - ))); + return Ok(CodexAuth::from_api_key(cached_token.access_token.as_str())); } } @@ -51,7 +49,7 @@ impl BearerTokenRefresher { }); access_token }; - Ok(Some(CodexAuth::from_api_key(access_token.as_str()))) + Ok(CodexAuth::from_api_key(access_token.as_str())) } async fn refresh(&self, _context: ExternalAuthRefreshContext) -> io::Result { @@ -66,7 +64,7 @@ impl BearerTokenRefresher { } impl ExternalAuth for BearerTokenRefresher { - fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { Box::pin(BearerTokenRefresher::resolve(self)) } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index ac8d20870d3..fa03168a766 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -211,16 +211,12 @@ pub struct ExternalAuthRefreshContext { /// Pluggable auth provider used by `AuthManager` for externally managed auth flows. /// -/// Implementations may either resolve auth eagerly via `resolve()` or provide refreshed -/// auth on demand via `refresh()`. +/// Implementations own the current auth value and any source-specific refresh mechanism. pub trait ExternalAuth: Send + Sync { - /// Returns cached or immediately available auth, if this provider owns the current value. - /// Returning `None` lets the manager use externally installed ChatGPT auth from its cache. - fn resolve(&self) -> ExternalAuthFuture<'_, Option> { - Box::pin(async { Ok(None) }) - } + /// Returns the provider's current auth value. + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth>; - /// Refreshes auth in response to a manager-driven refresh attempt. + /// Refreshes auth and makes the returned value current for future `resolve()` calls. fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth>; } @@ -2214,11 +2210,19 @@ impl AuthManager { } } - pub fn set_external_auth(&self, external_auth: Arc) { + pub async fn set_external_auth( + &self, + external_auth: Arc, + ) -> Result<(), RefreshTokenError> { + let auth = external_auth + .resolve() + .await + .map_err(RefreshTokenError::Transient)?; + self.apply_external_auth(auth).await?; if let Ok(mut guard) = self.external_auth.write() { *guard = Some(external_auth); } - self.set_resolved_external_auth(None); + Ok(()) } pub fn clear_external_auth(&self) { @@ -2334,17 +2338,13 @@ impl AuthManager { let external_auth = self.external_auth()?; match external_auth.resolve().await { - Ok(Some(auth)) => { - if let Err(err) = self.install_external_auth(auth.clone()).await { + Ok(auth) => { + if let Err(err) = self.apply_external_auth(auth.clone()).await { tracing::error!("Failed to install resolved external auth: {err}"); return None; } Some(auth) } - Ok(None) => { - self.set_resolved_external_auth(None); - None - } Err(err) => { self.set_resolved_external_auth(None); tracing::error!("Failed to resolve external auth: {err}"); @@ -2456,6 +2456,7 @@ impl AuthManager { self.keyring_backend_kind, )?; // Always reload to clear any cached auth (even if file absent). + self.clear_external_auth(); self.reload().await; Ok(removed) } @@ -2475,6 +2476,7 @@ impl AuthManager { self.keyring_backend_kind, )?; // Always reload to clear any cached auth (even if file absent). + self.clear_external_auth(); self.reload().await; Ok(result) } @@ -2540,11 +2542,10 @@ impl AuthManager { .refresh(context) .await .map_err(RefreshTokenError::Transient)?; - self.install_external_auth(refreshed).await + self.apply_external_auth(refreshed).await } - /// Validates and installs auth returned by an [`ExternalAuth`] provider. - pub async fn install_external_auth(&self, auth: CodexAuth) -> Result<(), RefreshTokenError> { + async fn apply_external_auth(&self, auth: CodexAuth) -> Result<(), RefreshTokenError> { if auth.api_auth_mode() == AuthMode::ApiKey { self.set_resolved_external_auth(Some(auth)); return Ok(()); diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index fc8c100a4f6..f1678fe9f25 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -117,8 +117,8 @@ impl TestModelsEndpoint { struct TestExternalApiKeyAuth; impl ExternalAuth for TestExternalApiKeyAuth { - fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, Option> { - Box::pin(async { Ok(Some(CodexAuth::from_api_key("test-external-api-key"))) }) + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(CodexAuth::from_api_key("test-external-api-key")) }) } fn refresh( @@ -133,6 +133,10 @@ impl ExternalAuth for TestExternalApiKeyAuth { struct TestUnresolvedExternalApiKeyAuth; impl ExternalAuth for TestUnresolvedExternalApiKeyAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Err(std::io::Error::other("unresolved test auth")) }) + } + fn refresh( &self, _context: ExternalAuthRefreshContext, @@ -850,7 +854,10 @@ async fn refresh_available_models_skips_network_when_external_api_key_overrides_ let codex_home = tempdir().expect("temp dir"); let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - auth_manager.set_external_auth(Arc::new(TestExternalApiKeyAuth)); + auth_manager + .set_external_auth(Arc::new(TestExternalApiKeyAuth)) + .await + .expect("external API key auth should resolve"); let endpoint = TestAuthAwareModelsEndpoint::new( Some(Arc::clone(&auth_manager)), vec![vec![remote_model( @@ -890,7 +897,10 @@ async fn refresh_available_models_uses_cached_chatgpt_when_external_api_key_is_u let codex_home = tempdir().expect("temp dir"); let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - auth_manager.set_external_auth(Arc::new(TestUnresolvedExternalApiKeyAuth)); + auth_manager + .set_external_auth(Arc::new(TestUnresolvedExternalApiKeyAuth)) + .await + .expect_err("unresolved external auth should be rejected"); let endpoint = TestAuthAwareModelsEndpoint::new( Some(Arc::clone(&auth_manager)), vec![vec![remote_model( @@ -1005,7 +1015,10 @@ async fn static_manager_reads_latest_auth_mode() { vec!["chatgpt-only", "api-model"] ); - auth_manager.set_external_auth(Arc::new(TestExternalApiKeyAuth)); + auth_manager + .set_external_auth(Arc::new(TestExternalApiKeyAuth)) + .await + .expect("external API key auth should resolve"); let api_models = manager.list_models(RefreshStrategy::Online).await; assert_eq!( From 77f76f151fdde5978ca0a4d8b7402d1ca12e6568 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 08:54:18 -0700 Subject: [PATCH 03/10] fix: retain message processor timeout import --- codex-rs/app-server/src/message_processor.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index c7f2df9a6a2..ca122920bd6 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -80,6 +80,7 @@ use tokio::sync::Semaphore; use tokio::sync::broadcast; use tokio::sync::watch; use tokio::time::Duration; +use tokio::time::timeout; use tokio_util::sync::CancellationToken; use tracing::Instrument; From 544653ff3b48e777fbf0247437ceb9a6e114a417 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 09:04:46 -0700 Subject: [PATCH 04/10] refactor: tighten external auth state --- codex-rs/app-server/src/external_auth.rs | 19 ++-- codex-rs/app-server/tests/suite/v2/account.rs | 21 ++++ codex-rs/login/src/auth/auth_tests.rs | 46 --------- codex-rs/login/src/auth/manager.rs | 98 ++++++++++--------- codex-rs/models-manager/src/manager.rs | 3 - 5 files changed, 82 insertions(+), 105 deletions(-) diff --git a/codex-rs/app-server/src/external_auth.rs b/codex-rs/app-server/src/external_auth.rs index 48d3499659b..d00777f7101 100644 --- a/codex-rs/app-server/src/external_auth.rs +++ b/codex-rs/app-server/src/external_auth.rs @@ -17,27 +17,19 @@ use crate::outgoing_message::OutgoingMessageSender; const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); -#[derive(Clone)] pub(crate) struct ExternalAuthBridge { outgoing: Arc, - auth: Arc>, + auth: RwLock, } impl ExternalAuthBridge { pub(crate) fn new(outgoing: Arc, auth: CodexAuth) -> Self { Self { outgoing, - auth: Arc::new(RwLock::new(auth)), + auth: RwLock::new(auth), } } - fn current_auth(&self) -> std::io::Result { - self.auth - .read() - .map(|auth| auth.clone()) - .map_err(|_| std::io::Error::other("external auth lock is poisoned")) - } - async fn refresh(&self, context: ExternalAuthRefreshContext) -> std::io::Result { let reason = match context.reason { ExternalAuthRefreshReason::Unauthorized => ChatgptAuthTokensRefreshReason::Unauthorized, @@ -89,7 +81,12 @@ impl ExternalAuthBridge { impl ExternalAuth for ExternalAuthBridge { fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { - Box::pin(async { self.current_auth() }) + Box::pin(async { + self.auth + .read() + .map(|auth| auth.clone()) + .map_err(|_| std::io::Error::other("external auth lock is poisoned")) + }) } fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 3cd98e2e52d..25106047a81 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -327,6 +327,27 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { } ); + let logout_id = mcp.send_logout_account_request().await?; + let logout_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(logout_id)), + ) + .await??; + let _: LogoutAccountResponse = to_response(logout_resp)?; + + let get_id = mcp + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let get_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(get_id)), + ) + .await??; + let account: GetAccountResponse = to_response(get_resp)?; + assert_eq!(account.account, None); + Ok(()) } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 380a17abcf1..77d8c4fdfb2 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1015,52 +1015,6 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None); } -struct TestExternalChatgptAuth(CodexAuth); - -impl ExternalAuth for TestExternalChatgptAuth { - fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { - Box::pin(async { Ok(self.0.clone()) }) - } - - fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { - Box::pin(async { Err(std::io::Error::other("refresh should not be called")) }) - } -} - -#[tokio::test] -#[serial(codex_auth_env)] -async fn external_chatgpt_auth_resolves_from_its_provider() { - let codex_home = tempdir().unwrap(); - let params = AuthFileParams { - openai_api_key: None, - chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), - }; - let access_token = fake_jwt_for_auth_file_params(¶ms).unwrap(); - let auth = - CodexAuth::from_external_chatgpt_tokens(&access_token, WORKSPACE_ID_ALLOWED, Some("pro")) - .unwrap(); - let manager = AuthManager::shared( - codex_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*forced_chatgpt_workspace_id*/ None, - /*chatgpt_base_url*/ None, - AuthKeyringBackendKind::default(), - /*auth_route_config*/ None, - ) - .await; - manager - .set_external_auth(Arc::new(TestExternalChatgptAuth(auth))) - .await - .unwrap(); - - let auth = manager.auth().await.expect("external auth should resolve"); - - assert_eq!(auth.api_auth_mode(), AuthMode::ChatgptAuthTokens); - assert_eq!(auth.get_account_id().as_deref(), Some(WORKSPACE_ID_ALLOWED)); -} - #[tokio::test] async fn external_bearer_only_auth_manager_uses_cached_provider_token() { let script = ProviderAuthScript::new(&["provider-token", "next-token"]).unwrap(); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index fa03168a766..86233361023 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1492,6 +1492,11 @@ struct CachedAuth { permanent_refresh_failure: Option, } +struct ExternalAuthState { + provider: Arc, + last_resolved_auth: Option, +} + #[derive(Clone)] struct AuthScopedRefreshFailure { auth: CodexAuth, @@ -1755,8 +1760,7 @@ pub struct AuthManager { refresh_lock: Semaphore, agent_identity_lock: Semaphore, agent_identity_bootstrap_cooldown: Mutex, - external_auth: RwLock>>, - resolved_external_auth: RwLock>, + external_auth: RwLock>, auth_route_config: Option, } @@ -1859,7 +1863,6 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), - resolved_external_auth: RwLock::new(None), auth_route_config, } } @@ -1886,7 +1889,6 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), - resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } @@ -1912,7 +1914,6 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), - resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } @@ -1946,7 +1947,6 @@ impl AuthManager { agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), - resolved_external_auth: RwLock::new(None), auth_route_config: None, }) } @@ -1969,17 +1969,24 @@ impl AuthManager { refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), - external_auth: RwLock::new(Some( - Arc::new(BearerTokenRefresher::new(config)) as Arc - )), - resolved_external_auth: RwLock::new(None), + external_auth: RwLock::new(Some(ExternalAuthState { + provider: Arc::new(BearerTokenRefresher::new(config)) as Arc, + last_resolved_auth: None, + })), auth_route_config: None, }) } /// Current cached auth (clone) without attempting a refresh. pub fn auth_cached(&self) -> Option { - self.resolved_external_auth() + self.external_auth + .read() + .ok() + .and_then(|state| { + state + .as_ref() + .and_then(|state| state.last_resolved_auth.clone()) + }) .or_else(|| self.inner.read().ok().and_then(|c| c.auth.clone())) } @@ -2218,10 +2225,14 @@ impl AuthManager { .resolve() .await .map_err(RefreshTokenError::Transient)?; - self.apply_external_auth(auth).await?; - if let Ok(mut guard) = self.external_auth.write() { - *guard = Some(external_auth); - } + self.apply_external_auth(&auth).await?; + let mut state = self.external_auth.write().map_err(|_| { + RefreshTokenError::Transient(std::io::Error::other("external auth lock is poisoned")) + })?; + *state = Some(ExternalAuthState { + provider: external_auth, + last_resolved_auth: Some(auth), + }); Ok(()) } @@ -2229,7 +2240,6 @@ impl AuthManager { if let Ok(mut guard) = self.external_auth.write() { *guard = None; } - self.set_resolved_external_auth(None); } pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option>) { @@ -2310,28 +2320,25 @@ impl AuthManager { self.external_auth .read() .ok() - .and_then(|guard| guard.as_ref().cloned()) - } - - fn resolved_external_auth(&self) -> Option { - self.resolved_external_auth - .read() - .ok() - .and_then(|guard| guard.clone()) + .and_then(|state| state.as_ref().map(|state| Arc::clone(&state.provider))) } - fn set_resolved_external_auth(&self, auth: Option) { - if let Ok(mut guard) = self.resolved_external_auth.write() { - *guard = auth; + fn cache_external_auth(&self, provider: &Arc, auth: Option) { + if let Ok(mut state) = self.external_auth.write() + && let Some(state) = state.as_mut() + && Arc::ptr_eq(&state.provider, provider) + { + state.last_resolved_auth = auth; } } fn has_external_api_key_auth(&self) -> bool { - self.has_external_auth() - && self - .resolved_external_auth() + self.external_auth.read().ok().is_some_and(|state| { + state .as_ref() + .and_then(|state| state.last_resolved_auth.as_ref()) .is_some_and(CodexAuth::is_api_key_auth) + }) } async fn resolve_external_auth(&self) -> Option { @@ -2339,14 +2346,15 @@ impl AuthManager { match external_auth.resolve().await { Ok(auth) => { - if let Err(err) = self.apply_external_auth(auth.clone()).await { + if let Err(err) = self.apply_external_auth(&auth).await { tracing::error!("Failed to install resolved external auth: {err}"); return None; } + self.cache_external_auth(&external_auth, Some(auth.clone())); Some(auth) } Err(err) => { - self.set_resolved_external_auth(None); + self.cache_external_auth(&external_auth, None); tracing::error!("Failed to resolve external auth: {err}"); None } @@ -2542,12 +2550,13 @@ impl AuthManager { .refresh(context) .await .map_err(RefreshTokenError::Transient)?; - self.apply_external_auth(refreshed).await + self.apply_external_auth(&refreshed).await?; + self.cache_external_auth(&external_auth, Some(refreshed)); + Ok(()) } - async fn apply_external_auth(&self, auth: CodexAuth) -> Result<(), RefreshTokenError> { + async fn apply_external_auth(&self, auth: &CodexAuth) -> Result<(), RefreshTokenError> { if auth.api_auth_mode() == AuthMode::ApiKey { - self.set_resolved_external_auth(Some(auth)); return Ok(()); } if !auth.is_external_chatgpt_tokens() { @@ -2558,14 +2567,6 @@ impl AuthManager { ), ))); } - if self - .inner - .read() - .ok() - .is_some_and(|cached| Self::auths_equal_for_refresh(cached.auth.as_ref(), Some(&auth))) - { - return Ok(()); - } let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); let account_id = auth.get_account_id().ok_or_else(|| { RefreshTokenError::Transient(std::io::Error::other( @@ -2577,10 +2578,18 @@ impl AuthManager { { return Err(RefreshTokenError::Transient(std::io::Error::other( format!( - "external auth refresh returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" + "external auth returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" ), ))); } + if self + .inner + .read() + .ok() + .is_some_and(|cached| Self::auths_equal_for_refresh(cached.auth.as_ref(), Some(auth))) + { + return Ok(()); + } let auth_dot_json = auth.get_current_auth_json().ok_or_else(|| { RefreshTokenError::Transient(std::io::Error::other( "external ChatGPT auth tokens are missing auth state", @@ -2594,7 +2603,6 @@ impl AuthManager { ) .map_err(RefreshTokenError::Transient)?; self.reload().await; - self.set_resolved_external_auth(None); Ok(()) } diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 9ab09469cc4..67b939051b6 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -86,9 +86,6 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { ) -> ModelsManagerFuture<'_, Vec> { Box::pin( async move { - if let Some(auth_manager) = self.auth_manager() { - let _ = auth_manager.auth().await; - } let catalog = self.raw_model_catalog(refresh_strategy).await; self.build_available_models(catalog.models) } From 8393b8ec85bd87d2bd7306797f7cba24367918e6 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 09:22:31 -0700 Subject: [PATCH 05/10] refactor: commit external auth in one step --- codex-rs/login/src/auth/manager.rs | 139 ++++++++++++++--------------- 1 file changed, 67 insertions(+), 72 deletions(-) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 86233361023..deb8f5c9887 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -2225,14 +2225,7 @@ impl AuthManager { .resolve() .await .map_err(RefreshTokenError::Transient)?; - self.apply_external_auth(&auth).await?; - let mut state = self.external_auth.write().map_err(|_| { - RefreshTokenError::Transient(std::io::Error::other("external auth lock is poisoned")) - })?; - *state = Some(ExternalAuthState { - provider: external_auth, - last_resolved_auth: Some(auth), - }); + self.commit_external_auth(&external_auth, &auth).await?; Ok(()) } @@ -2323,15 +2316,6 @@ impl AuthManager { .and_then(|state| state.as_ref().map(|state| Arc::clone(&state.provider))) } - fn cache_external_auth(&self, provider: &Arc, auth: Option) { - if let Ok(mut state) = self.external_auth.write() - && let Some(state) = state.as_mut() - && Arc::ptr_eq(&state.provider, provider) - { - state.last_resolved_auth = auth; - } - } - fn has_external_api_key_auth(&self) -> bool { self.external_auth.read().ok().is_some_and(|state| { state @@ -2345,16 +2329,19 @@ impl AuthManager { let external_auth = self.external_auth()?; match external_auth.resolve().await { - Ok(auth) => { - if let Err(err) = self.apply_external_auth(&auth).await { + Ok(auth) => match self.commit_external_auth(&external_auth, &auth).await { + Ok(()) => Some(auth), + Err(err) => { tracing::error!("Failed to install resolved external auth: {err}"); - return None; + None } - self.cache_external_auth(&external_auth, Some(auth.clone())); - Some(auth) - } + }, Err(err) => { - self.cache_external_auth(&external_auth, None); + if let Ok(mut state) = self.external_auth.write() + && let Some(state) = state.as_mut() + { + state.last_resolved_auth = None; + } tracing::error!("Failed to resolve external auth: {err}"); None } @@ -2550,59 +2537,67 @@ impl AuthManager { .refresh(context) .await .map_err(RefreshTokenError::Transient)?; - self.apply_external_auth(&refreshed).await?; - self.cache_external_auth(&external_auth, Some(refreshed)); + self.commit_external_auth(&external_auth, &refreshed) + .await?; Ok(()) } - async fn apply_external_auth(&self, auth: &CodexAuth) -> Result<(), RefreshTokenError> { - if auth.api_auth_mode() == AuthMode::ApiKey { - return Ok(()); - } - if !auth.is_external_chatgpt_tokens() { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth returned unsupported auth mode {:?}", - auth.api_auth_mode() - ), - ))); - } - let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); - let account_id = auth.get_account_id().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other( - "external ChatGPT auth tokens are missing an account id", - )) - })?; - if let Some(expected_workspace_ids) = forced_chatgpt_workspace_id.as_deref() - && !expected_workspace_ids.contains(&account_id) - { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" - ), - ))); - } - if self - .inner - .read() - .ok() - .is_some_and(|cached| Self::auths_equal_for_refresh(cached.auth.as_ref(), Some(auth))) - { - return Ok(()); + async fn commit_external_auth( + &self, + provider: &Arc, + auth: &CodexAuth, + ) -> Result<(), RefreshTokenError> { + if auth.api_auth_mode() != AuthMode::ApiKey { + if !auth.is_external_chatgpt_tokens() { + return Err(RefreshTokenError::Transient(std::io::Error::other( + format!( + "external auth returned unsupported auth mode {:?}", + auth.api_auth_mode() + ), + ))); + } + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); + let account_id = auth.get_account_id().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other( + "external ChatGPT auth tokens are missing an account id", + )) + })?; + if let Some(expected_workspace_ids) = forced_chatgpt_workspace_id.as_deref() + && !expected_workspace_ids.contains(&account_id) + { + return Err(RefreshTokenError::Transient(std::io::Error::other( + format!( + "external auth returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" + ), + ))); + } + let is_cached = self.inner.read().ok().is_some_and(|cached| { + Self::auths_equal_for_refresh(cached.auth.as_ref(), Some(auth)) + }); + if !is_cached { + let auth_dot_json = auth.get_current_auth_json().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other( + "external ChatGPT auth tokens are missing auth state", + )) + })?; + save_auth( + &self.codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + ) + .map_err(RefreshTokenError::Transient)?; + self.reload().await; + } } - let auth_dot_json = auth.get_current_auth_json().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other( - "external ChatGPT auth tokens are missing auth state", - )) + + let mut state = self.external_auth.write().map_err(|_| { + RefreshTokenError::Transient(std::io::Error::other("external auth lock is poisoned")) })?; - save_auth( - &self.codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - AuthKeyringBackendKind::default(), - ) - .map_err(RefreshTokenError::Transient)?; - self.reload().await; + *state = Some(ExternalAuthState { + provider: Arc::clone(provider), + last_resolved_auth: Some(auth.clone()), + }); Ok(()) } From 9f8d177f3cae47181ec8a9ada0eb069d6ac297b5 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 09:36:32 -0700 Subject: [PATCH 06/10] fix: preserve external auth 401 recovery --- codex-rs/login/src/auth/auth_tests.rs | 2 +- codex-rs/login/src/auth/manager.rs | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 77d8c4fdfb2..0b0bfc882d1 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1069,11 +1069,11 @@ async fn unauthorized_recovery_uses_external_refresh_for_bearer_manager() { let mut auth_config = script.auth_config(); auth_config.refresh_interval_ms = 0; let manager = AuthManager::external_bearer_only(auth_config); + let mut recovery = manager.unauthorized_recovery(); let initial_token = manager .auth() .await .and_then(|auth| auth.api_key().map(str::to_string)); - let mut recovery = manager.unauthorized_recovery(); assert!(recovery.has_next()); assert_eq!(recovery.mode_name(), "external"); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index deb8f5c9887..043de5c99d6 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1583,11 +1583,7 @@ impl UnauthorizedRecovery { fn new(manager: Arc) -> Self { let cached_auth = manager.auth_cached(); let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id); - let mode = if manager.has_external_api_key_auth() - || cached_auth - .as_ref() - .is_some_and(CodexAuth::is_external_chatgpt_tokens) - { + let mode = if manager.has_external_auth() { UnauthorizedRecoveryMode::External } else { UnauthorizedRecoveryMode::Managed From ce1f9e453b76731998b5c95dca44e41ea8637f06 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 09:51:03 -0700 Subject: [PATCH 07/10] test: exercise external auth refresh in cloud config --- codex-rs/cloud-config/src/service_tests.rs | 147 +++++++++++---------- 1 file changed, 80 insertions(+), 67 deletions(-) diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index dfd4fe50ced..0896b4628cc 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -19,11 +19,14 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; use codex_login::auth::AgentIdentityAuth; use codex_login::auth::AgentIdentityAuthRecord; +use codex_login::auth::ExternalAuth; +use codex_login::auth::ExternalAuthRefreshContext; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::VecDeque; use std::future::pending; use std::path::Path; +use std::sync::RwLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::tempdir; @@ -142,26 +145,20 @@ fn chatgpt_auth_json_with_last_refresh( refresh_token: &str, last_refresh: &str, ) -> serde_json::Value { - chatgpt_auth_json_with_mode( - plan_type, - chatgpt_user_id, - account_id, - access_token, - refresh_token, - last_refresh, - /*auth_mode*/ None, - ) + let fake_jwt = fake_chatgpt_jwt(plan_type, chatgpt_user_id, b"sig"); + json!({ + "OPENAI_API_KEY": null, + "tokens": { + "id_token": fake_jwt, + "access_token": access_token, + "refresh_token": refresh_token, + "account_id": account_id, + }, + "last_refresh": last_refresh, + }) } -fn chatgpt_auth_json_with_mode( - plan_type: &str, - chatgpt_user_id: Option<&str>, - account_id: Option<&str>, - access_token: &str, - refresh_token: &str, - last_refresh: &str, - auth_mode: Option<&str>, -) -> serde_json::Value { +fn fake_chatgpt_jwt(plan_type: &str, chatgpt_user_id: Option<&str>, signature: &[u8]) -> String { let header = json!({ "alg": "none", "typ": "JWT" }); let auth_payload = json!({ "chatgpt_plan_type": plan_type, @@ -174,23 +171,8 @@ fn chatgpt_auth_json_with_mode( }); let header_b64 = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).expect("header")); let payload_b64 = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).expect("payload")); - let signature_b64 = URL_SAFE_NO_PAD.encode(b"sig"); - let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - - let mut auth_json = json!({ - "OPENAI_API_KEY": null, - "tokens": { - "id_token": fake_jwt, - "access_token": access_token, - "refresh_token": refresh_token, - "account_id": account_id, - }, - "last_refresh": last_refresh, - }); - if let Some(auth_mode) = auth_mode { - auth_json["auth_mode"] = serde_json::Value::String(auth_mode.to_string()); - } - auth_json + let signature_b64 = URL_SAFE_NO_PAD.encode(signature); + format!("{header_b64}.{payload_b64}.{signature_b64}") } fn test_bundle() -> CloudConfigBundle { @@ -319,6 +301,39 @@ struct UnauthorizedBundleClient { request_count: AtomicUsize, } +struct TestExternalChatgptAuth { + current: RwLock, + refreshed: CodexAuth, + refresh_count: AtomicUsize, +} + +impl ExternalAuth for TestExternalChatgptAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + self.current + .read() + .map(|auth| auth.clone()) + .map_err(|_| std::io::Error::other("external auth lock is poisoned")) + }) + } + + fn refresh( + &self, + _context: ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + let refreshed = self.refreshed.clone(); + *self + .current + .write() + .map_err(|_| std::io::Error::other("external auth lock is poisoned"))? = + refreshed.clone(); + self.refresh_count.fetch_add(1, Ordering::SeqCst); + Ok(refreshed) + }) + } +} + impl BundleClient for UnauthorizedBundleClient { async fn get_bundle(&self, _auth: &CodexAuth) -> Result { self.request_count.fetch_add(1, Ordering::SeqCst); @@ -855,26 +870,13 @@ async fn get_bundle_surfaces_auth_recovery_message() { } #[tokio::test] -async fn get_bundle_unauthorized_without_recovery_uses_generic_message() { +async fn get_bundle_refreshes_external_auth_after_unauthorized() { let auth_home = tempdir().expect("tempdir"); - write_auth_json( - auth_home.path(), - chatgpt_auth_json_with_mode( - "enterprise", - Some("user-12345"), - Some("account-12345"), - "test-access-token", - "test-refresh-token", - "2025-01-01T00:00:00Z", - Some("chatgptAuthTokens"), - ), - ) - .expect("write auth"); let auth_manager = Arc::new( AuthManager::new( auth_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, + AuthCredentialsStoreMode::Ephemeral, /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), @@ -882,11 +884,32 @@ async fn get_bundle_unauthorized_without_recovery_uses_generic_message() { ) .await, ); + let initial_auth = CodexAuth::from_external_chatgpt_tokens( + &fake_chatgpt_jwt("enterprise", Some("user-12345"), b"initial"), + "account-12345", + Some("enterprise"), + ) + .expect("initial external auth"); + let refreshed_token = fake_chatgpt_jwt("enterprise", Some("user-12345"), b"refreshed"); + let refreshed_auth = CodexAuth::from_external_chatgpt_tokens( + &refreshed_token, + "account-12345", + Some("enterprise"), + ) + .expect("refreshed external auth"); + let external_auth = Arc::new(TestExternalChatgptAuth { + current: RwLock::new(initial_auth), + refreshed: refreshed_auth, + refresh_count: AtomicUsize::new(0), + }); + auth_manager + .set_external_auth(external_auth.clone()) + .await + .expect("set external auth"); - let fetcher = Arc::new(UnauthorizedBundleClient { - message: - "GET https://chatgpt.com/backend-api/wham/config/bundle failed: 401; content-type=text/html; body=nope" - .to_string(), + let fetcher = Arc::new(TokenBundleClient { + expected_token: refreshed_token, + bundle: test_bundle(), request_count: AtomicUsize::new(0), }); let codex_home = tempdir().expect("tempdir"); @@ -897,19 +920,9 @@ async fn get_bundle_unauthorized_without_recovery_uses_generic_message() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let err = service - .load_startup_bundle() - .await - .expect_err("cloud config bundle should fail closed"); - assert_eq!( - err, - CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Auth, - Some(401), - CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE, - ) - ); - assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); + assert_eq!(service.load_startup_bundle().await, Ok(Some(test_bundle()))); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 2); + assert_eq!(external_auth.refresh_count.load(Ordering::SeqCst), 1); } #[tokio::test] From b31594dc16891c255a85c9b6429a28e0939b30a5 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 13:29:17 -0700 Subject: [PATCH 08/10] Simplify external auth ownership --- codex-rs/login/src/auth/manager.rs | 237 +++++++++++------------------ 1 file changed, 87 insertions(+), 150 deletions(-) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 043de5c99d6..c8a77efde88 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1492,11 +1492,6 @@ struct CachedAuth { permanent_refresh_failure: Option, } -struct ExternalAuthState { - provider: Arc, - last_resolved_auth: Option, -} - #[derive(Clone)] struct AuthScopedRefreshFailure { auth: CodexAuth, @@ -1553,14 +1548,9 @@ 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 auth sources, UnauthorizedRecovery retries once. -// -// - External ChatGPT auth tokens (`chatgptAuthTokens`) are refreshed by asking -// the parent app for new tokens through the configured -// `ExternalAuth`, persisting them in the ephemeral auth store, and -// reloading the cached auth snapshot. -// - External bearer auth sources for custom model providers rerun the provider -// auth command without touching disk. +// For external auth sources, UnauthorizedRecovery retries once by asking the +// configured provider to refresh and caching the returned auth through the same +// path used by other auth sources. pub struct UnauthorizedRecovery { manager: Arc, step: UnauthorizedRecoveryStep, @@ -1719,9 +1709,7 @@ impl UnauthorizedRecovery { }); } UnauthorizedRecoveryStep::ExternalRefresh => { - self.manager - .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) - .await?; + self.manager.refresh_token_from_authority().await?; self.step = UnauthorizedRecoveryStep::Done; return Ok(UnauthorizedRecoveryStepResult { auth_state_changed: Some(true), @@ -1756,7 +1744,7 @@ pub struct AuthManager { refresh_lock: Semaphore, agent_identity_lock: Semaphore, agent_identity_bootstrap_cooldown: Mutex, - external_auth: RwLock>, + external_auth: RwLock>>, auth_route_config: Option, } @@ -1965,25 +1953,19 @@ impl AuthManager { refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), - external_auth: RwLock::new(Some(ExternalAuthState { - provider: Arc::new(BearerTokenRefresher::new(config)) as Arc, - last_resolved_auth: None, - })), + external_auth: RwLock::new(Some( + Arc::new(BearerTokenRefresher::new(config)) as Arc + )), auth_route_config: None, }) } /// Current cached auth (clone) without attempting a refresh. pub fn auth_cached(&self) -> Option { - self.external_auth + self.inner .read() .ok() - .and_then(|state| { - state - .as_ref() - .and_then(|state| state.last_resolved_auth.clone()) - }) - .or_else(|| self.inner.read().ok().and_then(|c| c.auth.clone())) + .and_then(|cached| cached.auth.clone()) } /// Subscribes to cached auth changes that can affect request recovery. @@ -2006,8 +1988,9 @@ impl AuthManager { /// a guarded reload and then refreshes only if the on-disk auth is unchanged. #[instrument(level = "trace", skip_all)] pub async fn auth(&self) -> Option { - if let Some(auth) = self.resolve_external_auth().await { - return Some(auth); + if self.has_external_auth() { + self.reload().await; + return self.auth_cached(); } let auth = self.auth_cached()?; @@ -2083,11 +2066,10 @@ impl AuthManager { .await } - /// Force a reload of the auth information from auth.json. Returns - /// whether the auth value changed. + /// Reloads auth from the active source. Returns whether the auth value changed. pub async fn reload(&self) -> bool { tracing::info!("Reloading auth"); - let new_auth = self.load_auth_from_storage().await; + let new_auth = self.load_auth().await; self.set_cached_auth(new_auth) } @@ -2103,7 +2085,7 @@ impl AuthManager { } }; - let new_auth = self.load_auth_from_storage().await; + let new_auth = self.load_auth().await; let new_account_id = new_auth.as_ref().and_then(CodexAuth::get_account_id); if new_account_id.as_deref() != Some(expected_account_id) { @@ -2176,7 +2158,17 @@ impl AuthManager { } } - async fn load_auth_from_storage(&self) -> Option { + async fn load_auth(&self) -> Option { + if let Some(external_auth) = self.external_auth() { + return match self.resolve_external_auth(&external_auth).await { + Ok(auth) => Some(auth), + Err(err) => { + tracing::error!("Failed to resolve external auth: {err}"); + None + } + }; + } + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); load_auth( &self.codex_home, @@ -2217,17 +2209,19 @@ impl AuthManager { &self, external_auth: Arc, ) -> Result<(), RefreshTokenError> { - let auth = external_auth - .resolve() - .await - .map_err(RefreshTokenError::Transient)?; - self.commit_external_auth(&external_auth, &auth).await?; + let auth = self.resolve_external_auth(&external_auth).await?; + *self.external_auth.write().map_err(|_| { + RefreshTokenError::Transient(std::io::Error::other("external auth lock is poisoned")) + })? = Some(external_auth); + self.set_cached_auth(Some(auth)); Ok(()) } pub fn clear_external_auth(&self) { - if let Ok(mut guard) = self.external_auth.write() { - *guard = None; + if let Ok(mut external_auth) = self.external_auth.write() + && external_auth.take().is_some() + { + self.set_cached_auth(None); } } @@ -2309,46 +2303,33 @@ impl AuthManager { self.external_auth .read() .ok() - .and_then(|state| state.as_ref().map(|state| Arc::clone(&state.provider))) + .and_then(|external_auth| external_auth.as_ref().map(Arc::clone)) } fn has_external_api_key_auth(&self) -> bool { - self.external_auth.read().ok().is_some_and(|state| { - state + self.has_external_auth() + && self + .auth_cached() .as_ref() - .and_then(|state| state.last_resolved_auth.as_ref()) .is_some_and(CodexAuth::is_api_key_auth) - }) } - async fn resolve_external_auth(&self) -> Option { - let external_auth = self.external_auth()?; - - match external_auth.resolve().await { - Ok(auth) => match self.commit_external_auth(&external_auth, &auth).await { - Ok(()) => Some(auth), - Err(err) => { - tracing::error!("Failed to install resolved external auth: {err}"); - None - } - }, - Err(err) => { - if let Ok(mut state) = self.external_auth.write() - && let Some(state) = state.as_mut() - { - state.last_resolved_auth = None; - } - tracing::error!("Failed to resolve external auth: {err}"); - None - } - } + async fn resolve_external_auth( + &self, + external_auth: &Arc, + ) -> Result { + let auth = external_auth + .resolve() + .await + .map_err(RefreshTokenError::Transient)?; + self.validate_external_auth(&auth)?; + Ok(auth) } - /// 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 - /// can assume that some other instance already refreshed it. If the persisted - /// token is the same as the cached, then ask the token authority to refresh. + /// Attempt to refresh the token by first performing a guarded reload from + /// the active auth source. If the loaded token differs from the cached token, + /// we can assume that the source already refreshed it. Otherwise, ask the + /// token authority to refresh. pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> { let _refresh_guard = self.refresh_lock.acquire().await.map_err(|_| { RefreshTokenError::Permanent(RefreshTokenFailedError::new( @@ -2386,9 +2367,8 @@ impl AuthManager { } /// Attempt to refresh the current auth token from the authority that issued - /// the token. On success, reloads the auth state from disk so other components - /// observe refreshed token. If the token refresh fails, returns the error to - /// the caller. + /// it and update the shared cache. If the token refresh fails, returns the + /// error to the caller. pub async fn refresh_token_from_authority(&self) -> Result<(), RefreshTokenError> { let _refresh_guard = self.refresh_lock.acquire().await.map_err(|_| { RefreshTokenError::Permanent(RefreshTokenFailedError::new( @@ -2411,24 +2391,26 @@ impl AuthManager { } let attempted_auth = auth.clone(); - let result = match auth { - CodexAuth::ChatgptAuthTokens(_) => { - self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) - .await - } - CodexAuth::Chatgpt(chatgpt_auth) => { - let token_data = chatgpt_auth.current_token_data().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other( - "Token data is not available.", - )) - })?; - self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token) - .await + let result = if self.has_external_auth() { + self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await + } else { + match auth { + CodexAuth::Chatgpt(chatgpt_auth) => { + let token_data = chatgpt_auth.current_token_data().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other( + "Token data is not available.", + )) + })?; + self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token) + .await + } + CodexAuth::ApiKey(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockApiKey(_) => Ok(()), } - CodexAuth::ApiKey(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) - | CodexAuth::BedrockApiKey(_) => Ok(()), }; if let Err(RefreshTokenError::Permanent(error)) = &result { self.record_permanent_refresh_failure_if_unchanged(&attempted_auth, error); @@ -2533,67 +2515,22 @@ impl AuthManager { .refresh(context) .await .map_err(RefreshTokenError::Transient)?; - self.commit_external_auth(&external_auth, &refreshed) - .await?; + self.validate_external_auth(&refreshed)?; + self.set_cached_auth(Some(refreshed)); Ok(()) } - async fn commit_external_auth( - &self, - provider: &Arc, - auth: &CodexAuth, - ) -> Result<(), RefreshTokenError> { - if auth.api_auth_mode() != AuthMode::ApiKey { - if !auth.is_external_chatgpt_tokens() { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth returned unsupported auth mode {:?}", - auth.api_auth_mode() - ), - ))); - } - let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); - let account_id = auth.get_account_id().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other( - "external ChatGPT auth tokens are missing an account id", - )) - })?; - if let Some(expected_workspace_ids) = forced_chatgpt_workspace_id.as_deref() - && !expected_workspace_ids.contains(&account_id) - { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" - ), - ))); - } - let is_cached = self.inner.read().ok().is_some_and(|cached| { - Self::auths_equal_for_refresh(cached.auth.as_ref(), Some(auth)) - }); - if !is_cached { - let auth_dot_json = auth.get_current_auth_json().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other( - "external ChatGPT auth tokens are missing auth state", - )) - })?; - save_auth( - &self.codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - AuthKeyringBackendKind::default(), - ) - .map_err(RefreshTokenError::Transient)?; - self.reload().await; - } + fn validate_external_auth(&self, auth: &CodexAuth) -> Result<(), RefreshTokenError> { + if let Some(account_id) = auth.get_account_id() + && let Some(expected_workspace_ids) = self.forced_chatgpt_workspace_id() + && !expected_workspace_ids.contains(&account_id) + { + return Err(RefreshTokenError::Transient(std::io::Error::other( + format!( + "external auth returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" + ), + ))); } - - let mut state = self.external_auth.write().map_err(|_| { - RefreshTokenError::Transient(std::io::Error::other("external auth lock is poisoned")) - })?; - *state = Some(ExternalAuthState { - provider: Arc::clone(provider), - last_resolved_auth: Some(auth.clone()), - }); Ok(()) } From 7077b30baa961bfb6ce69f441991a63d458a5a3e Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 13:45:45 -0700 Subject: [PATCH 09/10] codex: fix CI failure on PR #31421 --- codex-rs/login/src/auth/manager.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index c8a77efde88..15043a1d088 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -2221,7 +2221,7 @@ impl AuthManager { if let Ok(mut external_auth) = self.external_auth.write() && external_auth.take().is_some() { - self.set_cached_auth(None); + self.set_cached_auth(/*new_auth*/ None); } } From 3ae9aabe0e44beb7efedfb9ce6a547e6d81e5d35 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Tue, 7 Jul 2026 13:55:15 -0700 Subject: [PATCH 10/10] codex: address PR review feedback (#31421) --- .../app-server/tests/suite/v2/app_list.rs | 91 ++++++++++++++++++- codex-rs/login/src/auth/manager.rs | 27 +++++- 2 files changed, 114 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 5fad04929f7..a9a74802e2d 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -8,7 +8,9 @@ use std::time::Duration; use anyhow::Result; use anyhow::bail; use app_test_support::ChatGptAuthFixture; +use app_test_support::ChatGptIdTokenClaims; use app_test_support::TestAppServer; +use app_test_support::encode_id_token; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Json; @@ -29,6 +31,7 @@ use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadStartParams; @@ -156,6 +159,89 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { Ok(()) } +#[tokio::test] +async fn list_apps_uses_external_chatgpt_auth() -> Result<()> { + let access_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("external@example.com") + .plan_type("pro") + .chatgpt_account_id("account-123"), + )?; + let connectors = vec![AppInfo { + id: "beta".to_string(), + name: "Beta".to_string(), + description: Some("Beta connector".to_string()), + logo_url: None, + logo_url_dark: None, + icon_assets: None, + icon_dark_assets: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + }]; + let tools = vec![connector_tool("beta", "Beta App")?]; + let (server_url, server_handle, _) = start_apps_server_with_delays_and_control_inner( + connectors, + tools, + Duration::ZERO, + Duration::ZERO, + /*workspace_plugins_enabled*/ true, + &access_token, + ) + .await?; + + let codex_home = TempDir::new()?; + write_connectors_config(codex_home.path(), &server_url)?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + let login_id = mcp + .send_chatgpt_auth_tokens_login_request( + access_token, + "account-123".to_string(), + Some("pro".to_string()), + ) + .await?; + let login_response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(login_id)), + ) + .await??; + assert_eq!( + to_response::(login_response)?, + LoginAccountResponse::ChatgptAuthTokens {} + ); + + let request_id = mcp + .send_apps_list_request(AppsListParams { + limit: None, + cursor: None, + thread_id: None, + force_refetch: true, + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let AppsListResponse { data, next_cursor } = to_response(response)?; + + assert_eq!(data.len(), 1); + assert_eq!(data[0].id, "beta"); + assert!(data[0].is_accessible); + assert!(next_cursor.is_none()); + + server_handle.abort(); + let _ = server_handle.await; + Ok(()) +} + #[tokio::test] async fn list_apps_returns_empty_when_workspace_codex_plugins_disabled() -> Result<()> { let connectors = vec![AppInfo { @@ -1557,6 +1643,7 @@ async fn start_apps_server_with_workspace_plugins_enabled( Duration::ZERO, Duration::ZERO, workspace_plugins_enabled, + "chatgpt-token", ) .await?; Ok((server_url, server_handle)) @@ -1574,6 +1661,7 @@ async fn start_apps_server_with_delays_and_control( directory_delay, tools_delay, /*workspace_plugins_enabled*/ true, + "chatgpt-token", ) .await } @@ -1584,13 +1672,14 @@ async fn start_apps_server_with_delays_and_control_inner( directory_delay: Duration, tools_delay: Duration, workspace_plugins_enabled: bool, + expected_bearer: &str, ) -> Result<(String, JoinHandle<()>, AppsServerControl)> { let response = Arc::new(StdMutex::new( json!({ "apps": connectors, "next_token": null }), )); let tools = Arc::new(StdMutex::new(tools)); let state = AppsServerState { - expected_bearer: "Bearer chatgpt-token".to_string(), + expected_bearer: format!("Bearer {expected_bearer}"), expected_account_id: "account-123".to_string(), response: response.clone(), directory_delay, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 15043a1d088..dfc46b8f853 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -2213,8 +2213,7 @@ impl AuthManager { *self.external_auth.write().map_err(|_| { RefreshTokenError::Transient(std::io::Error::other("external auth lock is poisoned")) })? = Some(external_auth); - self.set_cached_auth(Some(auth)); - Ok(()) + self.commit_external_auth(auth) } pub fn clear_external_auth(&self) { @@ -2516,7 +2515,29 @@ impl AuthManager { .await .map_err(RefreshTokenError::Transient)?; self.validate_external_auth(&refreshed)?; - self.set_cached_auth(Some(refreshed)); + self.commit_external_auth(refreshed)?; + Ok(()) + } + + fn commit_external_auth(&self, auth: CodexAuth) -> Result<(), RefreshTokenError> { + if auth.is_external_chatgpt_tokens() { + let auth_dot_json = auth.get_current_auth_json().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other( + "external ChatGPT auth tokens are missing auth state", + )) + })?; + // App/connectors paths still construct independent AuthManagers from Config. Mirror + // external ChatGPT auth into the process-local store so those managers see it too. + save_auth( + &self.codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + ) + .map_err(RefreshTokenError::Transient)?; + } + + self.set_cached_auth(Some(auth)); Ok(()) }