Made message handling concurrent

Author: jerome3o-anthropic

src/mcp/server/lowlevel/server.py

@@ -70,6 +70,7 @@ async def main():
 from collections.abc import Awaitable, Callable
 from typing import Any, Sequence
 
+import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import AnyUrl
 
@@ -434,6 +435,29 @@ async def handler(req: types.CompleteRequest):
 
         return decorator
 
+    async def _handle_message(
+        self,
+        message: RequestResponder[types.ClientRequest, types.ServerResult]
+        | types.ClientNotification
+        | Exception,
+        session: ServerSession,
+        raise_exceptions: bool = False,
+    ):
+        with warnings.catch_warnings(record=True) as w:
+            match message:
+                case (
+                    RequestResponder(request=types.ClientRequest(root=req)) as responder
+                ):
+                    with responder:
+                        await self._handle_request(
+                            message, req, session, raise_exceptions
+                        )
+                case types.ClientNotification(root=notify):
+                    await self._handle_notification(notify)
+
+            for warning in w:
+                logger.info(f"Warning: {warning.category.__name__}: {warning.message}")
+
     async def run(
         self,
         read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
@@ -445,30 +469,19 @@ async def run(
         # in-process servers.
         raise_exceptions: bool = False,
     ):
-        with warnings.catch_warnings(record=True) as w:
-            async with ServerSession(
-                read_stream, write_stream, initialization_options
-            ) as session:
+        async with ServerSession(
+            read_stream, write_stream, initialization_options
+        ) as session:
+            async with anyio.create_task_group() as tg:
                 async for message in session.incoming_messages:
                     logger.debug(f"Received message: {message}")
 
-                    match message:
-                        case (
-                            RequestResponder(
-                                request=types.ClientRequest(root=req)
-                            ) as responder
-                        ):
-                            with responder:
-                                await self._handle_request(
-                                    message, req, session, raise_exceptions
-                                )
-                        case types.ClientNotification(root=notify):
-                            await self._handle_notification(notify)
-
-                    for warning in w:
-                        logger.info(
-                            f"Warning: {warning.category.__name__}: {warning.message}"
-                        )
+                    tg.start_soon(
+                        self._handle_message,
+                        message,
+                        session,
+                        raise_exceptions,
+                    )
 
     async def _handle_request(
         self,