vrr/agent-zero
2
0
Fork 0
mirror of https://github.com/agent0ai/agent-zero.git synced 2026-07-10 01:18:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

Fix MCP timeouts from hanging agent loop

Browse source 
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.
This commit is contained in:
Alessandro 2026-06-18 13:32:10 +02:00
parent 352320ef81
commit 8fda0ee69c
3 changed files with 264 additions and 81 deletions
Download patch file Download diff file Expand all files Collapse all files

204
helpers/mcp_handler.py
View file

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

9
helpers/mcp_handler.py.dox.md
View file

@ -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

132
tests/test_mcp_handler_multimodal.py
View file

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