diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e8cf46c3d5..c8c772fd7c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4183,6 +4183,7 @@ dependencies = [ "codex-agent-identity", "codex-api", "codex-aws-auth", + "codex-features", "codex-feedback", "codex-http-client", "codex-keyring-store", @@ -4192,12 +4193,16 @@ dependencies = [ "codex-otel", "codex-protocol", "codex-response-debug-context", + "codex-runtime-agent-instance", "codex-secrets", "codex-utils-cargo-bin", "codex-utils-process", "codex-utils-redacted-string", + "codex-utils-rustls-provider", "http 1.4.0", "pretty_assertions", + "rcgen", + "rustls", "serde", "serde_json", "sha2 0.10.9", diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index df64477120..3308777b5a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -71,11 +71,8 @@ use codex_api::build_session_headers; use codex_api::create_text_param_for_request; use codex_api::response_create_client_metadata; use codex_features::RuntimeAuthMode; -use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; -use codex_http_client::OPENAI_ORGANIZATION_HEADER; -use codex_http_client::OPENAI_PROJECT_HEADER; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::RefreshTokenError; @@ -88,7 +85,6 @@ use codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC; use codex_otel::current_span_w3c_trace_context; use codex_protocol::ResponseItemId; use codex_protocol::auth::AuthMode; -use codex_runtime_agent_instance::ActorInfo; use codex_runtime_agent_instance::RuntimeAuthManager; use codex_protocol::ThreadId; @@ -249,7 +245,7 @@ struct CurrentClientSetup { redirect_policy: ClientRedirectPolicy, api_auth: SharedAuthProvider, agent_identity_telemetry: Option, - runtime: Option<(String, ActorInfo)>, + runtime_id: Option, } impl CurrentClientSetup { @@ -676,19 +672,12 @@ impl ModelClient { mode: RuntimeAuthMode, ) -> Result { let info = self.state.provider.info(); - if !info.requires_openai_auth - || info.env_key.is_some() - || info.experimental_bearer_token.is_some() - || info.auth.is_some() - || info.aws.is_some() - || !self - .auth_manager() - .is_some_and(|auth| manager.uses_credentials(&auth)) - { - return Err(CodexErr::UnsupportedOperation( - "runtime auth requires the model provider's selected OpenAI credential".to_string(), - )); - } + codex_model_provider::validate_runtime_destination(info.base_url.as_deref())?; + codex_model_provider::validate_runtime_provider( + info, + self.auth_manager().as_ref(), + &manager, + )?; self.runtime_auth = Some((manager, mode)); Ok(self) } @@ -825,9 +814,7 @@ impl ModelClient { ) -> Result { // Create the media call over HTTP first, then retain matching auth so realtime can attach // the server-side control WebSocket to the call id from that HTTP response. - let mut client_setup = self - .provider_client_setup(ClientRouting::ConfiguredProvider) - .await?; + let mut client_setup = self.realtime_client_setup().await?; if let Some(header_value) = self.generate_attestation_header_for().await { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } @@ -856,9 +843,7 @@ impl ModelClient { &self, mut extra_headers: ApiHeaderMap, ) -> Result { - let client_setup = self - .provider_client_setup(ClientRouting::ConfiguredProvider) - .await?; + let client_setup = self.realtime_client_setup().await?; if let Some(header_value) = self.generate_attestation_header_for().await { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } @@ -1214,11 +1199,12 @@ impl ModelClient { /// lockstep when auth/provider resolution changes. async fn current_client_setup(&self, routing: ClientRouting) -> Result { if let Some((manager, mode)) = &self.runtime_auth { - let selected = manager.auth().await.map_err(|err| { - self.state - .provider - .map_api_error(codex_api::AuthError::from(err).into()) - })?; + let selected = codex_model_provider::select_runtime_auth( + manager, + &self.state.provider.info().build_header_map()?, + ) + .await + .map_err(|error| self.state.provider.map_api_error(error.into()))?; let auth = selected.auth().clone(); let auth_revision = Some(selected.change_state().generation); let auth_owner_generation = Some(selected.change_state().owner_generation); @@ -1239,16 +1225,19 @@ impl ModelClient { ClientRouting::ConfiguredProvider => (api_provider, ClientRedirectPolicy::Default), }; let runtime = selected.runtime(); - let actor = runtime.actor().info().clone(); - validate_runtime_selectors(&api_provider.headers, &actor)?; - let runtime = Some((runtime.id().to_string(), actor)); + codex_model_provider::validate_runtime_selectors( + &api_provider.headers, + runtime.actor().info(), + )?; + let runtime_id = Some(runtime.id().to_string()); let http_client_factory = selected.http_client_factory().clone(); let resolved_auth = self .state .provider .compose_auth(ResolvedProviderAuth { - auth: selected.into_auth_provider(*mode), + auth: selected.clone().into_auth_provider(*mode), agent_identity_telemetry: None, + runtime: Some(selected), }) .await?; tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication"); @@ -1261,42 +1250,54 @@ impl ModelClient { redirect_policy, api_auth: resolved_auth.auth, agent_identity_telemetry: None, - runtime, + runtime_id, }); } - self.provider_client_setup(routing).await + self.provider_client_setup(routing, &self.state.provider) + .await } // Realtime keeps provider auth until its runtime integration lands later in the stack. - async fn provider_client_setup(&self, routing: ClientRouting) -> Result { + async fn realtime_client_setup(&self) -> Result { + let provider = if self.state.provider.runtime_auth().is_some() { + create_model_provider( + self.state.provider.info().clone(), + self.state.provider.auth_manager(), + ) + } else { + Arc::clone(&self.state.provider) + }; + self.provider_client_setup(ClientRouting::ConfiguredProvider, &provider) + .await + } + + async fn provider_client_setup( + &self, + routing: ClientRouting, + provider: &SharedModelProvider, + ) -> Result { // Capture before resolving credentials so an account switch during setup cannot label // an old connection with the new owner's revision. let auth_owner_generation = self.auth_owner_generation(); - let auth_changes = self - .state - .provider + let auth_changes = provider .auth_manager() .map(|manager| manager.auth_change_receiver()); loop { let revision = auth_changes.as_ref().map(|changes| *changes.borrow()); - let auth = self.state.provider.auth().await; + let auth = provider.auth().await; let (api_provider, redirect_policy) = match routing { ClientRouting::Workspace => { - let resolved = self - .state - .provider + let resolved = provider .responses_api_provider(&self.state.workspace_routing) .await?; (resolved.provider, resolved.redirect_policy) } ClientRouting::ConfiguredProvider => ( - self.state.provider.api_provider().await?, + provider.api_provider().await?, ClientRedirectPolicy::Default, ), }; - let resolved_auth = self - .state - .provider + let resolved_auth = provider .api_auth_for_scope(ProviderAuthScope { agent_identity_policy: self.agent_identity_policy, session_source: self.state.session_source.clone(), @@ -1308,7 +1309,7 @@ impl ModelClient { .await?; // Command-backed bearer refreshes cannot change workspace routing. Keep the captured // revisions so a refresh during setup still invalidates cached WebSocket state. - if self.state.provider.info().auth.is_none() { + if provider.info().auth.is_none() { if self.auth_owner_generation() != auth_owner_generation { return Err(std::io::Error::other( "account changed while preparing model request", @@ -1329,7 +1330,7 @@ impl ModelClient { redirect_policy, api_auth: resolved_auth.auth, agent_identity_telemetry: resolved_auth.agent_identity_telemetry, - runtime: None, + runtime_id: None, }); } } @@ -1735,7 +1736,7 @@ impl ModelClientSession { auth_context, request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), responses_headers: &responses_headers, - runtime_id: client_setup.runtime.map(|(id, _)| id), + runtime_id: client_setup.runtime_id, }) .await { @@ -2164,7 +2165,7 @@ impl ModelClientSession { api_auth: client_setup.api_auth, http_client_factory: client_setup.http_client_factory, auth_owner_generation: client_setup.auth_owner_generation, - runtime_id: client_setup.runtime.map(|(id, _)| id), + runtime_id: client_setup.runtime_id, responses_metadata: &websocket_metadata, auth_context: request_auth_context, request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), @@ -2877,31 +2878,6 @@ impl AuthRequestTelemetryContext { } } -fn validate_runtime_selectors(headers: &ApiHeaderMap, actor: &ActorInfo) -> Result<()> { - let (account, organization, project) = match actor { - ActorInfo::Chatgpt { account_id, .. } => (Some(account_id.as_str()), None, None), - ActorInfo::Platform { - organization_id, - project_id, - .. - } => (None, Some(organization_id.as_str()), project_id.as_deref()), - }; - for (name, expected) in [ - (CHATGPT_ACCOUNT_ID_HEADER, account), - (OPENAI_ORGANIZATION_HEADER, organization), - (OPENAI_PROJECT_HEADER, project), - ] { - if let Some(value) = headers.get(name) - && value.to_str().ok() != expected - { - return Err(CodexErr::UnsupportedOperation(format!( - "configured {name} does not match the registered runtime" - ))); - } - } - Ok(()) -} - struct WebsocketConnectParams<'a> { runtime_id: Option, session_telemetry: &'a SessionTelemetry, diff --git a/codex-rs/core/src/client_runtime_auth_tests.rs b/codex-rs/core/src/client_runtime_auth_tests.rs index 2255a35b04..a1906c9504 100644 --- a/codex-rs/core/src/client_runtime_auth_tests.rs +++ b/codex-rs/core/src/client_runtime_auth_tests.rs @@ -128,11 +128,16 @@ async fn selected_runtime_transport_stays_revoked(mode: RuntimeAuthMode) { server.verify().await; } -#[test_case::test_case(RuntimeAuthMode::DualHeader; "dual_header")] -#[test_case::test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] +#[test_case::test_case(RuntimeAuthMode::DualHeader, false; "dual_header")] +#[test_case::test_case(RuntimeAuthMode::AssertionOnly, false; "assertion_only")] +#[test_case::test_case(RuntimeAuthMode::DualHeader, true; "runtime_provider_dual_header")] +#[test_case::test_case(RuntimeAuthMode::AssertionOnly, true; "runtime_provider_assertion_only")] #[tokio::test] -async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAuthMode) { - let (client, _) = model_client_with_counting_attestation(/*include_attestation*/ true); +async fn realtime_keeps_provider_auth_until_runtime_integration( + mode: RuntimeAuthMode, + runtime_provider: bool, +) { + let (mut client, _) = model_client_with_counting_attestation(/*include_attestation*/ true); let registration = MockServer::start().await; let manager = client.auth_manager().unwrap(); let expected_authorization = format!( @@ -140,7 +145,7 @@ async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAut manager.auth().await.unwrap().get_token().unwrap() ); let runtime = Arc::new(RuntimeAuthManager::new( - manager, + manager.clone(), AgentHostAuthorization::InstallationId { installation_id: uuid::Uuid::nil(), }, @@ -151,6 +156,15 @@ async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAut }, format!("{}/register", registration.uri()).parse().unwrap(), )); + if runtime_provider { + let info = client.state.provider.info().clone(); + Arc::get_mut(&mut client.state).unwrap().provider = + codex_model_provider::ProviderAuthContext { + credentials: manager, + runtime_auth: Some((Arc::clone(&runtime), mode)), + } + .model_provider(info); + } let client = client.with_runtime_auth(runtime, mode).unwrap(); let sideband = client .realtime_sideband_headers(http::HeaderMap::new()) 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 b9014324f0..5eb1ea55ed 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -192,6 +192,10 @@ async fn http_requests_keep_matching_runtime_and_credentials_across_retries_and_ let mut provider = ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); provider.supports_websockets = false; + provider.http_headers = Some(std::collections::HashMap::from([ + ("openai-organization".into(), "org-a".into()), + ("openai-project".into(), "project-a".into()), + ])); provider.request_max_retries = Some(1); if gateway { provider.gateway_oauth = Some(gateway_config( @@ -213,6 +217,11 @@ async fn http_requests_keep_matching_runtime_and_credentials_across_retries_and_ .filter(|r| r.url.path() == "/register") .collect(); assert_eq!(bootstraps.len(), 2); + for registration in &bootstraps { + let body: Value = registration.body_json().unwrap(); + assert_eq!(body["actor_auth"]["organization_id"], "org-a"); + assert_eq!(body["actor_auth"]["project_id"], "project-a"); + } assert!( bootstraps .iter() @@ -227,6 +236,8 @@ async fn http_requests_keep_matching_runtime_and_credentials_across_retries_and_ .iter() .zip(["key-a", "key-a", "key-other-user"]) .map(|(r, credential)| { + assert_eq!(r.headers["openai-organization"], "org-a"); + assert_eq!(r.headers["openai-project"], "project-a"); assert_eq!( r.headers .get("x-gateway-auth") @@ -305,6 +316,8 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( )); } provider.http_headers = Some(std::collections::HashMap::from([ + ("openai-organization".into(), "org-a".into()), + ("openai-project".into(), "project-a".into()), ("authorization".into(), "Bearer stale-credential".into()), ("x-openai-agent-assertion".into(), "stale-assertion".into()), ("x-openai-fedramp".into(), "true".into()), @@ -384,11 +397,18 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( assert_eq!(body["client_metadata"]["x-codex-turn-state"], Value::Null); assert_eq!(body["input"].as_array().unwrap().len(), index + 2); } + for request in registration.received_requests().await.unwrap() { + let body: Value = request.body_json().unwrap(); + assert_eq!(body["actor_auth"]["organization_id"], "org-a"); + assert_eq!(body["actor_auth"]["project_id"], "project-a"); + } server.shutdown().await; } +#[test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] #[tokio::test] -async fn mismatched_resource_selector_never_sends_model_request() { +async fn mismatched_organization_never_sends_model_request(mode: RuntimeAuthMode) { skip_if_no_network!(); let server = MockServer::start().await; Mock::given(path("/register")) @@ -399,12 +419,16 @@ async fn mismatched_resource_selector_never_sends_model_request() { ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); provider.supports_websockets = false; provider.http_headers = - Some([(String::from("OpenAI-Project"), "another-project".into())].into()); - let (harness, _) = runtime_harness(&server, provider, RuntimeAuthMode::DualHeader).await; + Some([(String::from("OpenAI-Organization"), "another-org".into())].into()); + let (harness, _) = runtime_harness(&server, provider, mode).await; let error = complete(&mut harness.client.new_session(), &harness) .await .unwrap_err(); - assert!(error.to_string().contains("openai-project does not match")); + assert!( + error + .to_string() + .contains("openai-organization does not match") + ); let requests = server.received_requests().await.unwrap(); assert_eq!(requests.len(), 1); assert_eq!(requests[0].url.path(), "/register"); diff --git a/codex-rs/http-client/src/client_builder.rs b/codex-rs/http-client/src/client_builder.rs index 86c7cbbcda..7de37b90ff 100644 --- a/codex-rs/http-client/src/client_builder.rs +++ b/codex-rs/http-client/src/client_builder.rs @@ -178,7 +178,9 @@ impl HttpClientBuilder { request_url: &str, route_class: ClientRouteClass, ) -> Result { - if http_client_factory.network_policy().is_managed() { + if http_client_factory.network_policy().is_managed() + || http_client_factory.protects_sensitive_headers() + { return Ok(crate::RouteAwareClientPool::with_builder( http_client_factory.clone(), route_class, diff --git a/codex-rs/http-client/src/lib.rs b/codex-rs/http-client/src/lib.rs index 94b7a88f70..db66e6b18b 100644 --- a/codex-rs/http-client/src/lib.rs +++ b/codex-rs/http-client/src/lib.rs @@ -10,6 +10,7 @@ mod custom_ca; mod error; mod native_roots; mod network_policy; +mod openai_endpoints; mod openai_headers; mod outbound_proxy; mod request; @@ -47,6 +48,9 @@ pub use crate::network_policy::NetworkPolicy; pub use crate::network_policy::NetworkPolicyController; pub use crate::network_policy::NetworkPolicyDenied; pub use crate::network_policy::NetworkPolicyRevision; +pub use crate::openai_endpoints::OpenAiDestination; +pub use crate::openai_endpoints::classify_openai_api_url; +pub use crate::openai_endpoints::is_openai_api_url; pub use crate::openai_headers::CHATGPT_ACCOUNT_ID_HEADER; pub use crate::openai_headers::OPENAI_ORGANIZATION_HEADER; pub use crate::openai_headers::OPENAI_PROJECT_HEADER; diff --git a/codex-rs/http-client/src/network_policy.rs b/codex-rs/http-client/src/network_policy.rs index 4c3ce3e5f0..b1956bf2b9 100644 --- a/codex-rs/http-client/src/network_policy.rs +++ b/codex-rs/http-client/src/network_policy.rs @@ -214,6 +214,30 @@ impl NetworkPolicy { self } + /// Bind a caller's restrictions to a captured credential policy without widening either. + /// Different managed controllers or account bindings cannot authorize the same request. + pub fn with_account_binding(mut self, credentials: &Self) -> Option { + if let Some(state) = &credentials.state { + if let Some(caller) = &self.state { + if !Arc::ptr_eq(caller, state) { + return None; + } + } else { + self.state = Some(Arc::clone(state)); + } + } + if let (Some(caller), Some(credentials)) = (self.account, credentials.account) + && caller != credentials + { + return None; + } + self.account = self.account.or(credentials.account); + if let Some(endpoints) = &credentials.endpoints { + self = self.restrict_to_endpoints(endpoints.as_ref().clone()); + } + Some(self) + } + /// Whether this transport participates in the application policy lifecycle. pub fn is_managed(&self) -> bool { self.state.is_some() diff --git a/codex-rs/http-client/src/network_policy_tests.rs b/codex-rs/http-client/src/network_policy_tests.rs index dd97871365..d3ce0bc80c 100644 --- a/codex-rs/http-client/src/network_policy_tests.rs +++ b/codex-rs/http-client/src/network_policy_tests.rs @@ -7,6 +7,81 @@ use reqwest::Url; use super::*; +#[test] +fn credential_binding_keeps_both_endpoint_restrictions_and_account_revocation() { + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let a = Url::parse("https://example.com/a").unwrap(); + let b = Url::parse("https://example.com/b").unwrap(); + let c = Url::parse("https://example.com/c").unwrap(); + let credentials = policy + .clone() + .for_current_account() + .restrict_to_endpoints(BTreeSet::from([b.clone(), c.clone()])); + let bound = policy + .clone() + .restrict_to_endpoints(BTreeSet::from([a.clone(), b.clone()])) + .with_account_binding(&credentials) + .unwrap(); + assert!(matches!( + bound.acquire(&a), + Err(NetworkPolicyDenied::Destination) + )); + assert!(matches!( + bound.acquire(&c), + Err(NetworkPolicyDenied::Destination) + )); + let permit = bound.acquire(&b).unwrap(); + + policy.invalidate(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + assert!(matches!( + bound.acquire(&b), + Err(NetworkPolicyDenied::Revoked) + )); + assert_eq!(permit.check(), Err(NetworkPolicyDenied::Revoked)); +} + +#[test] +fn credential_binding_cannot_replace_a_callers_policy_or_revive_its_account() { + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + let allowed = Url::parse("https://example.com/allowed").unwrap(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let caller = policy + .clone() + .for_current_account() + .restrict_to_endpoints(BTreeSet::from([allowed.clone()])); + let bound = caller + .clone() + .with_account_binding(&NetworkPolicy::unmanaged()) + .unwrap(); + assert!(bound.acquire(&allowed).is_ok()); + assert!(matches!( + bound.acquire(&Url::parse("https://example.com/other").unwrap()), + Err(NetworkPolicyDenied::Destination) + )); + assert!( + caller + .clone() + .with_account_binding(&NetworkPolicyController::default().policy()) + .is_none() + ); + + policy.invalidate(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + assert!( + caller + .with_account_binding(&policy.for_current_account()) + .is_none() + ); + assert!(matches!( + bound.acquire(&allowed), + Err(NetworkPolicyDenied::Revoked) + )); +} + #[test] fn restricted_policy_matches_exact_secure_hosts() { let controller = NetworkPolicyController::default(); diff --git a/codex-rs/http-client/src/openai_endpoints.rs b/codex-rs/http-client/src/openai_endpoints.rs new file mode 100644 index 0000000000..ba7fcf217b --- /dev/null +++ b/codex-rs/http-client/src/openai_endpoints.rs @@ -0,0 +1,155 @@ +//! Classify OpenAI API destinations without treating an unsupported origin as another provider. + +use crate::is_allowed_chatgpt_host; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OpenAiDestination { + Supported, + Unsupported, + Other, +} + +pub fn classify_openai_api_url(destination: &str) -> OpenAiDestination { + let Ok(url) = reqwest::Url::parse(destination) else { + return OpenAiDestination::Other; + }; + let Some(host) = url.host_str() else { + return OpenAiDestination::Other; + }; + let host = host.strip_suffix('.').unwrap_or(host); + let openai_owned = [ + "openai.com", + "openai.org", + "openai-staging.com", + "chatgpt.com", + "chatgpt-staging.com", + ] + .iter() + .any(|domain| { + let host = host.trim_end_matches('.'); + host == *domain + || host + .strip_suffix(*domain) + .is_some_and(|prefix| prefix.ends_with('.')) + }); + if !openai_owned { + return OpenAiDestination::Other; + } + let supported_host = matches!( + host, + "api.openai.com" + | "gov.api.openai.com" + | "us.api.openai.com" + | "eu.api.openai.com" + | "au.api.openai.com" + | "ca.api.openai.com" + | "jp.api.openai.com" + | "in.api.openai.com" + | "sg.api.openai.com" + | "kr.api.openai.com" + | "gb.api.openai.com" + | "ae.api.openai.com" + | "mtls.api.openai.com" + | "mtls-us.api.openai.com" + | "mtls-eu.api.openai.com" + | "api.openai.org" + | "api.openai-staging.com" + ) || is_allowed_chatgpt_host(host); + if supported_host + && url.scheme() == "https" + && url.port_or_known_default() == Some(443) + && url.username().is_empty() + && url.password().is_none() + { + OpenAiDestination::Supported + } else { + OpenAiDestination::Unsupported + } +} + +pub fn is_openai_api_url(destination: &str) -> bool { + classify_openai_api_url(destination) == OpenAiDestination::Supported +} + +#[cfg(test)] +mod tests { + use super::OpenAiDestination; + use super::classify_openai_api_url; + use super::is_openai_api_url; + + #[test] + fn unsupported_openai_origins_are_distinct_from_other_providers() { + for destination in [ + "https://unknown.api.openai.com/v1", + "https://api.openai.com../v1", + "https://api.openai.com:8443/v1", + "http://api.openai.com/v1", + "https://user@api.openai.com/v1", + "https://other.openai.org/v1", + ] { + assert_eq!( + classify_openai_api_url(destination), + OpenAiDestination::Unsupported, + "{destination}" + ); + } + for destination in [ + "https://api.openai.com.example.test/v1", + "https://example.test/v1", + "http://127.0.0.1:8080/v1", + ] { + assert_eq!( + classify_openai_api_url(destination), + OpenAiDestination::Other, + "{destination}" + ); + } + } + + #[test] + fn runtime_credentials_require_a_trusted_https_origin() { + for destination in [ + "https://api.openai.com/v1", + "https://api.openai.com.:443/v1", + "https://gov.api.openai.com/v1", + "https://eu.api.openai.com/v1", + "https://us.api.openai.com/v1", + "https://au.api.openai.com/v1", + "https://ca.api.openai.com/v1", + "https://jp.api.openai.com/v1", + "https://in.api.openai.com/v1", + "https://sg.api.openai.com/v1", + "https://kr.api.openai.com/v1", + "https://gb.api.openai.com/v1", + "https://ae.api.openai.com/v1", + "https://mtls.api.openai.com/v1", + "https://mtls-us.api.openai.com/v1", + "https://mtls-eu.api.openai.com/v1", + "https://api.openai.org/v1", + "https://api.openai-staging.com/v1", + "https://chatgpt.com/backend-api", + "https://chat.openai.com/backend-api", + "https://us.chatgpt.com/backend-api", + "https://chatgpt-staging.com/backend-api", + "https://us.chatgpt-staging.com/backend-api", + ] { + assert!(is_openai_api_url(destination), "{destination}"); + } + for destination in [ + "http://api.openai.com/v1", + "https://api.openai.com:8443/v1", + "https://user@api.openai.com/v1", + "https://user:password@chatgpt.com/backend-api", + "https://api.openai.com.example.test/v1", + "https://eu.api.openai.com.example.test/v1", + "https://unknown.api.openai.com/v1", + "https://evilchatgpt.com/backend-api", + "https://chatgpt.com.example.test/backend-api", + "https://foo.chat.openai.com/backend-api", + "http://127.0.0.1:8080/v1", + "https://example.test/v1", + ] { + assert!(!is_openai_api_url(destination), "{destination}"); + } + } +} diff --git a/codex-rs/http-client/src/outbound_proxy.rs b/codex-rs/http-client/src/outbound_proxy.rs index c8831f60ea..81e2b1fea2 100644 --- a/codex-rs/http-client/src/outbound_proxy.rs +++ b/codex-rs/http-client/src/outbound_proxy.rs @@ -167,6 +167,7 @@ impl fmt::Debug for OutboundProxyRoute { pub struct HttpClientFactory { outbound_proxy_policy: OutboundProxyPolicy, system_proxy_fallback: bool, + protect_sensitive_headers: bool, chatgpt_cookie_store: Option>, network_policy: NetworkPolicy, } @@ -175,6 +176,7 @@ impl PartialEq for HttpClientFactory { fn eq(&self, other: &Self) -> bool { self.outbound_proxy_policy == other.outbound_proxy_policy && self.system_proxy_fallback == other.system_proxy_fallback + && self.protect_sensitive_headers == other.protect_sensitive_headers && self.network_policy == other.network_policy && self .chatgpt_cookie_store @@ -194,6 +196,7 @@ impl fmt::Debug for HttpClientFactory { f.debug_struct("HttpClientFactory") .field("outbound_proxy_policy", &self.outbound_proxy_policy) .field("system_proxy_fallback", &self.system_proxy_fallback) + .field("protect_sensitive_headers", &self.protect_sensitive_headers) .field("network_policy", &self.network_policy) .finish() } @@ -205,11 +208,24 @@ impl HttpClientFactory { Self { outbound_proxy_policy, system_proxy_fallback: false, + protect_sensitive_headers: false, chatgpt_cookie_store: None, network_policy: NetworkPolicy::unmanaged(), } } + /// Follows redirects while keeping headers marked sensitive on the original origin. + /// Reqwest's default redirect handling only recognizes standard credential headers. + pub fn with_sensitive_header_protection(mut self) -> Self { + self.protect_sensitive_headers = true; + self + } + + /// Whether redirect handling removes sensitive headers when the origin changes. + pub fn protects_sensitive_headers(&self) -> bool { + self.protect_sensitive_headers + } + /// Carries the account/configuration owner's application policy into every transport. pub fn with_network_policy(mut self, policy: NetworkPolicy) -> Self { self.network_policy = policy; 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 dee6fea261..f327b1045f 100644 --- a/codex-rs/http-client/src/route_aware_client_pool.rs +++ b/codex-rs/http-client/src/route_aware_client_pool.rs @@ -564,6 +564,7 @@ impl RouteAwareClientPool { && (self.http_client_factory.outbound_proxy_policy() == OutboundProxyPolicy::RespectSystemProxy || self.http_client_factory.network_policy().is_managed() + || self.http_client_factory.protects_sensitive_headers() || self.rustls_clients.is_some()) } diff --git a/codex-rs/http-client/src/route_aware_client_pool/execution.rs b/codex-rs/http-client/src/route_aware_client_pool/execution.rs index cb5361914b..b96f9a4b95 100644 --- a/codex-rs/http-client/src/route_aware_client_pool/execution.rs +++ b/codex-rs/http-client/src/route_aware_client_pool/execution.rs @@ -211,7 +211,12 @@ impl RouteAwareClientPool { if redirects >= MAX_REDIRECTS { return Err(RouteAwareRequestError::TooManyRedirects); } - remove_sensitive_headers(next_request.headers_mut(), ¤t_url, &next_request_url); + remove_sensitive_headers( + next_request.headers_mut(), + ¤t_url, + &next_request_url, + self.http_client_factory.protects_sensitive_headers(), + ); insert_referer(next_request.headers_mut(), ¤t_url, &next_request_url); request = next_request; redirects += 1; 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 711bc6b0e5..0c4a453976 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 @@ -23,6 +23,57 @@ use super::*; use crate::OutboundProxyPolicy; use crate::request_draft::RequestDraft; +#[tokio::test] +async fn redirects_honor_sensitive_header_protection_without_changing_default_behavior() { + for policy in [ + OutboundProxyPolicy::ReqwestDefault, + OutboundProxyPolicy::RespectSystemProxy, + ] { + for protect_sensitive_headers in [false, true] { + let (target, target_server) = spawn_response_server(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok".into(), + ]); + let (source, source_server) = spawn_response_server(vec![ + "HTTP/1.1 307 Temporary Redirect\r\nLocation: /same\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".into(), + format!("HTTP/1.1 307 Temporary Redirect\r\nLocation: http://{target}/final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"), + ]); + let url = format!("http://{source}/start"); + let factory = HttpClientFactory::new(policy); + let factory = if protect_sensitive_headers { + factory.with_sensitive_header_protection() + } else { + factory + }; + let client = factory.build_client(&url, ClientRouteClass::Api).unwrap(); + let mut proof = HeaderValue::from_static("private-proof"); + proof.set_sensitive(/*val*/ true); + let response = client + .get(&url) + .header("authorization", "Bearer private-token") + .header("x-custom-proof", proof) + .header("x-public", "retained") + .send() + .await + .unwrap(); + assert_eq!(response.text().await.unwrap(), "ok"); + let source_requests = source_server.join().unwrap(); + assert_eq!(source_requests.len(), 2); + for request in source_requests { + assert!(request.contains("authorization: Bearer private-token\r\n")); + assert!(request.contains("x-custom-proof: private-proof\r\n")); + } + let target_requests = target_server.join().unwrap(); + assert_eq!(target_requests.len(), 1); + assert!(!target_requests[0].contains("private-token")); + assert_eq!( + target_requests[0].contains("private-proof"), + !protect_sensitive_headers + ); + assert!(target_requests[0].contains("x-public: retained\r\n")); + } + } +} + #[tokio::test] async fn initial_url_auth_is_built_once_after_routing_and_retained_only_on_same_origin_redirects() { let (first_proxy, first_server) = spawn_response_server(vec![ diff --git a/codex-rs/http-client/src/route_aware_redirect.rs b/codex-rs/http-client/src/route_aware_redirect.rs index 77af745725..23f0a3f26b 100644 --- a/codex-rs/http-client/src/route_aware_redirect.rs +++ b/codex-rs/http-client/src/route_aware_redirect.rs @@ -7,9 +7,10 @@ //! for the new URL without giving [`crate::RouteAwareClientPool`] an opportunity to resolve the //! redirect destination. //! -//! The pool therefore disables reqwest redirects for `RespectSystemProxy`, observes each redirect -//! response, and uses these helpers to construct the next request before resolving its route. The -//! helpers preserve the relevant reqwest method/body behavior while applying stricter +//! The pool disables reqwest redirects when route or sensitive-header policy requires checking +//! each hop. It observes each redirect response and uses these helpers to construct the next +//! request before resolving its route. The helpers preserve the relevant reqwest +//! method/body behavior while applying stricter //! origin-sensitive credential and `Referer` rules at the boundary between requests. use std::time::Duration; @@ -99,12 +100,28 @@ pub(super) fn remove_sensitive_headers( headers: &mut HeaderMap, previous: &reqwest::Url, next: &reqwest::Url, + protect_sensitive_headers: bool, ) { if !same_origin(previous, next) { for header in [AUTHORIZATION, COOKIE, PROXY_AUTHORIZATION, WWW_AUTHENTICATE] { headers.remove(header); } headers.remove("cookie2"); + if protect_sensitive_headers { + let sensitive: Vec<_> = headers + .keys() + .filter(|name| { + headers + .get_all(*name) + .iter() + .any(http::HeaderValue::is_sensitive) + }) + .cloned() + .collect(); + for name in sensitive { + headers.remove(name); + } + } } } diff --git a/codex-rs/http-client/src/route_aware_redirect_tests.rs b/codex-rs/http-client/src/route_aware_redirect_tests.rs index d9e8f6053f..f7e1d88cf8 100644 --- a/codex-rs/http-client/src/route_aware_redirect_tests.rs +++ b/codex-rs/http-client/src/route_aware_redirect_tests.rs @@ -219,7 +219,12 @@ fn redirect_credentials_are_retained_only_for_the_same_origin() { (COOKIE, HeaderValue::from_static("session=secret")), ]); - remove_sensitive_headers(&mut headers, &previous, &next); + remove_sensitive_headers( + &mut headers, + &previous, + &next, + /*protect_sensitive_headers*/ false, + ); assert_eq!( ( diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index 28a952ea2d..7d524a0e7c 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -223,7 +223,9 @@ pub fn create_client_with_chatgpt_cookies(http_client_factory: &HttpClientFactor let builder = default_http_client_builder() .with_chatgpt_cookies(http_client_factory) .without_request_logging(); - if !http_client_factory.network_policy().is_managed() { + if !http_client_factory.network_policy().is_managed() + && !http_client_factory.protects_sensitive_headers() + { return build_default_client(builder); } let mut pool = RouteAwareClientPool::with_builder( @@ -274,7 +276,9 @@ fn create_client_for_route_with_builder( route_class: ClientRouteClass, builder: HttpClientBuilder, ) -> Result { - if http_client_factory.network_policy().is_managed() { + if http_client_factory.network_policy().is_managed() + || http_client_factory.protects_sensitive_headers() + { let mut pool = codex_http_client::RouteAwareClientPool::with_builder( http_client_factory.clone(), route_class, @@ -392,6 +396,7 @@ pub async fn create_transport_for_routes_async( OutboundProxyPolicy::ReqwestDefault ); if !http_client_factory.network_policy().is_managed() + && !http_client_factory.protects_sensitive_headers() && (transport_default_proxy || sandboxed) { return ReqwestTransport::from_http_client(create_client()); diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index a03dfef900..9bf505dfa9 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -434,7 +434,8 @@ other non-default provider fields are not supported" } } - fn build_header_map(&self) -> CodexResult { + /// Resolve configured request headers before choosing credentials or a runtime. + pub fn build_header_map(&self) -> CodexResult { let capacity = self.http_headers.as_ref().map_or(0, HashMap::len) + self.env_http_headers.as_ref().map_or(0, HashMap::len); let mut headers = HeaderMap::with_capacity(capacity); diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index ffda92383e..4d0a9aadfd 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -20,7 +20,9 @@ codex-agent-identity = { workspace = true } codex-aws-auth = { workspace = true } codex-http-client = { workspace = true } codex-feedback = { workspace = true } +codex-features = { workspace = true } codex-login = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-keyring-store = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } @@ -40,7 +42,10 @@ url = { workspace = true } base64 = { workspace = true } codex-utils-cargo-bin = { workspace = true } codex-utils-redacted-string = { workspace = true } +codex-utils-rustls-provider = { workspace = true } pretty_assertions = { workspace = true } +rcgen = { workspace = true } +rustls = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } wiremock = { workspace = true } diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index e542b9d70e..4fcdd36a9f 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -10,6 +10,7 @@ use codex_api::AuthError; use codex_api::AuthHeadersFuture; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; +use codex_features::RuntimeAuthMode; use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; use codex_login::AuthHeaders; use codex_login::AuthManager; @@ -20,6 +21,8 @@ use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; use codex_protocol::protocol::SessionSource; +use codex_runtime_agent_instance::RuntimeAuth; +use codex_runtime_agent_instance::RuntimeAuthManager; use http::HeaderMap; use http::HeaderValue; @@ -28,6 +31,24 @@ use crate::bearer_auth_provider::BearerAuthProvider; const BEDROCK_API_KEY_UNSUPPORTED_MESSAGE: &str = "Bedrock API key auth is only supported by the Amazon Bedrock model provider"; +/// Credential ownership and optional OpenAI runtime auth for one harness/thread context. +/// Publish these together so auxiliary clients cannot select another thread's credentials. +#[derive(Clone, Debug)] +pub struct ProviderAuthContext { + pub credentials: Arc, + pub runtime_auth: Option<(Arc, RuntimeAuthMode)>, +} + +impl ProviderAuthContext { + pub fn model_provider(&self, info: ModelProviderInfo) -> crate::SharedModelProvider { + crate::provider::create_model_provider_with_runtime( + info, + Some(Arc::clone(&self.credentials)), + self.runtime_auth.clone(), + ) + } +} + #[derive(Clone, Debug)] pub struct ProviderAuthScope { pub agent_identity_policy: AgentIdentityAuthPolicy, @@ -55,6 +76,7 @@ impl AgentIdentitySessionFallback { pub struct ResolvedProviderAuth { pub auth: SharedAuthProvider, pub agent_identity_telemetry: Option, + pub runtime: Option, } impl ResolvedProviderAuth { @@ -62,6 +84,7 @@ impl ResolvedProviderAuth { Self { auth, agent_identity_telemetry: None, + runtime: None, } } @@ -70,6 +93,7 @@ impl ResolvedProviderAuth { Self { auth: Arc::new(AgentIdentityAuthProvider { auth }), agent_identity_telemetry: Some(agent_identity_telemetry), + runtime: None, } } } diff --git a/codex-rs/model-provider/src/combined_auth.rs b/codex-rs/model-provider/src/combined_auth.rs index e4518a8922..6165dce5bf 100644 --- a/codex-rs/model-provider/src/combined_auth.rs +++ b/codex-rs/model-provider/src/combined_auth.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use crate::auth::ResolvedProviderAuth; use codex_api::AuthError; +use codex_api::AuthHeaderTelemetry; use codex_api::AuthHeadersFuture; use codex_api::AuthProvider; use codex_api::AuthProviderFuture; @@ -57,16 +58,6 @@ pub(crate) async fn compose_auth( )) })?; let (name, value) = gateway_header(config, &token)?; - if resolved - .auth - .to_auth_headers() - .map_err(|error| codex_api::map_api_error(error.into()))? - .contains_key(&name) - { - return Err(CodexErr::InvalidRequest(format!( - "The model provider already sets the {name} header. Choose a different header or cookie in gateway_oauth.delivery for the gateway token." - ))); - } resolved.auth = Arc::new(CombinedAuth { primary: resolved.auth, name, @@ -125,19 +116,15 @@ impl CombinedAuth { } impl AuthProvider for CombinedAuth { + fn auth_header_telemetry(&self) -> std::result::Result { + self.primary.auth_header_telemetry() + } + fn add_auth_headers(&self, headers: &mut HeaderMap) -> std::result::Result<(), AuthError> { self.primary.add_auth_headers(headers)?; self.insert_gateway_header(headers) } - fn apply_auth(&self, request: codex_http_client::Request) -> AuthProviderFuture<'_> { - Box::pin(async move { - let mut request = self.primary.apply_auth(request).await?; - self.insert_gateway_header(&mut request.headers)?; - Ok(request) - }) - } - fn apply_auth_headers(&self, headers: HeaderMap) -> AuthHeadersFuture<'_> { Box::pin(async move { let mut headers = self.primary.apply_auth_headers(headers).await?; @@ -145,4 +132,12 @@ impl AuthProvider for CombinedAuth { Ok(headers) }) } + + fn apply_auth(&self, request: codex_http_client::Request) -> AuthProviderFuture<'_> { + Box::pin(async move { + let mut request = self.primary.apply_auth(request).await?; + self.insert_gateway_header(&mut request.headers)?; + Ok(request) + }) + } } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 4b13035963..f054a8f606 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -8,6 +8,7 @@ mod gcp; mod models_endpoint; mod models_identity; mod provider; +mod runtime_auth; mod shared_state; pub mod test_support; mod workspace_routing; @@ -20,6 +21,7 @@ pub use amazon_bedrock::is_amazon_bedrock_gov_cloud_region; pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use amazon_bedrock::resolve_amazon_bedrock_region; pub use auth::AgentIdentitySessionFallback; +pub use auth::ProviderAuthContext; pub use auth::ProviderAuthScope; pub use auth::ResolvedProviderAuth; pub use auth::auth_provider_from_auth; @@ -48,6 +50,11 @@ pub use provider::create_model_provider; pub use provider::provider_uses_first_party_auth_path; pub use provider::resolve_api_provider; +pub use runtime_auth::select_runtime_auth; +pub use runtime_auth::validate_runtime_destination; +pub use runtime_auth::validate_runtime_provider; +pub use runtime_auth::validate_runtime_selectors; + #[cfg(test)] #[path = "workspace_routing_tests.rs"] mod workspace_routing_tests; diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 31bb98636e..5a0f3652de 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -12,6 +12,7 @@ use codex_api::ReqwestTransport; use codex_api::TransportError; use codex_api::auth_header_telemetry; use codex_api::map_api_error; +use codex_features::RuntimeAuthMode; use codex_feedback::FeedbackRequestTags; use codex_feedback::emit_feedback_request_tags_with_auth_env; use codex_http_client::ClientRouteClass; @@ -34,6 +35,7 @@ use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CoreResult; use codex_response_debug_context::extract_response_debug_context; use codex_response_debug_context::telemetry_transport_error_message; +use codex_runtime_agent_instance::RuntimeAuthManager; use http::HeaderMap; use tokio::time::timeout; @@ -55,6 +57,7 @@ pub(crate) struct OpenAiModelsEndpoint { auth_manager: Option>, gateway_auth_manager: Option, String>>, transport_builder: Arc, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, } impl OpenAiModelsEndpoint { @@ -69,6 +72,7 @@ impl OpenAiModelsEndpoint { ClientRedirectPolicy::Default }; Self { + runtime_auth: None, provider_info, auth_manager, gateway_auth_manager, @@ -76,6 +80,15 @@ impl OpenAiModelsEndpoint { } } + /// Use the same runtime owner as the harness's model sessions. + pub(crate) fn with_runtime_auth( + mut self, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, + ) -> Self { + self.runtime_auth = runtime_auth; + self + } + async fn auth(&self) -> Option { match self.auth_manager.as_ref() { Some(auth_manager) => auth_manager.auth().await, @@ -95,7 +108,61 @@ impl OpenAiModelsEndpoint { client_version: &str, http_client_factory: HttpClientFactory, ) -> CoreResult { - let auth = self.auth().await; + timeout( + MODELS_REFRESH_TIMEOUT, + self.fetch_models(client_version, http_client_factory), + ) + .await + .map_err(|_| CodexErr::RequestTimeout)? + } + + async fn fetch_models( + &self, + client_version: &str, + http_client_factory: HttpClientFactory, + ) -> CoreResult { + let provider_info = self.provider_info.clone(); + let auth_manager = self.auth_manager.clone(); + let runtime_auth = self.runtime_auth.clone(); + // A timed-out catalog fetch must not discard a rotated refresh token before + // it is saved. Keep credential refresh and registration owned until they finish. + let (auth, selected_runtime_auth) = tokio::spawn(async move { + let selected_runtime_auth = match runtime_auth { + Some((runtime, mode)) => { + crate::validate_runtime_destination( + provider_info + .model_catalog_url + .as_deref() + .map(String::as_str) + .or(provider_info.base_url.as_deref()), + )?; + crate::validate_runtime_provider( + &provider_info, + auth_manager.as_ref(), + &runtime, + )?; + Some(( + crate::runtime_auth::select_runtime_auth( + &runtime, + &provider_info.build_header_map()?, + ) + .await + .map_err(|error| map_api_error(error.into()))?, + mode, + )) + } + None => None, + }; + let auth = match &selected_runtime_auth { + Some((selected_runtime_auth, _)) => Some(selected_runtime_auth.auth().clone()), + None => match auth_manager { + Some(auth_manager) => auth_manager.auth().await, + None => None, + }, + }; + Ok::<_, CodexErr>((auth, selected_runtime_auth)) + }) + .await??; let metric_auth_mode = if self.has_provider_api_key() || auth.as_ref().is_some_and(CodexAuth::is_api_key_auth) { @@ -125,10 +192,25 @@ impl OpenAiModelsEndpoint { } .to_string(); } + let http_client_factory = match &selected_runtime_auth { + Some((selected_runtime_auth, _)) => selected_runtime_auth + .bind_http_client_factory(http_client_factory) + .map_err(|error| map_api_error(error.into()))?, + None => http_client_factory, + }; + let primary = if let Some((selected_runtime_auth, mode)) = selected_runtime_auth { + crate::validate_runtime_selectors( + &api_provider.headers, + selected_runtime_auth.runtime().actor().info(), + )?; + selected_runtime_auth.into_auth_provider(mode) + } else { + resolve_provider_auth(auth.as_ref(), &self.provider_info)? + }; let resolved = compose_auth( &self.provider_info, self.gateway_auth_manager.as_ref(), - ResolvedProviderAuth::new(resolve_provider_auth(auth.as_ref(), &self.provider_info)?), + ResolvedProviderAuth::new(primary), ) .await?; let api_auth = resolved.auth; @@ -156,7 +238,7 @@ impl OpenAiModelsEndpoint { agent_identity_telemetry, auth_env: self.auth_env(), }); - let (models, etag) = timeout(MODELS_REFRESH_TIMEOUT, async { + let (models, etag) = { let transport = self .transport_builder .build(http_client_factory, request_url.clone()) @@ -186,9 +268,7 @@ impl OpenAiModelsEndpoint { } map_api_error(error) }) - }) - .await - .map_err(|_| CodexErr::RequestTimeout)??; + }?; if self.provider_info.model_catalog_url.is_some() { let mut slugs = HashSet::new(); if models @@ -232,6 +312,7 @@ impl ModelsEndpointClient for OpenAiModelsEndpoint { provider_info, auth_manager: self.auth_manager.clone(), gateway_auth_manager: self.gateway_auth_manager.clone(), + runtime_auth: self.runtime_auth.clone(), transport_builder: Arc::clone(&self.transport_builder), })) } @@ -551,6 +632,7 @@ mod tests { }, auth_manager: Some(auth.clone()), gateway_auth_manager: None, + runtime_auth: None, transport_builder: capture.clone(), }); let provider = crate::provider::create_model_provider( @@ -625,6 +707,7 @@ mod tests { }, auth_manager: auth_manager.clone(), gateway_auth_manager: None, + runtime_auth: None, transport_builder: capture.clone(), }); let provider = crate::provider::create_model_provider( @@ -717,6 +800,7 @@ mod tests { provider_info: ModelProviderInfo::create_openai_provider(Some(server.uri())), auth_manager: None, gateway_auth_manager: None, + runtime_auth: None, transport_builder: Arc::new(RecordingTransportBuilder { observed_request: Arc::clone(&observed_request), }), @@ -763,6 +847,7 @@ mod tests { provider_info, auth_manager: None, gateway_auth_manager: None, + runtime_auth: None, transport_builder: Arc::new(RecordingTransportBuilder { observed_request: Arc::new(Mutex::new(None)), }), @@ -1055,3 +1140,11 @@ mod tests { assert_eq!(destination.received_requests().await.unwrap().len(), 0); } } + +#[cfg(test)] +#[path = "models_endpoint_runtime_auth_tests.rs"] +mod runtime_auth_tests; + +#[cfg(test)] +#[path = "models_endpoint_refresh_tests.rs"] +mod refresh_tests; diff --git a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs new file mode 100644 index 0000000000..6bebf20ce9 --- /dev/null +++ b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs @@ -0,0 +1,200 @@ +//! A catalog deadline must not interrupt rotating OAuth credentials before persistence. +#![cfg(test)] + +use super::*; +use base64::Engine; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; +use codex_login::AuthRouteConfig; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test] +async fn catalog_timeout_preserves_refreshed_tokens_with_and_without_runtime_auth() { + for mode in [ + None, + Some(RuntimeAuthMode::DualHeader), + Some(RuntimeAuthMode::AssertionOnly), + ] { + let server = MockServer::start().await; + let home = tempfile::tempdir().unwrap(); + let token = |expires| { + let payload = json!({ + "iss": server.uri(), "session_id": "session-a", "exp": expires, + "https://api.openai.com/auth": { + "chatgpt_user_id": "user-a", "chatgpt_account_id": "workspace-a", + "chatgpt_plan_type": "plus", + }, + }); + format!( + "e30.{}.signature", + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload.to_string()) + ) + }; + let old_token = token(0); + let new_token = token(chrono::Utc::now().timestamp() + 3600); + let stored = serde_json::from_value(json!({ + "auth_mode": "chatgpt", + "tokens": { + "id_token": old_token, "access_token": old_token, + "refresh_token": "refresh-before", "account_id": "workspace-a", + "oauth_provenance": {"issuer": server.uri(), "client_id": "test-client"}, + }, + "last_refresh": "2020-01-01T00:00:00Z", + })) + .unwrap(); + codex_login::save_auth( + home.path(), + &stored, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .unwrap(); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_partial_json( + json!({"refresh_token": "refresh-before"}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_delay(MODELS_REFRESH_TIMEOUT + Duration::from_secs(1)) + .set_body_json(json!({ + "id_token": new_token, "access_token": new_token, + "refresh_token": "refresh-after", + })), + ) + .up_to_n_times(1) + .with_priority(1) + .expect(1) + .mount(&server) + .await; + // Reusing the consumed refresh token must fail, as it would at the issuer. + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": {"code": "refresh_token_reused"}, + }))) + .with_priority(2) + .expect(0) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/register")) + .and(body_partial_json( + json!({"actor_auth": {"credential": new_token}}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", "agent_host_id": "host-a", + "actor": {"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"}, + }))) + .expect(u64::from(mode.is_some())) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"models": []}))) + .expect(1) + .mount(&server) + .await; + + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let credentials = Arc::new( + AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + AuthRouteConfig::from_http_client_factory(factory.clone()), + ) + .await, + ); + let runtime_auth = mode.map(|mode| { + ( + Arc::new(RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/register", server.uri()).parse().unwrap(), + )), + mode, + ) + }); + let mut info = ModelProviderInfo::create_openai_provider(Some(server.uri())); + info.env_http_headers = None; + let endpoint = OpenAiModelsEndpoint::new( + info, + Some(Arc::clone(&credentials)), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(runtime_auth); + let error = endpoint + .list_models("test", factory.clone()) + .await + .unwrap_err(); + assert!(matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::RequestTimeout + )); + + // Wait through the existing refresh lock; this must reuse the completed refresh. + let current = timeout(Duration::from_secs(5), credentials.auth()) + .await + .unwrap() + .unwrap(); + assert_eq!(current.get_token().unwrap(), new_token); + let saved = codex_login::load_auth_dot_json( + home.path(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .unwrap() + .unwrap() + .tokens + .unwrap(); + assert_eq!(saved.access_token, new_token); + assert_eq!(saved.refresh_token, "refresh-after"); + + endpoint.list_models("test", factory).await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let models = requests + .iter() + .find(|request| request.url.path() == "/models") + .unwrap(); + if mode == Some(RuntimeAuthMode::AssertionOnly) { + assert!( + models.headers["authorization"] + .to_str() + .unwrap() + .starts_with("AgentAssertion ") + ); + } else { + assert_eq!( + models.headers["authorization"], + format!("Bearer {new_token}") + ); + } + assert_eq!( + models.headers.contains_key("x-openai-agent-assertion"), + mode == Some(RuntimeAuthMode::DualHeader) + ); + server.verify().await; + } +} diff --git a/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs b/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs new file mode 100644 index 0000000000..1959823a3d --- /dev/null +++ b/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs @@ -0,0 +1,424 @@ +use super::*; +use base64::Engine; +use codex_features::RuntimeAuthMode; +use codex_http_client::OutboundProxyPolicy; +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 serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test] +async fn provider_auth_rejects_selectors_that_do_not_match_the_registered_actor() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + for selector in ["chatgpt-account-id", "openai-organization"] { + let server = MockServer::start().await; + let (auth, actor, expected) = if selector == "chatgpt-account-id" { + let claims = json!({"session_id": "session-a", "https://api.openai.com/auth": {"user_id": "user-a", "chatgpt_account_id": "workspace-a"}}); + let token = format!( + "e30.{}.signature", + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()) + ); + ( + CodexAuth::from_external_chatgpt_tokens( + &token, + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap(), + json!({"kind":"chatgpt", "user_id":"user-a", "account_id":"workspace-a"}), + "workspace-a", + ) + } else { + ( + CodexAuth::from_api_key("key-a"), + json!({"kind":"platform", "user_id":"user-a", "organization_id":"org-a", "project_id":"project-a"}), + "org-a", + ) + }; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", "agent_host_id": "host-a", "actor": actor, + }))) + .mount(&server) + .await; + let credentials = AuthManager::from_auth_for_testing(auth); + let context = crate::ProviderAuthContext { + runtime_auth: Some((runtime_manager(&server, Arc::clone(&credentials)), mode)), + credentials, + }; + let mut info = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + info.env_http_headers = None; + info.http_headers = Some(std::collections::HashMap::from([( + selector.into(), + expected.into(), + )])); + let provider = context.model_provider(info.clone()); + let auth = provider.api_auth().await.unwrap(); + assert!( + auth.to_auth_headers() + .unwrap() + .contains_key("authorization") + ); + + info.http_headers = Some(std::collections::HashMap::from([( + selector.into(), + "different".into(), + )])); + let provider = context.model_provider(info); + let error = provider + .api_auth() + .await + .err() + .expect("mismatched selector must be rejected"); + assert!(matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::UnsupportedOperation(_) + )); + assert!(error.to_string().contains(selector)); + assert!(!server.received_requests().await.unwrap().is_empty()); + } + } +} + +#[tokio::test] +async fn model_requests_keep_the_selected_project_while_reusing_one_runtime() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", + "agent_host_id": "host-a", + "actor": {"kind":"platform", "user_id":"user-a", "organization_id":"org-a", "project_id":"project-a"}, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"models":[]}))) + .expect(4) + .mount(&server) + .await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let runtime = runtime_manager(&server, Arc::clone(&credentials)); + let context = crate::ProviderAuthContext { + runtime_auth: Some((Arc::clone(&runtime), mode)), + credentials: Arc::clone(&credentials), + }; + let projects = [ + Some("project-a"), + Some("project-b"), + None, + Some("project-a"), + ]; + for project in projects { + let mut info = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + info.env_http_headers = None; + let mut headers = + std::collections::HashMap::from([("openai-organization".into(), "org-a".into())]); + if let Some(project) = project { + headers.insert("openai-project".into(), project.into()); + } + info.http_headers = Some(headers); + let auth = context + .model_provider(info.clone()) + .api_auth() + .await + .unwrap(); + assert!( + auth.to_auth_headers() + .unwrap() + .contains_key("authorization") + ); + + info.base_url = Some(server.uri()); + let endpoint = OpenAiModelsEndpoint::new( + info, + Some(Arc::clone(&credentials)), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(Some((Arc::clone(&runtime), mode))); + endpoint + .list_models( + "test", + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await + .unwrap(); + } + let requests = server.received_requests().await.unwrap(); + let models: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/models") + .collect(); + assert_eq!(models.len(), projects.len()); + for (request, project) in models.into_iter().zip(projects) { + assert_eq!( + request + .headers + .get("openai-project") + .map(|value| value.to_str().unwrap()), + project + ); + assert_eq!(request.headers["openai-organization"], "org-a"); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(request.headers["authorization"], "Bearer key-a"); + assert!(request.headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!( + request.headers["authorization"] + .to_str() + .unwrap() + .starts_with("AgentAssertion ") + ); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + } + } + server.verify().await; + } +} + +#[tokio::test] +async fn unsupported_openai_destinations_do_not_fall_back_to_bearer_auth() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let runtime = runtime_manager(&server, Arc::clone(&credentials)); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + for destination in [ + "https://unknown.api.openai.com/v1", + "http://api.openai.com/v1", + "https://api.openai.com:8443/v1", + ] { + let info = ModelProviderInfo::create_openai_provider(Some(destination.into())); + let provider = crate::provider::create_model_provider_with_runtime( + info.clone(), + Some(Arc::clone(&credentials)), + Some((Arc::clone(&runtime), mode)), + ); + let error = provider + .request_auth(/*scope*/ None) + .await + .err() + .expect("unsupported inference origin"); + assert!( + error + .to_string() + .contains("does not support this OpenAI endpoint") + ); + let mut catalog_info = info; + catalog_info.base_url = Some("https://api.openai.com/v1".into()); + catalog_info.model_catalog_url = Some(format!("{destination}/models").into()); + let catalog = OpenAiModelsEndpoint::new( + catalog_info, + Some(Arc::clone(&credentials)), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(Some((Arc::clone(&runtime), mode))); + let error = catalog + .list_models("0.0.0", factory.clone()) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("does not support this OpenAI endpoint") + ); + } + // An explicitly configured other provider retains its existing bearer auth. + let info = ModelProviderInfo::create_openai_provider(Some(server.uri())); + let provider = crate::provider::create_model_provider_with_runtime( + info, + Some(credentials), + Some((runtime, mode)), + ); + let resolved = provider.request_auth(/*scope*/ None).await.unwrap(); + let headers = resolved + .auth + .apply_auth_headers(http::HeaderMap::new()) + .await + .unwrap(); + assert_eq!(headers["authorization"], "Bearer key-a"); + assert!(!headers.contains_key("x-openai-agent-assertion")); + assert!(server.received_requests().await.unwrap().is_empty()); + } +} + +#[tokio::test] +async fn catalog_timeout_keeps_registration_alive_for_the_next_request() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(6)).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", + "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(method("GET")) + .and(path("/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"models":[]}))) + .expect(1) + .mount(&server) + .await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let runtime = runtime_manager(&server, Arc::clone(&credentials)); + let mut provider = ModelProviderInfo::create_openai_provider(Some(server.uri())); + provider.env_http_headers = None; + let endpoint = OpenAiModelsEndpoint::new( + provider, + Some(credentials), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(Some((runtime, RuntimeAuthMode::DualHeader))); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let result = timeout( + Duration::from_secs(10), + endpoint.list_models("0.0.0", factory.clone()), + ) + .await + .expect("catalog timeout must precede the registration client's timeout"); + assert!(matches!( + result.unwrap_err().details(), + codex_protocol::error::CodexErrorDetails::RequestTimeout + )); + endpoint.list_models("0.0.0", factory).await.unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/register") + .count(), + 1 + ); + server.verify().await; +} + +#[tokio::test] +async fn runtime_catalog_follows_redirects_without_forwarding_credentials_to_another_origin() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + for policy in [ + OutboundProxyPolicy::ReqwestDefault, + OutboundProxyPolicy::RespectSystemProxy, + ] { + let server = MockServer::start().await; + let target = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", + "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(method("GET")) + .and(path("/models")) + .respond_with(ResponseTemplate::new(307).insert_header("location", "/same-origin")) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/same-origin")) + .respond_with( + ResponseTemplate::new(302) + .insert_header("location", format!("{}/models", target.uri())), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"models": []}))) + .expect(1) + .mount(&target) + .await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let runtime = runtime_manager(&server, Arc::clone(&credentials)); + let mut info = ModelProviderInfo::create_openai_provider(Some(server.uri())); + info.env_http_headers = None; + let endpoint = OpenAiModelsEndpoint::new( + info, + Some(credentials), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(Some((runtime, mode))); + endpoint + .list_models("test", HttpClientFactory::new(policy)) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let initial = requests + .iter() + .find(|request| request.url.path() == "/models") + .unwrap(); + let same_origin = requests + .iter() + .find(|request| request.url.path() == "/same-origin") + .unwrap(); + assert_eq!( + same_origin.headers.get("authorization"), + initial.headers.get("authorization") + ); + assert_eq!( + same_origin.headers.get("x-openai-agent-assertion"), + initial.headers.get("x-openai-agent-assertion") + ); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(initial.headers["authorization"], "Bearer key-a"); + assert!(initial.headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!( + initial.headers["authorization"] + .to_str() + .unwrap() + .starts_with("AgentAssertion ") + ); + assert!(!initial.headers.contains_key("x-openai-agent-assertion")); + } + } + let redirected = target.received_requests().await.unwrap(); + assert_eq!(redirected.len(), 1); + assert!(!redirected[0].headers.contains_key("authorization")); + assert!( + !redirected[0] + .headers + .contains_key("x-openai-agent-assertion") + ); + server.verify().await; + target.verify().await; + } + } +} + +fn runtime_manager(server: &MockServer, credentials: Arc) -> Arc { + Arc::new(RuntimeAuthManager::new( + credentials, + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/register", server.uri()).parse().unwrap(), + )) +} diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index b8c786f059..5bd223b949 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -8,6 +8,7 @@ use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; use codex_api::TransportError; +use codex_features::RuntimeAuthMode; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::GatewayAuthManager; @@ -22,6 +23,7 @@ use codex_protocol::account::ProviderAccount; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::error::CodexErr; use codex_protocol::openai_models::ModelsResponse; +use codex_runtime_agent_instance::RuntimeAuthManager; use crate::ProviderCapabilities; #[cfg(test)] @@ -172,6 +174,11 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { /// manager throughout the codebase; that is a larger refactor than this change. fn auth_manager(&self) -> Option>; + /// Runtime authentication selected for this provider's inference destination. + fn runtime_auth(&self) -> Option<&(Arc, RuntimeAuthMode)> { + None + } + /// Returns the gateway credential manager shared with inference and model discovery. /// Hosts use this handle for explicit login; configured setup failures remain errors. fn gateway_auth_manager(&self) -> std::io::Result>> { @@ -337,6 +344,17 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { }) } + /// Resolve auth together with metadata needed to route and admit this request. + fn request_auth( + &self, + scope: Option, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + match scope { + Some(scope) => self.api_auth_for_scope(scope), + None => Box::pin(async { self.api_auth().await.map(ResolvedProviderAuth::new) }), + } + } + /// Creates the model manager implementation appropriate for this provider. fn models_manager( &self, @@ -399,10 +417,19 @@ pub fn provider_uses_first_party_auth_path(provider: &ModelProviderInfo) -> bool && provider.aws.is_none() } -/// Creates the default runtime model provider for configured provider metadata. +/// Creates the default model provider for configured provider metadata. +/// Harnesses with runtime identity use [`crate::ProviderAuthContext::model_provider`]. pub fn create_model_provider( provider_info: ModelProviderInfo, auth_manager: Option>, +) -> SharedModelProvider { + create_model_provider_with_runtime(provider_info, auth_manager, /*runtime_auth*/ None) +} + +pub(crate) fn create_model_provider_with_runtime( + provider_info: ModelProviderInfo, + auth_manager: Option>, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, ) -> SharedModelProvider { if provider_info.is_amazon_bedrock() { return Arc::new(AmazonBedrockModelProvider::new(provider_info, auth_manager)); @@ -410,20 +437,10 @@ pub fn create_model_provider( if provider_info.is_gcp() { return Arc::new(GcpModelProvider::new(provider_info, auth_manager)); } - let gateway_auth_manager = provider_info.gateway_oauth.as_ref().map(|config| { - provider_info.validate()?; - let manager = auth_manager - .as_ref() - .ok_or_else(|| "gateway_oauth requires auth runtime configuration".to_string())?; - crate::shared_state::process_shared_state() - .gateway_auth(config, &manager.runtime_config()) - .map_err(|_| "failed to create provider OAuth HTTP client".to_string()) - }); - let auth_manager = auth_manager_for_provider(auth_manager, &provider_info); Arc::new(ConfiguredModelProvider::new( provider_info, auth_manager, - gateway_auth_manager, + runtime_auth, )) } @@ -434,6 +451,8 @@ struct ConfiguredModelProvider { auth_manager: Option>, // Construct eagerly; report setup failures when auth is requested because the factory is infallible. gateway_auth_manager: Option, String>>, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, + catalog_runtime_auth: Option<(Arc, RuntimeAuthMode)>, } enum ModelsCacheConfig { @@ -446,15 +465,78 @@ impl ConfiguredModelProvider { fn new( info: ModelProviderInfo, auth_manager: Option>, - gateway_auth_manager: Option, String>>, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, ) -> Self { + let gateway_auth_manager = info.gateway_oauth.as_ref().map(|config| { + info.validate()?; + let manager = auth_manager + .as_ref() + .ok_or_else(|| "gateway_oauth requires auth runtime configuration".to_string())?; + crate::shared_state::process_shared_state() + .gateway_auth(config, &manager.runtime_config()) + .map_err(|_| "failed to create provider OAuth HTTP client".to_string()) + }); + let auth_manager = auth_manager_for_provider(auth_manager, &info); + let runtime_for_destination = |destination: Option<&str>| { + runtime_auth.clone().filter(|_| { + destination.is_none_or(|url| { + codex_http_client::classify_openai_api_url(url) + != codex_http_client::OpenAiDestination::Other + }) + }) + }; + let inference_runtime_auth = runtime_for_destination(info.base_url.as_deref()); + // Catalog overrides are absolute URLs and can have a different origin from inference. + let catalog_runtime_auth = runtime_for_destination( + info.model_catalog_url + .as_deref() + .map(String::as_str) + .or(info.base_url.as_deref()), + ); Self { info, auth_manager, gateway_auth_manager, + runtime_auth: inference_runtime_auth, + catalog_runtime_auth, } } + async fn resolve_auth( + &self, + scope: Option, + ) -> codex_protocol::error::Result { + let resolved = if let Some((runtime, mode)) = &self.runtime_auth { + crate::validate_runtime_destination(self.info.base_url.as_deref())?; + crate::validate_runtime_provider(&self.info, self.auth_manager.as_ref(), runtime)?; + let headers = self.info.build_header_map()?; + let selected = crate::runtime_auth::select_runtime_auth(runtime, &headers) + .await + .map_err(|error| self.map_api_error(error.into()))?; + crate::validate_runtime_selectors(&headers, selected.runtime().actor().info())?; + ResolvedProviderAuth { + auth: selected.clone().into_auth_provider(*mode), + agent_identity_telemetry: None, + runtime: Some(selected), + } + } else { + let auth = self.auth().await; + match scope { + Some(scope) if provider_uses_first_party_auth_path(&self.info) => { + resolve_provider_auth_for_scope( + self.auth_manager.clone(), + auth.as_ref(), + &self.info, + scope, + ) + .await? + } + _ => ResolvedProviderAuth::new(resolve_provider_auth(auth.as_ref(), &self.info)?), + } + }; + self.compose_auth(resolved).await + } + fn create_models_manager( &self, config_model_catalog: Option, @@ -466,11 +548,14 @@ impl ConfiguredModelProvider { model_catalog, )); } - let endpoint = Arc::new(OpenAiModelsEndpoint::new( - self.info.clone(), - self.auth_manager.clone(), - self.gateway_auth_manager.clone(), - )); + let endpoint = Arc::new( + OpenAiModelsEndpoint::new( + self.info.clone(), + self.auth_manager.clone(), + self.gateway_auth_manager.clone(), + ) + .with_runtime_auth(self.catalog_runtime_auth.clone()), + ); let auth_manager = self.auth_manager.clone(); let manager = match cache { ModelsCacheConfig::Disk { codex_home } => { @@ -539,16 +624,22 @@ impl ModelProvider for ConfiguredModelProvider { }) } + fn runtime_auth(&self) -> Option<&(Arc, RuntimeAuthMode)> { + self.runtime_auth.as_ref() + } + + fn request_auth( + &self, + scope: Option, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + Box::pin(self.resolve_auth(scope)) + } + fn api_auth( &self, ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { Box::pin(async move { - let auth = self.auth().await; - let primary = resolve_provider_auth(auth.as_ref(), &self.info)?; - Ok(self - .compose_auth(ResolvedProviderAuth::new(primary)) - .await? - .auth) + Ok(self.resolve_auth(/*scope*/ None).await?.auth) }) } @@ -583,22 +674,7 @@ impl ModelProvider for ConfiguredModelProvider { &self, scope: ProviderAuthScope, ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { - Box::pin(async move { - let resolved = if provider_uses_first_party_auth_path(&self.info) { - let auth = self.auth().await; - resolve_provider_auth_for_scope( - self.auth_manager.clone(), - auth.as_ref(), - &self.info, - scope, - ) - .await? - } else { - let auth = self.auth().await; - ResolvedProviderAuth::new(resolve_provider_auth(auth.as_ref(), &self.info)?) - }; - self.compose_auth(resolved).await - }) + Box::pin(self.resolve_auth(Some(scope))) } fn account_state(&self) -> ProviderAccountResult { @@ -1601,3 +1677,20 @@ printf '%s\n' '{"AccessKeyId":"exported","SecretAccessKey":"secret"}' } } } + +pub(crate) mod test_support { + use super::*; + + /// Use an injected runtime with a local transport fixture. Production callers use the + /// destination-aware `ProviderAuthContext::model_provider` factory. + pub fn runtime_model_provider( + info: ModelProviderInfo, + credentials: Arc, + runtime: (Arc, RuntimeAuthMode), + ) -> SharedModelProvider { + let mut provider = + ConfiguredModelProvider::new(info, Some(credentials), Some(runtime.clone())); + provider.runtime_auth = Some(runtime); + Arc::new(provider) + } +} diff --git a/codex-rs/model-provider/src/runtime_auth.rs b/codex-rs/model-provider/src/runtime_auth.rs new file mode 100644 index 0000000000..3a817f428a --- /dev/null +++ b/codex-rs/model-provider/src/runtime_auth.rs @@ -0,0 +1,103 @@ +//! Keep model requests and catalog refreshes bound to the same registered actor. + +use codex_api::AuthError; +use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; +use codex_http_client::OPENAI_ORGANIZATION_HEADER; +use codex_http_client::OPENAI_PROJECT_HEADER; +use codex_http_client::OpenAiDestination; +use codex_http_client::classify_openai_api_url; +use codex_login::AuthManager; +use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::error::CodexErr; +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::RuntimeAuthManager; +use http::HeaderMap; +use std::sync::Arc; + +/// Select the runtime for the resource scope carried by the model request. +pub async fn select_runtime_auth( + manager: &RuntimeAuthManager, + headers: &HeaderMap, +) -> Result { + let selector = |name| { + headers + .get(name) + .map(|value| { + value + .to_str() + .map(str::to_owned) + .map_err(|_| AuthError::Build(format!("invalid {name} resource selector"))) + }) + .transpose() + }; + let selection = 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. +pub fn validate_runtime_provider( + info: &ModelProviderInfo, + credentials: Option<&Arc>, + runtime: &RuntimeAuthManager, +) -> CodexResult<()> { + if !info.requires_openai_auth + || info.env_key.is_some() + || info.experimental_bearer_token.is_some() + || info.auth.is_some() + || info.aws.is_some() + || !credentials.is_some_and(|auth| runtime.uses_credentials(auth)) + { + return Err(CodexErr::UnsupportedOperation( + "runtime auth requires the model provider's selected OpenAI credential".to_string(), + )); + } + Ok(()) +} + +/// An unrecognized OpenAI origin must not silently fall back to legacy auth. +pub fn validate_runtime_destination(destination: Option<&str>) -> CodexResult<()> { + if destination.is_some_and(|url| classify_openai_api_url(url) == OpenAiDestination::Unsupported) + { + return Err(CodexErr::UnsupportedOperation( + "runtime auth does not support this OpenAI endpoint".to_string(), + )); + } + Ok(()) +} + +/// 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<()> { + let validate = |name: &str, expected: Option<&str>| { + if let Some(value) = headers.get(name) + && value.to_str().ok() != expected + { + return Err(CodexErr::UnsupportedOperation(format!( + "configured {name} does not match the registered runtime" + ))); + } + Ok(()) + }; + match actor { + ActorInfo::Chatgpt { account_id, .. } => { + validate(CHATGPT_ACCOUNT_ID_HEADER, Some(account_id))?; + validate(OPENAI_ORGANIZATION_HEADER, None)?; + validate(OPENAI_PROJECT_HEADER, None) + } + ActorInfo::Platform { + organization_id, .. + } => { + validate(CHATGPT_ACCOUNT_ID_HEADER, None)?; + validate(OPENAI_ORGANIZATION_HEADER, Some(organization_id)) + } + } +} diff --git a/codex-rs/model-provider/src/shared_state_tests.rs b/codex-rs/model-provider/src/shared_state_tests.rs index a98a26269f..df75e49c83 100644 --- a/codex-rs/model-provider/src/shared_state_tests.rs +++ b/codex-rs/model-provider/src/shared_state_tests.rs @@ -132,7 +132,7 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { HeaderName::from_static(name), HeaderValue::from_static(value), ); - // WebSocket handshakes use the synchronous header interface. + // Legacy consumers and asynchronous handshakes preserve the same gateway credential. assert_eq!(auth.to_auth_headers().unwrap(), expected); assert_eq!( auth.apply_auth_headers(HeaderMap::new()).await.unwrap(), @@ -177,6 +177,123 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { } } +#[tokio::test] +async fn gateway_preserves_runtime_auth_and_rejects_header_collisions() { + use codex_features::RuntimeAuthMode; + use codex_runtime_agent_instance::AgentHostAuthorization; + use codex_runtime_agent_instance::RuntimeAuthManager; + use codex_runtime_agent_instance::RuntimeRegistrationMetadata; + use codex_runtime_agent_instance::RuntimeRole; + + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let home = tempfile::tempdir().unwrap(); + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", + "agent_host_id": "host-a", + "actor": {"kind":"platform", "user_id":"user-a", "organization_id":"org-a", "project_id":null}, + }))) + .expect(1) + .mount(&server) + .await; + let primary = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("primary-token"), + home.path().to_path_buf(), + ); + let info = ModelProviderInfo { + gateway_oauth: Some(GatewayOAuthConfig { + authorization_url: format!("{}/authorize", server.uri()), + token_url: format!("{}/token", server.uri()), + client_id: "client".into(), + resource: None, + scopes: vec![], + redirect_port: None, + delivery: GatewayOAuthDelivery::Header { + name: "x-gateway-auth".into(), + scheme: "Bearer".into(), + }, + }), + ..ModelProviderInfo::create_openai_provider(Some(server.uri())) + }; + let _gateway = seed_gateway_auth( + &info, + &primary, + json!({"access_token": "gateway-token", "expires_at": i64::MAX}), + ); + let runtime = Arc::new(RuntimeAuthManager::new( + Arc::clone(&primary), + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/register", server.uri()).parse().unwrap(), + )); + let provider = crate::test_support::runtime_model_provider(info, primary, (runtime, mode)); + let auth = provider.api_auth().await.unwrap(); + let stale_headers = HeaderMap::from_iter([ + ( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer stale"), + ), + ( + HeaderName::from_static("x-openai-agent-assertion"), + HeaderValue::from_static("stale"), + ), + ( + HeaderName::from_static("chatgpt-account-id"), + HeaderValue::from_static("stale"), + ), + ( + HeaderName::from_static("x-caller"), + HeaderValue::from_static("retained"), + ), + ]); + let headers = auth + .apply_auth_headers(stale_headers.clone()) + .await + .unwrap(); + assert_eq!(headers["x-gateway-auth"], "Bearer gateway-token"); + assert_eq!(headers["x-caller"], "retained"); + assert!(!headers.contains_key("chatgpt-account-id")); + let proof_header = match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(headers["authorization"], "Bearer primary-token"); + "x-openai-agent-assertion" + } + RuntimeAuthMode::AssertionOnly => { + assert!(!headers.contains_key("x-openai-agent-assertion")); + assert!( + headers["authorization"] + .to_str() + .unwrap() + .starts_with("AgentAssertion 1.") + ); + "authorization" + } + }; + let mut request = codex_http_client::Request::new(http::Method::POST, server.uri()); + request.headers = stale_headers.clone(); + let signed = auth.apply_auth(request).await.unwrap(); + assert_eq!(signed.headers["x-gateway-auth"], "Bearer gateway-token"); + assert_ne!(headers[proof_header], signed.headers[proof_header]); + + let mut collision = stale_headers; + collision.insert("x-gateway-auth", HeaderValue::from_static("caller-token")); + assert!(auth.apply_auth_headers(collision.clone()).await.is_err()); + let mut request = codex_http_client::Request::new(http::Method::POST, server.uri()); + request.headers = collision; + assert!(auth.apply_auth(request).await.is_err()); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + server.verify().await; + } +} + #[tokio::test] async fn gateway_refresh_preserves_primary_auth_and_hides_issuer_errors() { for succeeds in [true, false] { diff --git a/codex-rs/model-provider/src/test_support.rs b/codex-rs/model-provider/src/test_support.rs index 2a3adab398..f61820a114 100644 --- a/codex-rs/model-provider/src/test_support.rs +++ b/codex-rs/model-provider/src/test_support.rs @@ -21,3 +21,5 @@ pub fn models_cache_entry( } pub use crate::shared_state::test_support::seed_gateway_auth; + +pub use crate::provider::test_support::runtime_model_provider; diff --git a/codex-rs/model-provider/tests/runtime_catalog.rs b/codex-rs/model-provider/tests/runtime_catalog.rs new file mode 100644 index 0000000000..d65745cbeb --- /dev/null +++ b/codex-rs/model-provider/tests/runtime_catalog.rs @@ -0,0 +1,305 @@ +//! Exercise catalog auth through the public provider factory and real HTTPS transport. +#![cfg(test)] + +use base64::Engine; +use codex_features::RuntimeAuthMode; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_model_provider::ProviderAuthContext; +use codex_model_provider_info::ModelProviderInfo; +use codex_models_manager::manager::RefreshStrategy; +use codex_models_manager::model_info::model_info_from_slug; +use codex_protocol::openai_models::ModelsResponse; +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 serde_json::Value; +use serde_json::json; +use std::io::Read; +use std::io::Write; +use std::sync::Arc; +use std::time::Duration; +use tokio::net::TcpListener; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const CHILD_SERVER: &str = "CODEX_RUNTIME_CATALOG_TEST_SERVER"; +const TEST_NAME: &str = "production_catalog_uses_runtime_auth_and_preserves_provider_isolation"; + +#[tokio::test] +async fn production_catalog_uses_runtime_auth_and_preserves_provider_isolation() { + if let Ok(server) = std::env::var(CHILD_SERVER) { + check_catalog(&server).await; + return; + } + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-a", "agent_host_id": "host-a", + "actor": {"kind":"platform", "user_id":"user-a", "organization_id":"org-a", "project_id":null}, + }))) + .expect(2).mount(&server).await; + let catalog = ModelsResponse { + models: vec![model_info_from_slug("runtime-catalog-marker")], + }; + for (url, count) in [("/backend-api/codex/models", 4), ("/models", 2)] { + Mock::given(method("GET")) + .and(path(url)) + .respond_with(ResponseTemplate::new(200).set_body_json(&catalog)) + .expect(count) + .mount(&server) + .await; + } + + // Scope proxy and CA configuration to a child process, preserving the real OpenAI + // URL classification without changing environment variables used by other tests. + let (ca, tls) = tls_config(); + let home = tempfile::tempdir().unwrap(); + let ca_path = home.path().join("catalog-ca.pem"); + std::fs::write(&ca_path, ca).unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = *server.address(); + let proxy = tokio::spawn(async move { + for _ in 0..4 { + let (stream, _) = tokio::time::timeout(Duration::from_secs(30), listener.accept()) + .await + .unwrap() + .unwrap(); + let stream = stream.into_std().unwrap(); + let tls = Arc::clone(&tls); + tokio::task::spawn_blocking(move || proxy_catalog(stream, tls, upstream)) + .await + .unwrap(); + } + }); + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .arg("--exact") + .arg(TEST_NAME) + .arg("--nocapture") + .kill_on_drop(/*kill_on_drop*/ true); + for key in [ + "HTTP_PROXY", + "http_proxy", + "HTTPS_PROXY", + "https_proxy", + "ALL_PROXY", + "all_proxy", + "NO_PROXY", + "no_proxy", + "OPENAI_API_KEY", + "OPENAI_PROJECT", + "OPENAI_ORGANIZATION", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + ] { + command.env_remove(key); + } + let output = tokio::time::timeout( + Duration::from_secs(30), + command + .env(CHILD_SERVER, server.uri()) + .env("HTTPS_PROXY", proxy_url) + .env("NO_PROXY", "127.0.0.1,localhost") + .env("SSL_CERT_FILE", ca_path) + .output(), + ) + .await + .unwrap() + .unwrap(); + if !output.status.success() { + proxy.abort(); + } + assert!( + output.status.success(), + "catalog child failed:\n{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + proxy.await.unwrap(); + + let requests = server.received_requests().await.unwrap(); + let catalogs: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/backend-api/codex/models") + .collect(); + assert_eq!(catalogs.len(), 4); + let mut proofs = std::collections::HashSet::new(); + for request in catalogs { + let authorization = request.headers["authorization"].to_str().unwrap(); + let proof = if let Some(proof) = request.headers.get("x-openai-agent-assertion") { + assert_eq!(authorization, "Bearer catalog-key"); + proof.to_str().unwrap() + } else { + authorization.strip_prefix("AgentAssertion ").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-a"); + assert!(proofs.insert(payload["id"].as_str().unwrap().to_owned())); + } + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/backend-api/codex/models" + && request.headers["originator"] == "codex_work_cca") + .count(), + 2 + ); + for request in requests + .iter() + .filter(|request| request.url.path() == "/models") + { + assert_eq!(request.headers["authorization"], "Bearer catalog-key"); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + server.verify().await; +} + +async fn check_catalog(server: &str) { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let credentials = + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("catalog-key")); + let runtime = Arc::new(RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{server}/register").parse().unwrap(), + )); + let context = ProviderAuthContext { + credentials, + runtime_auth: Some((runtime, mode)), + }; + let info = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + let provider = context.model_provider(info); + let selected = provider + .request_auth(/*scope*/ None) + .await + .unwrap() + .runtime + .unwrap(); + assert_eq!(selected.runtime().id(), "runtime-a"); + let models = provider.models_manager_without_cache(/*config_model_catalog*/ None); + models.set_api_key_model_discovery_enabled(/*enabled*/ true); + let work = models.for_originator("codex_work_cca").unwrap(); + for manager in [models, work] { + let result = manager + .raw_model_catalog( + RefreshStrategy::Online, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await; + assert!( + result + .models + .iter() + .any(|model| model.slug == "runtime-catalog-marker"), + "the returned catalog must contain the remote model" + ); + } + let mut other_info = ModelProviderInfo::create_openai_provider(Some(server.to_owned())); + other_info.model_catalog_url = Some(format!("{server}/models").into()); + let other = context.model_provider(other_info); + let models = other.models_manager_without_cache(/*config_model_catalog*/ None); + models.set_api_key_model_discovery_enabled(/*enabled*/ true); + let result = models + .raw_model_catalog( + RefreshStrategy::Online, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await; + assert!( + result + .models + .iter() + .any(|model| model.slug == "runtime-catalog-marker"), + "the other provider must also receive its remote catalog" + ); + } +} + +fn tls_config() -> (String, Arc) { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let mut ca_params = rcgen::CertificateParams::default(); + ca_params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained); + ca_params.key_usages = vec![rcgen::KeyUsagePurpose::KeyCertSign]; + let ca = rcgen::CertifiedIssuer::self_signed(ca_params, rcgen::KeyPair::generate().unwrap()) + .unwrap(); + let mut leaf = rcgen::CertificateParams::new(vec!["chatgpt.com".into()]).unwrap(); + leaf.extended_key_usages = vec![rcgen::ExtendedKeyUsagePurpose::ServerAuth]; + let key = rcgen::KeyPair::generate().unwrap(); + let cert = leaf.signed_by(&key, &ca).unwrap(); + let config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert( + vec![cert.der().clone()], + rustls::pki_types::PrivateKeyDer::from(key), + ) + .unwrap(); + (ca.pem(), Arc::new(config)) +} + +fn proxy_catalog( + mut stream: std::net::TcpStream, + config: Arc, + upstream: std::net::SocketAddr, +) { + stream.set_nonblocking(/*nonblocking*/ false).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(10))) + .unwrap(); + stream + .set_write_timeout(Some(Duration::from_secs(10))) + .unwrap(); + let connect = read_headers(&mut stream); + assert!(connect.starts_with(b"CONNECT chatgpt.com:443 ")); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .unwrap(); + let mut tls = rustls::StreamOwned::new(rustls::ServerConnection::new(config).unwrap(), stream); + let mut headers = read_headers(&mut tls); + assert!(headers.starts_with(b"GET /backend-api/codex/models?")); + headers.truncate(headers.len() - 2); + headers.extend_from_slice(b"connection: close\r\n\r\n"); + let mut origin = std::net::TcpStream::connect(upstream).unwrap(); + origin + .set_read_timeout(Some(Duration::from_secs(10))) + .unwrap(); + origin.write_all(&headers).unwrap(); + let mut response = Vec::new(); + origin.read_to_end(&mut response).unwrap(); + tls.write_all(&response).unwrap(); + tls.flush().unwrap(); +} + +fn read_headers(reader: &mut impl Read) -> Vec { + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + let mut byte = [0_u8]; + reader.read_exact(&mut byte).unwrap(); + headers.push(byte[0]); + assert!(headers.len() < 64 * 1024); + } + headers +} diff --git a/codex-rs/runtime-agent-instance/src/auth.rs b/codex-rs/runtime-agent-instance/src/auth.rs index fbfc663b2e..5cd88e3da9 100644 --- a/codex-rs/runtime-agent-instance/src/auth.rs +++ b/codex-rs/runtime-agent-instance/src/auth.rs @@ -63,6 +63,28 @@ impl RuntimeAuth { self.change_state } + /// Preserve the caller's transport settings and restrictions while binding requests to + /// this credential's account. Independent managed policies must not replace each other. + pub fn bind_http_client_factory( + &self, + factory: HttpClientFactory, + ) -> Result { + let account_policy = self.http_client_factory().network_policy(); + let policy = factory + .network_policy() + .clone() + .with_account_binding(account_policy) + .ok_or_else(|| { + AuthError::Build( + "runtime credentials and transport must share the application network policy" + .to_string(), + ) + })?; + Ok(factory + .with_network_policy(policy) + .with_sensitive_header_protection()) + } + /// Use this selection for a request and its retries; select auth again for the next request. pub fn into_auth_provider(self, mode: RuntimeAuthMode) -> codex_api::SharedAuthProvider { Arc::new(RuntimeRequestAuth { auth: self, mode }) @@ -146,6 +168,7 @@ impl From for AuthError { } RuntimeAuthError::Unavailable | RuntimeAuthError::Unsupported + | RuntimeAuthError::UnsupportedResourceSelection | RuntimeAuthError::Changed | RuntimeAuthError::RefreshFailed | RuntimeAuthError::KeyGeneration diff --git a/codex-rs/runtime-agent-instance/src/bootstrap.rs b/codex-rs/runtime-agent-instance/src/bootstrap.rs index 1e2804e262..0bb6672efb 100644 --- a/codex-rs/runtime-agent-instance/src/bootstrap.rs +++ b/codex-rs/runtime-agent-instance/src/bootstrap.rs @@ -8,9 +8,26 @@ use uuid::Uuid; #[derive(Clone, Debug, Serialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum AuthenticatedActorAuthorization { - Oauth { credential: RedactedString }, - PersonalAccessToken { credential: RedactedString }, - PlatformApiKey { credential: RedactedString }, + Oauth { + credential: RedactedString, + }, + PersonalAccessToken { + credential: RedactedString, + }, + PlatformApiKey { + credential: RedactedString, + #[serde(flatten)] + selection: PlatformResourceSelection, + }, +} + +/// API resource selectors to authenticate, not claims about the resulting actor. +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize)] +pub struct PlatformResourceSelection { + #[serde(skip_serializing_if = "Option::is_none")] + pub organization_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub project_id: Option, } #[derive(Clone, Debug, Serialize)] diff --git a/codex-rs/runtime-agent-instance/src/bootstrap_tests.rs b/codex-rs/runtime-agent-instance/src/bootstrap_tests.rs index 8a4635a7fc..e50011766b 100644 --- a/codex-rs/runtime-agent-instance/src/bootstrap_tests.rs +++ b/codex-rs/runtime-agent-instance/src/bootstrap_tests.rs @@ -22,6 +22,7 @@ fn registration_serializes_the_supported_actor_authorizations() { ), ( AuthenticatedActorAuthorization::PlatformApiKey { + selection: Default::default(), credential: "secret".into(), }, json!({"type": "platform_api_key", "credential": "secret"}), diff --git a/codex-rs/runtime-agent-instance/src/lib.rs b/codex-rs/runtime-agent-instance/src/lib.rs index aaf258b63d..31ff7ea7cc 100644 --- a/codex-rs/runtime-agent-instance/src/lib.rs +++ b/codex-rs/runtime-agent-instance/src/lib.rs @@ -12,10 +12,12 @@ pub use actor::AuthenticatedActor; pub use bootstrap::ActorInfo; pub use bootstrap::AgentHostAuthorization; pub use bootstrap::AuthenticatedActorAuthorization; +pub use bootstrap::PlatformResourceSelection; pub use bootstrap::RuntimeRegistrationMetadata; pub use bootstrap::RuntimeRegistrationRequest; pub use bootstrap::RuntimeRegistrationResponse; pub use bootstrap::RuntimeRole; +pub use codex_http_client::is_openai_api_url; pub use auth::RuntimeAuth; pub use manager::RuntimeAuthError; diff --git a/codex-rs/runtime-agent-instance/src/manager.rs b/codex-rs/runtime-agent-instance/src/manager.rs index 045d2c4317..aa526b679d 100644 --- a/codex-rs/runtime-agent-instance/src/manager.rs +++ b/codex-rs/runtime-agent-instance/src/manager.rs @@ -1,7 +1,8 @@ -//! Share one runtime for the current credential source, including bootstrap retries and refreshes. +//! Reuse one active runtime for the selected credential source and organization. use crate::AgentHostAuthorization; use crate::AuthenticatedActorAuthorization; +use crate::PlatformResourceSelection; use crate::RegistrationError; use crate::RuntimeAgentInstance; use crate::RuntimeAuth; @@ -25,22 +26,19 @@ pub struct RuntimeAuthManager { host: AgentHostAuthorization, metadata: RuntimeRegistrationMetadata, registration_url: Url, - state: Mutex>, + state: Mutex>, +} + +struct CachedRuntime { + source_generation: u64, + selection: PlatformResourceSelection, + state: RuntimeState, } enum RuntimeState { - Pending { - source_generation: u64, - registration: RuntimeRegistration, - }, - Ready { - source_generation: u64, - runtime: Arc, - }, - Denied { - source_generation: u64, - error: RegistrationError, - }, + Pending { registration: RuntimeRegistration }, + Ready { runtime: Arc }, + Denied { error: RegistrationError }, } impl std::fmt::Debug for RuntimeAuthManager { @@ -70,11 +68,20 @@ impl RuntimeAuthManager { } } + pub async fn auth(&self) -> Result { + self.auth_with_selection(&PlatformResourceSelection::default()) + .await + } + + /// 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(&self) -> Result { + pub async fn auth_with_selection( + &self, + selection: &PlatformResourceSelection, + ) -> Result { // This initializer's lock can span IO. No AuthManager lock is held across registration. let mut state = self.state.lock().await; let (auth, http_client_factory) = self @@ -83,7 +90,9 @@ impl RuntimeAuthManager { .await .ok_or(RuntimeAuthError::Unavailable)?; let mut recovery = self.credentials.unauthorized_recovery(); - let mut result = self.auth_once(&mut state, auth, http_client_factory).await; + let mut result = self + .auth_once(&mut state, auth, http_client_factory, selection) + .await; while matches!( &result, Err(RuntimeAuthError::Registration( @@ -105,46 +114,51 @@ impl RuntimeAuthManager { .auth_with_http_client_factory() .await .ok_or(RuntimeAuthError::Unavailable)?; - result = self.auth_once(&mut state, auth, http_client_factory).await; + result = self + .auth_once(&mut state, auth, http_client_factory, selection) + .await; } result } async fn auth_once( &self, - state: &mut Option, + state: &mut Option, auth: CodexAuth, http_client_factory: HttpClientFactory, + selection: &PlatformResourceSelection, ) -> Result { - let actor_auth = actor_authorization(&auth)?; + 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)?; - let same_source = match state.as_ref() { - Some( - RuntimeState::Pending { - source_generation, .. - } - | RuntimeState::Ready { - source_generation, .. - } - | RuntimeState::Denied { - source_generation, .. - }, - ) => *source_generation == change_state.source_generation, - None => false, - }; - if !same_source { - *state = Some(RuntimeState::Pending { - registration: RuntimeRegistration::new( - actor_auth.clone(), - self.host.clone(), - self.metadata.clone(), - ) - .map_err(|_| RuntimeAuthError::KeyGeneration)?, - source_generation: change_state.source_generation, - }); + // 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| { + cached.source_generation != change_state.source_generation + || cached.selection.organization_id != selection.organization_id + || (!matches!(cached.state, RuntimeState::Ready { .. }) + && cached.selection != *selection) + }) { + *state = None; } - let registration = match state.as_mut().ok_or(RuntimeAuthError::Unavailable)? { - RuntimeState::Ready { runtime, .. } => { + let cached = match state { + Some(cached) => cached, + None => state.insert(CachedRuntime { + source_generation: change_state.source_generation, + selection: selection.clone(), + state: RuntimeState::Pending { + registration: RuntimeRegistration::new( + actor_auth.clone(), + self.host.clone(), + self.metadata.clone(), + ) + .map_err(|_| RuntimeAuthError::KeyGeneration)?, + }, + }), + }; + let runtime_state = &mut cached.state; + let registration = match runtime_state { + RuntimeState::Ready { runtime } => { return Ok(RuntimeAuth { runtime: runtime.clone(), auth, @@ -152,8 +166,8 @@ impl RuntimeAuthManager { change_state, }); } - RuntimeState::Denied { error, .. } => return Err(error.clone().into()), - RuntimeState::Pending { registration, .. } => { + RuntimeState::Denied { error } => return Err(error.clone().into()), + RuntimeState::Pending { registration } => { registration.request.actor_auth = actor_auth; registration } @@ -171,10 +185,9 @@ impl RuntimeAuthManager { match result { Ok(info) => { let runtime = Arc::new(registration.finish(info)); - *state = Some(RuntimeState::Ready { + *runtime_state = RuntimeState::Ready { runtime: runtime.clone(), - source_generation: change_state.source_generation, - }); + }; Ok(RuntimeAuth { runtime, auth, @@ -188,10 +201,9 @@ impl RuntimeAuthManager { .. }, ) => { - *state = Some(RuntimeState::Denied { - source_generation: change_state.source_generation, + *runtime_state = RuntimeState::Denied { error: error.clone(), - }); + }; Err(error.into()) } Err(error) => Err(error.into()), @@ -212,13 +224,20 @@ impl RuntimeAuthManager { fn actor_authorization( auth: &CodexAuth, + selection: &PlatformResourceSelection, ) -> Result { + if !matches!(auth, CodexAuth::ApiKey(_)) && selection != &PlatformResourceSelection::default() { + return Err(RuntimeAuthError::UnsupportedResourceSelection); + } let credential = auth .get_token() .map_err(|_| RuntimeAuthError::Unsupported)? .into(); Ok(match auth { - CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey { credential }, + CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey { + credential, + selection: selection.clone(), + }, CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { AuthenticatedActorAuthorization::Oauth { credential } } @@ -238,6 +257,8 @@ pub enum RuntimeAuthError { Unavailable, #[error("this credential cannot register a runtime; use a supported OpenAI sign-in")] Unsupported, + #[error("organization and project selection requires an OpenAI API key")] + UnsupportedResourceSelection, #[error("credentials changed during runtime registration; try again")] Changed, #[error("runtime registration could not refresh actor credentials")] diff --git a/codex-rs/runtime-agent-instance/src/manager_tests.rs b/codex-rs/runtime-agent-instance/src/manager_tests.rs index 64695f52b4..cff21d2189 100644 --- a/codex-rs/runtime-agent-instance/src/manager_tests.rs +++ b/codex-rs/runtime-agent-instance/src/manager_tests.rs @@ -82,7 +82,9 @@ fn registered(request: &Request) -> ResponseTemplate { let actor = if body["actor_auth"]["type"] == "oauth" { json!({"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"}) } else { - json!({"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": "project-a"}) + json!({"kind": "platform", "user_id": "user-a", + "organization_id": body["actor_auth"]["organization_id"].as_str().unwrap_or("org-a"), + "project_id": body["actor_auth"]["project_id"].as_str().unwrap_or("project-a")}) }; ResponseTemplate::new(200).set_body_json(json!({ "runtime_agent_instance_id": request.headers["idempotency-key"].to_str().unwrap(), @@ -161,6 +163,85 @@ async fn concurrent_requests_share_registration_and_replacement_preserves_in_fli ); } +#[tokio::test] +async fn projects_share_one_runtime_and_organization_changes_replace_it() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .expect(3) + .mount(&server) + .await; + let initial = CodexAuth::from_api_key("key-a"); + let (manager, _, _) = setup(&server, initial.clone(), initial).await; + let mut selection = PlatformResourceSelection { + organization_id: Some("org-a".into()), + project_id: Some("project-a".into()), + }; + let first = manager.auth_with_selection(&selection).await.unwrap(); + for project in [Some("project-b"), None, Some("project-a")] { + selection.project_id = project.map(str::to_owned); + let next = manager.auth_with_selection(&selection).await.unwrap(); + assert!(Arc::ptr_eq(first.runtime(), next.runtime())); + } + selection.organization_id = Some("org-b".into()); + let other_org = manager.auth_with_selection(&selection).await.unwrap(); + assert_ne!(first.runtime().id(), other_org.runtime().id()); + selection.organization_id = Some("org-a".into()); + let replacement = manager.auth_with_selection(&selection).await.unwrap(); + assert_ne!(first.runtime().id(), replacement.runtime().id()); + assert_ne!(other_org.runtime().id(), replacement.runtime().id()); + server.verify().await; +} + +#[tokio::test] +async fn failed_registration_does_not_poison_another_project_or_reuse_its_idempotency_key() { + for status in [403, 503] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(body_partial_json( + json!({"actor_auth": {"project_id": "denied"}}), + )) + .respond_with(ResponseTemplate::new(status)) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .respond_with(registered) + .expect(1) + .mount(&server) + .await; + let initial = CodexAuth::from_api_key("key-a"); + let (manager, _, _) = setup(&server, initial.clone(), initial).await; + let mut selection = PlatformResourceSelection { + organization_id: Some("org-a".into()), + project_id: Some("denied".into()), + }; + for _ in 0..2 { + assert!(manager.auth_with_selection(&selection).await.is_err()); + } + let failed = server.received_requests().await.unwrap(); + assert_eq!(failed.len(), if status == 403 { 1 } else { 2 }); + if status == 503 { + assert_eq!(failed[0].body, failed[1].body); + assert_eq!( + failed[0].headers["idempotency-key"], + failed[1].headers["idempotency-key"] + ); + } + + selection.project_id = Some("allowed".into()); + let ready = manager.auth_with_selection(&selection).await.unwrap(); + assert_ne!( + ready.runtime().id(), + failed[0].headers["idempotency-key"].to_str().unwrap() + ); + selection.project_id = Some("denied".into()); + let reused = manager.auth_with_selection(&selection).await.unwrap(); + assert!(Arc::ptr_eq(ready.runtime(), reused.runtime())); + server.verify().await; + } +} + #[tokio::test] async fn policy_recovery_preserves_pending_registration() { let server = MockServer::start().await; diff --git a/codex-rs/runtime-agent-instance/src/registration_tests.rs b/codex-rs/runtime-agent-instance/src/registration_tests.rs index ed76e74c9e..a3a6af1cac 100644 --- a/codex-rs/runtime-agent-instance/src/registration_tests.rs +++ b/codex-rs/runtime-agent-instance/src/registration_tests.rs @@ -27,6 +27,7 @@ use wiremock::matchers::path; fn pending() -> RuntimeRegistration { RuntimeRegistration::new( AuthenticatedActorAuthorization::PlatformApiKey { + selection: Default::default(), credential: "actor-secret".into(), }, AgentHostAuthorization::InstallationId {