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"); + } +}