[Hooks] Allow PreToolUse hooks to make exec_command wait for long-running commands (waitUntilCompletion)

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(),