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 000000000000..d00777f71015 --- /dev/null +++ b/codex-rs/app-server/src/external_auth.rs @@ -0,0 +1,95 @@ +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); + +pub(crate) struct ExternalAuthBridge { + outgoing: Arc, + auth: RwLock, +} + +impl ExternalAuthBridge { + pub(crate) fn new(outgoing: Arc, auth: CodexAuth) -> Self { + Self { + outgoing, + auth: RwLock::new(auth), + } + } + + 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.auth + .read() + .map(|auth| auth.clone()) + .map_err(|_| std::io::Error::other("external auth lock is poisoned")) + }) + } + + 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 fa6b6cd811f7..91da0c1645c9 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 29ace7d99fcc..ca122920bd64 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,12 +70,7 @@ 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::auth::AuthMode as LoginAuthMode; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::StateDbHandle; @@ -95,7 +86,6 @@ 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( @@ -109,77 +99,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 auth_mode(&self) -> LoginAuthMode { - LoginAuthMode::Chatgpt - } - - 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, @@ -324,9 +243,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.rs b/codex-rs/app-server/src/request_processors.rs index f03351468faa..7f5ee69a09d2 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 9592563bece0..5eaf6e8d33f7 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; @@ -617,14 +618,19 @@ 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 + .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( self.auth_manager.clone(), self.config.chatgpt_base_url.clone(), diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 3cd98e2e52d6..25106047a815 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/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 5fad04929f7f..a9a74802e2d3 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/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index dfd4fe50ced8..0896b4628cc6 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] diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 77d8c4fdfb2e..0b0bfc882d13 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/external_bearer.rs b/codex-rs/login/src/auth/external_bearer.rs index 1dbd2c3eeefe..0a276543eabf 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; @@ -30,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() { @@ -39,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())); } } @@ -52,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 { @@ -67,11 +64,7 @@ impl BearerTokenRefresher { } impl ExternalAuth for BearerTokenRefresher { - fn auth_mode(&self) -> AuthMode { - AuthMode::ApiKey - } - - 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 06d125422c44..dfc46b8f8536 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -211,19 +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 { - /// Indicates which top-level auth mode this external provider supplies. - fn auth_mode(&self) -> AuthMode; + /// Returns the provider's current auth value. + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth>; - /// Returns cached or immediately available auth, if this provider can resolve it synchronously - /// from the caller's perspective. - fn resolve(&self) -> ExternalAuthFuture<'_, Option> { - Box::pin(async { Ok(None) }) - } - - /// 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>; } @@ -1555,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, @@ -1585,11 +1573,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 @@ -1725,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), @@ -1980,7 +1962,10 @@ impl AuthManager { /// 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.inner + .read() + .ok() + .and_then(|cached| cached.auth.clone()) } /// Subscribes to cached auth changes that can affect request recovery. @@ -2003,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_api_key_auth().await { - return Some(auth); + if self.has_external_auth() { + self.reload().await; + return self.auth_cached(); } let auth = self.auth_cached()?; @@ -2080,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) } @@ -2100,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) { @@ -2173,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, @@ -2210,15 +2205,22 @@ impl AuthManager { } } - pub fn set_external_auth(&self, external_auth: Arc) { - if let Ok(mut guard) = self.external_auth.write() { - *guard = Some(external_auth); - } + pub async fn set_external_auth( + &self, + external_auth: Arc, + ) -> Result<(), RefreshTokenError> { + 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.commit_external_auth(auth) } 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(/*new_auth*/ None); } } @@ -2300,48 +2302,33 @@ impl AuthManager { self.external_auth .read() .ok() - .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()) + .and_then(|external_auth| external_auth.as_ref().map(Arc::clone)) } fn has_external_api_key_auth(&self) -> bool { - self.external_auth_mode() == Some(AuthMode::ApiKey) + self.has_external_auth() + && self + .auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_api_key_auth) } - async fn resolve_external_api_key_auth(&self) -> Option { - if !self.has_external_api_key_auth() { - return None; - } - - 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" - ); - None - } - Ok(None) => None, - Err(err) => { - tracing::error!("Failed to resolve external API key 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( @@ -2379,9 +2366,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( @@ -2404,24 +2390,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); @@ -2440,6 +2428,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) } @@ -2459,23 +2448,18 @@ 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) } /// 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 +2501,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,41 +2514,44 @@ impl AuthManager { .refresh(context) .await .map_err(RefreshTokenError::Transient)?; - if external_auth.auth_mode() == AuthMode::ApiKey { - return Ok(()); - } - if refreshed.api_auth_mode() != AuthMode::ChatgptAuthTokens { - return Err(RefreshTokenError::Transient(std::io::Error::other( - "external auth refresh did not return ChatGPT auth tokens", - ))); + self.validate_external_auth(&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)?; } - 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() + + self.set_cached_auth(Some(auth)); + Ok(()) + } + + 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 refresh returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" + "external auth returned workspace {account_id:?}, expected one of {expected_workspace_ids:?}" ), ))); } - 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, - AuthCredentialsStoreMode::Ephemeral, - AuthKeyringBackendKind::default(), - ) - .map_err(RefreshTokenError::Transient)?; - self.reload().await; Ok(()) } diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 2e97dfac431b..f1678fe9f255 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -117,12 +117,8 @@ 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"))) }) + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(CodexAuth::from_api_key("test-external-api-key")) }) } fn refresh( @@ -137,8 +133,8 @@ impl ExternalAuth for TestExternalApiKeyAuth { struct TestUnresolvedExternalApiKeyAuth; impl ExternalAuth for TestUnresolvedExternalApiKeyAuth { - fn auth_mode(&self) -> AuthMode { - AuthMode::ApiKey + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Err(std::io::Error::other("unresolved test auth")) }) } fn refresh( @@ -858,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( @@ -898,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( @@ -1013,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!(