Convert ServerAsyncOperationManager into async context manager

Author: LucaButBoring

src/mcp/server/fastmcp/server.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations as _annotations
 
+import contextlib
 import inspect
 import re
 from collections.abc import AsyncIterator, Awaitable, Callable, Collection, Iterable, Sequence
@@ -845,14 +846,21 @@ def decorator(
 
         return decorator
 
+    @contextlib.asynccontextmanager
+    async def _stdio_lifespan(self) -> AsyncIterator[None]:
+        """Lifespan that manages stdio operations."""
+        async with self._async_operations.run():
+            yield
+
     async def run_stdio_async(self) -> None:
         """Run the server using stdio transport."""
         async with stdio_server() as (read_stream, write_stream):
-            await self._mcp_server.run(
-                read_stream,
-                write_stream,
-                self._mcp_server.create_initialization_options(),
-            )
+            async with self._stdio_lifespan():
+                await self._mcp_server.run(
+                    read_stream,
+                    write_stream,
+                    self._mcp_server.create_initialization_options(),
+                )
 
     async def run_sse_async(self, mount_path: str | None = None) -> None:
         """Run the server using SSE transport."""
@@ -910,6 +918,12 @@ def _normalize_path(self, mount_path: str, endpoint: str) -> str:
         # Combine paths
         return mount_path + endpoint
 
+    @contextlib.asynccontextmanager
+    async def _sse_lifespan(self) -> AsyncIterator[None]:
+        """Lifespan that manages SSE operations."""
+        async with self._async_operations.run():
+            yield
+
     def sse_app(self, mount_path: str | None = None) -> Starlette:
         """Return an instance of the SSE server app."""
         from starlette.middleware import Middleware
@@ -1040,7 +1054,16 @@ async def sse_endpoint(request: Request) -> Response:
         routes.extend(self._custom_starlette_routes)
 
         # Create Starlette app with routes and middleware
-        return Starlette(debug=self.settings.debug, routes=routes, middleware=middleware)
+        return Starlette(
+            debug=self.settings.debug, routes=routes, middleware=middleware, lifespan=lambda app: self._sse_lifespan()
+        )
+
+    @contextlib.asynccontextmanager
+    async def _streamable_http_lifespan(self) -> AsyncIterator[None]:
+        """Lifespan that manages Streamable HTTP operations."""
+        async with self.session_manager.run():
+            async with self._async_operations.run():
+                yield
 
     def streamable_http_app(self) -> Starlette:
         """Return an instance of the StreamableHTTP server app."""
@@ -1135,7 +1158,7 @@ def streamable_http_app(self) -> Starlette:
             debug=self.settings.debug,
             routes=routes,
             middleware=middleware,
-            lifespan=lambda app: self.session_manager.run(),
+            lifespan=lambda app: self._streamable_http_lifespan(),
         )
 
     async def list_prompts(self) -> list[MCPPrompt]:
@@ -1337,12 +1360,17 @@ async def log(
             logger_name: Optional logger name
             **extra: Additional structured data to include
         """
-        await self.request_context.session.send_log_message(
-            level=level,
-            data=message,
-            logger=logger_name,
-            related_request_id=self.request_id,
-        )
+        try:
+            await self.request_context.session.send_log_message(
+                level=level,
+                data=message,
+                logger=logger_name,
+                related_request_id=self.request_id,
+            )
+        except Exception:
+            # Session might be closed (e.g., client disconnected)
+            logger.warning(f"Failed to send log message to client (session closed?): {message}")
+            pass
 
     @property
     def client_id(self) -> str | None:

src/mcp/server/lowlevel/server.py

@@ -95,6 +95,7 @@ async def main():
 from mcp.types import Operation, RequestId
 
 logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
 
 LifespanResultT = TypeVar("LifespanResultT", default=Any)
 RequestT = TypeVar("RequestT", default=Any)
@@ -564,8 +565,11 @@ async def execute_async():
                                 logger.exception(f"Async execution failed for {tool_name}")
                                 self.async_operations.fail_operation(operation.token, str(e))
 
-                        # Dispatch in concurrency scope of the server to run between requests
-                        server_scope.start_soon(execute_async)
+                        # Start task directly in independent task group
+                        current_request_context = request_ctx.get()
+                        self.async_operations.start_task(
+                            operation.token, execute_async, current_request_context, request_ctx
+                        )
 
                         # Return operation result with immediate content
                         logger.info(f"Returning async operation result for {tool_name}")
@@ -866,26 +870,17 @@ async def run(
             )
 
             async with anyio.create_task_group() as tg:
-                tg.start_soon(self.async_operations.cleanup_loop)
-
-                try:
-                    async for message in session.incoming_messages:
-                        logger.debug("Received message: %s", message)
-
-                        tg.start_soon(
-                            self._handle_message,
-                            message,
-                            session,
-                            lifespan_context,
-                            raise_exceptions,
-                            tg,
-                        )
-                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 for message in session.incoming_messages:
+                    logger.debug("Received message: %s", message)
+
+                    tg.start_soon(
+                        self._handle_message,
+                        message,
+                        session,
+                        lifespan_context,
+                        raise_exceptions,
+                        tg,
+                    )
 
     async def _handle_message(
         self,

src/mcp/shared/async_operations.py

@@ -2,18 +2,22 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import secrets
 import time
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
 import anyio
+from anyio.abc import TaskGroup
 
 import mcp.types as types
 from mcp.types import AsyncOperationStatus
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class ClientAsyncOperation:
@@ -120,7 +124,7 @@ async def cleanup_loop(self) -> None:
             await anyio.sleep(self._cleanup_interval)
             count = self.cleanup_expired()
             if count > 0:
-                logging.debug(f"Cleaned up {count} expired operations")
+                logger.debug(f"Cleaned up {count} expired operations")
 
 
 class ClientAsyncOperationManager(BaseOperationManager[ClientAsyncOperation]):
@@ -145,6 +149,63 @@ def get_tool_name(self, token: str) -> str | None:
 class ServerAsyncOperationManager(BaseOperationManager[ServerAsyncOperation]):
     """Manages async tool operations with token-based tracking."""
 
+    def __init__(self, *, token_generator: Callable[[str | None], str] | None = None):
+        super().__init__(token_generator=token_generator)
+        self._task_group: TaskGroup | None = None
+        self._run_lock = anyio.Lock()
+        self._running = False
+
+    @contextlib.asynccontextmanager
+    async def run(self) -> AsyncIterator[None]:
+        """Run the async operations manager with its own task group."""
+        # Thread-safe check to ensure run() is only called once
+        async with self._run_lock:
+            if self._running:
+                raise RuntimeError("ServerAsyncOperationManager.run() is already running.")
+            self._running = True
+
+        async with anyio.create_task_group() as tg:
+            self._task_group = tg
+            logger.info("ServerAsyncOperationManager started")
+            # Start cleanup loop
+            tg.start_soon(self.cleanup_loop)
+            try:
+                yield
+            finally:
+                logger.info("ServerAsyncOperationManager shutting down")
+                # Stop cleanup loop gracefully
+                await self.stop_cleanup_loop()
+                # Cancel task group to stop all spawned tasks
+                tg.cancel_scope.cancel()
+                self._task_group = None
+                self._running = False
+
+    def start_task(
+        self,
+        token: str,
+        task_func: Callable[[], Awaitable[None]],
+        request_context: Any = None,
+        request_ctx_var: Any = None,
+    ) -> None:
+        """Start an async task immediately in the independent task group."""
+        if self._task_group is None:
+            raise RuntimeError("Task group not started. Call run() first.")
+
+        async def run_task_with_context():
+            context_token = None
+            try:
+                if request_context and request_ctx_var:
+                    context_token = request_ctx_var.set(request_context)
+                await task_func()
+            except Exception:
+                # Handle task failures gracefully
+                pass
+            finally:
+                if context_token and request_ctx_var:
+                    request_ctx_var.reset(context_token)
+
+        self._task_group.start_soon(run_task_with_context, name=f"lro_{token}")
+
     def create_operation(
         self,
         tool_name: str,

src/mcp/shared/memory.py

@@ -71,30 +71,31 @@ async def create_connected_server_and_client_session(
         server_read, server_write = server_streams
 
         # Create a cancel scope for the server task
-        async with anyio.create_task_group() as tg:
-            tg.start_soon(
-                lambda: server.run(
-                    server_read,
-                    server_write,
-                    server.create_initialization_options(),
-                    raise_exceptions=raise_exceptions,
+        async with server.async_operations.run():
+            async with anyio.create_task_group() as tg:
+                tg.start_soon(
+                    lambda: server.run(
+                        server_read,
+                        server_write,
+                        server.create_initialization_options(),
+                        raise_exceptions=raise_exceptions,
+                    )
                 )
-            )
-
-            try:
-                async with ClientSession(
-                    read_stream=client_read,
-                    write_stream=client_write,
-                    read_timeout_seconds=read_timeout_seconds,
-                    sampling_callback=sampling_callback,
-                    list_roots_callback=list_roots_callback,
-                    logging_callback=logging_callback,
-                    message_handler=message_handler,
-                    client_info=client_info,
-                    elicitation_callback=elicitation_callback,
-                    protocol_version=protocol_version,
-                ) as client_session:
-                    await client_session.initialize()
-                    yield client_session
-            finally:
-                tg.cancel_scope.cancel()
+
+                try:
+                    async with ClientSession(
+                        read_stream=client_read,
+                        write_stream=client_write,
+                        read_timeout_seconds=read_timeout_seconds,
+                        sampling_callback=sampling_callback,
+                        list_roots_callback=list_roots_callback,
+                        logging_callback=logging_callback,
+                        message_handler=message_handler,
+                        client_info=client_info,
+                        elicitation_callback=elicitation_callback,
+                        protocol_version=protocol_version,
+                    ) as client_session:
+                        await client_session.initialize()
+                        yield client_session
+                finally:
+                    tg.cancel_scope.cancel()

tests/server/fastmcp/test_integration.py

@@ -1013,7 +1013,6 @@ async def test_immediate_result_backward_compatibility(server_transport: str, se
                 await anyio.sleep(0.5)
             else:
                 pytest.fail("Async operation timed out")
-                await anyio.sleep(0.01)
 
 
 # Test async progress notifications