add resume logic to request/join call_tool functions

Author: davemssavage

src/mcp/client/session.py

@@ -9,7 +9,7 @@
 
 import mcp.types as types
 from mcp.shared.context import RequestContext
-from mcp.shared.message import SessionMessage
+from mcp.shared.message import ClientMessageMetadata, SessionMessage
 from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder, RequestStateManager
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
@@ -289,7 +289,11 @@ async def request_call_tool(
         arguments: dict[str, Any] | None = None,
         progress_callback: ProgressFnT | None = None,
     ) -> types.RequestId:
-        return await self.start_request(
+        write, read = anyio.create_memory_object_stream[str]()
+
+        metadata = ClientMessageMetadata(on_resumption_token_update=write.send)
+
+        request_id = await self.start_request(
             types.ClientRequest(
                 types.CallToolRequest(
                     method="tools/call",
@@ -300,18 +304,39 @@ async def request_call_tool(
                 )
             ),
             progress_callback=progress_callback,
+            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 | None:
+        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,12 +21,14 @@
     ClientNotification,
     ClientRequest,
     ClientResult,
+    EmptyResult,
     ErrorData,
     JSONRPCError,
     JSONRPCMessage,
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    PingRequest,
     ProgressNotification,
     RequestParams,
     ServerNotification,
@@ -165,6 +167,10 @@ class RequestStateManager(
 ):
     def new_request(self, request: SendRequestT) -> RequestId: ...
 
+    def update_resume_token(self, request_id: RequestId, token: str) -> None: ...
+
+    def get_resume_token(self, request_id: RequestId) -> str | None: ...
+
     def add_progress_callback(self, request_id: RequestId, progress_callback: ProgressFnT): ...
 
     async def send_progress(
@@ -204,11 +210,13 @@ class ImMemoryRequestStateManager(
         ],
     ]
     _progress_callbacks: dict[RequestId, list[ProgressFnT]]
+    _resume_tokens: dict[RequestId, str]
 
     def __init__(self):
         self._request_id = 0
         self._response_streams = {}
         self._progress_callbacks = {}
+        self._resume_tokens = {}
 
     def new_request(self, request: SendRequestT) -> RequestId:
         request_id = self._request_id
@@ -219,6 +227,12 @@ 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
+
+    def get_resume_token(self, request_id: RequestId) -> str | None:
+        return self._resume_tokens.get(request_id)
+
     def add_progress_callback(self, request_id: RequestId, progress_callback: ProgressFnT):
         progress_list = self._progress_callbacks.get(request_id)
         if progress_list is None:
@@ -289,6 +303,7 @@ 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
 
@@ -373,38 +388,17 @@ async def start_request(
         instead.
         """
         request_id = self._request_state_manager.new_request(request)
-
-        # 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,
+        return await self._send_request(
+            request_id=request_id, request=request, metadata=metadata, progress_callback=progress_callback
         )
 
-        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:
         """
@@ -420,6 +414,15 @@ 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):
@@ -433,6 +436,38 @@ 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:
         """
         Cancels a request previously started via start_request