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.

Author: Hugo-DDT

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"),

tests/unittests/tools/test_bash_tool.py

@@ -295,22 +295,36 @@ 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,
             "create_subprocess_exec",
             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()
     ] == []