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 <noreply@openai.com>
This commit is contained in:
starr-openai
2026-05-05 12:12:03 -07:00
committed by GitHub
Unverified
parent fb7e1eb6fc
commit 78421face0
13 changed files with 313 additions and 91 deletions
+23
View File
@@ -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<Option<&'a TurnEnvironment>, 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(
@@ -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<Vec<String>>,
}
#[derive(Debug, Deserialize)]
struct ExecCommandEnvironmentArgs {
#[serde(default)]
environment_id: Option<String>,
// 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<String>,
}
#[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
@@ -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 {
+3 -2
View File
@@ -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,
@@ -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<Environment>,
pub env: HashMap<String, String>,
pub exec_server_env_config: Option<ExecServerEnvConfig>,
pub explicit_env_overrides: HashMap<String, String>,
@@ -214,6 +217,10 @@ impl Approvable<UnifiedExecRequest> for UnifiedExecRuntime<'_> {
}
impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> 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<UnifiedExecRequest, UnifiedExecProcess> 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<UnifiedExecRequest, UnifiedExecProcess> 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<UnifiedExecRequest, UnifiedExecProcess> 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<UnifiedExecRequest, UnifiedExecProcess> 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 {
+4
View File
@@ -358,6 +358,10 @@ pub(crate) trait ToolRuntime<Req, Out>: Approvable<Req> + Sandboxable {
None
}
fn sandbox_cwd<'a>(&self, _req: &'a Req) -> Option<&'a AbsolutePathBuf> {
None
}
async fn run(
&mut self,
req: &Req,
+3 -1
View File
@@ -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<usize>,
pub workdir: Option<AbsolutePathBuf>,
pub cwd: AbsolutePathBuf,
pub environment: Arc<Environment>,
pub network: Option<NetworkProxy>,
pub tty: bool,
pub sandbox_permissions: SandboxPermissions,
@@ -371,10 +371,7 @@ impl UnifiedExecProcessManager {
request: ExecCommandRequest,
context: &UnifiedExecContext,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
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(),
@@ -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,