modelcontextprotocol · qdivan · Aug 14, 2026 · Aug 14, 2026
diff --git 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
@@ -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