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()