diff --git a/codex-rs/cli/src/exec_server_command.rs b/codex-rs/cli/src/exec_server_command.rs index 703fadc10c..ffb08f5312 100644 --- a/codex-rs/cli/src/exec_server_command.rs +++ b/codex-rs/cli/src/exec_server_command.rs @@ -431,7 +431,7 @@ async fn load_exec_server_remote_auth_provider( validate_api_key_remote_host(base_url)?; } - if auth_manager.is_workload_identity_selected() { + if auth.is_chatgpt_auth() || auth_manager.is_workload_identity_selected() { Ok(codex_model_provider::auth_provider_from_auth_manager( auth_manager, &auth, diff --git a/codex-rs/cli/src/exec_server_credential.rs b/codex-rs/cli/src/exec_server_credential.rs index 6142a57e35..86a33fdf96 100644 --- a/codex-rs/cli/src/exec_server_credential.rs +++ b/codex-rs/cli/src/exec_server_credential.rs @@ -121,6 +121,18 @@ impl AuthProvider for ExecutorCredentialProvider { headers.insert(AUTHORIZATION, authorization); Ok(()) } + + fn refresh_after_unauthorized( + &self, + rejected_headers: HeaderMap, + ) -> 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) + }) + }) + } } enum CredentialState { diff --git a/codex-rs/cli/src/exec_server_credential_tests.rs b/codex-rs/cli/src/exec_server_credential_tests.rs index 1772985cb3..8599f61b1b 100644 --- a/codex-rs/cli/src/exec_server_credential_tests.rs +++ b/codex-rs/cli/src/exec_server_credential_tests.rs @@ -54,19 +54,24 @@ async fn waits_for_initial_rotation_replaces_token_and_eof_clears_auth() { assert_eq!(headers[AUTHORIZATION], "Bearer first-scoped-token"); assert!(headers[AUTHORIZATION].is_sensitive()); assert!(!format!("{provider:?} {headers:?}").contains("first-scoped-token")); + assert!(!provider.refresh_after_unauthorized(headers.clone()).await); writer .write_all(&frame("env-a", "rotated-scoped-token", expires)) .expect("rotate"); 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"); assert_eq!(headers.len(), 1); assert_eq!(headers[AUTHORIZATION], "Bearer rotated-scoped-token"); assert!(headers[AUTHORIZATION].is_sensitive()); + assert!(!provider.refresh_after_unauthorized(headers.clone()).await); drop(writer); pipe.parent_closed.await.expect("reader must monitor EOF"); assert!(provider.resolve_auth_headers().await.is_err()); + assert!(!provider.refresh_after_unauthorized(headers).await); assert_eq!(provider.check_pipe_failure(), Ok(())); } @@ -113,6 +118,12 @@ async fn malformed_rotations_latch_failure_and_clear_existing_credential() { .await .expect_err("reject auth"); assert_eq!(pipe.provider.check_pipe_failure(), Err(error)); + assert!( + !pipe + .provider + .refresh_after_unauthorized(HeaderMap::new()) + .await + ); assert!(!auth_error.to_string().contains("do-not-log")); assert!(!auth_error.to_string().contains("initial-scoped-token")); let _ = writer.write_all(&frame("env-a", "later-scoped-token", expires)); @@ -158,6 +169,7 @@ async fn expiry_is_checked_again_on_every_request() { Err(CredentialError::Expired) ); assert!(provider.resolve_auth_headers().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, diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index b0b8e8da69..59c6512602 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -2,6 +2,9 @@ #[path = "exec_server/pid_namespace_tests.rs"] mod pid_namespace_tests; +#[path = "exec_server/auth_refresh_tests.rs"] +mod auth_refresh_tests; + use std::collections::HashMap; #[cfg(unix)] use std::io::BufRead as _; diff --git a/codex-rs/cli/tests/exec_server/auth_refresh_tests.rs b/codex-rs/cli/tests/exec_server/auth_refresh_tests.rs new file mode 100644 index 0000000000..193a6f9422 --- /dev/null +++ b/codex-rs/cli/tests/exec_server/auth_refresh_tests.rs @@ -0,0 +1,118 @@ +//! Verify authentication updates across the running CLI executor process. + +use std::process::Stdio; + +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::write_chatgpt_auth; +use codex_login::AuthCredentialsStoreMode; +use tempfile::TempDir; +use tokio::net::TcpListener; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::accept_hdr_async; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::path; + +#[tokio::test] +async fn noise_executor_refreshes_on_401_and_reloads_rotated_credentials() -> Result<()> { + let home = TempDir::new()?; + let registry = MockServer::start().await; + let listener = TcpListener::bind("127.0.0.1:0").await?; + std::fs::write( + home.path().join("config.toml"), + "cli_auth_credentials_store = 'file'\n", + )?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("old-token").account_id("account"), + AuthCredentialsStoreMode::File, + )?; + Mock::given(path("/refresh")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-token", "refresh_token": "new-refresh-token", + }))) + .expect(1) + .mount(®istry) + .await; + Mock::given(path("/cloud/environment/env/register")) + .and(header("authorization", "Bearer old-token")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(®istry) + .await; + for token in ["refreshed-token", "rotated-token"] { + Mock::given(path("/cloud/environment/env/register")) + .and(header("authorization", format!("Bearer {token}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "env", "url": format!("ws://{}", listener.local_addr()?), + "executor_registration_id": "registration", "security_profile": "noise_hybrid_ik_v1", + }))) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(®istry) + .await; + } + // A later registry rejection reloads the rotated token from storage, without + // another OAuth refresh or unconditional reloads on ordinary requests. + Mock::given(path("/cloud/environment/env/register")) + .and(header("authorization", "Bearer refreshed-token")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(®istry) + .await; + let mut child = tokio::process::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?) + .current_dir(home.path()) + .env("CODEX_HOME", home.path()) + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY") + .env_remove("CODEX_ACCESS_TOKEN") + .env( + "CODEX_REFRESH_TOKEN_URL_OVERRIDE", + format!("{}/refresh", registry.uri()), + ) + .env("NO_PROXY", "127.0.0.1,localhost") + .env("no_proxy", "127.0.0.1,localhost") + .args([ + "exec-server", + "--remote", + ®istry.uri(), + "--environment-id", + "env", + "--exit-on-stdin-close", + ]) + .stdin(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + let stdin = child.stdin.take(); + let (socket, _) = listener.accept().await?; + let mut websocket = accept_async(socket).await?; + // Update storage while the executor is alive, then force a new registration. + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("rotated-token").account_id("account"), + AuthCredentialsStoreMode::File, + )?; + websocket.close(None).await?; + drop(websocket); + let (socket, _) = listener.accept().await?; + let rejected = accept_hdr_async(socket, |_: &_, _: _| { + Err(http::Response::builder().status(401).body(None).unwrap()) + }) + .await; + assert!(rejected.is_err()); + let (socket, _) = listener.accept().await?; + let _websocket = accept_async(socket).await?; + drop(stdin); + let output = child.wait_with_output().await?; + anyhow::ensure!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + Ok(()) +} diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index 4e23badcff..3ce6867eb2 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -56,6 +56,13 @@ pub trait AuthProvider: Send + Sync { 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. + fn refresh_after_unauthorized(&self, _rejected_headers: HeaderMap) -> AuthRefreshFuture<'_> { + Box::pin(async { false }) + } + /// Applies auth to a complete outbound request and returns the request to send. /// /// The input `request` is moved into this method. Implementations may mutate @@ -82,6 +89,8 @@ pub type AuthProviderFuture<'a> = pub type AuthHeadersFuture<'a> = Pin> + Send + 'a>>; +pub type AuthRefreshFuture<'a> = Pin + Send + 'a>>; + /// Shared auth handle passed through API clients. pub type SharedAuthProvider = Arc; diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 3afbcc6ef1..90752102d8 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -26,6 +26,7 @@ pub use crate::auth::AuthHeaderTelemetry; pub use crate::auth::AuthHeadersFuture; pub use crate::auth::AuthProvider; pub use crate::auth::AuthProviderFuture; +pub use crate::auth::AuthRefreshFuture; pub use crate::auth::SharedAuthProvider; pub use crate::auth::auth_header_telemetry; pub use crate::common::AccessPrograms; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index ebb3fa60de..4ba7ec70b3 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -7,11 +7,13 @@ use codex_api::SharedAuthProvider; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_http_client::HttpResponse; +use codex_http_client::Request; use codex_http_client::RouteAwareClientPool; use futures::FutureExt; use http::HeaderMap; use http::HeaderName; use http::HeaderValue; +use http::Method; use http::StatusCode; use http::header::AUTHORIZATION; use serde::Deserialize; @@ -59,6 +61,7 @@ use direct::run_direct_environment; const ERROR_BODY_PREVIEW_BYTES: usize = 4096; const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; +mod auth; #[path = "remote/connection_diagnostics.rs"] mod connection_diagnostics; mod reconnect_backoff; @@ -176,19 +179,10 @@ impl EnvironmentRegistryClient { security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), executor_public_key: executor_public_key.clone(), }; - let (response, executor_credential) = timeout_at(deadline, async { - let (auth_headers, executor_credential) = self.resolve_auth().await?; - let response = self - .http_for_auth() - .post(url) - .headers(auth_headers) - .headers(current_trace_context_headers()) - .json(&body) - .send() - .await - .map_err(ExecServerError::EnvironmentRegistryRequest)?; - Ok((response, executor_credential)) - }) + let (response, executor_credential) = timeout_at( + deadline, + self.send_authenticated(Request::new(Method::POST, url).with_json(&body)), + ) .await .unwrap_or_else(|_| { Err(ExecServerError::EnvironmentRegistryRequest( @@ -266,18 +260,19 @@ impl EnvironmentRegistryClient { &format!("/cloud/environment/{environment_id}/connect"), ); let body = EnvironmentRegistryConnectRequest { harness_public_key }; - let (auth_headers, _) = self.resolve_auth().await?; - let response = self - .http_for_auth() - .post(url) - .headers(auth_headers) - .headers(current_trace_context_headers()) - .json(&body) - .timeout(self.connect_timeout) - .send() - .await?; let response: EnvironmentRegistryConnectResponse = - self.parse_json_response(response).await?; + timeout_at(tokio::time::Instant::now() + self.connect_timeout, async { + let (response, _) = self + .send_authenticated(Request::new(Method::POST, url).with_json(&body)) + .await?; + self.parse_json_response(response).await + }) + .await + .unwrap_or_else(|_| { + Err(ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + )) + })?; if response.environment_id != environment_id { return Err(ExecServerError::Protocol( "environment registry returned a different environment id".to_string(), @@ -420,18 +415,16 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { harness_public_key: harness_public_key.clone(), harness_key_authorization: authorization.to_string(), }; - let (auth_headers, executor_credential) = self.client.resolve_auth().await?; - let mut request = self - .client - .http_for_auth() - .post(url) - .headers(auth_headers) - .headers(current_trace_context_headers()) - .json(&body); + let mut request = Request::new(Method::POST, url).with_json(&body); if let Some(stream_id) = stream_id { - request = request.header("x-codex-noise-stream-id", stream_id.to_string()); + request.headers.insert( + "x-codex-noise-stream-id", + HeaderValue::from_str(&stream_id.to_string()).map_err(|error| { + ExecServerError::Protocol(format!("invalid Noise stream ID header: {error}")) + })?, + ); } - let response = request.send().await?; + let (response, executor_credential) = self.client.send_authenticated(request).await?; let status = response.status(); if !status.is_success() { // The request contains the short-lived authorization. Do not include diff --git a/codex-rs/exec-server/src/remote/auth.rs b/codex-rs/exec-server/src/remote/auth.rs new file mode 100644 index 0000000000..1f0bbe482c --- /dev/null +++ b/codex-rs/exec-server/src/remote/auth.rs @@ -0,0 +1,70 @@ +//! Sends registry requests with current credentials and at most one retry after a 401. +//! Transport failures and other HTTP errors are never replayed here. + +use codex_http_client::HttpResponse; +use codex_http_client::Request; +use http::HeaderValue; +use http::StatusCode; + +use super::EnvironmentRegistryClient; +use super::RegistryAuthMode; +use super::current_trace_context_headers; +use crate::ExecServerError; + +impl EnvironmentRegistryClient { + pub(super) async fn send_authenticated( + &self, + request: Request, + ) -> Result<(HttpResponse, Option), ExecServerError> { + let request = request + .into_prepared() + .map_err(ExecServerError::EnvironmentRegistryConfig)?; + let mut retried = false; + loop { + let (authenticated, credential) = + if self.auth_mode == RegistryAuthMode::ExecutorCredential { + let (headers, credential) = self.resolve_auth().await?; + let mut authenticated = request.clone(); + authenticated.headers.extend(headers); + (authenticated, credential) + } else { + ( + self.auth_provider + .apply_auth(request.clone()) + .await + .map_err(super::direct::direct_auth_error)?, + None, + ) + }; + let prepared = authenticated + .prepare_body_for_send() + .map_err(ExecServerError::EnvironmentRegistryConfig)?; + let rejected_headers = prepared.headers.clone(); + let mut outgoing = self + .http_for_auth() + .request(authenticated.method, authenticated.url) + .headers(prepared.headers) + .headers(current_trace_context_headers()) + .body(prepared.body.unwrap_or_default()); + if let Some(timeout) = authenticated.timeout { + outgoing = outgoing.timeout(timeout); + } + let response = outgoing.send().await?; + if !retried + && response.status() == StatusCode::UNAUTHORIZED + && self + .auth_provider + .refresh_after_unauthorized(rejected_headers) + .await + { + retried = true; + continue; + } + return Ok((response, credential)); + } + } +} + +#[cfg(test)] +#[path = "auth_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/remote/auth_tests.rs b/codex-rs/exec-server/src/remote/auth_tests.rs new file mode 100644 index 0000000000..4c61178303 --- /dev/null +++ b/codex-rs/exec-server/src/remote/auth_tests.rs @@ -0,0 +1,192 @@ +//! Exercise live authentication at the registry request boundary. + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_api::AuthProvider; +use http::HeaderMap; +use http::HeaderValue; +use http::Method; +use http::header::AUTHORIZATION; +use pretty_assertions::assert_eq; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; + +use super::*; +use crate::NoiseChannelIdentity; +use crate::relay::HarnessKeyValidator; +use crate::remote::RegistryHarnessKeyValidator; + +#[derive(Default)] +struct RotatingAuth { + refreshes: AtomicUsize, + stall: Option, + reached: tokio::sync::Notify, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum AuthStage { + Initial, + Refresh, + Retry, +} + +impl RotatingAuth { + async fn stall_at(&self, stage: AuthStage) { + if self.stall == Some(stage) { + self.reached.notify_one(); + std::future::pending::<()>().await; + } + } +} + +impl AuthProvider for RotatingAuth { + fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), codex_api::AuthError> { + let token = if self.refreshes.load(Ordering::SeqCst) == 0 { + "Bearer old" + } else { + "Bearer new" + }; + headers.insert(AUTHORIZATION, HeaderValue::from_static(token)); + Ok(()) + } + + fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { + Box::pin(async { + self.stall_at(if self.refreshes.load(Ordering::SeqCst) == 0 { + AuthStage::Initial + } else { + AuthStage::Retry + }) + .await; + self.to_auth_headers() + }) + } + + fn refresh_after_unauthorized(&self, rejected: HeaderMap) -> codex_api::AuthRefreshFuture<'_> { + assert!(rejected.contains_key(AUTHORIZATION)); + Box::pin(async { + self.stall_at(AuthStage::Refresh).await; + self.refreshes.fetch_add(1, Ordering::SeqCst); + true + }) + } +} + +#[tokio::test] +async fn connect_deadline_covers_initial_auth_refresh_and_retry() -> anyhow::Result<()> { + for stage in [AuthStage::Initial, AuthStage::Refresh, AuthStage::Retry] { + let server = MockServer::start().await; + Mock::given(wiremock::matchers::any()) + .respond_with(ResponseTemplate::new(401)) + .mount(&server) + .await; + let provider = Arc::new(RotatingAuth { + stall: Some(stage), + ..Default::default() + }); + let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; + client.connect_timeout = Duration::from_secs(60); + let key = NoiseChannelIdentity::generate()?.public_key(); + let task = tokio::spawn(async move { client.connect_environment("env", key).await }); + provider.reached.notified().await; + // Advance only after real socket I/O reached the selected auth stage. + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(61)).await; + let result = tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("connect must finish at its original deadline")?; + assert!( + matches!(result, Err(ExecServerError::EnvironmentRegistryRequest(error)) if error.is_timeout()) + ); + tokio::time::resume(); + } + Ok(()) +} + +#[tokio::test] +async fn noise_registry_operations_retry_with_new_credentials() -> anyhow::Result<()> { + for (operation, auth_mode) in [ + ("register", RegistryAuthMode::Legacy), + ("connect", RegistryAuthMode::Legacy), + ("validate", RegistryAuthMode::Legacy), + ("register", RegistryAuthMode::ExecutorCredential), + ("validate", RegistryAuthMode::ExecutorCredential), + ] { + let server = MockServer::start().await; + let key = NoiseChannelIdentity::generate()?.public_key(); + Mock::given(header("authorization", "Bearer old")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(header("authorization", "Bearer new")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "env", "url": "wss://relay.example/ws", + "executor_registration_id": "registration", + "security_profile": "noise_hybrid_ik_v1", + "executor_public_key": key, "harness_key_authorization": "authorization", + "valid": true, + }))) + .expect(1) + .mount(&server) + .await; + let provider = Arc::new(RotatingAuth::default()); + let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; + client.auth_mode = auth_mode; + match operation { + "register" => { + client.register_environment_with_retry("env", &key).await?; + } + "connect" => { + client.connect_environment("env", key).await?; + } + "validate" => { + RegistryHarnessKeyValidator { + client, + environment_id: "env".into(), + executor_registration_id: "registration".into(), + } + .validate_harness_key(&key, "authorization", Some(uuid::Uuid::new_v4())) + .await?; + } + _ => unreachable!(), + } + assert_eq!(provider.refreshes.load(Ordering::SeqCst), 1); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests[0].body, requests[1].body); + assert_eq!( + requests[0].url.path(), + format!("/cloud/environment/env/{operation}") + ); + assert_eq!( + requests[0].headers.get("x-codex-noise-stream-id"), + requests[1].headers.get("x-codex-noise-stream-id") + ); + } + Ok(()) +} + +#[tokio::test] +async fn only_401_retries_and_only_once() -> anyhow::Result<()> { + for (status, attempts, refreshes) in [(401, 2, 1), (403, 1, 0), (500, 1, 0)] { + let server = MockServer::start().await; + Mock::given(wiremock::matchers::any()) + .respond_with(ResponseTemplate::new(status)) + .expect(attempts) + .mount(&server) + .await; + let provider = Arc::new(RotatingAuth::default()); + let client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; + let (response, _) = client + .send_authenticated(Request::new(Method::POST, server.uri())) + .await?; + assert_eq!(response.status().as_u16(), status); + assert_eq!(provider.refreshes.load(Ordering::SeqCst), refreshes); + } + Ok(()) +} diff --git a/codex-rs/exec-server/src/remote/direct.rs b/codex-rs/exec-server/src/remote/direct.rs index f6155afe89..7c68f02a9b 100644 --- a/codex-rs/exec-server/src/remote/direct.rs +++ b/codex-rs/exec-server/src/remote/direct.rs @@ -328,7 +328,7 @@ fn websocket_diagnostic_url(url: &str) -> String { url.split(['?', '#']).next().unwrap_or(url).to_string() } -fn direct_auth_error(error: AuthError) -> ExecServerError { +pub(super) fn direct_auth_error(error: AuthError) -> ExecServerError { let message = format!("failed to resolve environment registry authentication: {error}"); match error { AuthError::Policy(denied) => ExecServerError::ApplicationNetworkPolicy(denied), diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index db0d611914..ff60c13865 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -2485,7 +2485,7 @@ impl AuthManager { let auth = self.auth_cached()?; if let Some(reason) = Self::proactive_refresh_reason(&auth) - && let Err(err) = self.refresh_token_with_reason(reason).await + && let Err(err) = self.refresh_token_with_reason(reason, Some(&auth)).await { tracing::error!("Failed to refresh token: {}", err); return Some(auth); @@ -3094,17 +3094,28 @@ impl AuthManager { Ok(auth) } - /// 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. + /// Refreshes external auth through its provider. For managed auth, first + /// performs a guarded reload: if the stored token differs from the cached + /// token, reuse it; otherwise, ask the token authority to refresh. pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> { - self.refresh_token_with_reason(RefreshReason::Forced).await + self.refresh_token_with_reason(RefreshReason::Forced, /*expected_auth*/ None) + .await + } + + /// Refreshes only if the expected credential is still current after acquiring + /// the existing refresh guard, so queued callers reuse a completed refresh. + pub async fn refresh_token_if_unchanged( + &self, + expected_auth: &CodexAuth, + ) -> Result<(), RefreshTokenError> { + self.refresh_token_with_reason(RefreshReason::Forced, Some(expected_auth)) + .await } async fn refresh_token_with_reason( &self, reason: RefreshReason, + expected_auth: Option<&CodexAuth>, ) -> Result<(), RefreshTokenError> { let _refresh_guard = self.refresh_lock.acquire().await.map_err(|_| { RefreshTokenError::Permanent(RefreshTokenFailedError::new( @@ -3113,6 +3124,14 @@ impl AuthManager { )) })?; let auth_before_reload = self.auth_cached(); + if let Some(expected_auth) = expected_auth + && !Self::auths_equal_for_refresh(auth_before_reload.as_ref(), Some(expected_auth)) + { + return Ok(()); + } + if self.has_external_auth() { + return self.refresh_token_from_authority_impl(reason).await; + } if auth_before_reload .as_ref() .is_some_and(|auth| auth.is_api_key_auth() || auth.is_personal_access_token_auth()) diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 38430cf978..e72fc13cc4 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -168,6 +168,51 @@ impl AuthProvider for AuthManagerAuthProvider { auth_provider_from_auth(&auth).to_auth_headers() }) } + + fn refresh_after_unauthorized( + &self, + rejected_headers: HeaderMap, + ) -> codex_api::AuthRefreshFuture<'_> { + Box::pin(async move { + let Some(auth) = self.current_auth() else { + return false; + }; + let Ok(headers) = auth_provider_from_auth(&auth).to_auth_headers() else { + return false; + }; + if headers.is_empty() { + return false; + } + if headers + .iter() + .any(|(name, value)| rejected_headers.get(name) != Some(value)) + && self.auth_manager.refresh_failure_for_auth(&auth).is_none() + { + return true; + } + if !self.auth_manager.unauthorized_recovery().has_next() { + return false; + } + let result = self.auth_manager.refresh_token_if_unchanged(&auth).await; + let Some(refreshed_auth) = self.current_auth() else { + return false; + }; + if self + .auth_manager + .refresh_failure_for_auth(&refreshed_auth) + .is_some() + { + return false; + } + let Ok(refreshed_headers) = auth_provider_from_auth(&refreshed_auth).to_auth_headers() + else { + return false; + }; + !refreshed_headers.is_empty() + && refreshed_headers != headers + && !matches!(result, Err(codex_login::RefreshTokenError::Policy(_))) + }) + } } // Some providers are meant to send no auth headers. Examples include local OSS @@ -343,6 +388,10 @@ pub fn auth_provider_from_auth_manager( }) } +#[cfg(test)] +#[path = "auth_refresh_tests.rs"] +mod auth_refresh_tests; + #[cfg(test)] mod tests { use codex_agent_identity::generate_agent_key_material; @@ -646,6 +695,10 @@ mod tests { .auth_cached() .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() + .await + .expect("initial auth headers"); assert_eq!( provider.to_auth_headers().unwrap().get(AUTHORIZATION), @@ -659,7 +712,22 @@ mod tests { /*chatgpt_plan_type*/ None, ) .expect("save reloaded auth"); - auth_manager.reload().await; + + // Ordinary requests retain the manager's cache until refresh or recovery runs. + assert_eq!( + provider + .resolve_auth_headers() + .await + .expect("cached auth headers"), + initial_headers + ); + assert_eq!( + tokio::join!( + provider.refresh_after_unauthorized(initial_headers.clone()), + provider.refresh_after_unauthorized(initial_headers), + ), + (true, true) + ); let resolved_headers = provider .resolve_auth_headers() @@ -677,6 +745,7 @@ mod tests { /*chatgpt_plan_type*/ None, ) .expect("save switched-account auth"); + assert!(!provider.refresh_after_unauthorized(resolved_headers).await); auth_manager.reload().await; assert!(matches!( @@ -689,6 +758,104 @@ mod tests { )); } + #[derive(Clone, Copy)] + enum RefreshFailure { + Transient, + Permanent, + Policy, + } + + struct FailingRefreshAuth { + auth: tokio::sync::watch::Receiver, + barrier: tokio::sync::Barrier, + failure: RefreshFailure, + } + + impl codex_login::ExternalAuth for FailingRefreshAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.auth.borrow().clone()) }) + } + + fn refresh( + &self, + _context: codex_login::ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + self.barrier.wait().await; + self.barrier.wait().await; + Err(std::io::Error::other("refresh service unavailable")) + }) + } + + fn classify_error(&self, error: std::io::Error) -> codex_login::RefreshTokenError { + match self.failure { + RefreshFailure::Transient => codex_login::RefreshTokenError::Transient(error), + RefreshFailure::Permanent => codex_login::auth::RefreshTokenFailedError::new( + codex_login::auth::RefreshTokenFailedReason::Other, + error.to_string(), + ) + .into(), + RefreshFailure::Policy => codex_login::RefreshTokenError::Policy( + codex_http_client::NetworkPolicyDenied::Revoked, + ), + } + } + } + + #[tokio::test] + async fn failed_refresh_retries_only_with_newer_credentials() { + for failure in [ + RefreshFailure::Transient, + RefreshFailure::Permanent, + RefreshFailure::Policy, + ] { + for (token, account, changed) in [ + ("Bearer old", "account", false), + ("Bearer new", "account", true), + ("Bearer new", "other-account", false), + ] { + let initial = CodexAuth::Headers(AuthHeaders::new(HeaderMap::from_iter([ + (AUTHORIZATION, HeaderValue::from_static("Bearer old")), + ( + http::HeaderName::from_static("chatgpt-account-id"), + HeaderValue::from_static("account"), + ), + ]))); + let (updates, auth) = tokio::sync::watch::channel(initial.clone()); + let external = Arc::new(FailingRefreshAuth { + auth, + barrier: tokio::sync::Barrier::new(2), + failure, + }); + let manager = AuthManager::from_auth_for_testing(initial.clone()); + manager.set_external_auth(external.clone()).await.unwrap(); + let provider = auth_provider_from_auth_manager(manager.clone(), &initial); + let rejected = provider.to_auth_headers().expect("initial auth headers"); + let recovery = provider.refresh_after_unauthorized(rejected); + let concurrent_update = async { + external.barrier.wait().await; + updates.send_replace(CodexAuth::Headers(AuthHeaders::new( + HeaderMap::from_iter([ + (AUTHORIZATION, HeaderValue::from_static(token)), + ( + http::HeaderName::from_static("chatgpt-account-id"), + HeaderValue::from_static(account), + ), + ]), + ))); + manager.reload().await; + external.barrier.wait().await; + }; + let (retry, ()) = tokio::join!(recovery, concurrent_update); + assert_eq!(retry, changed && !matches!(failure, RefreshFailure::Policy)); + if matches!(failure, RefreshFailure::Permanent) && token == "Bearer old" { + // Even a late rejection must not reuse a credential known to have failed. + assert!(!provider.refresh_after_unauthorized(HeaderMap::new()).await); + } + } + } + } + #[tokio::test] async fn first_party_run_scope_uses_agent_assertion_and_exposes_telemetry() { let auth = CodexAuth::AgentIdentity( diff --git a/codex-rs/model-provider/src/auth_refresh_tests.rs b/codex-rs/model-provider/src/auth_refresh_tests.rs new file mode 100644 index 0000000000..2f5b3ad23a --- /dev/null +++ b/codex-rs/model-provider/src/auth_refresh_tests.rs @@ -0,0 +1,263 @@ +//! Verifies that concurrent proactive and unauthorized requests share one OAuth refresh. + +use super::auth_provider_from_auth_manager; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthHeaders; +use codex_login::AuthKeyringBackendKind; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::RefreshTokenError; +use codex_login::auth::RefreshTokenFailedError; +use codex_login::auth::RefreshTokenFailedReason; +use http::HeaderMap; +use http::HeaderValue; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; +use tempfile::TempDir; +use tokio::sync::Barrier; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test] +async fn concurrent_and_late_rejections_share_one_refresh() { + check_concurrent_refresh(RefreshTrigger::Unauthorized).await; +} + +#[tokio::test] +async fn concurrent_proactive_requests_share_one_refresh() { + check_concurrent_refresh(RefreshTrigger::Proactive).await; +} + +struct RotatingExternalAuth { + refreshes: AtomicUsize, + barrier: Barrier, +} + +impl ExternalAuth for RotatingExternalAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + let token = if self.refreshes.load(Ordering::SeqCst) == 0 { + "Bearer initial" + } else { + "Bearer replacement" + }; + Ok(CodexAuth::Headers(AuthHeaders::new(HeaderMap::from_iter( + [(http::header::AUTHORIZATION, HeaderValue::from_static(token))], + )))) + }) + } + + fn refresh( + &self, + _context: codex_login::ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + if self.refreshes.fetch_add(1, Ordering::SeqCst) != 0 { + return Err(std::io::Error::other("redundant refresh rejected")); + } + self.barrier.wait().await; + self.resolve().await + }) + } + + fn classify_error(&self, error: std::io::Error) -> RefreshTokenError { + RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + error.to_string(), + )) + } +} + +#[tokio::test] +async fn concurrent_external_rejections_share_one_refresh() { + let external = Arc::new(RotatingExternalAuth { + refreshes: AtomicUsize::new(0), + barrier: Barrier::new(2), + }); + let initial = external.resolve().await.unwrap(); + let manager = AuthManager::from_auth_for_testing(initial.clone()); + manager.set_external_auth(external.clone()).await.unwrap(); + let provider = auth_provider_from_auth_manager(manager, &initial); + let rejected = provider.to_auth_headers().unwrap(); + let mut first = provider.refresh_after_unauthorized(rejected.clone()); + let mut second = provider.refresh_after_unauthorized(rejected.clone()); + std::future::poll_fn(|cx| { + assert!(first.as_mut().poll(cx).is_pending()); + assert!(second.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + external.barrier.wait().await; + assert_eq!(tokio::join!(first, second), (true, true)); + 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], + "Bearer replacement" + ); +} + +#[derive(Clone, Copy)] +enum RefreshTrigger { + Unauthorized, + Proactive, +} + +async fn check_concurrent_refresh(trigger: RefreshTrigger) { + let server = MockServer::start().await; + let home = TempDir::new().unwrap(); + let auth = json!({ + "tokens": { + "id_token": "e30.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lkIjoib3JpZ2luYWwtdXNlciIsImNoYXRncHRfYWNjb3VudF9pZCI6Im9yaWdpbmFsLWFjY291bnQifX0.sig", + "access_token": "initial-token", + "refresh_token": "initial-refresh-token", + "account_id": "original-account", + "oauth_provenance": { + "issuer": server.uri(), + "client_id": "concurrent-refresh-test" + } + }, + "last_refresh": match trigger { + RefreshTrigger::Unauthorized => "2099-01-01T00:00:00Z", + RefreshTrigger::Proactive => "2000-01-01T00:00:00Z", + } + }); + std::fs::write( + home.path().join("auth.json"), + serde_json::to_vec(&auth).unwrap(), + ) + .unwrap(); + let manager = Arc::new( + AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + codex_login::test_support::transport_default_auth_route_config(), + ) + .await, + ); + let provider = + auth_provider_from_auth_manager(manager.clone(), &manager.auth_cached().unwrap()); + let rejected_headers = provider.to_auth_headers().unwrap(); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_partial_json( + json!({"refresh_token": "initial-refresh-token"}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "replacement-token", + "refresh_token": "replacement-refresh-token" + }))) + .expect(1) + .mount(&server) + .await; + + // A redundant refresh could permanently fail even though the replacement works. + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_partial_json( + json!({"refresh_token": "replacement-refresh-token"}), + )) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + }))) + .with_priority(10) + .expect(0) + .mount(&server) + .await; + + let refresh = || async { + match trigger { + RefreshTrigger::Unauthorized => { + provider + .refresh_after_unauthorized(rejected_headers.clone()) + .await + } + RefreshTrigger::Proactive => { + let headers = provider.resolve_auth_headers().await.unwrap(); + assert_eq!( + headers[http::header::AUTHORIZATION], + "Bearer replacement-token" + ); + true + } + } + }; + let mut first = Box::pin(refresh()); + let mut second = Box::pin(refresh()); + let mut third = Box::pin(refresh()); + // Queue both waiters before the runtime can service the first OAuth request. + std::future::poll_fn(|cx| { + assert!(first.as_mut().poll(cx).is_pending()); + assert!(second.as_mut().poll(cx).is_pending()); + assert!(third.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert_eq!(tokio::join!(first, second, third), (true, true, true)); + // A delayed 401 for the old token also reuses the completed refresh. + assert!( + provider + .refresh_after_unauthorized(rejected_headers.clone()) + .await + ); + let updated_headers = provider.resolve_auth_headers().await.unwrap(); + assert_eq!( + updated_headers[http::header::AUTHORIZATION], + "Bearer replacement-token" + ); + + // A rejection of the new credential must still be able to refresh again. + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_partial_json( + json!({"refresh_token": "replacement-refresh-token"}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "next-token", + "refresh_token": "next-refresh-token" + }))) + .expect(1) + .mount(&server) + .await; + assert!(provider.refresh_after_unauthorized(updated_headers).await); + assert_eq!( + provider.resolve_auth_headers().await.unwrap()[http::header::AUTHORIZATION], + "Bearer next-token" + ); + // A permanent failure must not prevent a later guarded reload of fresh credentials. + let rejected = provider.to_auth_headers().unwrap(); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_partial_json( + json!({"refresh_token": "next-refresh-token"}), + )) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({"error": "invalid_grant"}))) + .expect(1) + .mount(&server) + .await; + assert!(!provider.refresh_after_unauthorized(rejected.clone()).await); + let path = home.path().join("auth.json"); + let mut stored: serde_json::Value = + serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + stored["tokens"]["access_token"] = json!("reloaded-token"); + std::fs::write(path, serde_json::to_vec(&stored).unwrap()).unwrap(); + assert!(provider.refresh_after_unauthorized(rejected).await); + assert_eq!( + provider.to_auth_headers().unwrap()[http::header::AUTHORIZATION], + "Bearer reloaded-token" + ); + server.verify().await; +}