diff --git a/amplifier_module_tool_bash/__init__.py b/amplifier_module_tool_bash/__init__.py index 2145852..1180fb0 100644 --- a/amplifier_module_tool_bash/__init__.py +++ b/amplifier_module_tool_bash/__init__.py @@ -24,6 +24,44 @@ logger = logging.getLogger(__name__) +TIMEOUT_MIN_SECONDS = 1 +TIMEOUT_MAX_SECONDS = 3600 + + +def _validate_timeout_seconds(value: Any, *, source: str) -> int: + """Validate a timeout value (seconds) from either config or caller input. + + Requirements: + - must be an int (bool rejected) + - 1 <= value <= 3600 + + Raises: + TypeError: If the value is not an integer (including bool). + ValueError: If the integer is outside the supported range. + """ + + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError( + f"Invalid {source} timeout: timeout must be an integer number of seconds " + f"between {TIMEOUT_MIN_SECONDS} and {TIMEOUT_MAX_SECONDS} (got {value!r})." + ) + if value < TIMEOUT_MIN_SECONDS: + raise ValueError( + f"Invalid {source} timeout: timeout must be an integer number of seconds " + f"between {TIMEOUT_MIN_SECONDS} and {TIMEOUT_MAX_SECONDS} (got {value!r})." + ) + if value > TIMEOUT_MAX_SECONDS: + suggestion = "" + if value % 1000 == 0: + as_seconds = value // 1000 + if TIMEOUT_MIN_SECONDS <= as_seconds <= TIMEOUT_MAX_SECONDS: + suggestion = f" It looks like you passed milliseconds; did you mean {as_seconds} seconds?" + raise ValueError( + f"Invalid {source} timeout: timeout is specified in seconds and must be <= " + f"{TIMEOUT_MAX_SECONDS} (got {value!r}).{suggestion}" + ) + return value + def _read_ppid(pid: int) -> int | None: """Read a process's parent PID from /proc//stat (Linux only). @@ -730,6 +768,82 @@ def _arbitrate_windows_shell( if gitbash_exe: return gitbash_exe, False return None, False +async def _cleanup_process_tree( + process: asyncio.subprocess.Process, *, pgid: int | None, is_windows: bool +) -> None: + """Best-effort termination of a subprocess and its descendants. + + Mirrors the tool's timeout cleanup behavior: + - kill the process group (Unix) when available + - on Linux, also signal setsid()-detached descendants discovered via /proc + - wait briefly, then SIGKILL + - reap via communicate() + """ + + if pgid is not None and not is_windows: + # Walk /proc for descendants BEFORE killing anything. + descendant_pids = _find_descendant_pids(process.pid) + try: + # Send SIGTERM to process group first (graceful shutdown) + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + pass # Process group already gone + except PermissionError: + # Fall back to killing just the main process + process.kill() + + # Belt and suspenders: also signal any descendants that escaped the + # process group and wouldn't receive the killpg() above. + _signal_pids(descendant_pids, signal.SIGTERM) + + # Give processes a moment to clean up + await asyncio.sleep(0.5) + + # Force kill if still running + try: + os.killpg(pgid, signal.SIGKILL) + except ProcessLookupError: + pass # Already terminated + except PermissionError: + pass + _signal_pids(descendant_pids, signal.SIGKILL) + else: + # Windows or no pgid: kill just the main process + process.kill() + + # Reap / close pipes (best-effort) + try: + await asyncio.wait_for(process.communicate(), timeout=5) + except TimeoutError: + pass # Best effort cleanup + + +async def _await_process_tree_cleanup( + process: asyncio.subprocess.Process, *, pgid: int | None, is_windows: bool +) -> None: + """Run process-tree cleanup to completion despite repeated cancellation. + + If cancellation arrives while cleanup is running, defer propagation until + the bounded cleanup task finishes, then raise CancelledError. + """ + + cleanup_task = asyncio.create_task( + _cleanup_process_tree(process, pgid=pgid, is_windows=is_windows) + ) + cancellation_received = False + + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + cancellation_received = True + continue + except Exception as cleanup_error: # noqa: BLE001 + logger.error("Process cleanup failed: %s", cleanup_error) + break + + if cancellation_received: + raise asyncio.CancelledError() async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None = None): @@ -828,7 +942,9 @@ def __init__(self, config: dict[str, Any]): """ self.config = config self.require_approval = config.get("require_approval", True) - self.timeout = config.get("timeout", 30) + self.timeout = _validate_timeout_seconds( + config.get("timeout", 30), source="config" + ) self.working_dir = config.get("working_dir", ".") # Output limiting to prevent context overflow self.max_output_bytes = config.get( @@ -971,6 +1087,8 @@ def input_schema(self) -> dict: "command": {"type": "string", "description": "Bash command to execute"}, "timeout": { "type": "integer", + "minimum": TIMEOUT_MIN_SECONDS, + "maximum": TIMEOUT_MAX_SECONDS, "description": "Command timeout in seconds (default: 30). Increase for builds, tests, or monitoring. Use run_in_background for truly indefinite processes.", }, "run_in_background": { @@ -1010,7 +1128,18 @@ async def execute(self, input: dict[str, Any]) -> ToolResult: success=False, output=error_msg, error={"message": error_msg} ) - timeout = input.get("timeout", self.timeout) + if "timeout" in input: + try: + timeout = _validate_timeout_seconds( + input.get("timeout"), source="caller" + ) + except (TypeError, ValueError) as e: + error_msg = str(e) + return ToolResult( + success=False, output=error_msg, error={"message": error_msg} + ) + else: + timeout = self.timeout run_in_background = input.get("run_in_background", False) # Safety checks using profile-based validator @@ -1543,50 +1672,13 @@ async def _run_command( } except TimeoutError: - # Kill the entire process group (all children) on Unix - if pgid is not None and not is_windows: - # Walk /proc for descendants BEFORE killing anything. - # os.killpg() only reaches processes still in the original - # process group. A descendant that calls setsid (directly, - # or via a wrapper like tmux/incus/docker exec that manages - # its own session lifecycle) moves to a NEW process - # group/session, but its PPID chain back to `process.pid` - # is preserved -- setsid() only changes pgid/sid, it never - # reparents. Capturing descendants up front means we can - # still find them even if an intermediate process in the - # chain is killed first. - descendant_pids = _find_descendant_pids(process.pid) - try: - # Send SIGTERM to process group first (graceful shutdown) - os.killpg(pgid, signal.SIGTERM) - except ProcessLookupError: - pass # Process group already gone - except PermissionError: - # Fall back to killing just the main process - process.kill() - # Belt and suspenders: also signal any descendants that - # escaped the process group and wouldn't receive the - # killpg() above. - _signal_pids(descendant_pids, signal.SIGTERM) - - # Give processes a moment to clean up - await asyncio.sleep(0.5) - - # Force kill if still running - try: - os.killpg(pgid, signal.SIGKILL) - except ProcessLookupError: - pass # Already terminated - except PermissionError: - pass - _signal_pids(descendant_pids, signal.SIGKILL) - else: - # Windows or no pgid: kill just the main process - process.kill() - - # Clean up - try: - await asyncio.wait_for(process.communicate(), timeout=5) - except TimeoutError: - pass # Best effort cleanup + await _await_process_tree_cleanup(process, pgid=pgid, is_windows=is_windows) + raise + except asyncio.CancelledError: + # If our caller cancels the tool call, ensure we still clean up + # the spawned process tree (including the existing Linux /proc + # strategy for setsid()-detached descendants), then re-raise + # cancellation. The shared helper also defers repeated cancellation + # until the bounded cleanup task finishes. + await _await_process_tree_cleanup(process, pgid=pgid, is_windows=is_windows) raise diff --git a/tests/test_timeout_process_cleanup.py b/tests/test_timeout_process_cleanup.py index 4721cb5..f7641ee 100644 --- a/tests/test_timeout_process_cleanup.py +++ b/tests/test_timeout_process_cleanup.py @@ -1,4 +1,4 @@ -"""Regression test: timeout cleanup must kill setsid-detached descendants. +"""Regression tests for timeout and cancellation process cleanup. Root cause (confirmed via minimal reproduction outside this test): @@ -12,23 +12,33 @@ so `os.killpg()` never touches it, and it survives the timeout, running forever as an orphan reparented to PID 1. -These tests are Unix-only (they rely on setsid, /proc, and -os.killpg semantics that don't exist on Windows). +Process-group tests are Unix-only. Tests for setsid-detached descendants +are additionally Linux-only and require /proc plus the setsid executable. """ import asyncio import os import shlex +import shutil import signal import sys import pytest +import amplifier_module_tool_bash from amplifier_module_tool_bash import BashTool pytestmark = pytest.mark.skipif( sys.platform == "win32", reason="Unix process-group semantics only" ) +linux_detached_only = pytest.mark.skipif( + not ( + sys.platform.startswith("linux") + and os.path.isdir("/proc") + and shutil.which("setsid") is not None + ), + reason="Detached-descendant cleanup requires Linux /proc and setsid", +) def _pid_alive(pid: int) -> bool: @@ -51,9 +61,10 @@ def _force_kill(pid: int) -> None: pass -class TestTimeoutKillsSetsidDetachedDescendants: - """Timeout cleanup must reach descendants that escaped the process group.""" +class TestForegroundProcessCleanup: + """Foreground cleanup must terminate the command's process tree.""" + @linux_detached_only @pytest.mark.asyncio async def test_setsid_detached_child_is_killed_on_timeout(self, tmp_path): """A setsid-detached grandchild must not survive timeout cleanup. @@ -161,3 +172,288 @@ async def test_non_detached_sibling_still_killed(self, tmp_path): finally: if plain_pid is not None: _force_kill(plain_pid) + + @pytest.mark.asyncio + async def test_timeout_cleanup_survives_external_cancellation( + self, tmp_path, monkeypatch + ): + """Cancellation during timeout cleanup waits for cleanup, then propagates.""" + parent_marker = tmp_path / "parent.pid" + child_marker = tmp_path / "child.pid" + tool = BashTool({}) + cleanup_started = asyncio.Event() + cleanup_finished = asyncio.Event() + original_cleanup = amplifier_module_tool_bash._cleanup_process_tree + + async def tracked_cleanup(*args, **kwargs): + cleanup_started.set() + try: + await original_cleanup(*args, **kwargs) + finally: + cleanup_finished.set() + + monkeypatch.setattr( + amplifier_module_tool_bash, "_cleanup_process_tree", tracked_cleanup + ) + + # Ignore SIGTERM so cleanup stays in its bounded grace period long + # enough to cancel the outer task after the command timeout fires. + command = ( + f"trap '' TERM; echo $$ > {parent_marker}; " + f"bash -c 'trap \"\" TERM; echo $$ > {child_marker}; sleep 60' & " + "sleep 60" + ) + + parent_pid: int | None = None + child_pid: int | None = None + task = asyncio.create_task(tool._run_command(command, timeout=1)) + try: + for _ in range(50): + if ( + parent_marker.exists() + and child_marker.exists() + and parent_marker.read_text().strip() + and child_marker.read_text().strip() + ): + break + await asyncio.sleep(0.1) + + assert parent_marker.exists() and parent_marker.read_text().strip() + assert child_marker.exists() and child_marker.read_text().strip() + parent_pid = int(parent_marker.read_text().strip()) + child_pid = int(child_marker.read_text().strip()) + + await asyncio.wait_for(cleanup_started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert cleanup_finished.is_set(), ( + "Cancellation propagated before timeout cleanup completed" + ) + + for _ in range(20): + if not _pid_alive(parent_pid) and not _pid_alive(child_pid): + break + await asyncio.sleep(0.1) + + assert not _pid_alive(parent_pid), ( + f"Main process (pid {parent_pid}) survived cancelled timeout cleanup" + ) + assert not _pid_alive(child_pid), ( + f"Child process (pid {child_pid}) survived cancelled timeout cleanup" + ) + finally: + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + for pid in (parent_pid, child_pid): + if pid is not None: + _force_kill(pid) + + @pytest.mark.asyncio + async def test_external_cancellation_kills_process_group(self, tmp_path): + """Ordinary process-group cleanup must work across supported Unix systems.""" + parent_marker = tmp_path / "parent.pid" + child_marker = tmp_path / "child.pid" + tool = BashTool({}) + + command = ( + f"echo $$ > {parent_marker}; " + f"bash -c 'echo $$ > {child_marker}; sleep 60' & " + "sleep 60" + ) + + parent_pid: int | None = None + child_pid: int | None = None + task = asyncio.create_task(tool._run_command(command, timeout=120)) + try: + for _ in range(50): + if ( + parent_marker.exists() + and child_marker.exists() + and parent_marker.read_text().strip() + and child_marker.read_text().strip() + ): + break + await asyncio.sleep(0.1) + + assert parent_marker.exists() and parent_marker.read_text().strip() + assert child_marker.exists() and child_marker.read_text().strip() + parent_pid = int(parent_marker.read_text().strip()) + child_pid = int(child_marker.read_text().strip()) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert not _pid_alive(parent_pid), ( + f"Main process (pid {parent_pid}) survived cancellation cleanup" + ) + assert not _pid_alive(child_pid), ( + f"Child process (pid {child_pid}) survived cancellation cleanup" + ) + finally: + for pid in (parent_pid, child_pid): + if pid is not None: + _force_kill(pid) + + @pytest.mark.asyncio + async def test_repeated_cancellation_waits_for_cleanup(self, tmp_path, monkeypatch): + """A second cancellation must not interrupt the bounded cleanup task.""" + parent_marker = tmp_path / "parent.pid" + child_marker = tmp_path / "child.pid" + tool = BashTool({}) + cleanup_finished = asyncio.Event() + original_cleanup = amplifier_module_tool_bash._cleanup_process_tree + + async def tracked_cleanup(*args, **kwargs): + try: + await original_cleanup(*args, **kwargs) + finally: + cleanup_finished.set() + + monkeypatch.setattr( + amplifier_module_tool_bash, "_cleanup_process_tree", tracked_cleanup + ) + + # Ignore SIGTERM so cleanup cannot complete until its bounded grace + # period ends and it sends SIGKILL. This makes an early return caused + # by the second cancellation observable as still-live processes. + command = ( + f"trap '' TERM; echo $$ > {parent_marker}; " + f"bash -c 'trap \"\" TERM; echo $$ > {child_marker}; sleep 60' & " + "sleep 60" + ) + + parent_pid: int | None = None + child_pid: int | None = None + task = asyncio.create_task(tool._run_command(command, timeout=120)) + try: + for _ in range(50): + if ( + parent_marker.exists() + and child_marker.exists() + and parent_marker.read_text().strip() + and child_marker.read_text().strip() + ): + break + await asyncio.sleep(0.1) + + assert parent_marker.exists() and parent_marker.read_text().strip() + assert child_marker.exists() and child_marker.read_text().strip() + parent_pid = int(parent_marker.read_text().strip()) + child_pid = int(child_marker.read_text().strip()) + + task.cancel() + await asyncio.sleep(0) + assert not task.done(), ( + "First cancellation returned before cleanup finished" + ) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert cleanup_finished.is_set(), ( + "Outer task propagated repeated cancellation before cleanup completed" + ) + + # The main process has been reaped by communicate(); allow an + # orphaned child a brief moment to be reaped by the OS as well. + for _ in range(20): + if not _pid_alive(parent_pid) and not _pid_alive(child_pid): + break + await asyncio.sleep(0.1) + + assert not _pid_alive(parent_pid), ( + f"Main process (pid {parent_pid}) survived repeated cancellation" + ) + assert not _pid_alive(child_pid), ( + f"Child process (pid {child_pid}) survived repeated cancellation" + ) + finally: + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + for pid in (parent_pid, child_pid): + if pid is not None: + _force_kill(pid) + + @linux_detached_only + @pytest.mark.asyncio + async def test_external_cancellation_kills_setsid_detached_descendant( + self, tmp_path + ): + """If the coroutine is externally cancelled, the spawned process tree + (including setsid-detached descendants) must still be cleaned up.""" + + parent_marker = tmp_path / "parent.pid" + child_marker = tmp_path / "child.pid" + detached_marker = tmp_path / "detached.pid" + tool = BashTool({}) + + # Parent writes its own PID; it then spawns: + # - a normal background child (same process group) + # - a setsid-detached background child (new pgid/sid) + # Then sleeps so the process tree stays alive until we cancel. + command = ( + f"echo $$ > {parent_marker}; " + f"bash -c 'echo $$ > {child_marker}; sleep 60' & " + f"setsid bash -c 'echo $$ > {detached_marker}; sleep 60' & " + "sleep 60" + ) + + parent_pid: int | None = None + child_pid: int | None = None + detached_pid: int | None = None + + task = asyncio.create_task(tool._run_command(command, timeout=120)) + try: + # Wait for all PID markers to appear. + for _ in range(50): + if ( + parent_marker.exists() + and child_marker.exists() + and detached_marker.exists() + and parent_marker.read_text().strip() + and child_marker.read_text().strip() + and detached_marker.read_text().strip() + ): + break + await asyncio.sleep(0.1) + + assert parent_marker.exists() and parent_marker.read_text().strip() + assert child_marker.exists() and child_marker.read_text().strip() + assert detached_marker.exists() and detached_marker.read_text().strip() + + parent_pid = int(parent_marker.read_text().strip()) + child_pid = int(child_marker.read_text().strip()) + detached_pid = int(detached_marker.read_text().strip()) + + # Cancel the tool call while it is waiting for process completion. + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(0.5) + + assert not _pid_alive(parent_pid), ( + f"Main process (pid {parent_pid}) survived external cancellation cleanup" + ) + assert not _pid_alive(child_pid), ( + f"Child process (pid {child_pid}) survived external cancellation cleanup" + ) + assert not _pid_alive(detached_pid), ( + f"Detached setsid child (pid {detached_pid}) survived external cancellation cleanup" + ) + finally: + for pid in (parent_pid, child_pid, detached_pid): + if pid is not None: + _force_kill(pid) diff --git a/tests/test_validation.py b/tests/test_validation.py index 52f27e4..107e70e 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -3,8 +3,11 @@ Inherits authoritative tests from amplifier-core. """ +import pytest from amplifier_core.validation.structural import ToolStructuralTests +from amplifier_module_tool_bash import BashTool + class TestBashToolStructural(ToolStructuralTests): """Run standard tool structural tests for bash. @@ -12,3 +15,66 @@ class TestBashToolStructural(ToolStructuralTests): All tests from ToolStructuralTests run automatically. Add module-specific structural tests below if needed. """ + + +class TestTimeoutValidation: + def test_input_schema_includes_timeout_bounds(self): + tool = BashTool({}) + timeout_schema = tool.input_schema["properties"]["timeout"] + assert timeout_schema["minimum"] == 1 + assert timeout_schema["maximum"] == 3600 + + @pytest.mark.parametrize( + "bad_timeout", + [ + True, # bool must be rejected (bool is an int subclass in Python) + 0, + -1, + 3601, + 1.5, + "30", + None, + ], + ) + def test_config_timeout_rejected(self, bad_timeout): + with pytest.raises((TypeError, ValueError), match=r"seconds"): + BashTool({"timeout": bad_timeout}) + + @pytest.mark.parametrize( + "bad_timeout", + [ + True, + 0, + -1, + 3601, + 1.5, + "30", + None, + ], + ) + @pytest.mark.asyncio + async def test_caller_timeout_rejected(self, bad_timeout): + tool = BashTool({}) + result = await tool.execute({"command": "echo ok", "timeout": bad_timeout}) + assert result.success is False + assert isinstance(result.output, str) + assert "seconds" in result.output + + @pytest.mark.asyncio + async def test_caller_timeout_rejects_ms_looking_value_with_hint(self): + tool = BashTool({}) + result = await tool.execute({"command": "echo ok", "timeout": 1_200_000}) + assert result.success is False + assert isinstance(result.output, str) + assert "seconds" in result.output + # Regression: suggest a plausible seconds value if caller likely passed ms. + assert "1200" in result.output + + @pytest.mark.asyncio + async def test_execute_timeout_still_works_with_valid_value(self): + tool = BashTool({}) + result = await tool.execute({"command": "echo ok", "timeout": 5}) + assert result.success is True + assert isinstance(result.output, dict) + assert result.output["returncode"] == 0 + assert "ok" in result.output["stdout"]