[v1.x] fix: handle ClosedResourceError when transport closes mid-request (#2334)
Some checks failed
Main branch checks / checks (push) Has been cancelled
Some checks failed
Main branch checks / checks (push) Has been cancelled
Co-authored-by: Owen Devereaux <owendevereaux@users.noreply.github.com> Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>
This commit is contained in:
parent
2e9897e2b9
commit
6524782667
2 changed files with 208 additions and 25 deletions
|
|
@ -671,16 +671,23 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
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,
|
||||
|
|
@ -763,12 +770,18 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
response = await handler(req)
|
||||
except McpError as err: # pragma: no cover
|
||||
response = err.error
|
||||
except anyio.get_cancelled_exc_class(): # pragma: no cover
|
||||
logger.info(
|
||||
"Request %s cancelled - duplicate response suppressed",
|
||||
message.request_id,
|
||||
)
|
||||
return
|
||||
except anyio.get_cancelled_exc_class():
|
||||
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: # pragma: no cover
|
||||
if raise_exceptions:
|
||||
raise err
|
||||
|
|
@ -777,16 +790,24 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
# Reset the global state after we are done
|
||||
if token is not None: # pragma: no branch
|
||||
request_ctx.reset(token)
|
||||
|
||||
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")
|
||||
|
||||
async def _handle_notification(self, notify: Any):
|
||||
|
|
|
|||
|
|
@ -108,7 +108,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
|
|||
) -> None:
|
||||
"""Exit the context manager, performing cleanup and notifying completion."""
|
||||
try:
|
||||
if self._completed: # pragma: no branch
|
||||
if self._completed:
|
||||
self._on_complete(self)
|
||||
finally:
|
||||
self._entered = False
|
||||
|
|
@ -445,7 +445,9 @@ class BaseSession(
|
|||
finally:
|
||||
# after the read stream is closed, we need to send errors
|
||||
# to any pending requests
|
||||
for id, stream in self._response_streams.items():
|
||||
# Snapshot: stream.send() wakes the waiter, whose finally pops
|
||||
# from _response_streams before the next __next__() call.
|
||||
for id, stream in list(self._response_streams.items()):
|
||||
error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
|
||||
try:
|
||||
await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error))
|
||||
|
|
|
|||
Loading…
Reference in a new issue