From 65a5410f72dac1f4b40862c34ef0a8b71b12834a Mon Sep 17 00:00:00 2001 From: goodman-oai Date: Sat, 14 Nov 2026 15:31:19 +0000 Subject: [PATCH] Add trusted request-country metadata to Responses requests (#61535) ## What changed - Add `RequestCountry` for host-supplied country codes. Normalize two-letter ASCII codes to uppercase, omit the field by default, and serialize explicit unknown (`XX`) as `null` to prevent fallback to request provenance. - Carry `request_country` through thread defaults and captured step settings into Codex backend HTTP and WebSocket requests. Allow live country updates without enabling `step_model_switching`, and preserve WebSocket continuation reuse when the country changes. - Propagate the captured country to subagents, reviews, and compaction requests. ## Testing Add coverage for country validation and serialization, thread and live-turn update lifetimes, prewarm settings capture, agent inheritance and reloads, compaction and reviews, WebSocket country changes, and HTTP fallback. GitOrigin-RevId: 787cd1fa7ba5d465a4eb64a92f13c79e93d9bca7 --- .../thread_processor_tests.rs | 1 + .../src/request_processors/turn_processor.rs | 3 + codex-rs/codex-api/src/common.rs | 5 + .../src/endpoint/responses_websocket.rs | 23 ++++ codex-rs/codex-api/tests/clients.rs | 4 + codex-rs/core/src/agent/control/resume.rs | 3 +- codex-rs/core/src/agent/control/spawn.rs | 15 +++ codex-rs/core/src/agent/control_tests.rs | 117 +++++++++++----- codex-rs/core/src/agent/host/spawn.rs | 1 + codex-rs/core/src/agent/types.rs | 1 + codex-rs/core/src/client.rs | 10 ++ codex-rs/core/src/client_common_tests.rs | 4 + .../core/src/client_tool_metadata_tests.rs | 1 + codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/codex_thread.rs | 7 + codex-rs/core/src/compact.rs | 19 +-- .../core/src/compact_remote_v2_attempt.rs | 4 +- codex-rs/core/src/guardian/mod.rs | 5 + codex-rs/core/src/guardian/review_session.rs | 14 +- .../core/src/guardian/review_session_setup.rs | 1 + .../core/src/guardian/review_session_tests.rs | 6 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/responses_metadata.rs | 4 + .../core/src/session/mailbox_prewarm_tests.rs | 51 +++++++ codex-rs/core/src/session/mod.rs | 3 + codex-rs/core/src/session/session.rs | 12 +- codex-rs/core/src/session/startup_prewarm.rs | 1 + codex-rs/core/src/session/step_activation.rs | 9 +- .../core/src/session/step_activation_tests.rs | 57 ++++++++ codex-rs/core/src/session/step_settings.rs | 12 ++ .../core/src/session/step_settings_tests.rs | 6 + codex-rs/core/src/session/tests.rs | 13 +- .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/session/thread_settings.rs | 2 + codex-rs/core/src/session/turn.rs | 2 +- codex-rs/core/src/session/turn_context.rs | 5 + codex-rs/core/src/thread_manager.rs | 9 ++ codex-rs/core/src/thread_spawn.rs | 1 + .../src/tools/handlers/multi_agents/spawn.rs | 1 + .../src/tools/handlers/multi_agents_tests.rs | 14 +- codex-rs/core/tests/common/test_codex.rs | 9 ++ codex-rs/core/tests/suite/abort_tasks.rs | 35 +++++ codex-rs/core/tests/suite/agent_websocket.rs | 126 ++++++++++++++++++ codex-rs/core/tests/suite/compact_remote.rs | 16 +++ codex-rs/core/tests/suite/model_request.rs | 116 ++++++++++++++++ .../tests/suite/reasoning_effort_override.rs | 15 +++ codex-rs/core/tests/suite/review.rs | 34 +++-- codex-rs/core/tests/suite/scenarios.rs | 37 +++++ .../guardian-v2/src/async_scorer/sampler.rs | 1 + codex-rs/protocol/src/lib.rs | 3 + codex-rs/protocol/src/protocol.rs | 4 + codex-rs/protocol/src/request_country.rs | 59 ++++++++ .../protocol/src/request_country_tests.rs | 13 ++ 53 files changed, 854 insertions(+), 64 deletions(-) create mode 100644 codex-rs/core/src/session/mailbox_prewarm_tests.rs create mode 100644 codex-rs/protocol/src/request_country.rs create mode 100644 codex-rs/protocol/src/request_country_tests.rs diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 2fc26a2e4b..cd8bde8c13 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -732,6 +732,7 @@ mod thread_processor_behavior_tests { initial_turns_page: None, }; let config_snapshot = ThreadConfigSnapshot { + request_country: Default::default(), turn_extension_init: Default::default(), disabled_plugin_ids: Vec::new(), model: "gpt-5".to_string(), diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index c51fc93873..f4f50639fd 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -298,6 +298,7 @@ impl TurnRequestProcessor { Op::TurnSettings { turn_id: params.turn_id, update: TurnSettingsUpdate { + request_country: None, approvals_reviewer: params .approvals_reviewer .map(codex_app_server_protocol::ApprovalsReviewer::to_core), @@ -959,6 +960,7 @@ impl TurnRequestProcessor { if has_any_overrides { thread .preview_thread_settings_overrides(CodexThreadSettingsOverrides { + request_country: None, turn_extension_init: None, disabled_plugin_ids: disabled_plugin_ids.clone(), environments: environment_requests.clone(), @@ -985,6 +987,7 @@ impl TurnRequestProcessor { } Ok(codex_protocol::protocol::ThreadSettingsOverrides { + request_country: None, disabled_plugin_ids, environments: environment_requests, runtime_workspace_roots, diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 3942f3aa69..4e62543456 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -349,6 +349,8 @@ pub struct ResponsesApiRequest { pub client_metadata: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub access_programs: Option, + #[serde(skip_serializing_if = "codex_protocol::RequestCountry::is_omitted")] + pub request_country: codex_protocol::RequestCountry, } impl<'a> From<&'a ResponsesApiRequest> for ResponseCreateWsRequest<'a> { @@ -372,6 +374,7 @@ impl<'a> From<&'a ResponsesApiRequest> for ResponseCreateWsRequest<'a> { generate: None, client_metadata: request.client_metadata.clone(), access_programs: request.access_programs, + request_country: &request.request_country, } } } @@ -408,6 +411,8 @@ pub struct ResponseCreateWsRequest<'a> { pub client_metadata: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub access_programs: Option, + #[serde(skip_serializing_if = "codex_protocol::RequestCountry::is_omitted")] + pub request_country: &'a codex_protocol::RequestCountry, } pub fn response_create_client_metadata( diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 20493b6704..4ced74352d 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -1171,7 +1171,22 @@ mod tests { #[test] fn direct_serialization_preserves_websocket_request_payload() { + for (country_code, expected_country) in [ + (None, None), + (Some(" de "), Some(json!("DE"))), + (Some("xX"), Some(Value::Null)), + ] { + assert_serialized_request(country_code, expected_country); + } + } + + fn assert_serialized_request(country_code: Option<&str>, expected_country: Option) { let api_request = ResponsesApiRequest { + request_country: country_code + .map(codex_protocol::RequestCountry::from_iso_code) + .transpose() + .expect("country is valid") + .unwrap_or_default(), model: "gpt-test".to_string(), input: vec![ResponseItem::Message { status: None, @@ -1224,6 +1239,10 @@ mod tests { let mut expected_payload = serde_json::to_value(&api_request).expect("serialize responses API request"); + assert_eq!( + expected_payload.get("request_country"), + expected_country.as_ref() + ); expected_payload["type"] = json!("response.create"); expected_payload["previous_response_id"] = json!("resp-1"); expected_payload["generate"] = json!(false); @@ -1235,6 +1254,10 @@ mod tests { let wire_payload = serde_json::from_str::(&request_text).expect("parse websocket request"); + assert_eq!( + wire_payload.get("request_country"), + expected_country.as_ref() + ); assert_eq!(wire_payload, expected_payload); } diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index fc96b300c5..a32de4fb8a 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -335,6 +335,7 @@ async fn responses_client_stream_request_preserves_item_ids() -> Result<()> { let transport = RecordingTransport::new(state.clone()); let client = ResponsesClient::new(transport, provider("openai"), Arc::new(NoAuth)); let request = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-test".into(), input: vec![ResponseItem::Message { status: None, @@ -392,6 +393,7 @@ async fn responses_client_stream_request_sends_routing_fields_ahead_of_large_inp // buffering a potentially multi-megabyte prompt first. let large_input = "x".repeat(2 * 1024 * 1024); let request = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-test".into(), input: vec![ResponseItem::Message { status: None, @@ -491,6 +493,7 @@ async fn streaming_client_retries_on_transport_error() -> Result<()> { provider.retry.max_attempts = 2; let request = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-test".into(), input: Vec::new(), tools: Some(empty_tools().into()), @@ -609,6 +612,7 @@ async fn azure_store_sends_ids_and_headers() -> Result<()> { let client = ResponsesClient::new(transport, provider("azure"), Arc::new(NoAuth)); let request = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-test".into(), input: vec![ResponseItem::Message { status: None, diff --git a/codex-rs/core/src/agent/control/resume.rs b/codex-rs/core/src/agent/control/resume.rs index f69d8f07a0..3fe7e54f21 100644 --- a/codex-rs/core/src/agent/control/resume.rs +++ b/codex-rs/core/src/agent/control/resume.rs @@ -82,8 +82,9 @@ impl LocalAgentControl { // stores without a graph, using the same identity restoration path. self.restore_agent_identities(&config, &state, vec![id]) .await; + let config = config.as_ref().clone(); let pin = self - .ensure_v2_agent_loaded(config.as_ref().clone(), id, Some(thread)) + .ensure_v2_agent_loaded(config, id, Some(thread)) .await?; thread = state.get_thread(id).await?; // Allow an idle ancestor to give up its slot while its child reloads at capacity. diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index f93bf91b5a..bb3e8f1494 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -546,6 +546,11 @@ impl LocalAgentControl { } else { None }; + let request_country = if let Some(parent) = &inherited_parent { + parent.config_snapshot().await.request_country + } else { + codex_protocol::RequestCountry::UNKNOWN + }; let prompt_cache_key = inherited_parent.as_ref().map(|parent| { parent .session @@ -603,6 +608,7 @@ impl LocalAgentControl { match state .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { config, + request_country, prompt_cache_key, initial_history, agent_control: self.clone(), @@ -765,6 +771,7 @@ impl LocalAgentControl { // Capture child defaults once. Copying conversation history must not change // the capabilities of the turn which requested the spawn. let mut start_options = StartThreadOptions { + request_country: options.request_country.clone(), session_source: session_source.clone(), thread_source: session_source.as_ref().map(|_| ThreadSource::Subagent), environments: options @@ -1213,9 +1220,17 @@ impl LocalAgentControl { let inherited_exec_policy = inherited_exec_policy_for_source(&state, Some(&session_source), &config).await; + let request_country = match session_source.parent_thread_id() { + Some(parent_thread_id) => match state.get_thread(parent_thread_id).await { + Ok(parent) => parent.config_snapshot().await.request_country, + Err(_) => codex_protocol::RequestCountry::UNKNOWN, + }, + None => Default::default(), + }; let resumed_thread = state .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { config: config.clone(), + request_country, prompt_cache_key: None, initial_history, agent_control: self.clone(), diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 32978b4f56..8634e21fbd 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -32,9 +32,16 @@ use crate::init_state_db; use crate::responses_metadata::CodexResponsesRequestKind; use crate::session::SessionSettingsUpdate; use crate::session::step_context::StepContext; +use crate::session::step_settings::StepSettingsUpdate; +use crate::session::tests::update_selected_settings_for_test; use crate::thread_manager::ForkSnapshot; use crate::thread_manager::StartThreadOptions; +use crate::tools::context::ToolCallSource; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; use crate::tools::handlers::multi_agents_common::thread_spawn_source; +use crate::tools::handlers::multi_agents_v2::SendMessageHandler; +use crate::tools::registry::ToolExecutor; use assert_matches::assert_matches; use codex_extension_api::ExtensionDataInit; use codex_extension_api::Instructions; @@ -52,6 +59,7 @@ use codex_login::CodexAuth; use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::model_info::model_info_from_slug; use codex_protocol::AgentPath; +use codex_protocol::RequestCountry; use codex_protocol::ResponseItemId; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -1045,19 +1053,25 @@ async fn send_input_submits_user_message() { let harness = AgentControlHarness::new().await; let (thread_id, thread) = harness.start_thread().await; - let submission_id = harness - .control - .send_input( - thread_id, - vec![UserInput::Text { + let control = thread + .session + .services + .local_agent_runtime + .control(thread.session.session_id()); + let receipt = control + .send(SendRequest { + caller: ThreadId::new(), + target: AgentTarget::Id(thread_id), + resume_config: harness.config.clone(), + input: AgentInput::UserInput(vec![UserInput::Text { text: "hello from tests".to_string(), text_elements: Vec::new(), - }], - Default::default(), - ) + }]), + start_options: Default::default(), + }) .await - .expect("send_input should succeed"); - assert!(!submission_id.is_empty()); + .expect("a live recipient does not require a registered sender"); + assert!(!receipt.submission_id.is_empty()); wait_for_recorded_user_message(thread.as_ref(), "hello from tests").await; } @@ -1291,6 +1305,7 @@ async fn spawn_v2_reload_test_child( text_input("hello child"), source, SpawnAgentOptions { + request_country: parent.config_snapshot().await.request_country, parent_thread_id: Some(parent.session.thread_id), ..Default::default() }, @@ -1304,6 +1319,8 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { let _ = config.features.enable(Feature::MultiAgentV2); let _ = config.features.enable(Feature::Sqlite); config.model = Some("gpt-5.6-sol".to_string()); + let germany = RequestCountry::from_iso_code("DE").unwrap(); + let united_states = RequestCountry::from_iso_code("US").unwrap(); config.multi_agent_v2.max_concurrent_threads_per_session = 3; config.permissions.allow_login_shell = true; config @@ -1316,6 +1333,7 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { let root = harness .manager .start_thread(StartThreadOptions { + request_country: germany.clone(), history_mode: Some(ThreadHistoryMode::Paginated), client_mcp_extensions: client_mcp_extensions.clone(), user_instructions: Some(LoadedUserInstructions { @@ -1333,6 +1351,7 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { }) .await .expect("start root thread"); + assert_eq!(root.thread.config_snapshot().await.request_country, germany); let control = root .thread .session @@ -1446,10 +1465,14 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { parent_turn.config.cwd.clone(), vec![selection], )), + step_settings: StepSettingsUpdate { + request_country: Some(united_states.clone()), + ..Default::default() + }, ..Default::default() }) .await - .expect("save parent environments"); + .expect("save parent environments and country"); parent_thread.session.mark_interrupted(); // The fixture has no task runner to finish the turn or consume child results. *parent_thread.session.active_turn.lock().await = None; @@ -1473,6 +1496,13 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { .get_thread(spawned_agent.thread_id) .await .expect("reloaded child thread should exist"); + assert_eq!( + reloaded_child.config_snapshot().await.request_country, + match route { + V2ReloadRoute::Sender => germany, + V2ReloadRoute::NestedParent => united_states, + }, + ); if matches!(route, V2ReloadRoute::Sender) { assert_eq!( harness @@ -1731,10 +1761,12 @@ async fn cold_resume_with_thread_instructions_preserves_lazy_v2_child_inheritanc let (home, mut config) = test_config().await; let _ = config.features.enable(Feature::MultiAgentV2); let _ = config.features.enable(Feature::Sqlite); + let admitted_country = RequestCountry::from_iso_code("DE").unwrap(); let harness = AgentControlHarness::new_with_config(home, config).await; let parent = harness .manager .start_thread(StartThreadOptions { + request_country: admitted_country.clone(), history_mode: Some(ThreadHistoryMode::Legacy), thread_instructions_provider: Some(Arc::new(TestThreadInstructionsProvider { text: "initial thread instructions".into(), @@ -1814,6 +1846,7 @@ async fn cold_resume_with_thread_instructions_preserves_lazy_v2_child_inheritanc }); let resumed_parent = resumed_manager .start_thread(StartThreadOptions { + request_country: admitted_country.clone(), initial_history, session_source: Some(SessionSource::Exec), thread_instructions_provider: Some(provider.clone()), @@ -1837,30 +1870,37 @@ async fn cold_resume_with_thread_instructions_preserves_lazy_v2_child_inheritanc assert_thread_not_loaded(&resumed_manager, worker_thread_id).await; if nested { - let turn = resumed_parent.thread.session.new_default_turn().await; - resumed_parent - .thread - .session - .services - .agent_control - .send(crate::SendRequest { - caller: parent_thread_id, - target: crate::AgentTarget::Id(target_thread_id), - resume_config: crate::agent::child_config::build_agent_resume_config(&turn) - .expect("capture resume config"), - input: crate::AgentInput::Message { - message: AgentMessage::Plaintext("hello after resume".to_string()), - mode: MessageDeliveryMode::QueueOnly, - }, - start_options: TurnStartOptions { - root_turn_id: turn.turn_metadata_state.root_turn_id(), - turn_trigger: turn.turn_metadata_state.current_turn_trigger(), - cyber_access_program: turn.cyber_access_program, - ..Default::default() + let session = &resumed_parent.thread.session; + let turn = session.new_default_turn().await; + assert_eq!(turn.request_country(), &admitted_country); + let mut issuing_step = StepContext::for_test(Arc::clone(&turn)); + let issuing_country = RequestCountry::from_iso_code("US").unwrap(); + update_selected_settings_for_test( + Arc::make_mut(&mut Arc::make_mut(&mut issuing_step).settings), + |selected| selected.request_country = issuing_country, + ); + SendMessageHandler + .handle(ToolInvocation { + session: Arc::clone(session), + turn, + step_context: issuing_step, + cancellation_token: CancellationToken::new(), + tracker: Default::default(), + call_id: "queue-only-cold-load".to_string(), + tool_name: codex_tools::ToolName::plain("send_message"), + source: ToolCallSource::DirectPlaintextMessage, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "target": target_thread_id.to_string(), + "message": "hello after resume", + }) + .to_string(), }, + native_call: None, + environments: Default::default(), }) .await - .expect("message should reload the grandchild"); + .expect("queue-only message should reload the grandchild"); assert_thread_not_loaded(&resumed_manager, worker_thread_id).await; } else { resumed_manager @@ -1872,6 +1912,14 @@ async fn cold_resume_with_thread_instructions_preserves_lazy_v2_child_inheritanc .get_thread(target_thread_id) .await .expect("resumed target should be loaded"); + assert_eq!( + resumed_target.config_snapshot().await.request_country, + if nested { + RequestCountry::UNKNOWN + } else { + admitted_country + }, + ); let inherited = resumed_target.session.inherited_instructions().await; let root_provider = resumed_parent .thread @@ -5122,6 +5170,7 @@ async fn resume_thread_subagent_restores_stored_metadata() { parent_thread .update_thread_settings(ThreadSettingsOverrides { + request_country: Some(RequestCountry::from_iso_code("US").unwrap()), environments: Some(TurnEnvironmentRequests::new( harness.config.cwd.clone(), Vec::new(), @@ -5205,6 +5254,10 @@ async fn resume_thread_subagent_restores_stored_metadata() { roots ); let resumed_snapshot = resumed_thread.config_snapshot().await; + assert_eq!( + resumed_snapshot.request_country, + RequestCountry::from_iso_code("US").unwrap() + ); let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: resumed_parent_thread_id, depth: resumed_depth, diff --git a/codex-rs/core/src/agent/host/spawn.rs b/codex-rs/core/src/agent/host/spawn.rs index b0cffc5a4d..40e96627e3 100644 --- a/codex-rs/core/src/agent/host/spawn.rs +++ b/codex-rs/core/src/agent/host/spawn.rs @@ -66,6 +66,7 @@ async fn prepare_spawn_inputs( } let exec_policy = inherited_exec_policy_for_parent(parent, &config).await; let mut thread_options = StartThreadOptions { + request_country: options.request_country, session_source: Some(source), thread_source: Some(ThreadSource::Subagent), initial_history, diff --git a/codex-rs/core/src/agent/types.rs b/codex-rs/core/src/agent/types.rs index 55699aae2e..77141692a7 100644 --- a/codex-rs/core/src/agent/types.rs +++ b/codex-rs/core/src/agent/types.rs @@ -33,6 +33,7 @@ pub enum SpawnProgramSource { #[derive(Clone, Debug, Default)] pub struct SpawnAgentOptions { + pub request_country: codex_protocol::RequestCountry, pub fork_parent_spawn_call_id: Option, pub fork_mode: Option, pub parent_thread_id: Option, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index d2fcbbe7f3..439d47de56 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -425,6 +425,7 @@ impl NonIncrementalReason { // This is intentionally not a `PartialEq` implementation: request equality includes `input` and // `client_metadata`, while websocket reuse compares input and late tool-result metadata separately. // Access programs are authorized per response, including continuations, without replaying input. +// Request country is also per response and travels on every continuation. // Keep the destructuring exhaustive so new request fields require an explicit reuse decision. fn responses_request_properties_mismatch( previous: &ResponsesApiRequest, @@ -447,6 +448,7 @@ fn responses_request_properties_mismatch( text: previous_text, client_metadata: _, access_programs: _, + request_country: _, } = previous; let ResponsesApiRequest { model: current_model, @@ -465,6 +467,7 @@ fn responses_request_properties_mismatch( text: current_text, client_metadata: _, access_programs: _, + request_country: _, } = current; // Stream options control delivery for this response, not the context @@ -1311,6 +1314,7 @@ impl ModelClient { text, client_metadata: Some(client_metadata), access_programs: None, + request_country: Default::default(), }; Ok(request) } @@ -2038,6 +2042,9 @@ impl ResponsesHttpExecutor<'_> { responses_metadata, include_internal, )?; + if self.client.uses_codex_backend(client_setup.auth.as_ref()) { + request.request_country = responses_metadata.request_country.clone(); + } self.client.set_guardian_metadata( &mut request.client_metadata, responses_metadata.parent_response_id.as_deref(), @@ -2243,6 +2250,9 @@ impl ModelClientSession { responses_metadata, include_internal, )?; + if self.client.uses_codex_backend(client_setup.auth.as_ref()) { + request.request_country = responses_metadata.request_country.clone(); + } if is_guardian_reviewer(&responses_headers) { request.service_tier = None; } diff --git a/codex-rs/core/src/client_common_tests.rs b/codex-rs/core/src/client_common_tests.rs index b9c9485d90..8732c1028a 100644 --- a/codex-rs/core/src/client_common_tests.rs +++ b/codex-rs/core/src/client_common_tests.rs @@ -113,6 +113,7 @@ fn request_copies_project_image_details_for_receiving_model( fn serializes_text_verbosity_when_set() { let input: Vec = vec![]; let req = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-5.4".to_string(), input, tools: Some(empty_tools().into()), @@ -161,6 +162,7 @@ fn serializes_text_schema_with_strict_format() { .expect("text controls"); let req = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-5.4".to_string(), input, tools: Some(empty_tools().into()), @@ -223,6 +225,7 @@ fn serializes_text_schema_with_non_strict_format() { fn omits_text_when_not_set() { let input: Vec = vec![]; let req = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-5.4".to_string(), input, tools: Some(empty_tools().into()), @@ -248,6 +251,7 @@ fn omits_text_when_not_set() { #[test] fn serializes_flex_service_tier_when_set() { let req = ResponsesApiRequest { + request_country: Default::default(), model: "gpt-5.4".to_string(), input: vec![], tools: Some(empty_tools().into()), diff --git a/codex-rs/core/src/client_tool_metadata_tests.rs b/codex-rs/core/src/client_tool_metadata_tests.rs index 47318a3937..fca3e6bf3c 100644 --- a/codex-rs/core/src/client_tool_metadata_tests.rs +++ b/codex-rs/core/src/client_tool_metadata_tests.rs @@ -43,6 +43,7 @@ fn request_with_metadata_and_source( output.set_tool_call_cell_id("cell"); output.mark_tool_calls_complete(); ResponsesApiRequest { + request_country: Default::default(), model: "test".to_string(), input: vec![output], tools: None, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 3a4e33b583..fc08c00706 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -106,6 +106,7 @@ pub(crate) async fn run_codex_thread_interactive( attached_thread: false, startup: Some(Arc::clone(&startup)), config, + request_country: parent_ctx.request_country().clone(), prompt_cache_key: parent_session .services .model_client diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 801670b023..01e0bd0929 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -28,6 +28,7 @@ use codex_history::RolloutItem; use codex_models_manager::ModelsManagerConfig; use codex_otel::SessionTelemetry; use codex_otel::current_span_w3c_trace_context; +use codex_protocol::RequestCountry; use codex_protocol::ThreadId; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; @@ -100,6 +101,8 @@ static LIVE_THREADS: Gauge = Gauge::new("core.threads.live"); #[derive(Clone, Debug)] pub struct ThreadConfigSnapshot { + /// Trusted host country inherited by future turns. + pub request_country: RequestCountry, /// Host data inherited by future turns; not included in persisted settings. pub turn_extension_init: codex_extension_api::ExtensionDataInit, pub model: String, @@ -160,6 +163,8 @@ impl ThreadConfigSnapshot { /// Thread settings overrides that app-server validates before starting a turn. #[derive(Clone, Default)] pub struct CodexThreadSettingsOverrides { + /// Omission preserves the runtime country. + pub request_country: Option, /// Replaces the data captured by future turns. Omission preserves it. pub turn_extension_init: Option, pub environments: Option, @@ -909,6 +914,7 @@ impl CodexThread { overrides: CodexThreadSettingsOverrides, ) -> SessionSettingsUpdate { let CodexThreadSettingsOverrides { + request_country, turn_extension_init, environments: environment_requests, runtime_workspace_roots, @@ -932,6 +938,7 @@ impl CodexThread { requested_approvals_reviewer, turn_extension_init, step_settings: StepSettingsUpdate { + request_country, model, effort, collaboration_mode, diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index b980c8ff88..0c81d89417 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -120,13 +120,14 @@ fn assemble_compaction_history( pub(crate) async fn run_inline_auto_compact_task( sess: Arc, - turn_context: Arc, + step_context: Arc, replacement_step_context: Arc, world_state: Arc, reason: CompactionReason, phase: CompactionPhase, ) -> CodexResult<()> { - let prompt = turn_context + let prompt = step_context + .turn .config .compact_prompt .as_deref() @@ -140,7 +141,7 @@ pub(crate) async fn run_inline_auto_compact_task( run_compact_task_inner( sess, - turn_context, + step_context, replacement_step_context, input, world_state, @@ -163,7 +164,7 @@ pub(crate) async fn run_compact_task( ) -> CodexResult<()> { run_compact_task_inner( sess.clone(), - Arc::clone(&step_context.turn), + Arc::clone(&step_context), step_context, input, world_state, @@ -180,12 +181,13 @@ pub(crate) async fn run_compact_task( async fn run_compact_task_inner( sess: Arc, - turn_context: Arc, + step_context: Arc, replacement_step_context: Arc, input: Vec, world_state: Arc, compaction_metadata: CompactionTurnMetadata, ) -> CodexResult<()> { + let turn_context = Arc::clone(&step_context.turn); let trigger = compaction_metadata.trigger(); let reason = compaction_metadata.reason(); let phase = compaction_metadata.phase(); @@ -216,7 +218,7 @@ async fn run_compact_task_inner( } let result = run_compact_task_inner_impl( Arc::clone(&sess), - Arc::clone(&turn_context), + step_context, replacement_step_context, input, world_state, @@ -266,12 +268,13 @@ async fn run_compact_task_inner( async fn run_compact_task_inner_impl( sess: Arc, - turn_context: Arc, + step_context: Arc, replacement_step_context: Arc, input: Vec, world_state: Arc, compaction_metadata: CompactionTurnMetadata, ) -> CodexResult { + let turn_context = Arc::clone(&step_context.turn); let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new()); sess.emit_turn_item_started(&turn_context, &compaction_item) .await; @@ -331,7 +334,7 @@ async fn run_compact_task_inner_impl( ..Default::default() }; let responses_metadata = sess - .compaction_responses_metadata(turn_context.as_ref(), compaction_metadata) + .compaction_responses_metadata(&step_context, compaction_metadata) .await; let attempt_result = drain_to_completed( &sess, diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index e0f2ab5cbc..096eddeba5 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -199,10 +199,10 @@ pub(super) async fn run_remote_compact_v2_attempt( let mut responses_metadata = match captured_window { Some(window) => { - sess.compaction_responses_metadata_for_window(turn_context, compaction_metadata, window) + sess.compaction_responses_metadata_for_window(step_context, compaction_metadata, window) } None => { - sess.compaction_responses_metadata(turn_context, compaction_metadata) + sess.compaction_responses_metadata(step_context, compaction_metadata) .await } }; diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 0669d326cd..c9610e6685 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -36,6 +36,7 @@ pub(crate) use score_authorization::ServingAssessment; #[cfg(test)] pub(crate) mod test_host; +use codex_protocol::RequestCountry; use codex_protocol::items::ModelInvocationContext; use codex_protocol::openai_models::CatalogModel; use std::sync::Arc; @@ -94,6 +95,7 @@ pub(crate) use codex_guardian_reviewer::GUARDIAN_MAX_NODE_REPL_TOOL_RESULT_TOKEN /// MCP elicitations use their live invocation's issuing step when available, otherwise turn inputs. #[derive(Clone)] pub(crate) struct GuardianReviewContext { + pub(crate) request_country: RequestCountry, pub(crate) serving_assessment: Option>, /// The latest response ID and its connection affinity when review was requested. pub(crate) parent_response: Option>, @@ -139,6 +141,7 @@ impl GuardianReviewContext { environments: &TurnEnvironmentSnapshot, ) -> Self { Self { + request_country: settings.request_country().clone(), serving_assessment: None, review_execution: None, parent_response: turn.extension_data.get::(), @@ -166,6 +169,7 @@ impl GuardianReviewContext { impl From<&Arc> for GuardianReviewContext { fn from(step: &Arc) -> Self { Self { + request_country: step.settings.request_country().clone(), serving_assessment: step.serving_assessment.clone(), review_execution: None, parent_response: step.turn.extension_data.get::(), @@ -185,6 +189,7 @@ impl From<&Arc> for GuardianReviewContext { impl From> for GuardianReviewContext { fn from(turn: Arc) -> Self { Self { + request_country: turn.initial_settings.request_country().clone(), serving_assessment: None, review_execution: None, parent_response: turn.extension_data.get::(), diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 41dbedd532..bdb84ff2fe 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -278,7 +278,19 @@ async fn run_review_on_session( match run_before_review_deadline( deadline, params.external_cancel.as_ref(), - Box::pin(ensure_guardian_node_repl_policy(review_session, params)), + Box::pin(async { + review_session + .session + .update_settings(crate::session::session::SessionSettingsUpdate { + step_settings: crate::session::step_settings::StepSettingsUpdate { + request_country: Some(params.parent_context.request_country.clone()), + ..Default::default() + }, + ..Default::default() + }) + .await?; + ensure_guardian_node_repl_policy(review_session, params).await + }), ) .await { diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index 17b7fa3d6b..88d245d297 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -111,6 +111,7 @@ impl PreparedGuardianContext { .model_client .responses_websocket_enabled(); let options = crate::StartThreadOptions { + request_country: self.context.request_country.clone(), prompt_cache_key: self.parent.services.model_client.prompt_cache_key_policy(), history_mode: Some(codex_protocol::protocol::ThreadHistoryMode::Paginated), internal_parent: Some(crate::thread_manager::InternalSessionParent { diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index d75deeb23e..8d130b3fa9 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -1132,6 +1132,8 @@ async fn compaction_and_review_submissions_share_execution_observer() { let (review_session, events, submissions) = test_review_session().await; let child_session = Arc::clone(&review_session.session); let mut params = test_review_params().await; + let expected_country = codex_protocol::RequestCountry::UNKNOWN; + params.parent_context.request_country = expected_country.clone(); params .spawn_config .features @@ -1164,6 +1166,10 @@ async fn compaction_and_review_submissions_share_execution_observer() { .expect("child submission"); if index == 0 { assert!(matches!(submission.op, Op::Compact)); + assert_eq!( + child_session.thread_config_snapshot().await.request_country, + expected_country + ); } else { let Op::TurnInput { reply, .. } = submission.op else { panic!("expected review input after compaction"); diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 9e3d4c4f20..b6ece36a24 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -18,6 +18,8 @@ pub use responses_headers::CodexResponsesHeaders; mod responses_metadata; mod responses_retry; pub(crate) mod session; +pub use codex_protocol::InvalidRequestCountry; +pub use codex_protocol::RequestCountry; pub use codex_protocol::turn_input::NotSubmittedReason; pub use codex_protocol::turn_input::PromptCacheTtl; pub use codex_protocol::turn_input::RecoverTurnRequest; diff --git a/codex-rs/core/src/responses_metadata.rs b/codex-rs/core/src/responses_metadata.rs index e939afc2da..fd83de8aa5 100644 --- a/codex-rs/core/src/responses_metadata.rs +++ b/codex-rs/core/src/responses_metadata.rs @@ -1,3 +1,4 @@ +use codex_protocol::RequestCountry; use std::collections::BTreeMap; use std::collections::HashMap; use std::io::ErrorKind; @@ -237,6 +238,8 @@ pub(crate) enum TurnToolSource { /// truth. #[derive(Clone, Debug)] pub struct CodexResponsesMetadata { + /// Trusted country captured for this request. + pub(crate) request_country: RequestCountry, /// The private reviewer forwards recorded items through its parent's raw-event subscription. pub(crate) guardian_recorded_items: bool, /// Guardian parent reference; projected only onto a Guardian request. @@ -288,6 +291,7 @@ impl CodexResponsesMetadata { window_id: String, ) -> Self { Self { + request_country: RequestCountry::default(), parent_response_id: None, parent_routing_affinity: None, guardian_recorded_items: false, diff --git a/codex-rs/core/src/session/mailbox_prewarm_tests.rs b/codex-rs/core/src/session/mailbox_prewarm_tests.rs new file mode 100644 index 0000000000..bb60755bc6 --- /dev/null +++ b/codex-rs/core/src/session/mailbox_prewarm_tests.rs @@ -0,0 +1,51 @@ +//! Checks mailbox turn admission and startup prewarm settings capture. +#![cfg(test)] + +use super::*; +use codex_protocol::openai_models::ReasoningEffort; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn prewarm_turn_keeps_its_captured_settings_after_defaults_change() { + let (session, _turn) = make_session_and_context().await; + let germany = RequestCountry::from_iso_code("DE").unwrap(); + let united_states = RequestCountry::from_iso_code("US").unwrap(); + session + .update_settings(SessionSettingsUpdate { + step_settings: StepSettingsUpdate { + request_country: Some(germany.clone()), + effort: Some(Some(ReasoningEffort::Low)), + ..Default::default() + }, + ..Default::default() + }) + .await + .expect("update thread settings"); + let prewarm = session + .new_startup_prewarm_turn_with_sub_id(INITIAL_SUBMIT_ID.to_owned()) + .await; + session + .update_settings(SessionSettingsUpdate { + step_settings: StepSettingsUpdate { + request_country: Some(united_states.clone()), + effort: Some(Some(ReasoningEffort::High)), + ..Default::default() + }, + ..Default::default() + }) + .await + .expect("update thread settings"); + + assert_eq!( + ( + prewarm.request_country(), + prewarm.initial_settings.reasoning_effort() + ), + (&germany, Some(&ReasoningEffort::Low)) + ); + let defaults = session.thread_config_snapshot().await; + assert_eq!( + (defaults.request_country, defaults.reasoning_effort), + (united_states, Some(ReasoningEffort::High)) + ); +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 916fe9b2e5..ce1297bcdc 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -511,6 +511,7 @@ pub(crate) struct SessionSpawnArgs { pub(crate) attached_thread: bool, pub(crate) startup: Option>, pub(crate) config: Config, + pub(crate) request_country: codex_protocol::RequestCountry, pub(crate) prompt_cache_key: PromptCacheKey, pub(crate) allow_provider_model_fallback: bool, pub(crate) instructions: SessionInstructions, @@ -705,6 +706,7 @@ impl Session { attached_thread, startup, config, + request_country, prompt_cache_key, allow_provider_model_fallback, instructions, @@ -1030,6 +1032,7 @@ impl Session { turn_extension_init, provider: provider_auth.model_provider(config.model_provider.clone()), step_settings: Arc::new(StepSettings { + request_country, collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: config.service_tier.clone(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a6e902ed03..97acf034fe 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -285,6 +285,7 @@ impl SessionConfiguration { .map(|config| config.permission_profile.clone()) .unwrap_or_else(|| self.permission_profile_state.snapshot()); ThreadConfigSnapshot { + request_country: self.step_settings.request_country.clone(), model: self.step_settings.collaboration_mode.model().to_string(), model_provider_id: self.original_config_do_not_use.model_provider_id.clone(), service_tier: self.step_settings.service_tier.clone(), @@ -694,6 +695,7 @@ impl Session { request_kind, ); ExecutionMetadata::from_settings(&step_context.settings).apply_to(&mut responses_metadata); + responses_metadata.request_country = step_context.settings.request_country().clone(); responses_metadata.guardian_recorded_items = self.parent_observer_event_tx.get().is_some(); if responses_metadata.subagent_header.as_deref() == Some(crate::guardian::GUARDIAN_REVIEWER_NAME) @@ -730,11 +732,11 @@ impl Session { // request remains turn-backed; local compaction does not have a finalized request inventory. pub(crate) async fn compaction_responses_metadata( &self, - turn_context: &TurnContext, + step_context: &StepContext, compaction_metadata: CompactionTurnMetadata, ) -> CodexResponsesMetadata { self.compaction_responses_metadata_for_window( - turn_context, + step_context, compaction_metadata, self.current_window().await, ) @@ -742,15 +744,17 @@ impl Session { pub(crate) fn compaction_responses_metadata_for_window( &self, - turn_context: &TurnContext, + step_context: &StepContext, compaction_metadata: CompactionTurnMetadata, (window_id, window_number, context_window_id): (String, u64, ContextWindowId), ) -> CodexResponsesMetadata { - let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( + let turn_context = &step_context.turn; + let mut responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( self.installation_id.clone(), window_id, CodexResponsesRequestKind::Compaction(compaction_metadata), ); + responses_metadata.request_country = step_context.settings.request_country().clone(); self.with_window_and_fork_metadata( turn_context, responses_metadata, diff --git a/codex-rs/core/src/session/startup_prewarm.rs b/codex-rs/core/src/session/startup_prewarm.rs index b5f754a759..52fb528f8b 100644 --- a/codex-rs/core/src/session/startup_prewarm.rs +++ b/codex-rs/core/src/session/startup_prewarm.rs @@ -1,6 +1,7 @@ //! Background model warmup shared by startup and idle-thread resume. //! One scheduled task owns the prepared client session until the next regular //! turn consumes it; shutdown and turn cancellation use the same handoff. +//! Warmup settings are captured when the worker constructs its turn. use std::collections::HashMap; use std::sync::Arc; diff --git a/codex-rs/core/src/session/step_activation.rs b/codex-rs/core/src/session/step_activation.rs index ff09dd5f88..db3d472a08 100644 --- a/codex-rs/core/src/session/step_activation.rs +++ b/codex-rs/core/src/session/step_activation.rs @@ -241,7 +241,9 @@ impl Session { || update.summary.is_some() || update.service_tier.is_some(); let requires_model_switching = updates_model_settings - || (update.approvals_reviewer.is_none() && update.environments.is_none()); + || (update.approvals_reviewer.is_none() + && update.environments.is_none() + && update.request_country.is_none()); if requires_model_switching && !self.features.enabled(Feature::StepModelSwitching) { return TurnSettingsUpdateOutcome::Rejected { reason: "turn settings updates require the step_model_switching feature" @@ -284,6 +286,7 @@ impl Session { }; } let TurnSettingsUpdate { + request_country, approvals_reviewer, environments: environment_requests, model, @@ -299,8 +302,10 @@ impl Session { }) .collect::>() }); - let updates_step_settings = updates_model_settings || approvals_reviewer.is_some(); + let updates_step_settings = + updates_model_settings || approvals_reviewer.is_some() || request_country.is_some(); let update = StepSettingsUpdate { + request_country, approvals_reviewer, model, effort, diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index 65c17e5331..2366295b76 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -33,6 +33,7 @@ use codex_models_manager::manager::ModelsManagerFuture; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::StaticModelsManager; use codex_models_manager::model_info::with_config_overrides; +use codex_protocol::RequestCountry; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::config_types::ReasoningSummary; @@ -885,6 +886,62 @@ enum ManagedAuthorizationChange { ApprovalsReviewer, } +#[tokio::test] +async fn country_updates_preserve_captured_steps_without_model_switching() { + let (mut session, turn) = make_session_and_context().await; + session + .features + .disable(Feature::StepModelSwitching) + .unwrap(); + let session = Arc::new(session); + let turn = Arc::new(turn); + session + .spawn_task( + Arc::clone(&turn), + Vec::new(), + HeldStepTask { + kind: TaskKind::Compact, + finish: Arc::new(Notify::new()), + }, + ) + .await; + let original = turn.next_step_settings.load_full(); + let defaults = desired_step_settings(&session).await; + for request_country in [ + RequestCountry::from_iso_code("DE").unwrap(), + RequestCountry::UNKNOWN, + ] { + let captured = turn.next_step_settings.load_full(); + let expected = captured.request_country().clone(); + assert_eq!( + session + .apply_turn_settings( + &turn.sub_id, + TurnSettingsUpdate { + request_country: Some(request_country.clone()), + ..Default::default() + }, + ) + .await, + TurnSettingsUpdateOutcome::Applied, + ); + assert_eq!( + turn.next_step_settings.load().request_country(), + &request_country + ); + assert_eq!(captured.request_country(), &expected); + } + let issuing_step = session + .capture_step_context(Arc::clone(&turn), &CancellationToken::new()) + .await + .unwrap(); + let review = crate::guardian::GuardianReviewContext::from(&issuing_step); + assert_eq!(review.request_country, RequestCountry::UNKNOWN); + assert_eq!(original.request_country(), &RequestCountry::default()); + assert_eq!(desired_step_settings(&session).await, defaults); + session.abort_all_tasks(TurnAbortReason::Replaced).await; +} + #[test_case(false; "reviewer allow-list")] #[test_case(true; "model-required review")] #[tokio::test] diff --git a/codex-rs/core/src/session/step_settings.rs b/codex-rs/core/src/session/step_settings.rs index 3a75062ab5..203d6efc60 100644 --- a/codex-rs/core/src/session/step_settings.rs +++ b/codex-rs/core/src/session/step_settings.rs @@ -8,6 +8,7 @@ use codex_features::Features; use codex_models_manager::ModelsManagerConfig; use codex_models_manager::manager::ModelsManager; use codex_otel::SessionTelemetry; +use codex_protocol::RequestCountry; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::Personality; @@ -23,6 +24,8 @@ use std::sync::Arc; /// a turn. A turn may contain several steps, each using its own captured settings. #[derive(Clone, Debug, PartialEq)] pub(crate) struct StepSettings { + /// Trusted host country captured for requests issued with these settings. + pub(crate) request_country: RequestCountry, pub(crate) collaboration_mode: CollaborationMode, /// `None` follows the pinned model's default. pub(crate) reasoning_summary: Option, @@ -53,6 +56,10 @@ pub(crate) struct ResolvedStepSettings { } impl ResolvedStepSettings { + pub(crate) fn request_country(&self) -> &RequestCountry { + &self.selected.request_country + } + /// Derives request values from selected settings and pinned metadata. pub(super) fn new( selected: Arc, @@ -215,6 +222,8 @@ impl ModelInfoOverrides { /// that full value for another. Merge the requested edits with each target. #[derive(Clone, Debug, Default, PartialEq)] pub(crate) struct StepSettingsUpdate { + /// Omission preserves the runtime country. + pub(crate) request_country: Option, pub(crate) model: Option, pub(crate) effort: Option>, /// A complete collaboration mode takes precedence over model and effort edits. @@ -256,6 +265,9 @@ impl StepSettings { constraints: &StepSettingsConstraints<'_>, ) -> ConstraintResult { let mut next = self.clone(); + if let Some(request_country) = &update.request_country { + next.request_country.clone_from(request_country); + } next.collaboration_mode = update.collaboration_mode.clone().unwrap_or_else(|| { self.collaboration_mode.with_updates( update.model.clone(), diff --git a/codex-rs/core/src/session/step_settings_tests.rs b/codex-rs/core/src/session/step_settings_tests.rs index 049ca1b428..8d24c01a5d 100644 --- a/codex-rs/core/src/session/step_settings_tests.rs +++ b/codex-rs/core/src/session/step_settings_tests.rs @@ -11,6 +11,7 @@ use codex_features::Feature; use codex_models_manager::manager::StaticModelsManager; use codex_models_manager::model_info::model_info_from_slug; use codex_prompts::render_model_instructions; +use codex_protocol::RequestCountry; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; use codex_protocol::models::BaseInstructionsProvenance; @@ -292,6 +293,7 @@ async fn environment_only_update_revalidates_existing_step_settings() { fn configured_settings() -> StepSettings { StepSettings { + request_country: RequestCountry::from_iso_code("DE").unwrap(), collaboration_mode: CollaborationMode { mode: ModeKind::Default, settings: Settings { @@ -323,6 +325,7 @@ fn sparse_patch_uses_the_settings_version_being_updated() { let constraints = step_settings_constraints(&requirements); let initial = configured_settings(); let tier_update = StepSettingsUpdate { + request_country: Some(RequestCountry::from_iso_code("US").unwrap()), service_tier: Some(Some("fast".to_string())), ..Default::default() }; @@ -341,6 +344,7 @@ fn sparse_patch_uses_the_settings_version_being_updated() { .apply(&tier_update, &constraints) .expect("a previously prepared sparse update should apply"); let expected = StepSettings { + request_country: RequestCountry::from_iso_code("US").unwrap(), service_tier: Some(ServiceTier::Fast.request_value().to_string()), ..latest }; @@ -383,6 +387,7 @@ fn collaboration_replacement_wins_and_effort_clear_remains_sparse() { .apply( &StepSettingsUpdate { effort: Some(None), + request_country: Some(RequestCountry::UNKNOWN), ..Default::default() }, &constraints, @@ -391,6 +396,7 @@ fn collaboration_replacement_wins_and_effort_clear_remains_sparse() { assert_eq!( cleared, StepSettings { + request_country: RequestCountry::UNKNOWN, collaboration_mode: replacement.with_updates( /*model*/ None, Some(None), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 80144a67db..3a7236db02 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2,10 +2,14 @@ use crate::McpRefreshReason; use crate::agent::LocalAgentControl; use codex_mcp::McpServerKey; use codex_protocol::ContextWindowId; +use codex_protocol::RequestCountry; use codex_thread_store::ThreadPersistenceMetadata; #[path = "compaction_replay_tests.rs"] mod compaction_replay_tests; +#[path = "mailbox_prewarm_tests.rs"] +mod mailbox_prewarm_tests; + #[path = "network_policy_tests.rs"] mod network_policy_tests; @@ -4326,6 +4330,7 @@ async fn set_rate_limits_retains_previous_credits() { provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None), environments: Vec::new(), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, @@ -4450,6 +4455,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None), environments: Vec::new(), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, @@ -5076,6 +5082,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati provider: create_model_provider(config.model_provider.clone(), /*auth_manager*/ None), environments: Vec::new(), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, @@ -6671,7 +6678,7 @@ async fn response_metadata_builders_capture_fresh_mcp_attribution() { .await; let compaction = session .compaction_responses_metadata( - &turn_context, + &step_context, CompactionTurnMetadata::new( CompactionTrigger::Auto, CompactionReason::ContextLimit, @@ -6746,6 +6753,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { Some(Arc::clone(&auth_manager)), ), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, @@ -7084,6 +7092,7 @@ async fn make_session_with_config_and_rx( Some(Arc::clone(&auth_manager)), ), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, @@ -7223,6 +7232,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( Some(Arc::clone(&auth_manager)), ), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, @@ -9095,6 +9105,7 @@ where Some(Arc::clone(&auth_manager)), ), step_settings: Arc::new(StepSettings { + request_country: RequestCountry::default(), collaboration_mode, reasoning_summary: config.model_reasoning_summary, service_tier: None, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 24df4e9d0a..eeedf62ffc 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1333,6 +1333,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { attached_thread: false, startup: None, config, + request_country: Default::default(), prompt_cache_key: PromptCacheKey::default(), allow_provider_model_fallback: false, instructions: Default::default(), diff --git a/codex-rs/core/src/session/thread_settings.rs b/codex-rs/core/src/session/thread_settings.rs index facc5973bb..6191799d3d 100644 --- a/codex-rs/core/src/session/thread_settings.rs +++ b/codex-rs/core/src/session/thread_settings.rs @@ -105,6 +105,7 @@ pub(super) fn prepare_update( turn_extension_init, } = overrides.into(); let ThreadSettingsOverrides { + request_country, environments: environment_requests, runtime_workspace_roots, profile_workspace_roots, @@ -125,6 +126,7 @@ pub(super) fn prepare_update( SessionSettingsUpdate { turn_extension_init, step_settings: StepSettingsUpdate { + request_country, model, effort, collaboration_mode, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8ebb2e9550..c3d3a36c48 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1779,7 +1779,7 @@ async fn run_auto_compact( ); run_inline_auto_compact_task( Arc::clone(sess), - Arc::clone(turn_context), + step_context, replacement_step_context, world_state, reason, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index ebf12d21dd..a21c2c95c9 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -603,6 +603,11 @@ impl TurnContext { Arc::clone(&self.next_step_settings.load().model_info) } + /// Country captured when this turn was admitted. Subsequent steps capture their own settings. + pub fn request_country(&self) -> &codex_protocol::RequestCountry { + self.initial_settings.request_country() + } + /// Legacy: returns the frozen initial-turn model metadata. /// Step-scoped consumers should use their captured `StepContext::settings`. pub(crate) fn model_info(&self) -> &Arc { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 33f4bc47ec..8c2fd45869 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -344,6 +344,8 @@ pub enum PromptCacheKey { pub struct StartThreadOptions { pub config: Config, + /// Trusted host country for the initial thread settings. + pub request_country: codex_protocol::RequestCountry, /// Request-body cache key policy for this thread and its children. Hosts must reapply it on cold resume. pub prompt_cache_key: PromptCacheKey, /// Host-owned provider for this root thread's additional instructions. @@ -441,6 +443,7 @@ impl StartThreadOptions { pub fn new(config: Config) -> Self { Self { config, + request_country: codex_protocol::RequestCountry::default(), prompt_cache_key: PromptCacheKey::default(), thread_instructions_provider: None, allow_provider_model_fallback: false, @@ -566,6 +569,7 @@ fn effective_originator_value( pub(crate) struct ResumeThreadWithHistoryOptions { pub(crate) config: Config, + pub(crate) request_country: codex_protocol::RequestCountry, pub(crate) prompt_cache_key: Option, pub(crate) initial_history: InitialHistory, pub(crate) agent_control: LocalAgentControl, @@ -1783,6 +1787,7 @@ impl ThreadManager { thread_extension_init.insert(reasoning_effort_pin); Ok(( StartThreadOptions { + request_country: settings.request_country, history_mode: Some(settings.history_mode), environments: Some(settings.environments.into_requests().environment_requests), inherited_environments: Some(environments), @@ -2409,6 +2414,7 @@ impl ThreadManagerState { ) -> CodexResult { let ResumeThreadWithHistoryOptions { config, + request_country, prompt_cache_key, initial_history, agent_control, @@ -2426,6 +2432,7 @@ impl ThreadManagerState { let prompt_cache_key = prompt_cache_key.unwrap_or(inherited_prompt_cache_key); let thread_source = initial_history.get_resumed_thread_source(); let options = StartThreadOptions { + request_country, prompt_cache_key, initial_history, session_source: Some(session_source), @@ -2556,6 +2563,7 @@ impl ThreadManagerState { }; let StartThreadOptions { mut config, + request_country, prompt_cache_key, thread_instructions_provider, allow_provider_model_fallback, @@ -2899,6 +2907,7 @@ impl ThreadManagerState { attached_thread, startup, config, + request_country, prompt_cache_key, allow_provider_model_fallback, instructions: startup_context.instructions, diff --git a/codex-rs/core/src/thread_spawn.rs b/codex-rs/core/src/thread_spawn.rs index 58dab95f05..c68c71a2de 100644 --- a/codex-rs/core/src/thread_spawn.rs +++ b/codex-rs/core/src/thread_spawn.rs @@ -198,6 +198,7 @@ pub(crate) async fn prepare( config, source: spawn_source, options: SpawnAgentOptions { + request_country: step_context.settings.request_country().clone(), fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.to_string()), fork_mode, parent_thread_id: Some(session.thread_id), diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 42ec6dffe9..28bf068798 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -127,6 +127,7 @@ async fn handle_spawn_agent( /*task_name*/ None, )?, options: SpawnAgentOptions { + request_country: step_context.settings.request_country().clone(), fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), fork_mode: fork_mode.clone(), parent_thread_id: Some(session.thread_id), diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index e8db90c790..189f021201 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -396,7 +396,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { turn.provider = create_model_provider(provider_info, turn.auth_manager.clone()); turn.config = Arc::new(config); - let invocation = invocation( + let mut invocation = invocation( Arc::new(session), Arc::new(turn), "spawn_agent", @@ -405,6 +405,14 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { "agent_type": "explorer" })), ); + let settings = Arc::make_mut( + &mut Arc::get_mut(&mut invocation.step_context) + .expect("unshared invocation step") + .settings, + ); + update_selected_settings_for_test(settings, |selected| { + selected.request_country = codex_protocol::RequestCountry::UNKNOWN; + }); let output = SpawnAgentHandler::default() .handle(invocation) .await @@ -425,6 +433,10 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { .expect("spawned agent thread should exist") .config_snapshot() .await; + assert_eq!( + snapshot.request_country, + codex_protocol::RequestCountry::UNKNOWN + ); assert_eq!(snapshot.approval_policy, AskForApproval::OnRequest); assert_eq!(snapshot.model_provider_id, "ollama"); } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 2ce2fccba9..25cc2365a1 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -409,6 +409,7 @@ impl TestAuth { } pub struct TestCodexBuilder { + request_country: codex_protocol::RequestCountry, apps_tools_cache: codex_core::CodexAppsToolsCache, config_mutators: Vec>, thread_manager_configurer: Option ThreadManager + Send>>, @@ -434,6 +435,12 @@ pub struct TestCodexBuilder { } impl TestCodexBuilder { + /// Set the initial country when creating a new thread. + pub fn with_request_country(mut self, request_country: codex_protocol::RequestCountry) -> Self { + self.request_country = request_country; + self + } + pub fn with_apps_tools_cache(mut self, cache: codex_core::CodexAppsToolsCache) -> Self { self.apps_tools_cache = cache; self @@ -997,6 +1004,7 @@ impl TestCodexBuilder { None }; Box::pin(thread_manager.start_thread(StartThreadOptions { + request_country: self.request_country.clone(), history_mode: self.history_mode, client_mcp_extensions: client_mcp_extensions(), environments: environment_requests, @@ -1585,6 +1593,7 @@ fn function_call_output<'a>(bodies: &'a [Value], call_id: &str) -> &'a Value { pub fn test_codex() -> TestCodexBuilder { TestCodexBuilder { + request_country: Default::default(), apps_tools_cache: Default::default(), thread_manager_configurer: None, config_mutators: vec![Box::new(|config| { diff --git a/codex-rs/core/tests/suite/abort_tasks.rs b/codex-rs/core/tests/suite/abort_tasks.rs index 8c3419edc8..e0507ba86e 100644 --- a/codex-rs/core/tests/suite/abort_tasks.rs +++ b/codex-rs/core/tests/suite/abort_tasks.rs @@ -13,10 +13,12 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use codex_protocol::RequestCountry; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use core_test_support::ThreadIdle; use core_test_support::responses::ev_assistant_message; @@ -491,6 +493,12 @@ async fn graceful_suspend_preserves_a_finishing_child_result( } let server = start_mock_server().await; let response = mount_sse_once(&server, sse(vec![ev_completed("received")])).await; + let germany = RequestCountry::from_iso_code("DE")?; + let initial_country = if scenario == ChildHandoff::GrandchildFinishes { + germany.clone() + } else { + RequestCountry::default() + }; let replacement = test_codex() .with_model("gpt-5.6-sol") .with_config(|config| { @@ -503,7 +511,23 @@ async fn graceful_suspend_preserves_a_finishing_child_result( .restart_with_auto_env(&server, &source) .await?; let root = replacement.codex.clone(); + if !initial_country.is_omitted() { + root.update_thread_settings(ThreadSettingsOverrides { + request_country: Some(initial_country.clone()), + ..Default::default() + }) + .await?; + } + assert_eq!( + root.config_snapshot().await.request_country, + initial_country + ); if scenario == ChildHandoff::ParentEvicted { + root.update_thread_settings(ThreadSettingsOverrides { + request_country: Some(germany.clone()), + ..Default::default() + }) + .await?; // The cold worker has one child slot. Load the grandchild through its real // parent and release the parent's slot before admitting the grandchild. let leaf_id = @@ -527,6 +551,8 @@ async fn graceful_suspend_preserves_a_finishing_child_result( leaf.thread.session_source().parent_thread_id(), Some(worker_id) ); + assert_eq!(root.config_snapshot().await.request_country, germany); + assert_eq!(leaf.thread.config_snapshot().await.request_country, germany); assert!( replacement .thread_manager @@ -537,10 +563,19 @@ async fn graceful_suspend_preserves_a_finishing_child_result( return Ok(()); } let parent = if scenario == ChildHandoff::GrandchildFinishes { + root.update_thread_settings(ThreadSettingsOverrides { + request_country: Some(RequestCountry::default()), + ..Default::default() + }) + .await?; let loaded = replacement .thread_manager .load_suspended_agent(&root, worker_id) .await?; + assert_eq!( + loaded.thread.config_snapshot().await.request_country, + RequestCountry::default() + ); assert_matches!(loaded.thread.recover_turn_if_idle(RecoverTurnRequest { turn_id: worker_turn.clone(), thread_settings: Default::default(), trace: None, prompt_cache_ttl: None, cyber_access_program: None }).await?, codex_core::StartIfIdleSubmission::Started { turn_id, .. } if turn_id == worker_turn); Arc::clone(&loaded.thread) diff --git a/codex-rs/core/tests/suite/agent_websocket.rs b/codex-rs/core/tests/suite/agent_websocket.rs index 214167fa19..4118b4b4f6 100644 --- a/codex-rs/core/tests/suite/agent_websocket.rs +++ b/codex-rs/core/tests/suite/agent_websocket.rs @@ -6,6 +6,7 @@ use codex_core::AgentServiceTier; use codex_core::AgentTurnAdmission; use codex_core::CodexThreadSettingsOverrides; use codex_core::InheritedSessionId; +use codex_core::RequestCountry; use codex_core::SteerSubmission; use codex_core::ThreadManager; use codex_core::TurnInputRequest; @@ -926,6 +927,131 @@ async fn websocket_first_turn_uses_startup_prewarm_and_create( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_country_changes_on_one_connection() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: ["warm", "first", "refreshed", "cleared", "omitted"] + .into_iter() + .map(|id| vec![ev_response_created(id), ev_completed(id)]) + .collect(), + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: false, + }]) + .await; + let test = test_codex() + .with_model("gpt-5.2") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_request_country(RequestCountry::UNKNOWN) + .with_config(|config| { + config + .features + .enable(Feature::ResponsesWebsocketsV2) + .unwrap(); + }) + .build_with_websocket_server(&server) + .await?; + let warmup = server + .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0) + .await + .body_json(); + assert_eq!(warmup["generate"].as_bool(), Some(false)); + assert_eq!(warmup.get("request_country"), Some(&Value::Null)); + + for request_country in [ + Some(RequestCountry::from_iso_code("DE")?), + Some(RequestCountry::from_iso_code("US")?), + Some(RequestCountry::UNKNOWN), + Some(RequestCountry::default()), + ] { + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "hello".to_string(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + request_country, + ..Default::default() + }), + ) + .await?; + wait_for_successful_turn(&test.codex).await; + } + assert_eq!(server.handshakes().len(), 1); + let requests: Vec<_> = server + .single_connection() + .iter() + .map(core_test_support::responses::WebSocketRequest::body_json) + .collect(); + assert_eq!( + requests + .iter() + .map(|request| request.get("request_country")) + .collect::>(), + [ + Some(&Value::Null), + Some(&json!("DE")), + Some(&json!("US")), + Some(&Value::Null), + None, + ] + ); + assert_eq!(requests[1].get("generate"), None); + assert_eq!( + requests + .iter() + .map(|request| request.get("previous_response_id")) + .collect::>(), + [ + None, + Some(&json!("warm")), + Some(&json!("first")), + Some(&json!("refreshed")), + Some(&json!("cleared")) + ] + ); + server.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn request_country_survives_websocket_to_http_fallback() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = core_test_support::responses::start_mock_server().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/v1/responses")) + .respond_with(wiremock::ResponseTemplate::new(426)) + .expect(1) + .mount(&server) + .await; + let fallback = core_test_support::responses::mount_sse_once( + &server, + core_test_support::responses::sse_completed("fallback"), + ) + .await; + let test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_request_country(RequestCountry::UNKNOWN) + .with_config(|config| { + config + .features + .enable(Feature::ResponsesWebsocketsV2) + .unwrap(); + config.model_provider.supports_websockets = true; + }) + .build_with_auto_env(&server) + .await?; + test.submit_text_turn("Continue over HTTP.").await?; + let request = fallback.single_request(); + assert_eq!( + request.body_json().get("request_country"), + Some(&Value::Null) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn websocket_first_turn_handles_handshake_delay_with_startup_prewarm() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 71ca48832d..b49cccb801 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -5,6 +5,7 @@ use anyhow::Context; use anyhow::Result; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_core::RequestCountry; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_extension_api::ExtensionRegistryBuilder; @@ -312,6 +313,7 @@ async fn remote_compact_v2_retains_metadata_from_resumed_history() -> Result<()> let builder = || test_codex().with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let initial = builder() + .with_request_country(RequestCountry::from_iso_code("DE")?) .with_pre_build_hook(|home| { fs::write( home.join("config.toml"), @@ -334,12 +336,26 @@ async fn remote_compact_v2_retains_metadata_from_resumed_history() -> Result<()> let resumed = builder() .resume(&server, home, rollout_path.clone()) .await?; + resumed + .codex + .update_thread_settings(codex_protocol::protocol::ThreadSettingsOverrides { + request_country: Some(RequestCountry::UNKNOWN), + ..Default::default() + }) + .await?; resumed.codex.submit(Op::Compact).await?; wait_for_turn_complete(&resumed.codex).await; resumed.submit_turn("continue after compaction").await?; resumed.codex.shutdown_and_wait().await?; let requests = response_mock.requests(); + assert_eq!( + requests + .iter() + .map(|request| request.body_json().get("request_country").cloned()) + .collect::>(), + vec![Some(json!("DE")), Some(Value::Null), Some(Value::Null)] + ); let compact_request = &requests[1]; assert_eq!( compact_request.inputs_of_type("compaction_trigger").len(), diff --git a/codex-rs/core/tests/suite/model_request.rs b/codex-rs/core/tests/suite/model_request.rs index fb4131d5aa..019bb969ff 100644 --- a/codex-rs/core/tests/suite/model_request.rs +++ b/codex-rs/core/tests/suite/model_request.rs @@ -1,6 +1,8 @@ //! Output gates allow inference immediately, hold complete output, and cancel with the turn. use anyhow::Result; +use codex_core::RequestCountry; use codex_core::TurnInputRequest; +use codex_core::TurnInputSubmission; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::GuardianAssessment; @@ -11,6 +13,7 @@ use codex_extension_api::ModelResponseInterceptor; use codex_extension_api::ModelResponseStream; use codex_extension_api::ResponseEvent; use codex_features::Feature; +use codex_login::CodexAuth; use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; use codex_protocol::dynamic_tools::DynamicToolResponse; @@ -18,6 +21,9 @@ use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::protocol::TurnSettingsUpdate; +use codex_protocol::protocol::TurnSettingsUpdateOutcome; use codex_protocol::user_input::UserInput; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -142,6 +148,116 @@ impl ModelRequestContributor for BudgetCache { } } +#[tokio::test] +#[test_case::test_case(false; "future settings only")] +#[test_case::test_case(true; "live country update without model switching")] +async fn country_settings_follow_step_and_turn_lifetimes(live_update: bool) -> Result<()> { + const CALL_ID: &str = "country-check"; + let server = responses::start_mock_server().await; + let first = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("first"), + responses::ev_function_call(CALL_ID, "check_status", "{}"), + ev_completed("first"), + ]), + ) + .await; + let test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + config + .features + .disable(Feature::StepModelSwitching) + .expect("disable step model switching"); + }) + .build_with_auto_env(&server) + .await?; + let mut options = test.start_thread_options().await; + options.request_country = RequestCountry::from_iso_code("DE")?; + options.dynamic_tools = vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "check_status".into(), + description: "Check status.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })]; + let thread = test.thread_manager.start_thread(options).await?.thread; + let message = || { + TurnInputRequest::user_input(vec![UserInput::Text { + text: "Check status.".into(), + text_elements: Vec::new(), + }]) + }; + let TurnInputSubmission::Started { turn_id, .. } = + thread.start_or_steer_turn(message()).await? + else { + panic!("expected a new turn"); + }; + wait_for_event( + &thread, + |event| matches!(event, EventMsg::DynamicToolCallRequest(call) if call.call_id == CALL_ID), + ) + .await; + thread + .update_thread_settings(ThreadSettingsOverrides { + request_country: Some(RequestCountry::from_iso_code("US")?), + ..Default::default() + }) + .await?; + if live_update { + let (reply, receiver) = tokio::sync::oneshot::channel(); + thread + .submit(Op::TurnSettings { + turn_id: turn_id.clone(), + update: TurnSettingsUpdate { + request_country: Some(RequestCountry::from_iso_code("JP")?), + ..Default::default() + }, + reply, + }) + .await?; + assert_eq!(receiver.await?, TurnSettingsUpdateOutcome::Applied); + } + let continuation = + responses::mount_sse_once(&server, responses::sse_completed("continued")).await; + thread + .submit(Op::DynamicToolResponse { + id: CALL_ID.into(), + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "ready".into(), + }], + success: true, + }, + }) + .await?; + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + let next = responses::mount_sse_once(&server, responses::sse_completed("next")).await; + thread.start_or_steer_turn(message()).await?; + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + let cleared = responses::mount_sse_once(&server, responses::sse_completed("cleared")).await; + let mut input = message(); + input.thread_settings.request_country = Some(RequestCountry::UNKNOWN); + thread.start_or_steer_turn(input).await?; + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + assert_eq!( + [ + first.single_request(), + continuation.single_request(), + next.single_request(), + cleared.single_request() + ] + .map(|request| request.body_json().get("request_country").cloned()), + [ + Some(json!("DE")), + Some(json!(if live_update { "JP" } else { "DE" })), + Some(json!("US")), + Some(json!(null)) + ], + ); + Ok(()) +} + #[tokio::test] #[test_case::test_case(false; "budget available")] #[test_case::test_case(true; "budget exhausted")] diff --git a/codex-rs/core/tests/suite/reasoning_effort_override.rs b/codex-rs/core/tests/suite/reasoning_effort_override.rs index 305f8f4140..f45beee841 100644 --- a/codex-rs/core/tests/suite/reasoning_effort_override.rs +++ b/codex-rs/core/tests/suite/reasoning_effort_override.rs @@ -3,6 +3,7 @@ use codex_config::test_support::CloudConfigBundleFixture; use codex_core::ForkSnapshot; use codex_core::RecoverTurnRequest; +use codex_core::RequestCountry; use codex_core::StartIfIdleSubmission; use codex_core::StartThreadOptions; use codex_core::SuspendTurnOutcome; @@ -193,11 +194,13 @@ async fn prediction_fork_pins_effort_from_parent_with_overrides_disabled( ) .await; let mut test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_model_info_override("gpt-5.4", |model| { model.use_responses_lite = true; model.supports_reasoning_effort_updates = true; model.default_reasoning_level = Some(ReasoningEffort::Low); }) + .with_request_country(RequestCountry::from_iso_code("US")?) .with_config(|config| { config .features @@ -213,6 +216,7 @@ async fn prediction_fork_pins_effort_from_parent_with_overrides_disabled( &test.codex, ThreadSettingsOverrides { effort: Some(parent_effort), + request_country: Some(RequestCountry::from_iso_code("DE")?), ..Default::default() }, ) @@ -251,6 +255,17 @@ async fn prediction_fork_pins_effort_from_parent_with_overrides_disabled( let requests = mock.requests(); assert_eq!(requests.len(), 3); + assert_eq!( + requests + .iter() + .map(|request| request.body_json().get("request_country").cloned()) + .collect::>(), + vec![ + Some(Value::from("US")), + Some(Value::from("DE")), + Some(Value::from("DE")) + ] + ); assert_eq!(requests[0].body_json()["reasoning"]["effort"], "high"); assert_eq!(requests[1].body_json()["reasoning"]["effort"], "low"); assert!(effort_updates(&requests[1]).is_empty()); diff --git a/codex-rs/core/tests/suite/review.rs b/codex-rs/core/tests/suite/review.rs index 36372c7fd2..3af9090354 100644 --- a/codex-rs/core/tests/suite/review.rs +++ b/codex-rs/core/tests/suite/review.rs @@ -1,5 +1,6 @@ use codex_core::CodexThread; use codex_core::REVIEW_PROMPT; +use codex_core::RequestCountry; use codex_core::TurnInputRequest; use codex_core::config::Config; use codex_core::config::Constrained; @@ -915,6 +916,7 @@ async fn review_resolves_inherited_summary_preferences() -> anyhow::Result<()> { let (server, request_log) = start_responses_server_with_sse(completed_sse(), /*expected_requests*/ 2).await; let test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_model_info_override("gpt-5.2", |model| { model.default_reasoning_summary = ReasoningSummary::Auto; }) @@ -930,17 +932,19 @@ async fn review_resolves_inherited_summary_preferences() -> anyhow::Result<()> { .await?; // First follow the review model's default, then use a preference updated on the thread. - for summary in [None, Some(ReasoningSummary::Concise)] { - if let Some(summary) = summary { - core_test_support::submit_thread_settings( - &test.codex, - ThreadSettingsOverrides { - summary: Some(summary), - ..Default::default() - }, - ) - .await?; - } + for (summary, request_country) in [ + (None, RequestCountry::from_iso_code("DE")?), + (Some(ReasoningSummary::Concise), RequestCountry::UNKNOWN), + ] { + core_test_support::submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + summary, + request_country: Some(request_country), + ..Default::default() + }, + ) + .await?; let stored_settings = test.codex.thread_settings_snapshot().await; test.codex .submit(Op::Review { @@ -958,6 +962,14 @@ async fn review_resolves_inherited_summary_preferences() -> anyhow::Result<()> { .await; assert_eq!(test.codex.thread_settings_snapshot().await, stored_settings); } + assert_eq!( + request_log + .requests() + .iter() + .map(|request| request.body_json().get("request_country").cloned()) + .collect::>(), + vec![Some(serde_json::json!("DE")), Some(serde_json::Value::Null)] + ); let actual = request_log .requests() .iter() diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 820072d363..bb2ff29e2d 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -9,6 +9,7 @@ mod multi_environment_mcp; use super::image_rollout::RecordingFileAttachmentStore; use super::permissions_messages::PermissionPathScenario; use codex_attachment_store::AttachmentStoreErrorKind; +use codex_core::RequestCountry; use codex_protocol_test_support::ContentItemTestExt; use std::collections::HashMap; use std::collections::HashSet; @@ -3773,6 +3774,7 @@ async fn async_compaction_image_checks_do_not_block_turn_completion() -> Result< let test = test_codex() .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_image_store(image_store.clone()) + .with_request_country(RequestCountry::from_iso_code("DE")?) .with_config(move |config| { configure_scenario_catalog(config); config.model_provider.base_url = Some(base_url); @@ -3797,8 +3799,26 @@ async fn async_compaction_image_checks_do_not_block_turn_completion() -> Result< }, ])) .await?; + let turn_id = wait_for_event_match(&test.codex, |event| match event { + EventMsg::TurnStarted(event) => Some(event.turn_id.clone()), + _ => None, + }) + .await; // Release the sampling lookup, then expire the file before its response finishes. barrier.wait().await; + streaming.wait_for_request_count(/*count*/ 1).await; + let (reply, outcome) = oneshot::channel(); + test.codex + .submit(Op::TurnSettings { + turn_id, + update: TurnSettingsUpdate { + request_country: Some(RequestCountry::from_iso_code("US").unwrap()), + ..Default::default() + }, + reply, + }) + .await?; + assert_eq!(outcome.await?, TurnSettingsUpdateOutcome::Applied); image_store .resolve_errors .lock() @@ -3814,6 +3834,12 @@ async fn async_compaction_image_checks_do_not_block_turn_completion() -> Result< .await; assert_eq!(streaming.requests().await.len(), 1); + test.codex + .update_thread_settings(codex_protocol::protocol::ThreadSettingsOverrides { + request_country: Some(RequestCountry::from_iso_code("JP").unwrap()), + ..Default::default() + }) + .await?; // The turn completed without releasing the optional compaction lookup. barrier.wait().await; tokio::time::timeout( @@ -3828,6 +3854,17 @@ async fn async_compaction_image_checks_do_not_block_turn_completion() -> Result< .map(|body| serde_json::from_slice::(body)) .collect::>>()?; assert_eq!(requests.len(), 2); + // Compaction keeps the sampling step's country despite live and future settings updates. + assert_eq!( + requests + .iter() + .map(|request| request.get("request_country")) + .collect::>(), + vec![ + Some(&serde_json::json!("DE")), + Some(&serde_json::json!("DE")) + ] + ); assert_eq!( *image_store.resolves.lock().expect("resolves lock"), HashMap::from([("file_expiring".to_string(), 2)]) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index 7ad1faf6dd..495b153b4c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -225,6 +225,7 @@ impl LunaSampler { auth_owner: Option, ) -> execution::SamplingExecution { let api_request = ResponsesApiRequest { + request_country: Default::default(), model: MODEL.to_owned(), input, tools: None, diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs index 3c725f1c5d..a92d5cfaa7 100644 --- a/codex-rs/protocol/src/lib.rs +++ b/codex-rs/protocol/src/lib.rs @@ -3,6 +3,9 @@ mod agent_path; pub mod auth; mod context_window_id; mod item_id_generator; +mod request_country; +pub use request_country::InvalidRequestCountry; +pub use request_country::RequestCountry; mod response_item_id; mod response_usage; mod sanitized_git_url; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 5d05de3864..97bbc7a561 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -560,6 +560,8 @@ pub struct ConversationSpeechParams { /// task kind. Child sessions and consumers of frozen initial settings are unchanged. #[derive(Debug, Clone, Default, PartialEq)] pub struct TurnSettingsUpdate { + /// Trusted host country. `None` preserves the current setting. Published for the named live turn's subsequent steps. + pub request_country: Option, /// Changes the reviewer for subsequent approval requests, not pending reviews. pub approvals_reviewer: Option, /// Replaces the selection for subsequent steps, without changing future turns. @@ -601,6 +603,8 @@ pub enum TurnSettingsUpdateOutcome { }), )] pub struct ThreadSettingsOverrides { + /// Trusted host country. `None` preserves the current setting. Inherited by future turns on acceptance. + pub request_country: Option, /// Updated fallback `cwd` and environments supplied together as a complete pair. pub environments: Option, diff --git a/codex-rs/protocol/src/request_country.rs b/codex-rs/protocol/src/request_country.rs new file mode 100644 index 0000000000..d2201e64d6 --- /dev/null +++ b/codex-rs/protocol/src/request_country.rs @@ -0,0 +1,59 @@ +//! Trusted runtime country for the current Responses request. + +use serde::Serialize; +use serde::Serializer; + +/// Country selected by a trusted host. The default leaves provenance to the server. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RequestCountry(Country); + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +enum Country { + #[default] + Omitted, + Unknown, + Known(String), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +#[error("request country must contain exactly two ASCII letters")] +pub struct InvalidRequestCountry; + +impl RequestCountry { + /// Explicitly unknown country. The server must not fall back to request provenance. + pub const UNKNOWN: Self = Self(Country::Unknown); + + /// Normalizes a two-letter country code. The special code `XX` is explicit unknown. + pub fn from_iso_code(code: &str) -> Result { + let code = code.trim(); + if code.len() != 2 || !code.bytes().all(|byte| byte.is_ascii_alphabetic()) { + return Err(InvalidRequestCountry); + } + let code = code.to_ascii_uppercase(); + if code == "XX" { + return Ok(Self::UNKNOWN); + } + Ok(Self(Country::Known(code))) + } + + pub fn known_country(&self) -> Option<&str> { + match &self.0 { + Country::Known(code) => Some(code), + Country::Omitted | Country::Unknown => None, + } + } + + pub fn is_omitted(&self) -> bool { + matches!(self.0, Country::Omitted) + } +} + +impl Serialize for RequestCountry { + fn serialize(&self, serializer: S) -> Result { + self.known_country().serialize(serializer) + } +} + +#[cfg(test)] +#[path = "request_country_tests.rs"] +mod tests; diff --git a/codex-rs/protocol/src/request_country_tests.rs b/codex-rs/protocol/src/request_country_tests.rs new file mode 100644 index 0000000000..58d5cf525a --- /dev/null +++ b/codex-rs/protocol/src/request_country_tests.rs @@ -0,0 +1,13 @@ +#![cfg(test)] + +use super::*; + +#[test] +fn request_country_rejects_malformed_values() { + for malformed in ["", "D", "USA", "D1", "T1", "DÉ", "D E", "DE\nX"] { + assert_eq!( + RequestCountry::from_iso_code(malformed), + Err(InvalidRequestCountry), + ); + } +}