fix(tools): bound ExecuteBashTool stdout/stderr capture

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"<no {stream_name} captured>"
+    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,68 +335,78 @@ 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 "<no stdout captured>"
-            ),
-            "stderr": (
-                stderr.decode(errors="replace")
-                if stderr
-                else "<no stderr captured>"
-            ),
-            "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:
             _kill_process_group(process)
         except ProcessLookupError:
           pass
       return {
-          "stdout": (
-              stdout.decode(errors="replace")
-              if stdout
-              else "<no stdout captured>"
-          ),
-          "stderr": (
-              stderr.decode(errors="replace")
-              if stderr
-              else "<no stderr captured>"
-          ),
+          "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 "<no stdout captured>"
-      )
-      stderr_res = (
-          stderr.decode(errors="replace") if stderr else "<no stderr captured>"
-      )
-
+      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]: