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:
yorkhellen
2026-08-12 14:37:27 +08:00
committed by chengyongru
co-authored by TRAE CLI
parent abfcdd481a
commit d64b84604c
3 changed files with 65 additions and 20 deletions
+4 -3
View File
@@ -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(
+18 -6
View File
@@ -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
+43 -11
View File
@@ -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):