fix: JSONRPCDispatcher coerces string response/progress IDs to int for correlation

Matches BaseSession._normalize_request_id and the TypeScript SDK: a peer
that echoes the request ID as a JSON string still resolves the waiter.
Applied at both lookup sites (_resolve_pending and the progress-token
match). Parity prep for the PR6 e2e suite.

Author: maxisbey

src/mcp/shared/jsonrpc_dispatcher.py

@@ -76,6 +76,21 @@
 `TransportContext(kind="jsonrpc", can_send_request=True)` when not supplied."""
 
 
+def _coerce_id(request_id: RequestId) -> RequestId:
+    """Coerce a string request ID to int when it's a valid int literal.
+
+    `_allocate_id` only ever produces ``int`` keys for ``_pending``, but a peer
+    may echo the ID back as a JSON string. The TypeScript SDK and `BaseSession`
+    both perform this coercion at lookup time so the response still correlates.
+    """
+    if isinstance(request_id, str):
+        try:
+            return int(request_id)
+        except ValueError:
+            pass
+    return request_id
+
+
 @dataclass(slots=True)
 class _Pending:
     """An outbound request awaiting its response."""
@@ -409,7 +424,7 @@ def _dispatch_notification(
         if msg.method == "notifications/progress":
             match msg.params:
                 case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if (
-                    pending := self._pending.get(token)
+                    pending := self._pending.get(_coerce_id(token))
                 ) is not None and pending.on_progress is not None:
                     total = msg.params.get("total")
                     message = msg.params.get("message")
@@ -428,7 +443,7 @@ def _dispatch_notification(
         self._spawn(on_notify, dctx, msg.method, msg.params, sender_ctx=sender_ctx)
 
     def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None:
-        pending = self._pending.get(request_id) if request_id is not None else None
+        pending = self._pending.get(_coerce_id(request_id)) if request_id is not None else None
         if pending is None:
             logger.debug("dropping response for unknown/late request id %r", request_id)
             return

tests/shared/test_jsonrpc_dispatcher.py

@@ -18,6 +18,7 @@
 from mcp.shared.exceptions import MCPError
 from mcp.shared.jsonrpc_dispatcher import (  # pyright: ignore[reportPrivateUsage]
     JSONRPCDispatcher,
+    _coerce_id,
     _outbound_metadata,
     _Pending,
 )
@@ -29,6 +30,7 @@
     INVALID_PARAMS,
     ErrorData,
     JSONRPCError,
+    JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
     Tool,
@@ -511,6 +513,87 @@ def test_outbound_metadata_with_resumption_token_returns_client_metadata():
     assert _outbound_metadata(None, {}) is None
 
 
+@pytest.mark.anyio
+async def test_response_with_string_id_correlates_to_int_keyed_pending_request():
+    """A peer that echoes the request ID as a JSON string still resolves the waiter."""
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
+    on_request, on_notify = echo_handlers(Recorder())
+    try:
+        async with anyio.create_task_group() as tg:
+            await tg.start(client.run, on_request, on_notify)
+            with anyio.fail_after(5):
+
+                async def respond_stringly() -> None:
+                    out = await c2s_recv.receive()
+                    assert isinstance(out, SessionMessage)
+                    assert isinstance(out.message, JSONRPCRequest)
+                    rid = out.message.id
+                    assert isinstance(rid, int)
+                    await s2c_send.send(
+                        SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=str(rid), result={"ok": True}))
+                    )
+
+                tg.start_soon(respond_stringly)
+                result = await client.send_raw_request("ping", None)
+                assert result == {"ok": True}
+            tg.cancel_scope.cancel()
+    finally:
+        for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
+            s.close()
+
+
+@pytest.mark.anyio
+async def test_progress_with_string_token_reaches_callback_for_int_keyed_request():
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
+    on_request, on_notify = echo_handlers(Recorder())
+    seen: list[float] = []
+    try:
+        async with anyio.create_task_group() as tg:
+            await tg.start(client.run, on_request, on_notify)
+            with anyio.fail_after(5):
+
+                async def respond_with_string_token_progress() -> None:
+                    out = await c2s_recv.receive()
+                    assert isinstance(out, SessionMessage)
+                    assert isinstance(out.message, JSONRPCRequest)
+                    rid = out.message.id
+                    assert isinstance(rid, int)
+                    await s2c_send.send(
+                        SessionMessage(
+                            message=JSONRPCNotification(
+                                jsonrpc="2.0",
+                                method="notifications/progress",
+                                params={"progressToken": str(rid), "progress": 0.5},
+                            )
+                        )
+                    )
+                    await s2c_send.send(
+                        SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=rid, result={"ok": True}))
+                    )
+
+                async def on_progress(progress: float, total: float | None, message: str | None) -> None:
+                    seen.append(progress)
+
+                tg.start_soon(respond_with_string_token_progress)
+                result = await client.send_raw_request("ping", None, {"on_progress": on_progress})
+                assert result == {"ok": True}
+            tg.cancel_scope.cancel()
+    finally:
+        for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
+            s.close()
+    assert seen == [0.5]
+
+
+def test_coerce_id_passes_through_non_numeric_string_and_int():
+    assert _coerce_id("7") == 7
+    assert _coerce_id("not-an-int") == "not-an-int"
+    assert _coerce_id(42) == 42
+
+
 @pytest.mark.anyio
 async def test_jsonrpc_error_response_with_null_id_is_dropped():
     """Parse-error responses (id=null) have no waiter; they're logged and dropped."""