@samitouri / QOSami_Agents_Zero-1 / commits / 8fda0ee6

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.

Alessandro committed Jun 18, 2026 at 13:32 UTC 8fda0ee69c8b0a9a83f1a3eccc9f33a5da7cebc0

3 files changed +264 -81

helpers/mcp_handler.py

+126 -78

@@ -43,10 +43,13 @@ from pydantic import BaseModel, Field, Discriminator, Tag, PrivateAttr

43 from helpers import dirty_json, media_artifacts

44 from helpers.print_style import PrintStyle

45 from helpers.tool import Tool, Response

46 +from helpers.defer import DeferredTask

49 MCP_MEDIA_TOKENS_ESTIMATE = 1500

50 MAX_MCP_RESOURCE_TEXT_CHARS = 12_000

51 +MCP_SESSION_CLEANUP_TIMEOUT_SECONDS = 5.0

52 +MCP_OPERATION_TIMEOUT_GRACE_SECONDS = MCP_SESSION_CLEANUP_TIMEOUT_SECONDS + 2.0

53 DEFAULT_MCP_SERVERS_CONFIG = '{\n "mcpServers": {}\n}'

@@ -661,10 +664,10 @@ class MCPConfig(BaseModel):

664

665 @classmethod

666 def get_instance(cls) -> "MCPConfig":

664 - # with cls.__lock:

665 - if cls.__instance is None:

666 - cls.__instance = cls(servers_list=[], config_scope="global")

667 - return cls.__instance

667 + with cls.__lock:

668 + if cls.__instance is None:

669 + cls.__instance = cls(servers_list=[], config_scope="global")

670 + return cls.__instance

671

672 @classmethod

673 def clear_project_instances(cls):

@@ -791,35 +794,20 @@ class MCPConfig(BaseModel):

794

795 @classmethod

796 def update(cls, config_str: str) -> Any:

797 + servers_data = cls.parse_config_string(config_str)

798 + new_instance = cls(servers_list=servers_data, config_scope="global")

799 with cls.__lock:

795 - servers_data = cls.parse_config_string(config_str)

796 -

797 - # Initialize/update the singleton instance with the (potentially empty) list of server data

798 - instance = cls.get_instance()

799 - # Directly update the servers attribute of the existing instance or re-initialize carefully

800 - # For simplicity and to ensure __init__ logic runs if needed for setup:

801 - new_instance_data = {

802 - "servers": servers_data

803 - } # Prepare data for re-initialization or update

804 -

805 - # Option 1: Re-initialize the existing instance (if __init__ is idempotent for other fields)

806 - instance.__init__(servers_list=servers_data, config_scope="global")

800 + # Build and initialize outside the class lock so a slow or wedged MCP

801 + # server cannot freeze status reads, prompts, or later tool calls.

802 + instance = cls.__instance

803 + if instance is None:

804 + instance = new_instance

805 + cls.__instance = instance

806 + else:

807 + instance.servers = new_instance.servers

808 + instance.disconnected_servers = new_instance.disconnected_servers

809 + instance.config_scope = new_instance.config_scope

810 cls.__project_instances = {}

808 -

809 - # Option 2: Or, if __init__ has side effects we don't want to repeat,

810 - # and 'servers' is the primary thing 'update' changes:

811 - # instance.servers = [] # Clear existing servers first

812 - # for server_item_data in servers_data:

813 - # try:

814 - # if server_item_data.get("url", None):

815 - # instance.servers.append(MCPServerRemote(server_item_data))

816 - # else:

817 - # instance.servers.append(MCPServerLocal(server_item_data))

818 - # except Exception as e_init:

819 - # PrintStyle(background_color="grey", font_color="red", padding=True).print(

820 - # f"MCPConfig.update: Failed to create MCPServer from item '{server_item_data.get('name', 'Unknown')}': {e_init}"

821 - # )

822 -

811 cls.__initialized = True

812 return instance

813

@@ -1146,11 +1134,15 @@ class MCPConfig(BaseModel):

1134 ) -> CallToolResult:

1135 """Call a tool with the given input data"""

1136 server_name_part, tool_name_part = _split_qualified_tool_name(tool_name)

1137 + matched_server = None

1138 with self.__lock:

1139 for server in self.servers:

1140 if server.name == server_name_part and server.has_tool(tool_name_part):

1152 - return await server.call_tool(tool_name_part, input_data)

1141 + matched_server = server

1142 + break

1143 + if matched_server is None:

1144 raise ValueError(f"Tool {tool_name} not found")

1145 + return await matched_server.call_tool(tool_name_part, input_data)

1146

1147

1148 T = TypeVar("T")

@@ -1170,6 +1162,49 @@ class MCPClientBase(ABC):

1162 self.log: List[str] = []

1163 self.log_file: Optional[TextIO] = None

1164

1165 + def _operation_timeout_seconds(self, read_timeout_seconds: float) -> float:

1166 + try:

1167 + seconds = float(read_timeout_seconds)

1168 + except (TypeError, ValueError):

1169 + seconds = 60.0

1170 + if seconds <= 0:

1171 + seconds = 60.0

1172 + return seconds + MCP_OPERATION_TIMEOUT_GRACE_SECONDS

1173 +

1174 + def _operation_thread_name(self, operation_name: str) -> str:

1175 + server_name = normalize_name(str(getattr(self.server, "name", "") or "server"))

1176 + return f"MCPClient-{server_name[:32] or 'server'}-{operation_name}-{uuid.uuid4().hex[:8]}"

1177 +

1178 + async def _run_isolated_operation(

1179 + self,

1180 + operation_name: str,

1181 + operation: Callable[[], Awaitable[T]],

1182 + timeout_seconds: float,

1183 + ) -> T:

1184 + worker = DeferredTask(thread_name=self._operation_thread_name(operation_name))

1185 + timed_out = False

1186 + try:

1187 + return await asyncio.wait_for(

1188 + worker.execute_inside(operation),

1189 + timeout=timeout_seconds,

1190 + )

1191 + except asyncio.TimeoutError as exc:

1192 + timed_out = True

1193 + message = (

1194 + f"MCPClientBase ({self.server.name} - {operation_name}): "

1195 + f"operation did not finish within {timeout_seconds:.1f}s; "

1196 + "abandoning the isolated worker so Agent Zero can continue."

1197 + )

1198 + PrintStyle.warning(message)

1199 + with self.__lock:

1200 + self.error = message

1201 + raise TimeoutError(message) from exc

1202 + finally:

1203 + if timed_out:

1204 + worker.kill(terminate_thread=False)

1205 + else:

1206 + worker.kill(terminate_thread=True)

1207 +

1208 # Protected method

1209 @abstractmethod

1210 async def _create_stdio_transport(

@@ -1192,54 +1227,56 @@ class MCPClientBase(ABC):

1227 """

1228 operation_name = coro_func.__name__ # For logging

1229 # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name}): Creating new session for operation '{operation_name}'...")

1195 - # Store the original exception outside the async block

1230 original_exception = None

1231 + result: T | None = None

1232 + has_result = False

1233 + temp_stack = AsyncExitStack()

1234 try:

1198 - async with AsyncExitStack() as temp_stack:

1199 - try:

1200 -

1201 - stdio, write = await self._create_stdio_transport(temp_stack)

1202 - # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Transport created. Initializing session...")

1203 - session = await temp_stack.enter_async_context(

1204 - ClientSession(

1205 - stdio, # type: ignore

1206 - write, # type: ignore

1207 - read_timeout_seconds=timedelta(

1208 - seconds=read_timeout_seconds

1209 - ),

1210 - )

1211 - )

1212 - await session.initialize()

1213 -

1214 - result = await coro_func(session)

1235 + stdio, write = await self._create_stdio_transport(temp_stack)

1236 + # PrintStyle(font_color="cyan").print(f"MCPClientBase ({self.server.name} - {operation_name}): Transport created. Initializing session...")

1237 + session = await temp_stack.enter_async_context(

1238 + ClientSession(

1239 + stdio, # type: ignore

1240 + write, # type: ignore

1241 + read_timeout_seconds=timedelta(

1242 + seconds=read_timeout_seconds

1243 + ),

1244 + )

1245 + )

1246 + await session.initialize()

1247

1216 - return result

1217 - except Exception as e:

1218 - # Store the original exception and raise a dummy exception

1219 - excs = getattr(e, "exceptions", None) # Python 3.11+ ExceptionGroup

1220 - if excs:

1221 - original_exception = excs[0]

1222 - else:

1223 - original_exception = e

1224 - # Create a dummy exception to break out of the async block

1225 - raise RuntimeError("Dummy exception to break out of async block")

1248 + result = await coro_func(session)

1249 + has_result = True

1250 except Exception as e:

1227 - # Check if this is our dummy exception

1228 - if original_exception is not None:

1229 - e = original_exception

1230 - # We have the original exception stored

1251 + excs = getattr(e, "exceptions", None) # Python 3.11+ ExceptionGroup

1252 + if excs:

1253 + original_exception = excs[0]

1254 + else:

1255 + original_exception = e

1256 + try:

1257 + await asyncio.wait_for(

1258 + temp_stack.aclose(),

1259 + timeout=MCP_SESSION_CLEANUP_TIMEOUT_SECONDS,

1260 + )

1261 + except asyncio.TimeoutError:

1262 + PrintStyle.warning(

1263 + f"MCPClientBase ({self.server.name} - {operation_name}): "

1264 + f"session cleanup exceeded {MCP_SESSION_CLEANUP_TIMEOUT_SECONDS:.1f}s."

1265 + )

1266 + except Exception as cleanup_exception:

1267 + PrintStyle.warning(

1268 + f"MCPClientBase ({self.server.name} - {operation_name}): "

1269 + f"session cleanup failed: {type(cleanup_exception).__name__}: {cleanup_exception}"

1270 + )

1271 + if original_exception is not None:

1272 PrintStyle(

1273 background_color="#AA4455", font_color="white", padding=False

1274 ).print(

1234 - f"MCPClientBase ({self.server.name} - {operation_name}): Error during operation: {type(e).__name__}: {e}"

1275 + f"MCPClientBase ({self.server.name} - {operation_name}): Error during operation: {type(original_exception).__name__}: {original_exception}"

1276 )

1236 - raise e # Re-raise the original exception

1237 - # finally:

1238 - # PrintStyle(font_color="cyan").print(

1239 - # f"MCPClientBase ({self.server.name} - {operation_name}): Session and transport will be closed by AsyncExitStack."

1240 - # )

1241 - # This line should ideally be unreachable if the try/except/finally logic within the 'async with' is exhaustive.

1242 - # Adding it to satisfy linters that might not fully trace the raise/return paths through async context managers.

1277 + raise original_exception

1278 + if has_result:

1279 + return cast(T, result)

1280 raise RuntimeError(

1281 f"MCPClientBase ({self.server.name} - {operation_name}): _execute_with_session exited 'async with' block unexpectedly."

1282 )

@@ -1270,9 +1307,13 @@ class MCPClientBase(ABC):

1307 or current_settings.get("mcp_client_init_timeout", 10)

1308 or 10

1309 )

1273 - await self._execute_with_session(

1274 - list_tools_op,

1275 - read_timeout_seconds=init_timeout,

1310 + await self._run_isolated_operation(

1311 + "update_tools",

1312 + lambda: self._execute_with_session(

1313 + list_tools_op,

1314 + read_timeout_seconds=init_timeout,

1315 + ),

1316 + timeout_seconds=self._operation_timeout_seconds(init_timeout),

1317 )

1318 except Exception as e:

1319 # e = eg.exceptions[0]

@@ -1339,10 +1380,17 @@ class MCPClientBase(ABC):

1380 return response

1381

1382 try:

1342 - return await self._execute_with_session(

1343 - call_tool_op,

1344 - read_timeout_seconds=tool_timeout,

1383 + response = await self._run_isolated_operation(

1384 + "call_tool",

1385 + lambda: self._execute_with_session(

1386 + call_tool_op,

1387 + read_timeout_seconds=tool_timeout,

1388 + ),

1389 + timeout_seconds=self._operation_timeout_seconds(tool_timeout),

1390 )

1391 + with self.__lock:

1392 + self.error = ""

1393 + return response

1394 except Exception as e:

1395 # Error logged by _execute_with_session. Re-raise a specific error for the caller.

1396 PrintStyle(

helpers/mcp_handler.py.dox.md

+6 -3

@@ -67,7 +67,7 @@

67 - `_split_qualified_tool_name(tool_name: str) -> tuple[str, str]`: Split `server.tool` names while preserving dots inside MCP tool names.

68 - `_normalize_disabled_tools(value: Any) -> list[str]`: Normalize the optional per-server disabled tool list.

69 - `initialize_mcp(mcp_servers_config: str)`

70 -- Notable constants/configuration names: `DEFAULT_MCP_SERVERS_CONFIG`, `MCP_MEDIA_TOKENS_ESTIMATE`, `MAX_MCP_RESOURCE_TEXT_CHARS`, `T`.

70 +- 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`.

72 ## Runtime Contracts

@@ -81,13 +81,15 @@

81 - MCP tool names are qualified as `server_name.tool_name`; server names are normalized without dots, and the tool portion may contain dots.

82 - 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.

83 - Server-specific `init_timeout` and `tool_timeout` override global MCP client timeout settings for list-tools and call-tool operations.

84 +- 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.

85 +- 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.

86 - Server status marks initialized server objects with cached initialization errors as disconnected, even if the config object exists.

87 - Observed side-effect areas: filesystem writes, network calls, WebSocket state, settings/state persistence, secret handling.

86 -- 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`.

88 +- 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`.

90 ## Key Concepts

90 -- 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`.

92 +- 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`.

93 - Keep request/response, tool, or helper semantics documented here at the same time as source changes.

95 ## Work Guidance

@@ -95,6 +97,7 @@

97 - Preserve public helper APIs used by core code and plugins unless every caller is updated.

98 - Keep path, auth, secret, persistence, network, and subprocess behavior explicit and bounded.

99 - Prefer adding cohesive helper functions here only when behavior is reused across modules.

100 +- Keep MCP timeout and cleanup changes covered by deterministic tests that do not require real MCP servers or network credentials.

101

102 ## Verification

103

tests/test_mcp_handler_multimodal.py

+132

@@ -54,6 +54,20 @@ class _FakeCallToolResult(SimpleNamespace):

54 pass

57 +class _TrackingLock:

58 + def __init__(self):

59 + self.held = False

60 +

61 + def __enter__(self):

62 + assert self.held is False

63 + self.held = True

64 + return self

65 +

66 + def __exit__(self, exc_type, exc, tb):

67 + self.held = False

68 + return False

69 +

70 +

71 @pytest.fixture

72 def mcp_handler_module(monkeypatch, tmp_path):

73 monkeypatch.delitem(sys.modules, "helpers.mcp_handler", raising=False)

@@ -118,6 +132,10 @@ def mcp_handler_module(monkeypatch, tmp_path):

132 def stream(self, *args, **kwargs):

133 return self

134

135 + @staticmethod

136 + def warning(*args, **kwargs):

137 + return None

138 +

139 def _fake_get_abs_path(*parts):

140 return str(tmp_path.joinpath(*parts))

141

@@ -194,6 +212,57 @@ def test_mcp_config_preserves_dotted_tool_names(mcp_handler_module):

212 assert called == [("alpha.beta", {"value": 7})]

213

214

215 +def test_mcp_config_call_tool_releases_config_lock_before_await(

216 + mcp_handler_module, monkeypatch

217 +):

218 + module, _tmp_path = mcp_handler_module

219 + lock = _TrackingLock()

220 + observed_lock_state: list[bool] = []

221 +

222 + monkeypatch.setattr(module.MCPConfig, "_MCPConfig__lock", lock, raising=False)

223 +

224 + class _FakeServer:

225 + name = "server"

226 + description = "Fake MCP server"

227 + type = "stdio"

228 + scope = "global"

229 +

230 + def has_tool(self, tool_name):

231 + return tool_name == "run"

232 +

233 + async def call_tool(self, tool_name, input_data):

234 + observed_lock_state.append(lock.held)

235 + await asyncio.sleep(0)

236 + return _FakeCallToolResult(content=[], isError=False)

237 +

238 + config = module.MCPConfig(servers_list=[])

239 + config.servers = [_FakeServer()]

240 +

241 + asyncio.run(config.call_tool("server.run", {}))

242 +

243 + assert observed_lock_state == [False]

244 +

245 +

246 +def test_mcp_config_update_initializes_outside_config_lock(

247 + mcp_handler_module, monkeypatch

248 +):

249 + module, _tmp_path = mcp_handler_module

250 + lock = _TrackingLock()

251 + observed_lock_state: list[bool] = []

252 + original_init = module.MCPConfig.__init__

253 +

254 + def tracking_init(self, *args, **kwargs):

255 + observed_lock_state.append(lock.held)

256 + original_init(self, *args, **kwargs)

257 +

258 + monkeypatch.setattr(module.MCPConfig, "_MCPConfig__lock", lock, raising=False)

259 + monkeypatch.setattr(module.MCPConfig, "__init__", tracking_init)

260 +

261 + module.MCPConfig.update('{"mcpServers": {}}')

262 +

263 + assert observed_lock_state[-1] is False

264 +

265 +

266 def test_mcp_status_marks_servers_with_errors_disconnected(mcp_handler_module):

267 module, _tmp_path = mcp_handler_module

268

@@ -310,6 +379,69 @@ def test_mcp_client_call_tool_uses_server_tool_timeout(mcp_handler_module, monke

379 assert call_timeouts[0].total_seconds() == 7

380

381

382 +def test_mcp_session_cleanup_timeout_does_not_mask_success(

383 + mcp_handler_module, monkeypatch

384 +):

385 + module, _tmp_path = mcp_handler_module

386 + monkeypatch.setattr(module, "MCP_SESSION_CLEANUP_TIMEOUT_SECONDS", 0.01)

387 +

388 + class _HangingTransport:

389 + async def __aenter__(self):

390 + return "stdio", "write"

391 +

392 + async def __aexit__(self, exc_type, exc, tb):

393 + await asyncio.sleep(60)

394 +

395 + class _FakeSession:

396 + def __init__(self, *args, **kwargs):

397 + pass

398 +

399 + async def __aenter__(self):

400 + return self

401 +

402 + async def __aexit__(self, exc_type, exc, tb):

403 + return False

404 +

405 + async def initialize(self):

406 + pass

407 +

408 + class _FakeClient(module.MCPClientBase):

409 + async def _create_stdio_transport(self, current_exit_stack):

410 + return await current_exit_stack.enter_async_context(_HangingTransport())

411 +

412 + async def operation(_session):

413 + return "ok"

414 +

415 + monkeypatch.setattr(module, "ClientSession", _FakeSession)

416 + client = _FakeClient(SimpleNamespace(name="server"))

417 +

418 + assert asyncio.run(client._execute_with_session(operation)) == "ok"

419 +

420 +

421 +def test_mcp_isolated_operation_timeout_returns_control(mcp_handler_module):

422 + module, _tmp_path = mcp_handler_module

423 +

424 + class _FakeClient(module.MCPClientBase):

425 + async def _create_stdio_transport(self, current_exit_stack):

426 + raise AssertionError("transport should not be used")

427 +

428 + async def never_finishes():

429 + await asyncio.sleep(60)

430 +

431 + client = _FakeClient(SimpleNamespace(name="server"))

432 +

433 + with pytest.raises(TimeoutError):

434 + asyncio.run(

435 + client._run_isolated_operation(

436 + "wedged",

437 + never_finishes,

438 + timeout_seconds=0.01,

439 + )

440 + )

441 +

442 + assert "operation did not finish" in client.error

443 +

444 +

445 def test_mcp_image_content_becomes_history_image_attachment(mcp_handler_module, monkeypatch):

446 module, _tmp_path = mcp_handler_module

447 agent, log, tool_results, messages, updates, warnings = _agent_recorder()