Refactor code to send resume as part of join call rather than it, this results in the response being consumed prior to the join, also added a capability that identifies whether the server/transport supports resumption that is passed during initialisation

Author: davemssavage

src/mcp/client/session.py

@@ -135,6 +135,7 @@ def __init__(
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
+        self._resumable = False
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
@@ -172,24 +173,12 @@ async def initialize(self) -> types.InitializeResult:
         if result.protocolVersion not in SUPPORTED_PROTOCOL_VERSIONS:
             raise RuntimeError(f"Unsupported protocol version from the server: {result.protocolVersion}")
 
+        self._resumable = result.capabilities.resume and result.capabilities.resume.resumable
+
         await self.send_notification(
             types.ClientNotification(types.InitializedNotification(method="notifications/initialized"))
         )
 
-        resume_token = await self._request_state_manager.get_resume_token()
-        if resume_token:
-            metadata = ClientMessageMetadata(resumption_token=resume_token)
-            timeout = None
-            if self._session_read_timeout_seconds is not None:
-                timeout = self._session_read_timeout_seconds.total_seconds()
-
-            await self.send_request(
-                request=types.PingRequest(method="ping"),  # type: ignore
-                result_type=types.EmptyResult,
-                request_read_timeout_seconds=None if timeout is None else timedelta(seconds=timeout),
-                metadata=metadata,
-            )
-
         return result
 
     async def send_ping(self) -> types.EmptyResult:
@@ -303,21 +292,58 @@ async def request_call_tool(
         arguments: dict[str, Any] | None = None,
         progress_callback: ProgressFnT | None = None,
     ) -> types.RequestId:
-        metadata = ClientMessageMetadata(on_resumption_token_update=self._request_state_manager.update_resume_token)
-
-        return await self.start_request(
-            types.ClientRequest(
-                types.CallToolRequest(
-                    method="tools/call",
-                    params=types.CallToolRequestParams(
-                        name=name,
-                        arguments=arguments,
+        if self._resumable:
+            send_stream, receive_stream = anyio.create_memory_object_stream[str](1)
+
+            async def close() -> None:
+                await send_stream.aclose()
+                await receive_stream.aclose()
+
+            self._exit_stack.push_async_callback(close)
+
+            with send_stream, receive_stream:
+
+                async def send_token(token: str):
+                    try:
+                        await send_stream.send(token)
+                    except anyio.BrokenResourceError as e:
+                        raise e
+
+                metadata = ClientMessageMetadata(on_resumption_token_update=send_token)
+
+                request_id = await self.start_request(
+                    types.ClientRequest(
+                        types.CallToolRequest(
+                            method="tools/call",
+                            params=types.CallToolRequestParams(
+                                name=name,
+                                arguments=arguments,
+                            ),
+                        )
                     ),
+                    progress_callback=progress_callback,
+                    metadata=metadata,
                 )
-            ),
-            progress_callback=progress_callback,
-            metadata=metadata,
-        )
+
+                await anyio.lowlevel.checkpoint()
+
+                token = await receive_stream.receive()
+                await self._request_state_manager.update_resume_token(request_id, token)
+
+                return request_id
+        else:
+            return await self.start_request(
+                types.ClientRequest(
+                    types.CallToolRequest(
+                        method="tools/call",
+                        params=types.CallToolRequestParams(
+                            name=name,
+                            arguments=arguments,
+                        ),
+                    )
+                ),
+                progress_callback=progress_callback,
+            )
 
     async def join_call_tool(
         self,

src/mcp/client/streamable_http.py

@@ -29,6 +29,7 @@
     JSONRPCRequest,
     JSONRPCResponse,
     RequestId,
+    ResumeCapability,
 )
 
 logger = logging.getLogger(__name__)
@@ -136,18 +137,26 @@ def _maybe_extract_session_id_from_response(
     def _maybe_extract_protocol_version_from_message(
         self,
         message: JSONRPCMessage,
-    ) -> None:
+    ) -> JSONRPCMessage:
         """Extract protocol version from initialization response message."""
         if isinstance(message.root, JSONRPCResponse) and message.root.result:
             try:
                 # Parse the result as InitializeResult for type safety
                 init_result = InitializeResult.model_validate(message.root.result)
                 self.protocol_version = str(init_result.protocolVersion)
                 logger.info(f"Negotiated protocol version: {self.protocol_version}")
+                if init_result.capabilities.resume is None:
+                    # resumeablity is predicated on the server and the transport
+                    # this assumes that if the server hasn't explicitly configured
+                    # that streamable http transports are resumeable
+                    init_result.capabilities.resume = ResumeCapability(resumable=True)
+                    message.root.result = init_result.model_dump()
             except Exception as exc:
                 logger.warning(f"Failed to parse initialization response as InitializeResult: {exc}")
                 logger.warning(f"Raw result: {message.root.result}")
 
+        return message
+
     async def _handle_sse_event(
         self,
         sse: ServerSentEvent,
@@ -183,7 +192,10 @@ async def _handle_sse_event(
 
             except Exception as exc:
                 logger.exception("Error parsing SSE message")
-                await read_stream_writer.send(exc)
+                try:
+                    await read_stream_writer.send(exc)
+                except anyio.BrokenResourceError:
+                    pass
                 return False
         else:
             logger.warning(f"Unknown SSE event: {sse.event}")
@@ -303,7 +315,7 @@ async def _handle_json_response(
 
             # Extract protocol version from initialization response
             if is_initialization:
-                self._maybe_extract_protocol_version_from_message(message)
+                message = self._maybe_extract_protocol_version_from_message(message)
 
             session_message = SessionMessage(message)
             await read_stream_writer.send(session_message)
@@ -333,7 +345,10 @@ async def _handle_sse_response(
                     break
         except Exception as e:
             logger.exception("Error reading SSE stream:")
-            await ctx.read_stream_writer.send(e)
+            try:
+                await ctx.read_stream_writer.send(e)
+            except anyio.ClosedResourceError:
+                pass
 
     async def _handle_unexpected_content_type(
         self,