diff --git a/AGENTS.md b/AGENTS.md index ff36b0e19c..f5065f29b1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,6 +17,7 @@ In the codex-rs folder where the rust code lives: - Do not add these comments for string or char literals unless the comment adds real clarity; those literals are intentionally exempt from the lint. - If you add one of these comments, the parameter name must exactly match the callee signature. - When possible, make `match` statements exhaustive and avoid wildcard arms. +- Newly added traits should include doc comments that explain their role and how implementations are expected to use them. - When writing tests, prefer comparing the equality of entire objects over fields one by one. - When making a change that adds or changes an API, ensure that the documentation in the `docs/` folder is up to date if applicable. - If you change `ConfigToml` or nested config types, run `just write-config-schema` to update `codex-rs/core/config.schema.json`. diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index fbb3f451bf..47be1e1e59 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -71,9 +71,10 @@ 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::AuthMode as LoginAuthMode; +use codex_login::auth::ExternalAuth; use codex_login::auth::ExternalAuthRefreshContext; use codex_login::auth::ExternalAuthRefreshReason; -use codex_login::auth::ExternalAuthRefresher; use codex_login::auth::ExternalAuthTokens; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; @@ -103,7 +104,11 @@ impl ExternalAuthRefreshBridge { } #[async_trait] -impl ExternalAuthRefresher for ExternalAuthRefreshBridge { +impl ExternalAuth for ExternalAuthRefreshBridge { + fn auth_mode(&self) -> LoginAuthMode { + LoginAuthMode::Chatgpt + } + async fn refresh( &self, context: ExternalAuthRefreshContext, @@ -211,7 +216,7 @@ impl MessageProcessor { enable_codex_api_key_env, rpc_transport, } = args; - let auth_manager = AuthManager::shared_with_external_chatgpt_auth_refresher( + let auth_manager = AuthManager::shared_with_external_auth( config.codex_home.clone(), enable_codex_api_key_env, config.cli_auth_credentials_store_mode, @@ -290,7 +295,7 @@ impl MessageProcessor { } pub(crate) fn clear_runtime_references(&self) { - self.auth_manager.clear_external_chatgpt_auth_refresher(); + self.auth_manager.clear_external_auth(); } pub(crate) async fn process_request( diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 53580c72d3..3e230ca707 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -283,6 +283,8 @@ async fn external_bearer_only_auth_manager_uses_cached_provider_token() { assert_eq!(first.as_deref(), Some("provider-token")); assert_eq!(second.as_deref(), Some("provider-token")); + assert_eq!(manager.auth_mode(), Some(crate::AuthMode::ApiKey)); + assert_eq!(manager.get_api_auth_mode(), Some(ApiAuthMode::ApiKey)); } #[tokio::test] @@ -448,42 +450,8 @@ struct AuthFileParams { } fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result { + let fake_jwt = fake_jwt_for_auth_file_params(¶ms)?; let auth_file = get_auth_file(codex_home); - // Create a minimal valid JWT for the id_token field. - #[derive(Serialize)] - struct Header { - alg: &'static str, - typ: &'static str, - } - let header = Header { - alg: "none", - typ: "JWT", - }; - let mut auth_payload = serde_json::json!({ - "chatgpt_user_id": "user-12345", - "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_account_id) = params.chatgpt_account_id { - let org_value = serde_json::Value::String(chatgpt_account_id); - auth_payload["chatgpt_account_id"] = org_value; - } - - let payload = serde_json::json!({ - "email": "user@example.com", - "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 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}"); - let auth_json_data = json!({ "OPENAI_API_KEY": params.openai_api_key, "tokens": { @@ -498,6 +466,42 @@ fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result Ok(fake_jwt) } +fn fake_jwt_for_auth_file_params(params: &AuthFileParams) -> std::io::Result { + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + + let header = Header { + alg: "none", + typ: "JWT", + }; + let mut auth_payload = serde_json::json!({ + "chatgpt_user_id": "user-12345", + "user_id": "user-12345", + }); + + if let Some(chatgpt_plan_type) = params.chatgpt_plan_type.as_ref() { + auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type.clone()); + } + + if let Some(chatgpt_account_id) = params.chatgpt_account_id.as_ref() { + auth_payload["chatgpt_account_id"] = serde_json::Value::String(chatgpt_account_id.clone()); + } + + let payload = serde_json::json!({ + "email": "user@example.com", + "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 header_b64 = b64(&serde_json::to_vec(&header)?); + let payload_b64 = b64(&serde_json::to_vec(&payload)?); + let signature_b64 = b64(b"sig"); + Ok(format!("{header_b64}.{payload_b64}.{signature_b64}")) +} + async fn build_config( codex_home: &Path, forced_login_method: Option, diff --git a/codex-rs/login/src/auth/external_bearer.rs b/codex-rs/login/src/auth/external_bearer.rs index 4cc94e6ec5..18a94889e3 100644 --- a/codex-rs/login/src/auth/external_bearer.rs +++ b/codex-rs/login/src/auth/external_bearer.rs @@ -1,3 +1,7 @@ +use super::manager::ExternalAuth; +use super::manager::ExternalAuthRefreshContext; +use super::manager::ExternalAuthTokens; +use async_trait::async_trait; use codex_protocol::config_types::ModelProviderAuthInfo; use std::fmt; use std::io; @@ -10,47 +14,61 @@ use tokio::process::Command; use tokio::sync::Mutex; #[derive(Clone)] -pub(crate) struct ExternalBearerAuth { +pub(crate) struct BearerTokenRefresher { state: Arc, } -impl ExternalBearerAuth { +impl BearerTokenRefresher { pub(crate) fn new(config: ModelProviderAuthInfo) -> Self { Self { state: Arc::new(ExternalBearerAuthState::new(config)), } } +} - pub(crate) async fn resolve_access_token(&self) -> io::Result { - let mut cached = self.state.cached_token.lock().await; - 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()); - } +#[async_trait] +impl ExternalAuth for BearerTokenRefresher { + fn auth_mode(&self) -> crate::AuthMode { + crate::AuthMode::ApiKey + } + 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() + && cached_token.fetched_at.elapsed() < self.state.config.refresh_interval() + { + 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() } } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index d13106cdd3..632d92b3db 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; @@ -143,11 +143,21 @@ pub struct ExternalAuthRefreshContext { } #[async_trait] -pub trait ExternalAuthRefresher: Send + Sync { +/// Pluggable auth provider used by `AuthManager` for externally managed auth flows. +/// +/// Implementations may either resolve auth eagerly via `resolve()` or provide refreshed +/// credentials on demand via `refresh()`. +pub trait ExternalAuth: Send + Sync { + /// Indicates which top-level auth mode this external provider supplies. + fn auth_mode(&self) -> crate::AuthMode; + + /// Returns cached or immediately available auth, if this provider can resolve it synchronously + /// from the caller's perspective. async fn resolve(&self) -> std::io::Result> { Ok(None) } + /// Refreshes auth in response to a manager-driven refresh attempt. async fn refresh( &self, context: ExternalAuthRefreshContext, @@ -853,27 +863,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,7 +917,7 @@ 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 +// `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. @@ -954,7 +943,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_bearer_auth() + let mode = if manager.has_external_api_key_auth() || cached_auth .as_ref() .is_some_and(CodexAuth::is_external_chatgpt_tokens) @@ -976,7 +965,7 @@ impl UnauthorizedRecovery { } pub fn has_next(&self) -> bool { - if self.manager.has_external_bearer_auth() { + if self.manager.has_external_api_key_auth() { return !matches!(self.step, UnauthorizedRecoveryStep::Done); } @@ -989,9 +978,7 @@ impl UnauthorizedRecovery { return false; } - if self.mode == UnauthorizedRecoveryMode::External - && !self.manager.has_external_chatgpt_auth_refresher() - { + if self.mode == UnauthorizedRecoveryMode::External && !self.manager.has_external_auth() { return false; } @@ -999,7 +986,7 @@ impl UnauthorizedRecovery { } pub fn unavailable_reason(&self) -> &'static str { - if self.manager.has_external_bearer_auth() { + if self.manager.has_external_api_key_auth() { return if matches!(self.step, UnauthorizedRecoveryStep::Done) { "recovery_exhausted" } else { @@ -1016,10 +1003,8 @@ impl UnauthorizedRecovery { return "not_chatgpt_auth"; } - if self.mode == UnauthorizedRecoveryMode::External - && !self.manager.has_external_chatgpt_auth_refresher() - { - return "no_external_refresher"; + if self.mode == UnauthorizedRecoveryMode::External && !self.manager.has_external_auth() { + return "no_external_auth"; } if matches!(self.step, UnauthorizedRecoveryStep::Done) { @@ -1112,7 +1097,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 +1104,26 @@ pub struct AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode, forced_chatgpt_workspace_id: RwLock>, refresh_lock: AsyncMutex<()>, - external_auth: RwLock>, + external_auth: RwLock>>, +} + +impl Debug for AuthManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AuthManager") + .field("codex_home", &self.codex_home) + .field("inner", &self.inner) + .field("enable_codex_api_key_env", &self.enable_codex_api_key_env) + .field( + "auth_credentials_store_mode", + &self.auth_credentials_store_mode, + ) + .field( + "forced_chatgpt_workspace_id", + &self.forced_chatgpt_workspace_id, + ) + .field("has_external_auth", &self.has_external_auth()) + .finish_non_exhaustive() + } } impl AuthManager { @@ -1200,7 +1203,9 @@ 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_auth: RwLock::new(Some( + Arc::new(BearerTokenRefresher::new(config)) as Arc + )), }) } @@ -1223,7 +1228,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_api_key_auth().await { return Some(auth); } @@ -1346,13 +1351,13 @@ impl AuthManager { } } - pub fn set_external_chatgpt_auth_refresher(&self, refresher: Arc) { + pub fn set_external_auth(&self, external_auth: Arc) { if let Ok(mut guard) = self.external_auth.write() { - *guard = Some(ExternalAuth::ChatgptRefresher(refresher)); + *guard = Some(external_auth); } } - pub fn clear_external_chatgpt_auth_refresher(&self) { + pub fn clear_external_auth(&self) { if let Ok(mut guard) = self.external_auth.write() { *guard = None; } @@ -1371,12 +1376,8 @@ impl AuthManager { .and_then(|guard| guard.clone()) } - 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) + pub fn has_external_auth(&self) -> bool { + self.external_auth().is_some() } pub fn is_external_chatgpt_auth_active(&self) -> bool { @@ -1402,18 +1403,18 @@ impl AuthManager { )) } - pub fn shared_with_external_chatgpt_auth_refresher( + pub fn shared_with_external_auth( codex_home: PathBuf, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, - refresher: Arc, + external_auth: Arc, ) -> Arc { let manager = Self::shared( codex_home, enable_codex_api_key_env, auth_credentials_store_mode, ); - manager.set_external_chatgpt_auth_refresher(refresher); + manager.set_external_auth(external_auth); manager } @@ -1421,26 +1422,35 @@ impl AuthManager { UnauthorizedRecovery::new(Arc::clone(self)) } - fn external_auth(&self) -> Option { + fn external_auth(&self) -> Option> { self.external_auth .read() .ok() - .and_then(|guard| guard.clone()) + .and_then(|guard| guard.as_ref().cloned()) } - fn has_external_bearer_auth(&self) -> bool { - matches!(self.external_auth(), Some(ExternalAuth::Bearer(_))) + fn external_auth_mode(&self) -> Option { + self.external_auth() + .as_ref() + .map(|external_auth| external_auth.auth_mode()) } - async fn resolve_external_bearer_auth(&self) -> Option { - let ExternalAuth::Bearer(bearer_auth) = self.external_auth()? else { + fn has_external_api_key_auth(&self) -> bool { + self.external_auth_mode() == Some(crate::AuthMode::ApiKey) + } + + async fn resolve_external_api_key_auth(&self) -> Option { + if !self.has_external_api_key_auth() { return None; - }; + } - match bearer_auth.resolve_access_token().await { - Ok(access_token) => Some(CodexAuth::from_api_key(&access_token)), + let external_auth = self.external_auth()?; + + match external_auth.resolve().await { + Ok(Some(tokens)) => Some(CodexAuth::from_api_key(&tokens.access_token)), + Ok(None) => None, Err(err) => { - tracing::error!("Failed to resolve external bearer auth: {err}"); + tracing::error!("Failed to resolve external API key auth: {err}"); None } } @@ -1534,14 +1544,14 @@ impl AuthManager { } pub fn get_api_auth_mode(&self) -> Option { - if self.has_external_bearer_auth() { + if self.has_external_api_key_auth() { return Some(ApiAuthMode::ApiKey); } self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) } pub fn auth_mode(&self) -> Option { - if self.has_external_bearer_auth() { + if self.has_external_api_key_auth() { return Some(crate::AuthMode::ApiKey); } self.auth_cached().as_ref().map(CodexAuth::auth_mode) @@ -1573,20 +1583,12 @@ 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(external_auth) = self.external_auth() else { return Err(RefreshTokenError::Transient(std::io::Error::other( - "external auth refresher is not configured", + "external auth is not configured", ))); }; - + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); let previous_account_id = self .auth_cached() .as_ref() @@ -1596,7 +1598,13 @@ impl AuthManager { previous_account_id, }; - let refreshed = refresher.refresh(context).await?; + let refreshed = external_auth + .refresh(context) + .await + .map_err(RefreshTokenError::Transient)?; + if external_auth.auth_mode() == crate::AuthMode::ApiKey { + return Ok(()); + } let Some(chatgpt_metadata) = refreshed.chatgpt_metadata() else { return Err(RefreshTokenError::Transient(std::io::Error::other( "external auth refresh did not return ChatGPT metadata", diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index a8033cf888..1ea1c0cb4b 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -22,10 +22,10 @@ pub use auth::AuthManager; pub use auth::CLIENT_ID; pub use auth::CODEX_API_KEY_ENV_VAR; pub use auth::CodexAuth; +pub use auth::ExternalAuth; pub use auth::ExternalAuthChatgptMetadata; pub use auth::ExternalAuthRefreshContext; pub use auth::ExternalAuthRefreshReason; -pub use auth::ExternalAuthRefresher; pub use auth::ExternalAuthTokens; pub use auth::OPENAI_API_KEY_ENV_VAR; pub use auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;