diff --git a/backend/api/server.py b/backend/api/server.py index 64b7d89..ba14413 100644 --- a/backend/api/server.py +++ b/backend/api/server.py @@ -8,9 +8,11 @@ from __future__ import annotations import asyncio import json +import os import queue import re import shutil +import signal import threading import time import venv @@ -270,7 +272,7 @@ class RunProcessRegistry: for process in processes: try: if process.poll() is None: - process.terminate() + self._signal_process_tree(process, signal.SIGTERM) except OSError: continue deadline = time.monotonic() + 1.0 @@ -279,11 +281,35 @@ class RunProcessRegistry: remaining = max(0.0, deadline - time.monotonic()) process.wait(timeout=remaining) except Exception: + self._signal_process_tree(process, signal.SIGKILL) try: - process.kill() - except OSError: + process.wait(timeout=1.0) + except Exception: pass + @staticmethod + def _signal_process_tree(process: Any, sig: signal.Signals) -> None: + try: + if os.name == 'posix': + pgid = os.getpgid(process.pid) + if pgid == process.pid: + os.killpg(pgid, sig) + return + if sig == signal.SIGTERM: + process.terminate() + else: + process.kill() + except ProcessLookupError: + return + except OSError: + try: + if sig == signal.SIGTERM: + process.terminate() + else: + process.kill() + except OSError: + pass + @dataclass class RunRecord: diff --git a/src/agent_tools.py b/src/agent_tools.py index 363cdb1..1f65992 100644 --- a/src/agent_tools.py +++ b/src/agent_tools.py @@ -6,6 +6,7 @@ import os import re import selectors import shutil +import signal import subprocess import sys import time @@ -2087,9 +2088,9 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) - stderr=subprocess.PIPE, text=True, env=_build_subprocess_env(context), + **_subprocess_isolation_kwargs(), ) _register_child_process(context, process) - _register_child_process(context, process) if stdin and process.stdin is not None: process.stdin.write(stdin) process.stdin.close() @@ -2108,7 +2109,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) - except subprocess.TimeoutExpired: if process is not None: _terminate_process(process) - stdout, stderr = process.communicate() + stdout, stderr = _communicate_after_stop(process) else: stdout, stderr = '', '' payload = [ @@ -2200,19 +2201,67 @@ def _unregister_child_process( registry.discard(process) +def _subprocess_isolation_kwargs() -> dict[str, Any]: + if os.name == 'posix': + return {'start_new_session': True} + creation_flags = getattr(subprocess, 'CREATE_NEW_PROCESS_GROUP', 0) + if creation_flags: + return {'creationflags': creation_flags} + return {} + + def _terminate_process(process: subprocess.Popen[Any]) -> None: if process.poll() is not None: return + _signal_process_tree(process, signal.SIGTERM) + try: + process.wait(timeout=1) + return + except (OSError, subprocess.TimeoutExpired): + pass + _signal_process_tree(process, signal.SIGKILL) try: - process.terminate() process.wait(timeout=1) except (OSError, subprocess.TimeoutExpired): - try: + pass + + +def _signal_process_tree(process: subprocess.Popen[Any], sig: signal.Signals) -> None: + try: + if os.name == 'posix': + pgid = os.getpgid(process.pid) + if pgid == process.pid: + os.killpg(pgid, sig) + return + if sig == signal.SIGTERM: + process.terminate() + else: process.kill() + except ProcessLookupError: + return + except OSError: + try: + if sig == signal.SIGTERM: + process.terminate() + else: + process.kill() except OSError: pass +def _communicate_after_stop(process: subprocess.Popen[str]) -> tuple[str, str]: + try: + stdout, stderr = process.communicate(timeout=1) + return stdout or '', stderr or '' + except subprocess.TimeoutExpired: + _terminate_process(process) + try: + stdout, stderr = process.communicate(timeout=1) + return stdout or '', stderr or '' + except (OSError, subprocess.TimeoutExpired): + return '', '' + + def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext) -> str: _ensure_process_execution_allowed(context, 'Python package management') action = _require_string(arguments, 'action') @@ -3633,7 +3682,9 @@ def _stream_bash( text=True, bufsize=1, env=_build_subprocess_env(context), + **_subprocess_isolation_kwargs(), ) + _register_child_process(context, process) except (ToolPermissionError, ToolExecutionError, OSError, subprocess.SubprocessError) as exc: yield ToolStreamUpdate( kind='result', @@ -3663,7 +3714,7 @@ def _stream_bash( timeout_error = ( f'Command timed out after {context.command_timeout_seconds:.1f}s: {command}' ) - process.kill() + _terminate_process(process) break events = selector.select(timeout=min(remaining, 0.1)) if not events and process.poll() is not None: diff --git a/tests/test_python_exec_tool.py b/tests/test_python_exec_tool.py index c7ce5e9..69c0ac6 100644 --- a/tests/test_python_exec_tool.py +++ b/tests/test_python_exec_tool.py @@ -1,7 +1,12 @@ from __future__ import annotations +import os +import re +import signal +import subprocess import sys import tempfile +import time from pathlib import Path from unittest import TestCase @@ -179,3 +184,54 @@ class PythonExecToolTests(TestCase): self.assertTrue(result.ok) self.assertIn('a|b', result.content) self.assertEqual(result.metadata.get('mode'), 'script') + + def test_python_exec_timeout_kills_spawned_children(self) -> None: + if os.name != 'posix': + self.skipTest('process group cleanup test is POSIX-only') + + with tempfile.TemporaryDirectory() as tmp_dir: + config = AgentRuntimeConfig( + cwd=Path(tmp_dir), + permissions=AgentPermissions(allow_shell_commands=True), + ) + context = build_tool_context(config) + child_code = ( + 'import subprocess, sys, time\n' + 'child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"])\n' + 'print(child.pid, flush=True)\n' + 'time.sleep(60)\n' + ) + result = execute_tool( + default_tool_registry(), + 'python_exec', + { + 'code': child_code, + 'timeout_seconds': 1, + 'max_output_chars': 2000, + }, + context, + ) + + self.assertTrue(result.ok) + self.assertTrue(result.metadata.get('timed_out')) + match = re.search(r'\[stdout\]\s*(\d+)', result.content) + self.assertIsNotNone(match, result.content) + child_pid = int(match.group(1)) + time.sleep(0.3) + if self._pid_is_running(child_pid): + try: + os.kill(child_pid, signal.SIGKILL) + except ProcessLookupError: + pass + self.fail(f'spawned child process is still running: {child_pid}') + + @staticmethod + def _pid_is_running(pid: int) -> bool: + status = subprocess.run( + ['ps', '-p', str(pid), '-o', 'stat='], + capture_output=True, + text=True, + check=False, + ) + state = status.stdout.strip() + return bool(state) and not state.startswith('Z')