From 6cedc87ed25a60fe4996df01aae40f0e414f8bb4 Mon Sep 17 00:00:00 2001 From: FU-max-boop Date: Wed, 12 Aug 2026 15:35:22 +0800 Subject: [PATCH] Fix completed request cancellation cleanup --- src/mcp/shared/session.py | 6 +++- tests/shared/test_session.py | 53 +++++++++++++++++++++++++++++++++++- 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/mcp/shared/session.py b/src/mcp/shared/session.py index 35a83fcf1b..93f5f77bf7 100644 --- a/src/mcp/shared/session.py +++ b/src/mcp/shared/session.py @@ -114,7 +114,11 @@ def __exit__( self._entered = False if not self._cancel_scope: # pragma: no cover raise RuntimeError("No active cancel scope") - self._cancel_scope.__exit__(exc_type, exc_val, exc_tb) + try: + self._cancel_scope.__exit__(exc_type, exc_val, exc_tb) + except BaseException as exc: + if not (self._completed and isinstance(exc, anyio.get_cancelled_exc_class())): + raise async def respond(self, response: SendResultT | ErrorData) -> None: """Send a response for this request. diff --git a/tests/shared/test_session.py b/tests/shared/test_session.py index f4010141d8..11edf26e6d 100644 --- a/tests/shared/test_session.py +++ b/tests/shared/test_session.py @@ -1,5 +1,5 @@ from collections.abc import AsyncGenerator -from typing import Any +from typing import Any, cast import anyio import pytest @@ -10,6 +10,7 @@ from mcp.shared.exceptions import McpError from mcp.shared.memory import create_client_server_memory_streams, create_connected_server_and_client_session from mcp.shared.message import SessionMessage +from mcp.shared.session import RequestResponder from mcp.types import ( CancelledNotification, CancelledNotificationParams, @@ -25,6 +26,18 @@ ) +class _CancelScopeThatRaisesOnExit: + cancel_called = True + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: object | None, + ) -> None: + raise anyio.get_cancelled_exc_class()() + + @pytest.fixture def mcp_server() -> Server: return Server(name="test server") @@ -128,6 +141,44 @@ async def make_request(client_session: ClientSession): await ev_cancelled.wait() +@pytest.mark.anyio +async def test_completed_request_responder_suppresses_cancel_scope_exit() -> None: + completed: list[Any] = [] + responder = RequestResponder( + request_id=1, + request_meta=None, + request=types.ClientRequest(types.PingRequest()), + session=cast(Any, object()), + on_complete=completed.append, + ) + responder._completed = True # type: ignore[reportPrivateUsage] + responder._cancel_scope = cast( # type: ignore[reportPrivateUsage] + anyio.CancelScope, _CancelScopeThatRaisesOnExit() + ) + + responder.__exit__(None, None, None) + + assert completed == [responder] + assert not responder._entered # type: ignore[reportPrivateUsage] + + +@pytest.mark.anyio +async def test_incomplete_request_responder_propagates_cancel_scope_exit() -> None: + responder = RequestResponder( + request_id=1, + request_meta=None, + request=types.ClientRequest(types.PingRequest()), + session=cast(Any, object()), + on_complete=lambda _: None, + ) + responder._cancel_scope = cast( # type: ignore[reportPrivateUsage] + anyio.CancelScope, _CancelScopeThatRaisesOnExit() + ) + + with pytest.raises(anyio.get_cancelled_exc_class()): + responder.__exit__(None, None, None) + + @pytest.mark.anyio async def test_response_id_type_mismatch_string_to_int(): """