diff --git a/strix/tools/mcp/agent_tools.py b/strix/tools/mcp/agent_tools.py index 75ee3d73..168bdcb1 100644 --- a/strix/tools/mcp/agent_tools.py +++ b/strix/tools/mcp/agent_tools.py @@ -25,6 +25,7 @@ automatically. from __future__ import annotations import asyncio +import contextlib import json from typing import TYPE_CHECKING, Any @@ -47,6 +48,7 @@ def _registry_from_ctx(ctx: RunContextWrapper) -> McpRegistry | None: _NO_CONNECTIONS = "No MCP connections are configured for this run." +_MCP_DISCOVERY_WAIT_SECONDS = 30.0 def _unknown_connection(connection: str, registry: McpRegistry) -> str: @@ -66,7 +68,9 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]: Read-only. Returns one entry per connection with its ``id`` (the exact name you pass to the other MCP tools), ``name``, ``description``, and - ``tool_count``. The response does not include tool schemas. Call + ``tool_count`` (null while the catalog is unknown). Waits at most 30 seconds + for catalogs; unfinished catalogs keep loading in the background. + The response does not include tool schemas. Call ``search_mcp_tools`` next. Then call ``get_mcp_tool_schema`` for one selected tool before you call ``call_mcp``. Returns an empty ``connections`` list when the run has no MCP connections. @@ -82,10 +86,14 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]: and (entry.session is None or not entry.session.is_dead) ] if pending: - await asyncio.gather( - *(entry.ensure_catalog() for entry in pending), - return_exceptions=True, - ) + with contextlib.suppress(TimeoutError): + await asyncio.wait_for( + asyncio.gather( + *(entry.ensure_catalog() for entry in pending), + return_exceptions=True, + ), + timeout=_MCP_DISCOVERY_WAIT_SECONDS, + ) dead_by_name = {status.name: status.dead for status in registry.statuses()} return { "connections": [ @@ -93,7 +101,7 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]: "id": summary.name, "name": summary.name, "description": summary.purpose, - "tool_count": summary.tool_count, + "tool_count": summary.tool_count if summary.state == "catalog_ready" else None, "dead": dead_by_name.get(summary.name, False), "state": summary.state, } diff --git a/strix/tools/mcp/registry.py b/strix/tools/mcp/registry.py index 6e87c958..1c2cc5b1 100644 --- a/strix/tools/mcp/registry.py +++ b/strix/tools/mcp/registry.py @@ -76,6 +76,11 @@ 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``. @@ -216,6 +221,8 @@ 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) diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index bb7475a5..fcbac374 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -43,6 +43,7 @@ from strix.tools.mcp import ( resolve_mcp_call, search_mcp_tools, ) +from strix.tools.mcp import agent_tools as mcp_agent_tools from strix.tools.mcp import client as mcp_client from strix.tools.mcp import registry as mcp_registry_mod from strix.tools.mcp import session as mcp_session_mod @@ -484,6 +485,116 @@ async def test_list_mcps_loads_catalog_for_unwarmed_connection( await registry.close() +@pytest.mark.asyncio +@pytest.mark.parametrize("slow_stage", ["connect", "catalog"]) +@pytest.mark.parametrize("warmup", [False, True]) +async def test_list_mcps_returns_healthy_connections_while_slow_catalog_keeps_loading( + monkeypatch: pytest.MonkeyPatch, + slow_stage: str, + warmup: bool, +) -> None: + release = asyncio.Event() + started = asyncio.Event() + catalog_calls = 0 + + class _SlowServer(FakeMCPServer): + async def connect(self) -> None: + if slow_stage == "connect": + started.set() + await release.wait() + + async def list_tools(self, run_context: Any = None, agent: Any = None) -> list[MCPTool]: + nonlocal catalog_calls + catalog_calls += 1 + if slow_stage == "catalog": + started.set() + await release.wait() + return await super().list_tools(run_context, agent) + + servers = { + "healthy": FakeMCPServer("healthy", [_mcp_tool("read_file")]), + "slow": _SlowServer("slow", [_mcp_tool("query")]), + } + monkeypatch.setattr( + mcp_client, "_build_server", lambda config: _built_server(servers[config.name]) + ) + monkeypatch.setattr(mcp_agent_tools, "_MCP_DISCOVERY_WAIT_SECONDS", 0.01) + registry = McpRegistry() + for name in servers: + registry.register(McpConnectionRequest(config=_config(name, None))) + if warmup: + registry.start_warmup() + entry = registry.get("slow") + assert entry is not None + try: + out = await asyncio.wait_for(list_mcps.on_invoke_tool(_ctx(registry), "{}"), timeout=1) + assert started.is_set() + assert out["connections"][0]["tool_count"] == 1 + assert out["connections"][0]["state"] == "catalog_ready" + assert out["connections"][1]["tool_count"] is None + assert out["connections"][1]["state"] in {"connecting", "catalog_loading"} + assert entry._catalog_task is not None + assert not entry._catalog_task.done() + + repeated = await asyncio.wait_for(list_mcps.on_invoke_tool(_ctx(registry), "{}"), timeout=1) + assert repeated == out + release.set() + await asyncio.wait_for(entry.ensure_catalog(), timeout=1) + finished = await list_mcps.on_invoke_tool(_ctx(registry), "{}") + assert finished["connections"][1]["tool_count"] == 1 + assert finished["connections"][1]["state"] == "catalog_ready" + assert catalog_calls == 1 + finally: + release.set() + await registry.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_after_deadline", [False, True]) +async def test_list_mcps_retains_and_cleans_up_background_catalog_tasks( + monkeypatch: pytest.MonkeyPatch, + fail_after_deadline: bool, +) -> None: + release = asyncio.Event() + finished = asyncio.Event() + loop_errors: list[dict[str, Any]] = [] + + async def load_catalog() -> list[MCPTool]: + await release.wait() + raise RuntimeError("catalog unavailable") + + registry = McpRegistry() + entry = registry.add( + name="blocked", server=FakeMCPServer("blocked", []), purpose=None, tool_count=0 + ) + assert entry.session is not None + monkeypatch.setattr(entry.session, "list_tools", load_catalog) + monkeypatch.setattr(mcp_agent_tools, "_MCP_DISCOVERY_WAIT_SECONDS", 0.01) + loop = asyncio.get_running_loop() + previous_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _loop, context: loop_errors.append(context)) + try: + out = await asyncio.wait_for(list_mcps.on_invoke_tool(_ctx(registry), "{}"), timeout=1) + assert out["connections"][0]["tool_count"] is None + assert entry._catalog_task is not None + assert not entry._catalog_task.done() + if fail_after_deadline: + entry._catalog_task.add_done_callback(lambda _task: finished.set()) + release.set() + await asyncio.wait_for(finished.wait(), timeout=1) + assert entry.state == "connected" + # Drop the finished task without awaiting it; its error must already be consumed. + entry._catalog_task = None + else: + task = entry._catalog_task + await registry.close() + assert task.cancelled() + assert loop_errors == [] + finally: + await registry.close() + loop.set_exception_handler(previous_handler) + + @pytest.mark.asyncio async def test_list_mcps_empty_without_a_registry() -> None: assert await list_mcps.on_invoke_tool(_ctx(None), "{}") == {"connections": []}