diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 6ae73784ff66..29ace7d99fcc 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -74,10 +74,10 @@ 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_login::auth::ExternalAuthTokens; use codex_protocol::ThreadId; use codex_protocol::auth::AuthMode as LoginAuthMode; use codex_protocol::protocol::SessionSource; @@ -121,10 +121,7 @@ impl ExternalAuthRefreshBridge { } } - async fn refresh( - &self, - context: ExternalAuthRefreshContext, - ) -> std::io::Result { + 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, @@ -162,11 +159,11 @@ impl ExternalAuthRefreshBridge { let response: ChatgptAuthTokensRefreshResponse = serde_json::from_value(result).map_err(std::io::Error::other)?; - Ok(ExternalAuthTokens::chatgpt( - response.access_token, - response.chatgpt_account_id, - response.chatgpt_plan_type, - )) + CodexAuth::from_external_chatgpt_tokens( + response.access_token.as_str(), + response.chatgpt_account_id.as_str(), + response.chatgpt_plan_type.as_deref(), + ) } } @@ -178,7 +175,7 @@ impl ExternalAuth for ExternalAuthRefreshBridge { fn refresh( &self, context: ExternalAuthRefreshContext, - ) -> codex_login::ExternalAuthFuture<'_, ExternalAuthTokens> { + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { Box::pin(ExternalAuthRefreshBridge::refresh(self, context)) } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 9509cab45662..77d8c4fdfb2e 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1015,19 +1015,6 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None); } -#[test] -fn external_auth_tokens_without_chatgpt_metadata_cannot_seed_chatgpt_auth() { - let err = AuthDotJson::from_external_tokens(&ExternalAuthTokens::access_token_only( - "test-access-token", - )) - .expect_err("bearer-only external auth should not seed ChatGPT auth"); - - assert_eq!( - err.to_string(), - "external auth tokens are missing ChatGPT metadata" - ); -} - #[tokio::test] async fn 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 f30d52f434c1..1dbd2c3eeefe 100644 --- a/codex-rs/login/src/auth/external_bearer.rs +++ b/codex-rs/login/src/auth/external_bearer.rs @@ -1,7 +1,7 @@ +use super::manager::CodexAuth; use super::manager::ExternalAuth; use super::manager::ExternalAuthFuture; use super::manager::ExternalAuthRefreshContext; -use super::manager::ExternalAuthTokens; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ModelProviderAuthInfo; use std::fmt; @@ -30,7 +30,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() { @@ -39,8 +39,8 @@ impl BearerTokenRefresher { None => true, }; if should_use_cached_token { - return Ok(Some(ExternalAuthTokens::access_token_only( - cached_token.access_token.clone(), + return Ok(Some(CodexAuth::from_api_key( + cached_token.access_token.as_str(), ))); } } @@ -52,20 +52,17 @@ impl BearerTokenRefresher { }); access_token }; - Ok(Some(ExternalAuthTokens::access_token_only(access_token))) + Ok(Some(CodexAuth::from_api_key(access_token.as_str()))) } - async fn refresh( - &self, - _context: ExternalAuthRefreshContext, - ) -> io::Result { + 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(ExternalAuthTokens::access_token_only(access_token)) + Ok(CodexAuth::from_api_key(access_token.as_str())) } } @@ -74,14 +71,11 @@ impl ExternalAuth for BearerTokenRefresher { AuthMode::ApiKey } - fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + fn resolve(&self) -> ExternalAuthFuture<'_, Option> { Box::pin(BearerTokenRefresher::resolve(self)) } - fn refresh( - &self, - context: ExternalAuthRefreshContext, - ) -> ExternalAuthFuture<'_, ExternalAuthTokens> { + fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { Box::pin(BearerTokenRefresher::refresh(self, context)) } } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index a2d005dfa65f..06d125422c44 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -198,45 +198,6 @@ pub enum RefreshTokenError { Transient(#[from] std::io::Error), } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ExternalAuthTokens { - pub access_token: String, - pub chatgpt_metadata: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ExternalAuthChatgptMetadata { - pub account_id: String, - pub plan_type: Option, -} - -impl ExternalAuthTokens { - pub fn access_token_only(access_token: impl Into) -> Self { - Self { - access_token: access_token.into(), - chatgpt_metadata: None, - } - } - - pub fn chatgpt( - access_token: impl Into, - chatgpt_account_id: impl Into, - chatgpt_plan_type: Option, - ) -> Self { - Self { - access_token: access_token.into(), - chatgpt_metadata: Some(ExternalAuthChatgptMetadata { - account_id: chatgpt_account_id.into(), - plan_type: chatgpt_plan_type, - }), - } - } - - pub fn chatgpt_metadata(&self) -> Option<&ExternalAuthChatgptMetadata> { - self.chatgpt_metadata.as_ref() - } -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ExternalAuthRefreshReason { Unauthorized, @@ -251,22 +212,19 @@ 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 -/// credentials on demand via `refresh()`. +/// 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. - fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + fn resolve(&self) -> ExternalAuthFuture<'_, Option> { Box::pin(async { Ok(None) }) } /// Refreshes auth in response to a manager-driven refresh attempt. - fn refresh( - &self, - context: ExternalAuthRefreshContext, - ) -> ExternalAuthFuture<'_, ExternalAuthTokens>; + fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth>; } pub type ExternalAuthFuture<'a, T> = Pin> + Send + 'a>>; @@ -747,6 +705,24 @@ impl CodexAuth { Self::Chatgpt(ChatgptAuth { state, storage }) } + /// Constructs in-memory ChatGPT auth from externally managed tokens. + pub fn from_external_chatgpt_tokens( + access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, + ) -> std::io::Result { + let auth_dot_json = AuthDotJson::from_external_access_token( + access_token, + chatgpt_account_id, + chatgpt_plan_type, + )?; + let state = ChatgptAuthState { + auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), + client: create_client(), + }; + Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) + } + pub fn from_api_key(api_key: &str) -> Self { Self::ApiKey(ApiKeyAuth { api_key: api_key.to_owned(), @@ -1456,26 +1432,23 @@ fn refresh_token_endpoint() -> String { } impl AuthDotJson { - fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { - let Some(chatgpt_metadata) = external.chatgpt_metadata() else { - return Err(std::io::Error::other( - "external auth tokens are missing ChatGPT metadata", - )); - }; + fn from_external_access_token( + access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, + ) -> std::io::Result { let mut token_info = - parse_chatgpt_jwt_claims(&external.access_token).map_err(std::io::Error::other)?; - token_info.chatgpt_account_id = Some(chatgpt_metadata.account_id.clone()); - token_info.chatgpt_plan_type = chatgpt_metadata - .plan_type - .as_deref() + parse_chatgpt_jwt_claims(access_token).map_err(std::io::Error::other)?; + token_info.chatgpt_account_id = Some(chatgpt_account_id.to_string()); + token_info.chatgpt_plan_type = chatgpt_plan_type .map(InternalPlanType::from_raw_value) .or(token_info.chatgpt_plan_type) .or(Some(InternalPlanType::Unknown("unknown".to_string()))); let tokens = TokenData { id_token: token_info, - access_token: external.access_token.clone(), + access_token: access_token.to_string(), refresh_token: String::new(), - account_id: Some(chatgpt_metadata.account_id.clone()), + account_id: Some(chatgpt_account_id.to_string()), }; Ok(Self { @@ -1489,19 +1462,6 @@ impl AuthDotJson { }) } - fn from_external_access_token( - access_token: &str, - chatgpt_account_id: &str, - chatgpt_plan_type: Option<&str>, - ) -> std::io::Result { - let external = ExternalAuthTokens::chatgpt( - access_token, - chatgpt_account_id, - chatgpt_plan_type.map(str::to_string), - ); - Self::from_external_tokens(&external) - } - pub(super) fn resolved_mode(&self) -> AuthMode { if let Some(mode) = self.auth_mode { return mode; @@ -2361,7 +2321,14 @@ impl AuthManager { let external_auth = self.external_auth()?; match external_auth.resolve().await { - Ok(Some(tokens)) => Some(CodexAuth::from_api_key(&tokens.access_token)), + 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" + ); + None + } Ok(None) => None, Err(err) => { tracing::error!("Failed to resolve external API key auth: {err}"); @@ -2567,23 +2534,30 @@ impl AuthManager { if external_auth.auth_mode() == AuthMode::ApiKey { return Ok(()); } - let Some(chatgpt_metadata) = refreshed.chatgpt_metadata() else { + if refreshed.api_auth_mode() != AuthMode::ChatgptAuthTokens { return Err(RefreshTokenError::Transient(std::io::Error::other( - "external auth refresh did not return ChatGPT metadata", + "external auth refresh did not return ChatGPT auth tokens", ))); - }; + } + let account_id = refreshed.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(&chatgpt_metadata.account_id) + && !expected_workspace_ids.contains(&account_id) { return Err(RefreshTokenError::Transient(std::io::Error::other( format!( - "external auth refresh returned workspace {:?}, expected one of {:?}", - chatgpt_metadata.account_id, expected_workspace_ids, + "external auth refresh returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" ), ))); } - let auth_dot_json = - AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; + let auth_dot_json = refreshed.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, diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 7ea415fcc569..8a08ddeadb80 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -30,11 +30,9 @@ pub use auth::CODEX_ACCESS_TOKEN_ENV_VAR; pub use auth::CODEX_API_KEY_ENV_VAR; pub use auth::CodexAuth; pub use auth::ExternalAuth; -pub use auth::ExternalAuthChatgptMetadata; pub use auth::ExternalAuthFuture; pub use auth::ExternalAuthRefreshContext; pub use auth::ExternalAuthRefreshReason; -pub use auth::ExternalAuthTokens; pub use auth::OPENAI_API_KEY_ENV_VAR; pub use auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; pub use auth::REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR; diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index e7d1691a3a20..2e97dfac431b 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -7,7 +7,6 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::ExternalAuth; use codex_login::ExternalAuthRefreshContext; -use codex_login::ExternalAuthTokens; use codex_login::TokenData; use codex_protocol::auth::AuthMode; use codex_protocol::openai_models::ModelsResponse; @@ -122,23 +121,15 @@ impl ExternalAuth for TestExternalApiKeyAuth { AuthMode::ApiKey } - fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, Option> { - Box::pin(async { - Ok(Some(ExternalAuthTokens::access_token_only( - "test-external-api-key", - ))) - }) + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, Option> { + Box::pin(async { Ok(Some(CodexAuth::from_api_key("test-external-api-key"))) }) } fn refresh( &self, _context: ExternalAuthRefreshContext, - ) -> codex_login::ExternalAuthFuture<'_, ExternalAuthTokens> { - Box::pin(async { - Ok(ExternalAuthTokens::access_token_only( - "test-external-api-key", - )) - }) + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(CodexAuth::from_api_key("test-external-api-key")) }) } } @@ -153,7 +144,7 @@ impl ExternalAuth for TestUnresolvedExternalApiKeyAuth { fn refresh( &self, _context: ExternalAuthRefreshContext, - ) -> codex_login::ExternalAuthFuture<'_, ExternalAuthTokens> { + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { Box::pin(async { Err(std::io::Error::other("unresolved test auth")) }) } }