From f5a900416b6e1d8023e7c2ce11885232ffa8b806 Mon Sep 17 00:00:00 2001 From: alex s <46074070+bearsyankees@users.noreply.github.com> Date: Thu, 8 Oct 2026 12:42:39 -0400 Subject: [PATCH 1/3] fix(mcp): load catalogs before agents discover connections (#1492) * fix(mcp): load catalogs before agents discover connections * fix(mcp): bound discovery waits without cancelling catalog loads * fix(mcp): keep discovery waiters alive across deadlines --- strix/agents/prompts/scope.jinja | 2 +- strix/core/runner.py | 33 +++-- strix/tools/mcp/agent_tools.py | 27 +++- strix/tools/mcp/registry.py | 10 +- tests/test_mcp_client.py | 231 ++++++++++++++++++++++++++++++- tests/test_runner_mcp.py | 125 +++++++++++++++++ tests/test_runner_root_prompt.py | 18 ++- 7 files changed, 421 insertions(+), 25 deletions(-) diff --git a/strix/agents/prompts/scope.jinja b/strix/agents/prompts/scope.jinja index 4211fd53..c9dc5c72 100644 --- a/strix/agents/prompts/scope.jinja +++ b/strix/agents/prompts/scope.jinja @@ -21,7 +21,7 @@ MCP CONNECTIONS (available this run): {% if system_prompt_context.mcp_connections %} - Connected this run (search one to find relevant tools): {% for connection in system_prompt_context.mcp_connections %} - - {{ connection.name }} ({{ connection.tool_count }} tools){% if connection.purpose %}: {{ connection.purpose }}{% endif %} + - {{ connection.name }}{% if connection.state == "unavailable" %} (unavailable right now){% elif connection.tool_count is not none %} ({{ connection.tool_count }} tools){% endif %}{% if connection.purpose %}: {{ connection.purpose }}{% endif %} {% endfor %} {% endif %} - Reach for a connection whenever the target itself cannot give you information a connection could: its database schema and access policies, real deployment or infrastructure configuration, known issues or prior findings, or server logs. In those cases call list_mcps early to see what is available, and prefer a connection's authoritative data over inferring from the target's responses. Do not wait to be told a connection exists. diff --git a/strix/core/runner.py b/strix/core/runner.py index 048a59de..58124be2 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -68,6 +68,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +_MCP_PROMPT_WARMUP_TIMEOUT_SECONDS = 30.0 + StreamEventSink = Callable[[str, Any], None] # Receives the run's MCP connection roster as a list of non-secret status dicts @@ -93,6 +95,19 @@ def _mcp_roster_payload(registry: McpRegistry) -> list[dict[str, Any]]: ] +def _mcp_prompt_roster(registry: McpRegistry) -> list[dict[str, Any]]: + """The MCP roster rendered in agent prompts, with only verified counts.""" + return [ + { + "name": summary.name, + "purpose": summary.purpose, + "tool_count": summary.tool_count if summary.state == "catalog_ready" else None, + "state": summary.state, + } + for summary in registry.summaries() + ] + + def _record_mcp_connections(connection_names: list[str]) -> None: """Record which MCP servers this run configured, for the interfaces. @@ -438,19 +453,13 @@ async def run_strix_scan( _record_mcp_connections(mcp_registry.names()) report( f"MCP: configured {len(mcp_registry)} connection(s); " - "warming them in the background" + "connecting and listing tools" ) scope_context["mcp_available"] = True - scope_context["mcp_connections"] = [ - { - "name": summary.name, - "purpose": summary.purpose, - "tool_count": summary.tool_count, - } - for summary in mcp_registry.summaries() - ] + scope_context["mcp_connections"] = _mcp_prompt_roster(mcp_registry) def _emit_mcp_status() -> None: + scope_context["mcp_connections"] = _mcp_prompt_roster(mcp_registry) roster = _mcp_roster_payload(mcp_registry) _persist_mcp_status(roster) if mcp_status_sink is not None: @@ -461,7 +470,11 @@ async def run_strix_scan( mcp_registry.set_status_sink(_emit_mcp_status) _emit_mcp_status() - mcp_registry.start_warmup(max_concurrency=6) + warmup_task = mcp_registry.start_warmup(max_concurrency=6) + await asyncio.wait( + {warmup_task}, + timeout=_MCP_PROMPT_WARMUP_TIMEOUT_SECONDS, + ) except Exception: logger.exception("Failed to configure user MCP servers; continuing without them") diff --git a/strix/tools/mcp/agent_tools.py b/strix/tools/mcp/agent_tools.py index 7fa2c392..1033d334 100644 --- a/strix/tools/mcp/agent_tools.py +++ b/strix/tools/mcp/agent_tools.py @@ -24,6 +24,7 @@ automatically. from __future__ import annotations +import asyncio import json from typing import TYPE_CHECKING, Any @@ -46,6 +47,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: @@ -61,11 +63,13 @@ def _format_tool(tool: MCPTool) -> str: @function_tool(timeout=60) async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]: - """List the MCP connections available this run, so you can discover them. + """List available MCP connections, loading each connection's catalog on first use. 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. @@ -73,6 +77,23 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]: registry = _registry_from_ctx(ctx) if registry is None or not registry: return {"connections": []} + pending = [ + entry + for name in registry.names() + if (entry := registry.get(name)) is not None + and entry.state != "catalog_ready" + and (entry.session is None or not entry.session.is_dead) + ] + if pending: + await asyncio.wait( + { + 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": [ @@ -80,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 65652ab1..6e87c958 100644 --- a/strix/tools/mcp/registry.py +++ b/strix/tools/mcp/registry.py @@ -432,19 +432,19 @@ class McpRegistry: entry.set_status_sink(sink) def start_warmup(self, *, max_concurrency: int = 6) -> asyncio.Task[None]: - """Connect every configured entry in the background with a fixed bound.""" + """Connect and list every configured entry with a fixed concurrency bound.""" if self._warmup_task is not None: return self._warmup_task async def warm() -> None: semaphore = asyncio.Semaphore(max(1, max_concurrency)) - async def connect(entry: McpConnectionEntry) -> None: + async def load_catalog(entry: McpConnectionEntry) -> None: async with semaphore: - with contextlib.suppress(McpConnectionUnavailableError): - await entry.ensure_connected() + with contextlib.suppress(Exception): + await entry.ensure_catalog() - await asyncio.gather(*(connect(entry) for entry in self._entries.values())) + await asyncio.gather(*(load_catalog(entry) for entry in self._entries.values())) self._warmup_task = asyncio.create_task(warm(), name="mcp-warmup") return self._warmup_task diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 56b96887..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 @@ -421,8 +422,18 @@ def test_registry_summaries() -> None: @pytest.mark.asyncio async def test_list_mcps_returns_connections_with_ids_and_descriptions() -> None: registry = McpRegistry() - registry.add(name="fs", server=FakeMCPServer("fs", []), purpose="local files", tool_count=2) - registry.add(name="db", server=FakeMCPServer("db", []), purpose=None, tool_count=1) + registry.add( + name="fs", + server=FakeMCPServer("fs", [_mcp_tool("read_file"), _mcp_tool("write_file")]), + purpose="local files", + tool_count=2, + ) + registry.add( + name="db", + server=FakeMCPServer("db", [_mcp_tool("query")]), + purpose=None, + tool_count=1, + ) out = await list_mcps.on_invoke_tool(_ctx(registry), "{}") @@ -437,7 +448,7 @@ async def test_list_mcps_returns_connections_with_ids_and_descriptions() -> None "description": "local files", "tool_count": 2, "dead": False, - "state": "connected", + "state": "catalog_ready", }, { "id": "db", @@ -445,10 +456,143 @@ async def test_list_mcps_returns_connections_with_ids_and_descriptions() -> None "description": None, "tool_count": 1, "dead": False, - "state": "connected", + "state": "catalog_ready", }, ] } + await registry.close() + + +@pytest.mark.asyncio +async def test_list_mcps_loads_catalog_for_unwarmed_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server = FakeMCPServer("fs", [_mcp_tool("read_file"), _mcp_tool("write_file")]) + monkeypatch.setattr( + mcp_client, + "_build_server", + lambda _config: _built_server(server), + ) + registry = McpRegistry() + entry = registry.register(McpConnectionRequest(config=_config("fs", None))) + + out = await list_mcps.on_invoke_tool(_ctx(registry), "{}") + + assert entry.state == "catalog_ready" + assert entry.tool_count == 2 + assert out["connections"][0]["tool_count"] == 2 + assert out["connections"][0]["state"] == "catalog_ready" + 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 @@ -733,7 +877,37 @@ async def test_registry_warmup_bounds_parallel_connections( release.set() await warmup - assert [summary.state for summary in registry.summaries()] == ["connected"] * 4 + assert [summary.state for summary in registry.summaries()] == ["catalog_ready"] * 4 + assert [summary.tool_count for summary in registry.summaries()] == [1] * 4 + await registry.close() + + +@pytest.mark.asyncio +async def test_registry_warmup_swallows_catalog_listing_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FailingListServer(FakeMCPServer): + async def list_tools( + self, + run_context: Any = None, + agent: Any = None, + ) -> list[MCPTool]: + raise RuntimeError("catalog unavailable") + + monkeypatch.setattr( + mcp_client, + "_build_server", + lambda _config: _built_server(_FailingListServer("db", [])), + ) + registry = McpRegistry() + entry = registry.register(McpConnectionRequest(config=_config("db", None))) + + warmup = registry.start_warmup() + await warmup + + assert warmup.done() + assert entry.state == "connected" + assert entry.tool_count == 0 await registry.close() @@ -1012,6 +1186,53 @@ def test_prompt_renders_named_connection_inventory() -> None: assert "read the app's schema" in prompt +def test_prompt_omits_unloaded_tool_count() -> None: + prompt = render_system_prompt( + system_prompt_context={ + "mcp_available": True, + "mcp_connections": [ + {"name": "pending_conn", "purpose": None, "tool_count": None, "state": "connected"} + ], + } + ) + + connection_line = next(line for line in prompt.splitlines() if "pending_conn" in line) + assert "(0 tools)" not in connection_line + assert "tools)" not in connection_line + + +def test_prompt_renders_loaded_tool_count() -> None: + prompt = render_system_prompt( + system_prompt_context={ + "mcp_available": True, + "mcp_connections": [ + {"name": "ready_conn", "purpose": None, "tool_count": 8, "state": "catalog_ready"} + ], + } + ) + + assert "ready_conn (8 tools)" in prompt + + +def test_prompt_renders_unavailable_connection_without_tool_count() -> None: + prompt = render_system_prompt( + system_prompt_context={ + "mcp_available": True, + "mcp_connections": [ + { + "name": "unavailable_conn", + "purpose": None, + "tool_count": None, + "state": "unavailable", + } + ], + } + ) + + assert "unavailable_conn (unavailable right now)" in prompt + assert "unavailable_conn (0 tools)" not in prompt + + def test_prompt_inventory_is_gated_on_availability() -> None: """The block is gated on ``mcp_available``; an ``mcp_connections`` payload without it renders nothing, so a stale or spoofed list cannot leak names.""" diff --git a/tests/test_runner_mcp.py b/tests/test_runner_mcp.py index 25a5c888..62d6fab1 100644 --- a/tests/test_runner_mcp.py +++ b/tests/test_runner_mcp.py @@ -8,6 +8,7 @@ routine does the connecting. from __future__ import annotations +import importlib import types from typing import Any @@ -17,10 +18,21 @@ from agents import ModelSettings import strix.tools.mcp as mcp_pkg import strix.tools.notes.tools as notes_tools import strix.tools.todo.tools as todo_tools +from strix.agents import factory from strix.core import runner from strix.core.agents import AgentCoordinator from strix.runtime import session_manager from strix.tools.mcp import McpConnectionConfig, McpConnectionRequest +from strix.tools.mcp import client as mcp_client + + +_test_mcp_client = importlib.import_module("tests.test_mcp_client") +FakeMCPServer: Any = _test_mcp_client.FakeMCPServer +_mcp_tool: Any = _test_mcp_client._mcp_tool + + +def _built_server(server: Any) -> Any: + return mcp_client.BuiltMcpServer(server, None) def _settings() -> Any: @@ -62,6 +74,11 @@ def _wire_runner(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None: monkeypatch.setattr(runner, "build_strix_agent", lambda **_k: object()) monkeypatch.setattr(runner, "make_child_factory", lambda **_k: lambda **_kk: object()) monkeypatch.setattr(runner, "open_agent_session", lambda _root_id, _db: object()) + monkeypatch.setattr( + mcp_client, + "_build_server", + lambda config: _built_server(FakeMCPServer(config.name, [])), + ) async def _run_agent_loop(**_kwargs: Any) -> None: return None @@ -156,6 +173,11 @@ async def test_roster_is_persisted_even_without_a_status_sink( "load_user_mcp_configs", lambda: [McpConnectionConfig(name="local_fs", transport="stdio", command="npx")], ) + monkeypatch.setattr( + mcp_client, + "_build_server", + lambda _config: _built_server(FakeMCPServer("local_fs", [])), + ) persisted: list[list[dict[str, Any]]] = [] @@ -182,3 +204,106 @@ async def test_roster_is_persisted_even_without_a_status_sink( "state": "configured", } ] + + +@pytest.mark.asyncio +async def test_warmup_listing_failure_does_not_interrupt_scan( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + _wire_runner(monkeypatch, tmp_path) + + server = FakeMCPServer("db", []) + + async def _failing_list_tools( + run_context: Any = None, + agent: Any = None, + ) -> list[Any]: + del run_context, agent + raise RuntimeError("catalog unavailable") + + server.list_tools = _failing_list_tools + + monkeypatch.setattr( + mcp_client, + "_build_server", + lambda _config: _built_server(server), + ) + root_builds: list[dict[str, Any]] = [] + monkeypatch.setattr(runner, "build_strix_agent", lambda **kwargs: root_builds.append(kwargs)) + + await runner.run_strix_scan( + scan_config={"targets": [], "scan_mode": "deep"}, + scan_id="scan-mcp-list-failure", + image="img", + coordinator=AgentCoordinator(), + mcp_connection_requests=[ + McpConnectionRequest( + config=McpConnectionConfig(name="db", url="https://mcp.example.com") + ) + ], + ) + + assert len(root_builds) == 1 + + +@pytest.mark.asyncio +async def test_warmed_tool_counts_reach_root_and_child_agent_builds( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + _wire_runner(monkeypatch, tmp_path) + server = FakeMCPServer("db", [_mcp_tool(f"tool_{index}") for index in range(3)]) + monkeypatch.setattr(mcp_client, "_build_server", lambda _config: _built_server(server)) + + root_builds: list[dict[str, Any]] = [] + + def _build_root(**kwargs: Any) -> object: + root_builds.append(kwargs) + return object() + + monkeypatch.setattr( + runner, + "build_strix_agent", + _build_root, + ) + child_builds: list[dict[str, Any]] = [] + + def _build_child(**kwargs: Any) -> object: + child_builds.append(kwargs) + return object() + + monkeypatch.setattr( + factory, + "build_strix_agent", + _build_child, + ) + real_make_child_factory = factory.make_child_factory + child_factory_capture: dict[str, Any] = {} + + def _make_child_factory(**kwargs: Any) -> Any: + child_factory_capture["context"] = kwargs["system_prompt_context"] + child_factory_capture["factory"] = real_make_child_factory(**kwargs) + return child_factory_capture["factory"] + + monkeypatch.setattr(runner, "make_child_factory", _make_child_factory) + request = McpConnectionRequest( + config=McpConnectionConfig(name="db", url="https://mcp.example.com", notes="database") + ) + + await runner.run_strix_scan( + scan_config={"targets": [], "scan_mode": "deep"}, + scan_id="scan-mcp-prompt-counts", + image="img", + coordinator=AgentCoordinator(), + mcp_connection_requests=[request], + ) + + root_context = root_builds[0]["system_prompt_context"] + assert root_context["mcp_connections"][0]["tool_count"] == 3 + assert child_factory_capture["context"] is root_context + assert child_factory_capture["context"]["mcp_connections"][0]["tool_count"] == 3 + + child_factory_capture["factory"](name="child", skills=[]) + assert child_builds[0]["system_prompt_context"] is child_factory_capture["context"] + assert child_builds[0]["system_prompt_context"]["mcp_connections"][0]["tool_count"] == 3 diff --git a/tests/test_runner_root_prompt.py b/tests/test_runner_root_prompt.py index e25d09ab..a564a3ea 100644 --- a/tests/test_runner_root_prompt.py +++ b/tests/test_runner_root_prompt.py @@ -6,6 +6,7 @@ flow through to the root agent's ``build_strix_agent`` call. from __future__ import annotations +import importlib import os import types from typing import Any @@ -27,6 +28,11 @@ from strix.core.inputs import make_model_settings from strix.runtime import session_manager from strix.tools.load_skill.tool import load_skill from strix.tools.mcp import BearerAuth, McpConnectionConfig, McpConnectionRequest +from strix.tools.mcp import client as mcp_client + + +_test_mcp_client = importlib.import_module("tests.test_mcp_client") +FakeMCPServer: Any = _test_mcp_client.FakeMCPServer def _make_rate_limit_error() -> RateLimitError: @@ -84,6 +90,11 @@ def _patch_engine_scaffold( monkeypatch.setattr(runner, "build_root_task", lambda _scan_config: "task") monkeypatch.setattr(runner, "build_scope_context", lambda _scan_config: scope_context) monkeypatch.setattr(runner, "make_model_settings", lambda *_args, **_kwargs: ModelSettings()) + monkeypatch.setattr( + mcp_client, + "_build_server", + lambda config: mcp_client.BuiltMcpServer(FakeMCPServer(config.name, []), None), + ) captured: dict[str, Any] = {} @@ -222,7 +233,12 @@ async def test_mcp_available_flag_set_when_a_connection_attaches( assert kwargs["system_prompt_context"]["mcp_available"] is True # The named inventory names each connected server for the prompt. assert kwargs["system_prompt_context"]["mcp_connections"] == [ - {"name": "fs", "purpose": "local files", "tool_count": 0} + { + "name": "fs", + "purpose": "local files", + "tool_count": 0, + "state": "catalog_ready", + } ] From 62b496430da5df5e9e79a190e7af6f92529883a3 Mon Sep 17 00:00:00 2001 From: ian-at-strix Date: Thu, 8 Oct 2026 19:11:50 -0400 Subject: [PATCH 2/3] feat(agents): let callers choose the top-level agent's finish tool (#1496) --- strix/agents/factory.py | 4 ++- strix/core/agents.py | 1 + strix/core/execution.py | 12 +++++---- strix/core/hooks.py | 13 ++++++---- strix/core/runner.py | 6 +++++ strix/tools/agents_graph/tools.py | 3 ++- tests/test_agent_tool_registration.py | 9 +++++++ tests/test_execution.py | 36 +++++++++++++++++++++++++++ tests/test_hooks.py | 16 ++++++++++++ 9 files changed, 88 insertions(+), 12 deletions(-) diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 3b9dd616..4cd9e037 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -675,6 +675,7 @@ def build_strix_agent( extra_tools: Sequence[Tool] | None = None, instructions_override: str | None = None, supports_images: bool = True, + finish_tool: Tool = finish_scan, ) -> SandboxAgent[Any]: """Build a SandboxAgent for either root or child use. @@ -687,6 +688,7 @@ def build_strix_agent( registered via ``register_agent_tools``. instructions_override: Use this verbatim as the system prompt instead of rendering the built-in scan prompt. + finish_tool: The tool that ends the run, given to the root agent. """ if instructions_override is not None: instructions = instructions_override @@ -707,7 +709,7 @@ def build_strix_agent( # Yielding to the user is only meaningful when one is attached. agent_tools.append(wait_for_user) if is_root: - tools: list[Tool] = [*_BASE_TOOLS, *agent_tools, finish_scan] + tools: list[Tool] = [*_BASE_TOOLS, *agent_tools, finish_tool] else: tools = [*_BASE_TOOLS, *agent_tools, agent_finish] _ensure_unique_tool_names(tools) diff --git a/strix/core/agents.py b/strix/core/agents.py index efd7cb60..7e738be1 100644 --- a/strix/core/agents.py +++ b/strix/core/agents.py @@ -71,6 +71,7 @@ class AgentCoordinator: self._lock = asyncio.Lock() self._snapshot_path: Path | None = None self.is_shutting_down = False + self.root_finish_tool = "finish_scan" self._budget_stopped = False self._reserve_stopped = False self._budget_paused = False diff --git a/strix/core/execution.py b/strix/core/execution.py index 5997e3db..87a2bf02 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -302,7 +302,7 @@ async def _run_agent_loop( raise SubagentBudgetReservedError("scan reached the sub-agent budget reserve") if reserve_stopped and start_parked and interactive and context.get("parent_id") is None: - await coordinator.send(agent_id, _reserve_notice()) + await coordinator.send(agent_id, _reserve_notice(coordinator.root_finish_tool)) if not (start_parked and interactive): with contextlib.suppress(BudgetPausedError): @@ -642,6 +642,7 @@ async def _run_until_lifecycle( input_data = await _append_tool_required_message( session=session, context=context, + root_finish_tool=coordinator.root_finish_tool, attempt=recoveries, limit=recovery_limit, interactive=interactive, @@ -988,8 +989,9 @@ async def _append_tool_required_message( limit: int, interactive: bool, silent_yield: bool = False, + root_finish_tool: str = "finish_scan", ) -> list[dict[str, str]]: - finish_tool = "finish_scan" if context.get("parent_id") is None else "agent_finish" + finish_tool = root_finish_tool if context.get("parent_id") is None else "agent_finish" if silent_yield: message = ( "You called wait_for_user without having written anything to the user since " @@ -1107,7 +1109,7 @@ async def notify_parent_on_terminal( ) -def _reserve_notice() -> dict[str, Any]: +def _reserve_notice(finish_tool: str) -> dict[str, Any]: return { "from": "system", "type": "budget_reserve_stop", @@ -1117,7 +1119,7 @@ def _reserve_notice() -> dict[str, Any]: "sub-agent is being force-stopped as soon as its in-flight turn completes, and " "none will send a completion report. Their confirmed vulnerabilities are " "already filed as they were found. Do not wait on any sub-agents and do not " - "spawn new ones — wrap up now and call finish_scan." + f"spawn new ones — wrap up now and call {finish_tool}." ), } @@ -1126,7 +1128,7 @@ async def _notify_root_on_budget_reserve(coordinator: AgentCoordinator) -> None: root = await coordinator.claim_reserve_notification() if root is None: return - await coordinator.send(root, _reserve_notice()) + await coordinator.send(root, _reserve_notice(coordinator.root_finish_tool)) async def _notify_parent_on_exit( diff --git a/strix/core/hooks.py b/strix/core/hooks.py index eae4d7c7..06650be9 100644 --- a/strix/core/hooks.py +++ b/strix/core/hooks.py @@ -104,16 +104,16 @@ _ROOT_DIRECTIVES: tuple[str, ...] = ( ( "As the root agent, begin planning your wind-down of the whole scan: avoid " "starting large new lines of investigation, and keep your required objectives on " - "track so you can call finish_scan comfortably before the limit." + "track so you can call {finish_tool} comfortably before the limit." ), ( "As the root agent, prioritize wrapping up the whole scan now: stop opening new " "lines of investigation, close out only what is essential, and move toward calling " - "finish_scan to compile and deliver the final report." + "{finish_tool} to compile and deliver the final report." ), ( "As the root agent, STOP all other work on the whole scan and finish immediately: " - "secure your findings and call finish_scan now — anything left unfinished when the " + "secure your findings and call {finish_tool} now — anything left unfinished when the " "limit is hit is discarded." ), ) @@ -138,8 +138,11 @@ _SUBAGENT_DIRECTIVES: tuple[str, ...] = ( def _wrapup_directive(context: RunContextWrapper[dict[str, Any]], stage: int) -> str: is_root = context.context.get("parent_id") is None - directives = _ROOT_DIRECTIVES if is_root else _SUBAGENT_DIRECTIVES - return directives[stage] + if not is_root: + return _SUBAGENT_DIRECTIVES[stage] + coordinator = coordinator_from_context(context.context) + finish_tool = coordinator.root_finish_tool if coordinator is not None else "finish_scan" + return _ROOT_DIRECTIVES[stage].format(finish_tool=finish_tool) def _urgency(stage: int) -> str: diff --git a/strix/core/runner.py b/strix/core/runner.py index 58124be2..f9bb4940 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -49,6 +49,7 @@ from strix.report.state import get_global_report_state from strix.runtime import session_manager from strix.telemetry import set_scan_phase from strix.telemetry.logging import set_scan_id, setup_scan_logging +from strix.tools.finish.tool import finish_scan from strix.tools.output_store import ( WORKSPACE_SPILL_DIR, configure_spill_writer, @@ -58,6 +59,7 @@ from strix.tools.output_store import ( if TYPE_CHECKING: from agents.memory import SQLiteSession from agents.result import RunResultBase + from agents.tool import Tool from strix.runtime.status import StatusSink from strix.tools.mcp import ( @@ -215,6 +217,7 @@ async def run_strix_scan( status_sink: StatusSink | None = None, mcp_connection_requests: list[McpConnectionRequest] | None = None, mcp_status_sink: McpStatusSink | None = None, + root_finish_tool: Tool = finish_scan, ) -> RunResultBase | None: """Run or resume one Strix scan against a sandbox. @@ -237,6 +240,7 @@ async def run_strix_scan( command-line default) it reads ``~/.strix/mcp-servers.json`` itself. Either way the engine does the connecting, so the caller passes inert configs plus metadata and never live sessions. + ``root_finish_tool`` is the tool the root agent ends the run with. """ def report(phase: str) -> None: @@ -293,6 +297,7 @@ async def run_strix_scan( coordinator = AgentCoordinator() coordinator.set_snapshot_path(agents_path) coordinator.set_budget_policy(budget_policy) + coordinator.root_finish_tool = root_finish_tool.name from strix.tools.coverage.tools import hydrate_coverage_from_disk from strix.tools.notes.tools import hydrate_notes_from_disk @@ -503,6 +508,7 @@ async def run_strix_scan( system_prompt_context=root_context, instructions_override=root_instructions, supports_images=supports_images, + finish_tool=root_finish_tool, ) if not is_resume: diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index 60c88525..a26c1dc4 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -683,7 +683,8 @@ async def agent_finish( { "success": False, "error": ( - "agent_finish is for subagents. Root/main agents must call finish_scan instead" + "agent_finish is for subagents. Root/main agents must call " + f"{coordinator.root_finish_tool} instead" ), }, ensure_ascii=False, diff --git a/tests/test_agent_tool_registration.py b/tests/test_agent_tool_registration.py index 489fea70..e459670b 100644 --- a/tests/test_agent_tool_registration.py +++ b/tests/test_agent_tool_registration.py @@ -59,6 +59,15 @@ def test_registered_tools_appear_before_lifecycle_tool() -> None: assert child_names[-2:] == ["extra", "agent_finish"] +def test_root_ends_with_the_given_finish_tool() -> None: + root = factory.build_strix_agent(is_root=True, finish_tool=_tool("finish_pr_review")) + + root_names = [t.name for t in root.tools] + + assert root_names[-1] == "finish_pr_review" + assert "finish_scan" not in root_names + + def test_per_call_extra_tools_stack_with_registry() -> None: factory.register_agent_tools(_tool("registered")) diff --git a/tests/test_execution.py b/tests/test_execution.py index c3e1776d..0daf0da2 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -132,6 +132,27 @@ async def test_reserve_stop_notifies_root_once(monkeypatch: pytest.MonkeyPatch) assert "finish_scan" in str(message["content"]) +@pytest.mark.asyncio +async def test_reserve_notice_names_the_root_finish_tool(monkeypatch: pytest.MonkeyPatch) -> None: + coordinator = AgentCoordinator() + coordinator.root_finish_tool = "finish_pr_review" + await coordinator.register("root", "strix", parent_id=None) + await coordinator.register("child-a", "recon", parent_id="root") + + sent: list[dict[str, Any]] = [] + + async def _record(_target_agent_id: str, message: dict[str, Any]) -> bool: + sent.append(message) + return True + + monkeypatch.setattr(coordinator, "send", _record) + + await _notify_root_on_budget_reserve(coordinator) + + assert "call finish_pr_review" in str(sent[0]["content"]) + assert "finish_scan" not in str(sent[0]["content"]) + + @pytest.mark.asyncio async def test_concurrent_reserve_claims_yield_single_root() -> None: coordinator = AgentCoordinator() @@ -1295,6 +1316,21 @@ async def test_interactive_nudge_offers_waiting_without_repeating() -> None: assert "do not repeat it" in items[0]["content"] +@pytest.mark.asyncio +async def test_nudge_names_the_root_finish_tool() -> None: + items = await execution._append_tool_required_message( + session=None, + context={"parent_id": None}, + attempt=1, + limit=3, + interactive=False, + root_finish_tool="finish_pr_review", + ) + + assert "call finish_pr_review" in items[0]["content"] + assert "finish_scan" not in items[0]["content"] + + def _cycle_with_items( coordinator: AgentCoordinator, agent_id: str, diff --git a/tests/test_hooks.py b/tests/test_hooks.py index 16fce426..975a99b8 100644 --- a/tests/test_hooks.py +++ b/tests/test_hooks.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest +from strix.core.agents import AgentCoordinator from strix.core.hooks import ( BudgetExceededError, BudgetPausedError, @@ -335,6 +336,21 @@ async def test_budget_warning_root_directive_distinct_from_subagent() -> None: assert "confirmed" in sub +@pytest.mark.asyncio +async def test_root_directive_names_the_coordinators_finish_tool() -> None: + hooks = ReportUsageHooks(model="test-model", max_turns=100) + coordinator = AgentCoordinator() + coordinator.root_finish_tool = "finish_pr_review" + ctx = _make_warn_context(requests=85, parent_id=None) + ctx.context["coordinator"] = coordinator + + items: list[Any] = [] + await hooks.on_llm_start(ctx, MagicMock(), None, items) + + assert "finish_pr_review" in items[0]["content"] + assert "finish_scan" not in items[0]["content"] + + @pytest.mark.parametrize("parent_id", [None, "root-1"]) @pytest.mark.asyncio async def test_turn_warning_directive_escalates_per_stage(parent_id: str | None) -> None: From f7a32ae62d94872631bc8415073868d02a6ffdaf Mon Sep 17 00:00:00 2001 From: yoni-at-strix Date: Sat, 10 Oct 2026 20:31:50 -0400 Subject: [PATCH 3/3] docs(readme): add update instructions (#1501) --- README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/README.md b/README.md index f93b2376..5134c8bd 100644 --- a/README.md +++ b/README.md @@ -94,6 +94,18 @@ strix --target ./app-directory > [!NOTE] > First run automatically pulls the sandbox Docker image. Results are saved to `strix_runs/` +### Updating + +```bash +# Self-update the standalone binary to the latest release +strix --update + +# ...or re-run the installer +curl -sSL https://strix.ai/install | bash +``` + +Installed with pipx, uv, or pip? `strix --update` prints the matching upgrade command instead (`pipx upgrade strix-agent`, `uv tool upgrade strix-agent`, `pip install --upgrade strix-agent`). Strix also shows a notice with the right command when a new version is available. + --- ## Ways to Run Strix