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"]