aljaz/mcp-python-sdk
1
0
Fork 0
mirror of https://github.com/aljazceru/mcp-python-sdk.git synced 2025-12-19 14:54:24 +01:00
Code Issues Packages Projects Releases Wiki Activity

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
parent e5815bd162
commit fddba00723
1 changed files with 32 additions and 28 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, raise_exceptions: bool = False,
): ):
with warnings.catch_warnings(record=True) as w: with warnings.catch_warnings(record=True) as w:
async with self.lifespan(self) as lifespan_context: from contextlib import AsyncExitStack
async with ServerSession(
read_stream, write_stream, initialization_options
) as session:
async for message in session.incoming_messages:
logger.debug(f"Received message: {message}")
match message: async with AsyncExitStack() as stack:
case ( lifespan_context = await stack.enter_async_context(self.lifespan(self))
RequestResponder( session = await stack.enter_async_context(
request=types.ClientRequest(root=req) ServerSession(read_stream, write_stream, initialization_options)
) 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: async for message in session.incoming_messages:
logger.info( logger.debug(f"Received message: {message}")
"Warning: %s: %s",
warning.category.__name__, match message:
warning.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( async def _handle_request(
self, self,