From cc696c771066c7e76eeed3030b5b978eb439ec11 Mon Sep 17 00:00:00 2001 From: cyccolin Date: Thu, 8 May 2025 13:56:13 -0700 Subject: [PATCH 1/2] add support for fastMcp server to send progress notification with partialResults --- src/mcp/client/session.py | 3 +- src/mcp/server/fastmcp/server.py | 14 ++++++-- src/mcp/server/session.py | 2 ++ src/mcp/types.py | 25 ++++++++++++- tests/issues/test_176_progress_token.py | 47 +++++++++++++++++++++++-- 5 files changed, 84 insertions(+), 7 deletions(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 7bb8821f..ec88edbf 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -259,6 +259,7 @@ async def call_tool( name: str, arguments: dict[str, Any] | None = None, read_timeout_seconds: timedelta | None = None, + **meta ) -> types.CallToolResult: """Send a tools/call request.""" @@ -266,7 +267,7 @@ async def call_tool( types.ClientRequest( types.CallToolRequest( method="tools/call", - params=types.CallToolRequestParams(name=name, arguments=arguments), + params=types.CallToolRequestParams(name=name, arguments=arguments, _meta=meta), ) ), types.CallToolResult, diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index c31f29d4..10c04354 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -57,6 +57,7 @@ ImageContent, TextContent, ToolAnnotations, + PartialResult, ) from mcp.types import Prompt as MCPPrompt from mcp.types import PromptArgument as MCPPromptArgument @@ -952,13 +953,14 @@ def request_context(self) -> RequestContext[ServerSessionT, LifespanContextT]: return self._request_context async def report_progress( - self, progress: float, total: float | None = None + self, progress: float, total: float | None = None, partial_result: PartialResult | None = None ) -> None: """Report progress for the current operation. Args: progress: Current progress value e.g. 24 total: Optional total value e.g. 100 + partial_result: Optional partial result to include with the progress notification """ progress_token = ( @@ -967,11 +969,19 @@ async def report_progress( else None ) + partial_result = ( + partial_result + if self.request_context.meta and self.request_context.meta.partialResults + else None + ) + if progress_token is None: return + + print(f"partial_result={partial_result}") await self.request_context.session.send_progress_notification( - progress_token=progress_token, progress=progress, total=total + progress_token=progress_token, progress=progress, total=total, partial_result=partial_result ) async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]: diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index c769d1aa..59489ed0 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -279,6 +279,7 @@ async def send_progress_notification( progress: float, total: float | None = None, related_request_id: str | None = None, + partial_result: types.PartialResult | None = None, ) -> None: """Send a progress notification.""" await self.send_notification( @@ -289,6 +290,7 @@ async def send_progress_notification( progressToken=progress_token, progress=progress, total=total, + partialResult=partial_result, ), ) ), diff --git a/src/mcp/types.py b/src/mcp/types.py index 6ab7fba5..e6029ffb 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -47,7 +47,12 @@ class Meta(BaseModel): parameter is an opaque token that will be attached to any subsequent notifications. The receiver is not obligated to provide these notifications. """ - + partialResults: bool | None = None + """ + If true, the caller is requesting that results be streamed via progress notifications. + When this is set to true, the final response may be empty as the complete result + will have been delivered through progress notifications. + """ model_config = ConfigDict(extra="allow") meta: Meta | None = Field(alias="_meta", default=None) @@ -322,6 +327,19 @@ class PingRequest(Request[RequestParams | None, Literal["ping"]]): method: Literal["ping"] params: RequestParams | None = None +class PartialResult(BaseModel): + chunk: dict[str, Any] + """The partial result data chunk.""" + append: bool = False + """ + If true, this chunk should be appended to previously received chunks. + If false, this chunk replaces any previously received chunks. + """ + lastChunk: bool = False + """ + If true, this is the final chunk of the result. + No further chunks will be sent for this operation. + """ class ProgressNotificationParams(NotificationParams): """Parameters for progress notifications.""" @@ -338,6 +356,11 @@ class ProgressNotificationParams(NotificationParams): """ total: float | None = None """Total number of items to process (or total progress required), if known.""" + partialResult: PartialResult | None = None + """ + If present, contains a partial result chunk for the operation. + This is used to stream results incrementally while an operation is still in progress. + """ model_config = ConfigDict(extra="allow") diff --git a/tests/issues/test_176_progress_token.py b/tests/issues/test_176_progress_token.py index 7f9131a1..0eb9ad99 100644 --- a/tests/issues/test_176_progress_token.py +++ b/tests/issues/test_176_progress_token.py @@ -4,6 +4,7 @@ from mcp.server.fastmcp import Context from mcp.shared.context import RequestContext +from mcp.types import PartialResult pytestmark = pytest.mark.anyio @@ -39,11 +40,51 @@ async def test_progress_token_zero_first_call(): mock_session.send_progress_notification.call_count == 3 ), "All progress notifications should be sent" mock_session.send_progress_notification.assert_any_call( - progress_token=0, progress=0.0, total=10.0 + progress_token=0, progress=0.0, total=10.0, partial_result=None ) mock_session.send_progress_notification.assert_any_call( - progress_token=0, progress=5.0, total=10.0 + progress_token=0, progress=5.0, total=10.0, partial_result=None ) mock_session.send_progress_notification.assert_any_call( - progress_token=0, progress=10.0, total=10.0 + progress_token=0, progress=10.0, total=10.0, partial_result=None ) + +async def test_progress_token_with_partial_results(): + """Test that progress notifications work when progress_token is 0 on first call.""" + + # Create mock session with progress notification tracking + mock_session = AsyncMock() + mock_session.send_progress_notification = AsyncMock() + + # Create request context with progress token and partialResults as True + mock_meta = MagicMock() + mock_meta.progressToken = 0 + mock_meta.partialResults = True + request_context = RequestContext( + request_id="test-request", + session=mock_session, + meta=mock_meta, + lifespan_context=None, + ) + + # Create context with our mocks + ctx = Context(request_context=request_context, fastmcp=MagicMock()) + + # Test progress reporting + await ctx.report_progress(0, 10, PartialResult(chunk={"content": [{"type": "text", "text": "TestData1"}]}, append=False, lastChunk=False)) + await ctx.report_progress(5, 10) + await ctx.report_progress(10, 10, PartialResult(chunk={"content": [{"type": "text", "text": "TestData3"}]}, append=True, lastChunk=True)) + + # Verify progress notifications + assert ( + mock_session.send_progress_notification.call_count == 3 + ), "All progress notifications should be sent" + mock_session.send_progress_notification.assert_any_call( + progress_token=0, progress=0.0, total=10.0, partial_result=PartialResult(chunk={"content": [{"type": "text", "text": "TestData1"}]}, append=False, lastChunk=False) + ) + mock_session.send_progress_notification.assert_any_call( + progress_token=0, progress=5.0, total=10.0, partial_result=None + ) + mock_session.send_progress_notification.assert_any_call( + progress_token=0, progress=10.0, total=10.0, partial_result=PartialResult(chunk={"content": [{"type": "text", "text": "TestData3"}]}, append=True, lastChunk=True) + ) \ No newline at end of file From e26c2bef2bec6cb02804d1a48d515fe8d7bef911 Mon Sep 17 00:00:00 2001 From: cyccolin Date: Thu, 8 May 2025 14:57:15 -0700 Subject: [PATCH 2/2] code clean up --- src/mcp/server/fastmcp/server.py | 2 -- tests/issues/test_176_progress_token.py | 10 +++++----- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 10c04354..8b5165a9 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -977,8 +977,6 @@ async def report_progress( if progress_token is None: return - - print(f"partial_result={partial_result}") await self.request_context.session.send_progress_notification( progress_token=progress_token, progress=progress, total=total, partial_result=partial_result diff --git a/tests/issues/test_176_progress_token.py b/tests/issues/test_176_progress_token.py index 0eb9ad99..f91bcbe0 100644 --- a/tests/issues/test_176_progress_token.py +++ b/tests/issues/test_176_progress_token.py @@ -50,7 +50,7 @@ async def test_progress_token_zero_first_call(): ) async def test_progress_token_with_partial_results(): - """Test that progress notifications work when progress_token is 0 on first call.""" + """Test that progress notifications with partial results set to true""" # Create mock session with progress notification tracking mock_session = AsyncMock() @@ -58,7 +58,7 @@ async def test_progress_token_with_partial_results(): # Create request context with progress token and partialResults as True mock_meta = MagicMock() - mock_meta.progressToken = 0 + mock_meta.progressToken = "progress-token" mock_meta.partialResults = True request_context = RequestContext( request_id="test-request", @@ -80,11 +80,11 @@ async def test_progress_token_with_partial_results(): mock_session.send_progress_notification.call_count == 3 ), "All progress notifications should be sent" mock_session.send_progress_notification.assert_any_call( - progress_token=0, progress=0.0, total=10.0, partial_result=PartialResult(chunk={"content": [{"type": "text", "text": "TestData1"}]}, append=False, lastChunk=False) + progress_token="progress-token", progress=0.0, total=10.0, partial_result=PartialResult(chunk={"content": [{"type": "text", "text": "TestData1"}]}, append=False, lastChunk=False) ) mock_session.send_progress_notification.assert_any_call( - progress_token=0, progress=5.0, total=10.0, partial_result=None + progress_token="progress-token", progress=5.0, total=10.0, partial_result=None ) mock_session.send_progress_notification.assert_any_call( - progress_token=0, progress=10.0, total=10.0, partial_result=PartialResult(chunk={"content": [{"type": "text", "text": "TestData3"}]}, append=True, lastChunk=True) + progress_token="progress-token", progress=10.0, total=10.0, partial_result=PartialResult(chunk={"content": [{"type": "text", "text": "TestData3"}]}, append=True, lastChunk=True) ) \ No newline at end of file