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:
David Soria Parra
2025-02-11 12:26:32 +00:00
parent e5815bd162
commit fddba00723

View File

@@ -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}")