fix: cancel in-flight handlers when transport closes in server.run() (#2306)

This commit is contained in:
Max Isbey 2026-03-20 13:37:32 +00:00 • committed by GitHub
commit 92c693bb73
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -387,16 +387,23 @@ class Server(Generic[LifespanResultT]):
await stack.enter_async_context(task_support.run())
async with anyio.create_task_group() as tg:
async for message in session.incoming_messages:
logger.debug("Received message: %s", message)
try:
async for message in session.incoming_messages:
logger.debug("Received message: %s", message)
tg.start_soon(
self._handle_message,
message,
session,
lifespan_context,
raise_exceptions,
)
tg.start_soon(
self._handle_message,
message,
session,
lifespan_context,
raise_exceptions,
)
finally:
# Transport closed: cancel in-flight handlers. Without this the
# TG join waits for them, and when they eventually try to
# respond they hit a closed write stream (the session's
# _receive_loop closed it when the read stream ended).
tg.cancel_scope.cancel()
async def _handle_message(
self,
@ -470,16 +477,32 @@ class Server(Generic[LifespanResultT]):
except MCPError as err:
response = err.error
except anyio.get_cancelled_exc_class():
logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
return
if message.cancelled:
# Client sent CancelledNotification; responder.cancel() already
# sent an error response, so skip the duplicate.
logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
return
# Transport-close cancellation from the TG in run(); re-raise so the
# TG swallows its own cancellation.
raise
except Exception as err:
if raise_exceptions: # pragma: no cover
raise err
response = types.ErrorData(code=0, message=str(err))
await message.respond(response)
else: # pragma: no cover
await message.respond(types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found"))
response = types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found")
try:
await message.respond(response)
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
# Transport closed between handler unblocking and respond. Happens
# when _receive_loop's finally wakes a handler blocked on
# send_request: the handler runs to respond() before run()'s TG
# cancel fires, but after the write stream closed. Closed if our
# end closed (_receive_loop's async-with exit); Broken if the peer
# end closed first (streamable_http terminate()).
logger.debug("Response for %s dropped - transport closed", message.request_id)
return
logger.debug("Response sent")