diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 64bce20860..64294bd0b1 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -70,6 +70,7 @@ use codex_core::models_manager::collaboration_mode_presets::CollaborationModesCo use codex_exec_server::EnvironmentManager; use codex_features::Feature; use codex_feedback::CodexFeedback; +use codex_login::auth::ExternalChatGptAuthRefresher; use codex_login::auth::ExternalAuthRefreshContext; use codex_login::auth::ExternalAuthRefreshReason; use codex_login::auth::ExternalAuthRefresher; @@ -89,11 +90,11 @@ use tracing::Instrument; const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Clone)] -struct ExternalAuthRefreshBridge { +struct ExternalChatGptAuthRefreshBridge { outgoing: Arc, } -impl ExternalAuthRefreshBridge { +impl ExternalChatGptAuthRefreshBridge { fn map_reason(reason: ExternalAuthRefreshReason) -> ChatgptAuthTokensRefreshReason { match reason { ExternalAuthRefreshReason::Unauthorized => ChatgptAuthTokensRefreshReason::Unauthorized, @@ -102,7 +103,15 @@ impl ExternalAuthRefreshBridge { } #[async_trait] -impl ExternalAuthRefresher for ExternalAuthRefreshBridge { +impl ExternalAuthRefresher for ExternalChatGptAuthRefreshBridge { + fn auth_mode(&self) -> codex_app_server_protocol::AuthMode { + codex_app_server_protocol::AuthMode::ChatgptAuthTokens + } + + async fn resolve(&self) -> std::io::Result> { + Ok(None) + } + async fn refresh( &self, context: ExternalAuthRefreshContext, @@ -152,6 +161,8 @@ impl ExternalAuthRefresher for ExternalAuthRefreshBridge { } } +impl ExternalChatGptAuthRefresher for ExternalChatGptAuthRefreshBridge {} + pub(crate) struct MessageProcessor { outgoing: Arc, codex_message_processor: CodexMessageProcessor, @@ -210,7 +221,7 @@ impl MessageProcessor { config.codex_home.clone(), enable_codex_api_key_env, config.cli_auth_credentials_store_mode, - Arc::new(ExternalAuthRefreshBridge { + Arc::new(ExternalChatGptAuthRefreshBridge { outgoing: outgoing.clone(), }), ); diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 53580c72d3..afa7bcfdc3 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -447,8 +447,10 @@ struct AuthFileParams { chatgpt_account_id: Option, } -fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result { - let auth_file = get_auth_file(codex_home); +fn make_test_chatgpt_jwt( + chatgpt_account_id: Option<&str>, + chatgpt_plan_type: Option<&str>, +) -> std::io::Result { // Create a minimal valid JWT for the id_token field. #[derive(Serialize)] struct Header { @@ -464,13 +466,14 @@ fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result "user_id": "user-12345", }); - if let Some(chatgpt_plan_type) = params.chatgpt_plan_type { - auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type); + if let Some(chatgpt_plan_type) = chatgpt_plan_type { + auth_payload["chatgpt_plan_type"] = + serde_json::Value::String(chatgpt_plan_type.to_string()); } - if let Some(chatgpt_account_id) = params.chatgpt_account_id { - let org_value = serde_json::Value::String(chatgpt_account_id); - auth_payload["chatgpt_account_id"] = org_value; + if let Some(chatgpt_account_id) = chatgpt_account_id { + auth_payload["chatgpt_account_id"] = + serde_json::Value::String(chatgpt_account_id.to_string()); } let payload = serde_json::json!({ @@ -478,11 +481,19 @@ fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result "email_verified": true, "https://api.openai.com/auth": auth_payload, }); - let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b); + let b64 = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); let header_b64 = b64(&serde_json::to_vec(&header)?); let payload_b64 = b64(&serde_json::to_vec(&payload)?); let signature_b64 = b64(b"sig"); - let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + Ok(format!("{header_b64}.{payload_b64}.{signature_b64}")) +} + +fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result { + let auth_file = get_auth_file(codex_home); + let fake_jwt = make_test_chatgpt_jwt( + params.chatgpt_account_id.as_deref(), + params.chatgpt_plan_type.as_deref(), + )?; let auth_json_data = json!({ "OPENAI_API_KEY": params.openai_api_key, diff --git a/codex-rs/login/src/auth/external_bearer.rs b/codex-rs/login/src/auth/external_bearer.rs index 4cc94e6ec5..1b6eee3273 100644 --- a/codex-rs/login/src/auth/external_bearer.rs +++ b/codex-rs/login/src/auth/external_bearer.rs @@ -1,3 +1,8 @@ +use crate::auth::manager::ExternalAuthRefreshContext; +use crate::auth::manager::ExternalAuthRefresher; +use crate::auth::manager::ExternalAuthTokens; +use async_trait::async_trait; +use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ModelProviderAuthInfo; use std::fmt; use std::io; @@ -10,56 +15,67 @@ use tokio::process::Command; use tokio::sync::Mutex; #[derive(Clone)] -pub(crate) struct ExternalBearerAuth { - state: Arc, +pub(crate) struct BearerTokenRefresher { + state: Arc, } -impl ExternalBearerAuth { +impl BearerTokenRefresher { pub(crate) fn new(config: ModelProviderAuthInfo) -> Self { Self { - state: Arc::new(ExternalBearerAuthState::new(config)), + state: Arc::new(BearerTokenRefresherState::new(config)), } } +} - pub(crate) async fn resolve_access_token(&self) -> io::Result { +#[async_trait] +impl ExternalAuthRefresher for BearerTokenRefresher { + fn auth_mode(&self) -> ApiAuthMode { + ApiAuthMode::ApiKey + } + + async fn resolve(&self) -> io::Result> { let mut cached = self.state.cached_token.lock().await; - if let Some(cached_token) = cached.as_ref() + let access_token = if let Some(cached_token) = cached.as_ref() && cached_token.fetched_at.elapsed() < self.state.config.refresh_interval() { - return Ok(cached_token.access_token.clone()); - } + cached_token.access_token.clone() + } else { + let access_token = run_provider_auth_command(&self.state.config).await?; + *cached = Some(CachedExternalBearerToken { + access_token: access_token.clone(), + fetched_at: Instant::now(), + }); + access_token + }; + Ok(Some(ExternalAuthTokens::access_token_only(access_token))) + } + async fn refresh( + &self, + _context: ExternalAuthRefreshContext, + ) -> io::Result { let access_token = run_provider_auth_command(&self.state.config).await?; + let mut cached = self.state.cached_token.lock().await; *cached = Some(CachedExternalBearerToken { access_token: access_token.clone(), fetched_at: Instant::now(), }); - Ok(access_token) - } - - pub(crate) async fn refresh_after_unauthorized(&self) -> io::Result<()> { - let access_token = run_provider_auth_command(&self.state.config).await?; - let mut cached = self.state.cached_token.lock().await; - *cached = Some(CachedExternalBearerToken { - access_token, - fetched_at: Instant::now(), - }); - Ok(()) + Ok(ExternalAuthTokens::access_token_only(access_token)) } } -impl fmt::Debug for ExternalBearerAuth { +impl fmt::Debug for BearerTokenRefresher { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ExternalBearerAuth").finish_non_exhaustive() + f.debug_struct("BearerTokenRefresher").finish_non_exhaustive() } } -struct ExternalBearerAuthState { +struct BearerTokenRefresherState { config: ModelProviderAuthInfo, cached_token: Mutex>, } -impl ExternalBearerAuthState { +impl BearerTokenRefresherState { fn new(config: ModelProviderAuthInfo) -> Self { Self { config, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index d13106cdd3..3bdccca5ad 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -18,7 +18,7 @@ use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; -use super::external_bearer::ExternalBearerAuth; +use super::external_bearer::BearerTokenRefresher; use crate::auth::error::RefreshTokenFailedError; use crate::auth::error::RefreshTokenFailedReason; pub use crate::auth::storage::AuthCredentialsStoreMode; @@ -144,6 +144,8 @@ pub struct ExternalAuthRefreshContext { #[async_trait] pub trait ExternalAuthRefresher: Send + Sync { + fn auth_mode(&self) -> ApiAuthMode; + async fn resolve(&self) -> std::io::Result> { Ok(None) } @@ -154,6 +156,9 @@ pub trait ExternalAuthRefresher: Send + Sync { ) -> std::io::Result; } +#[async_trait] +pub trait ExternalChatGptAuthRefresher: ExternalAuthRefresher {} + impl RefreshTokenError { pub fn failed_reason(&self) -> Option { match self { @@ -853,27 +858,6 @@ struct AuthScopedRefreshFailure { error: RefreshTokenFailedError, } -#[derive(Clone)] -enum ExternalAuth { - Bearer(ExternalBearerAuth), - ChatgptRefresher(Arc), -} - -impl Debug for ExternalAuth { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Bearer(_) => f - .debug_tuple("ExternalAuth::Bearer") - .field(&"present") - .finish(), - Self::ChatgptRefresher(_) => f - .debug_tuple("ExternalAuth::ChatgptRefresher") - .field(&"present") - .finish(), - } - } -} - impl Debug for CachedAuth { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CachedAuth") @@ -928,8 +912,8 @@ enum UnauthorizedRecoveryMode { // // - External ChatGPT auth tokens (`chatgptAuthTokens`) are refreshed by asking // the parent app for new tokens through the configured -// `ExternalAuthRefresher`, persisting them in the ephemeral auth store, and -// reloading the cached auth snapshot. +// `ExternalChatGptAuthRefresher`, 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. pub struct UnauthorizedRecovery { @@ -1112,7 +1096,6 @@ impl UnauthorizedRecovery { /// External modifications to `auth.json` will NOT be observed until /// `reload()` is called explicitly. This matches the design goal of avoiding /// different parts of the program seeing inconsistent auth data mid‑run. -#[derive(Debug)] pub struct AuthManager { codex_home: PathBuf, inner: RwLock, @@ -1120,7 +1103,13 @@ pub struct AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode, forced_chatgpt_workspace_id: RwLock>, refresh_lock: AsyncMutex<()>, - external_auth: RwLock>, + external_refresher: RwLock>>, +} + +impl Debug for AuthManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AuthManager").finish_non_exhaustive() + } } impl AuthManager { @@ -1150,7 +1139,7 @@ impl AuthManager { auth_credentials_store_mode, forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), - external_auth: RwLock::new(None), + external_refresher: RwLock::new(None), } } @@ -1168,7 +1157,7 @@ impl AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), - external_auth: RwLock::new(None), + external_refresher: RwLock::new(None), }) } @@ -1185,7 +1174,7 @@ impl AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), - external_auth: RwLock::new(None), + external_refresher: RwLock::new(None), }) } @@ -1200,7 +1189,7 @@ impl AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), - external_auth: RwLock::new(Some(ExternalAuth::Bearer(ExternalBearerAuth::new(config)))), + external_refresher: RwLock::new(Some(Arc::new(BearerTokenRefresher::new(config)))), }) } @@ -1223,7 +1212,7 @@ 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 { - if let Some(auth) = self.resolve_external_bearer_auth().await { + if let Some(auth) = self.resolve_external_auth().await { return Some(auth); } @@ -1346,14 +1335,18 @@ impl AuthManager { } } - pub fn set_external_chatgpt_auth_refresher(&self, refresher: Arc) { - if let Ok(mut guard) = self.external_auth.write() { - *guard = Some(ExternalAuth::ChatgptRefresher(refresher)); + pub fn set_external_chatgpt_auth_refresher( + &self, + refresher: Arc, + ) { + if let Ok(mut guard) = self.external_refresher.write() { + let refresher: Arc = refresher; + *guard = Some(refresher); } } pub fn clear_external_chatgpt_auth_refresher(&self) { - if let Ok(mut guard) = self.external_auth.write() { + if let Ok(mut guard) = self.external_refresher.write() { *guard = None; } } @@ -1372,11 +1365,7 @@ impl AuthManager { } pub fn has_external_chatgpt_auth_refresher(&self) -> bool { - self.external_auth - .read() - .ok() - .map(|guard| matches!(guard.as_ref(), Some(ExternalAuth::ChatgptRefresher(_)))) - .unwrap_or(false) + self.external_refresher_auth_mode() == Some(ApiAuthMode::ChatgptAuthTokens) } pub fn is_external_chatgpt_auth_active(&self) -> bool { @@ -1406,7 +1395,7 @@ impl AuthManager { codex_home: PathBuf, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, - refresher: Arc, + refresher: Arc, ) -> Arc { let manager = Self::shared( codex_home, @@ -1421,26 +1410,59 @@ impl AuthManager { UnauthorizedRecovery::new(Arc::clone(self)) } - fn external_auth(&self) -> Option { - self.external_auth + fn external_refresher(&self) -> Option> { + self.external_refresher .read() .ok() - .and_then(|guard| guard.clone()) + .and_then(|guard| guard.as_ref().map(Arc::clone)) + } + + fn external_refresher_auth_mode(&self) -> Option { + self.external_refresher + .read() + .ok() + .and_then(|guard| guard.as_ref().map(|refresher| refresher.auth_mode())) } fn has_external_bearer_auth(&self) -> bool { - matches!(self.external_auth(), Some(ExternalAuth::Bearer(_))) + self.external_refresher_auth_mode() == Some(ApiAuthMode::ApiKey) } - async fn resolve_external_bearer_auth(&self) -> Option { - let ExternalAuth::Bearer(bearer_auth) = self.external_auth()? else { - return None; - }; - - match bearer_auth.resolve_access_token().await { - Ok(access_token) => Some(CodexAuth::from_api_key(&access_token)), + async fn resolve_external_auth(&self) -> Option { + let refresher = self.external_refresher()?; + let auth_mode = refresher.auth_mode(); + let resolved = match refresher.resolve().await { + Ok(resolved) => resolved, Err(err) => { - tracing::error!("Failed to resolve external bearer auth: {err}"); + tracing::error!("Failed to resolve external auth: {err}"); + None + } + }?; + + match auth_mode { + ApiAuthMode::ApiKey => Some(CodexAuth::from_api_key(&resolved.access_token)), + ApiAuthMode::ChatgptAuthTokens => { + let auth_dot_json = match AuthDotJson::from_external_tokens(&resolved) { + Ok(auth_dot_json) => auth_dot_json, + Err(err) => { + tracing::error!("Failed to resolve external ChatGPT auth: {err}"); + return None; + } + }; + match CodexAuth::from_auth_dot_json( + &self.codex_home, + auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + ) { + Ok(auth) => Some(auth), + Err(err) => { + tracing::error!("Failed to build external ChatGPT auth: {err}"); + None + } + } + } + ApiAuthMode::Chatgpt => { + tracing::error!("External auth refreshers do not support managed ChatGPT auth"); None } } @@ -1534,8 +1556,10 @@ impl AuthManager { } pub fn get_api_auth_mode(&self) -> Option { - if self.has_external_bearer_auth() { - return Some(ApiAuthMode::ApiKey); + if self.has_external_bearer_auth() + && let Some(auth_mode) = self.external_refresher_auth_mode() + { + return Some(auth_mode); } self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) } @@ -1573,19 +1597,13 @@ impl AuthManager { &self, reason: ExternalAuthRefreshReason, ) -> Result<(), RefreshTokenError> { - if let Some(ExternalAuth::Bearer(bearer_auth)) = self.external_auth() { - return bearer_auth - .refresh_after_unauthorized() - .await - .map_err(RefreshTokenError::Transient); - } - let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); - let Some(ExternalAuth::ChatgptRefresher(refresher)) = self.external_auth() else { + let Some(refresher) = self.external_refresher() else { return Err(RefreshTokenError::Transient(std::io::Error::other( "external auth refresher is not configured", ))); }; + let auth_mode = refresher.auth_mode(); let previous_account_id = self .auth_cached() @@ -1597,31 +1615,39 @@ impl AuthManager { }; let refreshed = refresher.refresh(context).await?; - 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() - && 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:?}", - chatgpt_metadata.account_id, - ), - ))); + match auth_mode { + ApiAuthMode::ApiKey => Ok(()), + ApiAuthMode::ChatgptAuthTokens => { + 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() + && 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:?}", + chatgpt_metadata.account_id, + ), + ))); + } + let auth_dot_json = AuthDotJson::from_external_tokens(&refreshed) + .map_err(RefreshTokenError::Transient)?; + save_auth( + &self.codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + ) + .map_err(RefreshTokenError::Transient)?; + self.reload(); + Ok(()) + } + ApiAuthMode::Chatgpt => Err(RefreshTokenError::Transient(std::io::Error::other( + "external auth refresher cannot refresh managed ChatGPT auth", + ))), } - let auth_dot_json = - AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; - save_auth( - &self.codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - ) - .map_err(RefreshTokenError::Transient)?; - self.reload(); - Ok(()) } // Refreshes ChatGPT OAuth tokens, persists the updated auth state, and diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index a8033cf888..261d17345c 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -22,6 +22,7 @@ pub use auth::AuthManager; pub use auth::CLIENT_ID; pub use auth::CODEX_API_KEY_ENV_VAR; pub use auth::CodexAuth; +pub use auth::ExternalChatGptAuthRefresher; pub use auth::ExternalAuthChatgptMetadata; pub use auth::ExternalAuthRefreshContext; pub use auth::ExternalAuthRefreshReason;