From 8fda0ee69c8b0a9a83f1a3eccc9f33a5da7cebc0 Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Thu, 18 Jun 2026 13:32:10 +0200 Subject: [PATCH] Fix MCP timeouts from hanging agent loop Release MCP config locks before awaited initialization or tool calls, isolate MCP session operations in disposable DeferredTask workers, and bound session cleanup so wedged transports cannot freeze later agent work. Add deterministic MCP regression coverage for lock scope, config update initialization, cleanup timeouts, and isolated operation timeouts. Update the helper DOX contract for the new concurrency and cleanup behavior. --- helpers/mcp_handler.py | 204 +++++++++++++++++---------- helpers/mcp_handler.py.dox.md | 9 +- tests/test_mcp_handler_multimodal.py | 132 +++++++++++++++++ 3 files changed, 264 insertions(+), 81 deletions(-) diff --git a/helpers/mcp_handler.py b/helpers/mcp_handler.py index 6aadfc80c..2b4e0d8b9 100644 --- a/helpers/mcp_handler.py +++ b/helpers/mcp_handler.py @@ -43,10 +43,13 @@ from pydantic import BaseModel, Field, Discriminator, Tag, PrivateAttr from helpers import dirty_json, media_artifacts from helpers.print_style import PrintStyle from helpers.tool import Tool, Response +from helpers.defer import DeferredTask MCP_MEDIA_TOKENS_ESTIMATE = 1500 MAX_MCP_RESOURCE_TEXT_CHARS = 12_000 +MCP_SESSION_CLEANUP_TIMEOUT_SECONDS = 5.0 +MCP_OPERATION_TIMEOUT_GRACE_SECONDS = MCP_SESSION_CLEANUP_TIMEOUT_SECONDS + 2.0 DEFAULT_MCP_SERVERS_CONFIG = '{\n "mcpServers": {}\n}' @@ -661,10 +664,10 @@ class MCPConfig(BaseModel): @classmethod def get_instance(cls) -> "MCPConfig": - # with cls.__lock: - if cls.__instance is None: - cls.__instance = cls(servers_list=[], config_scope="global") - return cls.__instance + with cls.__lock: + if cls.__instance is None: + cls.__instance = cls(servers_list=[], config_scope="global") + return cls.__instance @classmethod def clear_project_instances(cls): @@ -791,35 +794,20 @@ class MCPConfig(BaseModel): @classmethod def update(cls, config_str: str) -> Any: + servers_data = cls.parse_config_string(config_str) + new_instance = cls(servers_list=servers_data, config_scope="global") with cls.__lock: - servers_data = cls.parse_config_string(config_str) - - # Initialize/update the singleton instance with the (potentially empty) list of server data - instance = cls.get_instance() - # Directly update the servers attribute of the existing instance or re-initialize carefully - # For simplicity and to ensure __init__ logic runs if needed for setup: - new_instance_data = { - "servers": servers_data - } # Prepare data for re-initialization or update - - # Option 1: Re-initialize the existing instance (if __init__ is idempotent for other fields) - instance.__init__(servers_list=servers_data, config_scope="global") + # Build and initialize outside the class lock so a slow or wedged MCP + # server cannot freeze status reads, prompts, or later tool calls. + instance = cls.__instance + if instance is None: + instance = new_instance + cls.__instance = instance + else: + instance.servers = new_instance.servers + instance.disconnected_servers = new_instance.disconnected_servers + instance.config_scope = new_instance.config_scope cls.__project_instances = {} - - # Option 2: Or, if __init__ has side effects we don't want to repeat, - # and 'servers' is the primary thing 'update' changes: - # instance.servers = [] # Clear existing servers first - # for server_item_data in servers_data: - # try: - # if server_item_data.get("url", None): - # instance.servers.append(MCPServerRemote(server_item_data)) - # else: - # instance.servers.append(MCPServerLocal(server_item_data)) - # except Exception as e_init: - # PrintStyle(background_color="grey", font_color="red", padding=True).print( - # f"MCPConfig.update: Failed to create MCPServer from item '{server_item_data.get('name', 'Unknown')}': {e_init}" - # ) - cls.__initialized = True return instance @@ -1146,11 +1134,15 @@ class MCPConfig(BaseModel): ) -> CallToolResult: """Call a tool with the given input data""" server_name_part, tool_name_part = _split_qualified_tool_name(tool_name) + matched_server = None with self.__lock: for server in self.servers: if server.name == server_name_part and server.has_tool(tool_name_part): - return await server.call_tool(tool_name_part, input_data) + matched_server = server + break + if matched_server is None: raise ValueError(f"Tool {tool_name} not found") + return await matched_server.call_tool(tool_name_part, input_data) T = TypeVar("T") @@ -1170,6 +1162,49 @@ class MCPClientBase(ABC): self.log: List[str] = [] self.log_file: Optional[TextIO] = None + def _operation_timeout_seconds(self, read_timeout_seconds: float) -> float: + try: + seconds = float(read_timeout_seconds) + except (TypeError, ValueError): + seconds = 60.0 + if seconds <= 0: + seconds = 60.0 + return seconds + MCP_OPERATION_TIMEOUT_GRACE_SECONDS + + def _operation_thread_name(self, operation_name: str) -> str: + server_name = normalize_name(str(getattr(self.server, "name", "") or "server")) + return f"MCPClient-{server_name[:32] or 'server'}-{operation_name}-{uuid.uuid4().hex[:8]}" + + async def _run_isolated_operation( + self, + operation_name: str, + operation: Callable[[], Awaitable[T]], + timeout_seconds: float, + ) -> T: + worker = DeferredTask(thread_name=self._operation_thread_name(operation_name)) + timed_out = False + try: + return await asyncio.wait_for( + worker.execute_inside(operation), + timeout=timeout_seconds, + ) + except asyncio.TimeoutError as exc: + timed_out = True + message = ( + f"MCPClientBase ({self.server.name} - {operation_name}): " + f"operation did not finish within {timeout_seconds:.1f}s; " + "abandoning the isolated worker so Agent Zero can continue." + ) + PrintStyle.warning(message) + with self.__lock: + self.error = message + raise TimeoutError(message) from exc + finally: + if timed_out: + worker.kill(terminate_thread=False) + else: + worker.kill(terminate_thread=True) + # Protected method @abstractmethod async def _create_stdio_transport( @@ -1192,54 +1227,56 @@ class MCPClientBase(ABC): """ operation_name = coro_func.__name__ # For logging # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Creating new session for operation '{operation_name}'...") - # Store the original exception outside the async block original_exception = None + result: T | None = None + has_result = False + temp_stack = AsyncExitStack() try: - async with AsyncExitStack() as temp_stack: - try: + stdio, write = await self._create_stdio_transport(temp_stack) + # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Transport created. Initializing session...") + session = await temp_stack.enter_async_context( + ClientSession( + stdio, # type: ignore + write, # type: ignore + read_timeout_seconds=timedelta( + seconds=read_timeout_seconds + ), + ) + ) + await session.initialize() - stdio, write = await self._create_stdio_transport(temp_stack) - # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Transport created. Initializing session...") - session = await temp_stack.enter_async_context( - ClientSession( - stdio, # type: ignore - write, # type: ignore - read_timeout_seconds=timedelta( - seconds=read_timeout_seconds - ), - ) - ) - await session.initialize() - - result = await coro_func(session) - - return result - except Exception as e: - # Store the original exception and raise a dummy exception - excs = getattr(e, "exceptions", None) # Python 3.11+ ExceptionGroup - if excs: - original_exception = excs[0] - else: - original_exception = e - # Create a dummy exception to break out of the async block - raise RuntimeError("Dummy exception to break out of async block") + result = await coro_func(session) + has_result = True except Exception as e: - # Check if this is our dummy exception - if original_exception is not None: - e = original_exception - # We have the original exception stored + excs = getattr(e, "exceptions", None) # Python 3.11+ ExceptionGroup + if excs: + original_exception = excs[0] + else: + original_exception = e + try: + await asyncio.wait_for( + temp_stack.aclose(), + timeout=MCP_SESSION_CLEANUP_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + PrintStyle.warning( + f"MCPClientBase ({self.server.name} - {operation_name}): " + f"session cleanup exceeded {MCP_SESSION_CLEANUP_TIMEOUT_SECONDS:.1f}s." + ) + except Exception as cleanup_exception: + PrintStyle.warning( + f"MCPClientBase ({self.server.name} - {operation_name}): " + f"session cleanup failed: {type(cleanup_exception).__name__}: {cleanup_exception}" + ) + if original_exception is not None: PrintStyle( background_color="#AA4455", font_color="white", padding=False ).print( - f"MCPClientBase ({self.server.name} - {operation_name}): Error during operation: {type(e).__name__}: {e}" + f"MCPClientBase ({self.server.name} - {operation_name}): Error during operation: {type(original_exception).__name__}: {original_exception}" ) - raise e # Re-raise the original exception - # finally: - # PrintStyle(font_color="cyan").print( - # f"MCPClientBase ({self.server.name} - {operation_name}): Session and transport will be closed by AsyncExitStack." - # ) - # This line should ideally be unreachable if the try/except/finally logic within the 'async with' is exhaustive. - # Adding it to satisfy linters that might not fully trace the raise/return paths through async context managers. + raise original_exception + if has_result: + return cast(T, result) raise RuntimeError( f"MCPClientBase ({self.server.name} - {operation_name}): _execute_with_session exited 'async with' block unexpectedly." ) @@ -1270,9 +1307,13 @@ class MCPClientBase(ABC): or current_settings.get("mcp_client_init_timeout", 10) or 10 ) - await self._execute_with_session( - list_tools_op, - read_timeout_seconds=init_timeout, + await self._run_isolated_operation( + "update_tools", + lambda: self._execute_with_session( + list_tools_op, + read_timeout_seconds=init_timeout, + ), + timeout_seconds=self._operation_timeout_seconds(init_timeout), ) except Exception as e: # e = eg.exceptions[0] @@ -1339,10 +1380,17 @@ class MCPClientBase(ABC): return response try: - return await self._execute_with_session( - call_tool_op, - read_timeout_seconds=tool_timeout, + response = await self._run_isolated_operation( + "call_tool", + lambda: self._execute_with_session( + call_tool_op, + read_timeout_seconds=tool_timeout, + ), + timeout_seconds=self._operation_timeout_seconds(tool_timeout), ) + with self.__lock: + self.error = "" + return response except Exception as e: # Error logged by _execute_with_session. Re-raise a specific error for the caller. PrintStyle( diff --git a/helpers/mcp_handler.py.dox.md b/helpers/mcp_handler.py.dox.md index d645a0b16..906a6f7b1 100644 --- a/helpers/mcp_handler.py.dox.md +++ b/helpers/mcp_handler.py.dox.md @@ -67,7 +67,7 @@ - `_split_qualified_tool_name(tool_name: str) -> tuple[str, str]`: Split `server.tool` names while preserving dots inside MCP tool names. - `_normalize_disabled_tools(value: Any) -> list[str]`: Normalize the optional per-server disabled tool list. - `initialize_mcp(mcp_servers_config: str)` -- Notable constants/configuration names: `DEFAULT_MCP_SERVERS_CONFIG`, `MCP_MEDIA_TOKENS_ESTIMATE`, `MAX_MCP_RESOURCE_TEXT_CHARS`, `T`. +- Notable constants/configuration names: `DEFAULT_MCP_SERVERS_CONFIG`, `MCP_MEDIA_TOKENS_ESTIMATE`, `MAX_MCP_RESOURCE_TEXT_CHARS`, `MCP_SESSION_CLEANUP_TIMEOUT_SECONDS`, `MCP_OPERATION_TIMEOUT_GRACE_SECONDS`, `T`. ## Runtime Contracts @@ -81,13 +81,15 @@ - MCP tool names are qualified as `server_name.tool_name`; server names are normalized without dots, and the tool portion may contain dots. - Servers may define `disabled_tools` as a list of MCP tool names. Disabled tools are omitted from agent-facing prompts, status counts, `has_tool`, and calls, while detail views can still retrieve them through `get_all_tools()` with a `disabled` flag so users can re-enable them. - Server-specific `init_timeout` and `tool_timeout` override global MCP client timeout settings for list-tools and call-tool operations. +- MCP config locks must not be held across awaited server initialization or tool-call operations. Slow or wedged MCP servers must not block status reads, prompt construction, unrelated MCP servers, or later tool calls through the shared config lock. +- MCP client session work runs inside disposable isolated `DeferredTask` workers with an outer timeout. Normal `AsyncExitStack` cleanup is also bounded; if cleanup or transport shutdown does not finish, the operation reports failure or warning while Agent Zero keeps control of the agent loop. - Server status marks initialized server objects with cached initialization errors as disconnected, even if the config object exists. - Observed side-effect areas: filesystem writes, network calls, WebSocket state, settings/state persistence, secret handling. -- Imported dependency areas include: `abc`, `anyio.streams.memory`, `asyncio`, `contextlib`, `datetime`, `helpers`, `helpers.log`, `helpers.print_style`, `helpers.tool`, `httpx`, `json`, `mcp`, `mcp.client.sse`, `mcp.client.stdio`, `mcp.client.streamable_http`, `mcp.shared.message`. +- Imported dependency areas include: `abc`, `anyio.streams.memory`, `asyncio`, `contextlib`, `datetime`, `helpers`, `helpers.defer`, `helpers.log`, `helpers.print_style`, `helpers.tool`, `httpx`, `json`, `mcp`, `mcp.client.sse`, `mcp.client.stdio`, `mcp.client.streamable_http`, `mcp.shared.message`. ## Key Concepts -- Important called helpers/classes observed in the source: `TypeVar`, `name.strip.lower`, `re.sub`, `Field`, `PrivateAttr`, `threading.Lock`, `_split_qualified_tool_name`, `config_dict.lower`, `server_type.lower`, `MCPConfig.get_instance.is_initialized`, `MCPConfig.get_for_agent`, `projects.validate_project_name`, `projects.load_project_mcp_servers`, `settings.get_settings`, `self.agent.context.log.log`, `str.strip`, `media_artifacts.guess_extension`, `callable`, `self._content_item_dump`, `join`, `Response`, `self.get_log_object`, `self._raw_tool_response`, `additional.pop`, `self._coerce_media_token_estimate`. +- Important called helpers/classes observed in the source: `TypeVar`, `name.strip.lower`, `re.sub`, `Field`, `PrivateAttr`, `threading.Lock`, `DeferredTask`, `_split_qualified_tool_name`, `config_dict.lower`, `server_type.lower`, `MCPConfig.get_instance.is_initialized`, `MCPConfig.get_for_agent`, `projects.validate_project_name`, `projects.load_project_mcp_servers`, `settings.get_settings`, `self.agent.context.log.log`, `str.strip`, `media_artifacts.guess_extension`, `callable`, `self._content_item_dump`, `join`, `Response`, `self.get_log_object`, `self._raw_tool_response`, `additional.pop`, `self._coerce_media_token_estimate`. - Keep request/response, tool, or helper semantics documented here at the same time as source changes. ## Work Guidance @@ -95,6 +97,7 @@ - Preserve public helper APIs used by core code and plugins unless every caller is updated. - Keep path, auth, secret, persistence, network, and subprocess behavior explicit and bounded. - Prefer adding cohesive helper functions here only when behavior is reused across modules. +- Keep MCP timeout and cleanup changes covered by deterministic tests that do not require real MCP servers or network credentials. ## Verification diff --git a/tests/test_mcp_handler_multimodal.py b/tests/test_mcp_handler_multimodal.py index 421fa4224..040955e0e 100644 --- a/tests/test_mcp_handler_multimodal.py +++ b/tests/test_mcp_handler_multimodal.py @@ -54,6 +54,20 @@ class _FakeCallToolResult(SimpleNamespace): pass +class _TrackingLock: + def __init__(self): + self.held = False + + def __enter__(self): + assert self.held is False + self.held = True + return self + + def __exit__(self, exc_type, exc, tb): + self.held = False + return False + + @pytest.fixture def mcp_handler_module(monkeypatch, tmp_path): monkeypatch.delitem(sys.modules, "helpers.mcp_handler", raising=False) @@ -118,6 +132,10 @@ def mcp_handler_module(monkeypatch, tmp_path): def stream(self, *args, **kwargs): return self + @staticmethod + def warning(*args, **kwargs): + return None + def _fake_get_abs_path(*parts): return str(tmp_path.joinpath(*parts)) @@ -194,6 +212,57 @@ def test_mcp_config_preserves_dotted_tool_names(mcp_handler_module): assert called == [("alpha.beta", {"value": 7})] +def test_mcp_config_call_tool_releases_config_lock_before_await( + mcp_handler_module, monkeypatch +): + module, _tmp_path = mcp_handler_module + lock = _TrackingLock() + observed_lock_state: list[bool] = [] + + monkeypatch.setattr(module.MCPConfig, "_MCPConfig__lock", lock, raising=False) + + class _FakeServer: + name = "server" + description = "Fake MCP server" + type = "stdio" + scope = "global" + + def has_tool(self, tool_name): + return tool_name == "run" + + async def call_tool(self, tool_name, input_data): + observed_lock_state.append(lock.held) + await asyncio.sleep(0) + return _FakeCallToolResult(content=[], isError=False) + + config = module.MCPConfig(servers_list=[]) + config.servers = [_FakeServer()] + + asyncio.run(config.call_tool("server.run", {})) + + assert observed_lock_state == [False] + + +def test_mcp_config_update_initializes_outside_config_lock( + mcp_handler_module, monkeypatch +): + module, _tmp_path = mcp_handler_module + lock = _TrackingLock() + observed_lock_state: list[bool] = [] + original_init = module.MCPConfig.__init__ + + def tracking_init(self, *args, **kwargs): + observed_lock_state.append(lock.held) + original_init(self, *args, **kwargs) + + monkeypatch.setattr(module.MCPConfig, "_MCPConfig__lock", lock, raising=False) + monkeypatch.setattr(module.MCPConfig, "__init__", tracking_init) + + module.MCPConfig.update('{"mcpServers": {}}') + + assert observed_lock_state[-1] is False + + def test_mcp_status_marks_servers_with_errors_disconnected(mcp_handler_module): module, _tmp_path = mcp_handler_module @@ -310,6 +379,69 @@ def test_mcp_client_call_tool_uses_server_tool_timeout(mcp_handler_module, monke assert call_timeouts[0].total_seconds() == 7 +def test_mcp_session_cleanup_timeout_does_not_mask_success( + mcp_handler_module, monkeypatch +): + module, _tmp_path = mcp_handler_module + monkeypatch.setattr(module, "MCP_SESSION_CLEANUP_TIMEOUT_SECONDS", 0.01) + + class _HangingTransport: + async def __aenter__(self): + return "stdio", "write" + + async def __aexit__(self, exc_type, exc, tb): + await asyncio.sleep(60) + + class _FakeSession: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def initialize(self): + pass + + class _FakeClient(module.MCPClientBase): + async def _create_stdio_transport(self, current_exit_stack): + return await current_exit_stack.enter_async_context(_HangingTransport()) + + async def operation(_session): + return "ok" + + monkeypatch.setattr(module, "ClientSession", _FakeSession) + client = _FakeClient(SimpleNamespace(name="server")) + + assert asyncio.run(client._execute_with_session(operation)) == "ok" + + +def test_mcp_isolated_operation_timeout_returns_control(mcp_handler_module): + module, _tmp_path = mcp_handler_module + + class _FakeClient(module.MCPClientBase): + async def _create_stdio_transport(self, current_exit_stack): + raise AssertionError("transport should not be used") + + async def never_finishes(): + await asyncio.sleep(60) + + client = _FakeClient(SimpleNamespace(name="server")) + + with pytest.raises(TimeoutError): + asyncio.run( + client._run_isolated_operation( + "wedged", + never_finishes, + timeout_seconds=0.01, + ) + ) + + assert "operation did not finish" in client.error + + def test_mcp_image_content_becomes_history_image_attachment(mcp_handler_module, monkeypatch): module, _tmp_path = mcp_handler_module agent, log, tool_results, messages, updates, warnings = _agent_recorder()