From ccc162f4e2f2f61584960d4efa21b9f33e12085e Mon Sep 17 00:00:00 2001 From: arminanton <29869547+arminanton@users.noreply.github.com> Date: Wed, 17 Jun 2026 14:55:57 -0700 Subject: [PATCH 1/2] fix(tools): skip MCP keepalive during in-flight calls + fail orphaned calls on reconnect An MCP stdio session is a single JSON-RPC stream. The idle keepalive (list_tools/send_ping) could fire while a call_tool was in flight, wedging the stream so the call timed out -> false reconnect -> the SDK does not always fail the pending call when its streams close, so its run_coroutine_threadsafe future never resolves and the agent thread polls to the full tool_timeout (hours). The tree already serializes RPCs via _rpc_lock (mcp_tool.py:1193) but the keepalive loop does not participate: its list_tools (:1398) / send_ping (:1403) run outside the lock, nothing skips the cycle when a call is active, and a reconnect/shutdown never cancels a pending call. Fix: track in-flight call tasks; skip the keepalive when a call is active and otherwise run the probe under _rpc_lock; _fail_inflight_calls() cancels pending calls on reconnect/shutdown; _call() converts a deliberate teardown-cancel into a retryable RuntimeError so the agent self-heals on the rebuilt session. Builds on the existing _rpc_lock groundwork. Related to #30268 (partial mitigation of the post-sleep keepalive storm); complementary to #30694 and #26493. --- .../tools/test_mcp_keepalive_inflight_race.py | 122 ++++++++++++++++++ tools/mcp_tool.py | 92 +++++++++++-- 2 files changed, 202 insertions(+), 12 deletions(-) create mode 100644 tests/tools/test_mcp_keepalive_inflight_race.py diff --git a/tests/tools/test_mcp_keepalive_inflight_race.py b/tests/tools/test_mcp_keepalive_inflight_race.py new file mode 100644 index 000000000000..a535517ca4da --- /dev/null +++ b/tests/tools/test_mcp_keepalive_inflight_race.py @@ -0,0 +1,122 @@ +"""Regression tests for the MCP keepalive / in-flight tool-call race. + +Background +========== +An MCP stdio session is a SINGLE JSON-RPC stream. The idle keepalive +(``list_tools`` / ``send_ping``) could fire WHILE a normal ``call_tool`` was +in flight, wedging the stream so the in-flight call timed out. That timeout +triggered a false reconnect, and the SDK does not always fail the pending +``call_tool`` when its streams close — so its ``run_coroutine_threadsafe`` +future never resolved and the calling agent thread polled to the full +``tool_timeout`` (up to hours). + +The fix: + * the keepalive skips a cycle when a call is in flight (and otherwise runs + under the same ``_rpc_lock`` tool calls use, so the two can't overlap); + * a reconnect/shutdown teardown calls ``_fail_inflight_calls`` to cancel the + pending call tasks; and + * ``_call`` converts that deliberate cancellation into a clean, retryable + error so the agent recovers on the freshly rebuilt session. + +These tests exercise the in-flight bookkeeping and the teardown behavior +directly (no live MCP server required). +""" + +from __future__ import annotations + +import asyncio + + +def test_new_server_starts_with_empty_inflight_state(): + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("init-test") + assert server._inflight_tasks == set() + assert server._reconnecting is False + + +def test_fail_inflight_calls_is_noop_when_nothing_in_flight(): + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("noop-test") + # No in-flight tasks: must not flip the teardown flag (so a later genuine + # cancel isn't misread as a deliberate reconnect). + server._fail_inflight_calls("reconnect") + assert server._reconnecting is False + + +def test_fail_inflight_calls_cancels_pending_and_flags_teardown(): + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("cancel-test") + + async def drive(): + async def _long(): + await asyncio.sleep(3600) + + task = asyncio.create_task(_long()) + server._inflight_tasks.add(task) + await asyncio.sleep(0) # let the task start + + server._fail_inflight_calls("reconnect") + assert server._reconnecting is True + + # The pending task must have been cancelled. + try: + await asyncio.wait_for(task, timeout=1.0) + except asyncio.CancelledError: + return "cancelled" + except asyncio.TimeoutError: + return "still_running" + return "completed" + + assert asyncio.run(drive()) == "cancelled" + + +def test_inflight_task_tracking_add_and_discard(): + """The in-flight set tracks a running task and discards it on completion, + mirroring the add/finally-discard bookkeeping in ``_call``.""" + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("track-test") + + async def drive(): + async def _work(): + task = asyncio.current_task() + server._inflight_tasks.add(task) + try: + assert task in server._inflight_tasks + finally: + server._inflight_tasks.discard(task) + + await asyncio.create_task(_work()) + return server._inflight_tasks + + assert asyncio.run(drive()) == set() + + +def test_reconnecting_flag_distinguishes_deliberate_teardown(): + """``_reconnecting`` is the signal ``_call`` reads to convert a cancellation + into a retryable error vs. re-raising a genuine (external) cancel.""" + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("flag-test") + assert server._reconnecting is False + # Simulate what a teardown does when there IS an in-flight task. + + async def drive(): + async def _long(): + await asyncio.sleep(3600) + + task = asyncio.create_task(_long()) + server._inflight_tasks.add(task) + await asyncio.sleep(0) + server._fail_inflight_calls("shutdown") + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + return server._reconnecting + + assert asyncio.run(drive()) is True diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index db419196a471..bf03553757f1 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -1161,6 +1161,7 @@ class MCPServerTask: "_tools", "_error", "_config", "_sampling", "_registered_tool_names", "_auth_type", "_refresh_lock", "_rpc_lock", "_pending_refresh_tasks", + "_inflight_tasks", "_reconnecting", "initialize_result", ) @@ -1192,6 +1193,16 @@ def __init__(self, name: str): # transports for conservative per-server ordering. self._rpc_lock = asyncio.Lock() self._pending_refresh_tasks: set[asyncio.Task] = set() + # In-flight tool-call tasks (asyncio.Task running session.call_tool on + # the MCP loop). Tracked so a reconnect/shutdown can FAIL them cleanly + # instead of orphaning their run_coroutine_threadsafe futures. An + # orphaned future makes the calling agent thread poll to the full + # tool_timeout (hours). Also used to suppress the keepalive while a call + # is active (a busy server is provably alive). Single-loop access, so + # no lock needed. ``_reconnecting`` flags a deliberate teardown so the + # cancelled call surfaces a retryable error rather than a raw cancel. + self._inflight_tasks: set = set() + self._reconnecting: bool = False # Captures the ``InitializeResult`` returned by # ``await session.initialize()`` so downstream code can inspect the # server's real advertised capabilities (``.capabilities.resources``, @@ -1392,17 +1403,29 @@ async def _wait_for_lifecycle_event(self) -> str: # ``ping`` request for them instead — otherwise every # keepalive cycle would trigger a spurious reconnect. if self.session: + # CRITICAL: never keepalive while a tool call is in flight. + # The stdio transport is a SINGLE JSON-RPC stream; a + # concurrent list_tools/ping wedges the in-flight call, + # which then times out -> false reconnect -> the call is + # orphaned and the agent hangs to tool_timeout (root cause + # of multi-thousand-second hangs). A server actively serving + # a call is provably alive, so skip this cycle. We also wrap + # the keepalive in the SAME _rpc_lock that tool calls use, so + # a call starting concurrently can't overlap the keepalive. + if self._rpc_lock.locked() or self._inflight_tasks: + continue try: - if self._advertises_tools(): - await asyncio.wait_for( - self.session.list_tools(), - timeout=30.0, - ) - else: - await asyncio.wait_for( - self.session.send_ping(), - timeout=30.0, - ) + async with self._rpc_lock: + if self._advertises_tools(): + await asyncio.wait_for( + self.session.list_tools(), + timeout=30.0, + ) + else: + await asyncio.wait_for( + self.session.send_ping(), + timeout=30.0, + ) except Exception as exc: logger.warning( "MCP server '%s' keepalive failed, " @@ -1421,10 +1444,37 @@ async def _wait_for_lifecycle_event(self) -> str: pass if self._shutdown_event.is_set(): + self._fail_inflight_calls("shutdown") return "shutdown" self._reconnect_event.clear() + self._fail_inflight_calls("reconnect") return "reconnect" + def _fail_inflight_calls(self, reason: str) -> None: + """Cancel in-flight tool-call tasks before the session is torn down. + + The MCP session is about to close (reconnect/shutdown). Any pending + ``session.call_tool`` await would otherwise be orphaned: the SDK does + not always fail the call when its streams close, so the + ``run_coroutine_threadsafe`` future never resolves and the calling + agent thread polls to the full ``tool_timeout`` (up to hours). We flag + a deliberate teardown and cancel the tasks; ``_call`` converts that + cancellation into a clean, retryable error so the agent recovers and + the next call runs on the freshly rebuilt session (self-healing). + Runs on the MCP event loop, same as the call tasks, so no lock needed. + """ + if not self._inflight_tasks: + return + self._reconnecting = True + pending = [t for t in self._inflight_tasks if not t.done()] + if pending: + logger.warning( + "MCP server '%s': failing %d in-flight call(s) due to %s", + self.name, len(pending), reason, + ) + for task in pending: + task.cancel() + async def _run_stdio(self, config: dict): """Run the server using stdio transport.""" if not _MCP_AVAILABLE: @@ -2807,8 +2857,26 @@ def _handler(args: dict, **kwargs) -> str: }, ensure_ascii=False) async def _call(): - async with server._rpc_lock: - result = await server.session.call_tool(tool_name, arguments=args) + task = asyncio.current_task() + if task is not None: + server._inflight_tasks.add(task) + try: + async with server._rpc_lock: + result = await server.session.call_tool(tool_name, arguments=args) + except asyncio.CancelledError: + # A deliberate reconnect/shutdown teardown cancelled us + # (see _fail_inflight_calls). Convert to a clean, retryable + # error instead of propagating a raw cancellation, so the agent + # then retries on the freshly rebuilt session. + if getattr(server, "_reconnecting", False): + raise RuntimeError( + f"MCP server '{server_name}' reconnected during the " + f"call (transport reset); retry the tool." + ) from None + raise + finally: + if task is not None: + server._inflight_tasks.discard(task) # MCP CallToolResult has .content (list of content blocks) and .isError if result.isError: error_text = "" From 4a1fbe9e1651e22d19b93e939586a4164a67db89 Mon Sep 17 00:00:00 2001 From: arminanton <29869547+arminanton@users.noreply.github.com> Date: Sun, 21 Jun 2026 03:16:37 -0700 Subject: [PATCH 2/2] fix(mcp): reset _reconnecting flag on entry to healthy keepalive wait After a reconnect cycle, _reconnecting stayed True, so the next legitimate in-flight call could be falsely converted to a 'reconnected, retry' error. Clear the deliberate-teardown flag when entering a healthy wait state (session established + ready), matching the _reconnecting=True set on teardown. Picks up a refinement that post-dated the original branch. --- tools/mcp_tool.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 2ab02333987c..9d6829477949 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -1707,6 +1707,12 @@ async def _wait_for_lifecycle_event(self) -> str: float(self._config.get("keepalive_interval", _DEFAULT_KEEPALIVE_INTERVAL)), ) + # Entering a healthy wait state means the session is established and + # ready, so clear any lingering "deliberate teardown" flag from a prior + # cycle. New in-flight calls on this fresh session must not be treated + # as reconnect casualties. + self._reconnecting = False + shutdown_task = asyncio.create_task(self._shutdown_event.wait()) reconnect_task = asyncio.create_task(self._reconnect_event.wait()) try: