From 8f32b721518ed23248d81ecf653317182acb3951 Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Tue, 3 Nov 2026 23:08:35 +0000 Subject: [PATCH] Bind reusable backend clients to runtime authentication (#59255) ## Why Reusable clients need fresh authentication for each request while preserving the credential source and transport policy selected when the client was created. ## What changed - Add runtime-aware backend client construction and per-request authentication. Allow token refresh within the original credential source, reject credential replacement, and validate destination and account/organization selectors. - Let runtime authentication finish persisting rotated tokens when the requesting task is cancelled. - Apply asynchronous authentication to file creation and finalization requests, generating fresh assertions while keeping credentials off blob uploads. - Preserve direct sandbox routing without enabling legacy custom-CA fallback for protected clients that reject redirects. - Exclude login account metadata from model cache identities when provider bearer credentials override the login, and record authentication selection outcomes and duration. ## Testing Add regression coverage for credential refresh and replacement, organization selection, destination restrictions, transport policy preservation, cancellation during token rotation, credential stripping on redirects, and file upload authentication in both runtime authentication modes. GitOrigin-RevId: d045595e869dfcdfd9f705a6228cf0eb4c2d2c12 --- codex-rs/Cargo.lock | 7 + codex-rs/app-server/src/turn_cost_worker.rs | 3 +- codex-rs/backend-client/Cargo.toml | 5 + .../backend-client/src/analytics_session.rs | 3 +- .../src/analytics_session_tests.rs | 3 +- codex-rs/backend-client/src/client.rs | 124 +++- .../backend-client/src/client/analytics.rs | 2 +- .../src/client/chatgpt_turn_cost.rs | 2 +- .../backend-client/src/client/plan_history.rs | 4 +- codex-rs/backend-client/src/client/profile.rs | 4 +- .../src/client/rate_limit_resets.rs | 10 +- .../src/client/rate_limit_resets_tests.rs | 3 +- .../src/client/runtime_auth_tests.rs | 568 ++++++++++++++++++ .../backend-client/src/client/task_usage.rs | 2 +- .../backend-client/src/client/thread_usage.rs | 2 +- .../backend-client/src/client/turn_usage.rs | 3 +- .../src/client_request_tests.rs | 4 +- codex-rs/cloud-config/src/config_server.rs | 16 +- codex-rs/cloud-tasks-client/src/http.rs | 10 +- codex-rs/cloud-tasks/src/lib.rs | 2 +- codex-rs/codex-api/src/files.rs | 11 +- .../suite/client_websockets_runtime_auth.rs | 61 ++ .../src/route_aware_client_pool.rs | 23 +- .../src/route_aware_client_pool_tests.rs | 35 +- codex-rs/login/src/auth/default_client.rs | 43 +- .../login/src/auth/default_client_tests.rs | 58 ++ codex-rs/model-provider/src/auth.rs | 56 +- codex-rs/model-provider/src/client_auth.rs | 120 ++++ codex-rs/model-provider/src/lib.rs | 5 + .../src/models_endpoint_refresh_tests.rs | 96 ++- .../model-provider/src/models_identity.rs | 7 + .../src/models_identity_tests.rs | 27 + codex-rs/model-provider/src/request_setup.rs | 62 +- codex-rs/model-provider/src/runtime_auth.rs | 54 +- codex-rs/runtime-agent-instance/Cargo.toml | 2 + codex-rs/runtime-agent-instance/src/lib.rs | 3 + .../runtime-agent-instance/src/manager.rs | 75 ++- .../src/manager_tests.rs | 137 +++++ 38 files changed, 1512 insertions(+), 140 deletions(-) create mode 100644 codex-rs/backend-client/src/client/runtime_auth_tests.rs create mode 100644 codex-rs/model-provider/src/client_auth.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 02632ab011..bae3d87685 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2436,18 +2436,23 @@ name = "codex-backend-client" version = "0.0.0" dependencies = [ "anyhow", + "base64 0.22.1", "codex-api", "codex-backend-openapi-models", + "codex-features", "codex-http-client", "codex-login", "codex-model-provider", "codex-protocol", + "codex-runtime-agent-instance", "http 1.4.0", "pretty_assertions", "serde", "serde_json", + "tempfile", "tokio", "url", + "uuid", "wiremock", ] @@ -4773,6 +4778,7 @@ name = "codex-runtime-agent-instance" version = "0.0.0" dependencies = [ "base64 0.22.1", + "bytes", "codex-agent-identity", "codex-api", "codex-features", @@ -4780,6 +4786,7 @@ dependencies = [ "codex-login", "codex-utils-redacted-string", "ed25519-dalek", + "futures", "http 1.4.0", "pretty_assertions", "rand 0.9.3", diff --git a/codex-rs/app-server/src/turn_cost_worker.rs b/codex-rs/app-server/src/turn_cost_worker.rs index 1aa2c60c3d..197f17f95e 100644 --- a/codex-rs/app-server/src/turn_cost_worker.rs +++ b/codex-rs/app-server/src/turn_cost_worker.rs @@ -454,7 +454,8 @@ impl WorkerRuntime { .map_err(|error| RequestError::Other(error.into()))?; let endpoint = provider.url_for_path("analytics/codex/turn-costs"); let client = BackendClient::new(provider.base_url.clone(), http_client_factory) - .with_auth_provider(auth); + .with_auth_provider(auth) + .map_err(RequestError::Other)?; client .query_api_key_turn_costs_at(&endpoint, turn_ids, &provider.headers) .await diff --git a/codex-rs/backend-client/Cargo.toml b/codex-rs/backend-client/Cargo.toml index 8a8948a9a2..0de0d35336 100644 --- a/codex-rs/backend-client/Cargo.toml +++ b/codex-rs/backend-client/Cargo.toml @@ -20,13 +20,18 @@ http = { workspace = true } url = { workspace = true } codex-backend-openapi-models = { path = "../codex-backend-openapi-models" } codex-api = { workspace = true } +codex-features = { workspace = true } codex-http-client = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-protocol = { workspace = true } tokio = { workspace = true, features = ["time"] } [dev-dependencies] +base64 = { workspace = true } +uuid = { workspace = true } pretty_assertions = "1" +tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } wiremock = { workspace = true } diff --git a/codex-rs/backend-client/src/analytics_session.rs b/codex-rs/backend-client/src/analytics_session.rs index d2ba698b46..d6c00efae6 100644 --- a/codex-rs/backend-client/src/analytics_session.rs +++ b/codex-rs/backend-client/src/analytics_session.rs @@ -55,7 +55,8 @@ impl AnalyticsSession { .with_auth_provider(codex_model_provider::auth_provider_from_auth_manager( Arc::clone(&auth_manager), &auth, - )); + )) + .map_err(|_| "Couldn't initialize local Analytics authentication.")?; let mut session = Self { client, auth_manager, diff --git a/codex-rs/backend-client/src/analytics_session_tests.rs b/codex-rs/backend-client/src/analytics_session_tests.rs index c491cd88de..642ed72764 100644 --- a/codex-rs/backend-client/src/analytics_session_tests.rs +++ b/codex-rs/backend-client/src/analytics_session_tests.rs @@ -87,7 +87,8 @@ async fn assert_unauthorized_after_recovery(outcome: RefreshOutcome, expected_re .with_auth_provider(codex_model_provider::auth_provider_from_auth_manager( Arc::clone(&auth_manager), &auth, - )); + )) + .unwrap(); let session = AnalyticsSession { account: AnalyticsAccount { id: "account-a".into(), diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index c20b2694e2..a411ab9a20 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -9,7 +9,9 @@ use crate::types::RateLimitStatusPayload; use crate::types::TokenUsageProfile; use crate::types::TurnAttemptsSiblingTurnsResponse; use anyhow::Result; +use codex_api::AuthError; use codex_api::SharedAuthProvider; +use codex_features::RuntimeAuthMode; use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; @@ -19,12 +21,14 @@ use codex_http_client::RouteAwareClientPool; use codex_http_client::RouteAwareRequestError; use codex_login::CodexAuth; use codex_login::default_client::get_codex_user_agent; +use codex_model_provider::ClientAuth; use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitReachedType; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::SpendControlLimitSnapshot; +use codex_runtime_agent_instance::RuntimeAuthManager; use http::Method; use http::StatusCode; use http::header::CACHE_CONTROL; @@ -36,6 +40,7 @@ use http::header::USER_AGENT; use serde::Serialize; use serde::de::DeserializeOwned; use std::fmt; +use std::sync::Arc; use std::time::Duration; pub(crate) mod analytics; @@ -113,6 +118,15 @@ impl From for RequestError { } } +impl From for RequestError { + fn from(error: AuthError) -> Self { + match error { + AuthError::Policy(denied) => Self::Policy(denied), + other => Self::Other(other.into()), + } + } +} + impl From for RequestError { fn from(error: codex_http_client::HttpError) -> Self { match error { @@ -161,7 +175,7 @@ impl PathStyle { pub struct Client { base_url: String, http: RouteAwareClientPool, - auth_provider: SharedAuthProvider, + auth: ClientAuth, user_agent: Option, chatgpt_account_id: Option, chatgpt_account_is_fedramp: bool, @@ -173,7 +187,7 @@ impl fmt::Debug for Client { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Client") .field("base_url", &self.base_url) - .field("auth_provider", &"") + .field("auth", &"") .field("user_agent", &self.user_agent) .field("chatgpt_account_id", &self.chatgpt_account_id) .field( @@ -239,7 +253,7 @@ impl Client { Self { base_url, http, - auth_provider: codex_model_provider::unauthenticated_auth_provider(), + auth: ClientAuth::new(codex_model_provider::unauthenticated_auth_provider()), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -253,17 +267,47 @@ impl Client { auth: &CodexAuth, http_client_factory: HttpClientFactory, ) -> Self { - Self::new(base_url, http_client_factory) - .with_user_agent(get_codex_user_agent()) - .with_auth_provider(codex_model_provider::auth_provider_from_auth(auth)) + let mut client = + Self::new(base_url, http_client_factory).with_user_agent(get_codex_user_agent()); + client.auth = ClientAuth::new(codex_model_provider::auth_provider_from_auth(auth)); + client } - pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self { - self.auth_provider = auth; - self + /// Build a client using the captured credential and the request's organization/project headers. + pub async fn from_auth_with_runtime( + base_url: impl Into, + auth: &CodexAuth, + http_client_factory: HttpClientFactory, + runtime: Option<&(Arc, RuntimeAuthMode)>, + headers: &HeaderMap, + ) -> Result { + let base_url = base_url.into(); + let runtime = codex_model_provider::runtime_auth_for_url(runtime, &base_url)?; + Self::from_auth_with_selected_runtime(base_url, auth, http_client_factory, runtime, headers) + .await } - pub fn with_bearer_token(self, token: String) -> Self { + async fn from_auth_with_selected_runtime( + base_url: String, + auth: &CodexAuth, + http_client_factory: HttpClientFactory, + runtime: Option<&(Arc, RuntimeAuthMode)>, + headers: &HeaderMap, + ) -> Result { + let (auth, http_client_factory) = + ClientAuth::for_auth(auth, runtime, http_client_factory, headers).await?; + let mut client = Self::new(base_url, http_client_factory); + client.auth = auth; + Ok(client.with_user_agent(get_codex_user_agent())) + } + + /// Configure provider auth, returning an error for a runtime-bound client. + pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Result { + self.auth.set_provider(auth)?; + Ok(self) + } + + pub fn with_bearer_token(self, token: String) -> Result { self.with_auth_provider(std::sync::Arc::new( codex_model_provider::BearerAuthProvider::new(token), )) @@ -291,14 +335,19 @@ impl Client { self } - fn headers(&self) -> Result { - let mut h = HeaderMap::new(); + async fn headers(&self) -> Result { + self.headers_with(HeaderMap::new()).await + } + + async fn headers_with(&self, mut h: HeaderMap) -> Result { if let Some(ua) = &self.user_agent { h.insert(USER_AGENT, ua.clone()); } else { h.insert(USER_AGENT, HeaderValue::from_static("codex-cli")); } - self.auth_provider.add_auth_headers(&mut h)?; + + let selected = self.auth.resolve(&h).await?; + h = selected.auth.apply_auth_headers(h).await?; if let Some(acc) = &self.chatgpt_account_id && let Ok(hv) = HeaderValue::from_str(acc) { @@ -309,6 +358,9 @@ impl Client { { h.insert(name, HeaderValue::from_static("true")); } + selected + .validate_headers(&h) + .map_err(|error| RequestError::Other(error.into()))?; Ok(h) } @@ -322,6 +374,7 @@ impl Client { method: &str, url: &str, ) -> Result<(String, String)> { + self.auth.validate_destination(&self.base_url, url)?; let res = req.send().await?; let status = res.status(); let ct = res @@ -343,6 +396,9 @@ impl Client { method: &str, url: &str, ) -> std::result::Result<(String, String), RequestError> { + self.auth + .validate_destination(&self.base_url, url) + .map_err(|error| RequestError::Other(error.into()))?; let res = req.send().await?; let status = res.status(); let content_type = res @@ -368,7 +424,9 @@ impl Client { &self, url: &str, ) -> std::result::Result<(String, String), RequestError> { - let request = self.request(Method::GET, url).headers(self.headers()?); + let request = self + .request(Method::GET, url) + .headers(self.headers().await?); if !self.http.allows_system_proxy_fallback() { return self.exec_request_detailed(request, "GET", url).await; } @@ -395,7 +453,7 @@ impl Client { .http .clone() .with_outbound_proxy_policy(OutboundProxyPolicy::RespectSystemProxy); - let request = http.get(url).headers(self.headers()?); + let request = http.get(url).headers(self.headers().await?); self.exec_request_detailed(request, "GET", url).await } @@ -435,7 +493,9 @@ impl Client { pub async fn get_token_usage_profile(&self) -> Result { let url = self.token_usage_profile_url(); - let req = self.request(Method::GET, &url).headers(self.headers()?); + let req = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, ct) = self.exec_request(req, "GET", &url).await?; self.decode_json(&url, &ct, &body) } @@ -454,7 +514,7 @@ impl Client { let url = self.send_add_credits_nudge_email_url(); let req = self .request(Method::POST, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) .json(&SendAddCreditsNudgeEmailRequest { credit_type }); self.exec_request_detailed(req, "POST", &url).await?; @@ -469,7 +529,9 @@ impl Client { cursor: Option<&str>, ) -> Result { let url = self.list_tasks_url(limit, task_filter, environment_id, cursor)?; - let req = self.request(Method::GET, &url).headers(self.headers()?); + let req = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, ct) = self.exec_request(req, "GET", &url).await?; self.decode_json::(&url, &ct, &body) } @@ -521,7 +583,9 @@ impl Client { PathStyle::CodexApi => format!("{}/api/codex/tasks/{}", self.base_url, task_id), PathStyle::ChatGptApi => format!("{}/wham/tasks/{}", self.base_url, task_id), }; - let req = self.request(Method::GET, &url).headers(self.headers()?); + let req = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, ct) = self.exec_request(req, "GET", &url).await?; let parsed: CodeTaskDetailsResponse = self.decode_json(&url, &ct, &body)?; Ok((parsed, body, ct)) @@ -542,7 +606,9 @@ impl Client { self.base_url, task_id, turn_id ), }; - let req = self.request(Method::GET, &url).headers(self.headers()?); + let req = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, ct) = self.exec_request(req, "GET", &url).await?; self.decode_json::(&url, &ct, &body) } @@ -581,7 +647,7 @@ impl Client { let url = self.user_settings_url(); let req = self .request(Method::GET, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .header( CACHE_CONTROL, HeaderValue::from_static("no-cache, no-store"), @@ -597,7 +663,7 @@ impl Client { let url = self.workspace_messages_url(); let req = self .request(Method::GET, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .header(CACHE_CONTROL, HeaderValue::from_static("no-store")); let (body, ct) = self.exec_request_detailed(req, "GET", &url).await?; self.decode_json::(&url, &ct, &body) @@ -613,7 +679,7 @@ impl Client { }; let req = self .request(Method::POST, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) .json(&request_body); let (body, ct) = self.exec_request(req, "POST", &url).await?; @@ -1269,8 +1335,8 @@ mod tests { ); } - #[test] - fn authenticated_user_settings_client_uses_active_workspace_headers() { + #[tokio::test] + async fn authenticated_user_settings_client_uses_active_workspace_headers() { let auth = CodexAuth::from_external_chatgpt_tokens( "e30.e30.c2ln", "workspace-123", @@ -1282,7 +1348,7 @@ mod tests { &auth, HttpClientFactory::new(codex_http_client::OutboundProxyPolicy::ReqwestDefault), ); - let headers = client.headers().unwrap(); + let headers = client.headers().await.unwrap(); assert_eq!( [ @@ -1304,7 +1370,7 @@ mod tests { HttpClientFactory::new(codex_http_client::OutboundProxyPolicy::ReqwestDefault), ClientRouteClass::Api, ), - auth_provider: codex_model_provider::unauthenticated_auth_provider(), + auth: ClientAuth::new(codex_model_provider::unauthenticated_auth_provider()), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -1313,3 +1379,7 @@ mod tests { } } } + +#[cfg(test)] +#[path = "client/runtime_auth_tests.rs"] +mod runtime_auth_tests; diff --git a/codex-rs/backend-client/src/client/analytics.rs b/codex-rs/backend-client/src/client/analytics.rs index be92c63e25..c4b97c291c 100644 --- a/codex-rs/backend-client/src/client/analytics.rs +++ b/codex-rs/backend-client/src/client/analytics.rs @@ -96,7 +96,7 @@ impl Client { } let request = self .request(Method::GET, url.as_str()) - .headers(self.headers()?); + .headers(self.headers().await?); let (body, _) = self .exec_request_detailed(request, "GET", url.as_str()) .await?; diff --git a/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs b/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs index 2379d00ff3..ab3db6bb8c 100644 --- a/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs +++ b/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs @@ -57,7 +57,7 @@ impl Client { }; let request = self .request(Method::POST, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .json(&TurnCostsRequest { threads: threads .iter() diff --git a/codex-rs/backend-client/src/client/plan_history.rs b/codex-rs/backend-client/src/client/plan_history.rs index 4700123a25..aadf2a0eb0 100644 --- a/codex-rs/backend-client/src/client/plan_history.rs +++ b/codex-rs/backend-client/src/client/plan_history.rs @@ -66,7 +66,9 @@ impl Client { PathStyle::ChatGptApi => "wham", }; let url = format!("{}/{prefix}/usage/plan_limit_history?days=7", self.base_url); - let request = self.request(Method::GET, &url).headers(self.headers()?); + let request = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, _) = match self.exec_request_detailed(request, "GET", &url).await { Ok(response) => response, Err(error) if error.status().is_some_and(|status| status.as_u16() == 404) => { diff --git a/codex-rs/backend-client/src/client/profile.rs b/codex-rs/backend-client/src/client/profile.rs index 4117c8d5f7..0c868901c4 100644 --- a/codex-rs/backend-client/src/client/profile.rs +++ b/codex-rs/backend-client/src/client/profile.rs @@ -61,7 +61,9 @@ impl Client { /// Read the full profile using the same route as account token usage. pub async fn get_account_profile(&self) -> Result { let url = self.token_usage_profile_url(); - let request = self.request(Method::GET, &url).headers(self.headers()?); + let request = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, _) = self.exec_request_detailed(request, "GET", &url).await?; serde_json::from_str(&body) .map_err(|_| RequestError::Other(anyhow::anyhow!("Invalid account profile response."))) diff --git a/codex-rs/backend-client/src/client/rate_limit_resets.rs b/codex-rs/backend-client/src/client/rate_limit_resets.rs index 03e7a30481..70a347c20f 100644 --- a/codex-rs/backend-client/src/client/rate_limit_resets.rs +++ b/codex-rs/backend-client/src/client/rate_limit_resets.rs @@ -71,7 +71,9 @@ impl Client { supports_luna_reserve: bool, ) -> Result { let url = self.rate_limit_status_url(); - let mut req = self.request(Method::GET, &url).headers(self.headers()?); + let mut req = self + .request(Method::GET, &url) + .headers(self.headers().await?); if supports_luna_reserve { req = req.header("x-openai-codex-luna-reserve", HeaderValue::from_static("1")); } @@ -81,7 +83,9 @@ impl Client { pub async fn list_rate_limit_reset_credits(&self) -> Result { let url = self.rate_limit_reset_credits_url(); - let req = self.request(Method::GET, &url).headers(self.headers()?); + let req = self + .request(Method::GET, &url) + .headers(self.headers().await?); let (body, ct) = self.exec_request(req, "GET", &url).await?; self.decode_json(&url, &ct, &body) } @@ -111,7 +115,7 @@ impl Client { let url = self.consume_rate_limit_reset_credit_url(); let req = self .request(Method::POST, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) .json(&ConsumeRateLimitResetCreditRequest { redeem_request_id, diff --git a/codex-rs/backend-client/src/client/rate_limit_resets_tests.rs b/codex-rs/backend-client/src/client/rate_limit_resets_tests.rs index 4c4470fedf..fdee7e7c9c 100644 --- a/codex-rs/backend-client/src/client/rate_limit_resets_tests.rs +++ b/codex-rs/backend-client/src/client/rate_limit_resets_tests.rs @@ -5,6 +5,7 @@ use crate::types::ConsumeRateLimitResetCreditCode; use crate::types::RateLimitResetCreditDetails; use crate::types::RateLimitResetCreditsDetails; use crate::types::RateLimitResetCreditsSummary; +use codex_model_provider::ClientAuth; use pretty_assertions::assert_eq; use wiremock::Mock; use wiremock::MockServer; @@ -178,7 +179,7 @@ fn test_client(base_url: &str, path_style: PathStyle) -> Client { ), codex_http_client::ClientRouteClass::Api, ), - auth_provider: codex_model_provider::unauthenticated_auth_provider(), + auth: ClientAuth::new(codex_model_provider::unauthenticated_auth_provider()), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, diff --git a/codex-rs/backend-client/src/client/runtime_auth_tests.rs b/codex-rs/backend-client/src/client/runtime_auth_tests.rs new file mode 100644 index 0000000000..4cbe176dc2 --- /dev/null +++ b/codex-rs/backend-client/src/client/runtime_auth_tests.rs @@ -0,0 +1,568 @@ +//! Verify the backend client resolves a fresh runtime proof for each HTTP request. + +use super::*; +use codex_features::RuntimeAuthMode; +use codex_http_client::DestinationPolicy; +use codex_http_client::NetworkPolicyController; +use codex_http_client::NetworkPolicyDenied; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; +use codex_login::AuthManager; +use codex_login::AuthRouteConfig; +use codex_protocol::error::CodexErr; +use codex_protocol::error::CodexErrorDetails; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use uuid::Uuid; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +fn runtime_provider( + server: &MockServer, + credentials: Arc, + mode: RuntimeAuthMode, +) -> (Arc, RuntimeAuthMode) { + ( + Arc::new(RuntimeAuthManager::new( + credentials, + AgentHostAuthorization::InstallationId { + installation_id: Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/register", server.uri()).parse().unwrap(), + )), + mode, + ) +} + +#[tokio::test] +async fn account_runtime_auth_signs_each_request_and_strips_credentials_on_redirect() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let destination = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().unwrap(); + let organization = body["actor_auth"]["organization_id"].as_str().unwrap_or("org-a"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": format!("runtime-{organization}"), "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": organization, "project_id": null}, + })) + }) + .expect(2).mount(&server).await; + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"accounts": [], "account_ordering": []})), + ) + .expect(2) + .mount(&destination) + .await; + Mock::given(path("/turn-costs")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(307).insert_header("location", destination.uri())) + .expect(2) + .mount(&server) + .await; + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let auth = CodexAuth::from_api_key("key-a"); + let runtime = runtime_provider( + &server, + AuthManager::from_auth_for_testing(auth.clone()), + mode, + ); + let client = Client::from_auth_with_selected_runtime( + server.uri(), + &auth, + factory, + Some(&runtime), + &HeaderMap::new(), + ) + .await + .unwrap(); + client.get_accounts_check().await.unwrap(); + client.get_accounts_check().await.unwrap(); + assert!( + client + .clone() + .with_chatgpt_account_id("different-workspace") + .get_accounts_check() + .await + .is_err() + ); + let stale_headers = HeaderMap::from_iter([ + ( + HeaderName::from_static("openai-organization"), + HeaderValue::from_static("org-b"), + ), + ( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer stale"), + ), + ( + HeaderName::from_static("x-openai-agent-assertion"), + HeaderValue::from_static("stale"), + ), + ]); + client + .query_api_key_turn_costs_at( + &format!("{}/turn-costs", server.uri()), + &[], + &stale_headers, + ) + .await + .unwrap(); + assert!( + client + .query_api_key_turn_costs_at(&destination.uri(), &[], &stale_headers) + .await + .is_err() + ); + let redirected = destination.received_requests().await.unwrap(); + assert_eq!(redirected.len(), 2); + for request in redirected { + assert!(!request.headers.contains_key("authorization")); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + let requests = server.received_requests().await.unwrap(); + let cost_request = requests + .iter() + .find(|request| request.url.path() == "/turn-costs") + .unwrap(); + assert_eq!(cost_request.headers["openai-organization"], "org-b"); + let proofs: Vec<_> = requests + .iter() + .filter(|r| r.url.path() != "/register") + .map(|request| { + let authorization = request.headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + request.headers["x-openai-agent-assertion"] + .to_str() + .unwrap() + } + RuntimeAuthMode::AssertionOnly => { + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + authorization.strip_prefix("AgentAssertion ").unwrap() + } + } + }) + .collect(); + assert_eq!(proofs.len(), 3); + assert_ne!(proofs[0], proofs[1]); + assert_ne!(proofs[1], proofs[2]); + assert_ne!(proofs[2], "stale"); + } +} + +#[tokio::test] +async fn client_registers_in_requested_organization_when_default_is_denied() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().unwrap(); + if body["actor_auth"]["organization_id"] != "org-b" { + return ResponseTemplate::new(403); + } + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-org-b", + "agent_host_id": "host-a", + "actor": { + "kind": "platform", + "user_id": "user-a", + "organization_id": "org-b", + "project_id": null, + }, + })) + }) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/turn-costs")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})), + ) + .expect(1) + .mount(&server) + .await; + let auth = CodexAuth::from_api_key("key-a"); + let runtime = runtime_provider( + &server, + AuthManager::from_auth_for_testing(auth.clone()), + mode, + ); + let headers = HeaderMap::from_iter([( + HeaderName::from_static("openai-organization"), + HeaderValue::from_static("org-b"), + )]); + let client = Client::from_auth_with_selected_runtime( + server.uri(), + &auth, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(&runtime), + &headers, + ) + .await + .unwrap(); + client + .query_api_key_turn_costs_at(&format!("{}/turn-costs", server.uri()), &[], &headers) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[1].url.path(), "/turn-costs"); + assert_eq!(requests[1].headers["openai-organization"], "org-b"); + let authorization = requests[1].headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + assert!(requests[1].headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!(authorization.starts_with("AgentAssertion ")); + assert!(!requests[1].headers.contains_key("x-openai-agent-assertion")); + } + } + server.verify().await; + } +} + +#[tokio::test] +async fn unsupported_openai_backend_rejects_runtime_auth_before_registration() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let auth = CodexAuth::from_api_key("key-a"); + let runtime = runtime_provider( + &server, + AuthManager::from_auth_for_testing(auth.clone()), + mode, + ); + let result = Client::from_auth_with_runtime( + "https://future.openai.com/backend-api", + &auth, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(&runtime), + &HeaderMap::new(), + ) + .await; + let Err(error) = result else { + panic!("unsupported OpenAI backend accepted runtime auth"); + }; + assert!(matches!( + error.downcast_ref::().map(CodexErr::details), + Some(CodexErrorDetails::UnsupportedOperation(_)) + )); + assert!(server.received_requests().await.unwrap().is_empty()); + } +} + +#[tokio::test] +async fn custom_backend_keeps_its_credential_without_registering_a_runtime() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(path("/turn-costs")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})), + ) + .expect(1) + .mount(&server) + .await; + let auth = CodexAuth::from_api_key("custom-backend-key"); + let runtime = runtime_provider( + &server, + AuthManager::from_auth_for_testing(auth.clone()), + mode, + ); + let client = Client::from_auth_with_runtime( + server.uri(), + &auth, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(&runtime), + &HeaderMap::new(), + ) + .await + .unwrap(); + client + .query_api_key_turn_costs_at( + &format!("{}/turn-costs", server.uri()), + &[], + &HeaderMap::new(), + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].headers["authorization"], + "Bearer custom-backend-key" + ); + assert!(!requests[0].headers.contains_key("x-openai-agent-assertion")); + server.verify().await; + } +} + +#[tokio::test] +async fn reusable_runtime_client_keeps_the_selected_accounts_transport_policy() { + let server = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-policy", "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + }))) + .expect(1).mount(&server).await; + Mock::given(path("/turn-costs")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []}))) + .expect(2) + .mount(&server) + .await; + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + let home = tempfile::tempdir().unwrap(); + codex_login::login_with_api_key( + home.path(), + "key-a", + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::Direct, + ) + .unwrap(); + let credentials = Arc::new( + AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::Ephemeral, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + AuthRouteConfig::from_http_client_factory( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy.clone()), + ), + ) + .await, + ); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let auth = credentials.auth().await.unwrap(); + let runtime = runtime_provider(&server, credentials, RuntimeAuthMode::AssertionOnly); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let client = Client::from_auth_with_selected_runtime( + server.uri(), + &auth, + factory.clone(), + Some(&runtime), + &HeaderMap::new(), + ) + .await + .unwrap(); + let url = format!("{}/turn-costs", server.uri()); + client + .query_api_key_turn_costs_at(&url, &[], &HeaderMap::new()) + .await + .unwrap(); + + policy.invalidate(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let error = client + .query_api_key_turn_costs_at(&url, &[], &HeaderMap::new()) + .await + .unwrap_err(); + assert!(matches!( + error, + RequestError::Policy(NetworkPolicyDenied::Revoked) + )); + + let fresh = Client::from_auth_with_selected_runtime( + server.uri(), + &auth, + factory, + Some(&runtime), + &HeaderMap::new(), + ) + .await + .unwrap(); + fresh + .query_api_key_turn_costs_at(&url, &[], &HeaderMap::new()) + .await + .unwrap(); + + policy.invalidate(); + assert!(controller.publish( + policy.revision(), + DestinationPolicy::Restricted { + allowed_hosts: Default::default() + } + )); + let selectors = HeaderMap::from_iter([( + HeaderName::from_static("openai-organization"), + HeaderValue::from_static("org-b"), + )]); + let error = fresh + .query_api_key_turn_costs_at(&url, &[], &selectors) + .await + .unwrap_err(); + assert!( + matches!( + error, + RequestError::Policy(NetworkPolicyDenied::Destination) + ), + "{error}" + ); + server.verify().await; +} + +struct RotatingAuth(std::sync::RwLock); + +impl codex_login::ExternalAuth for RotatingAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.read().unwrap().clone()) }) + } + + fn refresh( + &self, + _: codex_login::ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +#[tokio::test] +async fn reusable_runtime_client_refreshes_credentials_without_accepting_a_new_session() { + use base64::Engine; + + let oauth = |session, token| { + let claims = serde_json::json!({ + "iss": "https://auth.openai.com/", "session_id": session, "jti": token, + "https://api.openai.com/auth": { + "user_id": "user-a", "chatgpt_account_id": "workspace-a" + } + }); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + CodexAuth::from_external_chatgpt_tokens( + &format!("header.{payload}.signature"), + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap() + }; + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-refresh", "agent_host_id": "host-a", + "actor": {"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"} + }))) + .mount(&server) + .await; + Mock::given(path("/api/codex/usage")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "plan_type": "plus", + "rate_limit": { + "allowed": true, "limit_reached": false, + "primary_window": { + "used_percent": 0, "limit_window_seconds": 300, + "reset_after_seconds": 60, "reset_at": 2000000000 + } + } + }))) + .expect(2) + .mount(&server) + .await; + let initial = oauth("login-a", "initial"); + let refreshed = oauth("login-a", "refreshed"); + let credentials = AuthManager::from_auth_for_testing(initial.clone()); + let source = Arc::new(RotatingAuth(std::sync::RwLock::new(initial.clone()))); + credentials.set_external_auth(source.clone()).await.unwrap(); + let runtime = runtime_provider(&server, credentials, mode); + let client = Client::from_auth_with_selected_runtime( + server.uri(), + &initial, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(&runtime), + &HeaderMap::new(), + ) + .await + .unwrap(); + + assert!( + client + .clone() + .with_bearer_token("replacement-bearer".into()) + .is_err() + ); + assert!( + client + .clone() + .with_auth_provider(Arc::new(codex_model_provider::BearerAuthProvider::new( + "replacement-provider".to_string(), + ))) + .is_err() + ); + client.get_rate_limits().await.unwrap(); + *source.0.write().unwrap() = refreshed.clone(); + client.get_rate_limits().await.unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|r| r.url.path() == "/register") + .count(), + 1 + ); + let requests: Vec<_> = requests + .iter() + .filter(|r| r.url.path() == "/api/codex/usage") + .collect(); + assert_eq!(requests.len(), 2); + let mut proofs = Vec::new(); + for (request, expected) in requests.iter().zip([initial, refreshed]) { + let authorization = request.headers["authorization"].to_str().unwrap(); + let proof = match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!( + authorization, + format!("Bearer {}", expected.get_token().unwrap()) + ); + request.headers["x-openai-agent-assertion"] + .to_str() + .unwrap() + } + RuntimeAuthMode::AssertionOnly => { + authorization.strip_prefix("AgentAssertion ").unwrap() + } + }; + let claims: serde_json::Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(proof.split('.').nth(2).unwrap()) + .unwrap(), + ) + .unwrap(); + assert_eq!(claims["runtime_agent_instance"], "runtime-refresh"); + proofs.push(proof); + } + assert_ne!(proofs[0], proofs[1]); + + *source.0.write().unwrap() = oauth("another-login", "replacement"); + assert!(client.get_rate_limits().await.is_err()); + } +} diff --git a/codex-rs/backend-client/src/client/task_usage.rs b/codex-rs/backend-client/src/client/task_usage.rs index cbb38cbe0f..f6aabf2c9e 100644 --- a/codex-rs/backend-client/src/client/task_usage.rs +++ b/codex-rs/backend-client/src/client/task_usage.rs @@ -103,7 +103,7 @@ impl Client { } let request = self .request(Method::POST, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .json(&Query { threads }); let (body, _) = self.exec_request_detailed(request, "POST", &url).await?; let response: TaskUsageResponse = serde_json::from_str(&body) diff --git a/codex-rs/backend-client/src/client/thread_usage.rs b/codex-rs/backend-client/src/client/thread_usage.rs index fb899a837a..371cec21b2 100644 --- a/codex-rs/backend-client/src/client/thread_usage.rs +++ b/codex-rs/backend-client/src/client/thread_usage.rs @@ -82,7 +82,7 @@ impl Client { let url = self.thread_usage_url(); let request = self .request(Method::POST, &url) - .headers(self.headers()?) + .headers(self.headers().await?) .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) .json(&ThreadUsageQueryRequest { thread_ids }); let (body, _) = self.exec_request_detailed(request, "POST", &url).await?; diff --git a/codex-rs/backend-client/src/client/turn_usage.rs b/codex-rs/backend-client/src/client/turn_usage.rs index ecc8fadd23..86559d93b0 100644 --- a/codex-rs/backend-client/src/client/turn_usage.rs +++ b/codex-rs/backend-client/src/client/turn_usage.rs @@ -85,8 +85,7 @@ impl Client { turn_ids: &[String], provider_headers: &HeaderMap, ) -> Result, RequestError> { - let mut headers = provider_headers.clone(); - headers.extend(self.headers()?); + let headers = self.headers_with(provider_headers.clone()).await?; let request = self .request(Method::POST, url) .headers(headers) diff --git a/codex-rs/backend-client/src/client_request_tests.rs b/codex-rs/backend-client/src/client_request_tests.rs index 0c0943699c..210ce43319 100644 --- a/codex-rs/backend-client/src/client_request_tests.rs +++ b/codex-rs/backend-client/src/client_request_tests.rs @@ -119,7 +119,8 @@ async fn migrated_requests_preserve_query_auth_and_json_body() { ) .with_auth_provider(Arc::new(codex_model_provider::BearerAuthProvider::new( "request-token".to_string(), - ))); + ))) + .unwrap(); let tasks = client .list_tasks( @@ -195,6 +196,7 @@ fn fallback_client(base_url: &str) -> Client { .with_auth_provider(Arc::new(codex_model_provider::BearerAuthProvider::new( "bootstrap-token".to_string(), ))) + .unwrap() .with_chatgpt_account_id("workspace-123") .with_user_agent("bootstrap-test") } diff --git a/codex-rs/cloud-config/src/config_server.rs b/codex-rs/cloud-config/src/config_server.rs index 4342273d2a..2dba54adf2 100644 --- a/codex-rs/cloud-config/src/config_server.rs +++ b/codex-rs/cloud-config/src/config_server.rs @@ -191,13 +191,15 @@ impl BundleSource for ConfigServerSource { }); } *identity = credentials.authorization_id; - fetch_bundle( - &self - .client - .clone() - .with_bearer_token(credentials.access_token), - ) - .await + let client = self + .client + .clone() + .with_bearer_token(credentials.access_token) + .map_err(|_| BundleRequestError::Fatal { + error: auth_failure(), + metric_reason: Some("auth"), + })?; + fetch_bundle(&client).await } async fn recover_auth( diff --git a/codex-rs/cloud-tasks-client/src/http.rs b/codex-rs/cloud-tasks-client/src/http.rs index d2edcdbe63..c70ae596e8 100644 --- a/codex-rs/cloud-tasks-client/src/http.rs +++ b/codex-rs/cloud-tasks-client/src/http.rs @@ -42,9 +42,13 @@ impl HttpClient { self } - pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self { - self.backend = self.backend.clone().with_auth_provider(auth); - self + pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Result { + self.backend = self + .backend + .clone() + .with_auth_provider(auth) + .map_err(|error| CloudTaskError::Msg(error.to_string()))?; + Ok(self) } pub fn with_chatgpt_account_id(mut self, account_id: impl Into) -> Self { diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index b0975fe929..c3f43f3e19 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -111,7 +111,7 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result } let auth_provider = codex_model_provider::auth_provider_from_auth(&auth); - http = http.with_auth_provider(auth_provider); + http = http.with_auth_provider(auth_provider)?; if let Some(acc) = auth.get_account_id() { append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}")); } diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index 0a37b2753a..3d2376bb52 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -8,6 +8,7 @@ use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; use codex_http_client::RouteAwareRequestError; use futures::Stream; +use http::HeaderMap; use http::Method; use http::StatusCode; use http::header::CONTENT_LENGTH; @@ -154,7 +155,7 @@ where "file_size": file_size_bytes, "use_case": OPENAI_FILE_USE_CASE, }); - let request = authorized_request(client_pool, auth, Method::POST, &create_url)?; + let request = authorized_request(client_pool, auth, Method::POST, &create_url).await?; let create_request = match hosted_upload { Some(context) => serde_json::json!({ "file_name": file_name.as_str(), @@ -365,7 +366,8 @@ where }; let finalize_started_at = Instant::now(); loop { - let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url)? + let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url) + .await? .json(&finalize_request) .send() .await @@ -432,14 +434,13 @@ where } } -fn authorized_request( +async fn authorized_request( client_pool: &RouteAwareClientPool, auth: &dyn AuthProvider, method: Method, url: &str, ) -> Result { - let mut headers = http::HeaderMap::new(); - auth.add_auth_headers(&mut headers)?; + let headers = auth.apply_auth_headers(HeaderMap::new()).await?; Ok(client_pool .request(method, url) diff --git a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs index 53041d393e..1ebdd5486e 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -499,3 +499,64 @@ async fn gateway_refresh_failure_blocks_runtime_inference(mode: RuntimeAuthMode) .all(|request| matches!(request.url.path(), "/register" | "/token")) ); } + +#[test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn model_runtime_auth_strips_credentials_on_cross_origin_redirects(mode: RuntimeAuthMode) { + skip_if_no_network!(); + let server = MockServer::start().await; + let destination = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(registered) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/v1/responses")) + .respond_with(ResponseTemplate::new(307).insert_header("location", destination.uri())) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(responses::sse(vec![ + ev_response_created("resp-redirected"), + ev_completed("resp-redirected"), + ])), + ) + .expect(1) + .mount(&destination) + .await; + let mut provider = + ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); + provider.supports_websockets = false; + let (harness, _) = runtime_harness(&server, provider, mode).await; + complete(&mut harness.client.new_session(), &harness) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let request = requests + .iter() + .find(|r| r.url.path() == "/v1/responses") + .unwrap(); + assertion_payload( + request.headers["authorization"].to_str().unwrap(), + request + .headers + .get("x-openai-agent-assertion") + .map(|v| v.to_str().unwrap().to_string()), + mode, + "key-a", + ); + let redirected = destination.received_requests().await.unwrap(); + assert_eq!(redirected.len(), 1); + assert_eq!(redirected[0].body, request.body); + assert!(!redirected[0].headers.contains_key("authorization")); + assert!( + !redirected[0] + .headers + .contains_key("x-openai-agent-assertion") + ); +} diff --git a/codex-rs/http-client/src/route_aware_client_pool.rs b/codex-rs/http-client/src/route_aware_client_pool.rs index 122a982434..fe39417e69 100644 --- a/codex-rs/http-client/src/route_aware_client_pool.rs +++ b/codex-rs/http-client/src/route_aware_client_pool.rs @@ -56,6 +56,7 @@ pub struct RouteAwareClientPool { pub(crate) client_builder: HttpClientBuilder, default_headers: HeaderMap, custom_ca_fallback: CustomCaFallback, + direct_proxy: bool, clients: Arc>>, client_build: Arc>, rustls_clients: Option, @@ -289,6 +290,7 @@ impl RouteAwareClientPool { client_builder, default_headers, custom_ca_fallback: CustomCaFallback::Disabled, + direct_proxy: false, clients: Arc::new(Mutex::new(HashMap::new())), client_build: Arc::default(), rustls_clients: None, @@ -357,6 +359,12 @@ impl RouteAwareClientPool { self } + /// Uses direct routing without changing destination policy or TLS configuration. + pub fn with_direct_proxy(mut self) -> Self { + self.direct_proxy = true; + self + } + /// Preserves the legacy sandbox client's direct routing and custom-CA fallback. pub fn with_legacy_direct_proxy_and_custom_ca_fallback(mut self) -> Self { self.custom_ca_fallback = CustomCaFallback::LegacyDirect; @@ -478,13 +486,14 @@ impl RouteAwareClientPool { F: FnOnce(String) -> Fut, Fut: Future>, { - let route = if self.custom_ca_fallback == CustomCaFallback::LegacyDirect { - OutboundProxyRoute::Direct - } else { - resolve_route(request_url.to_string()) - .await - .map_err(RouteAwareClientPoolError::Resolve)? - }; + let route = + if self.direct_proxy || self.custom_ca_fallback == CustomCaFallback::LegacyDirect { + OutboundProxyRoute::Direct + } else { + resolve_route(request_url.to_string()) + .await + .map_err(RouteAwareClientPoolError::Resolve)? + }; if let Some(rustls_clients) = self.rustls_clients.as_ref() && let Ok(url) = reqwest::Url::parse(request_url) && rustls_clients.requires_rustls(&url, &route) diff --git a/codex-rs/http-client/src/route_aware_client_pool_tests.rs b/codex-rs/http-client/src/route_aware_client_pool_tests.rs index 5d6512a2cf..d02f8814a1 100644 --- a/codex-rs/http-client/src/route_aware_client_pool_tests.rs +++ b/codex-rs/http-client/src/route_aware_client_pool_tests.rs @@ -377,7 +377,12 @@ async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() { .expect("invalid CA fixture should be written"); for ca_env in ["CODEX_CA_CERTIFICATE", "SSL_CERT_FILE"] { - for policy_name in ["reqwest-default", "respect-system-proxy"] { + for policy_name in [ + "reqwest-default", + "respect-system-proxy", + "strict-reqwest-default", + "strict-direct", + ] { let output = std::process::Command::new( std::env::current_exe().expect("test executable should be available"), ) @@ -403,18 +408,24 @@ async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() { }; let outbound_proxy_policy = match policy_name.as_str() { - "reqwest-default" => OutboundProxyPolicy::ReqwestDefault, + "reqwest-default" | "strict-reqwest-default" | "strict-direct" => { + OutboundProxyPolicy::ReqwestDefault + } "respect-system-proxy" => OutboundProxyPolicy::RespectSystemProxy, _ => panic!("unexpected test proxy policy: {policy_name}"), }; let pool = RouteAwareClientPool::with_chatgpt_cloudflare_cookies( HttpClientFactory::new(outbound_proxy_policy), ClientRouteClass::Other, - ) - .with_legacy_custom_ca_fallback(); + ); + let pool = match policy_name.as_str() { + "strict-reqwest-default" => pool, + "strict-direct" => pool.with_direct_proxy(), + _ => pool.with_legacy_custom_ca_fallback(), + }; - match outbound_proxy_policy { - OutboundProxyPolicy::ReqwestDefault => { + match policy_name.as_str() { + "reqwest-default" => { let (address, server) = spawn_response_server(vec![ "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string(), ]); @@ -428,13 +439,19 @@ async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() { let requests = server.join().expect("response server should finish"); assert_eq!(requests.len(), 1); } - OutboundProxyPolicy::RespectSystemProxy => { + _ => { let error = pool .client_for_url_with_resolver("http://127.0.0.1/update", |_| async { - Ok(OutboundProxyRoute::Direct) + if policy_name == "strict-direct" { + Err(io::Error::other( + "direct routing must not need proxy discovery", + )) + } else { + Ok(OutboundProxyRoute::Direct) + } }) .await - .expect_err("system-proxy routes should reject invalid custom CAs"); + .expect_err("strict routes should reject invalid custom CAs"); assert!(matches!(error, RouteAwareClientPoolError::Build(_))); } diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index 42d023e339..9d8af7f09a 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -263,11 +263,19 @@ pub fn create_client_for_route( route_class: ClientRouteClass, redirect_policy: ClientRedirectPolicy, ) -> Result { + let legacy_fallback = redirect_policy == ClientRedirectPolicy::Default + || !http_client_factory.protects_sensitive_headers(); let builder = match redirect_policy { ClientRedirectPolicy::Default => default_http_client_builder(), ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(), }; - create_client_for_route_with_builder(http_client_factory, request_url, route_class, builder) + create_client_for_route_with_builder( + http_client_factory, + request_url, + route_class, + builder, + legacy_fallback, + ) } fn create_client_for_route_with_builder( @@ -275,9 +283,11 @@ fn create_client_for_route_with_builder( request_url: &str, route_class: ClientRouteClass, builder: HttpClientBuilder, + legacy_fallback: bool, ) -> Result { if http_client_factory.network_policy().is_managed() || http_client_factory.protects_sensitive_headers() + || !legacy_fallback { let mut pool = codex_http_client::RouteAwareClientPool::with_builder( http_client_factory.clone(), @@ -285,11 +295,17 @@ fn create_client_for_route_with_builder( builder, ); if is_sandboxed() { - pool = pool.with_legacy_direct_proxy_and_custom_ca_fallback(); - } else if matches!( - http_client_factory.outbound_proxy_policy(), - OutboundProxyPolicy::ReqwestDefault - ) { + pool = if legacy_fallback { + pool.with_legacy_direct_proxy_and_custom_ca_fallback() + } else { + pool.with_direct_proxy() + }; + } else if legacy_fallback + && matches!( + http_client_factory.outbound_proxy_policy(), + OutboundProxyPolicy::ReqwestDefault + ) + { pool = pool.with_legacy_custom_ca_fallback(); } return Ok(pool.into_client()); @@ -324,6 +340,8 @@ pub async fn create_client_for_route_async( route_class: ClientRouteClass, redirect_policy: ClientRedirectPolicy, ) -> std::io::Result { + let legacy_fallback = redirect_policy == ClientRedirectPolicy::Default + || !http_client_factory.protects_sensitive_headers(); create_client_for_route_async_with_builder( http_client_factory, request_url, @@ -332,6 +350,7 @@ pub async fn create_client_for_route_async( ClientRedirectPolicy::Default => default_http_client_builder(), ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(), }, + legacy_fallback, ) .await } @@ -347,6 +366,7 @@ pub async fn create_client_for_route_without_request_logging_async( request_url, route_class, || default_http_client_builder().without_request_logging(), + /*legacy_fallback*/ true, ) .await } @@ -356,6 +376,7 @@ async fn create_client_for_route_async_with_builder( request_url: String, route_class: ClientRouteClass, builder: impl FnOnce() -> HttpClientBuilder + Send + 'static, + legacy_fallback: bool, ) -> std::io::Result { let permit = ROUTE_AWARE_CLIENT_BUILD_PERMIT .acquire() @@ -368,6 +389,7 @@ async fn create_client_for_route_async_with_builder( &request_url, route_class, builder(), + legacy_fallback, ) .map_err(std::io::Error::from) }) @@ -413,8 +435,13 @@ pub async fn create_transport_for_routes_async( }, ); if sandboxed { - pool = pool.with_legacy_direct_proxy_and_custom_ca_fallback(); - } else if transport_default_proxy { + pool = match redirect_policy { + ClientRedirectPolicy::Default => { + pool.with_legacy_direct_proxy_and_custom_ca_fallback() + } + ClientRedirectPolicy::Reject => pool.with_direct_proxy(), + }; + } else if transport_default_proxy && redirect_policy == ClientRedirectPolicy::Default { pool = pool.with_legacy_custom_ca_fallback(); } ReqwestTransport::from_route_aware_client_pool(pool) diff --git a/codex-rs/login/src/auth/default_client_tests.rs b/codex-rs/login/src/auth/default_client_tests.rs index 4a5d9410d8..5adbe75fed 100644 --- a/codex-rs/login/src/auth/default_client_tests.rs +++ b/codex-rs/login/src/auth/default_client_tests.rs @@ -355,3 +355,61 @@ fn test_macos() { .unwrap(); assert!(re.is_match(&user_agent)); } + +#[tokio::test] +async fn signed_reject_routes_preserve_invalid_ca_errors() { + const CHILD: &str = "CODEX_TEST_SIGNED_ROUTE_INVALID_CA"; + if std::env::var_os(CHILD).is_none() { + let dir = tempfile::tempdir().unwrap(); + let ca = dir.path().join("invalid.pem"); + std::fs::write(&ca, "invalid certificate").unwrap(); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .arg("signed_reject_routes_preserve_invalid_ca_errors") + .arg("--nocapture") + .env(CHILD, "1") + .env("CODEX_CA_CERTIFICATE", ca) + .env_remove("SSL_CERT_FILE") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(String::from_utf8_lossy(&output.stdout).contains("1 passed")); + return; + } + + let url = "https://127.0.0.1:9/"; + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_sensitive_header_protection(); + let sync_client = create_client_for_route( + &factory, + url, + ClientRouteClass::Api, + ClientRedirectPolicy::Reject, + ) + .unwrap(); + let async_client = create_client_for_route_async( + factory, + url.into(), + ClientRouteClass::Api, + ClientRedirectPolicy::Reject, + ) + .await + .unwrap(); + for client in [sync_client, async_client] { + let error = client.get(url).send().await.unwrap_err(); + assert!( + matches!( + error, + codex_http_client::HttpError::Route( + codex_http_client::RouteAwareClientPoolError::Build( + codex_http_client::BuildRouteAwareHttpClientError::CustomCa(_) + ) + ) + ), + "{error}" + ); + } +} diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 8d3e407d2e..84a6baa913 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -440,7 +440,61 @@ pub(crate) fn uses_fedramp_api_key( && auth.is_some_and(CodexAuth::is_fedramp_account)) } -/// Builds request-header auth for a first-party Codex auth snapshot. +/// Select the shared runtime without changing the identity of captured request data. +pub async fn auth_provider_for_request( + auth: &CodexAuth, + runtime: Option<&(Arc, RuntimeAuthMode)>, + headers: &HeaderMap, +) -> Result { + observe_auth_selection(runtime.map(|(_, mode)| *mode), "request", async { + match runtime { + Some((manager, mode)) => { + let manager = Arc::clone(manager); + let auth = auth.clone(); + let selection = crate::runtime_auth::resource_selection(headers)?; + let selected = crate::runtime_auth::finish_runtime_selection(async move { + manager.auth_for(&auth, &selection).await + }) + .await?; + Ok(ResolvedProviderAuth::for_runtime(selected, *mode)) + } + None => { + let mut resolved = ResolvedProviderAuth::new(auth_provider_from_auth(auth)); + resolved.credential = Some(auth.clone()); + Ok(resolved) + } + } + }) + .await +} + +/// Record bounded auth-mode and outcome labels without changing the selection result. +pub(crate) async fn observe_auth_selection( + mode: Option, + operation: &'static str, + selection: impl Future>, +) -> Result { + let started = std::time::Instant::now(); + let result = selection.await; + let mode = match mode { + None => "ordinary", + Some(RuntimeAuthMode::DualHeader) => "dual_header", + Some(RuntimeAuthMode::AssertionOnly) => "assertion_only", + }; + let _ = codex_otel::record_global_operation( + "codex.auth.selection", + "codex.auth.selection.duration", + started.elapsed(), + &[ + ("mode", mode), + ("operation", operation), + ("outcome", if result.is_ok() { "success" } else { "error" }), + ], + ); + result +} + +/// Builds request-header auth for a first-party Codex credential. pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { match auth { CodexAuth::AgentIdentity(auth) => { diff --git a/codex-rs/model-provider/src/client_auth.rs b/codex-rs/model-provider/src/client_auth.rs new file mode 100644 index 0000000000..3e229cb4a4 --- /dev/null +++ b/codex-rs/model-provider/src/client_auth.rs @@ -0,0 +1,120 @@ +//! Authentication for clients that reuse a connection across requests. + +use crate::ResolvedProviderAuth; +use codex_api::AuthError; +use codex_api::SharedAuthProvider; +use codex_features::RuntimeAuthMode; +use codex_http_client::HttpClientFactory; +use codex_login::CodexAuth; +use codex_protocol::error::Result as CodexResult; +use codex_runtime_agent_instance::RuntimeAuthManager; +use http::HeaderMap; +use std::sync::Arc; + +/// Keeps a reusable client's credential source fixed while allowing token refresh. +/// Each request resolves fresh auth using the same source as the client transport. +#[derive(Clone)] +pub struct ClientAuth { + source: ClientAuthSource, +} + +#[derive(Clone)] +enum ClientAuthSource { + Provider(SharedAuthProvider), + Runtime { + manager: Arc, + mode: RuntimeAuthMode, + source_generation: u64, + }, +} + +impl ClientAuth { + pub fn new(provider: SharedAuthProvider) -> Self { + Self { + source: ClientAuthSource::Provider(provider), + } + } + + /// Select the auth source and matching transport for a reusable client's request headers. + /// The caller selects runtime auth only for a supported OpenAI destination. + pub async fn for_auth( + auth: &CodexAuth, + runtime: Option<&(Arc, RuntimeAuthMode)>, + factory: HttpClientFactory, + headers: &HeaderMap, + ) -> Result<(Self, HttpClientFactory), AuthError> { + let resolved = crate::auth_provider_for_request(auth, runtime, headers).await?; + let factory = resolved.bind_http_client_factory(factory)?; + let source = match (runtime, resolved.runtime) { + (Some((manager, mode)), Some(selected)) => ClientAuthSource::Runtime { + manager: Arc::clone(manager), + mode: *mode, + source_generation: selected.change_state().source_generation, + }, + _ => ClientAuthSource::Provider(resolved.auth), + }; + Ok((Self { source }, factory)) + } + + /// Replace ordinary authentication; reject changes to a runtime-bound client. + pub fn set_provider(&mut self, provider: SharedAuthProvider) -> Result<(), AuthError> { + match &mut self.source { + ClientAuthSource::Provider(current) => { + *current = provider; + Ok(()) + } + ClientAuthSource::Runtime { .. } => Err(AuthError::Build( + "cannot replace credentials on a runtime-bound client".to_string(), + )), + } + } + + /// Resolve one request, preserving the source that authenticated the client. + pub async fn resolve(&self, headers: &HeaderMap) -> Result { + let mode = match &self.source { + ClientAuthSource::Provider(_) => None, + ClientAuthSource::Runtime { mode, .. } => Some(*mode), + }; + crate::auth::observe_auth_selection(mode, "client_request", async { + match &self.source { + ClientAuthSource::Provider(provider) => { + Ok(ResolvedProviderAuth::new(Arc::clone(provider))) + } + ClientAuthSource::Runtime { + manager, + mode, + source_generation, + } => { + let selected = + crate::select_runtime_auth_for_source(manager, *source_generation, headers) + .await?; + Ok(ResolvedProviderAuth::for_runtime(selected, *mode)) + } + } + }) + .await + } + + /// Allow OpenAI requests and the original origin, never an unrelated provider. + pub fn validate_destination(&self, base_url: &str, url: &str) -> CodexResult<()> { + if matches!(&self.source, ClientAuthSource::Provider(_)) { + return Ok(()); + } + crate::validate_runtime_destination(Some(url))?; + if codex_runtime_agent_instance::is_openai_api_url(url) { + return Ok(()); + } + let origin = |value| { + url::Url::parse(value) + .map(|url| url.origin()) + .map_err(std::io::Error::other) + }; + if origin(url)? != origin(base_url)? { + return Err(std::io::Error::other( + "runtime-authenticated clients cannot send requests to another provider", + ) + .into()); + } + Ok(()) + } +} diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 3b04f298cc..ba328161e5 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -4,6 +4,7 @@ mod bearer_auth_provider; mod capabilities; mod catalog_factory; mod chatgpt_base_url; +mod client_auth; mod combined_auth; mod gcp; mod models_endpoint; @@ -27,6 +28,7 @@ pub use auth::AgentIdentitySessionFallback; pub use auth::ProviderAuthContext; pub use auth::ProviderAuthScope; pub use auth::ResolvedProviderAuth; +pub use auth::auth_provider_for_request; pub use auth::auth_provider_from_auth; pub use auth::auth_provider_from_auth_manager; pub use auth::unauthenticated_auth_provider; @@ -36,6 +38,7 @@ pub use capabilities::ProviderCapabilities; pub use catalog_factory::ModelCatalogFactory; pub use chatgpt_base_url::ChatgptBaseUrl; pub use chatgpt_base_url::DEFAULT_CHATGPT_BASE_URL; +pub use client_auth::ClientAuth; pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; pub use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; @@ -57,6 +60,8 @@ pub use provider::resolve_api_provider; pub use request_preparation::prepare_provider_request; pub use request_setup::ResolvedProviderRequest; pub use request_setup::resolve_request; +pub use runtime_auth::runtime_auth_for_url; +pub use runtime_auth::select_runtime_auth_for_source; pub use runtime_auth::validate_runtime_destination; pub use runtime_auth::validate_runtime_provider; pub use runtime_auth::validate_runtime_selectors; diff --git a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs index 7a18508142..3aad6d169f 100644 --- a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs +++ b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs @@ -20,22 +20,43 @@ use wiremock::matchers::body_partial_json; use wiremock::matchers::method; use wiremock::matchers::path; +#[derive(Clone, Copy)] +enum Caller { + Catalog, + Startup, + CapturedCredential, + SourceGeneration, +} + #[tokio::test] async fn catalog_timeout_preserves_refreshed_tokens_with_and_without_runtime_auth() { - check_catalog_timeout_preserves_refreshed_tokens(/*refresh_after_auth_change*/ false).await; + check_timeout_preserves_refreshed_tokens(Caller::Catalog).await; } #[tokio::test] async fn catalog_startup_timeout_preserves_refreshed_tokens() { - check_catalog_timeout_preserves_refreshed_tokens(/*refresh_after_auth_change*/ true).await; + check_timeout_preserves_refreshed_tokens(Caller::Startup).await; } -async fn check_catalog_timeout_preserves_refreshed_tokens(refresh_after_auth_change: bool) { +#[tokio::test] +async fn cancelled_runtime_client_selection_persists_the_rotated_token() { + check_timeout_preserves_refreshed_tokens(Caller::CapturedCredential).await; + check_timeout_preserves_refreshed_tokens(Caller::SourceGeneration).await; +} + +async fn check_timeout_preserves_refreshed_tokens(caller: Caller) { for mode in [ None, Some(RuntimeAuthMode::DualHeader), Some(RuntimeAuthMode::AssertionOnly), ] { + if matches!( + caller, + Caller::CapturedCredential | Caller::SourceGeneration + ) && mode != Some(RuntimeAuthMode::DualHeader) + { + continue; + } let server = MockServer::start().await; let home = tempfile::tempdir().unwrap(); let token = |expires| { @@ -155,23 +176,60 @@ async fn check_catalog_timeout_preserves_refreshed_tokens(refresh_after_auth_cha Some(Arc::clone(&credentials)), /*gateway_auth_manager*/ None, ) - .with_runtime_auth(runtime_auth), + .with_runtime_auth(runtime_auth.clone()), ); - if refresh_after_auth_change { - let manager = OpenAiModelsManager::new_without_cache( - endpoint.clone(), - Some(Arc::clone(&credentials)), - ); - manager.refresh_after_auth_change(factory.clone()).await; - } else { - let error = endpoint - .list_models("test", factory.clone()) - .await - .unwrap_err(); - assert!(matches!( - error.details(), - codex_protocol::error::CodexErrorDetails::RequestTimeout - )); + match caller { + Caller::CapturedCredential => { + let auth = credentials.auth_cached().unwrap(); + assert!( + timeout( + MODELS_REFRESH_TIMEOUT, + crate::auth_provider_for_request( + &auth, + runtime_auth.as_ref(), + &http::HeaderMap::new() + ) + ) + .await + .is_err() + ); + } + Caller::SourceGeneration => { + let generation = credentials + .auth_change_state_receiver() + .borrow() + .source_generation; + let runtime = &runtime_auth.as_ref().unwrap().0; + assert!( + timeout( + MODELS_REFRESH_TIMEOUT, + crate::select_runtime_auth_for_source( + runtime, + generation, + &http::HeaderMap::new() + ) + ) + .await + .is_err() + ); + } + Caller::Startup => { + let manager = OpenAiModelsManager::new_without_cache( + endpoint.clone(), + Some(Arc::clone(&credentials)), + ); + manager.refresh_after_auth_change(factory.clone()).await; + } + Caller::Catalog => { + let error = endpoint + .list_models("test", factory.clone()) + .await + .unwrap_err(); + assert!(matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::RequestTimeout + )); + } } assert_eq!( server diff --git a/codex-rs/model-provider/src/models_identity.rs b/codex-rs/model-provider/src/models_identity.rs index 356edf9c8f..7095dadc63 100644 --- a/codex-rs/model-provider/src/models_identity.rs +++ b/codex-rs/model-provider/src/models_identity.rs @@ -16,6 +16,13 @@ pub(crate) fn identity( provider_info: &ModelProviderInfo, auth: Option<&CodexAuth>, ) -> CoreResult { + // Provider bearer overrides replace the login credential and its account metadata. + let auth = + if provider_info.env_key.is_some() || provider_info.experimental_bearer_token.is_some() { + None + } else { + auth + }; let mut api_provider = resolve_api_provider(provider_info, auth)?; let mut digest = Sha256::new(); // Length-prefix every field to avoid ambiguity between adjacent values. diff --git a/codex-rs/model-provider/src/models_identity_tests.rs b/codex-rs/model-provider/src/models_identity_tests.rs index 0dc9e9f190..0ad6a8e289 100644 --- a/codex-rs/model-provider/src/models_identity_tests.rs +++ b/codex-rs/model-provider/src/models_identity_tests.rs @@ -113,3 +113,30 @@ fn cache_identity_tracks_catalog_url() { assert_ne!(bundled, first); assert_ne!(first, second); } + +#[test] +fn provider_key_catalog_identity_does_not_depend_on_unused_login() { + let mut provider = + ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".into())); + // Keep this input stable while the residency test changes the process default. + provider.http_headers = Some(std::collections::HashMap::from([( + codex_login::default_client::RESIDENCY_HEADER_NAME.to_string(), + "us".into(), + )])); + provider.requires_openai_auth = false; + provider.experimental_bearer_token = Some("provider-key".into()); + let identity_without_login = identity(&provider, /*auth*/ None).unwrap(); + let selected_key = CodexAuth::from_api_key("provider-key"); + let unused_login = chatgpt_auth("user@example.com", "user", "workspace", "team", "token"); + for auth in [&selected_key, &unused_login] { + assert_eq!( + identity(&provider, Some(auth)).unwrap(), + identity_without_login + ); + } + provider.experimental_bearer_token = Some("different-provider-key".into()); + assert_ne!( + identity(&provider, /*auth*/ None).unwrap(), + identity_without_login + ); +} diff --git a/codex-rs/model-provider/src/request_setup.rs b/codex-rs/model-provider/src/request_setup.rs index c7e33a0152..26a0bbd009 100644 --- a/codex-rs/model-provider/src/request_setup.rs +++ b/codex-rs/model-provider/src/request_setup.rs @@ -33,31 +33,42 @@ pub(crate) async fn resolve_request_auth( destination: Option, ) -> codex_protocol::error::Result { tokio::spawn(async move { - if let Some((runtime, mode)) = runtime { - crate::validate_runtime_destination(destination.as_deref())?; - crate::validate_runtime_provider(&info, manager.as_ref(), &runtime)?; - let headers = info.build_header_map()?; - let selected = crate::runtime_auth::select_runtime_auth(&runtime, &headers) - .await - .map_err(|error| codex_api::map_api_error(error.into()))?; - let resolved = ResolvedProviderAuth::for_runtime(selected, mode); - resolved.validate_headers(&headers)?; - tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication"); - return Ok(resolved); - } + crate::auth::observe_auth_selection( + runtime.as_ref().map(|(_, mode)| *mode), + "model", + async move { + if let Some((runtime, mode)) = runtime { + crate::validate_runtime_destination(destination.as_deref())?; + crate::validate_runtime_provider(&info, manager.as_ref(), &runtime)?; + let headers = info.build_header_map()?; + let selected = crate::runtime_auth::select_runtime_auth(&runtime, &headers) + .await + .map_err(|error| codex_api::map_api_error(error.into()))?; + let resolved = ResolvedProviderAuth::for_runtime(selected, mode); + resolved.validate_headers(&headers)?; + tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication"); + return Ok(resolved); + } - let credential = match &manager { - Some(manager) => manager.auth().await, - None => None, - }; - let mut resolved = match scope { - Some(scope) if crate::provider_uses_first_party_auth_path(&info) => { - resolve_provider_auth_for_scope(manager, credential.as_ref(), &info, scope).await? - } - _ => ResolvedProviderAuth::new(resolve_provider_auth(credential.as_ref(), &info)?), - }; - resolved.credential = credential; - Ok(resolved) + let credential = match &manager { + Some(manager) => manager.auth().await, + None => None, + }; + let mut resolved = match scope { + Some(scope) if crate::provider_uses_first_party_auth_path(&info) => { + resolve_provider_auth_for_scope(manager, credential.as_ref(), &info, scope) + .await? + } + _ => ResolvedProviderAuth::new(resolve_provider_auth( + credential.as_ref(), + &info, + )?), + }; + resolved.credential = credential; + Ok(resolved) + }, + ) + .await }) .await? } @@ -85,9 +96,10 @@ impl ResolvedProviderRequest { } else { self.redirect_policy }; + let request_url = self.api_provider.url_for_path(endpoint); let client = create_client_for_route( &self.http_client_factory, - &self.api_provider.url_for_path(endpoint), + &request_url, ClientRouteClass::Api, redirect_policy, ) diff --git a/codex-rs/model-provider/src/runtime_auth.rs b/codex-rs/model-provider/src/runtime_auth.rs index 4c4b992f15..26a7be196c 100644 --- a/codex-rs/model-provider/src/runtime_auth.rs +++ b/codex-rs/model-provider/src/runtime_auth.rs @@ -1,6 +1,7 @@ //! Keep model requests and catalog refreshes bound to the same registered actor. use codex_api::AuthError; +use codex_features::RuntimeAuthMode; use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; use codex_http_client::OPENAI_ORGANIZATION_HEADER; use codex_http_client::OPENAI_PROJECT_HEADER; @@ -13,6 +14,7 @@ use codex_protocol::error::Result as CodexResult; use codex_runtime_agent_instance::ActorInfo; use codex_runtime_agent_instance::PlatformResourceSelection; use codex_runtime_agent_instance::RuntimeAuth; +use codex_runtime_agent_instance::RuntimeAuthError; use codex_runtime_agent_instance::RuntimeAuthManager; use http::HeaderMap; use std::sync::Arc; @@ -22,6 +24,39 @@ pub(crate) async fn select_runtime_auth( manager: &RuntimeAuthManager, headers: &HeaderMap, ) -> Result { + manager + .auth_with_selection(&resource_selection(headers)?) + .await + .map_err(AuthError::from) +} + +/// Select fresh request auth without allowing a reusable client's source to change. +pub async fn select_runtime_auth_for_source( + manager: &Arc, + source_generation: u64, + headers: &HeaderMap, +) -> Result { + let manager = Arc::clone(manager); + let selection = resource_selection(headers)?; + finish_runtime_selection( + async move { manager.auth_for_source(source_generation, &selection).await }, + ) + .await +} + +// Once token rotation begins, let it finish persisting even if the request is cancelled. +pub(crate) async fn finish_runtime_selection( + selection: impl Future> + Send + 'static, +) -> Result { + tokio::spawn(selection) + .await + .map_err(|error| AuthError::Build(format!("runtime auth task failed: {error}")))? + .map_err(AuthError::from) +} + +pub(crate) fn resource_selection( + headers: &HeaderMap, +) -> Result { let selector = |name| { headers .get(name) @@ -33,14 +68,10 @@ pub(crate) async fn select_runtime_auth( }) .transpose() }; - let selection = PlatformResourceSelection { + Ok(PlatformResourceSelection { organization_id: selector(OPENAI_ORGANIZATION_HEADER)?, project_id: selector(OPENAI_PROJECT_HEADER)?, - }; - manager - .auth_with_selection(&selection) - .await - .map_err(AuthError::from) + }) } /// A runtime may not replace an independently selected provider credential. @@ -74,6 +105,17 @@ pub fn validate_runtime_destination(destination: Option<&str>) -> CodexResult<() Ok(()) } +/// Select runtime auth for supported OpenAI URLs; reject unsupported OpenAI URLs. +pub fn runtime_auth_for_url<'a>( + runtime: Option<&'a (Arc, RuntimeAuthMode)>, + url: &str, +) -> CodexResult, RuntimeAuthMode)>> { + if runtime.is_some() { + validate_runtime_destination(Some(url))?; + } + Ok(runtime.filter(|_| classify_openai_api_url(url) == OpenAiDestination::Supported)) +} + /// Account and organization selectors must agree with the registered actor. /// API project access is authorized by the backend for each request. pub fn validate_runtime_selectors(headers: &HeaderMap, actor: &ActorInfo) -> CodexResult<()> { diff --git a/codex-rs/runtime-agent-instance/Cargo.toml b/codex-rs/runtime-agent-instance/Cargo.toml index 1c000678ab..53fd913ab3 100644 --- a/codex-rs/runtime-agent-instance/Cargo.toml +++ b/codex-rs/runtime-agent-instance/Cargo.toml @@ -30,6 +30,8 @@ uuid = { workspace = true, features = ["serde", "v4"] } zeroize = { workspace = true } [dev-dependencies] +bytes = { workspace = true } +futures = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/runtime-agent-instance/src/lib.rs b/codex-rs/runtime-agent-instance/src/lib.rs index 31ff7ea7cc..f28529e367 100644 --- a/codex-rs/runtime-agent-instance/src/lib.rs +++ b/codex-rs/runtime-agent-instance/src/lib.rs @@ -28,3 +28,6 @@ pub use registration::RegistrationErrorCode; pub use registration::RuntimeRegistrationClient; pub use runtime::RuntimeAgentInstance; pub use runtime::RuntimeRegistration; + +/// Harness-owned lazy runtime, shared before any request requires registration. +pub type RuntimeAuthSlot = std::sync::Arc>>; diff --git a/codex-rs/runtime-agent-instance/src/manager.rs b/codex-rs/runtime-agent-instance/src/manager.rs index 15d7b42c5e..e1bb6f204f 100644 --- a/codex-rs/runtime-agent-instance/src/manager.rs +++ b/codex-rs/runtime-agent-instance/src/manager.rs @@ -135,16 +135,25 @@ impl RuntimeAuthManager { } /// Select auth and a runtime for the requested organization and project. - #[expect( - clippy::await_holding_invalid_type, - reason = "runtime initialization serializes registration and refresh for one credential source" - )] pub async fn auth_with_selection( &self, selection: &PlatformResourceSelection, + ) -> Result { + self.select_auth(selection, /*expected_source*/ None).await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "runtime initialization serializes registration and refresh for one credential source" + )] + async fn select_auth( + &self, + selection: &PlatformResourceSelection, + expected_source: Option, ) -> Result { // This initializer's lock can span IO. No AuthManager lock is held across registration. let mut state = self.state.lock().await; + self.check_source(expected_source)?; let (auth, http_client_factory) = self .credentials .auth_with_http_client_factory() @@ -152,7 +161,13 @@ impl RuntimeAuthManager { .ok_or(RuntimeAuthError::Unavailable)?; let mut recovery = self.credentials.unauthorized_recovery(); let mut result = self - .auth_once(&mut state, auth, http_client_factory, selection) + .auth_once( + &mut state, + auth, + http_client_factory, + selection, + expected_source, + ) .await; while matches!( &result, @@ -164,6 +179,7 @@ impl RuntimeAuthManager { )) ) && recovery.has_next() { + self.check_source(expected_source)?; recovery.next().await.map_err(|error| match error { RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied), RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => { @@ -176,7 +192,13 @@ impl RuntimeAuthManager { .await .ok_or(RuntimeAuthError::Unavailable)?; result = self - .auth_once(&mut state, auth, http_client_factory, selection) + .auth_once( + &mut state, + auth, + http_client_factory, + selection, + expected_source, + ) .await; } result @@ -188,10 +210,12 @@ impl RuntimeAuthManager { auth: CodexAuth, http_client_factory: HttpClientFactory, selection: &PlatformResourceSelection, + expected_source: Option, ) -> Result { let actor_auth = actor_authorization(&auth, selection)?; let http_client_factory = http_client_factory.with_sensitive_header_protection(); let change_state = self.auth_state(&auth)?; + self.check_source(expected_source)?; // Registered runtimes can serve different projects. Pending and denied // registrations retain their original scope and idempotency key. if state.as_ref().is_some_and(|cached| { @@ -269,6 +293,45 @@ impl RuntimeAuthManager { } } + /// Resolve auth for work whose routing, cache key, or payload uses a captured credential. + /// A same-source refresh is allowed; a different source must not take over that work. + /// This binds credential identity, not when the caller captured it. Restoring the exact + /// same credential can select a new runtime after logout; a different source cannot. + pub async fn auth_for( + &self, + expected: &CodexAuth, + selection: &PlatformResourceSelection, + ) -> Result { + let source = self + .credentials + .auth_change_state_for(expected) + .ok_or(RuntimeAuthError::Changed)?; + self.auth_for_source(source.source_generation, selection) + .await + } + + /// Select another request within a reusable client's original credential source. + pub async fn auth_for_source( + &self, + source_generation: u64, + selection: &PlatformResourceSelection, + ) -> Result { + self.select_auth(selection, Some(source_generation)).await + } + + fn check_source(&self, expected: Option) -> Result<(), RuntimeAuthError> { + if expected.is_some_and(|expected| { + self.credentials + .auth_change_state_receiver() + .borrow() + .source_generation + != expected + }) { + return Err(RuntimeAuthError::Changed); + } + Ok(()) + } + fn auth_state(&self, auth: &CodexAuth) -> Result { let state = self .credentials diff --git a/codex-rs/runtime-agent-instance/src/manager_tests.rs b/codex-rs/runtime-agent-instance/src/manager_tests.rs index cff21d2189..6e4eee9b09 100644 --- a/codex-rs/runtime-agent-instance/src/manager_tests.rs +++ b/codex-rs/runtime-agent-instance/src/manager_tests.rs @@ -571,3 +571,140 @@ async fn registration_failure_reports_retryability() { ); } } + +#[tokio::test] +async fn captured_request_auth_allows_refresh_but_rejects_source_replacement() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let initial = oauth(Some("login-a"), "old"); + let refreshed = oauth(Some("login-a"), "new"); + let (manager, source, _) = setup(&server, initial.clone(), refreshed.clone()).await; + let runtime = manager.auth().await.unwrap().runtime().id().to_string(); + *source.current.write().unwrap() = refreshed.clone(); + let selected = manager + .auth_for(&initial, &PlatformResourceSelection::default()) + .await + .unwrap(); + let source_generation = selected.change_state().source_generation; + let request_auth = selected.into_auth_provider(RuntimeAuthMode::DualHeader); + let headers = request_auth + .apply_auth_headers(HeaderMap::new()) + .await + .unwrap(); + assert_eq!( + headers[AUTHORIZATION], + format!("Bearer {}", refreshed.get_token().unwrap()) + ); + let proof = headers["x-openai-agent-assertion"].to_str().unwrap(); + let payload: Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(proof.split('.').nth(2).unwrap()) + .unwrap(), + ) + .unwrap(); + assert_eq!(payload["runtime_agent_instance"], runtime); + *source.current.write().unwrap() = oauth(Some("another-login"), "replacement"); + assert!(matches!( + manager + .auth_for_source(source_generation, &PlatformResourceSelection::default()) + .await, + Err(RuntimeAuthError::Changed) + )); + assert!( + manager + .auth_for(&initial, &PlatformResourceSelection::default()) + .await + .is_err() + ); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn file_upload_runtime_auth_signs_control_requests_but_not_blob_uploads() { + use wiremock::matchers::path; + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let blob = MockServer::start().await; + Mock::given(path("/agent/register")) + .respond_with(registered) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/files")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "file_id": "file-a", "upload_url": blob.uri(), + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/files/file-a/uploaded")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"status": "retry"}))) + .up_to_n_times(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(path("/files/file-a/uploaded")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "status": "success", "download_url": blob.uri(), "file_name": "hello.txt", "file_size_bytes": 5, + }))).expect(1).mount(&server).await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&blob) + .await; + let initial = CodexAuth::from_api_key("key-a"); + let (manager, _, _) = setup(&server, initial.clone(), initial).await; + let provider = manager.auth().await.unwrap().into_auth_provider(mode); + let pool = + codex_http_client::RouteAwareClientPool::new_without_redirects_or_request_logging( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + codex_http_client::ClientRouteClass::Api, + ); + let uploaded = codex_api::upload_openai_file( + &server.uri(), + provider.as_ref(), + &pool, + "hello.txt".into(), + /*file_size_bytes*/ 5, + || async { + Ok(futures::stream::iter([Ok::<_, std::io::Error>( + bytes::Bytes::from_static(b"hello"), + )])) + }, + /*hosted_upload*/ None, + ) + .await + .unwrap(); + assert_eq!(uploaded.uri, "sediment://file-a"); + let requests = server.received_requests().await.unwrap(); + let mut proofs = std::collections::HashSet::new(); + for request in requests + .iter() + .filter(|r| r.url.path().starts_with("/files")) + { + let authorization = request.headers[AUTHORIZATION].to_str().unwrap(); + let proof = match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + request.headers["x-openai-agent-assertion"] + .to_str() + .unwrap() + } + RuntimeAuthMode::AssertionOnly => { + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + authorization.strip_prefix("AgentAssertion ").unwrap() + } + }; + assert!(proofs.insert(proof)); + } + assert_eq!(proofs.len(), 3); + let blobs = blob.received_requests().await.unwrap(); + assert_eq!(blobs.len(), 1); + assert!(!blobs[0].headers.contains_key(AUTHORIZATION)); + assert!(!blobs[0].headers.contains_key("x-openai-agent-assertion")); + assert_eq!(blobs[0].body, b"hello"); + } +}