Further updates, create lowlevel test for result caching and improve test logic for higher level session join test (currently skipped due to subtle bug in test)

davemssavage · davemssavage · commit bef1d72ae9e7 · 2025-06-03T05:36:08.000Z
diff --git a/src/mcp/server/lowlevel/result_cache.py b/src/mcp/server/lowlevel/result_cache.py
@@ -26,7 +26,7 @@ class InProgress:
     user: AuthenticatedUser | None = None
     future: Future[types.CallToolResult] | None = None
     sessions: dict[int, ServerSession] = field(default_factory=lambda: {})
-
+    session_progress: dict[int, types.ProgressToken | None] = field(default_factory=lambda: {})
 
 class ResultCache:
     """
@@ -72,7 +72,7 @@ async def __aexit__(
     ) -> bool | None:
         await anyio.to_thread.run_sync(lambda: self._portal_provider.__exit__)
 
-    async def add_call(
+    async def start_call(
         self,
         call: Callable[[types.CallToolRequest], Awaitable[types.ServerResult]],
         req: types.CallToolAsyncRequest,
@@ -101,6 +101,7 @@ async def call_tool():
 
         in_progress.user = user_context.get()
         in_progress.sessions[id(ctx.session)] = ctx.session
+        in_progress.session_progress[id(ctx.session)] = None if req.params.meta is None else req.params.meta.progressToken
         self._session_lookup[id(ctx.session)] = in_progress.token
         in_progress.future = self._portal.start_task_soon(call_tool)
         result = types.CallToolAsyncResult(
@@ -129,6 +130,7 @@ async def join_call(
                 logger.debug(f"Received join from {id(ctx.session)}")
                 self._session_lookup[id(ctx.session)] = req.params.token
                 in_progress.sessions[id(ctx.session)] = ctx.session
+                in_progress.session_progress[id(ctx.session)] = None if req.params.meta is None else req.params.meta.progressToken
                 return types.CallToolAsyncResult(token=req.params.token, accepted=True)
             else:
                 # TODO consider sending error via get result
@@ -167,10 +169,15 @@ async def get_result(self, req: types.GetToolAsyncResultRequest):
                     )
                 else:
                     # TODO add timeout to get async result
-                    # return isPending=True if timesout
-                    result = in_progress.future.result()
-                    logger.debug(f"Found result {result}")
-                    return result
+                    try:
+                        result = in_progress.future.result(1)
+                        logger.debug(f"Found result {result}")
+                        return result
+                    except TimeoutError:
+                        return types.CallToolResult(
+                            content=[],
+                            isPending=True,
+                        )
             else:
                 return types.CallToolResult(
                     content=[types.TextContent(type="text", text="Permission denied")],
@@ -180,6 +187,7 @@ async def get_result(self, req: types.GetToolAsyncResultRequest):
     async def notification_hook(
         self, session: ServerSession, notification: types.ServerNotification
     ):
+        logger.debug(f"received {notification} from {id(session)}")
         if type(notification.root) is types.ProgressNotification:
             # async with self._lock:
             async_token = self._session_lookup.get(id(session))
@@ -196,8 +204,12 @@ async def notification_hook(
                         logger.debug(f"Checking {session_id} == {id(session)}")
                         if not session_id == id(session):
                             logger.debug(f"Sending progress to {id(other_session)}")
+                            progress_token = in_progress.session_progress.get(id(other_session))
+                            assert progress_token is not None
                             await other_session.send_progress_notification(
-                                progress_token=1,
+                                # TODO this token is incorrect 
+                                # it needs to be collected from original request
+                                progress_token=progress_token,
                                 progress=notification.root.params.progress,
                                 total=notification.root.params.total,
                                 message=notification.root.params.message,
diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
@@ -432,7 +432,7 @@ async def handler(req: types.CallToolRequest):
 
             async def async_call_handler(req: types.CallToolAsyncRequest):
                 ctx = request_ctx.get()
-                result = await self.result_cache.add_call(handler, req, ctx)
+                result = await self.result_cache.start_call(handler, req, ctx)
                 return types.ServerResult(result)
 
             async def async_join_handler(req: types.JoinCallToolAsyncRequest):
diff --git a/tests/server/lowlevel/test_result_cache.py b/tests/server/lowlevel/test_result_cache.py
@@ -0,0 +1,121 @@
+import pytest
+from mcp import types
+from mcp.server.lowlevel.result_cache import ResultCache
+from unittest.mock import AsyncMock, Mock, patch
+from contextlib import AsyncExitStack
+
+@pytest.mark.anyio
+async def test_async_call():
+    """Tests basic async call"""
+    async def test_call(call: types.CallToolRequest) -> types.ServerResult:
+        return types.ServerResult(types.CallToolResult(
+            content=[types.TextContent(
+                type="text",
+                text="test"
+            )]
+        ))
+    async_call = types.CallToolAsyncRequest(
+        method="tools/async/call",
+        params=types.CallToolAsyncRequestParams(
+            name="test"
+        )
+    )
+
+    mock_session = AsyncMock()
+    mock_context = Mock()
+    mock_context.session = mock_session
+    result_cache = ResultCache(max_size=1, max_keep_alive=1)
+    async with AsyncExitStack() as stack:
+        await stack.enter_async_context(result_cache)
+        async_call_ref = await result_cache.start_call(test_call, async_call, mock_context)
+        assert async_call_ref.token is not None
+
+        result = await result_cache.get_result(types.GetToolAsyncResultRequest(
+            method="tools/async/get",
+            params=types.GetToolAsyncResultRequestParams(
+                token = async_call_ref.token
+            )
+        ))
+
+        assert not result.isError
+        assert not result.isPending
+        assert len(result.content) == 1
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "test"
+
+@pytest.mark.anyio
+async def test_async_join_call_progress():
+    """Tests basic async call"""
+    async def test_call(call: types.CallToolRequest) -> types.ServerResult:
+        return types.ServerResult(types.CallToolResult(
+            content=[types.TextContent(
+                type="text",
+                text="test"
+            )]
+        ))
+    async_call = types.CallToolAsyncRequest(
+        method="tools/async/call",
+        params=types.CallToolAsyncRequestParams(
+            name="test"
+        )
+    )
+
+    mock_session_1 = AsyncMock()
+    mock_context_1 = Mock()
+    mock_context_1.session = mock_session_1
+
+    mock_session_2 = AsyncMock()
+    mock_context_2 = Mock()
+
+    mock_context_2.session = mock_session_2
+    mock_session_2.send_progress_notification.result = None
+
+    result_cache = ResultCache(max_size=1, max_keep_alive=1)
+    async with AsyncExitStack() as stack:
+        await stack.enter_async_context(result_cache)
+        async_call_ref = await result_cache.start_call(test_call, async_call, mock_context_1)
+        assert async_call_ref.token is not None
+
+        await result_cache.join_call(
+            req=types.JoinCallToolAsyncRequest(
+                method="tools/async/join",
+                params=types.JoinCallToolRequestParams(
+                    token=async_call_ref.token,
+                    _meta = types.RequestParams.Meta(
+                        progressToken="test"
+                    )
+                )
+            ),
+            ctx=mock_context_2
+        )
+        assert async_call_ref.token is not None
+        await result_cache.notification_hook(
+            session=mock_session_1, 
+            notification=types.ServerNotification(types.ProgressNotification(
+                method="notifications/progress",
+                params=types.ProgressNotificationParams(
+                    progressToken="test",
+                    progress=1
+                )
+            )))
+
+        result = await result_cache.get_result(types.GetToolAsyncResultRequest(
+            method="tools/async/get",
+            params=types.GetToolAsyncResultRequestParams(
+                token = async_call_ref.token
+            )
+        ))
+
+        assert not result.isError
+        assert not result.isPending
+        assert len(result.content) == 1
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "test"
+        mock_context_1.send_progress_notification.assert_not_called()
+        mock_session_2.send_progress_notification.assert_called_with(
+            progress_token="test", 
+            progress=1.0, 
+            total=None, 
+            message=None, 
+            resource_uri = None
+        )
diff --git a/tests/shared/test_session.py b/tests/shared/test_session.py
