diff --git a/codex-rs/cli/src/exec_server_credential.rs b/codex-rs/cli/src/exec_server_credential.rs index 86a33fdf96..bd0719d775 100644 --- a/codex-rs/cli/src/exec_server_credential.rs +++ b/codex-rs/cli/src/exec_server_credential.rs @@ -128,9 +128,11 @@ impl AuthProvider for ExecutorCredentialProvider { ) -> codex_api::AuthRefreshFuture<'_> { Box::pin(async move { // The parent owns refresh; only retry if it already supplied a valid replacement. - self.resolve_auth_headers().await.is_ok_and(|headers| { - headers.get(AUTHORIZATION) != rejected_headers.get(AUTHORIZATION) - }) + self.apply_auth_headers(HeaderMap::new()) + .await + .is_ok_and(|headers| { + headers.get(AUTHORIZATION) != rejected_headers.get(AUTHORIZATION) + }) }) } } diff --git a/codex-rs/cli/src/exec_server_credential_tests.rs b/codex-rs/cli/src/exec_server_credential_tests.rs index 8599f61b1b..5a59e8bf26 100644 --- a/codex-rs/cli/src/exec_server_credential_tests.rs +++ b/codex-rs/cli/src/exec_server_credential_tests.rs @@ -42,14 +42,17 @@ async fn waits_for_initial_rotation_replaces_token_and_eof_clears_auth() { .await, Err(CredentialError::InitialTimeout) ); - assert!(provider.resolve_auth_headers().await.is_err()); + assert!(provider.apply_auth_headers(HeaderMap::new()).await.is_err()); assert!(provider.to_auth_headers().is_err()); writer .write_all(&frame("env-a", "first-scoped-token", expires)) .expect("write"); provider.wait_for_initial().await.expect("initial token"); next_update(&mut updates).await; - let headers = provider.resolve_auth_headers().await.expect("headers"); + let headers = provider + .apply_auth_headers(HeaderMap::new()) + .await + .expect("headers"); assert_eq!(headers.len(), 1); assert_eq!(headers[AUTHORIZATION], "Bearer first-scoped-token"); assert!(headers[AUTHORIZATION].is_sensitive()); @@ -62,7 +65,10 @@ async fn waits_for_initial_rotation_replaces_token_and_eof_clears_auth() { next_update(&mut updates).await; // A request sent before the parent update can retry with the replacement. assert!(provider.refresh_after_unauthorized(headers).await); - let headers = provider.resolve_auth_headers().await.expect("headers"); + let headers = provider + .apply_auth_headers(HeaderMap::new()) + .await + .expect("headers"); assert_eq!(headers.len(), 1); assert_eq!(headers[AUTHORIZATION], "Bearer rotated-scoped-token"); assert!(headers[AUTHORIZATION].is_sensitive()); @@ -70,7 +76,7 @@ async fn waits_for_initial_rotation_replaces_token_and_eof_clears_auth() { drop(writer); pipe.parent_closed.await.expect("reader must monitor EOF"); - assert!(provider.resolve_auth_headers().await.is_err()); + assert!(provider.apply_auth_headers(HeaderMap::new()).await.is_err()); assert!(!provider.refresh_after_unauthorized(headers).await); assert_eq!(provider.check_pipe_failure(), Ok(())); } @@ -114,7 +120,7 @@ async fn malformed_rotations_latch_failure_and_clear_existing_credential() { .expect("bad frame terminates without EOF"); let auth_error = pipe .provider - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .expect_err("reject auth"); assert_eq!(pipe.provider.check_pipe_failure(), Err(error)); @@ -152,7 +158,12 @@ async fn eof_before_a_complete_frame_never_authenticates() { drop(writer); pipe.parent_closed.await.expect("reader monitors EOF"); assert_eq!(pipe.provider.wait_for_initial().await, Err(expected)); - assert!(pipe.provider.resolve_auth_headers().await.is_err()); + assert!( + pipe.provider + .apply_auth_headers(HeaderMap::new()) + .await + .is_err() + ); } } @@ -168,14 +179,14 @@ async fn expiry_is_checked_again_on_every_request() { provider.wait_for_initial().await, Err(CredentialError::Expired) ); - assert!(provider.resolve_auth_headers().await.is_err()); + assert!(provider.apply_auth_headers(HeaderMap::new()).await.is_err()); assert!(!provider.refresh_after_unauthorized(HeaderMap::new()).await); updates.send_replace(CredentialState::Active { token: HeaderValue::from_static("new-scoped-token"), expires_at: expired + 60, }); let headers = provider - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .expect("rotation restores auth"); assert_eq!(headers[AUTHORIZATION], "Bearer new-scoped-token"); diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index 6d1ab4f513..c72dc2d7e6 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -36,7 +36,7 @@ pub trait AuthProvider: Send + Sync { /// /// Implementations should be cheap and non-blocking. This method is also /// used by telemetry and non-HTTP request paths. Providers that require async - /// work must return an error here and override `resolve_auth_headers`. + /// work must return an error here and override `apply_auth_headers`. /// Callers must not send the request if this returns an error. fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError>; @@ -60,15 +60,6 @@ pub trait AuthProvider: Send + Sync { }) } - /// Resolves auth headers for an outbound request. - /// - /// Unlike [`Self::to_auth_headers`], implementations may perform asynchronous work to refresh - /// credentials before returning. Header-only providers with static credentials can rely on the - /// default implementation. - fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> { - Box::pin(async { self.to_auth_headers() }) - } - /// Refreshes credentials after an explicit HTTP 401. Returns whether the caller /// may retry with new credentials. Callers must bound retries and pass the headers /// of the rejected request so concurrent failures can reuse an earlier refresh. @@ -77,10 +68,13 @@ pub trait AuthProvider: Send + Sync { } /// Authenticates a header-only request, including removing credentials it supersedes. - /// WebSocket handshakes must use the returned map rather than overlaying resolved headers. + /// + /// Implementations may refresh credentials asynchronously. Providers with static credentials + /// can rely on the default implementation. + /// Callers must use the returned map rather than overlaying resolved headers. fn apply_auth_headers(&self, mut headers: HeaderMap) -> AuthHeadersFuture<'_> { Box::pin(async move { - headers.extend(self.resolve_auth_headers().await?); + headers.extend(self.to_auth_headers()?); Ok(headers) }) } diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index 4410c15d75..664697da32 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -734,12 +734,13 @@ impl AuthProvider for AsyncHeaderAuth { )) } - fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { - Box::pin(async { - Ok(HeaderMap::from_iter([( + fn apply_auth_headers(&self, mut headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> { + Box::pin(async move { + headers.insert( http::header::AUTHORIZATION, HeaderValue::from_static("Bearer async-auth"), - )])) + ); + Ok(headers) }) } } diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index a9a42cef45..a43f5ac567 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -305,7 +305,7 @@ impl EnvironmentRegistryClient { async fn resolve_auth(&self) -> Result<(HeaderMap, Option), ExecServerError> { let mut headers = self .auth_provider - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .map_err(|error| { ExecServerError::EnvironmentRegistryAuth(format!( @@ -1132,9 +1132,8 @@ mod tests { Ok(()) } - fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { - Box::pin(async { - let mut headers = HeaderMap::new(); + fn apply_auth_headers(&self, mut headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> { + Box::pin(async move { let _ = headers.insert( http::header::AUTHORIZATION, HeaderValue::from_static("Bearer registry-token"), diff --git a/codex-rs/exec-server/src/remote/auth_tests.rs b/codex-rs/exec-server/src/remote/auth_tests.rs index 4c61178303..2b11e20d0d 100644 --- a/codex-rs/exec-server/src/remote/auth_tests.rs +++ b/codex-rs/exec-server/src/remote/auth_tests.rs @@ -55,15 +55,16 @@ impl AuthProvider for RotatingAuth { Ok(()) } - fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { - Box::pin(async { + fn apply_auth_headers(&self, mut headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> { + Box::pin(async move { self.stall_at(if self.refreshes.load(Ordering::SeqCst) == 0 { AuthStage::Initial } else { AuthStage::Retry }) .await; - self.to_auth_headers() + self.add_auth_headers(&mut headers)?; + Ok(headers) }) } diff --git a/codex-rs/exec-server/src/remote/registration_retry_tests.rs b/codex-rs/exec-server/src/remote/registration_retry_tests.rs index bae77ef078..81f642f077 100644 --- a/codex-rs/exec-server/src/remote/registration_retry_tests.rs +++ b/codex-rs/exec-server/src/remote/registration_retry_tests.rs @@ -47,9 +47,9 @@ impl AuthProvider for RegistrationAuthProvider { Ok(()) } - fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { + fn apply_auth_headers(&self, headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> { self.calls.fetch_add(1, Ordering::Relaxed); - Box::pin(async { Ok(HeaderMap::new()) }) + Box::pin(async move { Ok(headers) }) } } diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 2cafc0214e..ba50f628bf 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -156,7 +156,7 @@ impl AuthProvider for AuthManagerAuthProvider { auth_provider_from_auth(&auth).add_auth_headers(headers) } - fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> { + fn apply_auth_headers(&self, mut headers: HeaderMap) -> AuthHeadersFuture<'_> { Box::pin(async move { let auth = self .auth_manager @@ -166,7 +166,8 @@ impl AuthProvider for AuthManagerAuthProvider { .ok_or_else(|| { AuthError::Transient("managed authentication is unavailable".to_string()) })?; - auth_provider_from_auth(&auth).to_auth_headers() + headers.extend(auth_provider_from_auth(&auth).to_auth_headers()?); + Ok(headers) }) } @@ -751,7 +752,7 @@ mod tests { .expect("initial auth should be cached"); let provider = auth_provider_from_auth_manager(Arc::clone(&auth_manager), &expected_auth); let initial_headers = provider - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .expect("initial auth headers"); @@ -771,7 +772,7 @@ mod tests { // Ordinary requests retain the manager's cache until refresh or recovery runs. assert_eq!( provider - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .expect("cached auth headers"), initial_headers @@ -785,7 +786,7 @@ mod tests { ); let resolved_headers = provider - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .expect("managed auth headers should resolve"); assert_eq!( @@ -808,7 +809,7 @@ mod tests { Err(AuthError::Transient(_)) )); assert!(matches!( - provider.resolve_auth_headers().await, + provider.apply_auth_headers(HeaderMap::new()).await, Err(AuthError::Transient(_)) )); } diff --git a/codex-rs/model-provider/src/auth_refresh_tests.rs b/codex-rs/model-provider/src/auth_refresh_tests.rs index 2f5b3ad23a..7489edbb6d 100644 --- a/codex-rs/model-provider/src/auth_refresh_tests.rs +++ b/codex-rs/model-provider/src/auth_refresh_tests.rs @@ -101,7 +101,7 @@ async fn concurrent_external_rejections_share_one_refresh() { assert!(provider.refresh_after_unauthorized(rejected).await); assert_eq!(external.refreshes.load(Ordering::SeqCst), 1); assert_eq!( - provider.resolve_auth_headers().await.unwrap()[http::header::AUTHORIZATION], + provider.apply_auth_headers(HeaderMap::new()).await.unwrap()[http::header::AUTHORIZATION], "Bearer replacement" ); } @@ -186,7 +186,7 @@ async fn check_concurrent_refresh(trigger: RefreshTrigger) { .await } RefreshTrigger::Proactive => { - let headers = provider.resolve_auth_headers().await.unwrap(); + let headers = provider.apply_auth_headers(HeaderMap::new()).await.unwrap(); assert_eq!( headers[http::header::AUTHORIZATION], "Bearer replacement-token" @@ -213,7 +213,7 @@ async fn check_concurrent_refresh(trigger: RefreshTrigger) { .refresh_after_unauthorized(rejected_headers.clone()) .await ); - let updated_headers = provider.resolve_auth_headers().await.unwrap(); + let updated_headers = provider.apply_auth_headers(HeaderMap::new()).await.unwrap(); assert_eq!( updated_headers[http::header::AUTHORIZATION], "Bearer replacement-token" @@ -234,7 +234,7 @@ async fn check_concurrent_refresh(trigger: RefreshTrigger) { .await; assert!(provider.refresh_after_unauthorized(updated_headers).await); assert_eq!( - provider.resolve_auth_headers().await.unwrap()[http::header::AUTHORIZATION], + provider.apply_auth_headers(HeaderMap::new()).await.unwrap()[http::header::AUTHORIZATION], "Bearer next-token" ); // A permanent failure must not prevent a later guarded reload of fresh credentials. diff --git a/codex-rs/model-provider/src/combined_auth.rs b/codex-rs/model-provider/src/combined_auth.rs index 586a8bd7e6..da79c4ebea 100644 --- a/codex-rs/model-provider/src/combined_auth.rs +++ b/codex-rs/model-provider/src/combined_auth.rs @@ -120,19 +120,6 @@ impl AuthProvider for CombinedAuth { Ok(()) } - fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> { - Box::pin(async move { - let mut headers = self.primary.resolve_auth_headers().await?; - if headers.contains_key(&self.name) { - return Err(AuthError::Build( - "gateway OAuth conflicts with primary auth headers".into(), - )); - } - headers.insert(self.name.clone(), self.value.clone()); - Ok(headers) - }) - } - fn apply_auth(&self, request: codex_http_client::Request) -> AuthProviderFuture<'_> { Box::pin(async move { let mut request = self.primary.apply_auth(request).await?; diff --git a/codex-rs/model-provider/src/shared_state_tests.rs b/codex-rs/model-provider/src/shared_state_tests.rs index 6039fbb77d..a98a26269f 100644 --- a/codex-rs/model-provider/src/shared_state_tests.rs +++ b/codex-rs/model-provider/src/shared_state_tests.rs @@ -134,7 +134,10 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { ); // WebSocket handshakes use the synchronous header interface. assert_eq!(auth.to_auth_headers().unwrap(), expected); - assert_eq!(auth.resolve_auth_headers().await.unwrap(), expected); + assert_eq!( + auth.apply_auth_headers(HeaderMap::new()).await.unwrap(), + expected + ); assert_eq!(scoped.auth.to_auth_headers().unwrap(), expected); assert!(auth.to_auth_headers().unwrap()[name].is_sensitive()); let transport = ReqwestTransport::from_http_client( diff --git a/codex-rs/runtime-agent-instance/src/auth.rs b/codex-rs/runtime-agent-instance/src/auth.rs index fbccda89a3..fbfc663b2e 100644 --- a/codex-rs/runtime-agent-instance/src/auth.rs +++ b/codex-rs/runtime-agent-instance/src/auth.rs @@ -120,7 +120,7 @@ impl AuthProvider for RuntimeRequestAuth { Box::pin(async move { Ok(replace_auth_headers( headers, - self.resolve_auth_headers().await?, + self.auth.auth_headers(self.mode)?, )) }) } diff --git a/codex-rs/runtime-agent-instance/src/manager_tests.rs b/codex-rs/runtime-agent-instance/src/manager_tests.rs index 06b2d8038c..64695f52b4 100644 --- a/codex-rs/runtime-agent-instance/src/manager_tests.rs +++ b/codex-rs/runtime-agent-instance/src/manager_tests.rs @@ -134,7 +134,7 @@ async fn concurrent_requests_share_registration_and_replacement_preserves_in_fli first .clone() .into_auth_provider(RuntimeAuthMode::DualHeader) - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .unwrap()["authorization"], "Bearer key-a" @@ -143,7 +143,7 @@ async fn concurrent_requests_share_registration_and_replacement_preserves_in_fli replacement .clone() .into_auth_provider(RuntimeAuthMode::DualHeader) - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .unwrap()["authorization"], "Bearer key-b" @@ -230,7 +230,7 @@ async fn policy_recovery_preserves_pending_registration() { let headers = selected .clone() .into_auth_provider(RuntimeAuthMode::DualHeader) - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .unwrap(); assert_eq!(headers["authorization"], "Bearer key-a"); @@ -450,14 +450,14 @@ async fn new_request_uses_refreshed_credentials_without_replacing_the_runtime() let provider = runtime .clone() .into_auth_provider(RuntimeAuthMode::DualHeader); - let first = provider.resolve_auth_headers().await.unwrap(); + let first = provider.apply_auth_headers(HeaderMap::new()).await.unwrap(); let refreshed = oauth(Some("login-a"), "new"); *source.current.write().unwrap() = refreshed.clone(); let next_runtime = manager.auth().await.unwrap(); let next = next_runtime .clone() .into_auth_provider(RuntimeAuthMode::DualHeader) - .resolve_auth_headers() + .apply_auth_headers(HeaderMap::new()) .await .unwrap(); assert_eq!( @@ -466,7 +466,7 @@ async fn new_request_uses_refreshed_credentials_without_replacing_the_runtime() ); assert_ne!(first[AUTHORIZATION], next[AUTHORIZATION]); assert_eq!( - provider.resolve_auth_headers().await.unwrap()[AUTHORIZATION], + provider.apply_auth_headers(HeaderMap::new()).await.unwrap()[AUTHORIZATION], first[AUTHORIZATION] ); assert!(Arc::ptr_eq(runtime.runtime(), next_runtime.runtime()));