From ea0b42c43ac35027663035deabafb2bc41878627 Mon Sep 17 00:00:00 2001 From: harjoth Date: Mon, 8 Jun 2026 11:37:23 -0700 Subject: [PATCH] Handle minimal MCP server fakes --- tests/tools/test_mcp_circuit_breaker.py | 7 +++++++ tools/mcp_tool.py | 17 ++++++++++++----- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/tools/test_mcp_circuit_breaker.py b/tests/tools/test_mcp_circuit_breaker.py index 116af9397cd..f356c94e128 100644 --- a/tests/tools/test_mcp_circuit_breaker.py +++ b/tests/tools/test_mcp_circuit_breaker.py @@ -74,6 +74,13 @@ def _install_stub_server(mcp_tool_module, name: str, call_tool_impl): server._reconnect_event = _ReconnectAdapter() server._ready = _ReadyAdapter() + # A bare MagicMock returns a truthy Mock for every method, so + # ``_is_recycled_stdio()`` would spuriously report this stub as a recycled + # stdio server and divert dead-session tool calls into the lazy-reconnect + # wait (which polls the test-frozen ``time.monotonic`` forever). Real + # non-recycled servers return False here; make the stub faithful so the + # dead-session path falls through to the graceful reconnect handler. + server._is_recycled_stdio.return_value = False mcp_tool_module._servers[name] = server mcp_tool_module._server_error_counts.pop(name, None) diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 114fc172fc1..cd999f2712f 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -3849,6 +3849,13 @@ def _get_connected_server_for_call(server_name: str) -> Optional[MCPServerTask]: return server +def _mark_server_call_started(server: Any) -> None: + """Record a user-visible MCP operation when the server supports it.""" + mark_tool_call = getattr(server, "mark_tool_call", None) + if callable(mark_tool_call): + mark_tool_call() + + 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. @@ -3925,7 +3932,7 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): }, ensure_ascii=False) async def _call(): - server.mark_tool_call() + _mark_server_call_started(server) async with server._rpc_lock: # Snapshot the agent's context so an elicitation callback # triggered during this call (fired on the MCP recv loop @@ -4046,7 +4053,7 @@ def _make_list_resources_handler(server_name: str, tool_timeout: float): }, ensure_ascii=False) async def _call(): - server.mark_tool_call() + _mark_server_call_started(server) async with server._rpc_lock: result = await server.session.list_resources() resources = [] @@ -4110,7 +4117,7 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float): return tool_error("Missing required parameter 'uri'") async def _call(): - server.mark_tool_call() + _mark_server_call_started(server) async with server._rpc_lock: result = await server.session.read_resource(uri) # read_resource returns ReadResourceResult with .contents list @@ -4164,7 +4171,7 @@ def _make_list_prompts_handler(server_name: str, tool_timeout: float): }, ensure_ascii=False) async def _call(): - server.mark_tool_call() + _mark_server_call_started(server) async with server._rpc_lock: result = await server.session.list_prompts() prompts = [] @@ -4234,7 +4241,7 @@ def _make_get_prompt_handler(server_name: str, tool_timeout: float): arguments = args.get("arguments", {}) async def _call(): - server.mark_tool_call() + _mark_server_call_started(server) async with server._rpc_lock: result = await server.session.get_prompt(name, arguments=arguments) # GetPromptResult has .messages list