From 473545449b5147b14c0dfca2f9f8fc29e1875f49 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=B6rje?= Date: Sat, 18 Jul 2026 00:38:27 +0200 Subject: [PATCH] fix(mcp): make transport classification cycle safe --- tests/tools/test_mcp_tool_session_expired.py | 35 +++++++--- tools/mcp_tool.py | 73 ++++++++++---------- 2 files changed, 62 insertions(+), 46 deletions(-) diff --git a/tests/tools/test_mcp_tool_session_expired.py b/tests/tools/test_mcp_tool_session_expired.py index f6aa9ee590bc..7dba0505528b 100644 --- a/tests/tools/test_mcp_tool_session_expired.py +++ b/tests/tools/test_mcp_tool_session_expired.py @@ -121,27 +121,28 @@ def test_is_session_expired_rejects_mixed_group_with_user_interruption(): assert _is_session_expired_error(exc) is False -def test_exception_tree_finds_interruption_beyond_recursion_limit(): - """Arbitrarily deep wrapper trees must not overflow Python's call stack.""" +def test_is_session_expired_finds_closed_resource_beyond_recursion_limit(): + """The full classifier must handle arbitrarily deep transport wrappers.""" import sys - from tools.mcp_tool import _exception_tree_contains_interruption + from anyio import ClosedResourceError + from tools.mcp_tool import _is_session_expired_error class NestedException(Exception): exceptions: tuple[BaseException, ...] - exc = InterruptedError("cancel") + exc = ClosedResourceError() for _ in range(sys.getrecursionlimit() + 100): wrapper = NestedException("wrapped") wrapper.exceptions = (exc,) exc = wrapper - assert _exception_tree_contains_interruption(exc) is True + assert _is_session_expired_error(exc) is True -def test_exception_tree_handles_cyclic_exceptions_graph(): - """Malformed exception graphs may contain cycles and must terminate safely.""" - from tools.mcp_tool import _exception_tree_contains_interruption +def test_is_session_expired_handles_cyclic_graph_without_transport_error(): + """A cyclic non-transport graph must terminate and classify false.""" + from tools.mcp_tool import _is_session_expired_error class CyclicException(Exception): exceptions: tuple[BaseException, ...] @@ -151,7 +152,23 @@ def test_exception_tree_handles_cyclic_exceptions_graph(): first.exceptions = (second,) second.exceptions = (first,) - assert _exception_tree_contains_interruption(first) is False + assert _is_session_expired_error(first) is False + + +def test_is_session_expired_finds_transport_error_in_cyclic_graph(): + """Cycle detection must not prevent scanning reachable transport errors.""" + from anyio import ClosedResourceError + from tools.mcp_tool import _is_session_expired_error + + class CyclicException(Exception): + exceptions: tuple[BaseException, ...] + + first = CyclicException("first") + second = CyclicException("second") + first.exceptions = (second, ClosedResourceError()) + second.exceptions = (first,) + + assert _is_session_expired_error(first) is True def test_is_session_expired_rejects_empty_message(): diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index c2b9503a84e2..ef8ed39145cb 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -3900,22 +3900,6 @@ _SESSION_EXPIRED_MARKERS: tuple = ( ) -def _exception_tree_contains_interruption(exc: BaseException) -> bool: - """Return whether ``exc`` or any nested exception is user cancellation.""" - stack = [exc] - seen: set[int] = set() - while stack: - current = stack.pop() - identity = id(current) - if identity in seen: - continue - seen.add(identity) - if isinstance(current, InterruptedError): - return True - stack.extend(getattr(current, "exceptions", ())) - return False - - def _is_session_expired_error(exc: BaseException) -> bool: """Return True if ``exc`` looks like an MCP transport session expiry. @@ -3930,35 +3914,50 @@ def _is_session_expired_error(exc: BaseException) -> bool: ``streamablehttp_client`` + ``ClientSession`` pair, which is exactly what ``MCPServerTask._reconnect_event`` triggers. """ - if _exception_tree_contains_interruption(exc): - return False - # AnyIO's stream exceptions are commonly message-less. In particular, # ``str(ClosedResourceError()) == ""``, so marker matching alone misses the - # exact failure emitted by both MCP stdio and HTTP transports. Match the - # SDK's transport-closed exception types before inspecting text. + # exact failure emitted by both MCP stdio and HTTP transports. try: from anyio import BrokenResourceError, ClosedResourceError, EndOfStream - if isinstance(exc, (BrokenResourceError, ClosedResourceError, EndOfStream)): - return True + transport_error_types = ( + BrokenResourceError, + ClosedResourceError, + EndOfStream, + ) except ImportError: # pragma: no cover - AnyIO is supplied by the MCP SDK - pass + transport_error_types = () - # AnyIO task groups can wrap the transport exception in an - # ExceptionGroup whose own string omits the message-less leaf details. - nested = getattr(exc, "exceptions", ()) - if nested and any(_is_session_expired_error(child) for child in nested): - return True + # ExceptionGroup trees can be arbitrarily deep or even cyclic when custom + # exceptions expose ``exceptions``. Traverse once, iteratively, and inspect + # every reachable node so user interruption always overrides transport + # markers or types found elsewhere in the graph. + stack = [exc] + seen: set[int] = set() + transport_error_found = False + while stack: + current = stack.pop() + identity = id(current) + if identity in seen: + continue + seen.add(identity) - # Exception messages vary across SDK versions + server - # implementations, so match on a small allow-list of stable - # substrings rather than exception type. Kept narrow to avoid - # false positives on unrelated server errors. - msg = str(exc).lower() - if not msg: - return False - return any(marker in msg for marker in _SESSION_EXPIRED_MARKERS) + if isinstance(current, InterruptedError): + return False + if isinstance(current, transport_error_types): + transport_error_found = True + + # Exception messages vary across SDK versions + server + # implementations, so match on a small allow-list of stable + # substrings rather than exception type. Kept narrow to avoid + # false positives on unrelated server errors. + msg = str(current).lower() + if msg and any(marker in msg for marker in _SESSION_EXPIRED_MARKERS): + transport_error_found = True + + stack.extend(getattr(current, "exceptions", ())) + + return transport_error_found def _handle_session_expired_and_retry(