From 684ebb68fd30f5055df796570bf486247e933976 Mon Sep 17 00:00:00 2001 From: embwl0x Date: Sat, 8 Aug 2026 17:25:05 -0600 Subject: [PATCH] fix(mcp): fail RPCs when transport sessions restart --- tests/tools/test_mcp_session_invalidation.py | 247 +++++++++++++++++++ tools/mcp_tool.py | 145 +++++++++-- 2 files changed, 372 insertions(+), 20 deletions(-) create mode 100644 tests/tools/test_mcp_session_invalidation.py diff --git a/tests/tools/test_mcp_session_invalidation.py b/tests/tools/test_mcp_session_invalidation.py new file mode 100644 index 000000000000..7a29c65e7924 --- /dev/null +++ b/tests/tools/test_mcp_session_invalidation.py @@ -0,0 +1,247 @@ +"""Regression coverage for RPC ownership across MCP transport rebuilds.""" + +import asyncio +import json +import threading +import time +from types import SimpleNamespace + +import pytest + + +@pytest.mark.asyncio +async def test_invalidating_session_cancels_owned_rpc() -> None: + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("generation-unit") + started = asyncio.Event() + cancelled = asyncio.Event() + + class HangingSession: + async def call_tool(self): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + old_session = HangingSession() + server.session = old_session + call = asyncio.create_task( + server._call_session_rpc( + "tools/call probe", lambda session: session.call_tool() + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + server._invalidate_session() + + with pytest.raises(RuntimeError, match="restarted while tools/call probe"): + await asyncio.wait_for(call, timeout=1) + assert cancelled.is_set() + assert server.session is None + + replacement = object() + server.session = replacement + assert ( + await server._call_session_rpc( + "tools/call probe", lambda session: asyncio.sleep(0, result=session) + ) + is replacement + ) + + +@pytest.mark.asyncio +async def test_dynamic_tool_refresh_uses_session_invalidation() -> None: + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("generation-refresh") + started = asyncio.Event() + cancelled = asyncio.Event() + + class HangingSession: + async def list_tools(self, **_kwargs): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + server.session = HangingSession() + refresh = asyncio.create_task(server._refresh_tools()) + await asyncio.wait_for(started.wait(), timeout=1) + + server._invalidate_session() + + with pytest.raises(RuntimeError, match="restarted while tools/list refresh"): + await asyncio.wait_for(refresh, timeout=1) + assert cancelled.is_set() + + +@pytest.mark.asyncio +async def test_invalidated_session_rejects_new_rpc_before_transport_clears() -> None: + from tools.mcp_tool import MCPServerTask + + server = MCPServerTask("generation-teardown-window") + session = object() + invoked = False + server.session = session + server._invalidate_session(clear=False) + + async def call(_session): + nonlocal invoked + invoked = True + + with pytest.raises( + RuntimeError, match="restarted while resources/read was starting" + ): + await server._call_session_rpc("resources/read", call) + assert server.session is session + assert invoked is False + + +def test_tool_handler_fails_fast_when_lifecycle_replaces_session() -> None: + from tools import mcp_tool + + mcp_tool._ensure_mcp_loop() + initial_ready = threading.Event() + call_started = threading.Event() + call_cancelled = threading.Event() + replacement_ready = threading.Event() + transport_count = 0 + + class HangingSession: + async def call_tool(self, *_args, **_kwargs): + call_started.set() + try: + await asyncio.Event().wait() + finally: + call_cancelled.set() + + class HealthySession: + async def call_tool(self, *_args, **_kwargs): + return SimpleNamespace( + isError=False, + content=[SimpleNamespace(type="text", text="healthy")], + structuredContent=None, + ) + + class LifecycleServer(mcp_tool.MCPServerTask): + async def _run_stdio(self, config: dict): + nonlocal transport_count + assert config["command"] == "unused" + transport_count += 1 + self.session = ( + HangingSession() if transport_count == 1 else HealthySession() + ) + self._ready.set() + if transport_count == 1: + initial_ready.set() + else: + replacement_ready.set() + return await self._wait_for_lifecycle_event() + + server = LifecycleServer("generation-integration") + mcp_tool._servers[server.name] = server + mcp_tool._server_error_counts.pop(server.name, None) + mcp_tool._server_breaker_opened_at.pop(server.name, None) + loop = mcp_tool._mcp_loop + assert loop is not None + run_future = asyncio.run_coroutine_threadsafe( + server.run({"command": "unused"}), loop + ) + + def reconnect_after_call_starts() -> None: + assert call_started.wait(2) + assert mcp_tool._signal_reconnect(server) + + reconnect_thread = threading.Thread(target=reconnect_after_call_starts) + try: + assert initial_ready.wait(2) + reconnect_thread.start() + handler = mcp_tool._make_tool_handler(server.name, "probe", 5.0) + started_at = time.monotonic() + result = json.loads(handler({})) + elapsed = time.monotonic() - started_at + reconnect_thread.join(timeout=2) + + assert elapsed < 2 + assert "restarted while tools/call probe" in result["error"] + assert call_cancelled.wait(1) + assert replacement_ready.wait(2) + assert json.loads(handler({})) == {"result": "healthy"} + finally: + if reconnect_thread.is_alive(): + reconnect_thread.join(timeout=2) + loop.call_soon_threadsafe(server._shutdown_event.set) + run_future.result(timeout=5) + mcp_tool._servers.pop(server.name, None) + mcp_tool._server_error_counts.pop(server.name, None) + mcp_tool._server_breaker_opened_at.pop(server.name, None) + + +@pytest.mark.parametrize( + ("factory_name", "arguments", "session_method", "operation"), + [ + ("_make_list_resources_handler", {}, "list_resources", "resources/list"), + ( + "_make_read_resource_handler", + {"uri": "file:///report.txt"}, + "read_resource", + "resources/read", + ), + ("_make_list_prompts_handler", {}, "list_prompts", "prompts/list"), + ( + "_make_get_prompt_handler", + {"name": "review"}, + "get_prompt", + "prompts/get", + ), + ], +) +def test_utility_handler_fails_fast_when_session_is_invalidated( + factory_name: str, + arguments: dict, + session_method: str, + operation: str, +) -> None: + from tools import mcp_tool + + mcp_tool._ensure_mcp_loop() + call_started = threading.Event() + + class HangingSession: + pass + + async def hang(*_args, **_kwargs): + call_started.set() + await asyncio.Event().wait() + + session = HangingSession() + setattr(session, session_method, hang) + server = mcp_tool.MCPServerTask(f"generation-{session_method}") + server.session = session + server._ready.set() + mcp_tool._servers[server.name] = server + loop = mcp_tool._mcp_loop + assert loop is not None + + def invalidate_after_call_starts() -> None: + assert call_started.wait(2) + loop.call_soon_threadsafe(server._invalidate_session) + + invalidator = threading.Thread(target=invalidate_after_call_starts) + try: + invalidator.start() + handler = getattr(mcp_tool, factory_name)(server.name, 5.0) + started_at = time.monotonic() + result = json.loads(handler(arguments)) + elapsed = time.monotonic() - started_at + invalidator.join(timeout=2) + + assert elapsed < 2 + assert f"restarted while {operation}" in result["error"] + finally: + if invalidator.is_alive(): + invalidator.join(timeout=2) + mcp_tool._servers.pop(server.name, None) diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 9ff5a1dea7b9..20ba7d366178 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -1958,6 +1958,7 @@ class MCPServerTask: "_registered_tool_names", "_auth_type", "_refresh_lock", "_rpc_lock", "_pending_refresh_tasks", "_pending_call_context", + "_session_invalidation_event", "_session_invalidation_owner", "_lifecycle_started_at", "_last_tool_call_at", "_idle_timeout_seconds", "_max_lifetime_seconds", "_recycled_reason", "initialize_result", "_ping_unsupported", @@ -2016,6 +2017,12 @@ def __init__(self, name: str): # gateway-platform attribution and routes the approval prompt # to the right surface (Telegram, Slack, etc.). self._pending_call_context: Optional[contextvars.Context] = None + # A handler can still be awaiting an RPC when keepalive or another + # lifecycle signal tears down the transport. Keep one event per live + # ClientSession so those calls fail as soon as their owning session is + # invalidated instead of waiting for the full tool timeout (#81995). + self._session_invalidation_event = asyncio.Event() + self._session_invalidation_owner: Optional[Any] = None now = time.monotonic() self._lifecycle_started_at: float = now self._last_tool_call_at: float = now @@ -2104,7 +2111,61 @@ def _next_stdio_recycle_deadline(self) -> Optional[float]: def _mark_stdio_recycled(self, reason: str) -> None: """Mark a stdio session dormant before its transport finishes closing.""" self._recycled_reason = reason - self.session = None + self._invalidate_session() + + def _capture_session_for_rpc(self) -> tuple[Any, asyncio.Event]: + """Return the current session and its lifecycle invalidation event.""" + session = self.session + if session is None: + raise RuntimeError(f"MCP server '{self.name}' is not connected") + if self._session_invalidation_owner is not session: + self._session_invalidation_event = asyncio.Event() + self._session_invalidation_owner = session + return session, self._session_invalidation_event + + def _invalidate_session(self, *, clear: bool = True) -> None: + """Wake RPCs owned by the current transport, optionally clearing it.""" + if ( + self.session is not None + and self._session_invalidation_owner is not self.session + ): + self._session_invalidation_event = asyncio.Event() + self._session_invalidation_owner = self.session + self._session_invalidation_event.set() + if clear: + self._session_invalidation_owner = None + self.session = None + + async def _call_session_rpc(self, operation: str, call) -> Any: + """Run one RPC until it completes or its transport is replaced.""" + session, invalidated = self._capture_session_for_rpc() + if invalidated.is_set(): + raise RuntimeError( + f"MCP server '{self.name}' restarted while {operation} was " + "starting; retry the call on the replacement session." + ) + rpc_task = asyncio.ensure_future(call(session)) + invalidated_task = asyncio.create_task(invalidated.wait()) + try: + done, _pending = await asyncio.wait( + {rpc_task, invalidated_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + # Prefer a response that completed in the same loop turn as the + # lifecycle signal; the RPC result is already authoritative. + if rpc_task in done: + return await rpc_task + raise RuntimeError( + f"MCP server '{self.name}' restarted while {operation} was " + "in flight; retry the call on the replacement session." + ) + finally: + for task in (rpc_task, invalidated_task): + if not task.done(): + task.cancel() + await asyncio.gather( + rpc_task, invalidated_task, return_exceptions=True + ) # ----- Dynamic tool discovery (notifications/tools/list_changed) ----- @@ -2223,8 +2284,11 @@ async def _refresh_tools(self): # 1. Fetch current tool list from server (follow nextCursor) async with self._rpc_lock: - new_mcp_tools = await _paginate_full_list( - self.session.list_tools, "tools", self.name + new_mcp_tools = await self._call_session_rpc( + "tools/list refresh", + lambda session: _paginate_full_list( + session.list_tools, "tools", self.name + ), ) # 2. Re-register with fresh tool list. Avoid nuke-and-repave for @@ -2449,8 +2513,10 @@ async def _wait_for_lifecycle_event(self) -> str: pass if self._shutdown_event.is_set(): + self._invalidate_session(clear=False) return "shutdown" self._reconnect_event.clear() + self._invalidate_session(clear=False) return "reconnect" async def _wait_for_reconnect_or_shutdown( @@ -3267,10 +3333,16 @@ async def run(self, config: dict): while True: try: - if self._is_http(): - lifecycle_reason = await self._run_http(config) - else: - lifecycle_reason = await self._run_stdio(config) + try: + if self._is_http(): + lifecycle_reason = await self._run_http(config) + else: + lifecycle_reason = await self._run_stdio(config) + finally: + # Transport context managers can also exit via exceptions, + # bypassing _wait_for_lifecycle_event. Always release RPCs + # tied to that generation before the reconnect loop moves on. + self._invalidate_session() # Transport returned cleanly. Two cases: # - _shutdown_event was set: exit the run loop entirely. # - _reconnect_event was set (auth recovery): loop back and @@ -3284,7 +3356,7 @@ async def run(self, config: dict): "waiting for lazy reconnect", self.name, self._recycled_reason, ) - self.session = None + self._invalidate_session() await self._wait_for_lazy_reconnect() if self._shutdown_event.is_set(): break @@ -3346,7 +3418,7 @@ async def run(self, config: dict): # _ready set here lets handler-side recovery mistake the stale # pre-reconnect session for a fresh one and retry too early. self._ready.clear() - self.session = None + self._invalidate_session() continue except asyncio.CancelledError: # Task was cancelled (shutdown, gateway restart, explicit @@ -3358,10 +3430,10 @@ async def run(self, config: dict): # restarted. Re-raise so the task's cancellation propagates # correctly to asyncio's task machinery and ``shutdown()``'s # ``await self._task`` completes. See #9930. - self.session = None + self._invalidate_session() raise except Exception as exc: - self.session = None + self._invalidate_session() # Unwrap anyio TaskGroup wrappers first: str(exc) on a # BaseExceptionGroup is "unhandled errors in a TaskGroup # (N sub-exceptions)" — useless in logs, and it hides the @@ -3581,7 +3653,7 @@ async def run(self, config: dict): if self._shutdown_event.is_set(): return finally: - self.session = None + self._invalidate_session() async def start(self, config: dict): """Create the background Task and wait until ready (or failed).""" @@ -3630,7 +3702,7 @@ async def shutdown(self): await asyncio.gather(*self._pending_refresh_tasks, return_exceptions=True) self._pending_refresh_tasks.clear() self._deregister_tools() - self.session = None + self._invalidate_session() def _deregister_tools(self) -> None: """Drop this server's tools from the global registry (idempotent). @@ -5183,6 +5255,15 @@ def _mark_server_call_started(server: Any) -> None: mark_tool_call() +async def _call_server_session_rpc(server: Any, operation: str, call) -> Any: + """Run an RPC with lifecycle invalidation for managed server tasks.""" + if isinstance(server, MCPServerTask): + return await server._call_session_rpc(operation, call) + # Compatibility for lightweight third-party/test server adapters that + # predate MCPServerTask's lifecycle generation tracking. + return await call(server.session) + + def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): """Return a sync handler that calls an MCP tool via the background loop. @@ -5267,7 +5348,13 @@ async def _call(): # it and detect the gateway platform / session for routing. server._pending_call_context = contextvars.copy_context() try: - result = await server.session.call_tool(tool_name, arguments=args) + result = await _call_server_session_rpc( + server, + f"tools/call {tool_name}", + lambda session: session.call_tool( + tool_name, arguments=args + ), + ) finally: server._pending_call_context = None # The RPC round-trip completed — the session is demonstrably @@ -5417,8 +5504,12 @@ def _handler(args: dict, **kwargs) -> str: async def _call(): _mark_server_call_started(server) async with server._rpc_lock: - all_resources = await _paginate_full_list( - server.session.list_resources, "resources", server_name + all_resources = await _call_server_session_rpc( + server, + "resources/list", + lambda session: _paginate_full_list( + session.list_resources, "resources", server_name + ), ) resources = [] for r in all_resources: @@ -5477,7 +5568,11 @@ def _handler(args: dict, **kwargs) -> str: async def _call(): _mark_server_call_started(server) async with server._rpc_lock: - result = await server.session.read_resource(uri) + result = await _call_server_session_rpc( + server, + "resources/read", + lambda session: session.read_resource(uri), + ) # read_resource returns ReadResourceResult with .contents list parts: List[str] = [] contents = result.contents if hasattr(result, "contents") else [] @@ -5534,8 +5629,12 @@ def _handler(args: dict, **kwargs) -> str: async def _call(): _mark_server_call_started(server) async with server._rpc_lock: - all_prompts = await _paginate_full_list( - server.session.list_prompts, "prompts", server_name + all_prompts = await _call_server_session_rpc( + server, + "prompts/list", + lambda session: _paginate_full_list( + session.list_prompts, "prompts", server_name + ), ) prompts = [] for p in all_prompts: @@ -5600,7 +5699,13 @@ def _handler(args: dict, **kwargs) -> str: async def _call(): _mark_server_call_started(server) async with server._rpc_lock: - result = await server.session.get_prompt(name, arguments=arguments) + result = await _call_server_session_rpc( + server, + "prompts/get", + lambda session: session.get_prompt( + name, arguments=arguments + ), + ) # GetPromptResult has .messages list messages = [] for msg in (result.messages if hasattr(result, "messages") else []):