Deliver environment hook decisions while output reads are queued (#58996)

## Why

A completed environment hook can be held up by a queued output-read RPC when its executor is busy, delaying the hook's decision.

## What changed

- Race pushed process events against output polling so a complete, ordered event stream can deliver the hook's output and exit status immediately. Retain polling as a fallback when events are incomplete or unavailable.
- Include launch-parameter preparation in the hook's timeout budget.

## Testing

Extend the executor ownership test to hold the owner's RPC slot while a `PreToolUse` hook finishes. Cover both allowing and blocking commands, checking that the decision arrives before the slot is released and that only the owning executor receives the hook.

GitOrigin-RevId: bc80b99a2a32379e3c2eafefe6db66a6a20f4048

codex-rs/core/src/hook_environment_executor.rs

@@ -9,7 +9,10 @@ use anyhow::bail;
 use codex_exec_server::ExecOutputStream;
 use codex_exec_server::ExecParams;
 use codex_exec_server::ExecProcess;
+use codex_exec_server::ExecProcessEvent;
+use codex_exec_server::ExecProcessEventReceiver;
 use codex_exec_server::ProcessId;
+use codex_exec_server::ProcessOutputChunk;
 use codex_hooks::EnvironmentHookOutput;
 use codex_hooks::EnvironmentHookSource;
 use codex_hooks::HookEnvironmentExecutor;
@@ -197,52 +200,31 @@ impl HookEnvironmentExecutor for EnvironmentCommandExecutor {
             };
             let argv = shell.derive_exec_args(&script, environment.config().allow_login_shell);
             // Keep network attribution registered until the hook and its output are complete.
-            let (params, _execution_network) = self
-                .launch_params(environment, argv, env)
+            let deadline = tokio::time::Instant::now() + timeout;
+            let (params, _execution_network) = tokio::time::timeout_at(deadline, self.launch_params(environment, argv, env))
                 .await
+                .context("environment command hook timed out during launch")?
                 .map_err(|error| match error {
                     ToolError::Rejected(message) => anyhow::anyhow!("{message}"),
                     ToolError::Codex(error) => anyhow::Error::new(error),
                 })?;
-            let deadline = tokio::time::Instant::now() + timeout;
             let backend = environment.environment.get_exec_backend();
             let process = tokio::time::timeout_at(deadline, backend.start(params))
                 .await
                 .context("environment command hook timed out during launch")??
                 .process;
             let mut guard = HookProcessGuard(Some(Arc::clone(&process)));
+            let events = process.subscribe_events();
             let output = tokio::time::timeout_at(deadline, async {
                 let written = process.write(input_json).await?;
                 if written.status != codex_exec_server::WriteStatus::Accepted {
                     bail!("hook process did not accept its input: {:?}", written.status);
                 }
-                let mut after_seq = None;
-                let mut stdout = Vec::new();
-                let mut stderr = Vec::new();
-                loop {
-                    let result = process.read(after_seq, Some(64 * 1024), Some(1000)).await?;
-                    if let Some(failure) = result.failure {
-                        bail!("{failure}");
-                    }
-                    // Closed describes the process, not this bounded page. Keep
-                    // reading retained output until we receive an empty final page.
-                    let drained = result.closed && result.chunks.is_empty();
-                    for chunk in result.chunks {
-                        let target = match chunk.stream {
-                            ExecOutputStream::Stdout | ExecOutputStream::Pty => &mut stdout,
-                            ExecOutputStream::Stderr => &mut stderr,
-                        };
-                        let remaining = MAX_HOOK_OUTPUT_BYTES.saturating_sub(target.len());
-                        target.extend_from_slice(&chunk.chunk.0[..remaining.min(chunk.chunk.0.len())]);
-                    }
-                    after_seq = result.next_seq.checked_sub(1).or(after_seq);
-                    if drained {
-                        return Ok(EnvironmentHookOutput {
-                            exit_code: result.exit_code.context("hook process closed without exit status")?,
-                            stdout: String::from_utf8_lossy(&stdout).into_owned(),
-                            stderr: String::from_utf8_lossy(&stderr).into_owned(),
-                        });
-                    }
+                // A finished hook can deliver its decision while the read RPC is still queued.
+                tokio::select! {
+                    biased;
+                    Some(output) = pushed_hook_output(events) => Ok(output),
+                    result = read_hook_output(process.as_ref()) => result,
                 }
             })
             .await
@@ -253,3 +235,77 @@ impl HookEnvironmentExecutor for EnvironmentCommandExecutor {
         .boxed()
     }
 }
+
+async fn pushed_hook_output(mut events: ExecProcessEventReceiver) -> Option<EnvironmentHookOutput> {
+    let mut output = HookOutputBuffer::default();
+    let mut exit_code = None;
+    let mut next_seq = 1;
+    loop {
+        match events.recv().await.ok()? {
+            ExecProcessEvent::Output(chunk) if chunk.seq == next_seq => output.push(chunk),
+            ExecProcessEvent::Exited {
+                seq,
+                exit_code: code,
+                ..
+            } if seq == next_seq => {
+                exit_code = Some(code);
+            }
+            ExecProcessEvent::Closed { seq } if seq == next_seq => {
+                return Some(output.finish(exit_code?));
+            }
+            ExecProcessEvent::Output(_)
+            | ExecProcessEvent::Exited { .. }
+            | ExecProcessEvent::Closed { .. }
+            | ExecProcessEvent::Failed(_) => return None,
+        }
+        next_seq = next_seq.checked_add(1)?;
+    }
+}
+
+async fn read_hook_output(process: &dyn ExecProcess) -> anyhow::Result<EnvironmentHookOutput> {
+    let mut output = HookOutputBuffer::default();
+    let mut after_seq = None;
+    loop {
+        let result = process.read(after_seq, Some(64 * 1024), Some(1000)).await?;
+        if let Some(failure) = result.failure {
+            bail!("{failure}");
+        }
+        // Closed describes the process, not this bounded page. Keep reading
+        // retained output until we receive an empty final page.
+        if result.closed && result.chunks.is_empty() {
+            let code = result
+                .exit_code
+                .context("hook process closed without exit status")?;
+            return Ok(output.finish(code));
+        }
+        for chunk in result.chunks {
+            output.push(chunk);
+        }
+        after_seq = result.next_seq.checked_sub(1).or(after_seq);
+    }
+}
+
+#[derive(Default)]
+struct HookOutputBuffer {
+    stdout: Vec<u8>,
+    stderr: Vec<u8>,
+}
+
+impl HookOutputBuffer {
+    fn push(&mut self, chunk: ProcessOutputChunk) {
+        let target = match chunk.stream {
+            ExecOutputStream::Stdout | ExecOutputStream::Pty => &mut self.stdout,
+            ExecOutputStream::Stderr => &mut self.stderr,
+        };
+        let remaining = MAX_HOOK_OUTPUT_BYTES.saturating_sub(target.len());
+        target.extend_from_slice(&chunk.chunk.0[..remaining.min(chunk.chunk.0.len())]);
+    }
+
+    fn finish(self, exit_code: i32) -> EnvironmentHookOutput {
+        EnvironmentHookOutput {
+            exit_code,
+            stdout: String::from_utf8_lossy(&self.stdout).into_owned(),
+            stderr: String::from_utf8_lossy(&self.stderr).into_owned(),
+        }
+    }
+}

codex-rs/core/tests/suite/hooks_executor.rs

@@ -487,8 +487,20 @@ async fn environment_hook_receives_prompt_and_reports_exit_status(
     Ok(())
 }
 
+#[test_case(HookRunStatus::Completed, 0; "allowed command runs on its owner")]
+#[test_case(HookRunStatus::Blocked, 2; "denied command stays blocked on its owner")]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn pre_tool_use_hook_runs_only_on_the_executor_that_owns_the_command() -> Result<()> {
+async fn pre_tool_use_hook_runs_only_on_the_executor_that_owns_the_command(
+    expected_status: HookRunStatus,
+    exit_code: i32,
+) -> Result<()> {
+    use codex_exec_server::HttpRequestParams;
+    use codex_protocol::models::PermissionProfile;
+    use tokio::io::AsyncWriteExt;
+
+    let hook_ready = TcpListener::bind("127.0.0.1:0").await?;
+    let port = hook_ready.local_addr()?.port();
+    let hold_request = TcpListener::bind("127.0.0.1:0").await?;
     let server = start_mock_server().await;
     let test = test_codex()
         .with_config(|config| {
@@ -529,15 +541,23 @@ async fn pre_tool_use_hook_runs_only_on_the_executor_that_owns_the_command() ->
         let home = reported.codex_home_dir.to_abs_path()?;
         input_files.insert(id, home.join("pre-tool-use-input.json"));
         let mut config = environment_config_for_selection(&test.config, &selection);
-        config.environment_hooks = toml::from_str(
-            r#"
-            [[PreToolUse]]
-            matcher = '^Bash$'
-            [[PreToolUse.hooks]]
-            type = 'command'
-            command = 'cat > "$CODEX_HOME/pre-tool-use-input.json"; echo denied-by-environment-hook >&2; exit 2'
-        "#,
-        )?;
+        config.allow_login_shell = false;
+        config.permission_profile = PermissionProfileSnapshot::legacy(PermissionProfile::Disabled);
+        config.environment_hooks = serde_json::from_value(json!({
+            "PreToolUse": [{"matcher": "^Bash$", "hooks": [{
+                "type": "command", "timeout": 10,
+                "command": format!(r#"cat > "$CODEX_HOME/pre-tool-use-input.json" || exit 1
+                    bash -c 'exec 3<>/dev/tcp/127.0.0.1/{port} && read -r <&3' || exit 1
+                    echo owner-verdict >&2
+                    exit {exit_code}"#),
+                "command_windows": format!(r#"ForEach-Object {{ [IO.File]::WriteAllText("$env:CODEX_HOME/pre-tool-use-input.json", $_) }}
+                    $client = [Net.Sockets.TcpClient]::new("127.0.0.1", {port})
+                    $null = $client.GetStream().ReadByte()
+                    $client.Close()
+                    [Console]::Error.WriteLine("owner-verdict")
+                    exit {exit_code}"#),
+            }]}]
+        }))?;
         selection.config = EnvironmentConfigState::Ready(config);
         selections.push(selection);
     }
@@ -558,20 +578,75 @@ async fn pre_tool_use_hook_runs_only_on_the_executor_that_owns_the_command() ->
         ],
     )
     .await;
-    test.submit_turn_with_environments("run the command in owner", Some(selections))
+    let thread = &test.codex;
+    thread
+        .start_or_steer_turn(
+            TurnInputRequest::user_input(vec![UserInput::Text {
+                text: "run the command in owner".into(),
+                text_elements: Vec::new(),
+            }])
+            .with_thread_settings(ThreadSettingsOverrides {
+                environments: Some(TurnEnvironmentRequests::new(
+                    test.config.cwd.clone(),
+                    selections
+                        .into_iter()
+                        .map(TurnEnvironmentSelection::into_request)
+                        .collect(),
+                )),
+                permission_profile: Some(PermissionProfile::Disabled),
+                ..Default::default()
+            }),
+        )
         .await?;
-
+    let completed = wait_for_event_match(thread, |event| match event {
+        EventMsg::HookCompleted(event) => Some(event.run.status),
+        _ => None,
+    });
+    tokio::pin!(completed);
+    let mut hook = tokio::select! {
+        socket = hook_ready.accept() => socket?.0,
+        status = &mut completed => anyhow::bail!("hook exited before the barrier: {status:?}"),
+    };
+    // Hold the owner's sole RPC slot, then release the hook. Its decision must
+    // reach the turn before the executor can service another output read.
+    let http = manager
+        .get_environment("owner")
+        .context("missing owner")?
+        .get_http_client();
+    let request = http.http_request(HttpRequestParams {
+        method: "GET".into(),
+        url: format!("http://{}/hold", hold_request.local_addr()?),
+        headers: Vec::new(),
+        body: None,
+        timeout_ms: None,
+        redirect_policy: Default::default(),
+        request_id: "hold-owner".into(),
+        stream_response: false,
+    });
+    tokio::pin!(request);
+    let mut held = tokio::select! {
+        socket = hold_request.accept() => socket?.0,
+        status = &mut completed => anyhow::bail!("hook exited while reserving the slot: {status:?}"),
+        result = &mut request => anyhow::bail!("executor did not hold the request: {result:?}"),
+    };
+    hook.write_all(b"go\n").await?;
+    let status = completed.await;
+    held.write_all(b"HTTP/1.1 204 No Content\r\nConnection: close\r\n\r\n")
+        .await?;
+    assert_eq!(request.await?.status, 204);
+    assert_eq!(status, expected_status);
+    wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await;
     let output = responses
         .function_call_output_text(call_id)
         .context("missing tool result")?;
-    assert!(
-        output.contains("Command blocked by PreToolUse hook: denied-by-environment-hook"),
-        "{output}"
-    );
-    assert!(
-        !test.workspace_path("command-ran").try_exists()?,
-        "the denied command ran"
-    );
+    let allowed = expected_status == HookRunStatus::Completed;
+    if !allowed {
+        assert!(
+            output.contains("Command blocked by PreToolUse hook: owner-verdict"),
+            "{output}"
+        );
+    }
+    assert_eq!(test.workspace_path("command-ran").try_exists()?, allowed);
     let owner_input =
         std::fs::read(&input_files["owner"]).context("owner did not receive the hook")?;
     let input: Value = serde_json::from_slice(&owner_input)?;