From e1538c7b18b6391aeac4fa6d4a7dfa0cae1dccdb Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Fri, 18 Sep 2026 11:52:31 +0200 Subject: [PATCH] Bound direct subscription cleanup without delaying remote exits --- docs/client/subscriptions.md | 4 + src/mcp/client/subscriptions.py | 66 +++++++----- tests/client/test_subscriptions.py | 167 +++++++++++++++++++++++++++++ 3 files changed, 209 insertions(+), 28 deletions(-) diff --git a/docs/client/subscriptions.md b/docs/client/subscriptions.md index bf5b0a36d8..90478bc32a 100644 --- a/docs/client/subscriptions.md +++ b/docs/client/subscriptions.md @@ -59,6 +59,10 @@ Requests run freely beside an open stream, from the watcher task or any other, o To stop watching, leave the block: there is no `unsubscribe` call. Cancelling the task that owns the block does that for you, and the SDK cancels the listen request the way the transport expects: over streamable HTTP, by closing that request's stream. A watcher that runs for the life of your app never returns on its own, so cancel it, or its task group's scope, at shutdown. +With a direct in-memory connection (`Client(server)` or a `ClientSession` using `DirectDispatcher`), exit waits up to five seconds for the listen task to finish, even if the caller is cancelled. This lets cooperative handler cleanup release its subscription slot before you open another subscription. The limit prevents an uncooperative handler from blocking exit indefinitely. + +With a stream-backed connection, exit cancels the listen task without waiting for its courtesy cancellation write. The session still owns that task and its cleanup. A slow transport does not delay each subscription's exit, and exit does not acknowledge that the remote server has finished its cleanup. + ## Streams end A stream ends in one of two ways, both ordinary control flow. A graceful server close ends the `async for`; an abrupt drop raises `SubscriptionLost`. diff --git a/src/mcp/client/subscriptions.py b/src/mcp/client/subscriptions.py index 27283909be..a6ba307223 100644 --- a/src/mcp/client/subscriptions.py +++ b/src/mcp/client/subscriptions.py @@ -18,6 +18,7 @@ import mcp_types as types from mcp_types.version import MODERN_PROTOCOL_VERSIONS +from mcp.shared.direct_dispatcher import DirectDispatcher from mcp.shared.dispatcher import CallOptions from mcp.shared.exceptions import MCPError from mcp.shared.subscriptions import ( @@ -241,42 +242,51 @@ async def listen( data = request.model_dump(by_alias=True, mode="json", exclude_none=True) opts: CallOptions = {"request_id": request_id} session._stamp(data, opts) # pyright: ignore[reportPrivateUsage] + dispatcher = session._dispatcher # pyright: ignore[reportPrivateUsage] driver_scope = anyio.CancelScope() + driver_done = anyio.Event() async def drive() -> None: # Deliberately no result timeout: the response arrives when the stream ends. - with driver_scope: - try: - await session._dispatcher.send_raw_request( # pyright: ignore[reportPrivateUsage] - data["method"], data.get("params"), opts - ) - except MCPError as error: - route.settle("lost", error=error) - return - except ValueError as error: - # A raw request id collided with our minted listen id: fail this subscription - # and release the route in this same slice, so it cannot consume the raw caller's ack. - session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] - route.settle("lost", error=MCPError(types.INTERNAL_ERROR, str(error))) - return - # A result, whatever its body, is the spec's graceful close; with no prior ack - # it opens the subscription already closed. - route.set_acked(types.SubscriptionFilter()) - route.settle("graceful") + try: + with driver_scope: + try: + await dispatcher.send_raw_request(data["method"], data.get("params"), opts) + except MCPError as error: + route.settle("lost", error=error) + return + except ValueError as error: + # A raw request id collided with our minted listen id: fail this subscription + # and release the route in this same slice, so it cannot consume the raw caller's ack. + session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] + route.settle("lost", error=MCPError(types.INTERNAL_ERROR, str(error))) + return + # A result, whatever its body, is the spec's graceful close; with no prior ack + # it opens the subscription already closed. + route.set_acked(types.SubscriptionFilter()) + route.settle("graceful") + finally: + driver_done.set() # Register the demux route before the request is written so the ack cannot race it. route = session._register_listen_route(request_id) # pyright: ignore[reportPrivateUsage] try: task_group.start_soon(drive) - with anyio.fail_after(session._session_read_timeout_seconds): # pyright: ignore[reportPrivateUsage] - await route.acked.wait() - if route.honored is None: - # Only reachable on failure paths: a graceful no-ack result acked an empty filter in drive(). - if route.error is not None: - raise route.error - raise SubscriptionLost(f"subscription {request_id!r} ended before it was acknowledged") - yield Subscription(route, request_id, route.honored, on_event) + try: + with anyio.fail_after(session._session_read_timeout_seconds): # pyright: ignore[reportPrivateUsage] + await route.acked.wait() + if route.honored is None: + # Only reachable on failure paths: a graceful no-ack result acked an empty filter in drive(). + if route.error is not None: + raise route.error + raise SubscriptionLost(f"subscription {request_id!r} ended before it was acknowledged") + yield Subscription(route, request_id, route.honored, on_event) + finally: + route.settle("local") + driver_scope.cancel() + # Only direct drivers own handler cleanup; remote courtesy writes remain session-owned. + if isinstance(dispatcher, DirectDispatcher): + with anyio.move_on_after(5, shield=True): + await driver_done.wait() finally: - route.settle("local") - driver_scope.cancel() session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] diff --git a/tests/client/test_subscriptions.py b/tests/client/test_subscriptions.py index c9877c4ab0..c2967f72cd 100644 --- a/tests/client/test_subscriptions.py +++ b/tests/client/test_subscriptions.py @@ -10,6 +10,7 @@ import mcp_types as types import pytest from mcp_types import SubscriptionFilter +from trio.testing import MockClock import mcp.client.subscriptions as subscriptions_module from mcp import Client, MCPError @@ -34,10 +35,17 @@ ) from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair from mcp.shared.dispatcher import CallOptions +from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher +from mcp.shared.message import SessionMessage pytestmark = pytest.mark.anyio +@pytest.fixture(autouse=True) +def _module_runner_lease() -> None: + """Opt out of the shared runner because the cleanup timeout test parametrizes `anyio_backend`.""" + + def _bus_server(bus: InMemorySubscriptionBus, *, max_subscriptions: int | None = None) -> Server[Any]: """A lowlevel server whose only feature is serving listen streams from `bus`.""" handler = ( @@ -214,6 +222,7 @@ async def cancelling_listen( await anext(sub) +@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"]) async def test_exiting_the_context_frees_the_server_slot(): """Leaving the block ends the subscription server-side: a one-slot handler admits a second listen.""" bus = InMemorySubscriptionBus() @@ -226,6 +235,162 @@ async def test_exiting_the_context_frees_the_server_slot(): assert second.subscription_id != first.subscription_id +@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"]) +async def test_a_cancelled_task_can_close_a_subscription_opened_by_another_task() -> None: + """SDK-defined: cross-task exit joins direct handler cleanup even when the closing task is cancelled.""" + handler = ListenHandler(InMemorySubscriptionBus()) + cleanup_started = anyio.Event() + release_cleanup = anyio.Event() + cleanup_finished = anyio.Event() + closed = anyio.Event() + + async def slow_cleanup( + ctx: ServerRequestContext, params: types.SubscriptionsListenRequestParams + ) -> types.SubscriptionsListenResult: + assert params.notifications.tools_list_changed is True + try: + return await handler(ctx, params) + finally: + with anyio.fail_after(5, shield=True): + cleanup_started.set() + await release_cleanup.wait() + cleanup_finished.set() + + server = Server("subs", on_subscriptions_listen=slow_cleanup) + with anyio.fail_after(5): + async with Client(server) as client, anyio.create_task_group() as tg: + subscription = client.listen(tools_list_changed=True) + await subscription.__aenter__() + + async def close() -> None: + with anyio.CancelScope() as scope: + scope.cancel() + try: + await anyio.Event().wait() + except anyio.get_cancelled_exc_class() as exc: + await subscription.__aexit__(type(exc), exc, exc.__traceback__) + assert cleanup_finished.is_set() + closed.set() + raise + + tg.start_soon(close) + try: + await cleanup_started.wait() + await anyio.wait_all_tasks_blocked() + assert not closed.is_set() + finally: + release_cleanup.set() + await closed.wait() + + +@pytest.mark.parametrize( + "anyio_backend", + [pytest.param(("trio", {"clock": MockClock(autojump_threshold=0)}), id="trio-mockclock")], +) +async def test_exiting_a_subscription_bounds_uncooperative_direct_handler_cleanup() -> None: + """SDK-defined: exit stops waiting after five seconds if a direct handler shields its cleanup.""" + handler = ListenHandler(InMemorySubscriptionBus()) + cleanup_started = anyio.Event() + release_cleanup = anyio.Event() + cleanup_finished = anyio.Event() + + async def shielded_cleanup( + ctx: ServerRequestContext, params: types.SubscriptionsListenRequestParams + ) -> types.SubscriptionsListenResult: + assert params.notifications.tools_list_changed is True + try: + return await handler(ctx, params) + finally: + # The watchdog must outlast the SDK's five-second cleanup cap. + with anyio.fail_after(10, shield=True): + cleanup_started.set() + await release_cleanup.wait() + cleanup_finished.set() + + server = Server("subs", on_subscriptions_listen=shielded_cleanup) + async with Client(server) as client: + subscription = client.listen(tools_list_changed=True) + with anyio.fail_after(5): + sub = await subscription.__aenter__() + try: + started = anyio.current_time() + await subscription.__aexit__(None, None, None) + assert anyio.current_time() - started == 5 # MockClock time, never wall-clock time. + assert cleanup_started.is_set() + assert not cleanup_finished.is_set() + with pytest.raises(StopAsyncIteration): + await anext(sub) + finally: + release_cleanup.set() + with anyio.fail_after(5): + await cleanup_finished.wait() + + +@pytest.mark.parametrize( + "anyio_backend", + [pytest.param(("trio", {"clock": MockClock(autojump_threshold=0)}), id="trio-mockclock")], +) +@pytest.mark.parametrize("cancelled", [False, True], ids=["normal-exit", "cancelled-exit"]) +async def test_sequential_remote_subscription_exits_do_not_wait_for_courtesy_writes( + monkeypatch: pytest.MonkeyPatch, cancelled: bool +) -> None: + """SDK-defined: remote exits leave courtesy writes to session-owned drivers, even when cancelled. + + Block the public stream `send` boundary; typed server handlers cannot wedge a client's transport write. + """ + server = Server("subs", on_subscriptions_listen=ListenHandler(InMemorySubscriptionBus())) + client_write, server_read = anyio.create_memory_object_stream[SessionMessage | Exception]() + server_write, client_read = anyio.create_memory_object_stream[SessionMessage | Exception]() + release_writes = anyio.Event() + attempted: list[types.RequestId] = [] + delivered: list[types.RequestId] = [] + send = client_write.send + + async def block_courtesy_write(item: SessionMessage | Exception) -> None: + assert isinstance(item, SessionMessage) + message = item.message + if isinstance(message, types.JSONRPCNotification) and message.method == "notifications/cancelled": + assert message.params is not None + request_id = message.params["requestId"] + attempted.append(request_id) + with anyio.fail_after(5): + await release_writes.wait() + await send(item) + delivered.append(request_id) + else: + await send(item) + + monkeypatch.setattr(client_write, "send", block_courtesy_write) + dispatcher = JSONRPCDispatcher(client_read, client_write) + with anyio.fail_after(5): + async with client_write, server_read, server_write, client_read, anyio.create_task_group() as tg: + tg.start_soon(server.run, server_read, server_write, server.create_initialization_options()) + async with ClientSession(dispatcher=dispatcher) as session: + await session.discover() + subscription_ids: list[types.RequestId] = [] + started = anyio.current_time() + try: + for _ in range(2): + subscription = listen(session, tools_list_changed=True) + sub = await subscription.__aenter__() + subscription_ids.append(sub.subscription_id) + with anyio.CancelScope() as scope: + if cancelled: + scope.cancel() + await subscription.__aexit__(None, None, None) + assert anyio.current_time() == started # MockClock time, never wall-clock time. + with pytest.raises(StopAsyncIteration): + await anext(sub) + await anyio.wait_all_tasks_blocked() + assert attempted == subscription_ids + assert delivered == [] + finally: + release_writes.set() + await anyio.wait_all_tasks_blocked() + assert set(delivered) == set(subscription_ids) + tg.cancel_scope.cancel() + + async def test_concurrent_subscriptions_demux_independently(): """Two open subscriptions each receive only their own filter's events.""" bus = InMemorySubscriptionBus() @@ -599,6 +764,8 @@ async def test_client_listen_installs_the_cache_eviction_barrier_exactly_when_a_ with anyio.fail_after(5): async with uncached_client.listen(tools_list_changed=True) as sub: # pragma: no branch assert sub._on_event is None # pyright: ignore[reportPrivateUsage] + await bus.publish(ToolsListChanged()) + assert await anext(sub) == ToolsListChanged() async def test_the_cache_eviction_barrier_maps_events_and_contains_store_faults(