From d5d26a743d6fc7485fc01247ad87904a27d0e1a0 Mon Sep 17 00:00:00 2001 From: Soham Kukreti Date: Wed, 16 Sep 2026 13:11:37 +0530 Subject: [PATCH] fix(mcp): order SSE handshake before early tool calls Gate the SSE inbox so initialize is processed before concurrent tools/call POSTs, fixing -32602 rejections on connect (#2233). --- deploy/docker/mcp_bridge.py | 39 ++++++++++++- tests/test_issue_2233_mcp_sse.py | 96 ++++++++++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 tests/test_issue_2233_mcp_sse.py diff --git a/deploy/docker/mcp_bridge.py b/deploy/docker/mcp_bridge.py index 653bc368b..24c14c628 100644 --- a/deploy/docker/mcp_bridge.py +++ b/deploy/docker/mcp_bridge.py @@ -247,7 +247,11 @@ async def ws_to_srv(): class _MCPSseApp: async def __call__(self, scope, receive, send): async with sse.connect_sse(scope, receive, send) as (read_stream, write_stream): - await mcp.run(read_stream, write_stream, init_opts) + # Reorder inbound POSTs so the MCP handshake reaches the session first. + c2s_send, c2s_recv = anyio.create_memory_object_stream(100) + async with anyio.create_task_group() as tg: + tg.start_soon(mcp.run, c2s_recv, write_stream, init_opts) + tg.start_soon(_sse_init_gate, read_stream, c2s_send) app.routes.append(Route(f"{base}/sse", endpoint=_MCPSseApp())) app.routes.append(Mount(f"{base}/messages", app=sse.handle_post_message)) @@ -263,6 +267,39 @@ async def _schema_endpoint(): # ── helpers ──────────────────────────────────────────────────── +async def _sse_init_gate(read_stream, c2s_send, timeout: float = 2.0) -> None: + """Park inbound SSE messages until initialize and notifications/initialized + have been forwarded, then flush them in arrival order. After timeout seconds, + flush regardless so a client that never completes the handshake cannot pin + the session. + """ + def rpc_method_name(m): + msg = getattr(m, "message", None) # stream may carry an Exception + return getattr(getattr(msg, "root", msg), "method", None) # mcp 1.x RootModel / 2.x plain union + + parked, want = [], ["initialize", "notifications/initialized"] + try: + with anyio.move_on_after(timeout) as scope: + async for msg in read_stream: + if want and rpc_method_name(msg) == want[0]: + await c2s_send.send(msg) + want.pop(0) + if not want: + break + else: + parked.append(msg) + if want and not scope.cancelled_caught: + return # client disconnected mid-handshake; nothing to flush + for msg in parked: + await c2s_send.send(msg) + async for msg in read_stream: + await c2s_send.send(msg) + except (anyio.EndOfStream, anyio.ClosedResourceError): + pass + finally: + with anyio.CancelScope(shield=True), suppress(Exception): + await c2s_send.aclose() + def _route_name(path: str) -> str: return re.sub(r"[/{}}]", "_", path).strip("_") diff --git a/tests/test_issue_2233_mcp_sse.py b/tests/test_issue_2233_mcp_sse.py new file mode 100644 index 000000000..54bad8998 --- /dev/null +++ b/tests/test_issue_2233_mcp_sse.py @@ -0,0 +1,96 @@ +""" +Issue #2233: MCP SSE rejects tool calls that are POSTed before initialize. + +Each SSE message is its own HTTP POST, so a fan-out client can land +tools/call in the session inbox ahead of the handshake and get JSON-RPC +-32602 ("Received request before initialization was complete"). +_sse_init_gate reorders the inbox so the handshake always goes first. +""" +import sys +from pathlib import Path + +import anyio +import pytest +from mcp.shared.message import SessionMessage +from mcp.types import JSONRPCMessage +from pydantic import TypeAdapter + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "deploy" / "docker")) +mcp_bridge = pytest.importorskip("mcp_bridge") + + +def _msg(method: str, id_: int | None = None) -> SessionMessage: + body = {"jsonrpc": "2.0", "method": method, "params": {}} + if id_ is not None: + body["id"] = id_ + return SessionMessage(TypeAdapter(JSONRPCMessage).validate_python(body)) + + +def _key(m: SessionMessage) -> tuple: + root = getattr(m.message, "root", m.message) + return (root.method, getattr(root, "id", None)) + + +async def _run_gate(inbound: list[SessionMessage]) -> list[tuple]: + """Send every message, close the inbox, return what the gate forwarded.""" + in_send, in_recv = anyio.create_memory_object_stream(100) + out_send, out_recv = anyio.create_memory_object_stream(100) + async with anyio.create_task_group() as tg: + tg.start_soon(mcp_bridge._sse_init_gate, in_recv, out_send) + async with in_send: + for m in inbound: + await in_send.send(m) + return [_key(m) async for m in out_recv] + + +def test_tool_calls_posted_before_initialize_are_reordered(): + # The race from #2233: three tool calls beat the handshake. + out = anyio.run(_run_gate, [ + _msg("tools/call", 1), + _msg("tools/call", 2), + _msg("tools/call", 3), + _msg("initialize", 0), + _msg("notifications/initialized"), + ]) + assert out == [ + ("initialize", 0), + ("notifications/initialized", None), + ("tools/call", 1), + ("tools/call", 2), + ("tools/call", 3), + ] + + +def test_well_behaved_client_is_passed_through_unchanged(): + seq = [ + _msg("initialize", 0), + _msg("notifications/initialized"), + _msg("tools/list", 1), + _msg("tools/call", 2), + ] + out = anyio.run(_run_gate, seq) + assert out == [_key(m) for m in seq] + + +def test_missing_handshake_is_flushed_after_timeout(): + # Inbox stays OPEN, so only the timeout can release the parked calls. + async def body(): + in_send, in_recv = anyio.create_memory_object_stream(100) + out_send, out_recv = anyio.create_memory_object_stream(100) + async with anyio.create_task_group() as tg: + tg.start_soon(mcp_bridge._sse_init_gate, in_recv, out_send, 0.2) + await in_send.send(_msg("tools/call", 1)) + await in_send.send(_msg("tools/call", 2)) + with anyio.fail_after(2): + got = [_key(await out_recv.receive()) for _ in range(2)] + await in_send.aclose() + return got + + assert anyio.run(body) == [("tools/call", 1), ("tools/call", 2)] + + +def test_disconnect_while_parked_drops_parked_messages(): + # Client hangs up before the handshake: parked calls must not be flushed + # into a session whose output stream is already gone. + out = anyio.run(_run_gate, [_msg("tools/call", 1), _msg("tools/call", 2)]) + assert out == []