Support progress in async tools

Author: LucaButBoring

examples/snippets/clients/async_tools_client.py

@@ -90,13 +90,23 @@ async def demonstrate_batch_processing(session: ClientSession):
     print("\n=== Batch Processing Demo ===")
 
     items = ["apple", "banana", "cherry", "date", "elderberry"]
-    result = await session.call_tool("batch_operation_tool", arguments={"items": items})
+
+    # Define progress callback
+    async def progress_callback(progress: float, total: float | None, message: str | None) -> None:
+        progress_pct = int(progress * 100) if progress else 0
+        total_str = f"/{int(total * 100)}%" if total else ""
+        message_str = f" - {message}" if message else ""
+        print(f"Progress: {progress_pct}{total_str}{message_str}")
+
+    result = await session.call_tool(
+        "batch_operation_tool", arguments={"items": items}, progress_callback=progress_callback
+    )
 
     if result.operation:
         token = result.operation.token
         print(f"Batch operation started with token: {token}")
 
-        # Poll for status with progress tracking
+        # Poll for status
         while True:
             status = await session.get_operation_status(token)
             print(f"Status: {status.status}")

src/mcp/shared/session.py

@@ -16,11 +16,14 @@
 from mcp.types import (
     CONNECTION_CLOSED,
     INVALID_PARAMS,
+    CallToolResult,
     CancelledNotification,
     ClientNotification,
     ClientRequest,
     ClientResult,
     ErrorData,
+    GetOperationPayloadRequest,
+    GetOperationPayloadResult,
     JSONRPCError,
     JSONRPCMessage,
     JSONRPCNotification,
@@ -177,6 +180,7 @@ class BaseSession(
     _request_id: int
     _in_flight: dict[RequestId, RequestResponder[ReceiveRequestT, SendResultT]]
     _progress_callbacks: dict[RequestId, ProgressFnT]
+    _operation_requests: dict[str, RequestId]
 
     def __init__(
         self,
@@ -196,6 +200,7 @@ def __init__(
         self._session_read_timeout_seconds = read_timeout_seconds
         self._in_flight = {}
         self._progress_callbacks = {}
+        self._operation_requests = {}
         self._exit_stack = AsyncExitStack()
 
     async def __aenter__(self) -> Self:
@@ -251,6 +256,7 @@ async def send_request(
             # Store the callback for this request
             self._progress_callbacks[request_id] = progress_callback
 
+        pop_progress: RequestId | None = request_id
         try:
             jsonrpc_request = JSONRPCRequest(
                 jsonrpc="2.0",
@@ -285,11 +291,28 @@ async def send_request(
             if isinstance(response_or_error, JSONRPCError):
                 raise McpError(response_or_error.error)
             else:
-                return result_type.model_validate(response_or_error.result)
+                result = result_type.model_validate(response_or_error.result)
+                if isinstance(result, CallToolResult) and result.operation is not None:
+                    # Store mapping of operation token to request ID for async operations
+                    self._operation_requests[result.operation.token] = request_id
+
+                    # Don't pop the progress function if we were given one
+                    pop_progress = None
+                elif isinstance(request, GetOperationPayloadRequest) and isinstance(result, GetOperationPayloadResult):
+                    # Checked request and result to ensure no error
+                    operation_token = request.params.token
+
+                    # Pop the progress function for the original request
+                    pop_progress = self._operation_requests[operation_token]
+
+                    # Pop the token mapping since we know we won't need it anymore
+                    self._operation_requests.pop(operation_token, None)
+                return result
 
         finally:
             self._response_streams.pop(request_id, None)
-            self._progress_callbacks.pop(request_id, None)
+            if pop_progress:
+                self._progress_callbacks.pop(pop_progress, None)
             await response_stream.aclose()
             await response_stream_reader.aclose()
 

tests/server/fastmcp/test_integration.py

@@ -760,6 +760,42 @@ async def test_async_tools(server_transport: str, server_url: str) -> None:
             with pytest.raises(Exception):  # Should raise error when trying to access expired operation
                 await session.get_operation_result(quick_token)
 
+            # Test batch operation with progress notifications
+            progress_received = False
+
+            async def progress_callback(progress: float, total: float | None, message: str | None) -> None:
+                nonlocal progress_received
+                progress_received = True
+                assert 0.0 <= progress <= 1.0  # Progress should be between 0 and 1
+
+            batch_result = await session.call_tool(
+                "batch_operation_tool",
+                {"items": ["apple", "banana", "cherry"]},
+                progress_callback=progress_callback,
+            )
+            assert batch_result.operation is not None
+            batch_token = batch_result.operation.token
+
+            while True:
+                status = await session.get_operation_status(batch_token)
+
+                if status.status == "completed":
+                    final_result = await session.get_operation_result(batch_token)
+                    assert not final_result.result.isError
+                    # Should have structured content with processed items
+                    if final_result.result.structuredContent:
+                        # Structured content is wrapped in {"result": [...]} for list return types
+                        assert isinstance(final_result.result.structuredContent, dict)
+                        assert "result" in final_result.result.structuredContent
+                        assert isinstance(final_result.result.structuredContent["result"], list)
+                        assert len(final_result.result.structuredContent["result"]) == 3
+                    break
+                elif status.status == "failed":
+                    pytest.fail(f"Batch operation failed: {status.error}")
+
+            # Assert that we received at least one progress notification
+            assert progress_received, "Should have received progress notifications during batch operation"
+
 
 # Test async tools example with legacy protocol
 @pytest.mark.anyio