From 535795f7d12495ee6ea18db3201d2f85c3606aef Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Fri, 14 Aug 2026 00:18:00 +0000 Subject: [PATCH] Centralize turn environment selection state (#38461) ## What changed - Store `TurnEnvironmentSelection` directly on each resolved `TurnEnvironment` instead of copying its environment ID, working directory, and workspace roots into separate fields. - Update environment resolution, tool execution, approvals, and related tests to access selection state through the retained value. GitOrigin-RevId: ab723d21e112b8af5bcaf1f772a5e376553312bc --- codex-rs/core/src/agents_md.rs | 4 +- codex-rs/core/src/agents_md_tests.rs | 9 ++- .../src/context/world_state/environment.rs | 2 +- codex-rs/core/src/environment_selection.rs | 27 ++++---- codex-rs/core/src/guardian/tests.rs | 16 ++--- codex-rs/core/src/mcp_openai_file.rs | 15 ++--- codex-rs/core/src/mcp_tool_call.rs | 2 +- codex-rs/core/src/mcp_tool_call_tests.rs | 9 ++- codex-rs/core/src/session/mcp.rs | 2 +- codex-rs/core/src/session/mcp_runtime.rs | 2 +- codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/tests.rs | 63 ++++++++++++------- codex-rs/core/src/session/turn.rs | 4 +- codex-rs/core/src/session/turn_context.rs | 30 +++------ .../core/src/tools/handlers/apply_patch.rs | 4 +- .../src/tools/handlers/extension_tools.rs | 4 +- codex-rs/core/src/tools/handlers/mod.rs | 2 +- codex-rs/core/src/tools/handlers/shell.rs | 2 +- .../src/tools/handlers/shell/shell_command.rs | 2 +- .../core/src/tools/handlers/shell_tests.rs | 9 ++- .../handlers/unified_exec/exec_command.rs | 8 +-- .../core/src/tools/handlers/view_image.rs | 7 ++- .../tools/handlers/wait_for_environment.rs | 2 +- codex-rs/core/src/tools/network_approval.rs | 6 +- .../core/src/tools/runtimes/apply_patch.rs | 2 +- .../src/tools/runtimes/apply_patch_tests.rs | 9 ++- codex-rs/core/src/tools/runtimes/shell.rs | 6 +- .../tools/runtimes/shell/unix_escalation.rs | 6 +- .../core/src/tools/runtimes/shell_tests.rs | 11 ++-- .../core/src/tools/runtimes/unified_exec.rs | 25 ++++---- codex-rs/core/src/tools/spec_plan_tests.rs | 13 ++-- 31 files changed, 159 insertions(+), 146 deletions(-) diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index ce5cc47a66..bd9e433721 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -65,7 +65,7 @@ pub(crate) async fn load_project_instructions( match read_agents_md( config, filesystem.as_ref(), - &turn_environment.environment_id, + &turn_environment.selection.environment_id, turn_environment.cwd(), remaining, ) @@ -80,7 +80,7 @@ pub(crate) async fn load_project_instructions( Ok(None) => {} Err(e) => { error!( - environment_id = turn_environment.environment_id, + environment_id = turn_environment.selection.environment_id, "error trying to find AGENTS.md docs: {e:#}" ); } diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index 9deefd61e4..53cf15a838 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -23,6 +23,7 @@ use codex_exec_server::RemoveOptions; use codex_extension_api::UserInstructions; use codex_features::Feature; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use core_test_support::PathBufExt; @@ -332,13 +333,15 @@ fn resolved_local_environments( .into_iter() .map(|(environment_id, cwd)| { TurnEnvironmentState::Ready(TurnEnvironment::new( - environment_id.to_string(), + TurnEnvironmentSelection { + environment_id: environment_id.to_string(), + cwd: PathUri::from_abs_path(&cwd), + workspace_roots: Vec::new(), + }, Arc::new( Environment::create_for_tests(/*exec_server_url*/ None) .expect("local environment"), ), - PathUri::from_abs_path(&cwd), - Vec::new(), /*shell*/ None, TurnEnvironmentConfig { allow_login_shell: true, diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index 6e5aa9908a..e7dffc1b31 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -342,7 +342,7 @@ fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap { let mut turn_environment = TurnEnvironment::new( - starting.selection.environment_id, + starting.selection, environment.environment, - starting.selection.cwd, - starting.selection.workspace_roots, environment.shell, starting.config, ); @@ -578,7 +576,7 @@ impl TurnEnvironmentSnapshot { self.turn_environments() .map(|environment| { ( - environment.environment_id.clone(), + environment.selection.environment_id.clone(), Some(Arc::clone(&environment.environment)), ) }) @@ -932,6 +930,7 @@ url = "ws://127.0.0.1:8765" resolved .primary() .expect("primary environment") + .selection .environment_id, "local" ); @@ -1344,10 +1343,8 @@ url = "ws://127.0.0.1:8765" .expect("inherited environment"), ); let inherited = TurnEnvironment::new( - selection.environment_id.clone(), + selection.clone(), Arc::clone(&inherited_environment), - selection.cwd.clone(), - Vec::new(), /*shell*/ None, test_environment_config(), ); @@ -1495,10 +1492,12 @@ url = "ws://127.0.0.1:8765" ); let remote = TurnEnvironmentSnapshot { environments: vec![TurnEnvironmentState::Ready(TurnEnvironment::new( - REMOTE_ENVIRONMENT_ID.to_string(), + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd_uri.clone(), + workspace_roots: Vec::new(), + }, remote_environment.clone(), - remote_cwd_uri.clone(), - Vec::new(), /*shell*/ None, test_environment_config(), ))], @@ -1506,10 +1505,12 @@ url = "ws://127.0.0.1:8765" let multiple = TurnEnvironmentSnapshot { environments: vec![ TurnEnvironmentState::Ready(TurnEnvironment::new( - REMOTE_ENVIRONMENT_ID.to_string(), + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd_uri, + workspace_roots: Vec::new(), + }, remote_environment, - remote_cwd_uri, - Vec::new(), /*shell*/ None, test_environment_config(), )), diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index b6a83ea449..7a8815db74 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -15,7 +15,6 @@ use crate::guardian::review::guardian_review_session_config; use crate::guardian::review::routes_approval_to_guardian_with_reviewer; use crate::session::session::Session; use crate::session::turn_context::TurnContext; -use crate::session::turn_context::TurnEnvironment; use crate::test_support; use codex_analytics::GuardianApprovalRequestSource; use codex_config::ConfigLayerStack; @@ -631,17 +630,10 @@ async fn build_guardian_prompt_includes_parent_turn_denied_reads() -> anyhow::Re }; environment.config.permission_profile = PermissionProfileSnapshot::legacy(environment_permission_profile); - *environment = TurnEnvironment::new( - environment.environment_id.clone(), - Arc::clone(&environment.environment), - environment.cwd().clone(), - vec![ - PathUri::from_abs_path(&workspace_root), - PathUri::from_abs_path(&second_workspace_root), - ], - environment.shell.clone(), - environment.config.clone(), - ); + environment.selection.workspace_roots = vec![ + PathUri::from_abs_path(&workspace_root), + PathUri::from_abs_path(&second_workspace_root), + ]; let session = Arc::new(session); let turn = Arc::new(turn); seed_guardian_parent_history(&session, &turn).await; diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index f86a740a5f..c6f3d37eb1 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -161,10 +161,10 @@ async fn build_uploaded_argument_value( .join(file_path) .map_err(|error| contextualize_error(error.to_string()))?; let additional_permissions = merge_permission_profiles( - sess.granted_session_permissions(&turn_environment.environment_id) + sess.granted_session_permissions(&turn_environment.selection.environment_id) .await .as_ref(), - sess.granted_turn_permissions(&turn_environment.environment_id) + sess.granted_turn_permissions(&turn_environment.selection.environment_id) .await .as_ref(), ); @@ -271,7 +271,6 @@ mod tests { use crate::environment_selection::TurnEnvironmentState; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; - use crate::session::turn_context::TurnEnvironment; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; @@ -285,14 +284,8 @@ mod tests { else { panic!("expected ready primary environment"); }; - *primary = TurnEnvironment::new( - primary.environment_id.clone(), - Arc::clone(&primary.environment), - PathUri::from_abs_path(&cwd), - Vec::new(), - primary.shell.clone(), - primary.config.clone(), - ); + primary.selection.cwd = PathUri::from_abs_path(&cwd); + primary.selection.workspace_roots.clear(); } #[tokio::test] diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 7e408810e3..56229b3e87 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -802,7 +802,7 @@ fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) if let Some(environment) = step_context .environments .turn_environments() - .find(|environment| environment.environment_id == environment_id) + .find(|environment| environment.selection.environment_id == environment_id) { return Some(environment.cwd().clone()); } diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 71a6653d1f..f2844830ad 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -29,6 +29,7 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpInvocation; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_path_uri::PathUri; use core_test_support::hooks::trusted_config_layer_stack; use core_test_support::responses::ev_assistant_message; @@ -1106,10 +1107,12 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul .environments .environments .push(TurnEnvironmentState::Ready(TurnEnvironment::new( - "remote".to_string(), + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd: secondary_cwd.clone(), + workspace_roots: Vec::new(), + }, environment, - secondary_cwd.clone(), - Vec::new(), /*shell*/ None, TurnEnvironmentConfig { allow_login_shell: true, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 89e75623fa..1690591239 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -392,7 +392,7 @@ impl Session { sandbox.windows_sandbox_private_desktop = config.permissions.windows_sandbox_private_desktop; sandbox.use_legacy_landlock = config.features.use_legacy_landlock(); - (environment.environment_id.clone(), sandbox) + (environment.selection.environment_id.clone(), sandbox) }) .collect::>() } else { diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index d106c4f2aa..580c83d040 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -170,7 +170,7 @@ impl Session { .turn_environments() .map(|environment| { ( - environment.environment_id.clone(), + environment.selection.environment_id.clone(), environment.cwd().clone(), ) }) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index de3fe88441..8cc3fe92a4 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3160,7 +3160,7 @@ impl Session { .turn_environments() .map(|environment| { ( - environment.environment_id.clone(), + environment.selection.environment_id.clone(), turn_context.file_system_sandbox_context( /*additional_permissions*/ None, environment, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0d231921c2..33ae7f1c9b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6640,10 +6640,12 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir .clone(); turn_context_mut.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( - "remote".to_string(), + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd: PathUri::from_abs_path(&environment_cwd), + workspace_roots: Vec::new(), + }, current_environment.environment, - PathUri::from_abs_path(&environment_cwd), - Vec::new(), current_environment.shell, current_environment.config, )); @@ -7237,10 +7239,12 @@ async fn primary_environment_uses_first_turn_environment() { .environments .environments .push(TurnEnvironmentState::Ready(TurnEnvironment::new( - "second".to_string(), + TurnEnvironmentSelection { + environment_id: "second".to_string(), + cwd: second_cwd_uri.clone(), + workspace_roots: Vec::new(), + }, Arc::clone(&first_environment.environment), - second_cwd_uri.clone(), - Vec::new(), /*shell*/ None, first_environment.config.clone(), ))); @@ -7250,14 +7254,15 @@ async fn primary_environment_uses_first_turn_environment() { .environments .primary() .expect("primary environment") + .selection .environment_id, - first_environment.environment_id + first_environment.selection.environment_id ); assert_eq!( turn_context .environments .turn_environments() - .find(|environment| environment.environment_id == "second") + .find(|environment| environment.selection.environment_id == "second") .expect("second environment") .cwd(), &second_cwd_uri @@ -8645,10 +8650,12 @@ async fn conflicting_ready_environment_root_ids_keep_first_location() { ) .expect("report environment ready"); turn_environments.push(TurnEnvironment::new( - environment_id.clone(), + TurnEnvironmentSelection { + environment_id: environment_id.clone(), + cwd: local_environment.cwd().clone(), + workspace_roots: local_environment.workspace_roots().to_vec(), + }, environment, - local_environment.cwd().clone(), - local_environment.workspace_roots().to_vec(), local_environment.shell.clone(), local_environment.config.clone(), )); @@ -8710,7 +8717,7 @@ async fn capability_discovery_uses_environment_permission_profile() { )); let expected_sandbox = turn_context .file_system_sandbox_context(/*additional_permissions*/ None, &environment); - let environment_id = environment.environment_id.clone(); + let environment_id = environment.selection.environment_id.clone(); turn_context.environments.environments[0] = TurnEnvironmentState::Ready(environment); let discovery = session @@ -8902,10 +8909,12 @@ async fn record_context_updates_emits_environment_item_for_cwd_changes() { .clone(); current_context.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( - environment.environment_id, + TurnEnvironmentSelection { + environment_id: environment.selection.environment_id, + cwd: PathUri::from_abs_path(&cwd), + workspace_roots: Vec::new(), + }, environment.environment, - PathUri::from_abs_path(&cwd), - Vec::new(), environment.shell, environment.config, )); @@ -8959,10 +8968,12 @@ async fn record_context_updates_use_environment_permission_profile_and_workspace PermissionProfileSnapshot::legacy(PermissionProfile::workspace_write()); current_context.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( - environment.environment_id, + TurnEnvironmentSelection { + environment_id: environment.selection.environment_id, + cwd, + workspace_roots: vec![PathUri::from_abs_path(&workspace_root)], + }, environment.environment, - cwd, - vec![PathUri::from_abs_path(&workspace_root)], environment.shell, environment_config, )); @@ -9032,10 +9043,12 @@ async fn record_context_updates_omits_environment_item_when_disabled() { .clone(); current_context.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( - environment.environment_id, + TurnEnvironmentSelection { + environment_id: environment.selection.environment_id, + cwd: PathUri::from_abs_path(&test_path_buf("/new-repo").abs()), + workspace_roots: Vec::new(), + }, environment.environment, - PathUri::from_abs_path(&test_path_buf("/new-repo").abs()), - Vec::new(), environment.shell, environment.config, )); @@ -9498,10 +9511,12 @@ async fn turn_context_item_stores_local_cwd() { .clone(); let cwd = PathUri::parse("file:///C:/windows").expect("Windows cwd URI"); turn_context.environments.environments[0] = TurnEnvironmentState::Ready(TurnEnvironment::new( - "remote".to_string(), + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd, + workspace_roots: Vec::new(), + }, environment.environment, - cwd, - Vec::new(), environment.shell, environment.config, )); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2cf00a015a..2422ba0766 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -588,7 +588,7 @@ async fn turn_diff_display_roots(step_context: &StepContext) -> Vec<(String, Pat .ok() .flatten() .unwrap_or_else(|| cwd.clone()); - display_roots.push((turn_environment.environment_id.clone(), root)); + display_roots.push((turn_environment.selection.environment_id.clone(), root)); } display_roots } @@ -897,7 +897,7 @@ async fn build_extension_turn_input_items( .turn_environments() .enumerate() .map(|(index, environment)| TurnInputEnvironment { - environment_id: environment.environment_id.clone(), + environment_id: environment.selection.environment_id.clone(), cwd: environment.cwd().clone(), is_primary: index == 0, }) diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index c5f25854c1..2a968e65af 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -42,10 +42,8 @@ pub(crate) struct TurnEnvironmentConfig { #[derive(Clone)] pub(crate) struct TurnEnvironment { - pub(crate) environment_id: String, + pub(crate) selection: TurnEnvironmentSelection, pub(crate) environment: Arc, - cwd: PathUri, - workspace_roots: Vec, pub(crate) shell: Option, pub(crate) config: TurnEnvironmentConfig, pub(crate) shell_snapshot: ShellSnapshotTask, @@ -53,18 +51,14 @@ pub(crate) struct TurnEnvironment { impl TurnEnvironment { pub(crate) fn new( - environment_id: String, + selection: TurnEnvironmentSelection, environment: Arc, - cwd: PathUri, - workspace_roots: Vec, shell: Option, config: TurnEnvironmentConfig, ) -> Self { Self { - environment_id, + selection, environment, - cwd, - workspace_roots, shell, config, shell_snapshot: futures::future::ready(None).boxed().shared(), @@ -72,7 +66,7 @@ impl TurnEnvironment { } pub(crate) fn shell_snapshot(&self, cwd: &AbsolutePathBuf) -> Option { - if self.cwd != PathUri::from_abs_path(cwd) { + if self.selection.cwd != PathUri::from_abs_path(cwd) { return None; } self.shell_snapshot @@ -82,11 +76,11 @@ impl TurnEnvironment { } pub(crate) fn cwd(&self) -> &PathUri { - &self.cwd + &self.selection.cwd } pub(crate) fn workspace_roots(&self) -> &[PathUri] { - &self.workspace_roots + &self.selection.workspace_roots } pub(crate) fn permission_profile(&self) -> &PermissionProfile { @@ -109,21 +103,17 @@ impl TurnEnvironment { } pub(crate) fn selection(&self) -> TurnEnvironmentSelection { - TurnEnvironmentSelection { - environment_id: self.environment_id.clone(), - cwd: self.cwd.clone(), - workspace_roots: self.workspace_roots.clone(), - } + self.selection.clone() } } impl std::fmt::Debug for TurnEnvironment { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("TurnEnvironment") - .field("environment_id", &self.environment_id) + .field("environment_id", &self.selection.environment_id) .field("environment", &self.environment) - .field("cwd", &self.cwd) - .field("workspace_roots", &self.workspace_roots) + .field("cwd", &self.selection.cwd) + .field("workspace_roots", &self.selection.workspace_roots) .field("shell", &self.shell) .field("config", &self.config) .finish_non_exhaustive() diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index a739e0e338..36ae51ba7b 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -282,7 +282,7 @@ async fn effective_patch_permissions( crate::tools::handlers::EffectiveAdditionalPermissions, codex_protocol::permissions::FileSystemSandboxPolicy, )> { - let environment_id = environment.environment_id.as_str(); + let environment_id = environment.selection.environment_id.as_str(); let file_paths = file_paths_for_action(action); let native_cwd = cwd.to_abs_path()?; let granted_permissions = merge_permission_profiles( @@ -565,7 +565,7 @@ async fn execute_verified_patch( let emitter = ToolEmitter::apply_patch_for_environment( changes.clone(), apply.auto_approved, - turn_environment.environment_id.clone(), + turn_environment.selection.environment_id.clone(), ); let event_ctx = ToolEventCtx::new( tool_ctx.session.as_ref(), diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index af137cb777..87dc39563d 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -150,7 +150,7 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { }; let additional_permissions = apply_granted_turn_permissions( invocation.session.as_ref(), - &environment.environment_id, + &environment.selection.environment_id, native_cwd.as_path(), SandboxPermissions::UseDefault, /*additional_permissions*/ None, @@ -161,7 +161,7 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { .turn .file_system_sandbox_context(additional_permissions, environment); environments.push(ToolEnvironment { - environment_id: environment.environment_id.clone(), + environment_id: environment.selection.environment_id.clone(), cwd: native_cwd, file_system: environment.environment.get_filesystem(), file_system_sandbox_context, diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index f64706e868..69af0c8152 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -175,7 +175,7 @@ fn resolve_tool_environment<'a>( |environment_id| { environments .turn_environments() - .find(|environment| environment.environment_id == environment_id) + .find(|environment| environment.selection.environment_id == environment_id) .map(Some) .ok_or_else(|| { FunctionCallError::RespondToModel(format!( diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 7fdd9a42b2..bde82e3dbb 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -95,7 +95,7 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result for WaitForEnvironmentHandler { let already_ready = step_context .environments .turn_environments() - .any(|environment| environment.environment_id == environment_id); + .any(|environment| environment.selection.environment_id == environment_id); if !already_ready { let Some(environment) = step_context .environments diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index ca6630916d..c41f863592 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -604,7 +604,7 @@ impl NetworkApprovalService { active_turn .as_ref() .and_then(|(turn_context, _)| turn_context.environments.primary()) - .map(|environment| environment.environment_id.clone()) + .map(|environment| environment.selection.environment_id.clone()) }) else { return NetworkDecision::deny(REASON_NOT_ALLOWED); }; @@ -666,7 +666,7 @@ impl NetworkApprovalService { turn_context .environments .turn_environments() - .find(|environment| environment.environment_id == environment_id) + .find(|environment| environment.selection.environment_id == environment_id) .map(TurnEnvironment::permission_profile) }); if !permission_profile.is_some_and(permission_profile_allows_network_approval_flow) { @@ -705,7 +705,7 @@ impl NetworkApprovalService { turn_context .environments .turn_environments() - .find(|environment| environment.environment_id == environment_id) + .find(|environment| environment.selection.environment_id == environment_id) .and_then(|environment| environment.cwd().to_abs_path().ok()) .unwrap_or_else(|| { #[allow(deprecated)] diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index ab4a5c7b07..bedb63fbca 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -74,7 +74,7 @@ impl ApplyPatchRuntime { fn build_approval_action(req: &ApplyPatchRequest, call_id: &str) -> ApprovalAction { ApprovalAction::ApplyPatch { id: call_id.to_string(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), cwd: req.action.cwd.clone(), files: req.file_paths.clone(), patch: req.action.patch.clone(), diff --git a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs index fa9fbc0304..5be81de54e 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -8,6 +8,7 @@ use codex_protocol::models::FileSystemPermissions; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::GranularApprovalConfig; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_sandboxing::SandboxManager; use codex_sandboxing::SandboxType; use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; @@ -18,10 +19,12 @@ use pretty_assertions::assert_eq; use std::collections::HashMap; fn test_turn_environment(environment_id: &str) -> crate::session::turn_context::TurnEnvironment { crate::session::turn_context::TurnEnvironment::new( - environment_id.to_string(), + TurnEnvironmentSelection { + environment_id: environment_id.to_string(), + cwd: PathUri::from_abs_path(&std::env::temp_dir().abs()), + workspace_roots: Vec::new(), + }, std::sync::Arc::new(codex_exec_server::Environment::default_for_tests()), - PathUri::from_abs_path(&std::env::temp_dir().abs()), - Vec::new(), /*shell*/ None, TurnEnvironmentConfig { allow_login_shell: true, diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 623bafb536..95c44c7382 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -129,7 +129,7 @@ impl Approvable for ShellRuntime { ) -> std::io::Result { Ok(ApprovalAction::Shell { id: call_id.to_string(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), command: req.command.clone(), hook_command: req.hook_command.clone(), cwd: PathUri::from_abs_path(&req.cwd), @@ -186,7 +186,7 @@ impl ToolRuntime for ShellRuntime { tty: None, }, command: req.hook_command.clone(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), permission_profile: req.turn_environment.permission_profile().clone(), }) } @@ -307,7 +307,7 @@ impl ToolRuntime for ShellRuntime { command, options, managed_network, - Some(&req.turn_environment.environment_id), + Some(&req.turn_environment.selection.environment_id), ) .map_err(ToolError::Codex)?; execute_env(env, Self::stdout_stream(ctx)) diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 17dcca68e5..499e054995 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -149,7 +149,7 @@ pub(super) async fn try_run_zsh_fork( command, options, managed_network_for_sandbox_permissions(req.network.as_ref(), req.sandbox_permissions), - Some(&req.turn_environment.environment_id), + Some(&req.turn_environment.selection.environment_id), ) .map_err(ToolError::Codex)?; let crate::sandboxing::ExecRequest { @@ -244,7 +244,7 @@ pub(super) async fn try_run_zsh_fork( session: Arc::clone(&ctx.session), review_context: GuardianReviewContext::from(&ctx.step_context), call_id: ctx.call_id.clone(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), source: GuardianCommandSource::Shell, tool_name: ctx.tool_name.clone(), approval_policy: ctx.step_context.turn.approval_policy(), @@ -331,7 +331,7 @@ pub(crate) async fn prepare_unified_exec_zsh_fork( session: Arc::clone(&ctx.session), review_context: GuardianReviewContext::from(&ctx.step_context), call_id: ctx.call_id.clone(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), source: GuardianCommandSource::UnifiedExec, tool_name: ctx.tool_name.clone(), approval_policy: ctx.step_context.turn.approval_policy(), diff --git a/codex-rs/core/src/tools/runtimes/shell_tests.rs b/codex-rs/core/src/tools/runtimes/shell_tests.rs index 6b923af88d..2b84524515 100644 --- a/codex-rs/core/src/tools/runtimes/shell_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell_tests.rs @@ -4,6 +4,7 @@ use crate::session::turn_context::TurnEnvironmentConfig; use crate::tools::approvals::ApprovalCacheKey; use codex_exec_server::Environment; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; use std::sync::Arc; @@ -15,10 +16,12 @@ async fn approval_key_uses_path_uri_and_includes_environment_id() { let mut request = ShellRequest { command: vec!["echo".to_string(), "hello".to_string()], turn_environment: TurnEnvironment::new( - "remote".to_string(), + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd: PathUri::from_abs_path(&cwd), + workspace_roots: Vec::new(), + }, Arc::new(Environment::default_for_tests()), - PathUri::from_abs_path(&cwd), - Vec::new(), /*shell*/ None, TurnEnvironmentConfig { allow_login_shell: true, @@ -61,7 +64,7 @@ async fn approval_key_uses_path_uri_and_includes_environment_id() { additional_permissions: request.additional_permissions.clone(), })] ); - request.turn_environment.environment_id = "other".to_string(); + request.turn_environment.selection.environment_id = "other".to_string(); let other_key = runtime .approval_action(&request, "call-1") .expect("build approval action") diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 94829bd819..12713ae69c 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -168,7 +168,7 @@ impl Approvable for UnifiedExecRuntime<'_> { ) -> std::io::Result { Ok(ApprovalAction::ExecCommand { id: call_id.to_string(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), command: req.command.clone(), hook_command: req.hook_command.clone(), cwd: req.cwd.clone(), @@ -237,7 +237,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt tty: Some(req.tty), }, command: req.hook_command.clone(), - environment_id: req.turn_environment.environment_id.clone(), + environment_id: req.turn_environment.selection.environment_id.clone(), permission_profile: req.turn_environment.permission_profile().clone(), }) } @@ -325,12 +325,12 @@ impl<'a> ToolRuntime for UnifiedExecRunt let prepared = network .prepare_for_optional_environment( env, - Some(&req.turn_environment.environment_id), + Some(&req.turn_environment.selection.environment_id), ) .map_err(|err| { ToolError::Codex(CodexErr::Io(io::Error::other(format!( "failed to prepare network proxy for environment `{}`: {err}", - req.turn_environment.environment_id + req.turn_environment.selection.environment_id )))) })?; (prepared.env, Some(prepared.sandbox_context), None) @@ -419,7 +419,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt command, options, managed_network, - Some(&req.turn_environment.environment_id), + Some(&req.turn_environment.selection.environment_id), ) .map_err(ToolError::Codex)?; exec_env.exec_server_env_config = req.exec_server_env_config.clone(); @@ -495,7 +495,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt attempt, managed_network, network_proxy_launch, - /*environment_id*/ Some(&req.turn_environment.environment_id), + /*environment_id*/ Some(&req.turn_environment.selection.environment_id), req.exec_server_env_config.clone(), windows_sandbox_proxy_settings_mode, req.tty, @@ -520,6 +520,7 @@ mod tests { use codex_exec_server::Environment; use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_protocol::models::PermissionProfile; + use codex_protocol::protocol::TurnEnvironmentSelection; use codex_tools::ZshForkConfig; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; @@ -530,10 +531,12 @@ mod tests { fn test_turn_environment(cwd: PathUri) -> TurnEnvironment { TurnEnvironment::new( - LOCAL_ENVIRONMENT_ID.to_string(), + TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd, + workspace_roots: Vec::new(), + }, Arc::new(Environment::default_for_tests()), - cwd, - Vec::new(), /*shell*/ None, TurnEnvironmentConfig { allow_login_shell: true, @@ -578,12 +581,12 @@ mod tests { proposed_execpolicy_amendment: None, }, ); - request.turn_environment.environment_id = "remote".to_string(); + request.turn_environment.selection.environment_id = "remote".to_string(); let original_key = runtime .approval_action(&request, "call-1") .expect("build approval action") .cache_keys(); - request.turn_environment.environment_id = "other".to_string(); + request.turn_environment.selection.environment_id = "other".to_string(); let other_key = runtime .approval_action(&request, "call-1") .expect("build approval action") diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 0d2aef0357..76606317d4 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -24,6 +24,7 @@ use codex_protocol::openai_models::WebSearchToolType; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_tools::DiscoverablePluginInfo; use codex_tools::DiscoverableTool; use codex_tools::ResponsesApiNamespaceTool; @@ -373,7 +374,7 @@ fn duplicate_primary_environment(turn: &mut TurnContext) { .primary() .expect("primary environment") .clone(); - second_environment.environment_id = "secondary".to_string(); + second_environment.selection.environment_id = "secondary".to_string(); turn.environments .environments .push(TurnEnvironmentState::Ready(second_environment)); @@ -752,7 +753,7 @@ async fn shell_command_is_not_registered_without_a_single_local_environment() { else { panic!("primary environment should be ready"); }; - environment.environment_id = "remote".to_string(); + environment.selection.environment_id = "remote".to_string(); environment.environment = Arc::new( codex_exec_server::Environment::create_for_tests(Some( "ws://127.0.0.1:1/remote-exec-server".to_string(), @@ -902,15 +903,17 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava .environments .push(TurnEnvironmentState::Ready( crate::session::turn_context::TurnEnvironment::new( - "remote".to_string(), + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd: remote_cwd, + workspace_roots: Vec::new(), + }, Arc::new( codex_exec_server::Environment::create_for_tests(Some( "ws://127.0.0.1:1/remote-exec-server".to_string(), )) .expect("remote test environment"), ), - remote_cwd, - Vec::new(), /*shell*/ None, crate::session::turn_context::TurnEnvironmentConfig { allow_login_shell: true,