diff --git a/codex-rs/core/src/hook_environment_executor.rs b/codex-rs/core/src/hook_environment_executor.rs index fa6b303405..c8ff65108a 100644 --- a/codex-rs/core/src/hook_environment_executor.rs +++ b/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 { + 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 { + 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, + stderr: Vec, +} + +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(), + } + } +} diff --git a/codex-rs/core/tests/suite/hooks_executor.rs b/codex-rs/core/tests/suite/hooks_executor.rs index 4852848c75..79e68ebe4d 100644 --- a/codex-rs/core/tests/suite/hooks_executor.rs +++ b/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)?;