fix: cancel in-flight handlers when transport closes in server.run() (#2306)
This commit is contained in:
parent
883d893097
commit
92c693bb73
1 changed files with 199 additions and 16 deletions
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue