From e444f63bef51dfa5fe7b7375c29f1ebf9de38432 Mon Sep 17 00:00:00 2001 From: 686f6c61 Date: Sun, 2 Aug 2026 10:57:55 +0200 Subject: [PATCH] fix(mcp): serialize keepalive probes with the per-server RPC lock Keepalive ping/list_tools no longer interleave with in-flight tool calls on the same ClientSession stream. If the lock is already held, skip the probe (the active RPC is liveness). Otherwise acquire the lock for the probe. Fixes #70218 Co-authored-by: Diego Gomez <22959713+0xquinto@users.noreply.github.com> --- tests/tools/test_mcp_capability_gating.py | 75 +++++++++++++++++++++++ tools/mcp_tool.py | 65 +++++++++++++------- 2 files changed, 118 insertions(+), 22 deletions(-) diff --git a/tests/tools/test_mcp_capability_gating.py b/tests/tools/test_mcp_capability_gating.py index a0fef278fe91..b1043188b30b 100644 --- a/tests/tools/test_mcp_capability_gating.py +++ b/tests/tools/test_mcp_capability_gating.py @@ -139,6 +139,81 @@ async def test_keepalive_uses_ping_legacy_fallback(self): task.session.list_tools.assert_not_called() + + async def test_keepalive_skips_an_active_tool_rpc(self): + """An active RPC proves liveness; keepalive must not wait or overlap.""" + import asyncio + from types import SimpleNamespace + from unittest.mock import AsyncMock + + task = MCPServerTask("test") + rpc_started = asyncio.Event() + release_rpc = asyncio.Event() + ping_started = asyncio.Event() + + async def active_tool_rpc(): + async with task._rpc_lock: + rpc_started.set() + await release_rpc.wait() + + async def send_ping(): + ping_started.set() + + task.session = SimpleNamespace( + list_tools=AsyncMock(), + send_ping=AsyncMock(side_effect=send_ping), + ) + + active_task = asyncio.create_task(active_tool_rpc()) + try: + await rpc_started.wait() + await asyncio.wait_for(task._keepalive_probe(), timeout=0.1) + + assert not ping_started.is_set() + task.session.send_ping.assert_not_awaited() + task.session.list_tools.assert_not_awaited() + finally: + release_rpc.set() + await active_task + + async def test_active_keepalive_serializes_a_user_rpc(self): + """A user RPC must wait while keepalive owns the shared session lock.""" + import asyncio + from types import SimpleNamespace + from unittest.mock import AsyncMock + + task = MCPServerTask("test") + ping_started = asyncio.Event() + release_ping = asyncio.Event() + rpc_started = asyncio.Event() + + async def send_ping(): + ping_started.set() + await release_ping.wait() + + async def user_rpc(): + async with task._rpc_lock: + rpc_started.set() + + task.session = SimpleNamespace( + list_tools=AsyncMock(), + send_ping=AsyncMock(side_effect=send_ping), + ) + + keepalive_task = asyncio.create_task(task._keepalive_probe()) + await ping_started.wait() + rpc_task = asyncio.create_task(user_rpc()) + try: + await asyncio.sleep(0) + assert not rpc_started.is_set() + finally: + release_ping.set() + await keepalive_task + await rpc_task + + assert rpc_started.is_set() + + class TestKeepaliveInterval: """The keepalive cadence is configurable so servers with short session TTLs (e.g. Unreal Engine editor MCP, ~15s) can refresh fast enough to keep diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index b3c356788a08..4a2aeab5c66b 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -2178,32 +2178,53 @@ async def _keepalive_probe(self) -> None: transport connection so a server that gains ping support after a reconnect is re-probed with the cheap path. + When a user RPC owns ``_rpc_lock``, that RPC and its own timeout are the + liveness detector; the periodic probe is intentionally skipped so a + long-running tool call cannot interleave with ``ping`` / ``list_tools`` + on the same ClientSession stream (#70218). + Raises on a genuine connection failure so the caller triggers a reconnect; returns normally when the session is alive. """ - if not self._ping_unsupported: - try: - await asyncio.wait_for(self.session.send_ping(), timeout=30.0) - return - except Exception as exc: - # Only a "method not found" means ping is unsupported. Any - # other error (timeout, closed transport, session expired) is - # a real liveness failure — propagate so we reconnect. - if not _is_method_not_found_error(exc): - raise - if not self._advertises_tools(): - # No ping, no tools → no cheaper probe to fall back to. - raise - self._ping_unsupported = True - logger.info( - "MCP server '%s': does not implement the optional 'ping' " - "utility (-32601); using 'list_tools' for keepalive on " - "this connection.", - self.name, - ) + session = self.session + if session is None: + raise RuntimeError("MCP keepalive requested without an active session") + + # A ClientSession is one JSON-RPC stream. Sending a keepalive while a + # tools/call is in flight can wedge the original request on some + # SDK/server combinations even when the ping itself succeeds. An + # active RPC already proves liveness, so skip this cycle rather than + # queueing a redundant probe behind a potentially long-running call. + if self._rpc_lock.locked(): + return + + # Every other client-initiated RPC uses this lock, so keepalive must + # too. The lock also closes the race with a call that starts after the + # locked() check above. + async with self._rpc_lock: + if not self._ping_unsupported: + try: + await asyncio.wait_for(session.send_ping(), timeout=30.0) + return + except Exception as exc: + # Only a "method not found" means ping is unsupported. Any + # other error (timeout, closed transport, session expired) is + # a real liveness failure — propagate so we reconnect. + if not _is_method_not_found_error(exc): + raise + if not self._advertises_tools(): + # No ping, no tools → no cheaper probe to fall back to. + raise + self._ping_unsupported = True + logger.info( + "MCP server '%s': does not implement the optional 'ping' " + "utility (-32601); using 'list_tools' for keepalive on " + "this connection.", + self.name, + ) - # Fallback probe for servers without ping support. - await asyncio.wait_for(self.session.list_tools(), timeout=30.0) + # Fallback probe for servers without ping support. + await asyncio.wait_for(session.list_tools(), timeout=30.0) def _mark_session_proven(self) -> None: """Record that the current session demonstrated real health.