Fully switch AsyncOperationManager to anyio

Author: LucaButBoring

src/mcp/client/session.py

@@ -1,6 +1,6 @@
 import logging
 from datetime import timedelta
-from typing import Any, Protocol
+from typing import Any, Protocol, Self
 
 import anyio
 import anyio.lowlevel
@@ -139,6 +139,12 @@ def __init__(
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
         self._operation_manager = ClientAsyncOperationManager()
 
+    async def __aenter__(self) -> Self:
+        await super().__aenter__()
+        self._task_group.start_soon(self._operation_manager.cleanup_loop)
+        self._exit_stack.push_async_callback(lambda: self._operation_manager.stop_cleanup_loop())
+        return self
+
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
         elicitation = (
@@ -176,15 +182,8 @@ async def initialize(self) -> types.InitializeResult:
 
         await self.send_notification(types.ClientNotification(types.InitializedNotification()))
 
-        # Start cleanup task for operations
-        await self._operation_manager.start_cleanup_task()
-
         return result
 
-    async def close(self) -> None:
-        """Clean up resources."""
-        await self._operation_manager.stop_cleanup_task()
-
     async def send_ping(self) -> types.EmptyResult:
         """Send a ping request."""
         return await self.send_request(

src/mcp/server/lowlevel/server.py

@@ -865,11 +865,10 @@ async def run(
                 )
             )
 
-            # Start async operations cleanup task
-            await self.async_operations.start_cleanup_task()
+            async with anyio.create_task_group() as tg:
+                tg.start_soon(self.async_operations.cleanup_loop)
 
-            try:
-                async with anyio.create_task_group() as tg:
+                try:
                     async for message in session.incoming_messages:
                         logger.debug("Received message: %s", message)
 
@@ -881,9 +880,12 @@ async def run(
                             raise_exceptions,
                             tg,
                         )
-            finally:
-                # Stop cleanup task
-                await self.async_operations.stop_cleanup_task()
+                finally:
+                    # Stop cleanup loop before task group exits
+                    await self.async_operations.stop_cleanup_loop()
+
+                    # Cancel all remaining tasks in the task group (cleanup loop and potentially LROs)
+                    tg.cancel_scope.cancel()
 
     async def _handle_message(
         self,

src/mcp/shared/async_operations.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
-import asyncio
+import logging
 import secrets
 import time
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
+import anyio
+
 import mcp.types as types
 from mcp.types import AsyncOperationStatus
 
@@ -66,9 +68,9 @@ class BaseOperationManager(Generic[OperationT]):
 
     def __init__(self, *, token_generator: Callable[[str | None], str] | None = None):
         self._operations: dict[str, OperationT] = {}
-        self._cleanup_task: asyncio.Task[None] | None = None
         self._cleanup_interval = 60  # Cleanup every 60 seconds
         self._token_generator = token_generator or self._default_token_generator
+        self._running = False
 
     def _default_token_generator(self, session_id: str | None = None) -> str:
         """Default token generation using random tokens."""
@@ -105,31 +107,20 @@ def cleanup_expired(self) -> int:
             self._remove_operation(token)
         return len(expired_tokens)
 
-    async def start_cleanup_task(self) -> None:
-        """Start the background cleanup task."""
-        if self._cleanup_task is None:
-            self._cleanup_task = asyncio.create_task(self._cleanup_loop())
-
-    async def stop_cleanup_task(self) -> None:
-        """Stop the background cleanup task."""
-        if self._cleanup_task:
-            self._cleanup_task.cancel()
-            try:
-                await self._cleanup_task
-            except asyncio.CancelledError:
-                pass
-            self._cleanup_task = None
-
-    async def _cleanup_loop(self) -> None:
+    async def stop_cleanup_loop(self) -> None:
+        self._running = False
+
+    async def cleanup_loop(self) -> None:
         """Background task to clean up expired operations."""
-        while True:
-            try:
-                await asyncio.sleep(self._cleanup_interval)
-                count = self.cleanup_expired()
-                if count > 0:
-                    print(f"Cleaned up {count} expired operations")
-            except asyncio.CancelledError:
-                break
+        if self._running:
+            return
+        self._running = True
+
+        while self._running:
+            await anyio.sleep(self._cleanup_interval)
+            count = self.cleanup_expired()
+            if count > 0:
+                logging.debug(f"Cleaned up {count} expired operations")
 
 
 class ClientAsyncOperationManager(BaseOperationManager[ClientAsyncOperation]):
@@ -292,32 +283,3 @@ def mark_input_completed(self, token: str) -> bool:
 
         operation.status = "working"
         return True
-
-    async def start_cleanup_task(self) -> None:
-        """Start the background cleanup task."""
-        if self._cleanup_task is not None:
-            return
-
-        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
-
-    async def stop_cleanup_task(self) -> None:
-        """Stop the background cleanup task."""
-        if self._cleanup_task is not None:
-            self._cleanup_task.cancel()
-            try:
-                await self._cleanup_task
-            except asyncio.CancelledError:
-                pass
-            self._cleanup_task = None
-
-    async def _cleanup_loop(self) -> None:
-        """Background cleanup loop."""
-        while True:
-            try:
-                await asyncio.sleep(self._cleanup_interval)
-                self.cleanup_expired_operations()
-            except asyncio.CancelledError:
-                break
-            except Exception:
-                # Log error but continue cleanup loop
-                pass

tests/shared/test_async_operations.py

@@ -5,8 +5,6 @@
 from typing import Any, cast
 from unittest.mock import Mock
 
-import pytest
-
 import mcp.types as types
 from mcp.shared.async_operations import ServerAsyncOperation, ServerAsyncOperationManager
 from mcp.types import AsyncOperationStatus
@@ -213,50 +211,6 @@ def test_concurrent_operations(self):
         removed_count = manager.cleanup_expired_operations()
         assert removed_count == 25 and len(manager._operations) == 25
 
-    @pytest.mark.anyio
-    async def test_cleanup_task_lifecycle(self):
-        """Test background cleanup task management."""
-        manager = ServerAsyncOperationManager()
-
-        await manager.start_cleanup_task()
-        assert manager._cleanup_task is not None and not manager._cleanup_task.done()
-
-        # Starting again should be no-op
-        await manager.start_cleanup_task()
-
-        await manager.stop_cleanup_task()
-        assert manager._cleanup_task is None
-
-    def test_dependency_injection_and_integration(self):
-        """Test AsyncOperationManager dependency injection and server integration."""
-        from mcp.server.fastmcp import FastMCP
-        from mcp.server.lowlevel import Server
-
-        # Test custom manager injection
-        custom_manager = ServerAsyncOperationManager()
-        operation = custom_manager.create_operation("shared_tool", {"data": "shared"}, session_id="session1")
-
-        # Test FastMCP integration
-        fastmcp = FastMCP("FastMCP", async_operations=custom_manager)
-        assert fastmcp._async_operations is custom_manager
-        assert fastmcp._async_operations.get_operation(operation.token) is operation
-
-        # Test lowlevel Server integration
-        lowlevel = Server("LowLevel", async_operations=custom_manager)
-        assert lowlevel.async_operations is custom_manager
-        assert lowlevel.async_operations.get_operation(operation.token) is operation
-
-        # Test default creation
-        default_fastmcp = FastMCP("Default")
-        default_server = Server("Default")
-        assert isinstance(default_fastmcp._async_operations, ServerAsyncOperationManager)
-        assert isinstance(default_server.async_operations, ServerAsyncOperationManager)
-        assert default_fastmcp._async_operations is not custom_manager
-
-        # Test shared manager between servers
-        new_op = fastmcp._async_operations.create_operation("new_tool", {}, session_id="session2")
-        assert lowlevel.async_operations.get_operation(new_op.token) is new_op
-
 
 class TestAsyncOperation:
     """Test AsyncOperation dataclass."""