From a8657a4a7b0baec1df2a090dfe33eedf4f73bd6e Mon Sep 17 00:00:00 2001
From: Atul Joshi <120785343+AtulJoshi1206@users.noreply.github.com>
Date: Sat, 19 Sep 2026 09:12:49 +0530
Subject: [PATCH] fix(mcp): cancel the in-flight tool call when the caller is
 cancelled

asyncio.wait does not own the futures it waits on, so _run_guarded left
session.call_tool() running when its caller was cancelled. McpTool then ran
its finally and released the session, dropping the pool's in-flight count to
zero and stamping _session_last_used while the orphaned call was still
reading the transport, which leaves _evict_idle_sessions free to close that
transport underneath it.

Cancel and drain the call when the wait itself is interrupted. The
transport-crash branch already did exactly this, so both paths now share one
_cancel_and_drain helper.
---
 .../browser/assets/config/runtime-config.json |  2 +-
 .../adk/tools/mcp_tool/session_context.py     | 37 +++++++++++++----
 .../tools/mcp_tool/test_session_context.py    | 41 +++++++++++++++++++
 3 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/src/google/adk/cli/browser/assets/config/runtime-config.json b/src/google/adk/cli/browser/assets/config/runtime-config.json
index 888614ad56..873e88b1f1 100644
--- a/src/google/adk/cli/browser/assets/config/runtime-config.json
+++ b/src/google/adk/cli/browser/assets/config/runtime-config.json
@@ -1,4 +1,4 @@
 {
   "backendUrl": "",
-  "telemetry": false
+  "telemetry": null
 }
diff --git a/src/google/adk/tools/mcp_tool/session_context.py b/src/google/adk/tools/mcp_tool/session_context.py
index e2e8475bd1..1fecb46893 100644
--- a/src/google/adk/tools/mcp_tool/session_context.py
+++ b/src/google/adk/tools/mcp_tool/session_context.py
@@ -38,6 +38,21 @@
 _T = TypeVar('_T')
 
 
+async def _cancel_and_drain(task: asyncio.Future[Any]) -> None:
+  """Cancels an in-flight call and waits for it to actually stop.
+
+  Returning before the call has unwound would let it keep reading the
+  transport after the caller believes it is finished.
+  """
+  task.cancel()
+  try:
+    await task
+  except BaseException:
+    # Includes the CancelledError just requested, and a cancellation
+    # delivered to this frame while the drain was in progress.
+    pass
+
+
 def _read_timeout(seconds: Optional[float]) -> Optional[float | timedelta]:
   """Converts a timeout in seconds to the type ``ClientSession`` expects.
 
@@ -273,10 +288,18 @@ async def _run_guarded(self, coro: Coroutine[Any, Any, _T]) -> _T:
 
     coro_task = asyncio.ensure_future(coro)
 
-    done, _ = await asyncio.wait(
-        [coro_task, self._task],
-        return_when=asyncio.FIRST_COMPLETED,
-    )
+    try:
+      done, _ = await asyncio.wait(
+          [coro_task, self._task],
+          return_when=asyncio.FIRST_COMPLETED,
+      )
+    except BaseException:
+      # asyncio.wait does not own what it waits on, so it leaves the call
+      # running when this frame is cancelled. The caller's `finally` then
+      # releases the session back to the pool while that call is still
+      # reading the transport, and the pool is free to evict it underneath.
+      await _cancel_and_drain(coro_task)
+      raise
 
     if coro_task in done:
       # If the coroutine itself raised, the exception propagates as-is
@@ -287,11 +310,7 @@ async def _run_guarded(self, coro: Coroutine[Any, Any, _T]) -> _T:
 
     # The background task finished first, indicating a transport crash.
     # Cancel the in-flight tool call and surface the original error.
-    coro_task.cancel()
-    try:
-      await coro_task
-    except BaseException:
-      pass
+    await _cancel_and_drain(coro_task)
 
     exc = self._task.exception() if not self._task.cancelled() else None
     raise ConnectionError(
diff --git a/tests/unittests/tools/mcp_tool/test_session_context.py b/tests/unittests/tools/mcp_tool/test_session_context.py
index 892c2413e4..6df00acbd8 100644
--- a/tests/unittests/tools/mcp_tool/test_session_context.py
+++ b/tests/unittests/tools/mcp_tool/test_session_context.py
@@ -900,6 +900,47 @@ async def kill_background_task():
       finally:
         await killer
 
+  @pytest.mark.asyncio
+  async def test_run_guarded_cancels_coro_when_caller_is_cancelled(self):
+    """Cancelling the caller also cancels the in-flight tool call.
+
+    asyncio.wait does not own the futures it waits on, so a cancelled
+    caller would leave the call running against a session that
+    McpTool._run_async_impl has already released back to the pool.
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      await session_context.start()
+
+      coro_started = asyncio.Event()
+      coro_was_cancelled = False
+
+      async def slow_coro():
+        nonlocal coro_was_cancelled
+        coro_started.set()
+        try:
+          await asyncio.sleep(300)
+          return 'should never reach here'
+        except asyncio.CancelledError:
+          coro_was_cancelled = True
+          raise
+
+      caller = asyncio.create_task(session_context._run_guarded(slow_coro()))
+      await coro_started.wait()
+      caller.cancel()
+
+      with pytest.raises(asyncio.CancelledError):
+        await caller
+
+      assert coro_was_cancelled is True
+
 
 class TestSessionContextFlagOffPreservesPreFixBehavior:
   """Pin down that flag=OFF reproduces pre-fix behavior exactly.