diff --git a/strix/tools/mcp/agent_tools.py b/strix/tools/mcp/agent_tools.py index 168bdcb1..1033d334 100644 --- a/strix/tools/mcp/agent_tools.py +++ b/strix/tools/mcp/agent_tools.py @@ -25,7 +25,6 @@ automatically. from __future__ import annotations import asyncio -import contextlib import json from typing import TYPE_CHECKING, Any @@ -86,14 +85,15 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]: and (entry.session is None or not entry.session.is_dead) ] if pending: - with contextlib.suppress(TimeoutError): - await asyncio.wait_for( + await asyncio.wait( + { asyncio.gather( *(entry.ensure_catalog() for entry in pending), return_exceptions=True, - ), - timeout=_MCP_DISCOVERY_WAIT_SECONDS, - ) + ) + }, + timeout=_MCP_DISCOVERY_WAIT_SECONDS, + ) dead_by_name = {status.name: status.dead for status in registry.statuses()} return { "connections": [ diff --git a/strix/tools/mcp/registry.py b/strix/tools/mcp/registry.py index 1c2cc5b1..6e87c958 100644 --- a/strix/tools/mcp/registry.py +++ b/strix/tools/mcp/registry.py @@ -76,11 +76,6 @@ McpConnectionState = Literal[ _RETRY_DELAY_SECONDS = 5.0 -def _consume_task_exception(task: asyncio.Task[Any]) -> None: - if not task.cancelled(): - task.exception() - - @dataclasses.dataclass class McpConnectionEntry: """One configured MCP connection a scan may reach, keyed by ``name``. @@ -221,8 +216,6 @@ class McpConnectionEntry: self._load_catalog(), name=f"mcp-catalog-{self.name}", ) - # A discovery deadline can detach the waiter while this task keeps loading. - self._catalog_task.add_done_callback(_consume_task_exception) task = self._catalog_task try: return await asyncio.shield(task)