mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Add enforce_residency to requirements (#10263)
Add `enforce_residency` to requirements.toml and thread it through to a header on `default_client`.
This commit is contained in:
@@ -24,6 +24,7 @@ use crate::config_loader::ConfigRequirements;
|
||||
use crate::config_loader::LoaderOverrides;
|
||||
use crate::config_loader::McpServerIdentity;
|
||||
use crate::config_loader::McpServerRequirement;
|
||||
use crate::config_loader::ResidencyRequirement;
|
||||
use crate::config_loader::Sourced;
|
||||
use crate::config_loader::load_config_layers_state;
|
||||
use crate::features::Feature;
|
||||
@@ -140,6 +141,11 @@ pub struct Config {
|
||||
|
||||
pub sandbox_policy: Constrained<SandboxPolicy>,
|
||||
|
||||
/// enforce_residency means web traffic cannot be routed outside of a
|
||||
/// particular geography. HTTP clients should direct their requests
|
||||
/// using backend-specific headers or URLs to enforce this.
|
||||
pub enforce_residency: Constrained<Option<ResidencyRequirement>>,
|
||||
|
||||
/// True if the user passed in an override or set a value in config.toml
|
||||
/// for either of approval_policy or sandbox_mode.
|
||||
pub did_user_set_custom_approval_policy_or_sandbox_mode: bool,
|
||||
@@ -1516,6 +1522,7 @@ impl Config {
|
||||
sandbox_policy: mut constrained_sandbox_policy,
|
||||
mcp_servers,
|
||||
exec_policy: _,
|
||||
enforce_residency,
|
||||
} = requirements;
|
||||
|
||||
constrained_approval_policy
|
||||
@@ -1538,6 +1545,7 @@ impl Config {
|
||||
cwd: resolved_cwd,
|
||||
approval_policy: constrained_approval_policy,
|
||||
sandbox_policy: constrained_sandbox_policy,
|
||||
enforce_residency,
|
||||
did_user_set_custom_approval_policy_or_sandbox_mode,
|
||||
forced_auto_mode_downgraded_on_windows,
|
||||
shell_environment_policy,
|
||||
@@ -3771,6 +3779,7 @@ model_verbosity = "high"
|
||||
model_provider: fixture.openai_provider.clone(),
|
||||
approval_policy: Constrained::allow_any(AskForApproval::Never),
|
||||
sandbox_policy: Constrained::allow_any(SandboxPolicy::new_read_only_policy()),
|
||||
enforce_residency: Constrained::allow_any(None),
|
||||
did_user_set_custom_approval_policy_or_sandbox_mode: true,
|
||||
forced_auto_mode_downgraded_on_windows: false,
|
||||
shell_environment_policy: ShellEnvironmentPolicy::default(),
|
||||
@@ -3855,6 +3864,7 @@ model_verbosity = "high"
|
||||
model_provider: fixture.openai_chat_completions_provider.clone(),
|
||||
approval_policy: Constrained::allow_any(AskForApproval::UnlessTrusted),
|
||||
sandbox_policy: Constrained::allow_any(SandboxPolicy::new_read_only_policy()),
|
||||
enforce_residency: Constrained::allow_any(None),
|
||||
did_user_set_custom_approval_policy_or_sandbox_mode: true,
|
||||
forced_auto_mode_downgraded_on_windows: false,
|
||||
shell_environment_policy: ShellEnvironmentPolicy::default(),
|
||||
@@ -3954,6 +3964,7 @@ model_verbosity = "high"
|
||||
model_provider: fixture.openai_provider.clone(),
|
||||
approval_policy: Constrained::allow_any(AskForApproval::OnFailure),
|
||||
sandbox_policy: Constrained::allow_any(SandboxPolicy::new_read_only_policy()),
|
||||
enforce_residency: Constrained::allow_any(None),
|
||||
did_user_set_custom_approval_policy_or_sandbox_mode: true,
|
||||
forced_auto_mode_downgraded_on_windows: false,
|
||||
shell_environment_policy: ShellEnvironmentPolicy::default(),
|
||||
@@ -4039,6 +4050,7 @@ model_verbosity = "high"
|
||||
model_provider: fixture.openai_provider.clone(),
|
||||
approval_policy: Constrained::allow_any(AskForApproval::OnFailure),
|
||||
sandbox_policy: Constrained::allow_any(SandboxPolicy::new_read_only_policy()),
|
||||
enforce_residency: Constrained::allow_any(None),
|
||||
did_user_set_custom_approval_policy_or_sandbox_mode: true,
|
||||
forced_auto_mode_downgraded_on_windows: false,
|
||||
shell_environment_policy: ShellEnvironmentPolicy::default(),
|
||||
|
||||
@@ -52,6 +52,7 @@ pub struct ConfigRequirements {
|
||||
pub sandbox_policy: Constrained<SandboxPolicy>,
|
||||
pub mcp_servers: Option<Sourced<BTreeMap<String, McpServerRequirement>>>,
|
||||
pub(crate) exec_policy: Option<Sourced<RequirementsExecPolicy>>,
|
||||
pub enforce_residency: Constrained<Option<ResidencyRequirement>>,
|
||||
}
|
||||
|
||||
impl Default for ConfigRequirements {
|
||||
@@ -61,6 +62,7 @@ impl Default for ConfigRequirements {
|
||||
sandbox_policy: Constrained::allow_any(SandboxPolicy::ReadOnly),
|
||||
mcp_servers: None,
|
||||
exec_policy: None,
|
||||
enforce_residency: Constrained::allow_any(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -84,6 +86,7 @@ pub struct ConfigRequirementsToml {
|
||||
pub allowed_sandbox_modes: Option<Vec<SandboxModeRequirement>>,
|
||||
pub mcp_servers: Option<BTreeMap<String, McpServerRequirement>>,
|
||||
pub rules: Option<RequirementsExecPolicyToml>,
|
||||
pub enforce_residency: Option<ResidencyRequirement>,
|
||||
}
|
||||
|
||||
/// Value paired with the requirement source it came from, for better error
|
||||
@@ -114,6 +117,7 @@ pub struct ConfigRequirementsWithSources {
|
||||
pub allowed_sandbox_modes: Option<Sourced<Vec<SandboxModeRequirement>>>,
|
||||
pub mcp_servers: Option<Sourced<BTreeMap<String, McpServerRequirement>>>,
|
||||
pub rules: Option<Sourced<RequirementsExecPolicyToml>>,
|
||||
pub enforce_residency: Option<Sourced<ResidencyRequirement>>,
|
||||
}
|
||||
|
||||
impl ConfigRequirementsWithSources {
|
||||
@@ -146,6 +150,7 @@ impl ConfigRequirementsWithSources {
|
||||
allowed_sandbox_modes,
|
||||
mcp_servers,
|
||||
rules,
|
||||
enforce_residency,
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -156,12 +161,14 @@ impl ConfigRequirementsWithSources {
|
||||
allowed_sandbox_modes,
|
||||
mcp_servers,
|
||||
rules,
|
||||
enforce_residency,
|
||||
} = self;
|
||||
ConfigRequirementsToml {
|
||||
allowed_approval_policies: allowed_approval_policies.map(|sourced| sourced.value),
|
||||
allowed_sandbox_modes: allowed_sandbox_modes.map(|sourced| sourced.value),
|
||||
mcp_servers: mcp_servers.map(|sourced| sourced.value),
|
||||
rules: rules.map(|sourced| sourced.value),
|
||||
enforce_residency: enforce_residency.map(|sourced| sourced.value),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -193,12 +200,19 @@ impl From<SandboxMode> for SandboxModeRequirement {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ResidencyRequirement {
|
||||
Us,
|
||||
}
|
||||
|
||||
impl ConfigRequirementsToml {
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.allowed_approval_policies.is_none()
|
||||
&& self.allowed_sandbox_modes.is_none()
|
||||
&& self.mcp_servers.is_none()
|
||||
&& self.rules.is_none()
|
||||
&& self.enforce_residency.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -211,6 +225,7 @@ impl TryFrom<ConfigRequirementsWithSources> for ConfigRequirements {
|
||||
allowed_sandbox_modes,
|
||||
mcp_servers,
|
||||
rules,
|
||||
enforce_residency,
|
||||
} = toml;
|
||||
|
||||
let approval_policy: Constrained<AskForApproval> = match allowed_approval_policies {
|
||||
@@ -298,11 +313,33 @@ impl TryFrom<ConfigRequirementsWithSources> for ConfigRequirements {
|
||||
None => None,
|
||||
};
|
||||
|
||||
let enforce_residency: Constrained<Option<ResidencyRequirement>> = match enforce_residency {
|
||||
Some(Sourced {
|
||||
value: residency,
|
||||
source: requirement_source,
|
||||
}) => {
|
||||
let required = Some(residency);
|
||||
Constrained::new(required, move |candidate| {
|
||||
if candidate == &required {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(ConstraintError::InvalidValue {
|
||||
field_name: "enforce_residency",
|
||||
candidate: format!("{candidate:?}"),
|
||||
allowed: format!("{required:?}"),
|
||||
requirement_source: requirement_source.clone(),
|
||||
})
|
||||
}
|
||||
})?
|
||||
}
|
||||
None => Constrained::allow_any(None),
|
||||
};
|
||||
Ok(ConfigRequirements {
|
||||
approval_policy,
|
||||
sandbox_policy,
|
||||
mcp_servers,
|
||||
exec_policy,
|
||||
enforce_residency,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -329,6 +366,7 @@ mod tests {
|
||||
allowed_sandbox_modes,
|
||||
mcp_servers,
|
||||
rules,
|
||||
enforce_residency,
|
||||
} = toml;
|
||||
ConfigRequirementsWithSources {
|
||||
allowed_approval_policies: allowed_approval_policies
|
||||
@@ -337,6 +375,8 @@ mod tests {
|
||||
.map(|value| Sourced::new(value, RequirementSource::Unknown)),
|
||||
mcp_servers: mcp_servers.map(|value| Sourced::new(value, RequirementSource::Unknown)),
|
||||
rules: rules.map(|value| Sourced::new(value, RequirementSource::Unknown)),
|
||||
enforce_residency: enforce_residency
|
||||
.map(|value| Sourced::new(value, RequirementSource::Unknown)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -350,6 +390,8 @@ mod tests {
|
||||
SandboxModeRequirement::WorkspaceWrite,
|
||||
SandboxModeRequirement::DangerFullAccess,
|
||||
];
|
||||
let enforce_residency = ResidencyRequirement::Us;
|
||||
let enforce_source = source.clone();
|
||||
|
||||
// Intentionally constructed without `..Default::default()` so adding a new field to
|
||||
// `ConfigRequirementsToml` forces this test to be updated.
|
||||
@@ -358,6 +400,7 @@ mod tests {
|
||||
allowed_sandbox_modes: Some(allowed_sandbox_modes.clone()),
|
||||
mcp_servers: None,
|
||||
rules: None,
|
||||
enforce_residency: Some(enforce_residency),
|
||||
};
|
||||
|
||||
target.merge_unset_fields(source.clone(), other);
|
||||
@@ -372,6 +415,7 @@ mod tests {
|
||||
allowed_sandbox_modes: Some(Sourced::new(allowed_sandbox_modes, source)),
|
||||
mcp_servers: None,
|
||||
rules: None,
|
||||
enforce_residency: Some(Sourced::new(enforce_residency, enforce_source)),
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -401,6 +445,7 @@ mod tests {
|
||||
allowed_sandbox_modes: None,
|
||||
mcp_servers: None,
|
||||
rules: None,
|
||||
enforce_residency: None,
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
@@ -438,6 +483,7 @@ mod tests {
|
||||
allowed_sandbox_modes: None,
|
||||
mcp_servers: None,
|
||||
rules: None,
|
||||
enforce_residency: None,
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
|
||||
@@ -37,6 +37,7 @@ pub use config_requirements::ConfigRequirementsToml;
|
||||
pub use config_requirements::McpServerIdentity;
|
||||
pub use config_requirements::McpServerRequirement;
|
||||
pub use config_requirements::RequirementSource;
|
||||
pub use config_requirements::ResidencyRequirement;
|
||||
pub use config_requirements::SandboxModeRequirement;
|
||||
pub use config_requirements::Sourced;
|
||||
pub use diagnostics::ConfigError;
|
||||
|
||||
@@ -451,6 +451,7 @@ async fn load_requirements_toml_produces_expected_constraints() -> anyhow::Resul
|
||||
&requirements_file,
|
||||
r#"
|
||||
allowed_approval_policies = ["never", "on-request"]
|
||||
enforce_residency = "us"
|
||||
"#,
|
||||
)
|
||||
.await?;
|
||||
@@ -465,7 +466,6 @@ allowed_approval_policies = ["never", "on-request"]
|
||||
.cloned(),
|
||||
Some(vec![AskForApproval::Never, AskForApproval::OnRequest])
|
||||
);
|
||||
|
||||
let config_requirements: ConfigRequirements = config_requirements_toml.try_into()?;
|
||||
assert_eq!(
|
||||
config_requirements.approval_policy.value(),
|
||||
@@ -480,6 +480,10 @@ allowed_approval_policies = ["never", "on-request"]
|
||||
.can_set(&AskForApproval::OnFailure)
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
config_requirements.enforce_residency.value(),
|
||||
Some(crate::config_loader::ResidencyRequirement::Us)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -503,6 +507,7 @@ allowed_approval_policies = ["on-request"]
|
||||
allowed_sandbox_modes: None,
|
||||
mcp_servers: None,
|
||||
rules: None,
|
||||
enforce_residency: None,
|
||||
},
|
||||
);
|
||||
load_requirements_toml(&mut config_requirements_toml, &requirements_file).await?;
|
||||
@@ -537,6 +542,7 @@ async fn load_config_layers_includes_cloud_requirements() -> anyhow::Result<()>
|
||||
allowed_sandbox_modes: None,
|
||||
mcp_servers: None,
|
||||
rules: None,
|
||||
enforce_residency: None,
|
||||
};
|
||||
let expected = requirements.clone();
|
||||
let cloud_requirements = CloudRequirementsLoader::new(async move { Some(requirements) });
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use crate::config_loader::ResidencyRequirement;
|
||||
use crate::spawn::CODEX_SANDBOX_ENV_VAR;
|
||||
use codex_client::CodexHttpClient;
|
||||
pub use codex_client::CodexRequestBuilder;
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest::header::HeaderValue;
|
||||
use std::sync::LazyLock;
|
||||
use std::sync::Mutex;
|
||||
@@ -24,6 +26,7 @@ use std::sync::RwLock;
|
||||
pub static USER_AGENT_SUFFIX: LazyLock<Mutex<Option<String>>> = 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";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Originator {
|
||||
@@ -31,6 +34,8 @@ pub struct Originator {
|
||||
pub header_value: HeaderValue,
|
||||
}
|
||||
static ORIGINATOR: LazyLock<RwLock<Option<Originator>>> = LazyLock::new(|| RwLock::new(None));
|
||||
static REQUIREMENTS_RESIDENCY: LazyLock<RwLock<Option<ResidencyRequirement>>> =
|
||||
LazyLock::new(|| RwLock::new(None));
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SetOriginatorError {
|
||||
@@ -74,6 +79,14 @@ pub fn set_default_originator(value: String) -> Result<(), SetOriginatorError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_default_client_residency_requirement(enforce_residency: Option<ResidencyRequirement>) {
|
||||
let Ok(mut guard) = REQUIREMENTS_RESIDENCY.write() else {
|
||||
tracing::warn!("Failed to acquire requirements residency lock");
|
||||
return;
|
||||
};
|
||||
*guard = enforce_residency;
|
||||
}
|
||||
|
||||
pub fn originator() -> Originator {
|
||||
if let Ok(guard) = ORIGINATOR.read()
|
||||
&& let Some(originator) = guard.as_ref()
|
||||
@@ -166,10 +179,17 @@ pub fn create_client() -> CodexHttpClient {
|
||||
}
|
||||
|
||||
pub fn build_reqwest_client() -> reqwest::Client {
|
||||
use reqwest::header::HeaderMap;
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("originator", originator().header_value);
|
||||
if let Ok(guard) = REQUIREMENTS_RESIDENCY.read()
|
||||
&& let Some(requirement) = guard.as_ref()
|
||||
&& !headers.contains_key(RESIDENCY_HEADER_NAME)
|
||||
{
|
||||
let value = match requirement {
|
||||
ResidencyRequirement::Us => HeaderValue::from_static("us"),
|
||||
};
|
||||
headers.insert(RESIDENCY_HEADER_NAME, value);
|
||||
}
|
||||
let ua = get_codex_user_agent();
|
||||
|
||||
let mut builder = reqwest::Client::builder()
|
||||
@@ -214,6 +234,8 @@ mod tests {
|
||||
async fn test_create_client_sets_default_headers() {
|
||||
skip_if_no_network!();
|
||||
|
||||
set_default_client_residency_requirement(Some(ResidencyRequirement::Us));
|
||||
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
@@ -256,6 +278,13 @@ mod tests {
|
||||
.get("user-agent")
|
||||
.expect("user-agent header missing");
|
||||
assert_eq!(ua_header.to_str().unwrap(), expected_ua);
|
||||
|
||||
let residency_header = headers
|
||||
.get(RESIDENCY_HEADER_NAME)
|
||||
.expect("residency header missing");
|
||||
assert_eq!(residency_header.to_str().unwrap(), "us");
|
||||
|
||||
set_default_client_residency_requirement(None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user