From bddf77023af2ab753eef009ecf04b81b6f669ccd Mon Sep 17 00:00:00 2001 From: DDT <141296842+Hugo-DDT@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:50:10 +0800 Subject: [PATCH 1/4] fix(tools): bound ExecuteBashTool stdout/stderr capture Replace process.communicate() with concurrent 64 KiB chunked reads of stdout and stderr into per-stream buffers capped by the new BashToolPolicy.max_output_bytes_per_stream (default 1 MiB, None disables). Excess output is drained and discarded so the child never blocks on a full pipe, and the returned text ends with a marker noting how many bytes were omitted. Timeout handling keeps draining the pipes after killing the process group so the subprocess is reaped. Fixes #7507 --- src/google/adk/tools/bash_tool.py | 125 +++++++++++++++++------- tests/unittests/tools/test_bash_tool.py | 106 +++++++++++++++++++- 2 files changed, 194 insertions(+), 37 deletions(-) diff --git a/src/google/adk/tools/bash_tool.py b/src/google/adk/tools/bash_tool.py index aa3571c168e..a066c68ff3c 100644 --- a/src/google/adk/tools/bash_tool.py +++ b/src/google/adk/tools/bash_tool.py @@ -58,6 +58,66 @@ 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_readers(readers: list[asyncio.Task[None]]) -> None: + """Waits briefly for readers to reach EOF, then cancels any stragglers.""" + _, pending = await asyncio.wait(readers, timeout=_DRAIN_GRACE_SECONDS) + for task in pending: + task.cancel() + await asyncio.gather(*readers, return_exceptions=True) + @dataclasses.dataclass(frozen=True) class BashToolPolicy: @@ -75,6 +135,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,6 +149,7 @@ 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: for prefix in self.allowed_command_prefixes: @@ -245,8 +312,9 @@ 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[None]] = [] try: process = await asyncio.create_subprocess_exec( *shlex.split(command), @@ -256,10 +324,18 @@ 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)), + ] + completion = asyncio.gather(*readers, process.wait()) try: - stdout, stderr = await asyncio.wait_for( - process.communicate(), timeout=self._policy.timeout_seconds + # Shield so a timeout does not cancel the readers: they must keep + # draining the pipes until EOF for the killed process to be reaped. + await asyncio.wait_for( + asyncio.shield(completion), + timeout=self._policy.timeout_seconds, ) except asyncio.TimeoutError: try: @@ -267,22 +343,15 @@ async def run_async( _kill_process_group(process) except ProcessLookupError: pass - stdout, stderr = await process.communicate() + await _finish_readers(readers) + await process.wait() 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 "" - ), + "stdout": stdout.render("stdout"), + "stderr": stderr.render("stderr"), "returncode": process.returncode, } finally: @@ -292,32 +361,18 @@ 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 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 "" - ) - + if readers: + await _finish_readers(readers) 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..b4400258104 100644 --- a/tests/unittests/tools/test_bash_tool.py +++ b/tests/unittests/tools/test_bash_tool.py @@ -297,7 +297,8 @@ async def test_timeout(self, workspace, tool_context_confirmed): tool = bash_tool.ExecuteBashTool(workspace=workspace) mock_process = mock.AsyncMock() mock_process.pid = 12345 - mock_process.communicate.return_value = (b"", b"") + mock_process.stdout = None + mock_process.stderr = None with ( mock.patch.object( asyncio, @@ -344,7 +345,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 +367,103 @@ 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 From a62ce9c2349c551b8982984ac84822f0b234eec4 Mon Sep 17 00:00:00 2001 From: DDT <141296842+Hugo-DDT@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:06:17 +0800 Subject: [PATCH 2/4] fix(tools): clean up bash drain on cancel; validate output limit Address review on #7521 / #7507: on CancelledError kill the process group, cancel completion and drain readers, gather them, then re-raise; reject max_output_bytes_per_stream < 0; add exact-at-limit and cancellation tests. --- src/google/adk/tools/bash_tool.py | 29 ++++++++++ tests/unittests/tools/test_bash_tool.py | 73 +++++++++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/src/google/adk/tools/bash_tool.py b/src/google/adk/tools/bash_tool.py index a066c68ff3c..cde3e126663 100644 --- a/src/google/adk/tools/bash_tool.py +++ b/src/google/adk/tools/bash_tool.py @@ -152,6 +152,13 @@ class BashToolPolicy: 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) @@ -315,6 +322,8 @@ async def run_async( stdout = _BoundedOutput(self._policy.max_output_bytes_per_stream) stderr = _BoundedOutput(self._policy.max_output_bytes_per_stream) readers: list[asyncio.Task[None]] = [] + process: Optional[asyncio.subprocess.Process] = None + completion: Optional[asyncio.Future[Any]] = None try: process = await asyncio.create_subprocess_exec( *shlex.split(command), @@ -365,6 +374,26 @@ async def run_async( "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 + if completion is not None and not completion.done(): + completion.cancel() + for task in readers: + task.cancel() + if readers: + await asyncio.gather(*readers, return_exceptions=True) + if process is not None and process.returncode is None: + try: + await asyncio.shield(process.wait()) + except Exception: # pylint: disable=broad-except + pass + raise except Exception as e: # pylint: disable=broad-except logger.exception("ExecuteBashTool execution failed") if readers: diff --git a/tests/unittests/tools/test_bash_tool.py b/tests/unittests/tools/test_bash_tool.py index b4400258104..97c30796d88 100644 --- a/tests/unittests/tools/test_bash_tool.py +++ b/tests/unittests/tools/test_bash_tool.py @@ -467,3 +467,76 @@ async def test_timeout_output_is_bounded_and_process_reaped( assert result["stdout"].startswith("y" * 100 + "\n...[stdout truncated") assert len(result["stdout"]) < 300 assert result["returncode"] == -signal.SIGKILL + + 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 == [] From c9aff13df59d842da265e50537c4742a26671a7c Mon Sep 17 00:00:00 2001 From: DDT <141296842+Hugo-DDT@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:06:07 +0800 Subject: [PATCH 3/4] fix(tools): retrieve cancelled bash completion future --- src/google/adk/tools/bash_tool.py | 4 +++ tests/unittests/tools/test_bash_tool.py | 36 +++++++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/src/google/adk/tools/bash_tool.py b/src/google/adk/tools/bash_tool.py index cde3e126663..98e7ca4e3fe 100644 --- a/src/google/adk/tools/bash_tool.py +++ b/src/google/adk/tools/bash_tool.py @@ -388,6 +388,10 @@ async def run_async( task.cancel() if readers: await asyncio.gather(*readers, return_exceptions=True) + if completion is not None: + # Retrieve the cancelled gather's result so asyncio does not report + # "_GatheringFuture exception was never retrieved". + await asyncio.gather(completion, return_exceptions=True) if process is not None and process.returncode is None: try: await asyncio.shield(process.wait()) diff --git a/tests/unittests/tools/test_bash_tool.py b/tests/unittests/tools/test_bash_tool.py index 97c30796d88..77c626454ef 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 @@ -540,3 +541,38 @@ async def test_cancellation_cleans_up_drain_and_wait( == "_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", "") + ] == [] From 28f23be959dc196de02780149e423f9898164d7a Mon Sep 17 00:00:00 2001 From: DDT <141296842+Hugo-DDT@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:07:08 +0800 Subject: [PATCH 4/4] fix(tools): use explicit waiter for bash timeout/cancel cleanup Drop gather+shield so a setsid descendant holding a pipe cannot leave process.wait() pending after timeout, and so Python 3.14 does not log CancelledError from shielded futures. Track readers and an explicit waiter task, finish all three on timeout/cancel, and cover the setsid timeout path with a regression test. --- src/google/adk/tools/bash_tool.py | 91 +++++++++++------------ tests/unittests/tools/test_bash_tool.py | 96 +++++++++++++++++++++++-- 2 files changed, 136 insertions(+), 51 deletions(-) diff --git a/src/google/adk/tools/bash_tool.py b/src/google/adk/tools/bash_tool.py index 98e7ca4e3fe..f16f1bf9a87 100644 --- a/src/google/adk/tools/bash_tool.py +++ b/src/google/adk/tools/bash_tool.py @@ -111,12 +111,14 @@ async def _drain_stream( output.append(chunk) -async def _finish_readers(readers: list[asyncio.Task[None]]) -> None: - """Waits briefly for readers to reach EOF, then cancels any stragglers.""" - _, pending = await asyncio.wait(readers, timeout=_DRAIN_GRACE_SECONDS) +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(*readers, return_exceptions=True) + await asyncio.gather(*tasks, return_exceptions=True) @dataclasses.dataclass(frozen=True) @@ -321,9 +323,9 @@ async def run_async( stdout = _BoundedOutput(self._policy.max_output_bytes_per_stream) stderr = _BoundedOutput(self._policy.max_output_bytes_per_stream) - readers: list[asyncio.Task[None]] = [] + readers: list[asyncio.Task] = [] + waiter: Optional[asyncio.Task] = None process: Optional[asyncio.subprocess.Process] = None - completion: Optional[asyncio.Future[Any]] = None try: process = await asyncio.create_subprocess_exec( *shlex.split(command), @@ -337,32 +339,40 @@ async def run_async( asyncio.create_task(_drain_stream(process.stdout, stdout)), asyncio.create_task(_drain_stream(process.stderr, stderr)), ] - - completion = asyncio.gather(*readers, process.wait()) + # 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: - # Shield so a timeout does not cancel the readers: they must keep - # draining the pipes until EOF for the killed process to be reaped. - await asyncio.wait_for( - asyncio.shield(completion), + _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 - await _finish_readers(readers) - await process.wait() - return { - "error": ( - f"Command timed out after {self._policy.timeout_seconds}" - " seconds." - ), - "stdout": stdout.render("stdout"), - "stderr": stderr.render("stderr"), - "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: @@ -382,26 +392,17 @@ async def run_async( _kill_process_group(process) except ProcessLookupError: pass - if completion is not None and not completion.done(): - completion.cancel() - for task in readers: + tracked = [t for t in (*readers, waiter) if t is not None] + for task in tracked: task.cancel() - if readers: - await asyncio.gather(*readers, return_exceptions=True) - if completion is not None: - # Retrieve the cancelled gather's result so asyncio does not report - # "_GatheringFuture exception was never retrieved". - await asyncio.gather(completion, return_exceptions=True) - if process is not None and process.returncode is None: - try: - await asyncio.shield(process.wait()) - except Exception: # pylint: disable=broad-except - pass + if tracked: + await asyncio.gather(*tracked, return_exceptions=True) raise except Exception as e: # pylint: disable=broad-except logger.exception("ExecuteBashTool execution failed") - if readers: - await _finish_readers(readers) + 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.render("stdout"), diff --git a/tests/unittests/tools/test_bash_tool.py b/tests/unittests/tools/test_bash_tool.py index 77c626454ef..71340c163ec 100644 --- a/tests/unittests/tools/test_bash_tool.py +++ b/tests/unittests/tools/test_bash_tool.py @@ -295,11 +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.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, @@ -307,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"}, @@ -469,6 +483,73 @@ async def test_timeout_output_is_bounded_and_process_reaped( 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) @@ -574,5 +655,8 @@ async def test_cancellation_retrieves_completion_exception( finally: loop.set_exception_handler(previous_handler) assert [ - r for r in reports if "never retrieved" in r.get("message", "") + r + for r in reports + if "never retrieved" in r.get("message", "") + or "shielded" in r.get("message", "").lower() ] == []