From bb768865698566b6a0e66f093ae1a5935c5a0d22 Mon Sep 17 00:00:00 2001 From: FU-max-boop Date: Sun, 17 May 2026 16:08:54 +0800 Subject: [PATCH] Fix completed request cancellation cleanup --- src/mcp/shared/session.py | 6 ++++- tests/shared/test_session.py | 52 ++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/mcp/shared/session.py b/src/mcp/shared/session.py index 243eef5ae6..e51ef51d36 100644 --- a/src/mcp/shared/session.py +++ b/src/mcp/shared/session.py @@ -117,7 +117,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 d7c6cc3b5f..c0ec0a2ae4 100644 --- a/tests/shared/test_session.py +++ b/tests/shared/test_session.py @@ -1,3 +1,5 @@ +from typing import Any, cast + import anyio import pytest @@ -23,6 +25,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.mark.anyio async def test_in_flight_requests_cleared_after_completion(): """Verify that _in_flight is empty after all requests complete.""" @@ -98,6 +112,44 @@ async def make_request(client: Client): 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.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.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(): """Test that responses with string IDs are correctly matched to requests sent with