Preserve one-shot exec when unified exec is disabled (#41393)

## Why

Managed configuration can disable resumable unified execution while leaving shell tools enabled. This should not remove command execution entirely or grant access to retained processes.

## What changed

- Keep `exec_command` available in completion-only mode when `unified_exec` is disabled.
- Remove `tty`, `yield_time_ms`, `session_id`, and `write_stdin` from that mode, and add a `timeout_ms` argument.
- Terminate one-shot commands on timeout or turn cancellation, reporting timed-out commands with exit code 124.

## Testing

Added integration coverage for tool availability, timeout cleanup, and interruption cleanup under managed configuration.

GitOrigin-RevId: a382d34f922fdc4ba52889272c6db985abd42e47
This commit is contained in:
jif
2026-08-28 19:04:54 +00:00
committed by copyberry
parent 4878401e8f
commit b836aecd4d
10 changed files with 411 additions and 56 deletions
@@ -36,6 +36,8 @@ pub(crate) struct ExecCommandArgs {
#[serde(default = "default_exec_yield_time_ms")]
yield_time_ms: u64,
#[serde(default)]
timeout_ms: Option<u64>,
#[serde(default)]
max_output_tokens: Option<usize>,
#[serde(default)]
sandbox_permissions: Option<SandboxPermissions>,
@@ -1,6 +1,8 @@
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS;
use crate::exec_policy::prompt_is_rejected_by_policy;
use crate::function_tool::FunctionCallError;
use crate::maybe_emit_implicit_skill_invocation;
@@ -35,6 +37,7 @@ use codex_sandboxing::SandboxManager;
use codex_sandboxing::SandboxType;
use codex_sandboxing::SandboxablePreference;
use codex_shell_command::shell_detect::detect_shell_type;
use codex_tools::JsonSchema;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_utils_output_truncation::approx_token_count;
@@ -57,13 +60,21 @@ pub(crate) struct ExecCommandHandlerOptions {
pub(crate) include_windows_shell_guidance: bool,
}
#[derive(Clone, Copy)]
enum ExecCommandLifetime {
Interactive,
OneShot,
}
pub struct ExecCommandHandler {
options: ExecCommandHandlerOptions,
lifetime: ExecCommandLifetime,
}
impl Default for ExecCommandHandler {
fn default() -> Self {
Self {
lifetime: ExecCommandLifetime::Interactive,
options: ExecCommandHandlerOptions {
allow_login_shell: false,
exec_permission_approvals_enabled: false,
@@ -77,7 +88,17 @@ impl Default for ExecCommandHandler {
impl ExecCommandHandler {
pub(crate) fn new(options: ExecCommandHandlerOptions) -> Self {
Self { options }
Self {
options,
lifetime: ExecCommandLifetime::Interactive,
}
}
pub(crate) fn one_shot(options: ExecCommandHandlerOptions) -> Self {
Self {
options,
lifetime: ExecCommandLifetime::OneShot,
}
}
}
@@ -87,7 +108,7 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
}
fn spec(&self) -> ToolSpec {
create_exec_command_tool_with_environment_id(
let spec = create_exec_command_tool_with_environment_id(
CommandToolOptions {
allow_login_shell: self.options.allow_login_shell,
exec_permission_approvals_enabled: self.options.exec_permission_approvals_enabled,
@@ -95,7 +116,11 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
self.options.include_environment_id,
self.options.include_shell_parameter,
self.options.include_windows_shell_guidance,
)
);
match self.lifetime {
ExecCommandLifetime::Interactive => spec,
ExecCommandLifetime::OneShot => one_shot_exec_command_spec(spec),
}
}
fn supports_parallel_tool_calls(&self) -> bool {
@@ -253,8 +278,9 @@ impl ExecCommandHandler {
let command_for_display = codex_shell_command::parse_command::shlex_join(&command);
let ExecCommandArgs {
tty,
mut tty,
yield_time_ms,
timeout_ms,
max_output_tokens,
sandbox_permissions: _,
additional_permissions,
@@ -262,6 +288,15 @@ impl ExecCommandHandler {
prefix_rule,
..
} = args;
let completion_timeout = match self.lifetime {
ExecCommandLifetime::Interactive => None,
ExecCommandLifetime::OneShot => {
tty = false;
Some(Duration::from_millis(
timeout_ms.unwrap_or(DEFAULT_EXEC_COMMAND_TIMEOUT_MS),
))
}
};
let exec_permission_approvals_enabled =
session.features().enabled(Feature::ExecPermissionApprovals);
@@ -356,32 +391,34 @@ impl ExecCommandHandler {
}
emit_unified_exec_tty_metric(&turn.session_telemetry, tty);
match manager
.exec_command(
ExecCommandRequest {
command,
shell_type,
hook_command: hook_command.clone(),
process_id,
yield_time_ms,
max_output_tokens,
cwd,
sandbox_cwd: native_environment_cwd,
turn_environment: turn_environment.clone(),
shell_mode,
network: context.step_context.turn.network.clone(),
tty,
sandbox_permissions: effective_additional_permissions.sandbox_permissions,
additional_permissions: normalized_additional_permissions,
additional_permissions_preapproved: effective_additional_permissions
.permissions_preapproved,
justification,
prefix_rule,
},
&context,
)
.await
{
let request = ExecCommandRequest {
command,
shell_type,
hook_command: hook_command.clone(),
process_id,
yield_time_ms,
max_output_tokens,
cwd,
sandbox_cwd: native_environment_cwd,
turn_environment: turn_environment.clone(),
shell_mode,
network: context.step_context.turn.network.clone(),
tty,
sandbox_permissions: effective_additional_permissions.sandbox_permissions,
additional_permissions: normalized_additional_permissions,
additional_permissions_preapproved: effective_additional_permissions
.permissions_preapproved,
justification,
prefix_rule,
};
let result = match completion_timeout {
Some(timeout) => {
UnifiedExecProcessManager::exec_command_to_completion(request, &context, timeout)
.await
}
None => manager.exec_command(request, &context).await,
};
match result {
Ok(response) => Ok(boxed_tool_output(response)),
Err(UnifiedExecError::SandboxDenied {
output,
@@ -415,6 +452,35 @@ impl ExecCommandHandler {
}
}
fn one_shot_exec_command_spec(spec: ToolSpec) -> ToolSpec {
let ToolSpec::Function(mut spec) = spec else {
unreachable!("exec_command has a function schema");
};
spec.description = spec.description.replacen(
"Runs a command in a PTY, returning output or a session ID for ongoing interaction.",
"Runs a command to completion and returns its output. The process is terminated on timeout or cancellation and cannot be resumed.",
1,
);
let properties = spec.parameters.properties.get_or_insert_default();
properties.remove("tty");
properties.remove("yield_time_ms");
properties.insert(
"timeout_ms".to_string(),
JsonSchema::number(Some(
"Maximum command runtime. Defaults to 10000 ms.".to_string(),
)),
);
if let Some(output_properties) = spec
.output_schema
.as_mut()
.and_then(|schema| schema.get_mut("properties"))
.and_then(serde_json::Value::as_object_mut)
{
output_properties.remove("session_id");
}
ToolSpec::Function(spec)
}
impl CoreToolRuntime for ExecCommandHandler {
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(payload, ToolPayload::Function { .. })
+11 -4
View File
@@ -1084,7 +1084,6 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, registry: &mut ToolRegistr
let environment_mode = tool_environment_mode(context.environments);
if !environment_mode.has_environment()
|| !features.enabled(Feature::ShellTool)
|| !features.enabled(Feature::UnifiedExec)
|| matches!(context.model_info.shell_type, ConfigShellToolType::Disabled)
{
return;
@@ -1093,7 +1092,7 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, registry: &mut ToolRegistr
let allow_login_shell = any_environment_allows_login_shell(context.environments);
let exec_permission_approvals_enabled = features.enabled(Feature::ExecPermissionApprovals);
let include_environment_id = matches!(environment_mode, ToolEnvironmentMode::Multiple);
registry.add(ExecCommandHandler::new(ExecCommandHandlerOptions {
let options = ExecCommandHandlerOptions {
allow_login_shell,
exec_permission_approvals_enabled,
include_environment_id,
@@ -1102,8 +1101,16 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, registry: &mut ToolRegistr
context.environments,
),
include_windows_shell_guidance: should_include_windows_shell_guidance(context.environments),
}));
registry.add(WriteStdinHandler);
};
if features.enabled(Feature::UnifiedExec) {
registry.add(ExecCommandHandler::new(options));
registry.add(WriteStdinHandler);
} else {
// Managed requirements are the only configuration path that can keep
// unified exec disabled. Preserve command execution without exposing a
// resumable process or write_stdin authority prohibited by policy.
registry.add(ExecCommandHandler::one_shot(options));
}
}
fn unified_exec_should_include_shell_parameter(
@@ -213,6 +213,7 @@ pub(crate) fn spawn_exit_watcher(
.await;
} else {
let exit_code = process.exit_code().unwrap_or(-1);
let timed_out = process.timed_out();
finish_and_track_measurements(
plugin_metrics_sidecar,
exit_code,
@@ -233,6 +234,7 @@ pub(crate) fn spawn_exit_watcher(
String::new(),
exit_code,
duration,
timed_out,
)
.await;
}
@@ -344,6 +346,7 @@ pub(crate) async fn emit_exec_end_for_unified_exec(
fallback_output: String,
exit_code: i32,
duration: Duration,
timed_out: bool,
) {
let aggregated_output = resolve_aggregated_output(&transcript, fallback_output).await;
let output = ExecToolCallOutput {
@@ -352,7 +355,7 @@ pub(crate) async fn emit_exec_end_for_unified_exec(
stderr: StreamOutput::new(String::new()),
aggregated_output: StreamOutput::new(aggregated_output),
duration,
timed_out: false,
timed_out,
};
let event_ctx = ToolEventCtx::new(
session_ref.as_ref(),
+2
View File
@@ -2,6 +2,7 @@
//!
//! Responsibilities
//! - Manages interactive processes (create, reuse, buffer output with caps).
//! - Supports completion-only calls that terminate on timeout or cancellation.
//! - Uses the shared ToolOrchestrator to handle approval, sandbox selection, and
//! retry semantics in a single, descriptive flow.
//! - Spawns the PTY from a sandbox-transformed `ExecRequest`; on sandbox denial,
@@ -49,6 +50,7 @@ use codex_core_plugins::PluginMetricsSidecar;
mod async_watcher;
mod errors;
mod head_tail_buffer;
mod oneshot;
mod process;
mod process_manager;
mod process_state;
+87
View File
@@ -0,0 +1,87 @@
//! One-shot execution over the shared unified-exec lifecycle.
//!
//! The caller never receives a resumable process. Timeout and cancellation
//! terminate the exact process handle published by the normal startup path.
use std::sync::Arc;
use std::sync::OnceLock;
use tokio::time::Duration;
use tokio::time::Instant;
use super::ExecCommandRequest;
use super::UnifiedExecContext;
use super::UnifiedExecError;
use super::UnifiedExecProcess;
use super::UnifiedExecProcessManager;
use crate::tools::context::ExecCommandToolOutput;
pub(super) struct Completion<'a> {
pub timeout: Duration,
pub timed_out: bool,
pub process: &'a OnceLock<Arc<UnifiedExecProcess>>,
}
impl UnifiedExecProcessManager {
pub(crate) async fn exec_command_to_completion(
request: ExecCommandRequest,
context: &UnifiedExecContext,
timeout: Duration,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
let context = UnifiedExecContext::new(
Arc::clone(&context.session),
Arc::clone(&context.step_context),
context.cancellation_token.child_token(),
context.call_id.clone(),
);
let _cancel_on_drop = context.cancellation_token.clone().drop_guard();
tokio::spawn(async move {
let manager = &context.session.services.unified_exec_manager;
let process_id = request.process_id;
if Instant::now().checked_add(timeout).is_none() {
manager.release_process_id(process_id).await;
return Err(UnifiedExecError::process_failed(
"timeout_ms is too large".into(),
));
}
let process = OnceLock::new();
let mut completion = Completion {
timeout,
timed_out: false,
process: &process,
};
let result = {
let mut execution =
Box::pin(manager.exec_command_inner(request, &context, Some(&mut completion)));
tokio::select! {
biased;
_ = context.cancellation_token.cancelled() => {
if let Some(process) = process.get() {
if !process.has_exited()
&& let Err(err) = process.terminate_confirmed().await
{
process.fail_and_terminate(err.to_string());
}
let _ = execution.await;
} else {
drop(execution);
manager.release_process_id(process_id).await;
}
Err(UnifiedExecError::process_failed("command cancelled".into()))
}
result = &mut execution => result,
}
};
result.map(|mut output| {
if completion.timed_out {
output.process_id = None;
}
output
})
})
.await
.map_err(|err| UnifiedExecError::process_failed(err.to_string()))?
}
}
+13
View File
@@ -97,6 +97,7 @@ pub(crate) struct UnifiedExecProcess {
state_rx: watch::Receiver<ProcessState>,
output_task: Option<JoinHandle<()>>,
sandbox_type: SandboxType,
timed_out: AtomicBool,
_spawn_lifecycle: Option<SpawnLifecycleHandle>,
}
@@ -137,6 +138,7 @@ impl UnifiedExecProcess {
state_rx,
output_task: None,
sandbox_type,
timed_out: AtomicBool::new(false),
_spawn_lifecycle: spawn_lifecycle,
}
}
@@ -195,6 +197,9 @@ impl UnifiedExecProcess {
}
pub(super) fn exit_code(&self) -> Option<i32> {
if self.timed_out() {
return Some(124);
}
let state = self.state_rx.borrow().clone();
match &self.process_handle {
ProcessHandle::Local(process_handle) => {
@@ -204,6 +209,14 @@ impl UnifiedExecProcess {
}
}
pub(super) fn mark_timed_out(&self) {
self.timed_out.store(true, Ordering::Release);
}
pub(super) fn timed_out(&self) -> bool {
self.timed_out.load(Ordering::Acquire)
}
fn finish_termination(&self) {
self.output.cancellation_token.cancel();
if let Some(output_task) = &self.output_task {
@@ -12,6 +12,8 @@ use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use super::oneshot::Completion;
use crate::codex_thread::BackgroundTerminalInfo;
use crate::exec_env::CODEX_PERMISSION_PROFILE_ENV_VAR;
use crate::exec_env::CODEX_THREAD_ID_ENV_VAR;
@@ -471,6 +473,16 @@ impl UnifiedExecProcessManager {
&self,
request: ExecCommandRequest,
context: &UnifiedExecContext,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
self.exec_command_inner(request, context, /*completion*/ None)
.await
}
pub(super) async fn exec_command_inner(
&self,
request: ExecCommandRequest,
context: &UnifiedExecContext,
mut completion: Option<&mut Completion<'_>>,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
let cwd = request.cwd.clone();
let process = self
@@ -490,6 +502,9 @@ impl UnifiedExecProcessManager {
permissions,
} = attempt;
let process = Arc::new(process);
if let Some(completion) = completion.as_ref() {
let _ = completion.process.set(Arc::clone(&process));
}
let network_denial_monitor = deferred_network_approval.as_ref().map(|deferred| {
terminate_process_on_network_denial(
Arc::clone(&process),
@@ -574,13 +589,28 @@ impl UnifiedExecProcessManager {
// For the initial exec_command call, we both stream output to events
// (via start_streaming_output above) and collect a snapshot here for
// the tool response body.
let deadline = start + Duration::from_millis(yield_time_ms);
let wait = completion.as_ref().map_or_else(
|| Duration::from_millis(yield_time_ms),
|completion| completion.timeout,
);
let deadline = start
.checked_add(wait)
.ok_or_else(|| UnifiedExecError::process_failed("timeout_ms is too large".into()))?;
let collected_output = Self::collect_output_until_deadline(
process.output_handles(),
Some(context.session.subscribe_elicitation_pause_state()),
deadline,
)
.await;
if let Some(completion) = completion.as_mut()
&& !process.has_exited()
{
completion.timed_out = true;
process.mark_timed_out();
if let Err(err) = process.terminate_confirmed().await {
process.fail_and_terminate(err.to_string());
}
}
let wall_time = Instant::now().saturating_duration_since(start);
let original_token_count = usize::try_from(approx_tokens_from_byte_count(
@@ -657,15 +687,20 @@ impl UnifiedExecProcessManager {
{
return Err(fail_process_with_message(entry.process.as_ref(), message));
}
process
.check_for_sandbox_denial_with_text(&text)
.await
.map_err(|err| {
err.with_output_collection_metadata(
original_token_count,
output_omitted_bytes,
)
})?;
if !completion
.as_ref()
.is_some_and(|completion| completion.timed_out)
{
process
.check_for_sandbox_denial_with_text(&text)
.await
.map_err(|err| {
err.with_output_collection_metadata(
original_token_count,
output_omitted_bytes,
)
})?;
}
let metrics_sidecar = entry
.plugin_metrics_sidecar
.as_ref()
@@ -725,6 +760,7 @@ impl UnifiedExecProcessManager {
text.clone(),
exit,
wall_time,
process.timed_out(),
)
.await;