From 78421face0047c16be2d8a6aa2b73e73b9a3107f Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 12:12:03 -0700 Subject: [PATCH] Route process tools to selected environments (#20647) ## Why When a turn exposes multiple selected environments, shell-style tools need a model-facing way to identify the intended target environment and handlers need to resolve that target before parsing cwd-relative permission fields or launching processes. This PR scopes that rollout to process tools. Filesystem-oriented tools such as `apply_patch`, `view_image`, and `list_dir` are intentionally left for follow-up slices. ## What Changed - Adds an `include_environment_id` option to shell-style tool schema builders. - Exposes optional `environment_id` on `shell`, `shell_command`, and `exec_command` only when `ToolEnvironmentMode::Multiple` is active. - Adds a shared handler helper that parses `environment_id` and `workdir` from JSON function-call arguments and returns the selected `Environment` plus effective absolute cwd. - Uses that helper in `shell`, `shell_command`, and `exec_command` handling so process execution uses the selected environment filesystem and cwd. - Changes `ExecCommandRequest` to carry a required resolved `cwd`, removing the process-manager fallback to the primary turn cwd for new exec commands. - Leaves `write_stdin` unchanged because it targets an existing process id, not a new environment. ## Testing - Added unit coverage for process-tool schema exposure, selected environment resolution, primary fallback, no-environment handling, unknown environment ids, and resolving cwd-relative permission paths against the selected environment cwd. - Added a remote-suite e2e coverage case for `exec_command` routing across explicit zero environments, one local environment, and local+remote environments. - Ran `just fmt` and `git diff --check`. --------- Co-authored-by: Codex --- codex-rs/core/src/tools/handlers/mod.rs | 23 +++ .../core/src/tools/handlers/unified_exec.rs | 53 ++++--- .../src/tools/handlers/unified_exec_tests.rs | 40 ----- codex-rs/core/src/tools/orchestrator.rs | 5 +- .../core/src/tools/runtimes/unified_exec.rs | 32 ++-- codex-rs/core/src/tools/sandboxing.rs | 4 + codex-rs/core/src/unified_exec/mod.rs | 4 +- .../core/src/unified_exec/process_manager.rs | 8 +- .../src/unified_exec/process_manager_tests.rs | 6 +- codex-rs/core/tests/suite/remote_env.rs | 137 ++++++++++++++++++ codex-rs/tools/src/local_tool.rs | 15 ++ codex-rs/tools/src/tool_registry_plan.rs | 16 +- .../tools/src/tool_registry_plan_tests.rs | 61 ++++++++ 13 files changed, 313 insertions(+), 91 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index f713baffbe..dea4b7dd8f 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -30,6 +30,8 @@ use std::path::Path; use crate::function_tool::FunctionCallError; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; +use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; pub(crate) use crate::tools::code_mode::CodeModeExecuteHandler; pub(crate) use crate::tools::code_mode::CodeModeWaitHandler; pub use apply_patch::ApplyPatchHandler; @@ -84,6 +86,27 @@ fn resolve_workdir_base_path( .map_or_else(|| default_cwd.clone(), |workdir| default_cwd.join(workdir))) } +fn resolve_tool_environment<'a>( + turn: &'a TurnContext, + environment_id: Option<&str>, +) -> Result, FunctionCallError> { + environment_id.map_or_else( + || Ok(turn.environments.primary()), + |environment_id| { + turn.environments + .turn_environments + .iter() + .find(|environment| environment.environment_id == environment_id) + .map(Some) + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!( + "unknown turn environment id `{environment_id}`" + )) + }) + }, + ) +} + /// Validates feature/policy constraints for `with_additional_permissions` and /// normalizes any path-based permissions. Errors if the request is invalid. pub(crate) fn normalize_and_validate_additional_permissions( diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 563ee8d835..f109f22ac6 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -13,7 +13,7 @@ use crate::tools::handlers::implicit_granted_permissions; use crate::tools::handlers::normalize_and_validate_additional_permissions; use crate::tools::handlers::parse_arguments; use crate::tools::handlers::parse_arguments_with_base_path; -use crate::tools::handlers::resolve_workdir_base_path; +use crate::tools::handlers::resolve_tool_environment; use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; @@ -67,6 +67,16 @@ pub(crate) struct ExecCommandArgs { prefix_rule: Option>, } +#[derive(Debug, Deserialize)] +struct ExecCommandEnvironmentArgs { + #[serde(default)] + environment_id: Option, + // Keep this raw until after environment selection; relative paths must be + // resolved against the selected environment cwd, not the process cwd. + #[serde(default)] + workdir: Option, +} + #[derive(Debug, Deserialize)] struct WriteStdinArgs { // The model is trained on `session_id`. @@ -196,27 +206,38 @@ impl ToolHandler for UnifiedExecHandler { } }; - let Some(turn_environment) = turn.environments.primary() else { - return Err(FunctionCallError::RespondToModel( - "unified exec is unavailable in this session".to_string(), - )); - }; - let fs = turn_environment.environment.get_filesystem(); - let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager; let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone()); let response = match tool_name.name.as_str() { "exec_command" => { - let cwd = resolve_workdir_base_path(&arguments, &context.turn.cwd)?; + let environment_args: ExecCommandEnvironmentArgs = parse_arguments(&arguments)?; + let Some(turn_environment) = resolve_tool_environment( + turn.as_ref(), + environment_args.environment_id.as_deref(), + )? + else { + return Err(FunctionCallError::RespondToModel( + "unified exec is unavailable in this session".to_string(), + )); + }; + let cwd = environment_args + .workdir + .as_deref() + .filter(|workdir| !workdir.is_empty()) + .map_or_else( + || turn_environment.cwd.clone(), + |workdir| turn_environment.cwd.join(workdir), + ); + let environment = Arc::clone(&turn_environment.environment); + let fs = environment.get_filesystem(); let args: ExecCommandArgs = parse_arguments_with_base_path(&arguments, &cwd)?; let hook_command = args.cmd.clone(); - let workdir = context.turn.resolve_path(args.workdir.clone()); maybe_emit_implicit_skill_invocation( session.as_ref(), context.turn.as_ref(), &hook_command, - &workdir, + &cwd, ) .await; let process_id = manager.allocate_process_id().await; @@ -230,7 +251,6 @@ impl ToolHandler for UnifiedExecHandler { let command_for_display = codex_shell_command::parse_command::shlex_join(&command); let ExecCommandArgs { - workdir, tty, yield_time_ms, max_output_tokens, @@ -248,7 +268,7 @@ impl ToolHandler for UnifiedExecHandler { let requested_additional_permissions = additional_permissions.clone(); let effective_additional_permissions = apply_granted_turn_permissions( context.session.as_ref(), - context.turn.cwd.as_path(), + cwd.as_path(), sandbox_permissions, additional_permissions, ) @@ -275,10 +295,6 @@ impl ToolHandler for UnifiedExecHandler { ))); } - let workdir = workdir.filter(|value| !value.is_empty()); - - let workdir = workdir.map(|dir| context.turn.resolve_path(Some(dir))); - let cwd = workdir.clone().unwrap_or(cwd); let normalized_additional_permissions = match implicit_granted_permissions( sandbox_permissions, requested_additional_permissions.as_ref(), @@ -339,7 +355,8 @@ impl ToolHandler for UnifiedExecHandler { process_id, yield_time_ms, max_output_tokens: Some(max_output_tokens), - workdir, + cwd, + environment, network: context.turn.network.clone(), tty, sandbox_permissions: effective_additional_permissions diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index 1bdd0b82f9..70b933bad0 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -1,17 +1,10 @@ use super::*; use crate::shell::default_user_shell; -use crate::tools::handlers::parse_arguments_with_base_path; -use crate::tools::handlers::resolve_workdir_base_path; -use codex_protocol::models::AdditionalPermissionProfile as PermissionProfile; -use codex_protocol::models::FileSystemPermissions; use codex_tools::UnifiedExecShellMode; use codex_tools::ZshForkConfig; use codex_utils_absolute_path::AbsolutePathBuf; -use core_test_support::PathExt; use pretty_assertions::assert_eq; -use std::fs; use std::sync::Arc; -use tempfile::tempdir; use crate::session::tests::make_session_and_context; use crate::tools::context::ExecCommandToolOutput; @@ -185,39 +178,6 @@ fn test_get_command_ignores_explicit_shell_in_zsh_fork_mode() -> anyhow::Result< Ok(()) } -#[test] -fn exec_command_args_resolve_relative_additional_permissions_against_workdir() -> anyhow::Result<()> -{ - let cwd = tempdir()?; - let workdir = cwd.path().join("nested"); - fs::create_dir_all(&workdir)?; - let expected_write = workdir.join("relative-write.txt"); - let json = r#"{ - "cmd": "echo hello", - "workdir": "nested", - "additional_permissions": { - "file_system": { - "write": ["./relative-write.txt"] - } - } - }"#; - - let base_path = resolve_workdir_base_path(json, &cwd.path().abs())?; - let args: ExecCommandArgs = parse_arguments_with_base_path(json, &base_path)?; - - assert_eq!( - args.additional_permissions, - Some(PermissionProfile { - file_system: Some(FileSystemPermissions::from_read_write_roots( - /*read*/ None, - Some(vec![expected_write.abs()]), - )), - ..Default::default() - }) - ); - Ok(()) -} - #[tokio::test] async fn exec_command_pre_tool_use_payload_uses_raw_command() { let payload = ToolPayload::Function { diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index dcb42c36c6..c618d778d6 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -227,12 +227,13 @@ impl ToolOrchestrator { // Platform-specific flag gating is handled by SandboxManager::select_initial. let use_legacy_landlock = turn_ctx.features.use_legacy_landlock(); + let sandbox_cwd = tool.sandbox_cwd(req).unwrap_or(&turn_ctx.cwd); let initial_attempt = SandboxAttempt { sandbox: initial_sandbox, permissions: &turn_ctx.permission_profile, enforce_managed_network: managed_network_active, manager: &self.sandbox, - sandbox_cwd: &turn_ctx.cwd, + sandbox_cwd, codex_linux_sandbox_exe: turn_ctx.codex_linux_sandbox_exe.as_ref(), use_legacy_landlock, windows_sandbox_level: turn_ctx.windows_sandbox_level, @@ -350,7 +351,7 @@ impl ToolOrchestrator { permissions: &turn_ctx.permission_profile, enforce_managed_network: managed_network_active, manager: &self.sandbox, - sandbox_cwd: &turn_ctx.cwd, + sandbox_cwd, codex_linux_sandbox_exe: None, use_legacy_landlock, windows_sandbox_level: turn_ctx.windows_sandbox_level, diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 1bc5d72127..42f311bfcb 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -37,6 +37,7 @@ use crate::unified_exec::NoopSpawnLifecycle; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcess; use crate::unified_exec::UnifiedExecProcessManager; +use codex_exec_server::Environment; use codex_network_proxy::NetworkProxy; use codex_protocol::error::CodexErr; use codex_protocol::error::SandboxErr; @@ -48,6 +49,7 @@ use codex_tools::UnifiedExecShellMode; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::collections::HashMap; +use std::sync::Arc; use tokio_util::sync::CancellationToken; /// Request payload used by the unified-exec runtime after approvals and @@ -58,6 +60,7 @@ pub struct UnifiedExecRequest { pub hook_command: String, pub process_id: i32, pub cwd: AbsolutePathBuf, + pub environment: Arc, pub env: HashMap, pub exec_server_env_config: Option, pub explicit_env_overrides: HashMap, @@ -214,6 +217,10 @@ impl Approvable for UnifiedExecRuntime<'_> { } impl<'a> ToolRuntime for UnifiedExecRuntime<'a> { + fn sandbox_cwd<'b>(&self, req: &'b UnifiedExecRequest) -> Option<&'b AbsolutePathBuf> { + Some(&req.cwd) + } + fn network_approval_spec( &self, req: &UnifiedExecRequest, @@ -252,11 +259,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt if let Some(network) = managed_network { network.apply_to_env(&mut env); } - let environment_is_remote = ctx - .turn - .environments - .primary() - .is_some_and(|turn_environment| turn_environment.environment.is_remote()); + let environment_is_remote = req.environment.is_remote(); let command = if environment_is_remote { base_command.to_vec() } else { @@ -293,14 +296,10 @@ impl<'a> ToolRuntime for UnifiedExecRunt .await? { Some(prepared) => { - let Some(turn_environment) = ctx.turn.environments.primary() else { + if req.environment.is_remote() { return Err(ToolError::Rejected( - "exec_command is unavailable in this session".to_string(), - )); - }; - if turn_environment.environment.is_remote() { - return Err(ToolError::Rejected( - "unified_exec zsh-fork is not supported when exec_server_url is configured".to_string(), + "unified_exec zsh-fork is not supported for remote environments" + .to_string(), )); } return self @@ -310,7 +309,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt &prepared.exec_request, req.tty, prepared.spawn_lifecycle, - turn_environment.environment.as_ref(), + req.environment.as_ref(), ) .await .map_err(|err| match err { @@ -338,18 +337,13 @@ impl<'a> ToolRuntime for UnifiedExecRunt .env_for(command, options, managed_network) .map_err(|err| ToolError::Codex(err.into()))?; exec_env.exec_server_env_config = req.exec_server_env_config.clone(); - let Some(turn_environment) = ctx.turn.environments.primary() else { - return Err(ToolError::Rejected( - "exec_command is unavailable in this session".to_string(), - )); - }; self.manager .open_session_with_exec_env( req.process_id, &exec_env, req.tty, Box::new(NoopSpawnLifecycle), - turn_environment.environment.as_ref(), + req.environment.as_ref(), ) .await .map_err(|err| match err { diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index 122cd00fad..c17247beb4 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -358,6 +358,10 @@ pub(crate) trait ToolRuntime: Approvable + Sandboxable { None } + fn sandbox_cwd<'a>(&self, _req: &'a Req) -> Option<&'a AbsolutePathBuf> { + None + } + async fn run( &mut self, req: &Req, diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 97b37e8d80..9b74baf64c 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -27,6 +27,7 @@ use std::collections::HashSet; use std::sync::Arc; use std::sync::Weak; +use codex_exec_server::Environment; use codex_network_proxy::NetworkProxy; use codex_protocol::models::AdditionalPermissionProfile; use codex_utils_absolute_path::AbsolutePathBuf; @@ -93,7 +94,8 @@ pub(crate) struct ExecCommandRequest { pub process_id: i32, pub yield_time_ms: u64, pub max_output_tokens: Option, - pub workdir: Option, + pub cwd: AbsolutePathBuf, + pub environment: Arc, pub network: Option, pub tty: bool, pub sandbox_permissions: SandboxPermissions, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index c67abc48d6..4f85b1ddf7 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -371,10 +371,7 @@ impl UnifiedExecProcessManager { request: ExecCommandRequest, context: &UnifiedExecContext, ) -> Result { - let cwd = request - .workdir - .clone() - .unwrap_or_else(|| context.turn.cwd.clone()); + let cwd = request.cwd.clone(); let process = self .open_session_with_sandbox(&request, cwd.clone(), context) .await; @@ -1012,7 +1009,7 @@ impl UnifiedExecProcessManager { approval_policy: context.turn.approval_policy.value(), permission_profile: context.turn.permission_profile(), file_system_sandbox_policy: &file_system_sandbox_policy, - sandbox_cwd: context.turn.cwd.as_path(), + sandbox_cwd: cwd.as_path(), sandbox_permissions: if request.additional_permissions_preapproved { crate::sandboxing::SandboxPermissions::UseDefault } else { @@ -1026,6 +1023,7 @@ impl UnifiedExecProcessManager { hook_command: request.hook_command.clone(), process_id: request.process_id, cwd, + environment: Arc::clone(&request.environment), env, exec_server_env_config: Some(exec_server_env_config), explicit_env_overrides: context.turn.shell_environment_policy.r#set.clone(), diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 0c5b714161..bde32c9ab4 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -175,7 +175,11 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { process_id: 123, yield_time_ms: 1000, max_output_tokens: None, - workdir: None, + cwd: turn.cwd.clone(), + environment: turn + .environments + .primary_environment() + .expect("primary environment"), network: None, tty: true, sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 36cc2e6812..0bd449188c 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -3,23 +3,52 @@ use anyhow::Result; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::FileSystemSandboxContext; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::RemoveOptions; +use codex_features::Feature; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::PathBufExt; +use core_test_support::PathExt; use core_test_support::get_remote_test_env; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::test_codex; use core_test_support::test_codex::test_env; use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::fs; use std::path::PathBuf; use std::process::Command; use std::time::SystemTime; use std::time::UNIX_EPOCH; +use tempfile::TempDir; +async fn unified_exec_test(server: &wiremock::MockServer) -> Result { + let mut builder = test_codex().with_config(|config| { + config.use_experimental_unified_exec_tool = true; + let result = config.features.enable(Feature::UnifiedExec); + assert!( + result.is_ok(), + "unified exec should enable for test: {result:?}", + ); + }); + builder.build_remote_aware(server).await +} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_test_env_can_connect_and_use_filesystem() -> Result<()> { @@ -121,6 +150,114 @@ fn remote_exec(script: &str) -> Result<()> { Ok(()) } +async fn exec_command_routing_output( + test: &TestCodex, + server: &wiremock::MockServer, + call_id: &str, + arguments: Value, + environments: Option>, +) -> Result { + let response_mock = mount_sse_sequence( + server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "exec_command", &serde_json::to_string(&arguments)?), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + test.submit_turn_with_environments("route exec command", environments) + .await?; + + response_mock + .function_call_output_text(call_id) + .with_context(|| format!("missing function_call_output for {call_id}")) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_command_routes_to_selected_remote_environment() -> Result<()> { + skip_if_no_network!(Ok(())); + let Some(_remote_env) = get_remote_test_env() else { + return Ok(()); + }; + + let server = start_mock_server().await; + let test = unified_exec_test(&server).await?; + let local_cwd = TempDir::new()?; + fs::write(local_cwd.path().join("marker.txt"), "local-routing")?; + let local_selection = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: local_cwd.path().abs(), + }; + let remote_cwd = PathBuf::from(format!( + "/tmp/codex-remote-routing-{}", + SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + )) + .abs(); + let remote_marker_name = "marker.txt"; + test.fs() + .create_directory( + &remote_cwd, + CreateDirectoryOptions { recursive: true }, + /*sandbox*/ None, + ) + .await?; + test.fs() + .write_file( + &remote_cwd.join(remote_marker_name), + b"remote-routing".to_vec(), + /*sandbox*/ None, + ) + .await?; + let remote_selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd.clone(), + }; + let multi_env_output = exec_command_routing_output( + &test, + &server, + "call-multi-env", + json!({ + "shell": "/bin/sh", + "cmd": format!("cat {remote_marker_name}"), + "login": false, + "yield_time_ms": 1_000, + "environment_id": REMOTE_ENVIRONMENT_ID, + }), + Some(vec![local_selection, remote_selection]), + ) + .await?; + assert!( + multi_env_output.contains("remote-routing"), + "unexpected multi-env output: {multi_env_output}", + ); + assert!( + !multi_env_output.contains("local-routing"), + "multi-env command should not route to local: {multi_env_output}", + ); + + test.fs() + .remove( + &remote_cwd, + RemoveOptions { + recursive: true, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_test_env_sandboxed_read_allows_readable_root() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/tools/src/local_tool.rs b/codex-rs/tools/src/local_tool.rs index ed4080d5f4..aeabdbfa30 100644 --- a/codex-rs/tools/src/local_tool.rs +++ b/codex-rs/tools/src/local_tool.rs @@ -17,6 +17,13 @@ pub struct ShellToolOptions { } pub fn create_exec_command_tool(options: CommandToolOptions) -> ToolSpec { + create_exec_command_tool_with_environment_id(options, /*include_environment_id*/ false) +} + +pub(crate) fn create_exec_command_tool_with_environment_id( + options: CommandToolOptions, + include_environment_id: bool, +) -> ToolSpec { let mut properties = BTreeMap::from([ ( "cmd".to_string(), @@ -63,6 +70,14 @@ pub fn create_exec_command_tool(options: CommandToolOptions) -> ToolSpec { )), ); } + if include_environment_id { + properties.insert( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional environment id from the block. If omitted, uses the primary environment.".to_string(), + )), + ); + } properties.extend(create_approval_parameters( options.exec_permission_approvals_enabled, )); diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index cba5a4bafa..687f99fcda 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -7,6 +7,7 @@ use crate::ShellToolOptions; use crate::SpawnAgentToolOptions; use crate::TOOL_SEARCH_DEFAULT_LIMIT; use crate::TOOL_SEARCH_TOOL_NAME; +use crate::ToolEnvironmentMode; use crate::ToolHandlerKind; use crate::ToolName; use crate::ToolRegistryPlan; @@ -27,7 +28,6 @@ use crate::create_close_agent_tool_v1; use crate::create_close_agent_tool_v2; use crate::create_code_mode_tool; use crate::create_create_goal_tool; -use crate::create_exec_command_tool; use crate::create_followup_task_tool; use crate::create_get_goal_tool; use crate::create_image_generation_tool; @@ -60,6 +60,7 @@ use crate::create_web_search_tool; use crate::create_write_stdin_tool; use crate::default_namespace_description; use crate::dynamic_tool_to_loadable_tool_spec; +use crate::local_tool::create_exec_command_tool_with_environment_id; use crate::mcp_tool_to_responses_api_tool; use crate::request_permissions_tool_description; use crate::request_user_input_tool_description; @@ -135,6 +136,8 @@ pub fn build_tool_registry_plan( } if config.environment_mode.has_environment() { + let include_environment_id = + matches!(config.environment_mode, ToolEnvironmentMode::Multiple); match &config.shell_type { ConfigShellToolType::Default => { plan.push_spec( @@ -154,10 +157,13 @@ pub fn build_tool_registry_plan( } ConfigShellToolType::UnifiedExec => { plan.push_spec( - create_exec_command_tool(CommandToolOptions { - allow_login_shell: config.allow_login_shell, - exec_permission_approvals_enabled, - }), + create_exec_command_tool_with_environment_id( + CommandToolOptions { + allow_login_shell: config.allow_login_shell, + exec_permission_approvals_enabled, + }, + include_environment_id, + ), /*supports_parallel_tool_calls*/ true, config.code_mode_enabled, ); diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index 264a2e5c7e..fe5507ddc7 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -19,6 +19,7 @@ use crate::ToolRegistryPlanDeferredTool; use crate::ToolRegistryPlanMcpTool; use crate::ToolsConfigParams; use crate::WaitAgentTimeoutOptions; +use crate::create_exec_command_tool; use crate::mcp_call_tool_result_output_schema; use crate::request_user_input_available_modes; use codex_app_server_protocol::AppInfo; @@ -159,6 +160,49 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { } } +#[test] +fn exec_command_spec_includes_environment_id_only_for_multiple_selected_environments() { + let model_info = model_info(); + let available_models = Vec::new(); + let mut features = Features::with_defaults(); + features.enable(Feature::UnifiedExec); + let config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + permission_profile: &PermissionProfile::Disabled, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + + let (single_environment_tools, _) = build_specs( + &config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_process_tool_environment_id( + &single_environment_tools, + "exec_command", + /*expected_present*/ false, + ); + + let multi_environment_config = config.with_environment_mode(ToolEnvironmentMode::Multiple); + let (multi_environment_tools, _) = build_specs( + &multi_environment_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_process_tool_environment_id( + &multi_environment_tools, + "exec_command", + /*expected_present*/ true, + ); +} + #[test] fn test_build_specs_collab_tools_enabled() { let model_info = model_info(); @@ -2424,6 +2468,23 @@ fn find_tool<'a>(tools: &'a [ConfiguredToolSpec], expected_name: &str) -> &'a Co .unwrap_or_else(|| panic!("expected tool {expected_name}")) } +fn assert_process_tool_environment_id( + tools: &[ConfiguredToolSpec], + expected_name: &str, + expected_present: bool, +) { + let tool = find_tool(tools, expected_name); + let ToolSpec::Function(ResponsesApiTool { parameters, .. }) = &tool.spec else { + panic!("expected function tool {expected_name}"); + }; + let (properties, _) = expect_object_schema(parameters); + assert_eq!( + properties.contains_key("environment_id"), + expected_present, + "{expected_name} environment_id parameter presence" + ); +} + fn find_namespace_function_tool<'a>( tools: &'a [ConfiguredToolSpec], expected_namespace: &str,