diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index a845085afc..5b1fcf9330 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -57,6 +57,7 @@ use codex_app_server_protocol::TurnStartedNotification; use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::TurnSteerResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_login::default_client::RESIDENCY_HEADER_NAME; use codex_protocol::protocol::CodexResponseHandoffMode; use codex_protocol::protocol::ConversationTextRole; use codex_protocol::protocol::RealtimeConversationVersion; @@ -84,6 +85,7 @@ use std::sync::mpsc; use std::time::Duration; use tempfile::TempDir; use test_case::test_case; +use test_case::test_matrix; use tokio::time::timeout; use uuid::Uuid; use wiremock::Match; @@ -1895,6 +1897,156 @@ async fn realtime_mode_uses_client_instructions_on_entry_and_exit() -> Result<() Ok(()) } +const PROVIDER_RESIDENCY_ENV_VAR: &str = "CODEX_TEST_REALTIME_RESIDENCY_HEADER"; + +#[derive(Clone, Copy, Debug)] +enum RealtimeTransport { + Websocket, + Webrtc, +} + +#[derive(Clone, Copy, Debug)] +enum ResidencyPolicy { + Required, + Unmanaged, +} + +#[derive(Clone, Copy, Debug)] +enum ProviderHeaders { + Static, + StaticAndEnvironment, +} + +#[test_matrix( + [RealtimeTransport::Websocket, RealtimeTransport::Webrtc], + [ResidencyPolicy::Required, ResidencyPolicy::Unmanaged], + [ProviderHeaders::Static, ProviderHeaders::StaticAndEnvironment] +)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn realtime_respects_managed_residency( + transport: RealtimeTransport, + residency: ResidencyPolicy, + provider_headers: ProviderHeaders, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let responses_server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let call_capture = RealtimeCallRequestCapture::new(); + Mock::given(method("POST")) + .and(path("/v1/realtime/calls")) + .and(call_capture.clone()) + .respond_with( + ResponseTemplate::new(200) + .insert_header("Location", "/v1/realtime/calls/rtc_residency") + .set_body_string("v=answer\r\n"), + ) + .mount(&responses_server) + .await; + let realtime_server = start_websocket_server_with_headers(vec![ + open_realtime_sideband_connection(vec![vec![session_updated("sess_residency")]]), + ]) + .await; + let responses_uri = responses_server.uri(); + let realtime_uri = realtime_server.uri(); + let codex_home = TempDir::new()?; + let mut config = MockResponsesConfig::new(&responses_uri) + .with_root_config(&format!( + "experimental_realtime_ws_base_url = \"{realtime_uri}\"\n\ + experimental_realtime_webrtc_call_base_url = \"{responses_uri}/v1\"\n\ + experimental_realtime_ws_startup_context = \"startup context\"" + )) + .with_extra_config("[realtime]\nversion = \"v1\"\ntype = \"conversational\"") + .with_extra_config( + "[model_providers.mock_provider.http_headers]\n\ + \"X-OpenAI-Internal-Codex-Residency\" = \"eu-static\"\n\ + \"x-provider-header\" = \"preserved\"", + ); + let configured_residency = match provider_headers { + ProviderHeaders::Static => "eu-static", + ProviderHeaders::StaticAndEnvironment => { + config = config.with_extra_config(&format!( + "[model_providers.mock_provider.env_http_headers]\n\ + \"x-openai-internal-codex-residency\" = \"{PROVIDER_RESIDENCY_ENV_VAR}\"" + )); + "eu-environment" + } + }; + config.write(codex_home.path())?; + let expected_residency = match residency { + ResidencyPolicy::Required => { + std::fs::write( + codex_home.path().join("requirements.toml"), + "enforce_residency = \"us\"\n", + )?; + "us" + } + ResidencyPolicy::Unmanaged => configured_residency, + }; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_env_overrides(&[(PROVIDER_RESIDENCY_ENV_VAR, Some("eu-environment"))]) + .build() + .await?; + mcp.initialize().await?; + login_with_api_key(&mut mcp, "sk-test-key").await?; + let thread_request = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams::default()) + .await?; + let thread: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_request)).await??; + let mut harness = RealtimeE2eHarness { + mcp, + _codex_home: codex_home, + main_loop_responses_server: responses_server, + realtime_server, + call_capture, + thread_id: thread.thread.id, + }; + + match transport { + RealtimeTransport::Websocket => { + harness.start_websocket_realtime().await?; + } + RealtimeTransport::Webrtc => { + harness.start_webrtc_realtime("v=offer\r\n").await?; + let request = harness.call_capture.single_request(); + assert_eq!( + request + .headers + .get_all(RESIDENCY_HEADER_NAME) + .iter() + .map(|value| value.to_str().expect("residency header should be text")) + .collect::>(), + vec![expected_residency] + ); + assert_eq!( + request + .headers + .get("x-provider-header") + .expect("unrelated provider header should be preserved"), + "preserved" + ); + } + } + // The WebRTC started event precedes the asynchronous sideband connection. + harness.sideband_outbound_request(/*request_index*/ 0).await; + let handshake = harness.realtime_server.single_handshake(); + assert_eq!( + ( + handshake.header(RESIDENCY_HEADER_NAME), + handshake.header("x-provider-header"), + ), + ( + Some(expected_residency.to_string()), + Some("preserved".to_string()), + ) + ); + + harness.shutdown().await; + Ok(()) +} + #[tokio::test] async fn realtime_webrtc_start_emits_sdp_notification() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/config/src/config_requirements.rs b/codex-rs/config/src/config_requirements.rs index 821bbaef2a..978175dbc9 100644 --- a/codex-rs/config/src/config_requirements.rs +++ b/codex-rs/config/src/config_requirements.rs @@ -1,6 +1,7 @@ use crate::ApplicationRequirementsToml; use codex_features::FeatureToml; use codex_model_provider_info::ModelProviderInfo; +pub use codex_model_provider_info::ResidencyRequirement; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::SandboxMode; @@ -1463,12 +1464,6 @@ impl From for SandboxModeRequirement { } } -#[derive(Deserialize, Debug, Clone, Copy, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum ResidencyRequirement { - Us, -} - impl ConfigRequirementsToml { pub fn apply_remote_sandbox_config(&mut self, hostname: Option<&str>) { let Some(remote_sandbox_config) = self.remote_sandbox_config.as_ref() else { diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index 0d6a0fcc99..d26a82fbc7 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -39,9 +39,10 @@ use crate::outbound_proxy::AuthRouteConfig; pub static USER_AGENT_SUFFIX: LazyLock>> = LazyLock::new(|| Mutex::new(None)); pub const DEFAULT_ORIGINATOR: &str = "codex_cli_rs"; pub const CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR: &str = "CODEX_INTERNAL_ORIGINATOR_OVERRIDE"; -pub const RESIDENCY_HEADER_NAME: &str = "x-openai-internal-codex-residency"; - -pub use codex_config::ResidencyRequirement; +pub use codex_model_provider_info::RESIDENCY_HEADER_NAME; +pub use codex_model_provider_info::ResidencyRequirement; +pub use codex_model_provider_info::read_managed_residency_requirement as read_default_client_residency_requirement; +pub use codex_model_provider_info::set_managed_residency_requirement as set_default_client_residency_requirement; #[derive(Debug, Clone)] pub struct Originator { @@ -49,8 +50,6 @@ pub struct Originator { pub header_value: HeaderValue, } static ORIGINATOR: LazyLock>> = LazyLock::new(|| RwLock::new(None)); -static REQUIREMENTS_RESIDENCY: LazyLock>> = - LazyLock::new(|| RwLock::new(None)); static ROUTE_AWARE_CLIENT_BUILD_PERMIT: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1); @@ -96,19 +95,6 @@ pub fn set_default_originator(value: String) -> Result<(), SetOriginatorError> { Ok(()) } -pub fn set_default_client_residency_requirement(enforce_residency: Option) { - let Ok(mut guard) = REQUIREMENTS_RESIDENCY.write() else { - tracing::warn!("Failed to acquire requirements residency lock"); - return; - }; - *guard = enforce_residency; -} - -/// Returns the current process-wide residency requirement. -pub fn read_default_client_residency_requirement() -> 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() @@ -352,10 +338,7 @@ pub fn default_headers() -> HeaderMap { if let Ok(user_agent) = HeaderValue::from_str(&get_codex_user_agent()) { headers.insert(USER_AGENT, user_agent); } - if let Ok(guard) = REQUIREMENTS_RESIDENCY.read() - && let Some(requirement) = guard.as_ref() - && !headers.contains_key(RESIDENCY_HEADER_NAME) - { + if let Some(requirement) = read_default_client_residency_requirement() { let value = match requirement { ResidencyRequirement::Us => HeaderValue::from_static("us"), }; diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 9c27b15cc4..7fa9ab956e 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -4,6 +4,9 @@ //! 1. Built-in defaults compiled into the binary so Codex works out-of-the-box. //! 2. User-defined entries inside `~/.codex/config.toml` under the `model_providers` //! key. These override or extend the defaults at runtime. +//! +//! API provider construction applies the process-wide managed residency policy also +//! used by default HTTP headers. use codex_api::Provider as ApiProvider; use codex_api::RetryConfig as ApiRetryConfig; @@ -24,8 +27,35 @@ use std::fmt; use std::num::NonZeroU64; use std::path::Component; use std::path::Path; +use std::sync::PoisonError; +use std::sync::RwLock; use std::time::Duration; +pub const RESIDENCY_HEADER_NAME: &str = "x-openai-internal-codex-residency"; + +#[derive(Deserialize, Debug, Clone, Copy, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum ResidencyRequirement { + Us, +} + +static REQUIREMENTS_RESIDENCY: RwLock> = RwLock::new(None); + +/// Sets the process-wide residency requirement loaded from managed configuration. +pub fn set_managed_residency_requirement(enforce_residency: Option) { + // Recover the stored policy if the lock is poisoned rather than silently disabling it. + *REQUIREMENTS_RESIDENCY + .write() + .unwrap_or_else(PoisonError::into_inner) = enforce_residency; +} + +/// Returns the current process-wide managed residency requirement. +pub fn read_managed_residency_requirement() -> Option { + *REQUIREMENTS_RESIDENCY + .read() + .unwrap_or_else(PoisonError::into_inner) +} + const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; @@ -367,6 +397,7 @@ other non-default provider fields are not supported" Ok(headers) } + /// Builds an API provider with managed residency taking precedence over configured headers. pub fn to_api_provider(&self, auth_mode: Option) -> CodexResult { let default_base_url = if matches!( auth_mode, @@ -387,7 +418,13 @@ other non-default provider fields are not supported" .clone() .unwrap_or_else(|| default_base_url.to_string()); - let headers = self.build_header_map()?; + let mut headers = self.build_header_map()?; + if let Some(requirement) = read_managed_residency_requirement() { + let value = match requirement { + ResidencyRequirement::Us => HeaderValue::from_static("us"), + }; + headers.insert(RESIDENCY_HEADER_NAME, value); + } let retry = ApiRetryConfig { max_attempts: self.request_max_retries(), base_delay: Duration::from_millis(200), diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs index 8e13869cb2..641942703d 100644 --- a/codex-rs/model-provider-info/src/model_provider_info_tests.rs +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -5,6 +5,54 @@ use pretty_assertions::assert_eq; use std::num::NonZeroU64; use tempfile::tempdir; +#[test] +fn test_api_provider_applies_current_managed_residency() { + let info = ModelProviderInfo { + http_headers: Some(maplit::hashmap! { + "X-OpenAI-Internal-Codex-Residency".to_string() => "eu".into(), + "x-provider-header".to_string() => "preserved".into(), + }), + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + let original_info = info.clone(); + let previous_requirement = read_managed_residency_requirement(); + set_managed_residency_requirement(Some(ResidencyRequirement::Us)); + let managed = info.to_api_provider(/*auth_mode*/ None); + set_managed_residency_requirement(/*enforce_residency*/ None); + let unmanaged = info.to_api_provider(/*auth_mode*/ None); + set_managed_residency_requirement(previous_requirement); + + assert_eq!( + managed.expect("managed provider should resolve").headers, + HeaderMap::from_iter([ + ( + HeaderName::from_static(RESIDENCY_HEADER_NAME), + HeaderValue::from_static("us") + ), + ( + HeaderName::from_static("x-provider-header"), + HeaderValue::from_static("preserved") + ), + ]) + ); + assert_eq!( + unmanaged + .expect("unmanaged provider should resolve") + .headers, + HeaderMap::from_iter([ + ( + HeaderName::from_static(RESIDENCY_HEADER_NAME), + HeaderValue::from_static("eu") + ), + ( + HeaderName::from_static("x-provider-header"), + HeaderValue::from_static("preserved") + ), + ]) + ); + assert_eq!(info, original_info); +} + #[test] fn test_deserialize_ollama_model_provider_toml() { let azure_provider_toml = r#" diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 5b7b7997fe..f5e5838d84 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -35,7 +35,6 @@ 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"; @@ -92,7 +91,6 @@ impl OpenAiModelsEndpoint { // Codex metadata is served by the Codex backend, not the public /v1/models API. api_provider.base_url = CHATGPT_CODEX_BASE_URL.to_string(); } - 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); diff --git a/codex-rs/model-provider/src/models_identity.rs b/codex-rs/model-provider/src/models_identity.rs index 2baa42f2ce..883835b67d 100644 --- a/codex-rs/model-provider/src/models_identity.rs +++ b/codex-rs/model-provider/src/models_identity.rs @@ -9,14 +9,12 @@ use sha2::Digest; use sha2::Sha256; use crate::auth::resolve_provider_auth; -use crate::provider::enforce_managed_residency; pub(crate) fn identity( provider_info: &ModelProviderInfo, auth: Option<&CodexAuth>, ) -> CoreResult { let mut provider = provider_info.to_api_provider(auth.map(CodexAuth::auth_mode))?; - enforce_managed_residency(&mut provider); let mut digest = Sha256::new(); // Length-prefix every field to avoid ambiguity between adjacent values. let mut field = |value: &[u8]| { diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 58b51d1bb3..edeb748dc9 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -13,9 +13,6 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::WorkspaceRoutingRequest; use codex_login::default_client::ClientRedirectPolicy; -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; @@ -24,7 +21,6 @@ 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::ResolvedResponsesProvider; use crate::amazon_bedrock::AmazonBedrockModelProvider; @@ -36,15 +32,6 @@ use crate::auth::resolve_provider_auth_for_scope; use crate::models_endpoint::OpenAiModelsEndpoint; use crate::workspace_routing::WorkspaceRoutingContext; -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 { @@ -230,11 +217,8 @@ 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; - let mut provider = self - .info() - .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; - enforce_managed_residency(&mut provider); - Ok(provider) + self.info() + .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode)) }) }