diff --git a/core/hooks.py b/core/hooks.py index 90834ab..e7d1c39 100644 --- a/core/hooks.py +++ b/core/hooks.py @@ -34,6 +34,7 @@ from typing import Any from config.project_settings import HOOKS_REL, find_hooks_root, is_trusted +from utils.proc import run_shell try: import yaml @@ -161,15 +162,7 @@ def run( if not hook.matches(tool_name): continue try: - proc = subprocess.run( - hook.command, - shell=True, - input=stdin, - capture_output=True, - text=True, - timeout=hook.timeout, - cwd=self.cwd, - ) + proc = run_shell(hook.command, input=stdin, timeout=hook.timeout, cwd=self.cwd) except subprocess.TimeoutExpired: feedback.append(f"hook '{hook.command}' timed out after {hook.timeout}s") continue diff --git a/tests/test_hooks.py b/tests/test_hooks.py index e252380..9c2b677 100644 --- a/tests/test_hooks.py +++ b/tests/test_hooks.py @@ -18,6 +18,13 @@ def run(self, arguments): return "ran" +def py(code: str) -> str: + """A hook command that runs Python code. Works in sh and in Windows cmd.exe, where `echo x >&2; + exit 2` is not valid (cmd has no `;` separator and no `sleep`). Use single quotes inside `code`.""" + exe = f'"{sys.executable}"' if " " in sys.executable else sys.executable + return f'{exe} -c "{code}"' + + def _executor(hooks): return ToolExecutor([EchoTool()], require_confirmation=False, hooks=hooks) @@ -29,25 +36,27 @@ def test_matcher(): def test_pre_hook_blocks_with_exit_2(): - h = HookManager({PRE_TOOL_USE: [Hook("echo nope >&2; exit 2")]}) + h = HookManager({PRE_TOOL_USE: [Hook(py("import sys; sys.stderr.write('nope'); sys.exit(2)"))]}) out = _executor(h).execute("echo", {}) assert "Blocked by PreToolUse hook" in out and "nope" in out def test_pre_hook_nonblocking_error_allows(): - h = HookManager({PRE_TOOL_USE: [Hook("exit 1")]}) + h = HookManager({PRE_TOOL_USE: [Hook(py("import sys; sys.exit(1)"))]}) assert _executor(h).execute("echo", {}) == "ran" def test_post_hook_feedback_appended(): - h = HookManager({POST_TOOL_USE: [Hook("echo lint failed >&2; exit 2")]}) + h = HookManager( + {POST_TOOL_USE: [Hook(py("import sys; sys.stderr.write('lint failed'); sys.exit(2)"))]} + ) out = _executor(h).execute("echo", {}) assert out.startswith("ran") and "lint failed" in out def test_hook_receives_json_payload(tmp_path): dump = tmp_path / "payload.json" - cmd = f"{sys.executable} -c \"import sys; open(r'{dump}','w').write(sys.stdin.read())\"" + cmd = py(f"import sys; open(r'{dump}','w').write(sys.stdin.read())") h = HookManager({POST_TOOL_USE: [Hook(cmd)]}, cwd=tmp_path) _executor(h).execute("echo", {"a": 1}) data = json.loads(dump.read_text()) @@ -58,8 +67,12 @@ def test_hook_receives_json_payload(tmp_path): def test_timeout_does_not_block(): - h = HookManager({PRE_TOOL_USE: [Hook("sleep 5", timeout=1)]}) + import time + + h = HookManager({PRE_TOOL_USE: [Hook(py("import time; time.sleep(30)"), timeout=1)]}) + started = time.monotonic() assert _executor(h).execute("echo", {}) == "ran" + assert time.monotonic() - started < 15, "the hook was not stopped by its timeout" def test_load_file(tmp_path: Path): diff --git a/tests/test_proc.py b/tests/test_proc.py new file mode 100644 index 0000000..5842faa --- /dev/null +++ b/tests/test_proc.py @@ -0,0 +1,103 @@ +"""Tests for run_shell: timeouts must stop the whole process tree (the Windows hang).""" + +import subprocess +import sys +import time + +import pytest + +from tools.shell import ShellTool +from utils import proc as proc_mod +from utils.proc import run_shell + + +def py(code: str) -> str: + exe = f'"{sys.executable}"' if " " in sys.executable else sys.executable + return f'{exe} -c "{code}"' + + +def test_captures_output_and_exit_code(): + r = run_shell(py("import sys; print('out'); sys.stderr.write('err'); sys.exit(3)"), timeout=30) + assert r.returncode == 3 and "out" in r.stdout and "err" in r.stderr + + +def test_input_is_passed_on_stdin(): + r = run_shell(py("import sys; print(sys.stdin.read().upper())"), input="hello", timeout=30) + assert "HELLO" in r.stdout + + +def test_runs_in_the_given_directory(tmp_path): + r = run_shell(py("import os; print(os.getcwd())"), cwd=tmp_path, timeout=30) + assert tmp_path.name in r.stdout + + +def test_utf8_output_decodes(tmp_path): + r = run_shell( + py("import sys; sys.stdout.buffer.write('caf\\u00e9'.encode('utf-8'))"), + timeout=30, + encoding="utf-8", + errors="replace", + ) + assert "café" in r.stdout + + +def test_timeout_returns_promptly_instead_of_waiting_for_the_command(): + started = time.monotonic() + with pytest.raises(subprocess.TimeoutExpired): + run_shell(py("import time; time.sleep(30)"), timeout=1) + assert time.monotonic() - started < 15, "the command was not stopped by its timeout" + + +def test_timeout_kills_grandchildren_too(tmp_path): + # parent.py starts child.py (which would write a marker after 2s) and then sleeps; killing only the + # shell/parent would let the child finish and write the marker. + marker = tmp_path / "marker.txt" + (tmp_path / "child.py").write_text( + f"import time\ntime.sleep(2)\nopen(r'{marker}', 'w').write('alive')\n" + ) + (tmp_path / "parent.py").write_text( + "import subprocess, sys, time\n" + "subprocess.Popen([sys.executable, 'child.py'])\n" + "time.sleep(30)\n" + ) + with pytest.raises(subprocess.TimeoutExpired): + run_shell(py("exec(open('parent.py').read())"), timeout=1, cwd=tmp_path) + time.sleep(3.5) # longer than the child's 2s sleep + assert not marker.exists(), "a grandchild survived the timeout" + + +def test_keyboard_interrupt_kills_the_child_and_propagates(monkeypatch): + spawned = [] + + class InterruptedPopen(subprocess.Popen): + def __init__(self, *a, **k): + super().__init__(*a, **k) + spawned.append(self) + self._first = True + + def communicate(self, *a, **k): + if self._first: + self._first = False + raise KeyboardInterrupt + return super().communicate(*a, **k) + + monkeypatch.setattr(proc_mod.subprocess, "Popen", InterruptedPopen) + with pytest.raises(KeyboardInterrupt): + run_shell(py("import time; time.sleep(30)"), timeout=60) + assert spawned[0].poll() is not None, "the child was left running after Ctrl+C" + + +def test_shell_tool_reports_a_timeout_and_returns(monkeypatch): + import tools.shell as shell_mod + + real = shell_mod.run_shell + monkeypatch.setattr(shell_mod, "run_shell", lambda cmd, **kw: real(cmd, **{**kw, "timeout": 1})) + started = time.monotonic() + out = ShellTool().run({"command": py("import time; time.sleep(30)")}) + assert "timed out" in out.lower() + assert time.monotonic() - started < 15 + + +def test_shell_tool_still_formats_output(): + out = ShellTool().run({"command": py("print('hi')")}) + assert out.startswith("STDOUT:") and "hi" in out diff --git a/tools/shell.py b/tools/shell.py index 09f9ac1..7cde74e 100644 --- a/tools/shell.py +++ b/tools/shell.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, Field from tools.base import Tool +from utils.proc import run_shell def normalize_windows_command(command: str) -> str: @@ -78,16 +79,7 @@ def run(self, arguments: dict[str, Any]) -> Any: if is_win: exec_cmd = normalize_windows_command(command) - result = subprocess.run( - exec_cmd, - shell=True, - capture_output=True, - text=True, - check=False, - timeout=120, - encoding="utf-8", - errors="replace", - ) + result = run_shell(exec_cmd, timeout=120, encoding="utf-8", errors="replace") output = "" if result.stdout: diff --git a/utils/proc.py b/utils/proc.py new file mode 100644 index 0000000..bbae6c3 --- /dev/null +++ b/utils/proc.py @@ -0,0 +1,77 @@ +"""Run shell commands with a timeout that really stops them. + +``subprocess.run(cmd, shell=True, timeout=...)`` only kills the *shell*. On Windows the command's own +process keeps running and keeps the output pipes open, so ``run`` blocks until it exits anyway (a +30-second sleep with a 1-second timeout took 30 seconds). On POSIX the same happens whenever the shell +does not exec its last command. This helper starts the command in its own process group / session and +kills the whole tree on timeout or on Ctrl+C. +""" + +import os +import signal +import subprocess +from typing import Any + +REAP_TIMEOUT = 5 # seconds to wait for pipes to close after the tree has been killed + + +def kill_process_tree(proc: subprocess.Popen) -> None: + """Kill ``proc`` and everything it started.""" + try: + if os.name == "nt": + # /T = the whole tree, /F = force. taskkill ships with every Windows install. + subprocess.run( + ["taskkill", "/F", "/T", "/PID", str(proc.pid)], + capture_output=True, + check=False, + timeout=REAP_TIMEOUT, + ) + else: + os.killpg(proc.pid, signal.SIGKILL) # the group id equals the pid (new session) + except (ProcessLookupError, PermissionError, OSError, subprocess.SubprocessError): + pass + try: + proc.kill() # belt and braces, e.g. if taskkill was unavailable + except OSError: + pass + + +def run_shell( + command: str, + *, + input: str | None = None, + timeout: float | None = None, + cwd: Any = None, + encoding: str | None = None, + errors: str | None = None, +) -> subprocess.CompletedProcess: + """``subprocess.run(command, shell=True, capture_output=True, text=True)`` that kills the whole + process tree on timeout or interrupt. Raises ``subprocess.TimeoutExpired`` like ``run`` does. + + stdin is inherited unless ``input`` is given, matching ``subprocess.run``. + """ + extra: dict[str, Any] = {} + if os.name != "nt": + extra["start_new_session"] = True + proc = subprocess.Popen( + command, + shell=True, + stdin=subprocess.PIPE if input is not None else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + encoding=encoding, + errors=errors, + cwd=cwd, + **extra, + ) + try: + stdout, stderr = proc.communicate(input=input, timeout=timeout) + except BaseException: # TimeoutExpired, KeyboardInterrupt, ... + kill_process_tree(proc) + try: + proc.communicate(timeout=REAP_TIMEOUT) + except Exception: + pass + raise + return subprocess.CompletedProcess(command, proc.returncode, stdout, stderr)