mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-12 23:29:16 +03:00
fix(exec): terminate one-shot process trees on cleanup
Run one-shot commands in their own process tree and terminate all descendants after timeout, cancellation, or unexpected failures. Co-authored-by: TRAE CLI <noreply@bytedance.com>
This commit is contained in:
committed by
chengyongru
co-authored by
TRAE CLI
parent
abfcdd481a
commit
d64b84604c
@@ -326,6 +326,7 @@ class ExecTool(Tool):
|
||||
prepared.env,
|
||||
prepared.shell_program,
|
||||
prepared.login,
|
||||
process_tree=True,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -334,10 +335,10 @@ class ExecTool(Tool):
|
||||
timeout=prepared.timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
await self._kill_process(process)
|
||||
await self._kill_process_tree(process)
|
||||
return ToolResult.error(f"Error: Command timed out after {prepared.timeout} seconds")
|
||||
except asyncio.CancelledError:
|
||||
await self._kill_process(process)
|
||||
await self._kill_process_tree(process)
|
||||
raise
|
||||
|
||||
# Safety-net reap: asyncio *should* have reaped the child via
|
||||
@@ -374,7 +375,7 @@ class ExecTool(Tool):
|
||||
# Kill and reap the child if it was spawned but an unexpected
|
||||
# error prevented communicate() from completing.
|
||||
if process is not None:
|
||||
await self._kill_process(process)
|
||||
await self._kill_process_tree(process)
|
||||
return ToolResult.error(f"Error executing command: {str(e)}")
|
||||
|
||||
async def _execute_session(
|
||||
|
||||
@@ -302,7 +302,9 @@ class TestPathAppendPlatform:
|
||||
captured_cmd = None
|
||||
captured_env = {}
|
||||
|
||||
async def capture_spawn(cmd, cwd, env, shell_program=None, login=True):
|
||||
async def capture_spawn(
|
||||
cmd, cwd, env, shell_program=None, login=True, *, process_tree=False,
|
||||
):
|
||||
nonlocal captured_cmd
|
||||
captured_cmd = cmd
|
||||
captured_env.update(env)
|
||||
@@ -331,7 +333,9 @@ class TestPathAppendPlatform:
|
||||
captured_cmd = None
|
||||
captured_env = {}
|
||||
|
||||
async def capture_spawn(cmd, cwd, env, shell_program=None, login=True, *, stdin=None):
|
||||
async def capture_spawn(
|
||||
cmd, cwd, env, shell_program=None, login=True, *, stdin=None, process_tree=False,
|
||||
):
|
||||
nonlocal captured_cmd
|
||||
captured_cmd = cmd
|
||||
captured_env.update(env)
|
||||
@@ -359,7 +363,9 @@ class TestPathAppendPlatform:
|
||||
captured_cmd = None
|
||||
captured_env = {}
|
||||
|
||||
async def capture_spawn(cmd, cwd, env, shell_program=None, login=True, *, stdin=None):
|
||||
async def capture_spawn(
|
||||
cmd, cwd, env, shell_program=None, login=True, *, stdin=None, process_tree=False,
|
||||
):
|
||||
nonlocal captured_cmd
|
||||
captured_cmd = cmd
|
||||
captured_env.update(env)
|
||||
@@ -389,7 +395,9 @@ class TestPathAppendPlatform:
|
||||
|
||||
captured_env = {}
|
||||
|
||||
async def capture_spawn(cmd, cwd, env, shell_program=None, login=True):
|
||||
async def capture_spawn(
|
||||
cmd, cwd, env, shell_program=None, login=True, *, process_tree=False,
|
||||
):
|
||||
captured_env.update(env)
|
||||
return mock_proc
|
||||
|
||||
@@ -412,7 +420,9 @@ class TestPathAppendPlatform:
|
||||
|
||||
captured_env = {}
|
||||
|
||||
async def capture_spawn(cmd, cwd, env, shell_program=None, login=True, *, stdin=None):
|
||||
async def capture_spawn(
|
||||
cmd, cwd, env, shell_program=None, login=True, *, stdin=None, process_tree=False,
|
||||
):
|
||||
captured_env.update(env)
|
||||
return mock_proc
|
||||
|
||||
@@ -558,7 +568,9 @@ class TestExecuteEndToEnd:
|
||||
mock_proc.returncode = 0
|
||||
captured_login = []
|
||||
|
||||
async def capture_spawn(cmd, cwd, env, shell_program=None, login=None, *, stdin=None):
|
||||
async def capture_spawn(
|
||||
cmd, cwd, env, shell_program=None, login=None, *, stdin=None, process_tree=False,
|
||||
):
|
||||
captured_login.append(login)
|
||||
return mock_proc
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shlex
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
@@ -112,20 +113,37 @@ async def test_execute_timeout_kills_and_reaps():
|
||||
mock_proc.pid = 1002
|
||||
mock_proc.returncode = None
|
||||
mock_proc.communicate.side_effect = asyncio.TimeoutError()
|
||||
mock_proc.kill = MagicMock()
|
||||
mock_proc.wait = AsyncMock(return_value=-9)
|
||||
|
||||
with (
|
||||
patch.object(ExecTool, "_spawn", return_value=mock_proc),
|
||||
patch.object(ExecTool, "_spawn", return_value=mock_proc) as spawn,
|
||||
patch.object(ExecTool, "_guard_command", return_value=None),
|
||||
patch("nanobot.agent.tools.shell._reap_pid") as reap,
|
||||
patch.object(ExecTool, "_kill_process_tree", new_callable=AsyncMock) as kill_tree,
|
||||
):
|
||||
tool = ExecTool(timeout=1)
|
||||
result = await tool.execute(command="sleep 99", timeout=1)
|
||||
|
||||
assert "timed out" in result.lower()
|
||||
mock_proc.kill.assert_called_once()
|
||||
reap.assert_called_with(1002)
|
||||
kill_tree.assert_awaited_once_with(mock_proc)
|
||||
assert spawn.await_args.kwargs["process_tree"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_cancellation_kills_process_tree():
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.pid = 1005
|
||||
mock_proc.returncode = None
|
||||
mock_proc.communicate.side_effect = asyncio.CancelledError()
|
||||
|
||||
with (
|
||||
patch.object(ExecTool, "_spawn", return_value=mock_proc) as spawn,
|
||||
patch.object(ExecTool, "_guard_command", return_value=None),
|
||||
patch.object(ExecTool, "_kill_process_tree", new_callable=AsyncMock) as kill_tree,
|
||||
):
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await ExecTool().execute(command="sleep 99")
|
||||
|
||||
kill_tree.assert_awaited_once_with(mock_proc)
|
||||
assert spawn.await_args.kwargs["process_tree"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -161,21 +179,35 @@ async def test_execute_exception_during_communicate_kills_live_process():
|
||||
mock_proc.pid = 1004
|
||||
mock_proc.returncode = None
|
||||
mock_proc.communicate.side_effect = OSError("pipe broken")
|
||||
mock_proc.kill = MagicMock()
|
||||
mock_proc.wait = AsyncMock(return_value=-1)
|
||||
|
||||
with (
|
||||
patch.object(ExecTool, "_spawn", return_value=mock_proc),
|
||||
patch.object(ExecTool, "_guard_command", return_value=None),
|
||||
patch("nanobot.agent.tools.shell._reap_pid") as reap,
|
||||
patch.object(ExecTool, "_kill_process_tree", new_callable=AsyncMock) as kill_tree,
|
||||
):
|
||||
tool = ExecTool()
|
||||
result = await tool.execute(command="broken")
|
||||
|
||||
assert "Error executing command" in result
|
||||
assert "pipe broken" in result
|
||||
mock_proc.kill.assert_called_once()
|
||||
reap.assert_called_with(1004)
|
||||
kill_tree.assert_awaited_once_with(mock_proc)
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason="requires Unix process groups")
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_timeout_kills_background_process_tree(tmp_path):
|
||||
"""A one-shot timeout must stop background descendants before they write."""
|
||||
marker = tmp_path / "child-survived"
|
||||
command = f"(sleep 2; touch {shlex.quote(str(marker))}) >/dev/null 2>&1 & sleep 30"
|
||||
|
||||
result = await ExecTool(working_dir=str(tmp_path), timeout=1).execute(
|
||||
command=command,
|
||||
timeout=1,
|
||||
)
|
||||
|
||||
assert "timed out" in result.lower()
|
||||
await asyncio.sleep(2.5)
|
||||
assert not marker.exists()
|
||||
|
||||
|
||||
def _mock_session_process(*, pid: int, returncode: int | None):
|
||||
|
||||
Reference in New Issue
Block a user