Enforce managed residency for model providers (#39645)

## Why

Model provider configuration could override the residency header required by
`enforce_residency`, so provider-backed requests did not reliably honor the
managed setting.

## What changed

- Apply the managed residency header after building a provider, making it
  authoritative for model requests and model discovery.
- Warn when a provider configures the residency header through
  `http_headers` or `env_http_headers`, while preserving the original provider
  configuration and unrelated headers.
- Cover HTTP and WebSocket requests, model discovery, and case-insensitive
  header detection with targeted tests.

GitOrigin-RevId: 461ef8989d5bd06ffdea694aee4265e392e07590
This commit is contained in:
Anton Panasenko authored and copyberry committed 2026-08-20 07:33:27 +00:00
1 parent 663da53823
commit af0e82c562
7 files changed
+371 -3

No files matched your search

+54 -1
View File
@@ -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::<ReqwestTransport>::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")
);
}
}