diff --git a/codex-rs/hooks/src/engine/command_runner.rs b/codex-rs/hooks/src/engine/command_runner.rs index 7f8d8be070..36914c6287 100644 --- a/codex-rs/hooks/src/engine/command_runner.rs +++ b/codex-rs/hooks/src/engine/command_runner.rs @@ -16,6 +16,7 @@ use async_channel::Sender; use codex_protocol::shell_environment::scrub_non_inheritable_env_vars; #[cfg(windows)] use codex_utils_pty::JobObject; +use futures::future::try_join; use tokio::io::AsyncWriteExt; use tokio::process::Command; use tokio::sync::Semaphore; @@ -269,27 +270,28 @@ pub(crate) async fn run_command( job: process_tree_job, }; - if let Some(mut stdin) = child.stdin.take() - && let Err(err) = stdin.write_all(input_json.as_bytes()).await - && err.kind() != ErrorKind::BrokenPipe - { - let _ = child.kill().await; - return finish_command_run( - started_at, - started, - CommandRunCompletion { - exit_code: None, - stdout: String::new(), - stderr: String::new(), - error: Some(format!("failed to write hook stdin: {err}")), - outcome: "stdin_error", - }, - ); - } + let stdin = child.stdin.take(); + let write_stdin = async { + if let Some(mut stdin) = stdin + && let Err(err) = stdin.write_all(input_json.as_bytes()).await + && err.kind() != ErrorKind::BrokenPipe + { + return Err(("stdin_error", format!("failed to write hook stdin: {err}"))); + } + Ok(()) + }; + let wait_with_output = async { + child + .wait_with_output() + .await + .map_err(|err| ("wait_error", err.to_string())) + }; let timeout_duration = Duration::from_secs(handler.timeout_sec); - match timeout(timeout_duration, child.wait_with_output()).await { - Ok(Ok(output)) => { + // Drain output while sending input so neither pipe can block the other, and + // include stdin writes in the deadline even when the hook never reads them. + match timeout(timeout_duration, try_join(write_stdin, wait_with_output)).await { + Ok(Ok(((), output))) => { // Successfully completed hooks may intentionally leave detached helpers running. #[cfg(windows)] if let Some(job) = process_tree_guard.job.as_ref() { @@ -308,15 +310,15 @@ pub(crate) async fn run_command( }, ) } - Ok(Err(err)) => finish_command_run( + Ok(Err((outcome, error))) => finish_command_run( started_at, started, CommandRunCompletion { exit_code: None, stdout: String::new(), stderr: String::new(), - error: Some(err.to_string()), - outcome: "wait_error", + error: Some(error), + outcome, }, ), Err(_) => finish_command_run( diff --git a/codex-rs/hooks/src/engine/command_runner_tests.rs b/codex-rs/hooks/src/engine/command_runner_tests.rs index f0a45e5b4a..644a9de6fb 100644 --- a/codex-rs/hooks/src/engine/command_runner_tests.rs +++ b/codex-rs/hooks/src/engine/command_runner_tests.rs @@ -238,6 +238,42 @@ async fn fast_exiting_hook_preserves_stdout_when_stdin_is_not_consumed() { assert_eq!(result.error, None); } +#[tokio::test] +async fn hook_drains_output_and_times_out_while_stdin_is_blocked() { + let temp = tempdir().expect("create temp dir"); + let mut handler = write_handler( + &temp, + r#"from pathlib import Path +import sys +import time + +sys.stdout.write("x" * 1024 * 1024) +sys.stdout.flush() +sys.stderr.write("x" * 1024 * 1024) +sys.stderr.flush() +Path("output-drained").touch() +time.sleep(30) +"#, + ); + handler.timeout_sec = 2; + let ConfiguredHandlerKind::Command { command, env, .. } = &handler.kind else { + panic!("expected command hook"); + }; + let input_json = format!(r#"{{"padding":"{}"}}"#, "x".repeat(1024 * 1024)); + let (runtime, _result_receiver) = runtime(); + + let result = timeout( + Duration::from_secs(10), + run_command(&runtime, &handler, command, env, &input_json, temp.path()), + ) + .await + .expect("the hook timeout must also cover blocked stdin writes"); + + assert!(temp.path().join("output-drained").exists()); + assert_eq!(result.exit_code, None); + assert_eq!(result.error, Some("hook timed out after 2s".to_string())); +} + #[tokio::test] async fn command_hook_does_not_expose_configured_noise_auth_token() { let temp = tempdir().expect("create temp dir");