Client: collapse _inproc_server/_transport into a single _connect closure

__post_init__ now resolves a single _connect closure from the shape of
the server argument alone (in-process vs URL vs Transport instance).
mode and raise_exceptions are passed to the closure at enter time so
they're read at the same moment __aenter__ reads them for the handshake
step. _build_session collapses to one line of logic; the
mutually-exclusive Optional fields and the assert that guarded them are
gone.

JSONRPCDispatcher.on_stream_exception is now public-mutable so
ClientSession can install its message_handler routing after the
dispatcher is built; the install only happens when no caller-supplied
hook is already set.

ClientSession.adopt() now clears the opposite result slot so at most one
of initialize_result/discover_result is non-None by construction.

Author: maxisbey

src/mcp/client/client.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Awaitable, Callable, Mapping
 from contextlib import AsyncExitStack
 from dataclasses import KW_ONLY, dataclass, field
 from typing import Any, Literal, TypeVar
@@ -21,6 +21,7 @@
 from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
 from mcp.shared.dispatcher import Dispatcher, ProgressFnT
 from mcp.shared.exceptions import MCPDeprecationWarning, MCPError
+from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
 from mcp.shared.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
 from mcp.types import (
     METHOD_NOT_FOUND,
@@ -50,6 +51,43 @@
 
 _T = TypeVar("_T")
 
+_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]]
+"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources
+are needed onto the exit stack and hand back the ``Dispatcher`` ``ClientSession`` will drive.
+``mode`` and ``raise_exceptions`` are passed at call time so they're read at the same moment
+``__aenter__`` reads them for the handshake step."""
+
+
+def _connect_transport(transport: Transport) -> _Connector:
+    """Connector for the stream-backed paths (URL, user-supplied ``Transport``)."""
+
+    async def connect(exit_stack: AsyncExitStack, _mode: ConnectMode, _raise_exceptions: bool) -> Dispatcher[Any]:
+        read_stream, write_stream = await exit_stack.enter_async_context(transport)
+        return JSONRPCDispatcher(read_stream, write_stream)
+
+    return connect
+
+
+def _connect_inproc(server: Server[Any]) -> _Connector:
+    """Connector for an in-process ``Server``: legacy mode drives the stream loop via
+    ``InMemoryTransport``; any other mode drives the modern per-request path through a
+    ``DirectDispatcher`` peer pair (no streams, no JSON-RPC framing, no initialize handshake)."""
+
+    async def connect(exit_stack: AsyncExitStack, mode: ConnectMode, raise_exceptions: bool) -> Dispatcher[Any]:
+        if mode == "legacy":
+            transport = InMemoryTransport(server, raise_exceptions=raise_exceptions)
+            read_stream, write_stream = await exit_stack.enter_async_context(transport)
+            return JSONRPCDispatcher(read_stream, write_stream)
+        lifespan_state = await exit_stack.enter_async_context(server.lifespan(server))
+        client_disp, server_disp = create_direct_dispatcher_pair()
+        tg = await exit_stack.enter_async_context(anyio.create_task_group())
+        exit_stack.callback(server_disp.close)
+        on_request = modern_on_request(server, lifespan_state, raise_exceptions=raise_exceptions)
+        await tg.start(server_disp.run, on_request, _no_inbound_client_notifications)
+        return client_disp
+
+    return connect
+
 
 def _connected(value: _T | None) -> _T:
     """Narrow a post-handshake session attribute from ``T | None`` to ``T``.
@@ -161,19 +199,9 @@ async def main():
     _entered: bool = field(init=False, default=False)
     _session: ClientSession | None = field(init=False, default=None)
     _exit_stack: AsyncExitStack | None = field(init=False, default=None)
-    _transport: Transport | None = field(init=False, default=None)
-    _inproc_server: Server[Any] | None = field(init=False, default=None)
+    _connect: _Connector = field(init=False, repr=False, compare=False)
 
     def __post_init__(self) -> None:
-        if isinstance(self.server, MCPServer):
-            self._inproc_server = self.server._lowlevel_server  # pyright: ignore[reportPrivateUsage]
-        elif isinstance(self.server, Server):
-            self._inproc_server = self.server
-        elif isinstance(self.server, str):
-            self._transport = streamable_http_client(self.server)
-        else:
-            self._transport = self.server
-
         if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS:
             hint = (
                 f" ({self.mode!r} is a handshake-era version — use mode='legacy')"
@@ -184,31 +212,20 @@ def __post_init__(self) -> None:
                 f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}"
             )
 
-    async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession:
-        """Set up the dispatcher/transport and return an un-entered ClientSession."""
-        dispatcher: Dispatcher[Any] | None
-        if self._inproc_server is not None and self.mode != "legacy":
-            # Modern in-process path: drive the server through a DirectDispatcher peer-pair
-            # with one `serve_one` per request — no streams, no initialize handshake.
-            lifespan_state = await exit_stack.enter_async_context(self._inproc_server.lifespan(self._inproc_server))
-            client_disp, server_disp = create_direct_dispatcher_pair()
-            tg = await exit_stack.enter_async_context(anyio.create_task_group())
-            exit_stack.callback(server_disp.close)
-            on_request = modern_on_request(self._inproc_server, lifespan_state, raise_exceptions=self.raise_exceptions)
-            await tg.start(server_disp.run, on_request, _no_inbound_client_notifications)
-            dispatcher = client_disp
-            read_stream = write_stream = None
+        srv = self.server
+        if isinstance(srv, MCPServer):
+            srv = srv._lowlevel_server  # pyright: ignore[reportPrivateUsage]
+        if isinstance(srv, Server):
+            self._connect = _connect_inproc(srv)
+        elif isinstance(srv, str):
+            self._connect = _connect_transport(streamable_http_client(srv))
         else:
-            if self._inproc_server is not None:
-                transport: Transport = InMemoryTransport(self._inproc_server, raise_exceptions=self.raise_exceptions)
-            else:
-                assert self._transport is not None
-                transport = self._transport
-            read_stream, write_stream = await exit_stack.enter_async_context(transport)
-            dispatcher = None
+            self._connect = _connect_transport(srv)
+
+    async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession:
+        """Enter the resolved connector and return an un-entered ClientSession."""
+        dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions)
         return ClientSession(
-            read_stream=read_stream,
-            write_stream=write_stream,
             dispatcher=dispatcher,
             read_timeout_seconds=self.read_timeout_seconds,
             sampling_callback=self.sampling_callback,

src/mcp/client/session.py

@@ -216,6 +216,14 @@ def __init__(
             if read_stream is not None or write_stream is not None:
                 raise ValueError("pass read_stream/write_stream or dispatcher, not both")
             self._dispatcher: Dispatcher[Any] = dispatcher
+            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(maxisbey): 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 removes this seam.
+                dispatcher.on_stream_exception = self._on_stream_exception
         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")
@@ -358,6 +366,9 @@ async def initialize(self) -> types.InitializeResult:
     def adopt(self, result: types.InitializeResult | types.DiscoverResult) -> None:
         """Install negotiated state from a result the caller already holds (no wire traffic).
 
+        Clears the opposite slot, so at most one of `initialize_result` /
+        `discover_result` is ever non-None.
+
         Raises:
             RuntimeError: `result` is a `DiscoverResult` whose `supported_versions`
                 shares nothing with this client's `MODERN_PROTOCOL_VERSIONS`.
@@ -374,10 +385,12 @@ def adopt(self, result: types.InitializeResult | types.DiscoverResult) -> None:
             capabilities = self._build_capabilities().model_dump(by_alias=True, mode="json", exclude_none=True)
             self._stamp = _make_modern_stamp(mutual[-1], client_info, capabilities)
             self._discover_result = result
+            self._initialize_result = None
             self._negotiated_version = mutual[-1]
         else:
             self._stamp = _make_handshake_stamp(result.protocol_version)
             self._initialize_result = result
+            self._discover_result = None
             self._negotiated_version = result.protocol_version
 
     async def discover(self) -> types.DiscoverResult:

src/mcp/shared/jsonrpc_dispatcher.py

@@ -75,7 +75,7 @@ def handler_exception_to_error_data(exc: BaseException) -> ErrorData | None:
     with empty ``data`` (no pydantic text on the wire). Returns ``None`` for
     any other exception so each caller applies its own catch-all -
     `JSONRPCDispatcher` currently pins ``code=0`` for v1 compat,
-    `to_jsonrpc_response` uses `INTERNAL_ERROR`.
+    the modern HTTP entry uses `INTERNAL_ERROR`.
     """
     if isinstance(exc, MCPError):
         return exc.error
@@ -268,7 +268,10 @@ def __init__(
         self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode
         self._raise_handler_exceptions = raise_handler_exceptions
         self._inline_methods = inline_methods
-        self._on_stream_exception = on_stream_exception
+        self.on_stream_exception = on_stream_exception
+        """Observer for ``Exception`` items on the read stream. Mutable so a session can
+        bind it after the dispatcher is built (e.g. ``ClientSession`` routing into
+        ``message_handler``); only consulted inside ``run()`` so pre-enter assignment is safe."""
 
         self._next_id = 0
         self._pending: dict[RequestId, _Pending] = {}
@@ -484,11 +487,11 @@ async def _dispatch(
         are awaited; any other `await` would head-of-line block the read loop.
         """
         if isinstance(item, Exception):
-            if self._on_stream_exception is None:
+            if self.on_stream_exception is None:
                 logger.debug("transport yielded exception: %r", item)
                 return
             try:
-                await self._on_stream_exception(item)
+                await self.on_stream_exception(item)
             except Exception:
                 logger.exception("on_stream_exception observer raised")
             return