diff --git a/src/google/adk/tools/bash_tool.py b/src/google/adk/tools/bash_tool.py index aa3571c168e..f16f1bf9a87 100644 --- a/src/google/adk/tools/bash_tool.py +++ b/src/google/adk/tools/bash_tool.py @@ -58,6 +58,68 @@ def _load_resource_module() -> Optional[_ResourceModule]: _resource = _load_resource_module() +_DEFAULT_MAX_OUTPUT_BYTES = 1024 * 1024 +_READ_CHUNK_SIZE = 64 * 1024 +# Grace period to drain pipes after the process group has been killed. +_DRAIN_GRACE_SECONDS = 1.0 + + +class _BoundedOutput: + """Accumulates a stream's output up to a byte limit.""" + + def __init__(self, limit: Optional[int]): + self._limit = limit + self._data = bytearray() + self.total_bytes = 0 + + def append(self, chunk: bytes) -> None: + self.total_bytes += len(chunk) + if self._limit is None: + self._data.extend(chunk) + return + remaining = self._limit - len(self._data) + if remaining > 0: + self._data.extend(chunk[:remaining]) + + @property + def truncated_bytes(self) -> int: + return self.total_bytes - len(self._data) + + def render(self, stream_name: str) -> str: + if not self.total_bytes: + return f"" + text = bytes(self._data).decode(errors="replace") + if self.truncated_bytes: + text += ( + f"\n...[{stream_name} truncated: {self.truncated_bytes} of" + f" {self.total_bytes} bytes omitted; limit is {self._limit}" + " bytes]" + ) + return text + + +async def _drain_stream( + stream: Optional[asyncio.StreamReader], output: _BoundedOutput +) -> None: + """Reads a stream to EOF in fixed-size chunks into a bounded buffer.""" + if stream is None: + return + while True: + chunk = await stream.read(_READ_CHUNK_SIZE) + if not chunk: + return + output.append(chunk) + + +async def _finish_tasks(tasks: list[asyncio.Task]) -> None: + """Waits briefly for tasks to finish, then cancels any stragglers.""" + if not tasks: + return + _, pending = await asyncio.wait(tasks, timeout=_DRAIN_GRACE_SECONDS) + for task in pending: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + @dataclasses.dataclass(frozen=True) class BashToolPolicy: @@ -75,6 +137,12 @@ class BashToolPolicy: Values for max_memory_bytes, max_file_size_bytes, and max_child_processes will be enforced upon the spawned subprocess. + + max_output_bytes_per_stream bounds how many bytes of stdout and of stderr + are retained in the agent process. Output beyond the limit is still read + (so the subprocess never blocks on a full pipe) but discarded, and the + returned text ends with a marker reporting how many bytes were truncated. + Set it to None to disable the limit. """ allowed_command_prefixes: tuple[str, ...] = ("*",) @@ -83,8 +151,16 @@ class BashToolPolicy: max_memory_bytes: Optional[int] = None max_file_size_bytes: Optional[int] = None max_child_processes: Optional[int] = None + max_output_bytes_per_stream: Optional[int] = _DEFAULT_MAX_OUTPUT_BYTES def __post_init__(self) -> None: + if ( + self.max_output_bytes_per_stream is not None + and self.max_output_bytes_per_stream < 0 + ): + raise ValueError( + "max_output_bytes_per_stream must be >= 0, or None to disable" + ) for prefix in self.allowed_command_prefixes: try: shlex.split(prefix) @@ -245,8 +321,11 @@ async def run_async( if os.name != "posix": return {"error": "ExecuteBashTool is only supported on POSIX systems."} - stdout = None - stderr = None + stdout = _BoundedOutput(self._policy.max_output_bytes_per_stream) + stderr = _BoundedOutput(self._policy.max_output_bytes_per_stream) + readers: list[asyncio.Task] = [] + waiter: Optional[asyncio.Task] = None + process: Optional[asyncio.subprocess.Process] = None try: process = await asyncio.create_subprocess_exec( *shlex.split(command), @@ -256,35 +335,44 @@ async def run_async( start_new_session=True, preexec_fn=lambda: _set_resource_limits(self._policy), ) - + readers = [ + asyncio.create_task(_drain_stream(process.stdout, stdout)), + asyncio.create_task(_drain_stream(process.stderr, stderr)), + ] + # Explicit waiter instead of gather+shield: on 3.14 shield logs + # CancelledError, and a gather's process.wait() can stay pending when + # a setsid descendant keeps a pipe open after timeout cancel. + waiter = asyncio.create_task(process.wait()) + tracked = [*readers, waiter] try: - stdout, stderr = await asyncio.wait_for( - process.communicate(), timeout=self._policy.timeout_seconds + _done, pending = await asyncio.wait( + tracked, + timeout=self._policy.timeout_seconds, + return_when=asyncio.ALL_COMPLETED, ) - except asyncio.TimeoutError: - try: - if process.pid: - _kill_process_group(process) - except ProcessLookupError: - pass - stdout, stderr = await process.communicate() - return { - "error": ( - f"Command timed out after {self._policy.timeout_seconds}" - " seconds." - ), - "stdout": ( - stdout.decode(errors="replace") - if stdout - else "" - ), - "stderr": ( - stderr.decode(errors="replace") - if stderr - else "" - ), - "returncode": process.returncode, - } + if pending: + try: + if process.pid: + _kill_process_group(process) + except ProcessLookupError: + pass + await _finish_tasks(tracked) + if process.returncode is None: + try: + await asyncio.wait_for( + process.wait(), timeout=_DRAIN_GRACE_SECONDS + ) + except asyncio.TimeoutError: + pass + return { + "error": ( + f"Command timed out after {self._policy.timeout_seconds}" + " seconds." + ), + "stdout": stdout.render("stdout"), + "stderr": stderr.render("stderr"), + "returncode": process.returncode, + } finally: try: if process.pid: @@ -292,32 +380,33 @@ async def run_async( except ProcessLookupError: pass return { - "stdout": ( - stdout.decode(errors="replace") - if stdout - else "" - ), - "stderr": ( - stderr.decode(errors="replace") - if stderr - else "" - ), + "stdout": stdout.render("stdout"), + "stderr": stderr.render("stderr"), "returncode": process.returncode, } + except asyncio.CancelledError: + # Cancellation must not leave drain/wait tasks pending when a + # descendant still holds a pipe open after the process group kill. + try: + if process is not None and process.pid: + _kill_process_group(process) + except ProcessLookupError: + pass + tracked = [t for t in (*readers, waiter) if t is not None] + for task in tracked: + task.cancel() + if tracked: + await asyncio.gather(*tracked, return_exceptions=True) + raise except Exception as e: # pylint: disable=broad-except logger.exception("ExecuteBashTool execution failed") - - stdout_res = ( - stdout.decode(errors="replace") if stdout else "" - ) - stderr_res = ( - stderr.decode(errors="replace") if stderr else "" - ) - + tracked = [t for t in (*readers, waiter) if t is not None] + if tracked: + await _finish_tasks(tracked) return { "error": f"Execution failed: {str(e)}", - "stdout": stdout_res, - "stderr": stderr_res, + "stdout": stdout.render("stdout"), + "stderr": stderr.render("stderr"), } def _detect_error_in_response(self, response: Any) -> Optional[str]: diff --git a/tests/unittests/tools/test_bash_tool.py b/tests/unittests/tools/test_bash_tool.py index c858bd214d7..71340c163ec 100644 --- a/tests/unittests/tools/test_bash_tool.py +++ b/tests/unittests/tools/test_bash_tool.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import gc import logging import signal import sys @@ -294,10 +295,28 @@ async def test_nonzero_returncode(self, workspace, tool_context_confirmed): @pytest.mark.asyncio async def test_timeout(self, workspace, tool_context_confirmed): - tool = bash_tool.ExecuteBashTool(workspace=workspace) + tool = bash_tool.ExecuteBashTool( + workspace=workspace, + policy=bash_tool.BashToolPolicy(timeout_seconds=0), + ) + exit_fut = asyncio.get_running_loop().create_future() mock_process = mock.AsyncMock() mock_process.pid = 12345 - mock_process.communicate.return_value = (b"", b"") + mock_process.stdout = None + mock_process.stderr = None + mock_process.returncode = None + + async def _wait(): + if mock_process.returncode is not None: + return mock_process.returncode + return await exit_fut + + def _killpg(pid, sig): + mock_process.returncode = -sig + if not exit_fut.done(): + exit_fut.set_result(-sig) + + mock_process.wait = _wait with ( mock.patch.object( asyncio, @@ -305,10 +324,7 @@ async def test_timeout(self, workspace, tool_context_confirmed): autospec=True, return_value=mock_process, ), - mock.patch.object( - asyncio, "wait_for", autospec=True, side_effect=asyncio.TimeoutError - ), - mock.patch("os.killpg") as mock_killpg, + mock.patch("os.killpg", side_effect=_killpg) as mock_killpg, ): result = await tool.run_async( args={"command": "python scripts/do_thing.py"}, @@ -344,7 +360,8 @@ async def test_resource_limits_set(self, workspace, tool_context_confirmed): tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) mock_process = mock.AsyncMock() mock_process.pid = None # Ensure finally block doesn't try to kill it - mock_process.communicate.return_value = (b"", b"") + mock_process.stdout = None + mock_process.stderr = None mock_exec = mock.AsyncMock(return_value=mock_process) with mock.patch("asyncio.create_subprocess_exec", mock_exec): @@ -365,3 +382,281 @@ async def test_resource_limits_set(self, workspace, tool_context_confirmed): mock_setrlimit.assert_any_call( resource.RLIMIT_FSIZE, (50 * 1024 * 1024, 50 * 1024 * 1024) ) + + +class TestOutputLimits: + + def test_default_limit(self): + assert bash_tool.BashToolPolicy().max_output_bytes_per_stream == ( + 1024 * 1024 + ) + + @pytest.mark.asyncio + async def test_small_output_unchanged( + self, workspace, tool_context_confirmed + ): + policy = bash_tool.BashToolPolicy(max_output_bytes_per_stream=1024) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + result = await tool.run_async( + args={"command": "python3 -c 'print(\"hello\")'"}, + tool_context=tool_context_confirmed, + ) + assert result["stdout"] == "hello\n" + assert result["returncode"] == 0 + + @pytest.mark.asyncio + async def test_large_stdout_is_truncated( + self, workspace, tool_context_confirmed + ): + policy = bash_tool.BashToolPolicy(max_output_bytes_per_stream=1000) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + total = 2 * 1024 * 1024 + result = await tool.run_async( + args={ + "command": ( + "python3 -c 'import sys;" + f' sys.stdout.buffer.write(b"x" * {total})\'' + ) + }, + tool_context=tool_context_confirmed, + ) + assert result["returncode"] == 0 + assert result["stdout"].startswith("x" * 1000 + "\n...[stdout truncated") + assert f"{total - 1000} of {total} bytes omitted" in result["stdout"] + assert len(result["stdout"]) < 1200 + + @pytest.mark.asyncio + async def test_stdout_and_stderr_are_bounded_independently( + self, workspace, tool_context_confirmed + ): + policy = bash_tool.BashToolPolicy(max_output_bytes_per_stream=500) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + script = ( + "import sys\n" + "for _ in range(64):\n" + ' sys.stdout.buffer.write(b"o" * 16384)\n' + ' sys.stderr.buffer.write(b"e" * 16384)\n' + ) + (workspace / "noisy.py").write_text(script) + result = await tool.run_async( + args={"command": "python3 noisy.py"}, + tool_context=tool_context_confirmed, + ) + assert result["returncode"] == 0 + assert result["stdout"].startswith("o" * 500 + "\n...[stdout truncated") + assert result["stderr"].startswith("e" * 500 + "\n...[stderr truncated") + assert len(result["stdout"]) < 700 + assert len(result["stderr"]) < 700 + + @pytest.mark.asyncio + async def test_limit_can_be_disabled(self, workspace, tool_context_confirmed): + policy = bash_tool.BashToolPolicy(max_output_bytes_per_stream=None) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + result = await tool.run_async( + args={ + "command": ( + 'python3 -c \'import sys; sys.stdout.buffer.write(b"x" *' + " 200000)'" + ) + }, + tool_context=tool_context_confirmed, + ) + assert result["stdout"] == "x" * 200000 + + @pytest.mark.asyncio + async def test_timeout_output_is_bounded_and_process_reaped( + self, workspace, tool_context_confirmed + ): + policy = bash_tool.BashToolPolicy( + timeout_seconds=1, max_output_bytes_per_stream=100 + ) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + (workspace / "forever.py").write_text( + 'import sys\nwhile True:\n sys.stdout.buffer.write(b"y" * 65536)\n' + ) + result = await tool.run_async( + args={"command": "python3 forever.py"}, + tool_context=tool_context_confirmed, + ) + assert "timed out" in result["error"] + assert result["stdout"].startswith("y" * 100 + "\n...[stdout truncated") + assert len(result["stdout"]) < 300 + assert result["returncode"] == -signal.SIGKILL + + @pytest.mark.asyncio + async def test_timeout_setsid_child_holding_pipe_cleans_up( + self, workspace, tool_context_confirmed + ): + """Timeout must not leave wait/drain tasks pending when a setsid child holds a pipe. + + A descendant that calls setsid() leaves the process group, so killpg does + not close its inherited stdout. Cleanup must cancel the explicit waiter and + readers instead of leaving process.wait() pending inside a gather. + """ + import time + + policy = bash_tool.BashToolPolicy( + timeout_seconds=1, max_output_bytes_per_stream=64 + ) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + (workspace / "setsid_hold.py").write_text( + "import os\n" + "import sys\n" + "import time\n" + "if os.fork() == 0:\n" + " os.setsid()\n" + " while True:\n" + ' sys.stdout.buffer.write(b"z" * 1024)\n' + " sys.stdout.buffer.flush()\n" + " time.sleep(0.05)\n" + "else:\n" + " time.sleep(60)\n" + ) + loop = asyncio.get_running_loop() + reports = [] + previous_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _, context: reports.append(context)) + try: + started = time.monotonic() + result = await tool.run_async( + args={"command": "python3 setsid_hold.py"}, + tool_context=tool_context_confirmed, + ) + elapsed = time.monotonic() - started + for _ in range(5): + await asyncio.sleep(0) + gc.collect() + await asyncio.sleep(0) + finally: + loop.set_exception_handler(previous_handler) + + assert "timed out" in result["error"] + # Must return after timeout + drain grace, not after the child's 60s sleep. + assert elapsed < 5.0 + leftover = [ + t + for t in asyncio.all_tasks() + if not t.done() + and t is not asyncio.current_task() + and getattr(getattr(t, "get_coro", lambda: None)(), "__name__", None) + in {"_drain_stream", "wait"} + ] + assert leftover == [] + noisy = [ + r + for r in reports + if "never retrieved" in r.get("message", "") + or "shielded" in r.get("message", "").lower() + ] + assert noisy == [] + + def test_rejects_negative_output_limit(self): + with pytest.raises(ValueError, match="max_output_bytes_per_stream"): + bash_tool.BashToolPolicy(max_output_bytes_per_stream=-1) + + def test_zero_output_limit_is_allowed(self): + policy = bash_tool.BashToolPolicy(max_output_bytes_per_stream=0) + assert policy.max_output_bytes_per_stream == 0 + + @pytest.mark.asyncio + async def test_output_exactly_at_limit_is_not_truncated( + self, workspace, tool_context_confirmed + ): + limit = 1000 + policy = bash_tool.BashToolPolicy(max_output_bytes_per_stream=limit) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + result = await tool.run_async( + args={ + "command": ( + "python3 -c 'import sys;" + f' sys.stdout.buffer.write(b"x" * {limit})\'' + ) + }, + tool_context=tool_context_confirmed, + ) + assert result["returncode"] == 0 + assert result["stdout"] == "x" * limit + assert "truncated" not in result["stdout"] + + @pytest.mark.asyncio + async def test_cancellation_cleans_up_drain_and_wait( + self, workspace, tool_context_confirmed + ): + """Cancelling run_async must not leave drain/wait tasks pending. + + A child that stays in the process group and keeps stdout open would + otherwise leave _drain_stream tasks hanging after CancelledError. + """ + policy = bash_tool.BashToolPolicy( + timeout_seconds=30, max_output_bytes_per_stream=64 + ) + tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy) + (workspace / "hold_pipe.py").write_text( + "import os\n" + "import sys\n" + "import time\n" + "if os.fork() == 0:\n" + " while True:\n" + ' sys.stdout.buffer.write(b"z" * 1024)\n' + " sys.stdout.buffer.flush()\n" + " time.sleep(0.05)\n" + "else:\n" + " time.sleep(60)\n" + ) + task = asyncio.create_task( + tool.run_async( + args={"command": "python3 hold_pipe.py"}, + tool_context=tool_context_confirmed, + ) + ) + await asyncio.sleep(0.4) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) + leftover_drains = [ + t + for t in asyncio.all_tasks() + if not t.done() + and getattr(getattr(t, "get_coro", lambda: None)(), "__name__", None) + == "_drain_stream" + ] + assert leftover_drains == [] + + @pytest.mark.asyncio + async def test_cancellation_retrieves_completion_exception( + self, workspace, tool_context_confirmed + ): + """Cancelling run_async must not report an unretrieved gather exception.""" + loop = asyncio.get_running_loop() + reports = [] + previous_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _, context: reports.append(context)) + try: + tool = bash_tool.ExecuteBashTool( + workspace=workspace, + policy=bash_tool.BashToolPolicy(timeout_seconds=30), + ) + task = asyncio.create_task( + tool.run_async( + args={"command": "sleep 30"}, + tool_context=tool_context_confirmed, + ) + ) + await asyncio.sleep(0.4) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + del task + for _ in range(3): + await asyncio.sleep(0) + gc.collect() + await asyncio.sleep(0) + finally: + loop.set_exception_handler(previous_handler) + assert [ + r + for r in reports + if "never retrieved" in r.get("message", "") + or "shielded" in r.get("message", "").lower() + ] == []