Handle cancellation notifications on async ops

Author: LucaButBoring

src/mcp/server/lowlevel/server.py

@@ -90,6 +90,7 @@ async def main():
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.session import RequestResponder
+from mcp.types import RequestId
 
 logger = logging.getLogger(__name__)
 
@@ -147,10 +148,14 @@ def __init__(
         self.instructions = instructions
         self.lifespan = lifespan
         self.async_operations = async_operations or AsyncOperationManager()
+        # Track request ID to operation token mapping for cancellation
+        self._request_to_operation: dict[RequestId, str] = {}
         self.request_handlers: dict[type, Callable[..., Awaitable[types.ServerResult]]] = {
             types.PingRequest: _ping_handler,
         }
-        self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
+        self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {
+            types.CancelledNotification: self._handle_cancelled_notification,
+        }
         self._tool_cache: dict[str, types.Tool] = {}
         logger.debug("Initializing server %r", name)
 
@@ -566,6 +571,10 @@ def _validate_operation_token(self, token: str) -> AsyncOperation:
         if operation.is_expired:
             raise McpError(types.ErrorData(code=-32602, message="Token expired"))
 
+        # Check if operation was cancelled - ignore subsequent requests
+        if operation.status == "canceled":
+            raise McpError(types.ErrorData(code=-32602, message="Operation was cancelled"))
+
         return operation
 
     def get_operation_status(self):
@@ -615,6 +624,23 @@ async def handler(req: types.GetOperationPayloadRequest):
 
         return decorator
 
+    def handle_cancelled_notification(self, request_id: RequestId) -> None:
+        """Handle cancellation notification for a request."""
+        # Check if this request ID corresponds to an async operation
+        if request_id in self._request_to_operation:
+            token = self._request_to_operation[request_id]
+            # Cancel the operation
+            if self.async_operations.cancel_operation(token):
+                logger.debug(f"Cancelled async operation {token} for request {request_id}")
+            # Clean up the mapping
+            del self._request_to_operation[request_id]
+
+    async def _handle_cancelled_notification(self, notification: types.CancelledNotification) -> None:
+        """Handle cancelled notification from client."""
+        request_id = notification.params.requestId
+        logger.debug(f"Received cancellation notification for request {request_id}")
+        self.handle_cancelled_notification(request_id)
+
     async def run(
         self,
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -695,7 +721,7 @@ async def _handle_request(
         if handler := self.request_handlers.get(type(req)):  # type: ignore
             logger.debug("Dispatching request of type %s", type(req).__name__)
 
-            token = None
+            context_token = None
             try:
                 # Extract request context from message metadata
                 request_data = None
@@ -704,7 +730,7 @@ async def _handle_request(
 
                 # Set our global state that can be retrieved via
                 # app.get_request_context()
-                token = request_ctx.set(
+                context_token = request_ctx.set(
                     RequestContext(
                         message.request_id,
                         message.request_meta,
@@ -714,6 +740,16 @@ async def _handle_request(
                     )
                 )
                 response = await handler(req)
+
+                # Track async operations for cancellation
+                if isinstance(req, types.CallToolRequest):
+                    result = response.root
+                    if isinstance(result, types.CallToolResult) and result.operation_result is not None:
+                        # This is an async operation, track the request ID to token mapping
+                        operation_token = result.operation_result.token
+                        self._request_to_operation[message.request_id] = operation_token
+                        logger.debug(f"Tracking async operation {operation_token} for request {message.request_id}")
+
             except McpError as err:
                 response = err.error
             except anyio.get_cancelled_exc_class():
@@ -728,8 +764,8 @@ async def _handle_request(
                 response = types.ErrorData(code=0, message=str(err), data=None)
             finally:
                 # Reset the global state after we are done
-                if token is not None:
-                    request_ctx.reset(token)
+                if context_token is not None:
+                    request_ctx.reset(context_token)
 
             await message.respond(response)
         else:

tests/server/test_cancellation_logic.py

@@ -0,0 +1,86 @@
+"""Tests for async operation cancellation logic."""
+
+import pytest
+
+import mcp.types as types
+from mcp.server.lowlevel.async_operations import AsyncOperationManager
+from mcp.server.lowlevel.server import Server
+from mcp.shared.exceptions import McpError
+
+
+class TestCancellationLogic:
+    """Test cancellation logic for async operations."""
+
+    def test_handle_cancelled_notification(self):
+        """Test handling of cancelled notifications."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create an operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+
+        # Track the operation with a request ID
+        request_id = "req_123"
+        server._request_to_operation[request_id] = operation.token
+
+        # Handle cancellation
+        server.handle_cancelled_notification(request_id)
+
+        # Verify operation was cancelled
+        cancelled_op = manager.get_operation(operation.token)
+        assert cancelled_op is not None
+        assert cancelled_op.status == "canceled"
+
+        # Verify mapping was cleaned up
+        assert request_id not in server._request_to_operation
+
+    def test_cancelled_notification_handler(self):
+        """Test the async cancelled notification handler."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create an operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+
+        # Track the operation with a request ID
+        request_id = "req_456"
+        server._request_to_operation[request_id] = operation.token
+
+        # Create cancelled notification
+        notification = types.CancelledNotification(params=types.CancelledNotificationParams(requestId=request_id))
+
+        # Handle the notification
+        import asyncio
+
+        asyncio.run(server._handle_cancelled_notification(notification))
+
+        # Verify operation was cancelled
+        cancelled_op = manager.get_operation(operation.token)
+        assert cancelled_op is not None
+        assert cancelled_op.status == "canceled"
+
+    def test_validate_operation_token_cancelled(self):
+        """Test that cancelled operations are rejected."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create and cancel an operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.cancel_operation(operation.token)
+
+        # Verify that accessing cancelled operation raises error
+        with pytest.raises(McpError) as exc_info:
+            server._validate_operation_token(operation.token)
+
+        assert exc_info.value.error.code == -32602
+        assert "cancelled" in exc_info.value.error.message.lower()
+
+    def test_nonexistent_request_id_cancellation(self):
+        """Test cancellation of non-existent request ID."""
+        server = Server("Test")
+
+        # Should not raise error for non-existent request ID
+        server.handle_cancelled_notification("nonexistent_request")
+
+        # Verify no operations were affected
+        assert len(server._request_to_operation) == 0