From 13d3b895a1e5769dcd6e4faddfe387c8258f1873 Mon Sep 17 00:00:00 2001 From: Jihaochang1 <358550956@qq.com> Date: Sat, 24 Oct 2026 12:30:05 +0800 Subject: [PATCH] fix: serialize progress callbacks and drain them before the request completes Each `notifications/progress` used to spawn its own task with nothing waiting for it, so callbacks for the same request could run concurrently, complete in any order, and still be running after `send_raw_request()` handed back the result (#3747). A callback that awaits anything (updating a UI, writing a log) also delayed the callbacks behind it arbitrarily. Each notification's task now awaits its predecessor's done-event before invoking the callback, so the callbacks for a request run one at a time in the order the notifications arrived, and `send_raw_request` drains the tail of the chain before the request completes - on the success and error-response paths alike. The callback stays shielded and exception-contained, the read loop never blocks on it, and the drain is bounded because the pending entry is removed first (progress arriving after the response is still dropped, as before). The same dispatcher serves client `call_tool(progress_callback=...)` and server `send_request(progress_callback=...)`, so both directions are fixed. Fixes #3747 --- src/mcp/shared/jsonrpc_dispatcher.py | 66 ++++++++++++++++++++- tests/interaction/lowlevel/test_progress.py | 38 ++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 223373da0f..dcb548c688 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -127,6 +127,13 @@ class _Pending: send: MemoryObjectSendStream[dict[str, Any] | ErrorData] receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData] on_progress: ProgressFnT | None = None + progress_chain: anyio.Event | None = None + """Done-event of the most recently spawned progress-callback task for this request. + + Each notification's task awaits its predecessor's done-event before invoking + the callback, so callbacks run one at a time in arrival order. `send_raw_request` + drains the tail before the request completes; see `_spawn_progress`. + """ @dataclass(slots=True) @@ -378,6 +385,7 @@ async def send_raw_request( # saw, while skipping it would leak a delivered request's handler. request_write_started = False timeout_armed = False + progress_chain: anyio.Event | None = None target = out_params.get("name") span_name = f"MCP send {method}{f' {target}' if isinstance(target, str) else ''}" @@ -453,10 +461,23 @@ async def send_raw_request( raise finally: # Remove the waiter on every path so a late response is dropped, not leaked. + # Captured before the pop: progress notifications still correlating to this + # pending have already been spawned by the sequential read loop, and the tail + # of their chain is what the drain below waits for. + progress_chain = pending.progress_chain self._pending.pop(pending_key, None) send.close() receive.close() + # Wait for the request's progress callbacks to finish, serialized behind + # `progress_chain`, so the request does not complete while an earlier + # progress callback is still running (#3747). Shielded so a caller + # cancellation cannot strand the chain mid-drain. Bounded: the pending + # entry is gone, so no further callbacks can be spawned for this request. + if progress_chain is not None: + with anyio.CancelScope(shield=True): + await progress_chain.wait() + if isinstance(outcome, ErrorData): raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data) return outcome @@ -653,8 +674,8 @@ def _dispatch_notification( ): total = msg.params.get("total") message = msg.params.get("message") - self._spawn( - _shielded_progress(pending.on_progress), + self._spawn_progress( + pending, float(progress), float(total) if isinstance(total, int | float) else None, message if isinstance(message, str) else None, @@ -675,6 +696,47 @@ def _dispatch_notification( ) self._spawn(_contained_notify(on_notify), dctx, msg.method, msg.params, sender_ctx=sender_ctx) + def _spawn_progress( + self, + pending: _Pending, + progress: float, + total: float | None, + message: str | None, + sender_ctx: contextvars.Context | None, + ) -> None: + """Deliver one progress notification to `pending.on_progress`, in arrival order. + + The read loop sees notifications one at a time, but spawning each callback + outright lets them overlap: callbacks for the same request then run + concurrently, finish in any order, and can still be running when + `send_raw_request` hands back the result (#3747). Instead each task first + awaits its predecessor's done-event, so the callbacks for a request run + one at a time in the order the notifications arrived. The newest + done-event is kept on `pending.progress_chain`, which `send_raw_request` + drains before the request completes. + + The callback stays shielded and exception-contained (via `_shielded_progress`), + so a slow or failing callback cannot stall the read loop or the task group; + the done-event is set in a `finally`, so the chain and the drain always + release even when the dispatcher is cancelled. + """ + previous = pending.progress_chain + done = anyio.Event() + on_progress = pending.on_progress + assert on_progress is not None # only called for a pending that carries a callback + callback = _shielded_progress(on_progress) + + async def run() -> None: + try: + if previous is not None: + await previous.wait() + await callback(progress, total, message) + finally: + done.set() + + pending.progress_chain = done + self._spawn(run, sender_ctx=sender_ctx) + def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None: pending = self._pending.get(coerce_request_id(request_id)) if request_id is not None else None if pending is None: diff --git a/tests/interaction/lowlevel/test_progress.py b/tests/interaction/lowlevel/test_progress.py index ed98764622..6c566488a0 100644 --- a/tests/interaction/lowlevel/test_progress.py +++ b/tests/interaction/lowlevel/test_progress.py @@ -277,3 +277,41 @@ async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestPara await client.call_tool("zigzag", {}, progress_callback=collect) assert received == snapshot([0.5, 0.3, 0.9]) + + +@requirement("protocol:progress:callback") +async def test_progress_callbacks_are_serialized_and_finish_before_the_result(connect: Connect) -> None: + """Progress callbacks for one request run one at a time, in order, and before the result. + + Regression test for #3747: each ``notifications/progress`` used to spawn its own + task, so callbacks for the same request ran concurrently, completed in any order, + and could still be running after ``call_tool()`` had returned. The existing + in-order test cannot catch this because its callback never awaits. + """ + events: list[str] = [] + + async def collect(progress: float, total: float | None, message: str | None) -> None: + # Reverse the sleep durations so overlapping callbacks are observable: + # notification 3 would finish first if they ran concurrently. + await anyio.sleep(0.1 * (4 - progress)) + events.append(f"progress {progress}") + + async def list_tools( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListToolsResult: + return types.ListToolsResult(tools=[types.Tool(name="download", input_schema={"type": "object"})]) + + async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult: + assert params.name == "download" + await ctx.session.report_progress(1.0, total=3.0) + await ctx.session.report_progress(2.0, total=3.0) + await ctx.session.report_progress(3.0, total=3.0) + return CallToolResult(content=[TextContent(text="downloaded")]) + + server = Server("downloader", on_list_tools=list_tools, on_call_tool=call_tool) + + async with connect(server) as client: + await client.call_tool("download", {}, progress_callback=collect) + events.append("call_tool returned") + + assert events == ["progress 1.0", "progress 2.0", "progress 3.0", "call_tool returned"]