From 452307449efda20de5efea77511102f6c1a2ba36 Mon Sep 17 00:00:00 2001 From: "hanzhi.421" Date: Tue, 22 Sep 2026 23:25:09 +0800 Subject: [PATCH] fix(skills): clean up bash process groups on timeout and cancellation --- tests/tools/test_skills_bash_tool.py | 195 ++++++++++++++++++++++++++ veadk/tools/skills_tools/bash_tool.py | 64 +++++++-- 2 files changed, 245 insertions(+), 14 deletions(-) create mode 100644 tests/tools/test_skills_bash_tool.py diff --git a/tests/tools/test_skills_bash_tool.py b/tests/tools/test_skills_bash_tool.py new file mode 100644 index 000000000..a361f9229 --- /dev/null +++ b/tests/tools/test_skills_bash_tool.py @@ -0,0 +1,195 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import importlib +import os +import signal +import subprocess +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +bash_module = importlib.import_module("veadk.tools.skills_tools.bash_tool") +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.skipif(os.name != "posix", reason="POSIX process groups"), +] + + +@pytest.fixture +def execution(tmp_path, monkeypatch): + monkeypatch.setattr(bash_module, "get_session_path", lambda **kwargs: tmp_path) + context = SimpleNamespace(session=SimpleNamespace(id="timeout-test")) + processes = [] + original = asyncio.create_subprocess_shell + + async def spawn(*args, **kwargs): + process = await original(*args, **kwargs) + processes.append(process) + return process + + monkeypatch.setattr(asyncio, "create_subprocess_shell", spawn) + + async def run(command, timeout=5): + return await bash_module.bash_tool( + command, "timeout regression", context, timeout + ) + + yield run, processes + # Independent cleanup also makes a regression failure safe to run locally. + for process in processes: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + + +def live_group(pgid): + rows = subprocess.check_output( + ["ps", "-eo", "pgid=,stat=,comm="], text=True + ).splitlines() + return [ + row + for row in rows + if row.split()[0] == str(pgid) and not row.split()[1].startswith("Z") + ] + + +async def test_pipeline_timeout_cleans_children_and_allows_next_call(execution): + run, processes = execution + result = await asyncio.wait_for(run("sleep 900 | cat", timeout=5), 12) + assert result == "Error: Command timed out after 5s" + assert not live_group(processes[0].pid) + assert ( + await asyncio.wait_for(run("echo TOOL_TIMEOUT_RECOVERED"), 7) + == "TOOL_TIMEOUT_RECOVERED" + ) + + +async def test_timeout_cleans_group_after_shell_has_already_exited(execution): + run, processes = execution + assert ( + await asyncio.wait_for(run("sleep 900 & exit 0", timeout=0.2), 7) + == "Error: Command timed out after 0.2s" + ) + assert processes[0].returncode == 0 + assert not live_group(processes[0].pid) + + +async def test_cancellation_cleans_pipeline_and_propagates(execution): + run, processes = execution + task = asyncio.create_task(run("sleep 900 | cat", timeout=600)) + for _ in range(100): + if processes: + break + await asyncio.sleep(0.01) + assert processes + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 7) + assert not live_group(processes[0].pid) + assert await run("echo recovered") == "recovered" + + +async def test_cancellation_during_spawn_still_cleans_children(execution, monkeypatch): + run, processes = execution + spawn = asyncio.create_subprocess_shell + started = asyncio.Event() + release = asyncio.Event() + + async def delayed_spawn(*args, **kwargs): + process = await spawn(*args, **kwargs) + started.set() + await release.wait() + return process + + monkeypatch.setattr(asyncio, "create_subprocess_shell", delayed_spawn) + task = asyncio.create_task(run("sleep 900 | cat", timeout=600)) + await asyncio.wait_for(started.wait(), 3) + task.cancel() + await asyncio.sleep(0) + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 7) + assert not live_group(processes[0].pid) + + +async def test_cleanup_wait_is_bounded(execution, monkeypatch): + run, _ = execution + + async def stuck(): + await asyncio.Event().wait() + + process = SimpleNamespace( + pid=12345, + communicate=AsyncMock(side_effect=stuck), + wait=AsyncMock(side_effect=stuck), + _transport=Mock(), + ) + monkeypatch.setattr( + asyncio, "create_subprocess_shell", AsyncMock(return_value=process) + ) + kill = Mock() + monkeypatch.setattr(os, "killpg", kill) + monkeypatch.setattr(bash_module, "_PROCESS_CLEANUP_TIMEOUT", 0.05) + assert ( + await asyncio.wait_for(run("test", timeout=0.05), 1) + == "Error: Command timed out after 0.05s" + ) + kill.assert_called_once_with(process.pid, signal.SIGKILL) + process._transport.close.assert_called_once() + + +async def test_process_lookup_race_preserves_timeout(execution, monkeypatch): + run, _ = execution + process = SimpleNamespace( + pid=12345, + communicate=AsyncMock(side_effect=asyncio.TimeoutError), + wait=AsyncMock(return_value=0), + ) + monkeypatch.setattr( + asyncio, "create_subprocess_shell", AsyncMock(return_value=process) + ) + monkeypatch.setattr(os, "killpg", Mock(side_effect=ProcessLookupError)) + assert await run("test") == "Error: Command timed out after 5s" + process.wait.assert_awaited_once() + + +@pytest.mark.parametrize( + "command,expected", + [ + ("", "Error: No command provided"), + ("printf ok", "ok"), + ("true", "Command completed successfully."), + ("printf error >&2; exit 7", "Command failed with exit code 7:\nerror"), + ("printf output; exit 2", "Command failed with exit code 2:\noutput"), + ("printf ok; printf WARNING >&2", "ok"), + ], +) +async def test_output_compatibility(execution, command, expected): + run, _ = execution + assert await run(command) == expected + + +async def test_spawn_failure_is_returned_as_tool_error(execution, monkeypatch): + run, _ = execution + monkeypatch.setattr( + asyncio, + "create_subprocess_shell", + AsyncMock(side_effect=OSError("spawn failed")), + ) + assert await run("test") == "Error executing command 'test': spawn failed" diff --git a/veadk/tools/skills_tools/bash_tool.py b/veadk/tools/skills_tools/bash_tool.py index 56c2b9afb..7104d4629 100644 --- a/veadk/tools/skills_tools/bash_tool.py +++ b/veadk/tools/skills_tools/bash_tool.py @@ -16,12 +16,57 @@ import asyncio import os +import signal from google.adk.tools import ToolContext from veadk.tools.skills_tools.session_path import get_session_path from veadk.utils.logger import get_logger logger = get_logger(__name__) +_PROCESS_CLEANUP_TIMEOUT = 5.0 + + +async def _run_command(command, working_dir, env, timeout): + process = None + completed = False + spawning = asyncio.create_task( + asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=working_dir, + env=env, + start_new_session=os.name == "posix", + ) + ) + try: + try: + process = await asyncio.shield(spawning) + except asyncio.CancelledError: + # Cancellation can arrive after fork but before the Process is + # returned. Retain ownership so that its children are cleaned up. + process = await spawning + raise + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout) + completed = True + return process.returncode, stdout, stderr + finally: + if process is not None and not completed: + try: + if os.name == "posix": + # Do this even if the shell has exited: its descendants + # can still hold stdout/stderr open and block communicate. + os.killpg(process.pid, signal.SIGKILL) + else: + process.kill() + except ProcessLookupError: + pass + try: + await asyncio.wait_for(process.wait(), _PROCESS_CLEANUP_TIMEOUT) + except asyncio.TimeoutError: + # A detached descendant may keep a pipe open. Process has no + # public close API; release the local transports on this path. + process._transport.close() async def bash_tool( @@ -61,6 +106,7 @@ async def bash_tool( Timeouts: - Default timeout is 600 seconds (10 minutes). - Adjust the 'timeout' parameter as needed for longer-running commands. + - Interrupted commands have up to 5 additional seconds for process cleanup. Args: command: Bash command to execute. Use && to chain commands. @@ -102,29 +148,19 @@ async def bash_tool( # Execute with local bash shell local_bash_command = f"{command}" - process = await asyncio.create_subprocess_shell( - local_bash_command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=working_dir, - env=env, # Pass the modified environment - ) - try: - stdout, stderr = await asyncio.wait_for( - process.communicate(), timeout=timeout + returncode, stdout, stderr = await _run_command( + local_bash_command, working_dir, env, timeout ) except asyncio.TimeoutError: - process.kill() - await process.wait() return f"Error: Command timed out after {timeout}s" stdout_str = stdout.decode("utf-8", errors="replace") if stdout else "" stderr_str = stderr.decode("utf-8", errors="replace") if stderr else "" # Handle command failure - if process.returncode != 0: - error_msg = f"Command failed with exit code {process.returncode}" + if returncode != 0: + error_msg = f"Command failed with exit code {returncode}" if stderr_str: error_msg += f":\n{stderr_str}" elif stdout_str: