diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 5bec9aaa12..f34a78e95f 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -75,6 +75,7 @@ mod remote_thread_store; mod request_permissions; mod request_user_input; mod request_validation; +mod residency; mod review; mod rollout_migration; mod safety_check_downgrade; diff --git a/codex-rs/app-server/tests/suite/v2/residency.rs b/codex-rs/app-server/tests/suite/v2/residency.rs new file mode 100644 index 0000000000..74cb932a67 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/residency.rs @@ -0,0 +1,189 @@ +use anyhow::Context; +use anyhow::Result; +use app_test_support::TestAppServer; +use app_test_support::write_models_cache; +use codex_app_server_protocol::ConfigWarningNotification; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::UserInput; +use codex_login::default_client::RESIDENCY_HEADER_NAME; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use serde_json::Value; +use std::path::Path; +use std::time::Duration; +use tempfile::TempDir; +use test_case::test_case; +use tokio::time::timeout; + +const READ_TIMEOUT: Duration = Duration::from_secs(60); +const PROVIDER_RESIDENCY_ENV_VAR: &str = "CODEX_TEST_RESIDENCY_HEADER"; + +#[derive(Clone, Copy)] +enum ModelTransport { + Http, + Websocket, +} + +fn write_provider_config( + codex_home: &Path, + model_base_url: &str, + transport: ModelTransport, +) -> Result<()> { + let supports_websockets = matches!(transport, ModelTransport::Websocket); + let config = format!( + r#"model = "gpt-5.4" +approval_policy = "never" +sandbox_mode = "danger-full-access" +model_provider = "custom-openai" + +[features] +shell_snapshot = false + +[model_providers.custom-openai] +name = "OpenAI" +base_url = "{model_base_url}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +requires_openai_auth = true +supports_websockets = {supports_websockets} + +[model_providers.custom-openai.http_headers] +"X-OpenAI-Internal-Codex-Residency" = "eu-static" +"x-provider-header" = "preserved" + +[model_providers.custom-openai.env_http_headers] +"x-openai-internal-codex-residency" = "{PROVIDER_RESIDENCY_ENV_VAR}" +"# + ); + std::fs::write(codex_home.join("config.toml"), config)?; + write_models_cache(codex_home)?; + Ok(()) +} + +#[test_case(ModelTransport::Http; "http")] +#[test_case(ModelTransport::Websocket; "websocket")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn managed_residency_overrides_provider_headers(transport: ModelTransport) -> Result<()> { + skip_if_no_network!(Ok(())); + + let http_server = responses::start_mock_server().await; + let http_responses = responses::mount_sse_once( + &http_server, + responses::sse(vec![ + responses::ev_response_created("http-response"), + responses::ev_assistant_message("http-message", "done"), + responses::ev_completed("http-response"), + ]), + ) + .await; + let websocket_server = responses::start_websocket_server(vec![vec![ + vec![ + responses::ev_response_created("prewarm-response"), + responses::ev_completed("prewarm-response"), + ], + vec![ + responses::ev_response_created("websocket-response"), + responses::ev_assistant_message("websocket-message", "done"), + responses::ev_completed("websocket-response"), + ], + ]]) + .await; + let model_base_url = match transport { + ModelTransport::Http => http_server.uri(), + ModelTransport::Websocket => websocket_server.uri().to_string(), + }; + + let codex_home = TempDir::new()?; + write_provider_config(codex_home.path(), &model_base_url, transport)?; + std::fs::write( + codex_home.path().join("requirements.toml"), + "enforce_residency = \"us\"\n", + )?; + + let mut app_server = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_env_overrides(&[ + ("OPENAI_API_KEY", Some("sk-test")), + ("CODEX_ACCESS_TOKEN", None), + (PROVIDER_RESIDENCY_ENV_VAR, Some("eu-environment")), + ]) + .build_initialized_with_timeout(READ_TIMEOUT) + .await?; + + let expected_summary = format!( + "Ignoring `{RESIDENCY_HEADER_NAME}` in `model_providers.custom-openai` because managed residency is required." + ); + let notification = timeout( + READ_TIMEOUT, + app_server.read_stream_until_matching_notification( + "managed residency configWarning", + |notification| { + notification.method == "configWarning" + && notification + .params + .as_ref() + .and_then(|params| params.get("summary")) + .and_then(Value::as_str) + == Some(expected_summary.as_str()) + }, + ), + ) + .await??; + let warning: ConfigWarningNotification = serde_json::from_value( + notification + .params + .context("managed residency warning should have parameters")?, + )?; + assert_eq!( + warning, + ConfigWarningNotification { + summary: expected_summary, + details: None, + path: None, + range: None, + } + ); + + let thread = app_server + .start_thread(ThreadStartParams::default()) + .await?; + let completed = timeout( + READ_TIMEOUT, + app_server.start_turn_and_wait_for_completion(TurnStartParams { + thread_id: thread.thread.id, + input: vec![UserInput::Text { + text: "hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }), + ) + .await??; + assert_eq!(completed.turn.status, TurnStatus::Completed); + + let (residency_header, provider_header) = match transport { + ModelTransport::Http => { + let request = http_responses.single_request(); + ( + request.header(RESIDENCY_HEADER_NAME), + request.header("x-provider-header"), + ) + } + ModelTransport::Websocket => { + let handshake = websocket_server.single_handshake(); + ( + handshake.header(RESIDENCY_HEADER_NAME), + handshake.header("x-provider-header"), + ) + } + }; + assert_eq!(residency_header, Some("us".to_string())); + assert_eq!(provider_header, Some("preserved".to_string())); + + websocket_server.shutdown().await; + Ok(()) +} diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 493b9ee48d..22e5724704 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -73,6 +73,7 @@ use codex_config::types::WindowsToml; use codex_exec_server::LOCAL_FS; use codex_features::Feature; use codex_features::FeaturesToml; +use codex_login::default_client::RESIDENCY_HEADER_NAME; use codex_login::test_support::auth_manager_from_optional_auth; use codex_model_provider::ProviderCapabilities; use codex_model_provider_info::LMSTUDIO_OSS_PROVIDER_ID; @@ -947,6 +948,88 @@ profile = "codex-bedrock" ); } +#[tokio::test] +async fn managed_residency_warns_about_provider_header_overrides_without_mutating_config() +-> std::io::Result<()> { + for enforce_residency in [None, Some(ResidencyRequirement::Us)] { + let cfg = toml::from_str::( + r#" +model_provider = "custom-openai" + +[model_providers.custom-openai] +name = "OpenAI" +http_headers = { "X-OpenAI-Internal-Codex-Residency" = "request-override", "x-provider-header" = "preserved" } +env_http_headers = { "x-openai-internal-codex-residency" = "CODEX_TEST_UNSET_RESIDENCY_HEADER", "x-provider-env-header" = "CODEX_TEST_UNSET_PROVIDER_HEADER" } +"#, + ) + .expect("existing provider configuration should remain loadable"); + let requirements = ConfigRequirements { + enforce_residency: ConstrainedWithSource::new( + Constrained::allow_only(enforce_residency), + enforce_residency.map(|_| RequirementSource::Unknown), + ), + ..Default::default() + }; + let requirements_toml = ConfigRequirementsToml { + enforce_residency, + ..Default::default() + }; + let config_layer_stack = + ConfigLayerStack::new(Vec::new(), requirements, requirements_toml)?; + let config = Config::load_config_with_layer_stack( + LOCAL_FS.as_ref(), + cfg, + ConfigOverrides::default(), + tempdir()?.abs(), + config_layer_stack, + ) + .await?; + + let static_headers = config + .model_provider + .http_headers + .as_ref() + .expect("static headers should remain configured"); + let environment_headers = config + .model_provider + .env_http_headers + .as_ref() + .expect("environment-backed headers should remain configured"); + assert_eq!( + static_headers + .get("X-OpenAI-Internal-Codex-Residency") + .map(String::as_str), + Some("request-override") + ); + assert_eq!( + environment_headers + .get(RESIDENCY_HEADER_NAME) + .map(String::as_str), + Some("CODEX_TEST_UNSET_RESIDENCY_HEADER") + ); + assert_eq!( + static_headers.get("x-provider-header").map(String::as_str), + Some("preserved") + ); + assert_eq!( + environment_headers + .get("x-provider-env-header") + .map(String::as_str), + Some("CODEX_TEST_UNSET_PROVIDER_HEADER") + ); + + let expected_warning = format!( + "Ignoring `{RESIDENCY_HEADER_NAME}` in `model_providers.custom-openai` because managed residency is required." + ); + assert_eq!( + config.startup_warnings.contains(&expected_warning), + enforce_residency.is_some() + ); + } + + Ok(()) +} + #[test] fn accepts_amazon_bedrock_aws_profile_override() { let cfg = toml::from_str::( diff --git a/codex-rs/core/src/config/requirements.rs b/codex-rs/core/src/config/requirements.rs index 512ec50d36..1ac2c26ef0 100644 --- a/codex-rs/core/src/config/requirements.rs +++ b/codex-rs/core/src/config/requirements.rs @@ -5,6 +5,7 @@ use codex_config::config_toml::ConfigToml; use codex_config::types::ApprovalsReviewer; use codex_config::types::FeedbackConfigToml; use codex_features::FeatureToml; +use codex_login::default_client::RESIDENCY_HEADER_NAME; use codex_utils_absolute_path::AbsolutePathBuf; use std::path::Path; @@ -47,6 +48,26 @@ pub(super) fn apply_to_config( requirements.feedback.as_ref(), startup_warnings, ); + if requirements.enforce_residency.value().is_some() { + for (provider_name, provider) in &config.model_providers { + let has_residency_header = [&provider.http_headers, &provider.env_http_headers] + .into_iter() + .flatten() + .any(|headers| { + headers + .keys() + .any(|name| name.eq_ignore_ascii_case(RESIDENCY_HEADER_NAME)) + }); + + if has_residency_header { + let warning = format!( + "Ignoring `{RESIDENCY_HEADER_NAME}` in `model_providers.{provider_name}` because managed residency is required." + ); + tracing::warn!(provider = provider_name, "{warning}"); + startup_warnings.push(warning); + } + } + } if let Some(requirement) = requirements.windows_sandbox_private_desktop.as_ref() { apply_exact_requirement( "windows.sandbox_private_desktop", diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index 985cdb0b43..cd4cbd5431 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -104,6 +104,11 @@ pub fn set_default_client_residency_requirement(enforce_residency: Option Option { + REQUIREMENTS_RESIDENCY.read().ok().and_then(|guard| *guard) +} + pub fn originator() -> Originator { if let Ok(guard) = ORIGINATOR.read() && let Some(originator) = guard.as_ref() diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 49ed6979b7..e0a01ee5f4 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -34,6 +34,7 @@ use tokio::time::timeout; use crate::auth::agent_identity_telemetry; use crate::auth::resolve_provider_auth; +use crate::provider::enforce_managed_residency; const MODELS_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); const MODELS_ENDPOINT: &str = "/models"; @@ -81,7 +82,8 @@ impl OpenAiModelsEndpoint { codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); let auth = self.auth().await; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); - let api_provider = self.provider_info.to_api_provider(auth_mode)?; + let mut api_provider = self.provider_info.to_api_provider(auth_mode)?; + enforce_managed_residency(&mut api_provider); let api_auth = resolve_provider_auth(auth.as_ref(), &self.provider_info)?; let request_url = ModelsClient::::request_url(&api_provider, client_version); @@ -280,13 +282,17 @@ mod tests { use super::*; use codex_http_client::OutboundProxyPolicy; + use codex_login::default_client::RESIDENCY_HEADER_NAME; + use codex_login::default_client::ResidencyRequirement; use codex_login::default_client::create_client; + use codex_login::default_client::set_default_client_residency_requirement; use codex_protocol::config_types::ModelProviderAuthInfo; use codex_protocol::openai_models::ModelsResponse; use pretty_assertions::assert_eq; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; + use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; use wiremock::matchers::query_param; @@ -390,4 +396,51 @@ mod tests { )) ); } + + #[tokio::test] + async fn model_discovery_enforces_managed_residency_over_provider_headers() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/models")) + .and(header(RESIDENCY_HEADER_NAME, "us")) + .respond_with( + ResponseTemplate::new(200).set_body_json(ModelsResponse { models: Vec::new() }), + ) + .expect(1) + .mount(&server) + .await; + + let mut provider_info = ModelProviderInfo::create_openai_provider(Some(server.uri())); + provider_info.http_headers = Some(std::collections::HashMap::from([( + RESIDENCY_HEADER_NAME.to_string(), + "eu".to_string(), + )])); + let endpoint = OpenAiModelsEndpoint { + provider_info, + auth_manager: None, + transport_builder: Arc::new(RecordingTransportBuilder { + observed_request: Arc::new(Mutex::new(None)), + }), + }; + + set_default_client_residency_requirement(Some(ResidencyRequirement::Us)); + endpoint + .list_models( + "0.0.0", + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await + .expect("managed residency model discovery should succeed"); + set_default_client_residency_requirement(/*enforce_residency*/ None); + + assert_eq!( + endpoint + .provider_info + .http_headers + .as_ref() + .and_then(|headers| headers.get(RESIDENCY_HEADER_NAME)) + .map(String::as_str), + Some("eu") + ); + } } diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 65f4204d56..192c0eb0c1 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -11,6 +11,9 @@ use codex_api::TransportError; use codex_api::is_azure_responses_provider; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::default_client::RESIDENCY_HEADER_NAME; +use codex_login::default_client::ResidencyRequirement; +use codex_login::default_client::read_default_client_residency_requirement; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::cache::ModelsCache; use codex_models_manager::manager::OpenAiModelsManager; @@ -19,6 +22,7 @@ use codex_models_manager::manager::StaticModelsManager; use codex_protocol::account::ProviderAccount; use codex_protocol::error::CodexErr; use codex_protocol::openai_models::ModelsResponse; +use http::HeaderValue; use crate::amazon_bedrock::AmazonBedrockModelProvider; use crate::auth::ProviderAuthScope; @@ -28,6 +32,15 @@ use crate::auth::resolve_provider_auth; use crate::auth::resolve_provider_auth_for_scope; use crate::models_endpoint::OpenAiModelsEndpoint; +pub(crate) fn enforce_managed_residency(provider: &mut Provider) { + if let Some(requirement) = read_default_client_residency_requirement() { + let value = match requirement { + ResidencyRequirement::Us => HeaderValue::from_static("us"), + }; + provider.headers.insert(RESIDENCY_HEADER_NAME, value); + } +} + /// Remote context-compaction protocols supported by a model provider. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RemoteCompactionSupport { @@ -203,8 +216,11 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { fn api_provider(&self) -> ModelProviderFuture<'_, codex_protocol::error::Result> { Box::pin(async move { let auth = self.auth().await; - self.info() - .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode)) + let mut provider = self + .info() + .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; + enforce_managed_residency(&mut provider); + Ok(provider) }) }