fix: Prevent retry_on_errors from retrying asyncio.CancelledError

The retry_on_errors decorator in mcp_session_manager.py now catches asyncio.CancelledError and re-raises it immediately, ensuring that cancellation requests are not suppressed or retried

Close #4009

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 852468382

Author: GWeale

src/google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import asyncio
+from collections import deque
 from contextlib import AsyncExitStack
 from datetime import timedelta
 import functools
@@ -30,7 +31,6 @@
 from typing import TextIO
 from typing import Union
 
-import anyio
 from mcp import ClientSession
 from mcp import StdioServerParameters
 from mcp.client.sse import sse_client
@@ -44,6 +44,30 @@
 logger = logging.getLogger('google_adk.' + __name__)
 
 
+def _has_cancelled_error_context(exc: BaseException) -> bool:
+  """Returns True if `exc` is/was caused by `asyncio.CancelledError`.
+
+  Cancellation can be translated into other exceptions during teardown (e.g.
+  connection errors) while still retaining the original cancellation in an
+  exception's context chain.
+  """
+
+  seen: set[int] = set()
+  queue = deque([exc])
+  while queue:
+    current = queue.popleft()
+    if id(current) in seen:
+      continue
+    seen.add(id(current))
+    if isinstance(current, asyncio.CancelledError):
+      return True
+    if current.__cause__ is not None:
+      queue.append(current.__cause__)
+    if current.__context__ is not None:
+      queue.append(current.__context__)
+  return False
+
+
 class StdioConnectionParams(BaseModel):
   """Parameters for the MCP Stdio connection.
 
@@ -119,6 +143,10 @@ def retry_on_errors(func):
   action once. The create_session method will handle creating a new session
   if the old one was disconnected.
 
+  Cancellation is not retried and must be allowed to propagate. In async
+  runtimes, cancellation may surface as `asyncio.CancelledError` or as another
+  exception while the task is cancelling.
+
   Args:
       func: The function to decorate.
 
@@ -131,6 +159,13 @@ async def wrapper(self, *args, **kwargs):
     try:
       return await func(self, *args, **kwargs)
     except Exception as e:
+      task = asyncio.current_task()
+      if task is not None:
+        cancelling = getattr(task, 'cancelling', None)
+        if cancelling is not None and cancelling() > 0:
+          raise
+      if _has_cancelled_error_context(e):
+        raise
       # If an error is thrown, we will retry the function to reconnect to the
       # server. create_session will handle detecting and replacing disconnected
       # sessions.

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

@@ -391,7 +391,8 @@ async def test_close_with_errors(self):
     assert "Close error 1" in error_output
 
 
-def test_retry_on_errors_decorator():
+@pytest.mark.asyncio
+async def test_retry_on_errors_decorator():
   """Test the retry_on_errors decorator."""
 
   call_count = 0
@@ -401,23 +402,77 @@ async def mock_function(self):
     nonlocal call_count
     call_count += 1
     if call_count == 1:
-      import anyio
-
-      raise anyio.ClosedResourceError("Resource closed")
+      raise ConnectionError("Resource closed")
     return "success"
 
-  @pytest.mark.asyncio
-  async def test_retry():
+  mock_self = Mock()
+  result = await mock_function(mock_self)
+
+  assert result == "success"
+  assert call_count == 2  # First call fails, second succeeds
+
+
+@pytest.mark.asyncio
+async def test_retry_on_errors_decorator_does_not_retry_cancelled_error():
+  """Test the retry_on_errors decorator does not retry cancellation."""
+
+  call_count = 0
+
+  @retry_on_errors
+  async def mock_function(self):
     nonlocal call_count
-    call_count = 0
+    call_count += 1
+    raise asyncio.CancelledError()
+
+  mock_self = Mock()
+  with pytest.raises(asyncio.CancelledError):
+    await mock_function(mock_self)
+
+  assert call_count == 1
 
-    mock_self = Mock()
-    result = await mock_function(mock_self)
 
-    assert result == "success"
-    assert call_count == 2  # First call fails, second succeeds
+@pytest.mark.asyncio
+async def test_retry_on_errors_decorator_does_not_retry_when_task_is_cancelling():
+  """Test the retry_on_errors decorator does not retry when cancelling."""
+
+  call_count = 0
+
+  @retry_on_errors
+  async def mock_function(self):
+    nonlocal call_count
+    call_count += 1
+    raise ConnectionError("Resource closed")
+
+  class _MockTask:
+
+    def cancelling(self):
+      return 1
+
+  mock_self = Mock()
+  with patch.object(asyncio, "current_task", return_value=_MockTask()):
+    with pytest.raises(ConnectionError):
+      await mock_function(mock_self)
+
+  assert call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_retry_on_errors_decorator_does_not_retry_exception_from_cancel():
+  """Test the retry_on_errors decorator does not retry exceptions on cancel."""
+
+  call_count = 0
+
+  @retry_on_errors
+  async def mock_function(self):
+    nonlocal call_count
+    call_count += 1
+    try:
+      raise asyncio.CancelledError()
+    except asyncio.CancelledError:
+      raise ConnectionError("Resource closed")
 
-  # Run the test
-  import asyncio
+  mock_self = Mock()
+  with pytest.raises(ConnectionError):
+    await mock_function(mock_self)
 
-  asyncio.run(test_retry())
+  assert call_count == 1