[verified] fix(client): release owned dispatcher exception hook

Author: qdivan

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():

tests/client/test_session.py

@@ -45,6 +45,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 +1178,70 @@ 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:
+        async def caller_hook(_exc: Exception) -> None:
+            pass
+
+        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_direct_dispatch_roots_list_reaches_callback_with_synthesized_request_id():
     """A server-initiated roots/list over dispatcher= reaches the registered callback and round-trips