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

src/mcp/server/lowlevel/server.py

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

src/mcp/shared/session.py

@@ -105,7 +105,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
@@ -418,7 +418,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))