mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 14:54:24 +01:00
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.
This commit is contained in:
@@ -470,10 +470,14 @@ 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:
|
||||
from contextlib import AsyncExitStack
|
||||
|
||||
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)
|
||||
)
|
||||
|
||||
async for message in session.incoming_messages:
|
||||
logger.debug(f"Received message: {message}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user