[Hooks] Allow PreToolUse hooks to make exec_command wait for long-running commands (waitUntilCompletion)
What variant of Codex are you using?
cli
What feature would you like to see?
## Summary
PreToolUse hooks can already rewrite an exec_command's command via
hookSpecificOutput.updatedInput.command, but there is no way for a hook to make the
resulting tool call block until the process exits. This makes hooks a poor fit for
long-running commands: security/audit wrappers, sandboxing shims, and scanning tools
that only produce their final, meaningful output at exit.
## Problem
Today exec_command always returns after yield_time_ms (default ~1s, and capped
around 30s on the initial call — see #22541). For anything longer, the model is bounced
into session-polling mode with write_stdin, which:
- costs extra round-trips and tokens for every long command,
- is unreliable in practice — models abandon yielded sessions and retry duplicate
commands (#33816),
- breaks hook use cases where the hook wraps the command with a scanner/auditor whose
verdict is only printed when the process finishes. The hook can rewrite the command,
but cannot say "wait for it to finish", so the model receives partial output and a
dangling session instead of the final result.
Related issues: #22541, #33816, #35713, #24175.
## Proposal
Let a PreToolUse hook opt a specific exec_command call into run-to-completion
semantics by returning a boolean waitUntilCompletion in updatedInput alongside
command:
``json``
{
"hookSpecificOutput": {
"hookEventName": "PreToolUse",
"updatedInput": {
"command": "my-security-wrapper --scan -- npm test",
"waitUntilCompletion": true
}
}
}
Behavior:
- When
waitUntilCompletion: true,exec_commandcollects output until the process
exits and returns the final output in a single tool response. No session is left
open and no write_stdin polling is needed.
- When the field is absent, behavior is completely unchanged (normal
yield_time_ms
semantics).
- A non-boolean
waitUntilCompletionfails the tool call with
hook returned non-boolean field 'waitUntilCompletion'.
- Output remains bounded by the existing head/tail buffer and
max_output_tokens;
only the wait becomes unbounded. Hooks are expected to wrap commands with
timeout(1) where a hard cap is desired.
## Implementation
I have a working patch (based on current main). It is intentionally small
(+118/-22 lines across 7 files, including tests):
ExecCommandArgsgainswait_for_completion: bool(defaultfalse).ExecCommandHandler::with_updated_hook_inputacceptswaitUntilCompletionfrom
the hook's updatedInput and rewrites the tool arguments with both cmd and
wait_for_completion (via the existing rewrite_function_arguments helper).
ExecCommandRequestcarries the flag through to the process manager.UnifiedExecProcessManager::collect_output_until_deadlinenow takes
Option<Instant> for the deadline; None means "wait until process exit" instead
of the yield_time_ms deadline. Pause-state deadline extension is skipped when
there is no deadline.
<details>
<summary>Full diff</summary>
```diff
diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs
index f6a5bfa542..e2a86a0258 100644
--- a/codex-rs/core/src/tools/handlers/unified_exec.rs
+++ b/codex-rs/core/src/tools/handlers/unified_exec.rs
@@ -36,6 +36,8 @@ pub(crate) struct ExecCommandArgs {
#[serde(default = "default_exec_yield_time_ms")]
yield_time_ms: u64,
#[serde(default)]
- wait_for_completion: bool,
- #[serde(default)]
max_output_tokens: Option<usize>,
#[serde(default)]
sandbox_permissions: Option<SandboxPermissions>,
diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
index f7a24e86d4..ffc0d8704e 100644
--- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
+++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
@@ -15,7 +15,7 @@ use crate::tools::handlers::parse_arguments;
use crate::tools::handlers::parse_arguments_with_base_path;
use crate::tools::handlers::resolve_sandbox_permissions;
use crate::tools::handlers::resolve_tool_environment;
-use crate::tools::handlers::rewrite_function_string_argument;
+use crate::tools::handlers::rewrite_function_arguments;
use crate::tools::handlers::updated_hook_command;
use crate::tools::hook_names::HookToolName;
use crate::tools::registry::CoreToolRuntime;
@@ -243,6 +243,7 @@ impl ExecCommandHandler {
let ExecCommandArgs {
tty,
yield_time_ms,
- wait_for_completion,
max_output_tokens,
sandbox_permissions: _,
additional_permissions,
@@ -354,6 +355,7 @@ impl ExecCommandHandler {
hook_command: hook_command.clone(),
process_id,
yield_time_ms,
- wait_for_completion,
max_output_tokens,
cwd,
sandbox_cwd: native_environment_cwd,
@@ -433,13 +435,28 @@ impl CoreToolRuntime for ExecCommandHandler {
"hook input rewrite received unsupported exec_command payload".to_string(),
));
};
- let command = updated_hook_command(&updated_input)?;
- let wait_for_completion = match updated_input.get("waitUntilCompletion") {
- None => None,
- Some(value) => Some(value.as_bool().ok_or_else(|| {
- FunctionCallError::RespondToModel(
- "hook returned non-boolean field
waitUntilCompletion".to_string(), - )
- })?),
- };
invocation.payload = ToolPayload::Function {
- arguments: rewrite_function_string_argument(
- &arguments,
- "exec_command",
- "cmd",
- updated_hook_command(&updated_input)?,
- )?,
- arguments: rewrite_function_arguments(&arguments, "exec_command", move |arguments| {
- arguments.insert(
- "cmd".to_string(),
- serde_json::Value::String(command.to_string()),
- );
- if let Some(wait_for_completion) = wait_for_completion {
- arguments.insert(
- "wait_for_completion".to_string(),
- serde_json::Value::Bool(wait_for_completion),
- );
- }
- })?,
};
Ok(invocation)
}
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 eee95fd43a..40eb184b2e 100644
--- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs
+++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs
@@ -329,6 +329,36 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() {
);
}
+#[tokio::test]
+async fn exec_command_hook_input_rewrite_preserves_native_wait() {
- let payload = ToolPayload::Function {
- arguments: serde_json::json!({ "cmd": "slow-analysis" }).to_string(),
- };
- let invocation = invocation_for_payload("exec_command", "call-native-wait", payload).await;
- let handler = ExecCommandHandler::default();
+
- let invocation = handler
- .with_updated_hook_input(
- invocation,
- serde_json::json!({
- "command": "rewritten-analysis",
- "waitUntilCompletion": true
- }),
- )
- .expect("hook input rewrite should succeed");
+
- assert_eq!(
- invocation.payload,
- ToolPayload::Function {
- arguments: serde_json::json!({
- "cmd": "rewritten-analysis",
- "wait_for_completion": true
- })
- .to_string()
- }
- );
+}
+
#[tokio::test]
async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_shot_commands() {
let payload = ToolPayload::Function {
diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs
index 77b23625c6..79ae5284fd 100644
--- a/codex-rs/core/src/unified_exec/mod.rs
+++ b/codex-rs/core/src/unified_exec/mod.rs
@@ -97,6 +97,7 @@ pub(crate) struct ExecCommandRequest {
pub hook_command: String,
pub process_id: i32,
pub yield_time_ms: u64,
- pub wait_for_completion: bool,
pub max_output_tokens: Option<usize>,
pub cwd: PathUri,
pub sandbox_cwd: PathUri,
diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs
index d097e783d9..d297c09d1a 100644
--- a/codex-rs/core/src/unified_exec/mod_tests.rs
+++ b/codex-rs/core/src/unified_exec/mod_tests.rs
@@ -157,7 +157,7 @@ async fn exec_command_with_tty(
let collected_output = UnifiedExecProcessManager::collect_output_until_deadline(
process.output_handles(),
Some(session.subscribe_elicitation_pause_state()),
- deadline,
- Some(deadline),
)
.await;
let wall_time = Instant::now().saturating_duration_since(started_at);
@@ -806,7 +806,7 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul
let collected = UnifiedExecProcessManager::collect_output_until_deadline(
process.output_handles(),
/pause_state/ None,
- Instant::now() + Duration::from_millis(2_500),
- Some(Instant::now() + Duration::from_millis(2_500)),
)
.await
.to_bytes_with_omission_marker();
diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs
index 8b2924ffbd..234161528c 100644
--- a/codex-rs/core/src/unified_exec/process_manager.rs
+++ b/codex-rs/core/src/unified_exec/process_manager.rs
@@ -560,7 +560,11 @@ 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 deadline = if request.wait_for_completion {
- None
- } else {
- Some(start + Duration::from_millis(yield_time_ms))
- };
let collected_output = Self::collect_output_until_deadline(
process.output_handles(),
Some(context.session.subscribe_elicitation_pause_state()),
@@ -824,7 +828,7 @@ impl UnifiedExecProcessManager {
let start = Instant::now();
let deadline = start + Duration::from_millis(yield_time_ms);
let collected_output =
- Self::collect_output_until_deadline(&output, pause_state, deadline).await;
- Self::collect_output_until_deadline(&output, pause_state, Some(deadline)).await;
let wall_time = Instant::now().saturating_duration_since(start);
let original_token_count = usize::try_from(approx_tokens_from_byte_count(
@@ -1308,7 +1312,7 @@ impl UnifiedExecProcessManager {
pub(super) async fn collect_output_until_deadline(
output: &OutputHandles,
mut pause_state: Option<watch::Receiver<bool>>,
- mut deadline: Instant,
- mut deadline: Option<Instant>,
) -> HeadTailBuffer {
const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50);
@@ -1323,12 +1327,14 @@ impl UnifiedExecProcessManager {
let mut exit_signal_received = cancellation_token.is_cancelled();
let mut post_exit_deadline: Option<Instant> = None;
loop {
- Self::extend_deadlines_while_paused(
- &mut pause_state,
- &mut deadline,
- &mut post_exit_deadline,
- )
- .await;
- if let Some(deadline) = deadline.as_mut() {
- Self::extend_deadlines_while_paused(
- &mut pause_state,
- deadline,
- &mut post_exit_deadline,
- )
- .await;
- }
let drained_output: HeadTailBuffer;
let has_drained_output: bool;
let mut wait_for_output = None;
@@ -1348,7 +1354,9 @@ impl UnifiedExecProcessManager {
{
break;
}
- let remaining = deadline.saturating_duration_since(Instant::now());
- let remaining = deadline
- .map(|deadline| deadline.saturating_duration_since(Instant::now()))
- .unwrap_or(Duration::MAX);
if remaining == Duration::ZERO {
break;
}
@@ -1390,7 +1398,7 @@ impl UnifiedExecProcessManager {
collected.push_buffer(drained_output);
exit_signal_received |= cancellation_token.is_cancelled();
- if Instant::now() >= deadline {
- if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
break;
}
}
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 4d3429e038..2e92854c1c 100644
--- a/codex-rs/core/src/unified_exec/process_manager_tests.rs
+++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs
@@ -299,7 +299,7 @@ async fn output_collection_stays_bounded_across_repeated_drains() {
let collect = UnifiedExecProcessManager::collect_output_until_deadline(
&output,
/pause_state/ None,
- Instant::now() + Duration::from_secs(5),
- Some(Instant::now() + Duration::from_secs(5)),
);
let produce = async {
for byte in [b'a', b'b', b'c'] {
@@ -336,6 +336,43 @@ async fn output_collection_stays_bounded_across_repeated_drains() {
assert_eq!(collected, expected);
}
+#[tokio::test]
+async fn output_collection_without_deadline_waits_for_exit() {
- let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
- let output_notify = Arc::new(Notify::new());
- let output_closed = Arc::new(AtomicBool::new(false));
- let output_closed_notify = Arc::new(Notify::new());
- let cancellation_token = CancellationToken::new();
- let output = OutputHandles {
- output_buffer: Arc::clone(&output_buffer),
- output_notify: Arc::clone(&output_notify),
- output_closed: Arc::clone(&output_closed),
- output_closed_notify: Arc::clone(&output_closed_notify),
- cancellation_token: cancellation_token.clone(),
- };
- let _exit_task = tokio::spawn(async move {
- tokio::time::sleep(Duration::from_millis(75)).await;
- output_closed.store(true, Ordering::Release);
- cancellation_token.cancel();
- output_closed_notify.notify_waiters();
- output_notify.notify_waiters();
- });
+
- let collected = tokio::time::timeout(
- Duration::from_secs(2),
- UnifiedExecProcessManager::collect_output_until_deadline(
- &output, /pause_state/ None, None,
- ),
- )
- .await
- .expect("collector without a deadline should wait for exit");
+
- assert_eq!(
- collected.to_bytes_with_omission_marker().as_slice(),
- b"" as &[u8]
- );
+}
+
#[tokio::test]
async fn output_collection_preserves_omissions_from_drained_buffer() {
let mut buffered_output = HeadTailBuffer::default();
@@ -414,6 +451,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() {
hook_command: "echo before".to_string(),
process_id: 123,
yield_time_ms: 1000,
- wait_for_completion: false,
max_output_tokens: None,
#[allow(deprecated)]
cwd: turn.cwd.clone().into(),
```
</details>
## Testing
The patch adds two unit tests:
exec_command_hook_input_rewrite_preserves_native_wait: a hook returning
waitUntilCompletion: true rewrites the tool arguments to
{"cmd": "...", "wait_for_completion": true}.
output_collection_without_deadline_waits_for_exit: the output collector with no
deadline blocks until the process signals exit instead of returning early.
Happy to open a PR with this change if the approach looks reasonable, and to adjust
the field naming or wire format to match whatever convention you prefer.
Additional information
_No response_
1 Comment
Potential duplicates detected. Please review them and close your issue if it is a duplicate.
Powered by Codex Action