Fix process tree cancellation for tools
This commit is contained in:
+27
-1
@@ -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,7 +281,31 @@ class RunProcessRegistry:
|
||||
remaining = max(0.0, deadline - time.monotonic())
|
||||
process.wait(timeout=remaining)
|
||||
except Exception:
|
||||
self._signal_process_tree(process, signal.SIGKILL)
|
||||
try:
|
||||
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
|
||||
|
||||
+55
-4
@@ -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):
|
||||
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:
|
||||
|
||||
@@ -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')
|
||||
|
||||
Reference in New Issue
Block a user