diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 2ac9b931ff..6f780e5de6 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1278,63 +1278,48 @@ impl ModelClient { routing: ClientRouting, provider: &SharedModelProvider, ) -> Result { - // Capture before resolving credentials so an account switch during setup cannot label - // an old connection with the new owner's revision. - let auth_owner_generation = self.auth_owner_generation(); - let auth_changes = provider - .auth_manager() - .map(|manager| manager.auth_change_receiver()); - loop { - let revision = auth_changes.as_ref().map(|changes| *changes.borrow()); - let auth = provider.auth().await; - let (api_provider, redirect_policy) = match routing { - ClientRouting::Workspace => { - let resolved = provider - .responses_api_provider(&self.state.workspace_routing) - .await?; - (resolved.provider, resolved.redirect_policy) - } - ClientRouting::ConfiguredProvider => ( - provider.api_provider().await?, - ClientRedirectPolicy::Default, - ), - }; - let resolved_auth = provider - .api_auth_for_scope(ProviderAuthScope { - agent_identity_policy: self.agent_identity_policy, - session_source: self.state.session_source.clone(), - agent_identity_session_fallback: self - .state - .agent_identity_session_fallback - .clone(), + let routing = &routing; + codex_model_provider::prepare_provider_request( + provider.as_ref(), + provider.info().auth.is_none(), + |state| async move { + let auth = provider.auth().await; + let (api_provider, redirect_policy) = match routing { + ClientRouting::Workspace => { + let resolved = provider + .responses_api_provider(&self.state.workspace_routing) + .await?; + (resolved.provider, resolved.redirect_policy) + } + ClientRouting::ConfiguredProvider => ( + provider.api_provider().await?, + ClientRedirectPolicy::Default, + ), + }; + let resolved_auth = provider + .api_auth_for_scope(ProviderAuthScope { + agent_identity_policy: self.agent_identity_policy, + session_source: self.state.session_source.clone(), + agent_identity_session_fallback: self + .state + .agent_identity_session_fallback + .clone(), + }) + .await?; + Ok(CurrentClientSetup { + auth, + auth_owner_generation: state.map(|state| state.owner_generation), + auth_revision: state.map(|state| state.generation), + api_provider, + http_client_factory: self.http_client_factory.clone(), + redirect_policy, + api_auth: resolved_auth.auth, + agent_identity_telemetry: resolved_auth.agent_identity_telemetry, + runtime_id: None, }) - .await?; - // Command-backed bearer refreshes cannot change workspace routing. Keep the captured - // revisions so a refresh during setup still invalidates cached WebSocket state. - if provider.info().auth.is_none() { - if self.auth_owner_generation() != auth_owner_generation { - return Err(std::io::Error::other( - "account changed while preparing model request", - ) - .into()); - } - // Rebuild the entire bundle after a credential-only refresh, including routing. - if auth_changes.as_ref().map(|changes| *changes.borrow()) != revision { - continue; - } - } - return Ok(CurrentClientSetup { - auth, - auth_owner_generation, - auth_revision: revision, - api_provider, - http_client_factory: self.http_client_factory.clone(), - redirect_policy, - api_auth: resolved_auth.auth, - agent_identity_telemetry: resolved_auth.agent_identity_telemetry, - runtime_id: None, - }); - } + }, + ) + .await } fn responses_headers(&self, auth: Option<&CodexAuth>, model: &str) -> ApiHeaderMap { diff --git a/codex-rs/login/src/auth/change_state.rs b/codex-rs/login/src/auth/change_state.rs index 4a09b84de9..cdb1ab3cf0 100644 --- a/codex-rs/login/src/auth/change_state.rs +++ b/codex-rs/login/src/auth/change_state.rs @@ -19,6 +19,27 @@ pub struct AuthChangeState { pub login_generation: u64, } +/// Whether request preparation can use the credentials it selected. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RequestAuthChange { + Unchanged, + CredentialsChanged, + OwnerChanged, +} + +impl AuthChangeState { + /// Compare a completed attempt with its original owner and starting credential revision. + pub fn compare_for_request(&self, expected_owner: u64, prepared: Self) -> RequestAuthChange { + if self.owner_generation != expected_owner { + RequestAuthChange::OwnerChanged + } else if self.generation != prepared.generation { + RequestAuthChange::CredentialsChanged + } else { + RequestAuthChange::Unchanged + } + } +} + pub(super) fn same_owner(previous: Option<&CodexAuth>, current: Option<&CodexAuth>) -> bool { let (Some(previous), Some(current)) = (previous, current) else { return false; diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index af0afd0e1d..e629be9e6a 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -24,6 +24,7 @@ pub use bedrock_access_keys::login_with_bedrock_access_keys; pub use bedrock_api_key::BedrockApiKeyAuth; pub use bedrock_api_key::login_with_bedrock_api_key; pub use change_state::AuthChangeState; +pub use change_state::RequestAuthChange; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 39fbb7ead1..a7cbdce3f6 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -58,6 +58,7 @@ 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; pub use auth::RefreshTokenError; +pub use auth::RequestAuthChange; pub use auth::UnauthorizedRecovery; pub use auth::default_client; pub use auth::enforce_login_restrictions; diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index f054a8f606..d24b696e68 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -8,6 +8,7 @@ mod gcp; mod models_endpoint; mod models_identity; mod provider; +mod request_preparation; mod runtime_auth; mod shared_state; pub mod test_support; @@ -50,6 +51,7 @@ pub use provider::create_model_provider; pub use provider::provider_uses_first_party_auth_path; pub use provider::resolve_api_provider; +pub use request_preparation::prepare_provider_request; pub use runtime_auth::select_runtime_auth; pub use runtime_auth::validate_runtime_destination; pub use runtime_auth::validate_runtime_provider; diff --git a/codex-rs/model-provider/src/request_preparation.rs b/codex-rs/model-provider/src/request_preparation.rs new file mode 100644 index 0000000000..359cfd7535 --- /dev/null +++ b/codex-rs/model-provider/src/request_preparation.rs @@ -0,0 +1,56 @@ +//! Keeps request preparation bound to the credential revision it started with. + +use std::future::Future; +use std::io; + +use codex_login::AuthChangeState; +use codex_login::RequestAuthChange; +use codex_protocol::error::Result; + +use crate::ModelProvider; + +/// Prepare a request without dispatching it. When the request uses managed auth, +/// repeat preparation after a credential refresh and fail if the owner changes. +pub async fn prepare_provider_request( + provider: &dyn ModelProvider, + check_auth: bool, + mut prepare: impl FnMut(Option) -> Fut, +) -> Result +where + Fut: Future>, +{ + let auth_state = || { + provider + .auth_manager() + .map(|manager| *manager.auth_change_state_receiver().borrow()) + }; + let mut before = auth_state(); + loop { + let prepared = prepare(before).await?; + if !check_auth { + return Ok(prepared); + } + // A managed provider can begin or cease exposing its auth manager during setup. + let after = auth_state(); + let change = match (before, after) { + (Some(before), Some(after)) => { + after.compare_for_request(before.owner_generation, before) + } + (None, None) => RequestAuthChange::Unchanged, + _ => RequestAuthChange::OwnerChanged, + }; + match change { + RequestAuthChange::Unchanged => return Ok(prepared), + RequestAuthChange::CredentialsChanged => before = after, + RequestAuthChange::OwnerChanged => { + return Err( + io::Error::other("account changed while preparing model request").into(), + ); + } + } + } +} + +#[cfg(test)] +#[path = "request_preparation_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider/src/request_preparation_tests.rs b/codex-rs/model-provider/src/request_preparation_tests.rs new file mode 100644 index 0000000000..6fe63f8185 --- /dev/null +++ b/codex-rs/model-provider/src/request_preparation_tests.rs @@ -0,0 +1,62 @@ +use super::prepare_provider_request; +use crate::create_model_provider; +use crate::provider_uses_first_party_auth_path; +use base64::Engine; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_model_provider_info::ModelProviderInfo; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; + +#[tokio::test] +async fn explicit_provider_token_survives_a_workspace_switch_during_setup() { + let home = TempDir::new().unwrap(); + let claims = json!({"https://api.openai.com/auth": {"chatgpt_user_id": "user-a"}}); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + let token = format!("header.{payload}.signature"); + let initial = CodexAuth::from_external_chatgpt_tokens( + &token, + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap(); + let manager = AuthManager::from_auth_for_testing_with_home(initial, home.path().to_path_buf()); + let mut info = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + info.experimental_bearer_token = Some("provider-token".into()); + let provider = create_model_provider(info, Some(manager.clone())); + + let mut attempts = 0; + let auth = prepare_provider_request( + provider.as_ref(), + provider_uses_first_party_auth_path(provider.info()), + |_| { + attempts += 1; + async { + let (_, auth) = provider.api_provider_and_auth().await?; + codex_login::auth::login_with_chatgpt_auth_tokens( + home.path(), + &token, + "workspace-b", + /*chatgpt_plan_type*/ None, + )?; + manager.reload().await; + Ok(auth) + } + }, + ) + .await + .unwrap(); + + assert_eq!( + manager.auth_cached().unwrap().get_account_id().as_deref(), + Some("workspace-b") + ); + assert_eq!(attempts, 1); + let mut headers = http::HeaderMap::new(); + auth.add_auth_headers(&mut headers).unwrap(); + assert_eq!( + headers[http::header::AUTHORIZATION], + "Bearer provider-token" + ); +}