diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 895339ca18..93cd280e5d 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -415,6 +415,7 @@ def __init__( self._negotiated_version: str | None = None self._stamp: Callable[[dict[str, Any], CallOptions], None] = _preconnect_stamp self._task_group: anyio.abc.TaskGroup | None = None + self._owned_stream_exception_hook: Callable[[Exception], Any] | None = None # subscriptions/listen demux routes; membership decides ack consumption (raw listens are never registered) self._listen_routes: dict[RequestId, ListenRoute] = {} if dispatcher is not None: @@ -424,11 +425,11 @@ def __init__( if isinstance(dispatcher, JSONRPCDispatcher) and dispatcher.on_stream_exception is None: # Route transport-level Exception items into message_handler — only # stream-backed dispatchers carry these; DirectDispatcher has none. - # Don't clobber a caller-supplied hook. - # TODO(L78): this leaves a bound-method ref on the dispatcher after the - # session exits (memory pin) and a second wrap of the same dispatcher would - # skip install. The Transport-as-Dispatcher rework (L77) removes this seam. - dispatcher.on_stream_exception = self._on_stream_exception + # Don't clobber a caller-supplied hook, and remember the exact bound + # method object so shutdown can remove only our own installation. + hook = self._on_stream_exception + dispatcher.on_stream_exception = hook + self._owned_stream_exception_hook = hook else: if read_stream is None or write_stream is None: raise ValueError("read_stream and write_stream are required when no dispatcher is given") @@ -465,6 +466,7 @@ async def __aenter__(self) -> Self: await task_group.__aexit__(None, None, None) finally: self._close_binding_queues() + self._remove_owned_stream_exception_hook() raise return self @@ -482,9 +484,18 @@ async def __aexit__( finally: self._close_binding_queues() self._settle_listen_routes_closed() + self._remove_owned_stream_exception_hook() await resync_tracer() return result + def _remove_owned_stream_exception_hook(self) -> None: + """Remove the stream hook only if this session still owns the dispatcher slot.""" + hook = self._owned_stream_exception_hook + if hook is not None and isinstance(self._dispatcher, JSONRPCDispatcher): + if self._dispatcher.on_stream_exception is hook: + self._dispatcher.on_stream_exception = None + self._owned_stream_exception_hook = None + def _close_binding_queues(self) -> None: # Unclosed memory object streams warn at garbage collection; close is idempotent. for send, receive in self._binding_queues.values(): diff --git a/tests/client/test_session.py b/tests/client/test_session.py index 6e5e0a14d8..3af077f016 100644 --- a/tests/client/test_session.py +++ b/tests/client/test_session.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator, Mapping from contextlib import AsyncExitStack, asynccontextmanager from typing import Any, cast +from unittest.mock import AsyncMock import anyio import anyio.abc @@ -45,6 +46,7 @@ from mcp.server import Server, ServerRequestContext from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair from mcp.shared.dispatcher import CallOptions, DispatchContext, OnNotify, OnNotifyIntercept, OnRequest +from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher from mcp.shared.message import SessionMessage from mcp.shared.subscriptions import SUBSCRIPTION_ID_META_KEY from mcp.shared.transport_context import TransportContext @@ -1177,6 +1179,90 @@ async def server_on_notify( assert notified == ["notifications/roots/list_changed"] +@pytest.mark.anyio +async def test_dispatcher_keyword_removes_its_stream_exception_hook_on_exit(): + """An injected stream dispatcher must not retain the exited session through its hook.""" + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](1) + try: + dispatcher = JSONRPCDispatcher(s2c_recv, c2s_send) + session = ClientSession(dispatcher=dispatcher) + + assert dispatcher.on_stream_exception is not None + + async with session: + pass + + assert dispatcher.on_stream_exception is None + finally: + s2c_send.close() + s2c_recv.close() + c2s_send.close() + c2s_recv.close() + + +@pytest.mark.anyio +async def test_dispatcher_keyword_reinstalls_stream_exception_hook_for_reused_dispatcher(): + """A dispatcher can be wrapped by another session after the first session exits.""" + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](1) + try: + dispatcher = JSONRPCDispatcher(s2c_recv, c2s_send) + async with ClientSession(dispatcher=dispatcher): + pass + assert dispatcher.on_stream_exception is None + + async with ClientSession(dispatcher=dispatcher): + assert dispatcher.on_stream_exception is not None + assert dispatcher.on_stream_exception is None + finally: + s2c_send.close() + s2c_recv.close() + c2s_send.close() + c2s_recv.close() + + +@pytest.mark.anyio +async def test_dispatcher_keyword_preserves_caller_stream_exception_hook_on_exit(): + """A hook supplied by the dispatcher owner remains installed after session shutdown.""" + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](1) + try: + caller_hook = AsyncMock() + dispatcher = JSONRPCDispatcher(s2c_recv, c2s_send, on_stream_exception=caller_hook) + async with ClientSession(dispatcher=dispatcher): + pass + + assert dispatcher.on_stream_exception is caller_hook + finally: + s2c_send.close() + s2c_recv.close() + c2s_send.close() + c2s_recv.close() + + +@pytest.mark.anyio +async def test_dispatcher_keyword_preserves_replacement_stream_exception_hook_on_exit(): + """A caller replacement made while a session is alive survives shutdown (SDK-defined ownership).""" + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1) + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](1) + try: + dispatcher = JSONRPCDispatcher(s2c_recv, c2s_send) + session = ClientSession(dispatcher=dispatcher) + assert dispatcher.on_stream_exception is not None + + caller_hook = AsyncMock() + async with session: + dispatcher.on_stream_exception = caller_hook + + assert dispatcher.on_stream_exception is caller_hook + finally: + s2c_send.close() + s2c_recv.close() + c2s_send.close() + c2s_recv.close() + + @pytest.mark.anyio async def test_direct_dispatch_roots_list_reaches_callback_with_synthesized_request_id(): """A server-initiated roots/list over dispatcher= reaches the registered callback and round-trips