modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: improve server context management with AsyncExitStack

Browse source 
Replace nested context managers with AsyncExitStack to ensure proper cleanup
order during server shutdown and make the code more maintainable.
This commit is contained in:
David Soria Parra 2025-02-11 12:26:32 +00:00
commit fddba00723
No known key found for this signature in database
1 changed files with 34 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

60
src/mcp/server/lowlevel/server.py
View file

@ -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,