add support for supplying a progress callback to call_tool requests

Author: davemssavage

src/mcp/client/session.py

@@ -35,6 +35,13 @@ async def __call__(
     ) -> None: ...
 
 
+class ProgressFnT(Protocol):
+    async def __call__(
+        self,
+        params: types.ProgressNotificationParams,
+    ) -> None: ...
+
+
 class MessageHandlerFnT(Protocol):
     async def __call__(
         self,
@@ -91,6 +98,9 @@ class ClientSession(
         types.ServerNotification,
     ]
 ):
+    _progress_id: int
+    _in_progress: dict[types.ProgressToken, ProgressFnT]
+
     def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -114,6 +124,8 @@ def __init__(
         self._list_roots_callback = list_roots_callback or _default_list_roots_callback
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
+        self._progress_id = 0
+        self._in_progress = {}
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability()
@@ -259,19 +271,37 @@ async def call_tool(
         name: str,
         arguments: dict[str, Any] | None = None,
         read_timeout_seconds: timedelta | None = None,
+        progress_callback: ProgressFnT | None = None,
     ) -> types.CallToolResult:
         """Send a tools/call request."""
 
-        return await self.send_request(
-            types.ClientRequest(
-                types.CallToolRequest(
-                    method="tools/call",
-                    params=types.CallToolRequestParams(name=name, arguments=arguments),
-                )
-            ),
-            types.CallToolResult,
-            request_read_timeout_seconds=read_timeout_seconds,
-        )
+        if progress_callback is None:
+            progress_id = None
+            call_params = types.CallToolRequestParams(name=name, arguments=arguments)
+        else:
+            progress_id = self._progress_id
+            self._progress_id = progress_id + 1
+
+            call_meta = types.RequestParams.Meta(progressToken=progress_id)
+            call_params = types.CallToolRequestParams(
+                name=name, arguments=arguments, _meta=call_meta
+            )
+            self._in_progress[progress_id] = progress_callback
+
+        try:
+            return await self.send_request(
+                types.ClientRequest(
+                    types.CallToolRequest(
+                        method="tools/call",
+                        params=call_params,
+                    )
+                ),
+                types.CallToolResult,
+                request_read_timeout_seconds=read_timeout_seconds,
+            )
+        finally:
+            if progress_id is not None:
+                self._in_progress.pop(progress_id, None)
 
     async def list_prompts(self) -> types.ListPromptsResult:
         """Send a prompts/list request."""
@@ -384,5 +414,9 @@ async def _received_notification(
         match notification.root:
             case types.LoggingMessageNotification(params=params):
                 await self._logging_callback(params)
+            case types.ProgressNotification(params=params):
+                if params.progressToken in self._in_progress:
+                    progress_callback = self._in_progress[params.progressToken]
+                    await progress_callback(params)
             case _:
                 pass

tests/client/test_session.py

@@ -250,3 +250,112 @@ async def mock_server():
 
     # Assert that the default client info was sent
     assert received_client_info == DEFAULT_CLIENT_INFO
+
+
+@pytest.mark.anyio
+async def test_client_session_progress():
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
+        SessionMessage
+    ](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
+        SessionMessage
+    ](1)
+
+    async def mock_server():
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
+        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
+        request = ClientRequest.model_validate(
+            jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
+        )
+        assert isinstance(request.root, types.CallToolRequest)
+        assert request.root.params.meta
+        assert request.root.params.meta.progressToken is not None
+
+        progress_token = request.root.params.meta.progressToken
+
+        notifications = [
+            types.ServerNotification(
+                root=types.ProgressNotification(
+                    params=types.ProgressNotificationParams(
+                        progressToken=progress_token, progress=1
+                    ),
+                    method="notifications/progress",
+                )
+            ),
+            types.ServerNotification(
+                root=types.ProgressNotification(
+                    params=types.ProgressNotificationParams(
+                        progressToken=progress_token, progress=2
+                    ),
+                    method="notifications/progress",
+                )
+            ),
+        ]
+        result = ServerResult(types.CallToolResult(content=[]))
+
+        async with server_to_client_send:
+            for notification in notifications:
+                await server_to_client_send.send(
+                    SessionMessage(
+                        JSONRPCMessage(
+                            types.JSONRPCNotification(
+                                jsonrpc="2.0",
+                                **notification.model_dump(
+                                    by_alias=True, mode="json", exclude_none=True
+                                ),
+                            )
+                        )
+                    )
+                )
+            await server_to_client_send.send(
+                SessionMessage(
+                    JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(
+                                by_alias=True, mode="json", exclude_none=True
+                            ),
+                        )
+                    )
+                )
+            )
+
+    # Create a message handler to catch exceptions
+    async def message_handler(
+        message: RequestResponder[types.ServerRequest, types.ClientResult]
+        | types.ServerNotification
+        | Exception,
+    ) -> None:
+        if isinstance(message, Exception):
+            raise message
+
+    progress_count = 0
+
+    async with (
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+            message_handler=message_handler,
+        ) as session,
+        anyio.create_task_group() as tg,
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        tg.start_soon(mock_server)
+
+        async def progress_callback(params: types.ProgressNotificationParams):
+            nonlocal progress_count
+            progress_count = progress_count + 1
+
+        result = await session.call_tool(
+            "tool_with_progress", progress_callback=progress_callback
+        )
+
+    # Assert the result
+    assert isinstance(result, types.CallToolResult)
+    assert len(result.content) == 0
+    assert progress_count == 2