Add tests for cancellation behaviour and implement

Author: davemssavage

src/mcp/server/lowlevel/result_cache.py

@@ -1,6 +1,6 @@
 import time
 from collections.abc import Awaitable, Callable
-from concurrent.futures import Future
+from concurrent.futures import CancelledError, Future
 from dataclasses import dataclass, field
 from logging import getLogger
 from types import TracebackType
@@ -168,7 +168,8 @@ async def cancel(self, notification: types.CancelToolAsyncNotification) -> None:
         if in_progress is not None:
             if in_progress.user == user_context.get():
                 # in_progress.task_group.cancel_scope.cancel()
-                del self._in_progress[notification.params.token]
+                assert in_progress.future is not None, "In progress future not found"
+                in_progress.future.cancel()
             else:
                 logger.warning(
                     "Permission denied for cancel notification received"
@@ -202,6 +203,11 @@ async def get_result(self, req: types.GetToolAsyncResultRequest):
                     result = in_progress.future.result(1)
                     logger.debug(f"Found result {result}")
                     return result
+                except CancelledError:
+                    return types.CallToolResult(
+                        content=[types.TextContent(type="text", text="cancelled")],
+                        isError=True,
+                    )
                 except TimeoutError:
                     return types.CallToolResult(
                         content=[],

tests/server/lowlevel/test_result_cache.py

@@ -1,6 +1,7 @@
 from contextlib import AsyncExitStack
 from unittest.mock import AsyncMock, Mock
 
+import anyio
 import pytest
 
 from mcp import types
@@ -124,6 +125,84 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
         )
 
 
+@pytest.mark.anyio
+async def test_async_cancel_in_progress():
+    """Tests basic async call"""
+
+    async def slow_call(call: types.CallToolRequest) -> types.ServerResult:
+        with anyio.move_on_after(10) as scope:
+            await anyio.sleep(10)
+
+        if scope.cancel_called:
+            return types.ServerResult(
+                types.CallToolResult(
+                    content=[
+                        types.TextContent(type="text", text="should be discarded")
+                    ],
+                    isError=True,
+                )
+            )
+        else:
+            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
+
+    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(
+            slow_call, async_call, mock_context_1
+        )
+        assert async_call_ref.token is not None
+
+        await result_cache.cancel(
+            notification=types.CancelToolAsyncNotification(
+                method="tools/async/cancel",
+                params=types.CancelToolAsyncNotificationParams(
+                    token=async_call_ref.token
+                ),
+            ),
+        )
+
+        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 result.isError
+        assert not result.isPending
+        assert len(result.content) == 1
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "cancelled"
+
+
 @pytest.mark.anyio
 async def test_async_call_keep_alive():
     """Tests async call keep alive"""