From 30d3411d603f12ca5bcdd2d71773d087f3191dba Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 5 Jan 2026 14:57:48 -0800 Subject: [PATCH] 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 PiperOrigin-RevId: 852468382 --- .../adk/tools/mcp_tool/mcp_session_manager.py | 37 ++++++++- .../mcp_tool/test_mcp_session_manager.py | 83 +++++++++++++++---- 2 files changed, 105 insertions(+), 15 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 722730e4..89f01457 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/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 runtime_checkable 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 @@ from pydantic import ConfigDict 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 @@ def retry_on_errors(func): 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. diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index cac156ae..7e18d9d4 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -391,7 +391,8 @@ class TestMCPSessionManager: 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 @@ def test_retry_on_errors_decorator(): 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() - result = await mock_function(mock_self) + mock_self = Mock() + with pytest.raises(asyncio.CancelledError): + await mock_function(mock_self) - assert result == "success" - assert call_count == 2 # First call fails, second succeeds + assert call_count == 1 - # Run the test - import asyncio - asyncio.run(test_retry()) +@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") + + mock_self = Mock() + with pytest.raises(ConnectionError): + await mock_function(mock_self) + + assert call_count == 1