send resume on init rather than part of join, refactor resume to be global to session rather than per request (read the spec)

Author: davemssavage

src/mcp/client/session.py

@@ -176,6 +176,20 @@ async def initialize(self) -> types.InitializeResult:
             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=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:
@@ -289,11 +303,9 @@ async def request_call_tool(
         arguments: dict[str, Any] | None = None,
         progress_callback: ProgressFnT | None = None,
     ) -> types.RequestId:
-        write, read = anyio.create_memory_object_stream[str]()
-
-        metadata = ClientMessageMetadata(on_resumption_token_update=write.send)
+        metadata = ClientMessageMetadata(on_resumption_token_update=self._request_state_manager.update_resume_token)
 
-        request_id = await self.start_request(
+        return await self.start_request(
             types.ClientRequest(
                 types.CallToolRequest(
                     method="tools/call",
@@ -307,36 +319,16 @@ async def request_call_tool(
             metadata=metadata,
         )
 
-        async def update_token() -> None:
-            try:
-                async for token in read:
-                    self._request_state_manager.update_resume_token(request_id, token)
-            except anyio.ClosedResourceError:
-                pass
-
-        async def close() -> None:
-            await write.aclose()
-            await read.aclose()
-
-        self._exit_stack.push_async_callback(update_token)
-        self._exit_stack.push_async_callback(close)
-
-        return request_id
-
     async def join_call_tool(
         self,
         request_id: types.RequestId,
         progress_callback: ProgressFnT | None = None,
         request_read_timeout_seconds: timedelta | None = None,
         done_on_timeout: bool = True,
     ) -> types.CallToolResult:
-        resume_token = self._request_state_manager.get_resume_token(request_id)
-        metadata = ClientMessageMetadata(resumption_token=resume_token)
-
         return await self.join_request(
             request_id,
             types.CallToolResult,
-            metadata=metadata,
             request_read_timeout_seconds=request_read_timeout_seconds,
             progress_callback=progress_callback,
             done_on_timeout=done_on_timeout,

src/mcp/shared/session.py

@@ -21,14 +21,12 @@
     ClientNotification,
     ClientRequest,
     ClientResult,
-    EmptyResult,
     ErrorData,
     JSONRPCError,
     JSONRPCMessage,
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
-    PingRequest,
     ProgressNotification,
     RequestParams,
     ServerNotification,
@@ -167,9 +165,9 @@ class RequestStateManager(
 ):
     def new_request(self, request: SendRequestT) -> RequestId: ...
 
-    def update_resume_token(self, request_id: RequestId, token: str) -> None: ...
+    async def update_resume_token(self, token: str) -> None: ...
 
-    def get_resume_token(self, request_id: RequestId) -> str | None: ...
+    async def get_resume_token(self) -> str | None: ...
 
     def add_progress_callback(self, request_id: RequestId, progress_callback: ProgressFnT): ...
 
@@ -210,13 +208,13 @@ class InMemoryRequestStateManager(
         ],
     ]
     _progress_callbacks: dict[RequestId, list[ProgressFnT]]
-    _resume_tokens: dict[RequestId, str]
+    _resume_token: str | None
 
     def __init__(self):
         self._request_id = 0
         self._response_streams = {}
         self._progress_callbacks = {}
-        self._resume_tokens = {}
+        self._resume_token = None
 
     def new_request(self, request: SendRequestT) -> RequestId:
         request_id = self._request_id
@@ -227,11 +225,11 @@ def new_request(self, request: SendRequestT) -> RequestId:
 
         return request_id
 
-    def update_resume_token(self, request_id: RequestId, token: str) -> None:
-        self._resume_tokens[request_id] = token
+    async def update_resume_token(self, token: str) -> None:
+        self._resume_token = token
 
-    def get_resume_token(self, request_id: RequestId) -> str | None:
-        return self._resume_tokens.get(request_id)
+    async def get_resume_token(self) -> str | None:
+        return self._resume_token
 
     def add_progress_callback(self, request_id: RequestId, progress_callback: ProgressFnT):
         progress_list = self._progress_callbacks.get(request_id)
@@ -303,7 +301,6 @@ async def close_request(self, request_id: RequestId) -> bool:
             await response_stream_reader.aclose()
 
         self._progress_callbacks.pop(request_id, None)
-        self._resume_tokens.pop(request_id, None)
 
         return response_stream is not None
 
@@ -388,17 +385,37 @@ async def start_request(
         instead.
         """
         request_id = self._request_state_manager.new_request(request)
-        return await self._send_request(
-            request_id=request_id, request=request, metadata=metadata, progress_callback=progress_callback
+        # Set up progress token if progress callback is provided
+        request_data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
+        if progress_callback is not None:
+            # Use request_id as progress token
+            if "params" not in request_data:
+                request_data["params"] = {}
+            if "_meta" not in request_data["params"]:
+                request_data["params"]["_meta"] = {}
+            request_data["params"]["_meta"]["progressToken"] = request_id
+            # Store the callback for this request
+            self._request_state_manager.add_progress_callback(request_id, progress_callback)
+
+        jsonrpc_request = JSONRPCRequest(
+            jsonrpc="2.0",
+            id=request_id,
+            **request_data,
         )
 
+        try:
+            await self._write_stream.send(SessionMessage(message=JSONRPCMessage(jsonrpc_request), metadata=metadata))
+            return request_id
+        except Exception as e:
+            await self._request_state_manager.close_request(request_id)
+            raise e
+
     async def join_request(
         self,
         request_id: RequestId,
         result_type: type[ReceiveResultT],
         request_read_timeout_seconds: timedelta | None = None,
         progress_callback: ProgressFnT | None = None,
-        metadata: MessageMetadata | None = None,
         done_on_timeout: bool = True,
     ) -> ReceiveResultT:
         """
@@ -414,15 +431,6 @@ async def join_request(
         elif self._session_read_timeout_seconds is not None:
             timeout = self._session_read_timeout_seconds.total_seconds()
 
-        if metadata:
-            # need to resend metadata - primary use case is client resume support
-            await self.send_request(
-                request=PingRequest(method="ping"),  # type: ignore
-                result_type=EmptyResult,
-                request_read_timeout_seconds=None if timeout is None else timedelta(seconds=timeout),
-                metadata=metadata,
-            )
-
         response_or_error = await self._request_state_manager.receive_response(request_id, timeout)
 
         if isinstance(response_or_error, JSONRPCError):
@@ -436,37 +444,6 @@ async def join_request(
             await self._request_state_manager.close_request(request_id)
             return result_type.model_validate(response_or_error.result)
 
-    async def _send_request(
-        self,
-        request_id: RequestId,
-        request: SendRequestT,
-        metadata: MessageMetadata = None,
-        progress_callback: ProgressFnT | None = None,
-    ):
-        # Set up progress token if progress callback is provided
-        request_data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
-        if progress_callback is not None:
-            # Use request_id as progress token
-            if "params" not in request_data:
-                request_data["params"] = {}
-            if "_meta" not in request_data["params"]:
-                request_data["params"]["_meta"] = {}
-            request_data["params"]["_meta"]["progressToken"] = request_id
-            # Store the callback for this request
-            self._request_state_manager.add_progress_callback(request_id, progress_callback)
-
-        jsonrpc_request = JSONRPCRequest(
-            jsonrpc="2.0",
-            id=request_id,
-            **request_data,
-        )
-
-        try:
-            await self._write_stream.send(SessionMessage(message=JSONRPCMessage(jsonrpc_request), metadata=metadata))
-            return request_id
-        except Exception as e:
-            await self._request_state_manager.close_request(request_id)
-            raise e
 
     async def cancel_request(self, request_id: RequestId) -> bool:
         """