@@ -120,19 +120,11 @@ class ClientSession:
120120 """Client half of an MCP connection, running on a `Dispatcher`.
121121
122122 Construct it over a transport's stream pair (or pass a pre-built
123- `dispatcher=` instead, e.g. a `DirectDispatcher` for in-process
124- embedding), enter it as an async context manager, then call
125- `initialize()`. The receive loop, request correlation, and per-request
126- concurrency live in the dispatcher; this class owns the MCP type layer:
127- typed requests, the initialize handshake, and routing server-initiated
128- traffic to the constructor callbacks.
129-
130- Transport-level `Exception` items reach `message_handler` only when the
131- session builds its own dispatcher from streams, where it wires the
132- dispatcher's `on_stream_exception` itself. Faults are delivered
133- concurrently in the session's task group, like notifications — never
134- inline in the read loop — so the handler may await session I/O, and one
135- that raises costs that delivery, not the connection.
123+ `dispatcher=`), enter as an async context manager, then call
124+ `initialize()`. The dispatcher owns the receive loop and request
125+ correlation; this class owns the typed MCP layer and the constructor
126+ callbacks. Transport `Exception` items reach `message_handler` only when
127+ the session builds its own dispatcher from a stream pair.
136128 """
137129
138130 def __init__ (
@@ -168,8 +160,7 @@ def __init__(
168160 else :
169161 if read_stream is None or write_stream is None :
170162 raise ValueError ("read_stream and write_stream are required when no dispatcher is given" )
171- # Built here (inert until run() starts in __aenter__) so notifications
172- # can be sent before entering the context manager, as before.
163+ # Built eagerly so notifications can be sent before entering the context manager.
173164 self ._dispatcher = JSONRPCDispatcher (
174165 read_stream , write_stream , on_stream_exception = self ._on_stream_exception
175166 )
@@ -180,20 +171,14 @@ async def __aenter__(self) -> Self:
180171 try :
181172 await self ._task_group .start (self ._dispatcher .run , self ._on_request , self ._on_notify )
182173 except BaseException :
183- # A cancellation landing here (e.g. the caller wrapped connect in
184- # `move_on_after`) would abandon the entered task group, and anyio
185- # later raises "exited non-innermost cancel scope" instead of a
186- # clean timeout. Unwind the group before propagating; cancelling
187- # its scope first keeps __aexit__ from blocking under the
188- # still-active cancellation.
174+ # Unwind the entered task group before propagating: a cancellation
175+ # landing here (e.g. `move_on_after` around connect) would abandon
176+ # it and anyio would later raise "exited non-innermost cancel scope".
189177 task_group = self ._task_group
190178 self ._task_group = None
191179 task_group .cancel_scope .cancel ()
192- # Shield the group's own scope (not a new one: scope exits must
193- # stay LIFO) so a pending outer cancellation cannot re-fire
194- # inside __aexit__; the join is prompt because the scope is
195- # cancelled. The original exception then propagates from the
196- # `raise`; a child error supersedes it, raised by __aexit__.
180+ # Shield the group's own scope (a new one would break LIFO exit)
181+ # so a pending outer cancellation cannot re-fire inside __aexit__.
197182 task_group .cancel_scope .shield = True
198183 await task_group .__aexit__ (None , None , None )
199184 raise
@@ -205,8 +190,7 @@ async def __aexit__(
205190 exc_val : BaseException | None ,
206191 exc_tb : TracebackType | None ,
207192 ) -> bool | None :
208- # Exit must not block: cancel the dispatcher and any in-flight
209- # callbacks rather than waiting for them.
193+ # Exit must not block: cancel the dispatcher and in-flight callbacks.
210194 assert self ._task_group is not None
211195 self ._task_group .cancel_scope .cancel ()
212196 result = await self ._task_group .__aexit__ (exc_type , exc_val , exc_tb )
@@ -223,18 +207,14 @@ async def send_request(
223207 ) -> ReceiveResultT :
224208 """Send a request and wait for its typed result.
225209
226- A per-request read timeout takes precedence over the session-level
227- one. `metadata` carries transport hints: `ClientMessageMetadata`
228- resumption fields (streamable HTTP), or a
229- `ServerMessageMetadata.related_request_id` to route the message onto
230- an originating request's stream.
210+ Args:
211+ metadata: Transport hints: `ClientMessageMetadata` resumption fields
212+ (streamable HTTP), or a `ServerMessageMetadata.related_request_id`
213+ routing the message onto the originating request's stream.
231214
232215 Raises:
233- MCPError: The server responded with an error, or the read timeout
234- elapsed, or the connection closed while sending or waiting.
235- RuntimeError: Called before entering the context manager. Raised
236- by the stream-built dispatcher; a user-supplied `dispatcher=`
237- may not enforce this.
216+ MCPError: Error response, read timeout, or connection closed.
217+ RuntimeError: Called before entering the context manager.
238218 """
239219 data = request .model_dump (by_alias = True , mode = "json" , exclude_none = True )
240220 method : str = data ["method" ]
@@ -253,12 +233,10 @@ async def send_request(
253233 elif isinstance (metadata , ServerMessageMetadata ):
254234 related_request_id = metadata .related_request_id
255235 if method == "initialize" :
256- # The spec forbids cancelling initialize; opt out of the
257- # dispatcher's courtesy cancel-on-abandon.
236+ # The spec forbids cancelling initialize.
258237 opts ["cancel_on_abandon" ] = False
259238 if related_request_id is not None and isinstance (self ._dispatcher , JSONRPCDispatcher ):
260- # Related-request routing is JSON-RPC stream plumbing; other
261- # dispatchers have no per-request streams to route onto.
239+ # Only JSON-RPC dispatchers have per-request streams to route onto.
262240 raw = await self ._dispatcher .send_raw_request (
263241 method , data .get ("params" ), opts , _related_request_id = related_request_id
264242 )
@@ -273,7 +251,7 @@ async def send_notification(
273251 ) -> None :
274252 """Send a one-way notification. Usable before entering the context manager."""
275253 data = notification .model_dump (by_alias = True , mode = "json" , exclude_none = True )
276- # `is not None`, not truthiness : request ids are opaque and 0 is valid.
254+ # `is not None`: request ids are opaque and 0 is valid.
277255 if related_request_id is not None and isinstance (self ._dispatcher , JSONRPCDispatcher ):
278256 await self ._dispatcher .notify (data ["method" ], data .get ("params" ), _related_request_id = related_request_id )
279257 else :
@@ -529,17 +507,8 @@ async def send_roots_list_changed(self) -> None:
529507 async def _on_request (
530508 self , dctx : DispatchContext [TransportContext ], method : str , params : Mapping [str , Any ] | None
531509 ) -> dict [str , Any ]:
532- """Answer a server-initiated request via the registered callbacks.
533-
534- An unknown method raises `MCPError` (METHOD_NOT_FOUND), which the
535- dispatcher puts on the wire as-is; malformed params for a known method
536- raise `ValidationError`, which the dispatcher answers with
537- INVALID_PARAMS; an `ErrorData` returned by a callback becomes the
538- error response.
539- """
510+ """Answer a server-initiated request via the registered callbacks."""
540511 if method not in _SERVER_REQUEST_METHODS :
541- # Unknown methods are METHOD_NOT_FOUND (-32601) per JSON-RPC 2.0,
542- # not validation failures (-32602).
543512 raise MCPError (code = types .METHOD_NOT_FOUND , message = "Method not found" , data = method )
544513 payload : dict [str , Any ] = {"method" : method }
545514 if params is not None :
@@ -577,25 +546,18 @@ async def _on_notify(
577546 logger .warning ("Failed to validate notification: %s" , payload , exc_info = True )
578547 return
579548 if isinstance (notification , types .CancelledNotification ):
580- # The dispatcher already applied the cancellation to the in-flight
581- # request; message_handler never sees it, so handlers matching
582- # exhaustively over ServerNotification need no arm for it.
549+ # The dispatcher already applied the cancellation; not surfaced to message_handler.
583550 return
584551 if isinstance (notification , types .LoggingMessageNotification ):
585552 await self ._logging_callback (notification .params )
586553 await self ._message_handler (notification )
587554
588555 async def _on_stream_exception (self , exc : Exception ) -> None :
589- """Spawn delivery of a transport-level fault (connection error, parse error) to message_handler .
556+ """Deliver a transport-level fault to message_handler via a spawned task .
590557
591- The dispatcher awaits this observer inline in its read loop, so the
592- handler must not run here: a slow handler would head-of-line block the
593- session, and one that awaits session I/O (e.g. sends a ping) would
594- deadlock against the parked loop. Spawn it instead, with the same
595- containment notification deliveries get.
558+ Running the handler inline would park the dispatcher's read loop and
559+ deadlock handlers that await session I/O.
596560 """
597- # The dispatcher only runs inside the task group entered in
598- # __aenter__, so the group is always live when it calls back here.
599561 assert self ._task_group is not None
600562 self ._task_group .start_soon (self ._deliver_stream_exception , exc )
601563
0 commit comments