add progress notificaiton to the client

Author: ihrpr

src/mcp/client/session.py

@@ -8,7 +8,7 @@
 import mcp.types as types
 from mcp.shared.context import RequestContext
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import BaseSession, RequestResponder
+from mcp.shared.session import BaseSession, ProgressCallbackFnT, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
 DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
@@ -270,18 +270,23 @@ async def call_tool(
         name: str,
         arguments: dict[str, Any] | None = None,
         read_timeout_seconds: timedelta | None = None,
+        progress_callback: ProgressCallbackFnT | None = None,
     ) -> types.CallToolResult:
-        """Send a tools/call request."""
+        """Send a tools/call request with optional progress callback support."""
 
         return await self.send_request(
             types.ClientRequest(
                 types.CallToolRequest(
                     method="tools/call",
-                    params=types.CallToolRequestParams(name=name, arguments=arguments),
+                    params=types.CallToolRequestParams(
+                        name=name,
+                        arguments=arguments,
+                    ),
                 )
             ),
             types.CallToolResult,
             request_read_timeout_seconds=read_timeout_seconds,
+            progress_callback=progress_callback,
         )
 
     async def list_prompts(self, cursor: str | None = None) -> types.ListPromptsResult:

src/mcp/server/fastmcp/server.py

@@ -963,7 +963,6 @@ async def report_progress(
             total: Optional total value e.g. 100
             message: Optional message e.g. Starting render...
         """
-
         progress_token = (
             self.request_context.meta.progressToken
             if self.request_context.meta

src/mcp/shared/session.py

@@ -3,7 +3,7 @@
 from contextlib import AsyncExitStack
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Protocol, TypeVar
 
 import anyio
 import httpx
@@ -24,6 +24,7 @@
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    ProgressNotification,
     RequestParams,
     ServerNotification,
     ServerRequest,
@@ -42,6 +43,22 @@
 RequestId = str | int
 
 
+class ProgressCallbackFnT(Protocol):
+    """Protocol for progress notification callbacks."""
+
+    def __call__(
+        self, progress: float, total: float | None, message: str | None
+    ) -> None:
+        """Called when progress updates are received.
+
+        Args:
+            progress: Current progress value
+            total: Total progress value (if known), None if indeterminate
+            message: Optional progress message
+        """
+        ...
+
+
 class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
     """Handles responding to MCP requests and manages request lifecycle.
 
@@ -169,6 +186,7 @@ class BaseSession(
     ]
     _request_id: int
     _in_flight: dict[RequestId, RequestResponder[ReceiveRequestT, SendResultT]]
+    _progress_callbacks: dict[RequestId, ProgressCallbackFnT]
 
     def __init__(
         self,
@@ -187,6 +205,7 @@ def __init__(
         self._receive_notification_type = receive_notification_type
         self._session_read_timeout_seconds = read_timeout_seconds
         self._in_flight = {}
+        self._progress_callbacks = {}
         self._exit_stack = AsyncExitStack()
 
     async def __aenter__(self) -> Self:
@@ -214,6 +233,7 @@ async def send_request(
         result_type: type[ReceiveResultT],
         request_read_timeout_seconds: timedelta | None = None,
         metadata: MessageMetadata = None,
+        progress_callback: ProgressCallbackFnT | None = None,
     ) -> ReceiveResultT:
         """
         Sends a request and wait for a response. Raises an McpError if the
@@ -231,15 +251,25 @@ async def send_request(
         ](1)
         self._response_streams[request_id] = response_stream
 
+        # 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._progress_callbacks[request_id] = progress_callback
+
         try:
             jsonrpc_request = JSONRPCRequest(
                 jsonrpc="2.0",
                 id=request_id,
-                **request.model_dump(by_alias=True, mode="json", exclude_none=True),
+                **request_data,
             )
 
-            # TODO: Support progress callbacks
-
             await self._write_stream.send(
                 SessionMessage(
                     message=JSONRPCMessage(jsonrpc_request), metadata=metadata
@@ -275,6 +305,7 @@ async def send_request(
 
         finally:
             self._response_streams.pop(request_id, None)
+            self._progress_callbacks.pop(request_id, None)
             await response_stream.aclose()
             await response_stream_reader.aclose()
 
@@ -333,7 +364,6 @@ async def _receive_loop(self) -> None:
                             by_alias=True, mode="json", exclude_none=True
                         )
                     )
-
                     responder = RequestResponder(
                         request_id=message.message.root.id,
                         request_meta=validated_request.root.params.meta
@@ -362,6 +392,18 @@ async def _receive_loop(self) -> None:
                             cancelled_id = notification.root.params.requestId
                             if cancelled_id in self._in_flight:
                                 await self._in_flight[cancelled_id].cancel()
+                        # Handle progress notifications
+                        elif isinstance(notification.root, ProgressNotification):
+                            progress_token = notification.root.params.progressToken
+                            # If there is a progress callback for this token,
+                            # call it with the progress information
+                            if progress_token in self._progress_callbacks:
+                                callback = self._progress_callbacks[progress_token]
+                                callback(
+                                    notification.root.params.progress,
+                                    notification.root.params.total,
+                                    notification.root.params.message,
+                                )
                         else:
                             await self._received_notification(notification)
                             await self._handle_incoming(notification)

tests/server/fastmcp/test_integration.py

@@ -13,14 +13,23 @@
 import anyio
 import pytest
 import uvicorn
+from pydantic import AnyUrl
 
+import mcp.types as types
 from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.server.fastmcp import FastMCP
-import mcp.types as types
-from mcp.types import InitializeResult, TextContent, TextResourceContents
-from pydantic import AnyUrl
+from mcp.shared.context import RequestContext
+from mcp.types import (
+    CreateMessageRequestParams,
+    CreateMessageResult,
+    GetPromptResult,
+    InitializeResult,
+    ReadResourceResult,
+    TextContent,
+    TextResourceContents,
+)
 
 
 @pytest.fixture
@@ -625,15 +634,6 @@ async def test_fastmcp_all_features_sse(
     comprehensive_server: None, comprehensive_server_url: str
 ) -> None:
     """Test all MCP features work correctly with SSE transport."""
-    from mcp.types import (
-        GetPromptResult,
-        ReadResourceResult,
-        CreateMessageResult,
-        CreateMessageRequestParams,
-        SamplingMessage,
-        TextContent,
-    )
-    from mcp.shared.context import RequestContext
 
     # Create notification collector
     collector = NotificationCollector()
@@ -690,20 +690,38 @@ async def message_handler(message):
             assert tool_result.content[0].text == "Echo: hello"
 
             # 2. Tool with context (logging and progress)
-            # Test with progress token to capture progress notifications
+            # Test progress callback functionality
+            progress_updates = []
+
+            def progress_callback(
+                progress: float, total: float | None, message: str | None
+            ) -> None:
+                """Collect progress updates for testing."""
+                progress_updates.append((progress, total, message))
+                print(f"Progress: {progress}/{total} - {message}")
+
+            params = {
+                "message": "test",
+                "steps": 3,
+            }
             tool_result = await session.call_tool(
                 "tool_with_context",
-                {
-                    "message": "test",
-                    "steps": 3,
-                    "_meta": {"progressToken": "sse_test_token"},
-                },
+                params,
+                progress_callback=progress_callback,
             )
             assert len(tool_result.content) == 1
-            assert len(collector.progress_notifications) > 0
             assert isinstance(tool_result.content[0], TextContent)
             assert "Processed 'test' in 3 steps" in tool_result.content[0].text
 
+            # Verify progress callback was called
+            assert len(progress_updates) == 3
+            for i, (progress, total, message) in enumerate(progress_updates):
+                expected_progress = (i + 1) / 3
+                assert abs(progress - expected_progress) < 0.01
+                assert total == 1.0
+                assert message is not None
+                assert f"step {i + 1} of 3" in message
+
             # Verify we received log messages from the tool
             # Note: Progress notifications require special handling in the MCP client
             # that's not implemented by default, so we focus on testing logging
@@ -832,15 +850,6 @@ async def test_fastmcp_all_features_streamable_http(
     comprehensive_streamable_http_server: None, comprehensive_http_server_url: str
 ) -> None:
     """Test all MCP features work correctly with StreamableHTTP transport."""
-    from mcp.types import (
-        GetPromptResult,
-        ReadResourceResult,
-        CreateMessageResult,
-        CreateMessageRequestParams,
-        SamplingMessage,
-        TextContent,
-    )
-    from mcp.shared.context import RequestContext
 
     # Create notification collector
     collector = NotificationCollector()
@@ -902,21 +911,33 @@ async def message_handler(message):
             assert tool_result.content[0].text == "Echo: hello"
 
             # 2. Tool with context (logging and progress)
-            # Test with progress token to capture progress notifications
+            # Test progress callback functionality over HTTP
+            progress_updates_http = []
+
+            def progress_callback_http(
+                progress: float, total: float | None, message: str | None
+            ) -> None:
+                """Collect progress updates for HTTP testing."""
+                progress_updates_http.append((progress, total, message))
+                print(f"HTTP Progress: {progress}/{total} - {message}")
+
             await session.call_tool(
                 "tool_with_context",
                 {
                     "message": "http_test",
                     "steps": 2,
                 },
+                progress_callback=progress_callback_http,
             )
 
-            # Verify we received progress notifications
-            assert len(collector.progress_notifications) > 0
-            assert any(
-                p.progressToken == "http_test_token"
-                for p in collector.progress_notifications
-            )
+            # Verify progress callback was called over HTTP
+            assert len(progress_updates_http) == 2
+            for i, (progress, total, message) in enumerate(progress_updates_http):
+                expected_progress = (i + 1) / 2
+                assert abs(progress - expected_progress) < 0.01
+                assert total == 1.0
+                assert message is not None
+                assert f"step {i + 1} of 2" in message
 
             # 3. Test sampling tool
             sampling_result = await session.call_tool(