From fddba007230bf9fb94a6e2595f0c67ef2f857815 Mon Sep 17 00:00:00 2001
From: David Soria Parra <davidsp@anthropic.com>
Date: Tue, 11 Feb 2025 12:26:32 +0000
Subject: [PATCH] refactor: improve server context management with
 AsyncExitStack

Replace nested context managers with AsyncExitStack to ensure proper cleanup
order during server shutdown and make the code more maintainable.
---
 src/mcp/server/lowlevel/server.py | 60 ++++++++++++++++---------------
 1 file changed, 32 insertions(+), 28 deletions(-)

diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
index 28942cf8..a4a8510c 100644
--- a/src/mcp/server/lowlevel/server.py
+++ b/src/mcp/server/lowlevel/server.py
@@ -470,36 +470,40 @@ class Server(Generic[LifespanResultT]):
         raise_exceptions: bool = False,
     ):
         with warnings.catch_warnings(record=True) as w:
-            async with self.lifespan(self) as lifespan_context:
-                async with ServerSession(
-                    read_stream, write_stream, initialization_options
-                ) as session:
-                    async for message in session.incoming_messages:
-                        logger.debug(f"Received message: {message}")
+            from contextlib import AsyncExitStack
 
-                        match message:
-                            case (
-                                RequestResponder(
-                                    request=types.ClientRequest(root=req)
-                                ) as responder
-                            ):
-                                with responder:
-                                    await self._handle_request(
-                                        message,
-                                        req,
-                                        session,
-                                        lifespan_context,
-                                        raise_exceptions,
-                                    )
-                            case types.ClientNotification(root=notify):
-                                await self._handle_notification(notify)
+            async with AsyncExitStack() as stack:
+                lifespan_context = await stack.enter_async_context(self.lifespan(self))
+                session = await stack.enter_async_context(
+                    ServerSession(read_stream, write_stream, initialization_options)
+                )
 
-                        for warning in w:
-                            logger.info(
-                                "Warning: %s: %s",
-                                warning.category.__name__,
-                                warning.message,
-                            )
+                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,
+                                    lifespan_context,
+                                    raise_exceptions,
+                                )
+                        case types.ClientNotification(root=notify):
+                            await self._handle_notification(notify)
+
+                    for warning in w:
+                        logger.info(
+                            "Warning: %s: %s",
+                            warning.category.__name__,
+                            warning.message,
+                        )
 
     async def _handle_request(
         self,