mirror of
https://github.com/usestrix/strix.git
synced 2026-10-16 04:28:23 +00:00
Merge branch 'usestrix:main' into feat/anthropic-workload-identity
This commit is contained in:
commit
c08256c3ee
16 changed files with 521 additions and 37 deletions
12
README.md
12
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/<run-name>`
|
||||
|
||||
### 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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
@ -68,6 +70,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 +97,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.
|
||||
|
||||
|
|
@ -200,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.
|
||||
|
||||
|
|
@ -222,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:
|
||||
|
|
@ -278,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
|
||||
|
|
@ -438,19 +458,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 +475,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")
|
||||
|
||||
|
|
@ -490,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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue