fix: JSONRPCDispatcher.run cancels in-flight handlers on read-stream EOF

Restores the Server.run() behaviour the dispatcher rework dropped: at
read-stream EOF the task group cancels in-flight handler tasks instead of
joining on them. Without this, a handler that outlives its caller (its
request timed out client-side, or the client disconnected mid-call) keeps
run() from returning forever, leaking the handler task and over SSE the
GET request that hosts the session.

Regression test parks a handler in sleep_forever(), EOFs the read stream,
asserts run() returns within fail_after(5). Confirmed to hang on the
unpatched code.

Author: maxisbey

src/mcp/shared/jsonrpc_dispatcher.py

@@ -356,24 +356,34 @@ async def run(
                 self._tg = tg
                 self._running = True
                 task_status.started()
-                async with self._read_stream, self._write_stream:
-                    try:
-                        async for item in self._read_stream:
-                            # Duck-typed: `_context_streams.ContextReceiveStream`
-                            # exposes `.last_context` (the sender's contextvars
-                            # snapshot per message). Plain memory streams don't.
-                            sender_ctx: contextvars.Context | None = getattr(self._read_stream, "last_context", None)
-                            self._dispatch(item, on_request, on_notify, sender_ctx)
-                    except anyio.ClosedResourceError:
-                        # The transport closed our receive end and we looped back
-                        # to `__anext__` on the now-closed stream (stateless SHTTP
-                        # teardown). Same as EOF.
-                        logger.debug("read stream closed by transport; treating as EOF")
-                # Read stream EOF: wake any blocked `send_raw_request` waiters now,
-                # *before* the task group joins, so handlers parked in
-                # `dctx.send_raw_request()` can unwind and the join doesn't deadlock.
-                self._running = False
-                self._fan_out_closed()
+                try:
+                    async with self._read_stream, self._write_stream:
+                        try:
+                            async for item in self._read_stream:
+                                # Duck-typed: `_context_streams.ContextReceiveStream`
+                                # exposes `.last_context` (the sender's contextvars
+                                # snapshot per message). Plain memory streams don't.
+                                sender_ctx: contextvars.Context | None = getattr(
+                                    self._read_stream, "last_context", None
+                                )
+                                self._dispatch(item, on_request, on_notify, sender_ctx)
+                        except anyio.ClosedResourceError:
+                            # The transport closed our receive end and we looped
+                            # back to `__anext__` on the now-closed stream
+                            # (stateless SHTTP teardown). Same as EOF.
+                            logger.debug("read stream closed by transport; treating as EOF")
+                    # Read stream EOF: wake any blocked `send_raw_request` waiters
+                    # (callers outside this task group) with CONNECTION_CLOSED.
+                    self._running = False
+                    self._fan_out_closed()
+                finally:
+                    # Transport closed: cancel in-flight handlers. Without this
+                    # the task-group join waits for them, and a handler that
+                    # outlives its caller (its request timed out client-side, or
+                    # the client disconnected mid-call) would keep `run()` from
+                    # returning forever. Same behaviour as `Server.run()` before
+                    # the dispatcher rework.
+                    tg.cancel_scope.cancel()
         finally:
             # Covers the cancel/crash paths where the inline fan-out above is
             # never reached. Idempotent.

tests/shared/test_jsonrpc_dispatcher.py

@@ -196,6 +196,48 @@ async def test_run_returns_cleanly_when_read_stream_receive_end_is_closed():
         s.close()
 
 
+@pytest.mark.anyio
+async def test_run_cancels_in_flight_handlers_when_read_stream_eofs():
+    """A handler that outlives its caller must not keep run() from returning.
+
+    Without the cancel-at-EOF, the task-group join would wait on this handler
+    forever (over SSE that leaks the handler task and the GET request hosting
+    the session).
+    """
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
+    handler_started = anyio.Event()
+    handler_cancelled = anyio.Event()
+
+    async def park(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        handler_started.set()
+        try:
+            await anyio.sleep_forever()
+        finally:
+            handler_cancelled.set()
+        raise NotImplementedError
+
+    async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
+        raise NotImplementedError
+
+    run_returned = anyio.Event()
+
+    async def drive() -> None:
+        await server.run(park, on_notify)
+        run_returned.set()
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(drive)
+        await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="x", params=None)))
+        with anyio.fail_after(5):
+            await handler_started.wait()
+            c2s_send.close()  # EOF the read side; run() must cancel the parked handler
+            await run_returned.wait()
+    assert handler_cancelled.is_set()
+    s2c_recv.close()
+
+
 @pytest.mark.anyio
 async def test_run_closes_write_stream_on_exit():
     """run() enters both streams; the write end is released on EOF."""