mirror of
https://hubproxy.babadafafafafa.cn/https://github.com/usestrix/strix.git
synced 2026-09-20 16:13:44 +08:00
Compare commits
5 Commits
fix/user-m
...
mcp-generi
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
eaf4853415 | ||
|
|
c220e1fd31 | ||
|
|
181e83d3d6 | ||
|
|
e3f95d6cfd | ||
|
|
cb5691d6e2 |
@@ -255,11 +255,6 @@ ignore = [
|
||||
"strix/tools/notes/tools.py" = ["PLC0415", "TC002"]
|
||||
"strix/tools/finish/tool.py" = ["PLC0415", "TC002"]
|
||||
"strix/tools/reporting/tool.py" = ["PLC0415", "TC002"]
|
||||
# Lazy imports of strix.tools.mcp.client avoid a circular import (client imports
|
||||
# the session module at module load).
|
||||
"strix/tools/mcp/session.py" = ["PLC0415"]
|
||||
# call_mcp is a chain of guard clauses that each return an error string.
|
||||
"strix/tools/mcp/agent_tools.py" = ["PLR0911"]
|
||||
"strix/tools/**/*.py" = [
|
||||
"ARG001", # Unused function argument (tools may have unused args for interface consistency)
|
||||
]
|
||||
|
||||
@@ -235,7 +235,7 @@ VALIDATION REQUIREMENTS:
|
||||
- CLOSURE DISCIPLINE: every candidate you open ends in exactly one explicit state — `confirmed` (working PoC, or a complete source→control→sink→impact trace that is reachable), `ruled_out` (you can name the SPECIFIC control, at a location, that runs on every attacker-reachable path before the sink), or `open_proof_gap` (plausible, unconfirmed, and you could NOT name such a control). "I moved on" is not a closure state. Silently dropping an uncertain candidate is mislabelling an `open_proof_gap` as `ruled_out` and is how real bugs get missed.
|
||||
- Missing information is NOT proof of safety: no caller found, can't tell if deployed/exposed, couldn't stand up the service, build failed — each is an `open_proof_gap`, never a reason to mark a candidate clean. Difficulty is a reason to defer, not to suppress.
|
||||
- COVERAGE: record every surface you assess with `record_coverage` (surface + risk area + outcome + evidence), including the ones that came back clean — a report that only lists findings cannot say what was reviewed and cleared. Use the `needs_follow_up` outcome for anything left in an `open_proof_gap` state, and carry the same items up in `agent_finish(open_items=[...])`. The ledger is shared and mutable: when you resolve a surface another agent left open — or find that a closed one is not — move that entry with `update_coverage` instead of recording a second one for the same surface. The root agent reconciles all of it via `list_coverage` before `finish_scan`.
|
||||
- THREAT MODEL: before you start testing, call `get_threat_model` on the target you were pointed at — it is the scan's shared answer to who the attacker is, where the trust boundaries sit, and what counts as critical here. It is scoped to this scan and nothing carries over from an earlier run, so `found: false` means no agent on this run has derived one yet. Read it instead of re-deriving trust boundaries yourself; where your testing disproves it — a boundary it calls trusted turns out to be attacker-reachable, a role it did not know about, a host or endpoint it never listed — record that with `amend_threat_model` so the agents after you inherit the correction. Amending is not optional politeness: a model nobody corrects turns the first agent's guesses into everyone's assumptions.
|
||||
- THREAT MODEL: before you start testing, call `get_threat_model` on the target you were pointed at — it is the scan's shared answer to who the attacker is, where the trust boundaries sit, and what counts as critical here, and it is cached per target rather than per scan. Read it instead of re-deriving trust boundaries yourself; where your testing disproves it — a boundary it calls trusted turns out to be attacker-reachable, a role it did not know about, a host or endpoint it never listed — record that with `amend_threat_model` so the agents after you inherit the correction. Amending is not optional politeness: a model nobody corrects turns the first agent's guesses into everyone's assumptions.
|
||||
- Before filing any report, run the counterevidence pass: argue the strongest case AGAINST the finding, record what you found in the `counterevidence` field, set `confidence` honestly (a static-only trace you couldn't execute is at best `medium`), and state what evidence would change the severity. See the counterevidence and severity-calibration knowledge above.
|
||||
- A vulnerability is ONLY considered reported when a reporting agent uses create_vulnerability_report (or create_dependency_report for known-CVE dependency/supply-chain findings) with full details. Mentions in agent_finish, finish_scan, or generic messages are NOT sufficient
|
||||
- Reporting and fixing are ONE step, not two: when source is available, the reporting agent derives the concrete fix and files it INLINE via create_vulnerability_report (`code_locations` with `fix_before`/`fix_after` + `fix_pr_body`) — the report is not complete without it. Do NOT report first and then spawn a separate downstream agent to re-derive and re-apply the same patch; that just re-does the analysis and wastes tokens. (Do not silently patch a finding WITHOUT filing a report — the report, with its embedded fix, is the deliverable.)
|
||||
|
||||
@@ -18,7 +18,7 @@ from agents import (
|
||||
)
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.fake_id import FAKE_RESPONSES_ID
|
||||
from agents.models.interface import Model, ModelProvider
|
||||
from agents.models.interface import Model
|
||||
from agents.models.multi_provider import MultiProvider
|
||||
from agents.models.openai_responses import OpenAIResponsesModel
|
||||
from agents.retry import (
|
||||
@@ -48,7 +48,7 @@ if TYPE_CHECKING:
|
||||
from agents.agent_output import AgentOutputSchemaBase
|
||||
from agents.handoffs import Handoff
|
||||
from agents.items import ModelResponse, TResponseInputItem, TResponseStreamEvent
|
||||
from agents.models.interface import ModelTracing
|
||||
from agents.models.interface import ModelProvider, ModelTracing
|
||||
from agents.retry import ModelRetryAdvice, ModelRetryAdviceRequest
|
||||
from agents.tool import Tool
|
||||
from agents.usage import Usage
|
||||
@@ -445,61 +445,12 @@ def _response_usage(usage: Usage | None) -> ResponseUsage | None:
|
||||
)
|
||||
|
||||
|
||||
class _CredentialedLitellmProvider(ModelProvider):
|
||||
"""LiteLLM route bound to one endpoint's credentials.
|
||||
|
||||
``LitellmProvider`` reads them from the process-wide LiteLLM globals, which
|
||||
belong to the main model; a secondary endpoint needs its own.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str | None, base_url: str | None) -> None:
|
||||
self._api_key = api_key
|
||||
self._base_url = base_url
|
||||
|
||||
def get_model(self, model_name: str | None) -> Model:
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
from agents.models.default_models import get_default_model
|
||||
|
||||
return LitellmModel(
|
||||
model=model_name or get_default_model(),
|
||||
api_key=self._api_key,
|
||||
base_url=self._base_url,
|
||||
)
|
||||
|
||||
|
||||
class StrixProvider(MultiProvider):
|
||||
"""Route any non-OpenAI prefix through LiteLLM with the prefix preserved,
|
||||
so users type ``deepseek/deepseek-chat`` rather than
|
||||
``litellm/deepseek/deepseek-chat``.
|
||||
|
||||
``api_key``/``base_url`` bind every route this provider resolves to one
|
||||
endpoint, for a secondary model (the dedupe judge) whose endpoint differs
|
||||
from the main model's process-wide defaults.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
openai_api_key=api_key,
|
||||
openai_base_url=base_url,
|
||||
# A custom endpoint is OpenAI-compatible, i.e. chat completions; the
|
||||
# global default is the main model's and may say otherwise.
|
||||
openai_use_responses=False if base_url else None,
|
||||
**kwargs,
|
||||
)
|
||||
self._override_api_key = api_key
|
||||
self._override_base_url = base_url
|
||||
|
||||
def _create_fallback_provider(self, prefix: str) -> ModelProvider:
|
||||
if prefix == "litellm" and (self._override_api_key or self._override_base_url):
|
||||
return _CredentialedLitellmProvider(self._override_api_key, self._override_base_url)
|
||||
return super()._create_fallback_provider(prefix)
|
||||
|
||||
def _resolve_prefixed_model(
|
||||
self,
|
||||
*,
|
||||
@@ -913,22 +864,6 @@ def is_claude_model(model_name: str) -> bool:
|
||||
return "claude" in (model_name or "").strip().lower()
|
||||
|
||||
|
||||
def routes_through_litellm(model_name: str | None) -> bool:
|
||||
"""Whether :class:`StrixProvider` sends this model through LiteLLM.
|
||||
|
||||
Bare names and the ``openai/``/``any-llm/`` prefixes are served by the SDK's
|
||||
own clients, which raise ``TypeError`` on request fields they do not know,
|
||||
so LiteLLM-only fields must not be attached there. A bare ``claude-...``
|
||||
name is exactly that case: an ``LLM_API_BASE`` pointing at an
|
||||
OpenAI-compatible gateway in front of Claude.
|
||||
"""
|
||||
name = (model_name or "").strip()
|
||||
if not name or codex.subscription_model(name):
|
||||
return False
|
||||
prefix, _, rest = name.partition("/")
|
||||
return bool(rest) and prefix.lower() not in {"openai", "any-llm"}
|
||||
|
||||
|
||||
def is_bedrock_route(model_name: str) -> bool:
|
||||
name = (model_name or "").strip().lower()
|
||||
return name.startswith("bedrock/") or "anthropic." in name
|
||||
|
||||
@@ -291,12 +291,6 @@ class AgentCoordinator:
|
||||
self.pending_counts[target_agent_id] = self.pending_counts.get(target_agent_id, 0) + 1
|
||||
if from_user:
|
||||
runtime.user_wake_required = False
|
||||
self.errors.pop(target_agent_id, None)
|
||||
self.wait_kinds.pop(target_agent_id, None)
|
||||
self.recovery_counts.pop(target_agent_id, None)
|
||||
self.idle_resume_counts.pop(target_agent_id, None)
|
||||
self._parent_notified.discard(target_agent_id)
|
||||
self.statuses[target_agent_id] = "waiting"
|
||||
runtime.wake.set()
|
||||
stream = runtime.stream
|
||||
interrupt_on_message = runtime.interrupt_on_message
|
||||
|
||||
@@ -18,7 +18,6 @@ from strix.config.models import (
|
||||
is_openrouter_model,
|
||||
model_supports_reasoning,
|
||||
request_timeout_extra_args,
|
||||
routes_through_litellm,
|
||||
)
|
||||
from strix.core.sessions import scrub_images_from_items
|
||||
|
||||
@@ -268,7 +267,7 @@ def make_model_settings(
|
||||
and model_supports_reasoning(model_name)
|
||||
):
|
||||
model_settings = model_settings.resolve(
|
||||
_reasoning_settings(reasoning_effort),
|
||||
_reasoning_settings(reasoning_effort, model_settings.extra_args),
|
||||
)
|
||||
if force_required_tool_choice and _accepts_required_tool_choice(model_name):
|
||||
model_settings = model_settings.resolve(ModelSettings(tool_choice="required"))
|
||||
@@ -294,19 +293,20 @@ def _request_headers(
|
||||
return headers or None
|
||||
|
||||
|
||||
def _reasoning_settings(effort: ReasoningEffort) -> ModelSettings:
|
||||
def _reasoning_settings(
|
||||
effort: ReasoningEffort,
|
||||
extra_args: dict[str, Any] | None,
|
||||
) -> ModelSettings:
|
||||
"""``max`` is not in the OpenAI SDK's ``Reasoning.effort`` enum, so send it as
|
||||
a raw body field instead — also keeping it clear of LiteLLM's DeepSeek mapping,
|
||||
which collapses every ``reasoning_effort`` level to plain thinking-enabled.
|
||||
Providers that don't support ``max`` reject the request.
|
||||
|
||||
It goes in ``extra_body``, the field every model implementation forwards as the
|
||||
request's ``extra_body``; the same value under ``extra_args`` collides with that
|
||||
keyword and raises before a request is ever sent.
|
||||
"""
|
||||
if effort != "max":
|
||||
return ModelSettings(reasoning=Reasoning(effort=effort))
|
||||
return ModelSettings(extra_body={"reasoning_effort": "max"})
|
||||
return ModelSettings(
|
||||
extra_args={**(extra_args or {}), "extra_body": {"reasoning_effort": "max"}},
|
||||
)
|
||||
|
||||
|
||||
def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None:
|
||||
@@ -317,13 +317,8 @@ def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None:
|
||||
it — elsewhere it leaks onto the wire and native Anthropic 400s). Unmapped
|
||||
Bedrock models get no points at all: Bedrock rejects the passed-through
|
||||
field outright.
|
||||
|
||||
The field is LiteLLM's own, consumed by its transform, so it only goes to
|
||||
routes LiteLLM serves. A bare ``claude-...`` name is served by the SDK's
|
||||
OpenAI client instead (a gateway in front of Claude), and that client raises
|
||||
``TypeError`` on request kwargs it does not know.
|
||||
"""
|
||||
if not is_claude_model(model_name) or not routes_through_litellm(model_name):
|
||||
if not is_claude_model(model_name):
|
||||
return None
|
||||
if is_bedrock_route(model_name) and not bedrock_route_supports_prompt_caching(model_name):
|
||||
return None
|
||||
|
||||
@@ -53,43 +53,18 @@ from strix.tools.output_store import (
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.mcp import MCPServer
|
||||
from agents.memory import SQLiteSession
|
||||
from agents.result import RunResultBase
|
||||
|
||||
from strix.runtime.status import StatusSink
|
||||
from strix.tools.mcp import (
|
||||
ConnectedMcpServer,
|
||||
McpConnectionRequest,
|
||||
McpRegistry,
|
||||
SupervisedMcpSession,
|
||||
)
|
||||
from strix.tools.mcp import ConnectedMcpServer, McpConnectionRequest
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
StreamEventSink = Callable[[str, Any], None]
|
||||
|
||||
# Receives the run's MCP connection roster as a list of non-secret status dicts
|
||||
# ({"name", "provider", "tool_count", "dead"}), once when the connections are
|
||||
# established and again each time a connection transitions to dead. An interface
|
||||
# can persist it, render it, or forward it on as connection status. Kept as a
|
||||
# snapshot of the whole roster (not a per-
|
||||
# connection delta) so every call carries a consistent, current picture.
|
||||
McpStatusSink = Callable[[list[dict[str, Any]]], None]
|
||||
|
||||
|
||||
def _mcp_roster_payload(registry: McpRegistry) -> list[dict[str, Any]]:
|
||||
"""The run's MCP roster as non-secret status dicts (name/provider/tool_count/dead)."""
|
||||
return [
|
||||
{
|
||||
"name": status.name,
|
||||
"provider": status.provider,
|
||||
"tool_count": status.tool_count,
|
||||
"dead": status.dead,
|
||||
}
|
||||
for status in registry.statuses()
|
||||
]
|
||||
|
||||
|
||||
def _mcp_startup_summary(connections: list[ConnectedMcpServer]) -> str:
|
||||
"""One user-facing line summarizing the MCP servers that connected."""
|
||||
@@ -116,21 +91,6 @@ def _record_mcp_connections(connections: list[ConnectedMcpServer]) -> None:
|
||||
report_state.record_mcp_connections([connection.name for connection in connections])
|
||||
|
||||
|
||||
def _persist_mcp_status(roster: list[dict[str, Any]]) -> None:
|
||||
"""Write the run's non-secret MCP connection status roster to run.json.
|
||||
|
||||
The viewer rebuilds its display by re-reading the run's files from disk, so
|
||||
it cannot see the in-memory ``mcp_status_sink`` the TUI consumes. Persisting
|
||||
the same non-secret roster (name / provider / tool_count / dead) gives the
|
||||
viewer a source it can poll. Runs regardless of whether an interface sink is
|
||||
attached, so the standalone / non-TUI CLI path records health too.
|
||||
"""
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
report_state.record_mcp_connection_status(roster)
|
||||
|
||||
|
||||
def _merge_root_prompt_context(
|
||||
scope_context: dict[str, Any],
|
||||
extra_system_prompt_context: dict[str, Any] | None,
|
||||
@@ -197,7 +157,6 @@ async def run_strix_scan(
|
||||
extra_system_prompt_context: dict[str, Any] | None = None,
|
||||
status_sink: StatusSink | None = None,
|
||||
mcp_connection_requests: list[McpConnectionRequest] | None = None,
|
||||
mcp_status_sink: McpStatusSink | None = None,
|
||||
) -> RunResultBase | None:
|
||||
"""Run or resume one Strix scan against a sandbox.
|
||||
|
||||
@@ -263,13 +222,11 @@ async def run_strix_scan(
|
||||
|
||||
from strix.tools.coverage.tools import hydrate_coverage_from_disk
|
||||
from strix.tools.notes.tools import hydrate_notes_from_disk
|
||||
from strix.tools.threat_model.tools import hydrate_threat_models_from_disk
|
||||
from strix.tools.todo.tools import hydrate_todos_from_disk
|
||||
|
||||
hydrate_todos_from_disk(state_dir)
|
||||
hydrate_notes_from_disk(state_dir)
|
||||
hydrate_coverage_from_disk(state_dir)
|
||||
hydrate_threat_models_from_disk(state_dir)
|
||||
|
||||
root_id: str | None = None
|
||||
if is_resume:
|
||||
@@ -338,7 +295,7 @@ async def run_strix_scan(
|
||||
configure_spill_writer(_spill_to_workspace)
|
||||
|
||||
sessions_to_close: list[SQLiteSession] = []
|
||||
mcp_sessions: list[SupervisedMcpSession] = []
|
||||
mcp_servers: list[MCPServer] = []
|
||||
|
||||
try:
|
||||
targets = scan_config.get("targets") or []
|
||||
@@ -407,7 +364,7 @@ async def run_strix_scan(
|
||||
mcp_requests = mcp_connection_requests
|
||||
if mcp_requests:
|
||||
connections = await attach_mcp_requests(mcp_requests, mcp_registry)
|
||||
mcp_sessions = [c.session for c in connections]
|
||||
mcp_servers = [c.server for c in connections]
|
||||
# Recorded even when nothing connected, so a resumed run does not
|
||||
# keep attributing tool calls to servers it no longer has.
|
||||
_record_mcp_connections(connections)
|
||||
@@ -428,31 +385,6 @@ async def run_strix_scan(
|
||||
}
|
||||
for summary in mcp_registry.summaries()
|
||||
]
|
||||
# Feed a non-secret connection roster (name / provider /
|
||||
# tool_count / dead) to two consumers: once now (all
|
||||
# currently healthy) and again whenever a connection later
|
||||
# dies. It is always persisted to run.json so the viewer,
|
||||
# which re-reads the run's files from disk, can render the
|
||||
# MCP connections panel and health without an in-memory
|
||||
# sink. When an interface sink is attached (the TUI backend,
|
||||
# or pro forwarding into the app's event stream) it also
|
||||
# receives the same snapshot. In-use is derived separately by
|
||||
# each interface from the connection-tagged tool-call events,
|
||||
# so it is not carried here.
|
||||
def _emit_mcp_status() -> None:
|
||||
roster = _mcp_roster_payload(mcp_registry)
|
||||
_persist_mcp_status(roster)
|
||||
if mcp_status_sink is not None:
|
||||
try:
|
||||
mcp_status_sink(roster)
|
||||
except Exception:
|
||||
logger.exception("MCP status sink failed")
|
||||
|
||||
for connection_name in mcp_registry.names():
|
||||
entry = mcp_registry.get(connection_name)
|
||||
if entry is not None:
|
||||
entry.session.set_on_dead(_emit_mcp_status)
|
||||
_emit_mcp_status()
|
||||
except Exception:
|
||||
logger.exception("Failed to connect user MCP servers; continuing without them")
|
||||
|
||||
@@ -647,9 +579,9 @@ async def run_strix_scan(
|
||||
for s in sessions_to_close:
|
||||
with contextlib.suppress(Exception):
|
||||
s.close()
|
||||
for mcp_session in mcp_sessions:
|
||||
for mcp_server in mcp_servers:
|
||||
with contextlib.suppress(Exception):
|
||||
await mcp_session.aclose()
|
||||
await mcp_server.cleanup() # type: ignore[no-untyped-call]
|
||||
with contextlib.suppress(Exception):
|
||||
await coordinator._maybe_snapshot()
|
||||
if cleanup_on_exit:
|
||||
|
||||
@@ -127,10 +127,12 @@ def _subscription_error_hint(exc: BaseException) -> str | None:
|
||||
|
||||
|
||||
async def warm_up_llm(show_model_warning: bool = True) -> None:
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import ModelTracing
|
||||
|
||||
from strix.config.models import (
|
||||
RECOMMENDED_MODEL_NAMES,
|
||||
StrixProvider,
|
||||
configure_sdk_model_defaults,
|
||||
is_known_openai_bare_model,
|
||||
is_recommended_or_frontier_model,
|
||||
@@ -207,11 +209,12 @@ async def warm_up_llm(show_model_warning: bool = True) -> None:
|
||||
logger.info("LLM warm-up succeeded for model %s", (llm.model or "").strip())
|
||||
|
||||
if settings.dedupe.model:
|
||||
from strix.report.dedupe import resolve_dedupe_model
|
||||
from strix.report.dedupe import _dedupe_extra_args
|
||||
|
||||
dedupe_model = settings.dedupe.model.strip()
|
||||
raw_model = dedupe_model
|
||||
deduper = resolve_dedupe_model(settings.dedupe, dedupe_model)
|
||||
deduper = StrixProvider().get_model(dedupe_model)
|
||||
deduper_extra = _dedupe_extra_args(settings.dedupe)
|
||||
# A dedicated dedupe model may route to another provider, which must
|
||||
# never receive the main endpoint's headers; it has its own
|
||||
# DEDUPE_LLM_EXTRA_HEADERS.
|
||||
@@ -223,6 +226,9 @@ async def warm_up_llm(show_model_warning: bool = True) -> None:
|
||||
extra_headers=settings.dedupe.extra_headers,
|
||||
has_tools=False,
|
||||
)
|
||||
if deduper_extra:
|
||||
merged = {**(deduper_settings.extra_args or {}), **deduper_extra}
|
||||
deduper_settings = deduper_settings.resolve(ModelSettings(extra_args=merged))
|
||||
await asyncio.wait_for(
|
||||
deduper.get_response(
|
||||
system_instructions="You are a helpful assistant.",
|
||||
|
||||
@@ -103,11 +103,6 @@ class TuiController:
|
||||
self.messages: list[dict[str, str]] = []
|
||||
self._next_message_id = 1
|
||||
self.error: str | None = None
|
||||
# The run's MCP connection roster (name / tool_count / dead), pushed by
|
||||
# the engine via the mcp_status_sink once the connections are established
|
||||
# and again each time one dies. Empty for a run with no MCP connections,
|
||||
# so the Go sidebar simply omits the panel. Non-secret by construction.
|
||||
self.mcp_connections: list[dict[str, Any]] = []
|
||||
self.viewer_status = "idle"
|
||||
self.viewer_url: str | None = None
|
||||
self._viewer_httpd: Any = None
|
||||
@@ -133,24 +128,6 @@ class TuiController:
|
||||
if scan_loop is not None:
|
||||
self.scan_loop = scan_loop
|
||||
|
||||
def set_mcp_connections(self, roster: list[dict[str, Any]]) -> None:
|
||||
"""Store the run's MCP connection roster and repaint.
|
||||
|
||||
``roster`` is the engine's non-secret status snapshot: one entry per
|
||||
connection carrying ``name``, ``tool_count``, and ``dead``. Called once
|
||||
when the connections are established (all healthy) and again whenever a
|
||||
connection dies (the same whole-roster snapshot, with that one now dead)."""
|
||||
self.mcp_connections = [
|
||||
{
|
||||
"name": str(entry.get("name", "")),
|
||||
"tool_count": int(entry.get("tool_count", 0) or 0),
|
||||
"dead": bool(entry.get("dead", False)),
|
||||
}
|
||||
for entry in roster
|
||||
if isinstance(entry, dict) and entry.get("name")
|
||||
]
|
||||
self.notify_changed()
|
||||
|
||||
def begin_preparation(self) -> None:
|
||||
"""Mark a directly-launched run as preparing behind the live TUI."""
|
||||
self.scan_state = "preparing"
|
||||
@@ -223,14 +200,6 @@ class TuiController:
|
||||
],
|
||||
"usage": terminal_projection(usage, max_string=256, max_items=20),
|
||||
"subscription": subscription,
|
||||
"connections": [
|
||||
{
|
||||
"name": terminal_projection(entry["name"], max_string=64),
|
||||
"tool_count": entry["tool_count"],
|
||||
"dead": entry["dead"],
|
||||
}
|
||||
for entry in self.mcp_connections[:32]
|
||||
],
|
||||
"viewer_status": self.viewer_status,
|
||||
"viewer_url": terminal_projection(self.viewer_url, max_string=1024),
|
||||
"error": terminal_projection(self.error, max_string=2 * 1024),
|
||||
@@ -411,7 +380,6 @@ class TuiController:
|
||||
delivered = await asyncio.wrap_future(future)
|
||||
if not delivered:
|
||||
raise RuntimeError("Message could not be delivered")
|
||||
self.live_view.upsert_agent(agent_id, status="waiting", error_message=None)
|
||||
return {"sent": True}
|
||||
|
||||
async def _stop_agent(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
@@ -60,9 +60,6 @@ class TuiLiveView(BaseLiveView):
|
||||
if error_message and current.get("error_message") != error_message:
|
||||
current["error_message"] = error_message
|
||||
changed = True
|
||||
elif error_message is None and "error_message" in current:
|
||||
current.pop("error_message", None)
|
||||
changed = True
|
||||
if changed:
|
||||
current["updated_at"] = now
|
||||
return changed
|
||||
|
||||
@@ -164,21 +164,19 @@ def bounded_state_projection(state: dict[str, Any]) -> dict[str, Any]:
|
||||
"scan_state": state["scan_state"],
|
||||
"targets": state["targets"][:4],
|
||||
"target_count": state["target_count"],
|
||||
"working_dir": terminal_projection(state.get("working_dir", ""), max_string=256),
|
||||
"pending_mount": terminal_projection(state.get("pending_mount", ""), max_string=256),
|
||||
"instruction": terminal_projection(state["instruction"], max_string=128),
|
||||
"scan_mode": state["scan_mode"],
|
||||
"max_budget_usd": state["max_budget_usd"],
|
||||
"max_turns": state["max_turns"],
|
||||
"scope_mode": state["scope_mode"],
|
||||
"diff_base": state["diff_base"],
|
||||
"provider": state["provider"],
|
||||
"model": state["model"],
|
||||
"model_warning": "",
|
||||
"caido_url": None,
|
||||
"messages": [],
|
||||
"usage": state["usage"],
|
||||
"subscription": state["subscription"],
|
||||
"connections": state.get("connections", [])[:32],
|
||||
"viewer_status": state["viewer_status"],
|
||||
"viewer_url": None,
|
||||
"error": terminal_projection(state["error"], max_string=256),
|
||||
|
||||
@@ -209,7 +209,7 @@ func (m *Model) ensureAgentVisible() {
|
||||
m.agentOffset = 0
|
||||
return
|
||||
}
|
||||
_, _, _, agentHeight := m.sidebarHeights()
|
||||
_, _, agentHeight := m.sidebarHeights()
|
||||
rows := max(1, agentHeight-4)
|
||||
row := selectedAgentRow(entries, m.selectedAgent)
|
||||
if row < m.agentOffset {
|
||||
@@ -221,7 +221,7 @@ func (m *Model) ensureAgentVisible() {
|
||||
}
|
||||
|
||||
func (m Model) agentPageSize() int {
|
||||
_, _, _, agentHeight := m.sidebarHeights()
|
||||
_, _, agentHeight := m.sidebarHeights()
|
||||
return max(1, agentHeight-4)
|
||||
}
|
||||
|
||||
|
||||
@@ -1,105 +0,0 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/charmbracelet/x/ansi"
|
||||
"github.com/usestrix/strix/tui/internal/protocol"
|
||||
)
|
||||
|
||||
func mcpModel(t *testing.T) Model {
|
||||
t.Helper()
|
||||
m := New(nil)
|
||||
m.width, m.height = 130, 40
|
||||
m.showSplash = false
|
||||
m.handleEnvelope(stateEnvelope(t, 1, protocol.Snapshot{
|
||||
ScanState: "running",
|
||||
Connections: []protocol.Connection{
|
||||
{Name: "supabase", ToolCount: 3, Dead: false},
|
||||
{Name: "vercel", ToolCount: 1, Dead: true},
|
||||
},
|
||||
}))
|
||||
return m
|
||||
}
|
||||
|
||||
func TestMcpPanelShowsHealthyAndOffline(t *testing.T) {
|
||||
m := mcpModel(t)
|
||||
out := ansi.Strip(m.mcpConnectionsView(40, 6))
|
||||
for _, want := range []string{"MCP Connections (2)", "supabase", "3 tools", "vercel", "offline"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Fatalf("panel missing %q:\n%s", want, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A roster longer than the panel height shows a window of rows rather than every
|
||||
// connection, while the header keeps the full count.
|
||||
func TestMcpPanelWindowsLargeRosterAndCountsAll(t *testing.T) {
|
||||
m := New(nil)
|
||||
m.width, m.height = 130, 40
|
||||
m.showSplash = false
|
||||
conns := make([]protocol.Connection, 0, 12)
|
||||
for i := 0; i < 12; i++ {
|
||||
conns = append(conns, protocol.Connection{Name: fmt.Sprintf("conn-%02d", i), ToolCount: 2})
|
||||
}
|
||||
m.snapshot.Connections = conns
|
||||
|
||||
// rows = 6 → one header line + five roster rows.
|
||||
out := ansi.Strip(m.mcpConnectionsView(40, 6))
|
||||
if !strings.Contains(out, "MCP Connections (12)") {
|
||||
t.Fatalf("header did not carry the full connection count:\n%s", out)
|
||||
}
|
||||
if !strings.Contains(out, "conn-00") {
|
||||
t.Fatalf("top of the roster was not rendered:\n%s", out)
|
||||
}
|
||||
if strings.Contains(out, "conn-11") {
|
||||
t.Fatalf("a roster past the panel height should be windowed, not fully drawn:\n%s", out)
|
||||
}
|
||||
if got := strings.Count(out, "\n") + 1; got != 6 {
|
||||
t.Fatalf("panel rendered %d lines, want 6 (header + five rows)", got)
|
||||
}
|
||||
|
||||
// Scrolling the roster brings the tail into view while the header count holds.
|
||||
m.mcpOffset = 7
|
||||
scrolled := ansi.Strip(m.mcpConnectionsView(40, 6))
|
||||
if !strings.Contains(scrolled, "conn-11") || !strings.Contains(scrolled, "MCP Connections (12)") {
|
||||
t.Fatalf("scrolled window did not reveal the tail with the count intact:\n%s", scrolled)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpPanelHeightReservedFromAgentBudget(t *testing.T) {
|
||||
m := mcpModel(t)
|
||||
_, _, mcpHeight, _ := m.sidebarHeights()
|
||||
if mcpHeight <= 0 {
|
||||
t.Fatalf("connections present but no panel height was reserved: %d", mcpHeight)
|
||||
}
|
||||
|
||||
empty := New(nil)
|
||||
empty.width, empty.height = 130, 40
|
||||
empty.showSplash = false
|
||||
empty.handleEnvelope(stateEnvelope(t, 1, protocol.Snapshot{ScanState: "running"}))
|
||||
if _, _, emptyHeight, _ := empty.sidebarHeights(); emptyHeight != 0 {
|
||||
t.Fatalf("no connections should leave the panel absent, got height %d", emptyHeight)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpInUseReadsRunningConnectionTaggedCalls(t *testing.T) {
|
||||
m := mcpModel(t)
|
||||
m.handleEnvelope(bootstrapEnvelope(t, "events", 1,
|
||||
protocol.Event{ID: "e1", Type: "tool", AgentID: "a1", Data: map[string]any{
|
||||
"tool_name": "call_mcp", "mcp_connection": "supabase", "status": "running",
|
||||
}},
|
||||
protocol.Event{ID: "e2", Type: "tool", AgentID: "a1", Data: map[string]any{
|
||||
"tool_name": "call_mcp", "mcp_connection": "vercel", "status": "completed",
|
||||
}},
|
||||
))
|
||||
inUse := m.mcpInUse()
|
||||
if !inUse["supabase"] {
|
||||
t.Fatalf("a running connection-tagged call should mark the connection in use")
|
||||
}
|
||||
if inUse["vercel"] {
|
||||
t.Fatalf("a completed call must not mark the connection in use")
|
||||
}
|
||||
}
|
||||
@@ -73,7 +73,6 @@ const (
|
||||
focusChat
|
||||
focusAgents
|
||||
focusVulnerabilities
|
||||
focusMcp
|
||||
)
|
||||
|
||||
type scrollbarTarget int
|
||||
@@ -83,7 +82,6 @@ const (
|
||||
scrollbarTrace
|
||||
scrollbarAgents
|
||||
scrollbarFindings
|
||||
scrollbarMcp
|
||||
)
|
||||
|
||||
type Model struct {
|
||||
@@ -111,7 +109,6 @@ type Model struct {
|
||||
selectedVuln int
|
||||
agentOffset int
|
||||
vulnOffset int
|
||||
mcpOffset int
|
||||
modalChoice int
|
||||
reportFocus string
|
||||
ready bool
|
||||
@@ -359,9 +356,7 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.resyncRequested[msg.collection] = false
|
||||
}
|
||||
} else if msg.command == "collection.resync" && msg.requestID != "" && msg.collection != "" {
|
||||
if m.resyncRequested[msg.collection] {
|
||||
m.resyncRequests[msg.requestID] = msg.collection
|
||||
}
|
||||
m.resyncRequests[msg.requestID] = msg.collection
|
||||
}
|
||||
case selectionCopiedMsg:
|
||||
text := "Copied to clipboard"
|
||||
|
||||
@@ -103,21 +103,6 @@ func bootstrapEnvelope(t *testing.T, collection string, revision int, items ...a
|
||||
return protocol.Envelope{Version: protocol.Version, Type: "collection_bootstrap", Payload: rawJSON(t, payload)}
|
||||
}
|
||||
|
||||
func TestStateSnapshotClearsNilError(t *testing.T) {
|
||||
model := New(nil)
|
||||
errText := "provider rejected"
|
||||
model.handleEnvelope(stateEnvelope(t, 1, protocol.Snapshot{ScanState: "failed", Error: &errText}))
|
||||
if model.errorText != errText {
|
||||
t.Fatalf("error was not installed: %q", model.errorText)
|
||||
}
|
||||
|
||||
model.handleEnvelope(stateEnvelope(t, 2, protocol.Snapshot{ScanState: "running"}))
|
||||
|
||||
if model.errorText != "" {
|
||||
t.Fatalf("nil snapshot error did not clear errorText: %q", model.errorText)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackendDisconnectBecomesFatalUnlessUserIsQuitting(t *testing.T) {
|
||||
model := New(nil)
|
||||
updated, cmd := model.Update(wireErrMsg{err: fmt.Errorf("socket closed")})
|
||||
@@ -175,27 +160,6 @@ func TestCollectionBootstrapChunksAndVersionedDelta(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentCollectionDeltaClearsErrorMessage(t *testing.T) {
|
||||
model := New(nil)
|
||||
failed := protocol.Agent{ID: "root", Name: "Strix", Status: "failed", ErrorMessage: "provider rejected"}
|
||||
model.handleEnvelope(bootstrapEnvelope(t, "agents", 1, failed))
|
||||
|
||||
resumed := protocol.Agent{ID: "root", Name: "Strix", Status: "waiting"}
|
||||
delta := protocol.CollectionDelta{
|
||||
Collection: "agents", BaseRevision: 1, Revision: 2, Cursor: 0, NextCursor: 1, Done: true,
|
||||
Operations: []protocol.CollectionOperation{{Op: "upsert", Item: rawJSON(t, resumed)}},
|
||||
}
|
||||
model.handleEnvelope(protocol.Envelope{Version: protocol.Version, Type: "collection_delta", Payload: rawJSON(t, delta)})
|
||||
|
||||
if len(model.snapshot.Agents) != 1 {
|
||||
t.Fatalf("agents were not retained: %#v", model.snapshot.Agents)
|
||||
}
|
||||
agent := model.snapshot.Agents[0]
|
||||
if agent.Status != "waiting" || agent.ErrorMessage != "" {
|
||||
t.Fatalf("agent error was not cleared: %#v", agent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectionMismatchRequestsOneResync(t *testing.T) {
|
||||
connection := &recordingConn{}
|
||||
model := New(newClient(connection))
|
||||
@@ -220,42 +184,6 @@ func TestCollectionMismatchRequestsOneResync(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailedResyncResultBeforeSentMsgRearmsResync(t *testing.T) {
|
||||
connection := &recordingConn{}
|
||||
model := New(newClient(connection))
|
||||
model.collectionRevisions["events"] = 4
|
||||
bad := protocol.CollectionDelta{
|
||||
Collection: "events", BaseRevision: 2, Revision: 3, Cursor: 0, NextCursor: 0, Done: true,
|
||||
}
|
||||
|
||||
cmd := model.handleEnvelope(protocol.Envelope{Version: protocol.Version, Type: "collection_delta", Payload: rawJSON(t, bad)})
|
||||
if cmd == nil {
|
||||
t.Fatal("revision mismatch did not request a resync")
|
||||
}
|
||||
sent, ok := cmd().(sentMsg)
|
||||
if !ok || sent.err != nil || sent.requestID == "" {
|
||||
t.Fatalf("resync send = %#v", sent)
|
||||
}
|
||||
|
||||
failed := protocol.CommandResult{
|
||||
OK: false,
|
||||
Command: "collection.resync",
|
||||
Error: &protocol.CommandError{Code: "command_failed", Message: "resync failed"},
|
||||
}
|
||||
model.handleEnvelope(protocol.Envelope{
|
||||
Version: protocol.Version, Type: "command_result", RequestID: sent.requestID, Payload: rawJSON(t, failed),
|
||||
})
|
||||
updated, _ := model.Update(sent)
|
||||
model = updated.(Model)
|
||||
|
||||
if model.resyncRequested["events"] {
|
||||
t.Fatal("failed resync result left resync suppressed")
|
||||
}
|
||||
if retry := model.handleEnvelope(protocol.Envelope{Version: protocol.Version, Type: "collection_delta", Payload: rawJSON(t, bad)}); retry == nil {
|
||||
t.Fatal("resync was not rearmed after failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentsCollectionPreservesSelectedIDAcrossUpsertsAndDeletes(t *testing.T) {
|
||||
model := New(nil)
|
||||
model.handleEnvelope(bootstrapEnvelope(t, "agents", 1,
|
||||
@@ -671,7 +599,7 @@ func TestVulnerabilityListSupportsWheelAndPageNavigation(t *testing.T) {
|
||||
})
|
||||
}
|
||||
_, _, chatWidth, _ := model.layout()
|
||||
_, _, _, agentHeight := model.sidebarHeights()
|
||||
_, _, agentHeight := model.sidebarHeights()
|
||||
pageItems := model.vulnerabilityPageItems()
|
||||
|
||||
updated, _ := model.updateMouse(tea.MouseMsg{
|
||||
@@ -937,20 +865,6 @@ func TestPanelPaddingResetsLeakingLineBackground(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFillBackgroundRestoresBaseForegroundAfterReset(t *testing.T) {
|
||||
const textFG = "\x1b[38;2;212;212;212m"
|
||||
view := "\x1b[38;2;167;139;250m◈ \x1b[0m\x1b[2mspawning\x1b[0m"
|
||||
filled := fillBackground(view)
|
||||
baseStyle := blackBG + textFG
|
||||
|
||||
if !strings.HasPrefix(filled, baseStyle) {
|
||||
t.Fatalf("frame does not set its base colors: %q", filled)
|
||||
}
|
||||
if got, want := strings.Count(filled, "\x1b[0m"+baseStyle), 2; got != want {
|
||||
t.Fatalf("base colors restored after %d resets, want %d: %q", got, want, filled)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMainTraceTreeAndFindingsRenderScrollbars(t *testing.T) {
|
||||
model := New(nil)
|
||||
model.width, model.height = 150, 35
|
||||
@@ -995,7 +909,7 @@ func TestMainScrollbarsSupportClickAndDrag(t *testing.T) {
|
||||
model.viewport.SetContent(model.viewportContent)
|
||||
showSidebar, _, chatWidth, chatHeight := model.layout()
|
||||
viewerHeight := model.viewerHeight()
|
||||
_, vulnHeight, _, agentHeight := model.sidebarHeights()
|
||||
_, vulnHeight, agentHeight := model.sidebarHeights()
|
||||
if !showSidebar {
|
||||
t.Fatal("test requires sidebar")
|
||||
}
|
||||
@@ -1039,61 +953,6 @@ func TestMainScrollbarsSupportClickAndDrag(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpRosterScrollsByKeyWheelAndScrollbar(t *testing.T) {
|
||||
model := New(nil)
|
||||
model.width, model.height = 150, 35
|
||||
model.ready = true
|
||||
conns := make([]protocol.Connection, 0, 12)
|
||||
for i := 0; i < 12; i++ {
|
||||
conns = append(conns, protocol.Connection{Name: fmt.Sprintf("conn-%02d", i), ToolCount: 2})
|
||||
}
|
||||
model.snapshot.Connections = conns
|
||||
|
||||
showSidebar, _, chatWidth, _ := model.layout()
|
||||
if !showSidebar {
|
||||
t.Fatal("test requires sidebar")
|
||||
}
|
||||
viewerHeight := model.viewerHeight()
|
||||
_, vulnHeight, mcpHeight, agentHeight := model.sidebarHeights()
|
||||
mcpTop := viewerHeight + agentHeight + vulnHeight
|
||||
bottom := model.clampMcpOffset(1 << 30)
|
||||
if bottom == 0 {
|
||||
t.Fatalf("a roster of %d should overflow the panel", len(conns))
|
||||
}
|
||||
|
||||
// Wheel over the panel focuses it and advances the window.
|
||||
updated, _ := model.updateMouse(tea.MouseMsg{
|
||||
X: chatWidth + 2, Y: mcpTop + 1, Button: tea.MouseButtonWheelDown,
|
||||
})
|
||||
model = updated.(Model)
|
||||
if model.focus != focusMcp || model.mcpOffset != 3 {
|
||||
t.Fatalf("wheel scroll did not focus and advance roster: focus=%v offset=%d", model.focus, model.mcpOffset)
|
||||
}
|
||||
|
||||
// Page down pins to the bottom; up steps back one.
|
||||
updated, _ = model.updateMain(tea.KeyMsg{Type: tea.KeyPgDown})
|
||||
model = updated.(Model)
|
||||
if model.mcpOffset != bottom {
|
||||
t.Fatalf("page down did not reach the roster bottom: offset=%d want=%d", model.mcpOffset, bottom)
|
||||
}
|
||||
updated, _ = model.updateMain(tea.KeyMsg{Type: tea.KeyUp})
|
||||
model = updated.(Model)
|
||||
if model.mcpOffset != bottom-1 {
|
||||
t.Fatalf("up did not step the roster back one: offset=%d want=%d", model.mcpOffset, bottom-1)
|
||||
}
|
||||
|
||||
// Clicking the scrollbar thumb captures it and moves the window.
|
||||
model.mcpOffset = 0
|
||||
updated, _ = model.updateMouse(tea.MouseMsg{
|
||||
X: model.width - 3, Y: mcpTop + mcpHeight - 2,
|
||||
Button: tea.MouseButtonLeft, Action: tea.MouseActionPress,
|
||||
})
|
||||
model = updated.(Model)
|
||||
if model.draggingScrollbar != scrollbarMcp || model.mcpOffset == 0 {
|
||||
t.Fatalf("mcp scrollbar click failed: drag=%v offset=%d", model.draggingScrollbar, model.mcpOffset)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminalSnapshotWithoutAgentsDoesNotKeepLoading(t *testing.T) {
|
||||
tests := []struct {
|
||||
state string
|
||||
|
||||
@@ -59,14 +59,6 @@ func (m Model) updateMain(key tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.ensureVulnerabilityVisible()
|
||||
return m, nil
|
||||
}
|
||||
if m.focus == focusMcp && len(m.snapshot.Connections) > 0 {
|
||||
delta := 1
|
||||
if key.String() == "up" {
|
||||
delta = -1
|
||||
}
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset + delta)
|
||||
return m, nil
|
||||
}
|
||||
case "enter", " ":
|
||||
if m.focus == focusVulnerabilities && len(m.snapshot.Vulnerabilities) > 0 {
|
||||
if key.String() == "enter" {
|
||||
@@ -102,10 +94,6 @@ func (m Model) updateMain(key tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.ensureVulnerabilityVisible()
|
||||
return m, nil
|
||||
}
|
||||
if m.focus == focusMcp && len(m.snapshot.Connections) > 0 {
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset - m.mcpPageSize())
|
||||
return m, nil
|
||||
}
|
||||
m.focus = focusChat
|
||||
m.input.Blur()
|
||||
m.followOutput = false
|
||||
@@ -117,10 +105,6 @@ func (m Model) updateMain(key tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.ensureVulnerabilityVisible()
|
||||
return m, nil
|
||||
}
|
||||
if m.focus == focusMcp && len(m.snapshot.Connections) > 0 {
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset + m.mcpPageSize())
|
||||
return m, nil
|
||||
}
|
||||
m.focus = focusChat
|
||||
m.input.Blur()
|
||||
m.viewport.HalfViewDown()
|
||||
@@ -163,10 +147,10 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
||||
}
|
||||
showSidebar, _, chatWidth, chatHeight := m.layout()
|
||||
viewerHeight := m.viewerHeight()
|
||||
_, vulnHeight, mcpHeight, agentHeight := m.sidebarHeights()
|
||||
_, vulnHeight, agentHeight := m.sidebarHeights()
|
||||
x, y := msg.X, msg.Y
|
||||
if m.updateMainScrollbarMouse(
|
||||
msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight,
|
||||
msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight,
|
||||
) {
|
||||
return m, nil
|
||||
}
|
||||
@@ -212,10 +196,6 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
||||
m.input.Blur()
|
||||
m.vulnOffset = max(0, m.vulnOffset-3)
|
||||
m.keepVulnerabilitySelectionInWindow()
|
||||
case mcpHeight > 0 && y < viewerHeight+agentHeight+vulnHeight+mcpHeight:
|
||||
m.focus = focusMcp
|
||||
m.input.Blur()
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset - 3)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
@@ -242,10 +222,6 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
||||
totalRows, _ := m.vulnerabilityScrollRows()
|
||||
m.vulnOffset = min(max(0, totalRows-m.vulnerabilityPageSize()), m.vulnOffset+3)
|
||||
m.keepVulnerabilitySelectionInWindow()
|
||||
case mcpHeight > 0 && y < viewerHeight+agentHeight+vulnHeight+mcpHeight:
|
||||
m.focus = focusMcp
|
||||
m.input.Blur()
|
||||
m.mcpOffset = m.clampMcpOffset(m.mcpOffset + 3)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
@@ -327,7 +303,7 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
||||
func (m *Model) updateMainScrollbarMouse(
|
||||
msg tea.MouseMsg,
|
||||
showSidebar bool,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight int,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
) bool {
|
||||
if msg.Action == tea.MouseActionRelease {
|
||||
if m.draggingScrollbar == scrollbarNone {
|
||||
@@ -337,18 +313,18 @@ func (m *Model) updateMainScrollbarMouse(
|
||||
return true
|
||||
}
|
||||
if msg.Action == tea.MouseActionMotion && m.draggingScrollbar != scrollbarNone {
|
||||
m.scrollFromMouse(m.draggingScrollbar, msg.Y, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
m.scrollFromMouse(m.draggingScrollbar, msg.Y, chatHeight, viewerHeight, agentHeight)
|
||||
return true
|
||||
}
|
||||
if msg.Action != tea.MouseActionPress || msg.Button != tea.MouseButtonLeft {
|
||||
return false
|
||||
}
|
||||
target := m.scrollbarAt(msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight)
|
||||
target := m.scrollbarAt(msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
if target == scrollbarNone {
|
||||
return false
|
||||
}
|
||||
m.draggingScrollbar = target
|
||||
m.scrollFromMouse(target, msg.Y, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
m.scrollFromMouse(target, msg.Y, chatHeight, viewerHeight, agentHeight)
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -365,9 +341,8 @@ func nearColumn(x, column int) bool {
|
||||
func (m Model) scrollbarAt(
|
||||
msg tea.MouseMsg,
|
||||
showSidebar bool,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight int,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
) scrollbarTarget {
|
||||
mcpTop := viewerHeight + agentHeight + vulnHeight
|
||||
switch {
|
||||
case nearColumn(msg.X, chatWidth-2) && msg.Y >= 1 && msg.Y < chatHeight-1 &&
|
||||
m.viewport.TotalLineCount() > m.viewport.VisibleLineCount():
|
||||
@@ -383,20 +358,13 @@ func (m Model) scrollbarAt(
|
||||
if totalRows > m.vulnerabilityPageSize() {
|
||||
return scrollbarFindings
|
||||
}
|
||||
// The roster scrolls below a fixed header, so its bar starts two rows into
|
||||
// the panel (border then header) rather than one.
|
||||
case showSidebar && mcpHeight > 0 && nearColumn(msg.X, m.width-3) &&
|
||||
msg.Y >= mcpTop+2 && msg.Y < mcpTop+mcpHeight-1:
|
||||
if len(m.snapshot.Connections) > m.mcpPageSize() {
|
||||
return scrollbarMcp
|
||||
}
|
||||
}
|
||||
return scrollbarNone
|
||||
}
|
||||
|
||||
func (m *Model) scrollFromMouse(
|
||||
target scrollbarTarget,
|
||||
y, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
y, chatHeight, viewerHeight, agentHeight int,
|
||||
) {
|
||||
switch target {
|
||||
case scrollbarTrace:
|
||||
@@ -422,13 +390,6 @@ func (m *Model) scrollFromMouse(
|
||||
// The offset is a row, so dragging moves the list continuously.
|
||||
m.vulnOffset = scrollbarOffset(y-viewerHeight-agentHeight-1, height, totalRows, height)
|
||||
m.keepVulnerabilitySelectionInWindow()
|
||||
case scrollbarMcp:
|
||||
height := m.mcpPageSize()
|
||||
total := len(m.snapshot.Connections)
|
||||
m.focus = focusMcp
|
||||
m.input.Blur()
|
||||
// The bar starts two rows into the panel (border then the fixed header).
|
||||
m.mcpOffset = scrollbarOffset(y-viewerHeight-agentHeight-vulnHeight-2, height, total, height)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -584,9 +545,6 @@ func (m *Model) cycleFocus(delta int) {
|
||||
if len(m.snapshot.Vulnerabilities) > 0 {
|
||||
available = append(available, focusVulnerabilities)
|
||||
}
|
||||
if len(m.snapshot.Connections) > 0 {
|
||||
available = append(available, focusMcp)
|
||||
}
|
||||
}
|
||||
idx := 0
|
||||
for i, focus := range available {
|
||||
|
||||
@@ -352,28 +352,21 @@ func (m Model) toastOverlay(view string) string {
|
||||
return strings.Join(bg, "\n")
|
||||
}
|
||||
|
||||
// Base frame colors are reapplied after full SGR resets so the TUI does not
|
||||
// inherit an unreadable foreground from the user's terminal profile.
|
||||
const (
|
||||
blackBG = "\x1b[48;2;0;0;0m"
|
||||
textFG = "\x1b[38;2;212;212;212m"
|
||||
baseFrameColors = blackBG + textFG
|
||||
)
|
||||
// blackBG is the SGR that selects a solid black background.
|
||||
const blackBG = "\x1b[48;2;0;0;0m"
|
||||
|
||||
// fillBackground paints the whole frame black like Textual's Screen background.
|
||||
// Bubble Tea has no screen compositor, so any cell the view does not explicitly
|
||||
// color shows the terminal's default background. lipgloss emits a full reset
|
||||
// (\x1b[0m) at the end of every styled span, which clears both foreground and
|
||||
// background. Reasserting only black made uncolored and faint text inherit the
|
||||
// terminal profile's foreground; light profiles therefore rendered that text
|
||||
// black-on-black. Reapply both base colors after each reset (and at the start).
|
||||
// Spans that set their own colors — inline code, selected rows, buttons — keep
|
||||
// them, because their color is emitted after the base style.
|
||||
// (\x1b[0m) at the end of every styled span, which also clears the background, so
|
||||
// we reassert black after each reset (and at the start). Spans that set their own
|
||||
// background — inline code, selected rows, buttons — keep it, because their color
|
||||
// is emitted before the reset.
|
||||
func fillBackground(view string) string {
|
||||
if view == "" {
|
||||
return view
|
||||
}
|
||||
return baseFrameColors + strings.ReplaceAll(view, "\x1b[0m", "\x1b[0m"+baseFrameColors)
|
||||
return blackBG + strings.ReplaceAll(view, "\x1b[0m", "\x1b[0m"+blackBG)
|
||||
}
|
||||
|
||||
func (m Model) splashView() string {
|
||||
@@ -507,7 +500,7 @@ func (m Model) mainView() string {
|
||||
func (m Model) sidebarView(width, height int) string {
|
||||
// Stats box height fits its content (auto, max 15); vulns panel max-height 12.
|
||||
statsBody := m.statsView()
|
||||
statsHeight, vulnHeight, mcpHeight, agentHeight := m.sidebarHeights()
|
||||
statsHeight, vulnHeight, agentHeight := m.sidebarHeights()
|
||||
agentBorder := dark
|
||||
if m.focus == focusAgents {
|
||||
agentBorder = green
|
||||
@@ -546,19 +539,11 @@ func (m Model) sidebarView(width, height int) string {
|
||||
)
|
||||
parts = append(parts, lipgloss.NewStyle().Width(width-2).Height(vulnRows).Border(lipgloss.RoundedBorder()).BorderForeground(vulnBorder).Padding(0, 1).Render(findings))
|
||||
}
|
||||
if mcpHeight > 0 {
|
||||
mcpBorder := dark
|
||||
if m.focus == focusMcp {
|
||||
mcpBorder = green
|
||||
}
|
||||
mcpRows := max(1, mcpHeight-2)
|
||||
parts = append(parts, lipgloss.NewStyle().Width(width-2).Height(mcpRows).Border(lipgloss.RoundedBorder()).BorderForeground(mcpBorder).Padding(0, 1).Render(m.mcpConnectionsView(width-4, mcpRows)))
|
||||
}
|
||||
parts = append(parts, lipgloss.NewStyle().Width(width-2).Height(statsHeight-2).Border(lipgloss.RoundedBorder()).BorderForeground(dark).Padding(0, 1).Render(statsBody))
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func (m Model) sidebarHeights() (statsHeight, vulnHeight, mcpHeight, agentHeight int) {
|
||||
func (m Model) sidebarHeights() (statsHeight, vulnHeight, agentHeight int) {
|
||||
// Measure the stats panel the way its box will render it: a long model name
|
||||
// wraps inside the sidebar, and counting only its newlines would size the
|
||||
// box short and push the whole frame past the bottom of the terminal.
|
||||
@@ -567,13 +552,7 @@ func (m Model) sidebarHeights() (statsHeight, vulnHeight, mcpHeight, agentHeight
|
||||
if len(m.snapshot.Vulnerabilities) > 0 {
|
||||
vulnHeight = min(12, len(m.vulnerabilityRows(m.vulnerabilityListWidth()))+2)
|
||||
}
|
||||
// One header line + one line per connection + the box border (2). Capped so a
|
||||
// long roster cannot crowd out the agent tree; a roster past the cap scrolls
|
||||
// inside the panel. Absent entirely when the run has no MCP connections.
|
||||
if len(m.snapshot.Connections) > 0 {
|
||||
mcpHeight = min(9, len(m.snapshot.Connections)+3)
|
||||
}
|
||||
agentHeight = max(3, m.height-m.viewerHeight()-statsHeight-vulnHeight-mcpHeight)
|
||||
agentHeight = max(3, m.height-m.viewerHeight()-statsHeight-vulnHeight)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -642,111 +621,6 @@ func (m Model) statsView() string {
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// mcpConnectionsView renders the sidebar MCP panel: a header carrying the total
|
||||
// connection count, then one row per connection with a status glyph and its tool
|
||||
// count (or "offline").
|
||||
// - a solid green dot marks an attached, idle connection;
|
||||
// - a green cycling quarter-circle (◐ ◓ ◑ ◒) marks a call running against it;
|
||||
// - a red dot plus "offline" marks a connection whose live session has died.
|
||||
//
|
||||
// The header stays fixed while the roster below it scrolls: when there are more
|
||||
// connections than the panel can show, the visible window is chosen by
|
||||
// m.mcpOffset and withVerticalScrollbar draws a thumb in the reserved last
|
||||
// column, exactly as the agent tree and findings list scroll.
|
||||
//
|
||||
// "In use" is derived from the connection-tagged tool-call events in the stream,
|
||||
// not carried on the connection roster, so a call in flight shows motion without
|
||||
// any extra backend signal. The quarter-circle rides the shared sweepFrame tick.
|
||||
func (m Model) mcpConnectionsView(width, rows int) string {
|
||||
conns := m.snapshot.Connections
|
||||
header := truncate(lipgloss.NewStyle().Foreground(dim).Render(
|
||||
fmt.Sprintf("MCP Connections (%d)", len(conns))), width)
|
||||
bodyRows := max(0, rows-1)
|
||||
if bodyRows == 0 {
|
||||
return header
|
||||
}
|
||||
inUse := m.mcpInUse()
|
||||
frames := []rune{'◐', '◓', '◑', '◒'}
|
||||
// Reserve the scrollbar column whether or not the bar is showing, so the
|
||||
// roster does not shift sideways as it grows past the panel.
|
||||
rosterWidth := max(1, width-1)
|
||||
start := windowStart(m.mcpOffset, len(conns), bodyRows)
|
||||
end := min(len(conns), start+bodyRows)
|
||||
lines := make([]string, 0, max(0, end-start))
|
||||
for i := start; i < end; i++ {
|
||||
conn := conns[i]
|
||||
var glyph, right string
|
||||
switch {
|
||||
case conn.Dead:
|
||||
glyph = lipgloss.NewStyle().Foreground(red).Render("●")
|
||||
right = lipgloss.NewStyle().Foreground(red).Render("offline")
|
||||
case inUse[conn.Name]:
|
||||
glyph = lipgloss.NewStyle().Foreground(green).Render(string(frames[m.sweepFrame%len(frames)]))
|
||||
right = lipgloss.NewStyle().Foreground(dim).Render(toolsLabel(conn.ToolCount))
|
||||
default:
|
||||
glyph = lipgloss.NewStyle().Foreground(green).Render("●")
|
||||
right = lipgloss.NewStyle().Foreground(dim).Render(toolsLabel(conn.ToolCount))
|
||||
}
|
||||
rightWidth := lipgloss.Width(right)
|
||||
name := truncate(lipgloss.NewStyle().Foreground(textColor).Render(conn.Name), max(1, rosterWidth-2-rightWidth-1))
|
||||
gap := max(1, rosterWidth-2-lipgloss.Width(name)-rightWidth)
|
||||
lines = append(lines, glyph+" "+name+strings.Repeat(" ", gap)+right)
|
||||
}
|
||||
roster := withVerticalScrollbar(
|
||||
strings.Join(lines, "\n"),
|
||||
width,
|
||||
bodyRows,
|
||||
len(conns),
|
||||
bodyRows,
|
||||
m.mcpOffset,
|
||||
m.scrollbarThumb(scrollbarMcp),
|
||||
)
|
||||
return header + "\n" + roster
|
||||
}
|
||||
|
||||
// mcpPageSize is how many connection rows the roster shows at once, below its
|
||||
// fixed header line.
|
||||
func (m Model) mcpPageSize() int {
|
||||
_, _, mcpHeight, _ := m.sidebarHeights()
|
||||
// mcpHeight = 2 (border) + header (1) + roster rows.
|
||||
return max(1, mcpHeight-3)
|
||||
}
|
||||
|
||||
// clampMcpOffset keeps the roster offset within the range that still shows a
|
||||
// full page of connections at the bottom.
|
||||
func (m Model) clampMcpOffset(offset int) int {
|
||||
return min(max(0, offset), max(0, len(m.snapshot.Connections)-m.mcpPageSize()))
|
||||
}
|
||||
|
||||
// mcpInUse is the set of MCP connections with a tool call currently running,
|
||||
// read off the connection-tagged tool events the model already holds. Each MCP
|
||||
// dispatch event carries the connection name (mcp_connection) and a status that
|
||||
// moves running -> completed as its own event is upserted, so a connection is
|
||||
// "in use" exactly while one of its events is still running.
|
||||
func (m Model) mcpInUse() map[string]bool {
|
||||
inUse := map[string]bool{}
|
||||
for _, event := range m.snapshot.Events {
|
||||
if event.Type != "tool" {
|
||||
continue
|
||||
}
|
||||
connection := render.StringValue(event.Data["mcp_connection"])
|
||||
if connection == "" {
|
||||
continue
|
||||
}
|
||||
if render.StringValue(event.Data["status"]) == "running" {
|
||||
inUse[connection] = true
|
||||
}
|
||||
}
|
||||
return inUse
|
||||
}
|
||||
|
||||
func toolsLabel(count int) string {
|
||||
if count == 1 {
|
||||
return "1 tool"
|
||||
}
|
||||
return fmt.Sprintf("%d tools", count)
|
||||
}
|
||||
|
||||
func numberValue(value any) int64 {
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
|
||||
@@ -134,7 +134,7 @@ func clampVulnerabilityOffset(offset, total, height int) int {
|
||||
}
|
||||
|
||||
func (m Model) vulnerabilityPageSize() int {
|
||||
_, vulnHeight, _, _ := m.sidebarHeights()
|
||||
_, vulnHeight, _ := m.sidebarHeights()
|
||||
return max(1, vulnHeight-2)
|
||||
}
|
||||
|
||||
|
||||
@@ -33,8 +33,6 @@ func (m *Model) handleEnvelope(envelope protocol.Envelope) tea.Cmd {
|
||||
m.stateRevision = update.Revision
|
||||
if m.snapshot.Error != nil {
|
||||
m.errorText = *m.snapshot.Error
|
||||
} else {
|
||||
m.errorText = ""
|
||||
}
|
||||
if m.snapshot.SetupMode {
|
||||
// The start screen is its own landing page; never sit on the
|
||||
@@ -81,10 +79,6 @@ func (m *Model) handleEnvelope(envelope protocol.Envelope) tea.Cmd {
|
||||
if collection := m.resyncRequests[envelope.RequestID]; collection != "" {
|
||||
m.resyncRequested[collection] = false
|
||||
delete(m.resyncRequests, envelope.RequestID)
|
||||
} else {
|
||||
for collection := range m.resyncRequested {
|
||||
m.resyncRequested[collection] = false
|
||||
}
|
||||
}
|
||||
}
|
||||
message := "Command failed"
|
||||
|
||||
@@ -32,17 +32,6 @@ type Agent struct {
|
||||
ErrorMessage string `json:"error_message"`
|
||||
}
|
||||
|
||||
// Connection is one MCP connection the run may reach, as the backend projects
|
||||
// it for the sidebar's MCP panel. Non-secret by construction: only the display
|
||||
// name, how many tools the connection offers, and whether its live session has
|
||||
// died (its reconnect-retry gave up). "In use" is not carried here; the client
|
||||
// derives it from the connection-tagged tool-call events in the event stream.
|
||||
type Connection struct {
|
||||
Name string `json:"name"`
|
||||
ToolCount int `json:"tool_count"`
|
||||
Dead bool `json:"dead"`
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
@@ -79,7 +68,6 @@ type Snapshot struct {
|
||||
Vulnerabilities []map[string]any `json:"-"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Subscription bool `json:"subscription"`
|
||||
Connections []Connection `json:"connections"`
|
||||
ViewerStatus string `json:"viewer_status"`
|
||||
ViewerURL *string `json:"viewer_url"`
|
||||
Error *string `json:"error"`
|
||||
|
||||
@@ -96,12 +96,14 @@ func TestCoverageDuplicateRejectionSurfacesTheError(t *testing.T) {
|
||||
requireContains(t, out, "/login", "already has coverage entry a1b2c3")
|
||||
}
|
||||
|
||||
func TestGetThreatModelRendersAmendments(t *testing.T) {
|
||||
func TestGetThreatModelRendersStalenessAndAmendments(t *testing.T) {
|
||||
out := ansi.Strip(Tool(tool("get_threat_model",
|
||||
map[string]any{"target": "https://app.example.com"},
|
||||
map[string]any{
|
||||
"success": true,
|
||||
"found": true,
|
||||
"success": true,
|
||||
"found": true,
|
||||
"stale": true,
|
||||
"cached_revision": "0123456789abcdef",
|
||||
"content": "# Overview\nMulti-tenant billing app.\n\n" +
|
||||
"## Trust Boundaries and Assumptions\n\n## Attack Surface\n",
|
||||
"amendments": []any{
|
||||
@@ -114,6 +116,7 @@ func TestGetThreatModelRendersAmendments(t *testing.T) {
|
||||
"completed")))
|
||||
requireContains(t, out,
|
||||
"Threat Model", "https://app.example.com",
|
||||
"stale", "01234567",
|
||||
"1 amendment(s)", "ReconAgent", "staging host shares the production database",
|
||||
"Multi-tenant billing app.", "Overview", "Trust Boundaries and Assumptions",
|
||||
)
|
||||
@@ -123,7 +126,7 @@ func TestGetThreatModelMissingModelIsExplicit(t *testing.T) {
|
||||
out := ansi.Strip(Tool(tool("get_threat_model",
|
||||
map[string]any{"target": "10.0.0.5"},
|
||||
map[string]any{"success": true, "found": false}, "completed")))
|
||||
requireContains(t, out, "No model derived for this target yet")
|
||||
requireContains(t, out, "No model cached for this target yet")
|
||||
}
|
||||
|
||||
func TestSaveThreatModelWarnsWhenAmendmentsAreCleared(t *testing.T) {
|
||||
@@ -131,10 +134,15 @@ func TestSaveThreatModelWarnsWhenAmendmentsAreCleared(t *testing.T) {
|
||||
map[string]any{"target": "app.example.com", "content": "# Overview\nA thing.\n"},
|
||||
map[string]any{
|
||||
"success": true,
|
||||
"revision": "unversioned",
|
||||
"amendments_cleared": 2,
|
||||
},
|
||||
"completed")))
|
||||
requireContains(t, out, "Threat Model Saved", "saved", "cleared 2 amendment(s)")
|
||||
// An unversioned target has no revision worth printing.
|
||||
if strings.Contains(out, "unversioned") {
|
||||
t.Fatalf("unversioned revision should not be rendered:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmendThreatModelRendersAddendum(t *testing.T) {
|
||||
|
||||
@@ -45,51 +45,3 @@ func renderMcpInspect(connection, status string) string {
|
||||
b.WriteString(style.Render(icon))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// renderMcpList renders list_mcps: the inventory of connections the run may
|
||||
// reach, not a call to any of them, so no connection leads and the event
|
||||
// carries no connection tag. Unlike the other MCP results, the names are worth
|
||||
// showing: Strix assembled them itself from the run's registered connections,
|
||||
// so they are short and never an outside server's payload.
|
||||
func renderMcpList(result any, status string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(mcpIcon + Dim().Render("Listing MCP servers") + "\n")
|
||||
for _, conn := range mcpConnectionEntries(result) {
|
||||
b.WriteString(" " + Col(Slate).Render(conn.name))
|
||||
if conn.dead {
|
||||
b.WriteString(Dim().Render(" · ") + Col(Red).Render("offline"))
|
||||
}
|
||||
b.WriteString("\n")
|
||||
}
|
||||
icon, style := statusIcon(status)
|
||||
b.WriteString(style.Render(icon))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// mcpListEntry is one connection read out of a list_mcps result: its display
|
||||
// name and whether its live session has died.
|
||||
type mcpListEntry struct {
|
||||
name string
|
||||
dead bool
|
||||
}
|
||||
|
||||
// mcpConnectionEntries reads the connections out of a list_mcps result, which is
|
||||
// {"connections": [{"name": ..., "dead": ...}, ...]}. Anything else (still
|
||||
// running, or a result bounded down to a string) yields no entries, and the
|
||||
// header plus status stand alone.
|
||||
func mcpConnectionEntries(result any) []mcpListEntry {
|
||||
resultMap, _ := result.(map[string]any)
|
||||
connections, _ := resultMap["connections"].([]any)
|
||||
var entries []mcpListEntry
|
||||
for _, raw := range connections {
|
||||
entry, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if name := strings.TrimSpace(StringValue(entry["name"])); name != "" {
|
||||
dead, _ := entry["dead"].(bool)
|
||||
entries = append(entries, mcpListEntry{name: name, dead: dead})
|
||||
}
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
||||
@@ -70,10 +70,6 @@ func Tool(data map[string]any) string {
|
||||
}
|
||||
|
||||
switch name {
|
||||
// list_mcps inventories every connection rather than touching one, so it is
|
||||
// the one MCP tool with no connection tag and routes by name like a built-in.
|
||||
case "list_mcps":
|
||||
return renderMcpList(result, status)
|
||||
case "exec_command":
|
||||
return renderExecCommand(args, result, status)
|
||||
case "write_stdin":
|
||||
|
||||
@@ -268,24 +268,6 @@ func TestMcpDescribeInspectsConnection(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpListMarksDeadConnectionsOffline(t *testing.T) {
|
||||
// list_mcps carries a per-connection dead flag; a dead connection reads as
|
||||
// offline in the inventory while a live one shows normally.
|
||||
result := map[string]any{
|
||||
"connections": []any{
|
||||
map[string]any{"name": "supabase", "tool_count": float64(3), "dead": false},
|
||||
map[string]any{"name": "vercel", "tool_count": float64(1), "dead": true},
|
||||
},
|
||||
}
|
||||
data := tool("list_mcps", nil, result, "completed")
|
||||
|
||||
out := ansi.Strip(Tool(data))
|
||||
requireContains(t, out, "Listing MCP servers", "supabase", "vercel", "offline")
|
||||
if strings.Count(out, "offline") != 1 {
|
||||
t.Fatalf("only the dead connection should read offline:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollapseToolShellPreviewAndExpand(t *testing.T) {
|
||||
lines := make([]string, 16)
|
||||
for i := range lines {
|
||||
|
||||
@@ -56,6 +56,9 @@ func renderThreatModel(name string, args map[string]any, result any) string {
|
||||
threatModelBody(&b, StringValue(args["addendum"]))
|
||||
default:
|
||||
b.WriteString("\n " + Col(Green).Render("✓ saved"))
|
||||
if revision := shortRevision(StringValue(m["revision"])); revision != "" {
|
||||
b.WriteString(Dim().Render(" at " + revision))
|
||||
}
|
||||
// Saving folds amendments away, so the count that vanished is worth
|
||||
// stating: it is the one destructive thing this tool does.
|
||||
if cleared, ok := NumericValue(m["amendments_cleared"]); ok && cleared > 0 {
|
||||
@@ -69,9 +72,15 @@ func renderThreatModel(name string, args map[string]any, result any) string {
|
||||
|
||||
func threatModelReadBody(b *strings.Builder, result map[string]any) {
|
||||
if !truthy(result["found"]) {
|
||||
b.WriteString("\n " + Dim().Render("No model derived for this target yet"))
|
||||
b.WriteString("\n " + Dim().Render("No model cached for this target yet"))
|
||||
return
|
||||
}
|
||||
if truthy(result["stale"]) {
|
||||
b.WriteString("\n " + Col(AmberY).Render("⚠ stale"))
|
||||
if cached := shortRevision(StringValue(result["cached_revision"])); cached != "" {
|
||||
b.WriteString(Dim().Render(" (written at " + cached + ")"))
|
||||
}
|
||||
}
|
||||
if amendments, ok := result["amendments"].([]any); ok && len(amendments) > 0 {
|
||||
b.WriteString("\n " + Col(Gold).Render("+ "+strconv.Itoa(len(amendments))+
|
||||
" amendment(s)") + Dim().Render(" — later statements win"))
|
||||
@@ -117,3 +126,13 @@ func threatModelBody(b *strings.Builder, content string) {
|
||||
b.WriteString("\n " + Dim().Render(strings.Join(headings, " · ")))
|
||||
}
|
||||
}
|
||||
|
||||
// shortRevision abbreviates a git sha; "unversioned" targets have no revision
|
||||
// worth showing.
|
||||
func shortRevision(revision string) string {
|
||||
revision = strings.TrimSpace(revision)
|
||||
if revision == "" || revision == "unversioned" {
|
||||
return ""
|
||||
}
|
||||
return firstN(revision, 8)
|
||||
}
|
||||
|
||||
@@ -103,7 +103,6 @@ class TuiLiveView:
|
||||
statuses = agents_data.get("statuses") or {}
|
||||
names = agents_data.get("names") or {}
|
||||
parent_of = agents_data.get("parent_of") or {}
|
||||
errors = agents_data.get("errors") or {}
|
||||
if not isinstance(statuses, dict):
|
||||
return
|
||||
for agent_id, status in statuses.items():
|
||||
@@ -114,7 +113,6 @@ class TuiLiveView:
|
||||
name=names.get(agent_id, agent_id) if isinstance(names, dict) else agent_id,
|
||||
parent_id=parent_of.get(agent_id) if isinstance(parent_of, dict) else None,
|
||||
status=str(status),
|
||||
error_message=errors.get(agent_id) if isinstance(errors, dict) else None,
|
||||
)
|
||||
# Ahead of the replayed history, so it opens the transcript.
|
||||
self.flush_user_instruction()
|
||||
|
||||
@@ -185,7 +185,6 @@ class GoTuiRuntime:
|
||||
max_turns=self.args.max_turns,
|
||||
max_budget_usd=self.args.max_budget_usd,
|
||||
event_sink=self.capture_event,
|
||||
mcp_status_sink=self.capture_mcp_status,
|
||||
)
|
||||
await self._sync_agent_state()
|
||||
if self.controller.scan_state == "running":
|
||||
@@ -211,15 +210,6 @@ class GoTuiRuntime:
|
||||
self.live_view.ingest_sdk_event(agent_id, event)
|
||||
self.controller.notify_changed()
|
||||
|
||||
def capture_mcp_status(self, roster: list[dict[str, Any]]) -> None:
|
||||
"""Receive the engine's MCP connection roster and hand it to the controller.
|
||||
|
||||
Runs on the scan's event loop (called from the runner at establishment
|
||||
and from a session's on-dead callback), the same loop that drives
|
||||
``capture_event``, so updating the controller and repainting here is
|
||||
safe. The controller renders it as the sidebar MCP connections panel."""
|
||||
self.controller.set_mcp_connections(roster)
|
||||
|
||||
async def _sync_agent_state(self) -> bool:
|
||||
parent_of, statuses, names, errors = await self.coordinator.graph_snapshot()
|
||||
changed = False
|
||||
@@ -258,9 +248,6 @@ class GoTuiRuntime:
|
||||
scan_state = "failed"
|
||||
if root_id is not None and errors.get(root_id):
|
||||
self.controller.error = errors[root_id]
|
||||
elif scan_state == "failed" and root_status in {"running", "waiting", "budget_paused"}:
|
||||
scan_state = "running"
|
||||
self.controller.error = None
|
||||
elif scan_state != "failed":
|
||||
if report_status == "completed":
|
||||
scan_state = "completed"
|
||||
|
||||
@@ -30,7 +30,6 @@ import {
|
||||
fetchTranscript,
|
||||
fetchVulnerabilities,
|
||||
forgetAuth,
|
||||
parseMcpConnectionStatus,
|
||||
type AuthStatus,
|
||||
type LoadedRun,
|
||||
type RunsPayload,
|
||||
@@ -170,27 +169,6 @@ export default function App() {
|
||||
const agentCount = run?.transcript.agents.length ?? 0;
|
||||
const verified = auth?.verified === true;
|
||||
|
||||
// The run's persisted MCP roster (from run.json via /api/run), plus the set of
|
||||
// connections with a tool call currently in flight. "In use" is derived here
|
||||
// from the connection-tagged tool events rather than carried on the roster:
|
||||
// an MCP dispatch event carries its connection name and a status that moves
|
||||
// running -> completed, so a connection is in use while one of its events is
|
||||
// still running. This mirrors the terminal UI's MCP panel exactly.
|
||||
const mcpConnections = useMemo(
|
||||
() => (run ? parseMcpConnectionStatus(run.raw) : []),
|
||||
[run]
|
||||
);
|
||||
const mcpInUse = useMemo(() => {
|
||||
const inUse = new Set<string>();
|
||||
for (const event of run?.transcript.events ?? []) {
|
||||
if (event.type !== "tool") continue;
|
||||
const connection = event.data?.mcp_connection;
|
||||
if (typeof connection !== "string" || !connection) continue;
|
||||
if (event.data?.status === "running") inUse.add(connection);
|
||||
}
|
||||
return inUse;
|
||||
}, [run]);
|
||||
|
||||
// Per-run guard for the default view: land on Agents while a scan is live,
|
||||
// Overview once it finishes. Applied at most once per run and never once the
|
||||
// user has navigated manually (userSetView flips the guard).
|
||||
@@ -273,8 +251,6 @@ export default function App() {
|
||||
}}
|
||||
issuesCount={run?.vulnerabilities.length ?? 0}
|
||||
agentCount={agentCount}
|
||||
mcpConnections={mcpConnections}
|
||||
mcpInUse={mcpInUse}
|
||||
runCount={runs?.count ?? 0}
|
||||
finished={run?.finished ?? false}
|
||||
verified={verified}
|
||||
|
||||
@@ -14,7 +14,6 @@ import { IoChatbubblesOutline } from "react-icons/io5";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { ctaUrl, trackCta } from "@/lib/cta";
|
||||
import { UpgradeModal } from "@/components/UpgradeModal";
|
||||
import type { McpConnectionStatus } from "@/data/serverSource";
|
||||
import type { View } from "@/App";
|
||||
|
||||
/**
|
||||
@@ -38,8 +37,6 @@ interface SidebarProps {
|
||||
onSelectView: (view: View) => void;
|
||||
issuesCount: number;
|
||||
agentCount: number;
|
||||
mcpConnections: McpConnectionStatus[];
|
||||
mcpInUse: Set<string>;
|
||||
runCount: number;
|
||||
finished: boolean;
|
||||
verified: boolean;
|
||||
@@ -64,8 +61,6 @@ export default function Sidebar({
|
||||
onSelectView,
|
||||
issuesCount,
|
||||
agentCount,
|
||||
mcpConnections,
|
||||
mcpInUse,
|
||||
runCount,
|
||||
finished,
|
||||
verified,
|
||||
@@ -251,9 +246,6 @@ export default function Sidebar({
|
||||
onClick={() => onSelectView("agents")}
|
||||
/>
|
||||
)}
|
||||
{mcpConnections.length > 0 && (
|
||||
<McpConnectionsPanel connections={mcpConnections} inUse={mcpInUse} />
|
||||
)}
|
||||
<NavItem
|
||||
icon={<History className="h-4 w-4" />}
|
||||
label="Past runs"
|
||||
@@ -429,84 +421,6 @@ function NavItem({ icon, label, active, onClick, count }: NavItemProps) {
|
||||
);
|
||||
}
|
||||
|
||||
// The quarter-circle sweep frames the terminal UI cycles for an in-use
|
||||
// connection, and the sub-second tick that advances them.
|
||||
const SWEEP_FRAMES = ["◐", "◓", "◑", "◒"] as const;
|
||||
const SWEEP_MS = 220;
|
||||
|
||||
/**
|
||||
* The MCP connections panel: a compact roster of the run's connected MCP
|
||||
* servers, matching the terminal UI's sidebar panel. A header carries the
|
||||
* total count; each row shows a status glyph, the connection name, and its
|
||||
* tool count (or "offline"):
|
||||
* - solid green dot: attached and idle;
|
||||
* - green cycling quarter-circle (◐◓◑◒): a tool call is running against it;
|
||||
* - red dot + "offline": the connection's live session has died.
|
||||
*
|
||||
* "In use" is derived by the caller from the connection-tagged tool events, not
|
||||
* carried on the roster, so a call in flight shows motion with no extra signal.
|
||||
* The roster scrolls within a bounded height so a long list never blows out the
|
||||
* rail, mirroring how the nav above it scrolls.
|
||||
*/
|
||||
function McpConnectionsPanel({
|
||||
connections,
|
||||
inUse,
|
||||
}: {
|
||||
connections: McpConnectionStatus[];
|
||||
inUse: Set<string>;
|
||||
}) {
|
||||
const anyInUse = connections.some((c) => !c.dead && inUse.has(c.name));
|
||||
const [frame, setFrame] = useState(0);
|
||||
|
||||
// Advance the sweep only while at least one connection is in use, so an idle
|
||||
// panel does no work.
|
||||
useEffect(() => {
|
||||
if (!anyInUse) return;
|
||||
const id = setInterval(() => setFrame((f) => (f + 1) % SWEEP_FRAMES.length), SWEEP_MS);
|
||||
return () => clearInterval(id);
|
||||
}, [anyInUse]);
|
||||
|
||||
return (
|
||||
<div className="mt-1">
|
||||
<div className="flex h-7 items-center px-2 text-[11px] font-medium text-[#666]">
|
||||
MCP Connections ({connections.length})
|
||||
</div>
|
||||
<div className="max-h-48 overflow-y-auto overflow-x-clip scrollbar-thin">
|
||||
{connections.map((conn) => {
|
||||
const busy = !conn.dead && inUse.has(conn.name);
|
||||
return (
|
||||
<div
|
||||
key={conn.name}
|
||||
className="flex h-7 items-center gap-2 rounded-md px-2"
|
||||
title={conn.provider ? `${conn.name} · ${conn.provider}` : conn.name}
|
||||
>
|
||||
<span
|
||||
className={cn(
|
||||
"w-3 flex-none text-center text-[11px] leading-none",
|
||||
conn.dead ? "text-red-400" : "text-emerald-400"
|
||||
)}
|
||||
aria-hidden="true"
|
||||
>
|
||||
{conn.dead ? "●" : busy ? SWEEP_FRAMES[frame] : "●"}
|
||||
</span>
|
||||
<span className="min-w-0 flex-1 truncate text-[13px] font-medium text-[#ededed]">
|
||||
{conn.name}
|
||||
</span>
|
||||
{conn.dead ? (
|
||||
<span className="flex-none text-[11px] text-red-400">offline</span>
|
||||
) : (
|
||||
<span className="flex-none text-[11px] tabular-nums text-[#666]">
|
||||
{conn.toolCount} {conn.toolCount === 1 ? "tool" : "tools"}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Overview icon: a dashboard grid glyph (16x16 viewBox).
|
||||
function ProjectsIcon() {
|
||||
return (
|
||||
|
||||
@@ -14,10 +14,6 @@ import type { ToolRendererProps } from "@/types/events";
|
||||
* never as markdown, since it came from a server outside Strix.
|
||||
*
|
||||
* The full result is still in the run's event data on disk either way.
|
||||
*
|
||||
* list_mcps is the other exception: its result is the engine's own inventory of
|
||||
* the run's connections (names and tool counts), short and assembled by Strix
|
||||
* rather than returned by an outside server, so it is shown inline.
|
||||
*/
|
||||
|
||||
/** Arguments one line each, as the terminal prints them. */
|
||||
@@ -29,52 +25,6 @@ function argLines(args: unknown): string[] {
|
||||
});
|
||||
}
|
||||
|
||||
/** One connection out of a list_mcps inventory. */
|
||||
interface McpListingEntry {
|
||||
name: string;
|
||||
toolCount: number | null;
|
||||
dead: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* The connections out of a list_mcps result, which is
|
||||
* `{"connections": [{id, name, description, tool_count}, ...]}`, sometimes
|
||||
* arriving JSON-encoded as a string. Anything else yields an empty list and the
|
||||
* row shows just the header and status. Unlike other MCP results this one is
|
||||
* safe to show: the engine assembled it from the run's own registered
|
||||
* connections, so it is short and never an outside server's payload. It still
|
||||
* renders as inert text.
|
||||
*/
|
||||
function listingEntries(result: unknown): McpListingEntry[] {
|
||||
let value = result;
|
||||
if (typeof value === "string") {
|
||||
try {
|
||||
value = JSON.parse(value);
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
const connections =
|
||||
value && typeof value === "object" && !Array.isArray(value)
|
||||
? (value as Record<string, unknown>).connections
|
||||
: null;
|
||||
if (!Array.isArray(connections)) return [];
|
||||
return connections.flatMap((entry) => {
|
||||
if (!entry || typeof entry !== "object" || Array.isArray(entry)) return [];
|
||||
const record = entry as Record<string, unknown>;
|
||||
const name =
|
||||
typeof record.name === "string" && record.name.trim()
|
||||
? record.name.trim()
|
||||
: typeof record.id === "string"
|
||||
? record.id.trim()
|
||||
: "";
|
||||
if (!name) return [];
|
||||
const toolCount = typeof record.tool_count === "number" ? record.tool_count : null;
|
||||
const dead = record.dead === true;
|
||||
return [{ name, toolCount, dead }];
|
||||
});
|
||||
}
|
||||
|
||||
const MAX_ERROR_CHARS = 600;
|
||||
|
||||
function errorText(result: unknown): string | null {
|
||||
@@ -100,17 +50,11 @@ export default function McpRenderer({
|
||||
// describe_mcp inspects a connection's catalog rather than calling a tool on
|
||||
// it, so the connection is the subject and there is no underlying tool.
|
||||
const inspecting = toolName === "describe_mcp";
|
||||
// list_mcps inventories every connection rather than touching one, so it
|
||||
// carries no connection at all and is routed here by name instead.
|
||||
const listing = toolName === "list_mcps";
|
||||
const entries = listing ? listingEntries(result) : [];
|
||||
|
||||
return (
|
||||
<div>
|
||||
<div className="flex items-center gap-2 flex-wrap">
|
||||
{listing ? (
|
||||
<span className="text-[13px] text-[#555]">Listing connected MCP servers</span>
|
||||
) : inspecting ? (
|
||||
{inspecting ? (
|
||||
<>
|
||||
<span className="text-[13px] text-[#555]">Inspecting MCP server</span>
|
||||
{mcpConnection && (
|
||||
@@ -138,26 +82,6 @@ export default function McpRenderer({
|
||||
</div>
|
||||
)}
|
||||
|
||||
{entries.length > 0 && (
|
||||
<div className="mt-1 font-mono text-[13px] leading-relaxed">
|
||||
{entries.map((entry) => (
|
||||
<div key={entry.name} className={`break-all${entry.dead ? " opacity-50" : ""}`}>
|
||||
<span className="text-teal-300">{entry.name}</span>
|
||||
{entry.dead ? (
|
||||
<span className="text-red-400/80"> · offline</span>
|
||||
) : (
|
||||
entry.toolCount !== null && (
|
||||
<span className="text-[#555]">
|
||||
{" "}
|
||||
· {entry.toolCount} {entry.toolCount === 1 ? "tool" : "tools"}
|
||||
</span>
|
||||
)
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-1 text-[13px]">
|
||||
{status === "running" && <span className="text-[#666]">Running</span>}
|
||||
{status === "completed" && <span className="text-emerald-400/80">✓ Done</span>}
|
||||
|
||||
@@ -16,6 +16,13 @@ const ACTION_LABELS: Record<string, { label: string; Icon: typeof Crosshair }> =
|
||||
amend_threat_model: { label: "Threat model amended", Icon: Plus },
|
||||
};
|
||||
|
||||
/** A git sha is noise past its first bytes, and "unversioned" is not a revision. */
|
||||
function shortRevision(revision: unknown): string {
|
||||
const value = typeof revision === "string" ? revision.trim() : "";
|
||||
if (!value || value === "unversioned") return "";
|
||||
return value.slice(0, 8);
|
||||
}
|
||||
|
||||
export default function ThreatModelRenderer({ toolName, args, result }: ToolRendererProps) {
|
||||
const action = ACTION_LABELS[toolName] ?? { label: "Threat model", Icon: Crosshair };
|
||||
const ActionIcon = action.Icon;
|
||||
@@ -52,15 +59,22 @@ export default function ThreatModelRenderer({ toolName, args, result }: ToolRend
|
||||
return (
|
||||
<div>
|
||||
{header}
|
||||
<div className="mt-1.5 text-[#555] text-xs">No model derived for this target yet</div>
|
||||
<div className="mt-1.5 text-[#555] text-xs">No model cached for this target yet</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
const rawAmendments = structured?.amendments;
|
||||
const amendments: Amendment[] = Array.isArray(rawAmendments) ? (rawAmendments as Amendment[]) : [];
|
||||
const cachedRevision = shortRevision(structured?.cached_revision);
|
||||
return (
|
||||
<div>
|
||||
{header}
|
||||
{structured?.stale === true && (
|
||||
<div className="mt-1.5 flex items-center gap-1.5 text-yellow-400/80 text-xs">
|
||||
<AlertTriangle className="w-3 h-3 shrink-0" />
|
||||
<span>stale{cachedRevision ? ` — written at ${cachedRevision}` : ""}</span>
|
||||
</div>
|
||||
)}
|
||||
{amendments.length > 0 && (
|
||||
<div className="mt-2">
|
||||
<span className="text-amber-400/70 text-xs font-semibold">
|
||||
@@ -105,10 +119,12 @@ export default function ThreatModelRenderer({ toolName, args, result }: ToolRend
|
||||
}
|
||||
|
||||
const cleared = (structured?.amendments_cleared as number | undefined) ?? 0;
|
||||
const revision = shortRevision(structured?.revision);
|
||||
const content = (args.content as string) ?? "";
|
||||
return (
|
||||
<div>
|
||||
{header}
|
||||
{revision && <div className="mt-1.5 text-[#666] font-mono text-xs">at {revision}</div>}
|
||||
{/* Saving folds amendments away — the one destructive thing this tool does. */}
|
||||
{cleared > 0 && (
|
||||
<div className="mt-1.5 flex items-center gap-1.5 text-yellow-400/80 text-xs">
|
||||
|
||||
@@ -93,9 +93,7 @@ const CATEGORY_META: Record<ToolCategory, CategoryMeta> = {
|
||||
threatModel: { renderer: ThreatModelRenderer, icon: Crosshair, color: "text-blue-400", match: /threat_model/ },
|
||||
telemetry: { renderer: FallbackRenderer, icon: Wrench, color: "text-[#555]" },
|
||||
// Tools from the user's own MCP servers. Resolved from the connection on the
|
||||
// event rather than from a tool name — except list_mcps, the engine's
|
||||
// inventory of every connection, which touches none and so carries no
|
||||
// connection to resolve from; it is the family's one name below.
|
||||
// event rather than from a tool name, so this family has no names below.
|
||||
mcp: { renderer: McpRenderer, icon: Plug, color: "text-teal-400" },
|
||||
};
|
||||
|
||||
@@ -130,7 +128,7 @@ const CATEGORY_TOOLS: Record<ToolCategory, readonly string[]> = {
|
||||
// Per-target threat model, shared across the agent tree
|
||||
threatModel: ["get_threat_model", "save_threat_model", "amend_threat_model"],
|
||||
telemetry: ["sandbox_error_details", "llm_error_details"],
|
||||
mcp: ["list_mcps"],
|
||||
mcp: [],
|
||||
};
|
||||
|
||||
/** Reverse index (tool name → family), built once from CATEGORY_TOOLS. */
|
||||
|
||||
@@ -39,39 +39,6 @@ export interface Transcript {
|
||||
events: TranscriptEvent[];
|
||||
}
|
||||
|
||||
/**
|
||||
* One MCP connection's non-secret status, as persisted to run.json by the
|
||||
* engine under `mcp_connection_status` and surfaced verbatim by GET /api/run.
|
||||
* Only name / provider / tool_count / dead ride here; never config, url, or
|
||||
* token. `dead` means the connection's live session gave up reconnecting.
|
||||
*/
|
||||
export interface McpConnectionStatus {
|
||||
name: string;
|
||||
provider: string | null;
|
||||
toolCount: number;
|
||||
dead: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Read the MCP connection roster out of a raw run record. Tolerates the field
|
||||
* being absent (older runs, or a run with no MCP) and any malformed entry,
|
||||
* yielding an empty list rather than throwing.
|
||||
*/
|
||||
export function parseMcpConnectionStatus(raw: Record<string, unknown>): McpConnectionStatus[] {
|
||||
const list = raw?.mcp_connection_status;
|
||||
if (!Array.isArray(list)) return [];
|
||||
return list.flatMap((entry) => {
|
||||
if (!entry || typeof entry !== "object" || Array.isArray(entry)) return [];
|
||||
const record = entry as Record<string, unknown>;
|
||||
const name = typeof record.name === "string" ? record.name.trim() : "";
|
||||
if (!name) return [];
|
||||
const provider = typeof record.provider === "string" && record.provider.trim() ? record.provider.trim() : null;
|
||||
const toolCount = typeof record.tool_count === "number" ? record.tool_count : 0;
|
||||
const dead = record.dead === true;
|
||||
return [{ name, provider, toolCount, dead }];
|
||||
});
|
||||
}
|
||||
|
||||
export interface LoadedRun {
|
||||
summary: ParsedRunSummary;
|
||||
/** Whole raw run record (for llm_usage, targets_info details, etc.). */
|
||||
|
||||
File diff suppressed because one or more lines are too long
10
strix/interface/viewer/static/assets/index-D0453ODW.css
Normal file
10
strix/interface/viewer/static/assets/index-D0453ODW.css
Normal file
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -6,8 +6,8 @@
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<meta name="color-scheme" content="dark" />
|
||||
<title>Strix Results</title>
|
||||
<script type="module" crossorigin src="./assets/index-Bpn8GiSb.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-qwPOPAGC.css">
|
||||
<script type="module" crossorigin src="./assets/index-CYf9nnT3.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-D0453ODW.css">
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
|
||||
@@ -14,7 +14,6 @@ from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
|
||||
|
||||
@@ -31,38 +30,12 @@ _lock = threading.Lock()
|
||||
_thread: threading.Thread | None = None
|
||||
|
||||
|
||||
def _purge_orphaned_modules(before: frozenset[str]) -> None:
|
||||
"""Remove submodules stranded by an import attempt that just failed.
|
||||
|
||||
When a package import fails partway (for example CPython's import-lock
|
||||
deadlock avoidance breaking a cross-thread cycle), the failed package is
|
||||
removed from ``sys.modules`` but submodules it already finished stay
|
||||
behind. A later import of one of those submodules then short-circuits on
|
||||
the cached entry without re-importing its parent, and re-entering the
|
||||
parent from inside a submodule crashes with "partially initialized
|
||||
module". Dropping the orphans (cached submodules whose ancestor package is
|
||||
gone) restores a clean slate, and touches nothing another thread imported
|
||||
successfully.
|
||||
"""
|
||||
added = set(sys.modules) - before
|
||||
for name in added:
|
||||
parent = name.rpartition(".")[0]
|
||||
while parent:
|
||||
if parent not in sys.modules:
|
||||
sys.modules.pop(name, None)
|
||||
logger.debug("Import warm-up purged orphaned module %r", name)
|
||||
break
|
||||
parent = parent.rpartition(".")[0]
|
||||
|
||||
|
||||
def _warm(modules: tuple[str, ...]) -> None:
|
||||
for name in modules:
|
||||
before = frozenset(sys.modules)
|
||||
try:
|
||||
importlib.import_module(name)
|
||||
except Exception: # noqa: BLE001 - a failed warm-up must never fail the run.
|
||||
logger.debug("Import warm-up for %r failed", name, exc_info=True)
|
||||
_purge_orphaned_modules(before)
|
||||
|
||||
|
||||
def start_import_warmup(modules: tuple[str, ...] = WARMUP_MODULES) -> threading.Thread:
|
||||
|
||||
@@ -1,26 +1,12 @@
|
||||
"""Report/finding helpers."""
|
||||
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from strix.report.dedupe import check_duplicate
|
||||
from strix.report.state import ReportState, get_global_report_state, set_global_report_state
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from strix.report.dedupe import check_duplicate
|
||||
|
||||
__all__ = [
|
||||
"ReportState",
|
||||
"check_duplicate",
|
||||
"get_global_report_state",
|
||||
"set_global_report_state",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
# check_duplicate pulls in the agents SDK import graph, so it resolves
|
||||
# lazily: importing this package must stay lightweight and never enter
|
||||
# that graph (the import warm-up thread may be walking it concurrently).
|
||||
if name == "check_duplicate":
|
||||
return import_module("strix.report.dedupe").check_duplicate
|
||||
raise AttributeError(name)
|
||||
|
||||
@@ -7,6 +7,7 @@ import logging
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import ModelTracing
|
||||
from openai.types.responses import ResponseOutputMessage
|
||||
|
||||
@@ -21,8 +22,6 @@ from strix.report.state import get_global_report_state
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.items import ModelResponse
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import Model
|
||||
|
||||
from strix.config.settings import DedupeSettings
|
||||
|
||||
@@ -30,11 +29,30 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _dedupe_extra_args(dedupe: DedupeSettings) -> dict[str, str]:
|
||||
"""Per-call credential + endpoint for the dedupe model.
|
||||
|
||||
Provider env vars and the global base URL are process-wide, so a
|
||||
shared-provider dedupe key or a distinct dedupe endpoint can't be installed
|
||||
globally without clobbering (or being clobbered by) the main model's
|
||||
config. Passing them per call keeps the two apart. Only applies when a
|
||||
dedicated dedupe model is configured.
|
||||
"""
|
||||
if not dedupe.model:
|
||||
return {}
|
||||
extra: dict[str, str] = {}
|
||||
if dedupe.api_key and dedupe.api_key.strip():
|
||||
extra["api_key"] = dedupe.api_key.strip()
|
||||
if dedupe.api_base and dedupe.api_base.strip():
|
||||
extra["api_base"] = dedupe.api_base.strip()
|
||||
return extra
|
||||
|
||||
|
||||
def _dedupe_model_settings(
|
||||
dedupe: DedupeSettings, model_name: str, request_timeout: float | None
|
||||
) -> ModelSettings:
|
||||
llm = load_settings().llm
|
||||
return make_model_settings(
|
||||
settings = make_model_settings(
|
||||
dedupe.reasoning_effort,
|
||||
model_name=model_name,
|
||||
force_required_tool_choice=False,
|
||||
@@ -46,21 +64,10 @@ def _dedupe_model_settings(
|
||||
extra_headers=dedupe.extra_headers if dedupe.model else llm.extra_headers,
|
||||
has_tools=False,
|
||||
)
|
||||
|
||||
|
||||
def resolve_dedupe_model(dedupe: DedupeSettings, model_name: str) -> Model:
|
||||
"""Resolve the dedupe model, bound to its own endpoint when it has one.
|
||||
|
||||
Credentials can't ride on the request: every model implementation already
|
||||
passes its own ``api_key``/``base_url``, so the same keys in ``extra_args``
|
||||
collide with them and raise before anything is sent. A provider bound to the
|
||||
dedupe endpoint keeps it apart from the main model's process-wide defaults.
|
||||
"""
|
||||
api_key = (dedupe.api_key or "").strip() if dedupe.model else ""
|
||||
api_base = (dedupe.api_base or "").strip() if dedupe.model else ""
|
||||
if not (api_key or api_base):
|
||||
return StrixProvider().get_model(model_name)
|
||||
return StrixProvider(api_key=api_key or None, base_url=api_base or None).get_model(model_name)
|
||||
extra = _dedupe_extra_args(dedupe)
|
||||
if extra:
|
||||
settings = settings.resolve(ModelSettings(extra_args=extra))
|
||||
return settings
|
||||
|
||||
|
||||
DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge.
|
||||
@@ -364,7 +371,7 @@ async def check_duplicate(
|
||||
|
||||
configure_sdk_model_defaults(settings)
|
||||
resolved_model = model_name.strip()
|
||||
model = resolve_dedupe_model(dedupe, resolved_model)
|
||||
model = StrixProvider().get_model(resolved_model)
|
||||
response = await model.get_response(
|
||||
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
||||
input=user_msg,
|
||||
|
||||
@@ -6,15 +6,18 @@ from collections.abc import Callable
|
||||
from datetime import UTC, datetime
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Optional, cast
|
||||
from typing import Any, Optional, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from agents.usage import Usage
|
||||
|
||||
from strix.config import codex
|
||||
from strix.config.loader import load_settings
|
||||
from strix.core.paths import run_dir_for, runtime_state_dir
|
||||
from strix.report.coverage import write_coverage
|
||||
from strix.report.pricing import resolve_litellm_model
|
||||
from strix.report.sarif import write_sarif
|
||||
from strix.report.usage import LLMUsageLedger
|
||||
from strix.report.writer import (
|
||||
read_run_record,
|
||||
write_executive_report,
|
||||
@@ -24,10 +27,6 @@ from strix.report.writer import (
|
||||
from strix.telemetry import posthog, scarf
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.usage import Usage
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_global_report_state: Optional["ReportState"] = None
|
||||
@@ -132,10 +131,6 @@ class ReportState:
|
||||
|
||||
self.scan_results: dict[str, Any] | None = None
|
||||
self.scan_config: dict[str, Any] | None = None
|
||||
# Imported here so importing this module never enters the agents SDK
|
||||
# package (which the warm-up thread may be initializing concurrently).
|
||||
from strix.report.usage import LLMUsageLedger
|
||||
|
||||
self._llm_usage = LLMUsageLedger()
|
||||
self._telemetry_llm_usage_baseline: dict[str, Any] = {}
|
||||
auth_mode = codex.auth_mode(load_settings().llm.model)
|
||||
@@ -343,7 +338,7 @@ class ReportState:
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
usage: "Usage | None",
|
||||
usage: Usage | None,
|
||||
agent_name: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> None:
|
||||
@@ -421,22 +416,6 @@ class ReportState:
|
||||
self.run_record["mcp_connections"] = names
|
||||
self.save_run_data()
|
||||
|
||||
def record_mcp_connection_status(self, status: list[dict[str, Any]]) -> None:
|
||||
"""Persist the run's non-secret MCP connection status roster.
|
||||
|
||||
``status`` is one entry per connection carrying only ``name``,
|
||||
``provider``, ``tool_count``, and ``dead`` (no config, url, token, or
|
||||
auth). Saved as soon as the run connects and rewritten each time a
|
||||
connection dies, so the viewer, which rebuilds its display by re-reading
|
||||
the run's files from disk, can show a live connections panel and health
|
||||
without any in-memory event sink. Kept separate from the
|
||||
``mcp_connections`` name list so neither field repurposes the other.
|
||||
"""
|
||||
if self.run_record.get("mcp_connection_status") == status:
|
||||
return
|
||||
self.run_record["mcp_connection_status"] = status
|
||||
self.save_run_data()
|
||||
|
||||
def set_scan_config(self, config: dict[str, Any]) -> None:
|
||||
self.scan_config = config
|
||||
self.run_record["status"] = "running"
|
||||
|
||||
@@ -27,7 +27,7 @@ Before spawning agents, analyze the target from the scan config/scope and any pr
|
||||
|
||||
## Establish the Threat Model
|
||||
|
||||
Every scan needs one shared answer to "who is the attacker here, and what are they attacking" — black-box or white-box. Without it, five agents derive five different answers and their findings cannot be reconciled. Call `get_threat_model` on the target (a host, a URL, or a repository path) before you spawn hunters; if no model exists yet, derive one and share it with `save_threat_model`. It lives for this scan only — nothing carries over from an earlier run, so every scan derives its own — but within the run every agent reads the same document, and a model written from source is read back by an agent testing the deployment.
|
||||
Every scan needs one shared answer to "who is the attacker here, and what are they attacking" — black-box or white-box. Without it, five agents derive five different answers and their findings cannot be reconciled. Call `get_threat_model` on the target (a host, a URL, or a repository path) before you spawn hunters; if nothing is cached, derive one and persist it with `save_threat_model`. It is cached per target, so a later scan of the same host or tree reads it back instead of paying for it twice, and a model written from source is read back by an agent testing the deployment.
|
||||
|
||||
**When the target includes a repository**, derive it up front: the code tells you the boundaries, entrypoints, and controls before you send a single request.
|
||||
|
||||
|
||||
@@ -23,12 +23,10 @@ from strix.tools.mcp.registry import (
|
||||
McpCallInfo,
|
||||
McpConnectionEntry,
|
||||
McpConnectionRequest,
|
||||
McpConnectionStatus,
|
||||
McpConnectionSummary,
|
||||
McpRegistry,
|
||||
resolve_mcp_call,
|
||||
)
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcpSession
|
||||
|
||||
|
||||
__all__ = [
|
||||
@@ -43,11 +41,8 @@ __all__ = [
|
||||
"McpConnectionConfig",
|
||||
"McpConnectionEntry",
|
||||
"McpConnectionRequest",
|
||||
"McpConnectionStatus",
|
||||
"McpConnectionSummary",
|
||||
"McpConnectionUnavailableError",
|
||||
"McpRegistry",
|
||||
"SupervisedMcpSession",
|
||||
"attach_mcp_requests",
|
||||
"call_mcp",
|
||||
"connect_mcp_servers",
|
||||
|
||||
@@ -26,10 +26,9 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.tools.mcp.client import _errored_tool_output
|
||||
from strix.tools.mcp.client import dispatch_mcp_call
|
||||
from strix.tools.mcp.naming import namespaced_tool_name
|
||||
from strix.tools.mcp.registry import MCP_REGISTRY_CONTEXT_KEY, McpRegistry
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -50,13 +49,6 @@ def _unknown_connection(connection: str, registry: McpRegistry) -> str:
|
||||
return f"Unknown MCP connection {connection!r}. Available connections: {available}."
|
||||
|
||||
|
||||
def _unavailable_connection(connection: str) -> str:
|
||||
return (
|
||||
f"MCP connection {connection!r} is unavailable: its live session failed and "
|
||||
"could not be reconnected, so it is unavailable for the rest of this run."
|
||||
)
|
||||
|
||||
|
||||
def _format_tool(tool: MCPTool) -> str:
|
||||
schema = json.dumps(tool.inputSchema or {"type": "object"}, indent=2, ensure_ascii=False)
|
||||
description = (tool.description or "").strip() or "(no description)"
|
||||
@@ -78,7 +70,6 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]:
|
||||
registry = _registry_from_ctx(ctx)
|
||||
if registry is None or not registry:
|
||||
return {"connections": []}
|
||||
dead_by_name = {status.name: status.dead for status in registry.statuses()}
|
||||
return {
|
||||
"connections": [
|
||||
{
|
||||
@@ -86,7 +77,6 @@ async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]:
|
||||
"name": summary.name,
|
||||
"description": summary.purpose,
|
||||
"tool_count": summary.tool_count,
|
||||
"dead": dead_by_name.get(summary.name, False),
|
||||
}
|
||||
for summary in registry.summaries()
|
||||
]
|
||||
@@ -112,10 +102,7 @@ async def describe_mcp(ctx: RunContextWrapper, connection: str) -> str:
|
||||
entry = registry.get(connection)
|
||||
if entry is None:
|
||||
return _unknown_connection(connection, registry)
|
||||
try:
|
||||
tools = await entry.session.list_tools()
|
||||
except McpConnectionUnavailableError:
|
||||
return _unavailable_connection(connection)
|
||||
tools = await entry.server.list_tools()
|
||||
if not tools:
|
||||
return f"MCP connection {connection!r} offers no tools."
|
||||
header = f"MCP connection {connection!r} offers {len(tools)} tool(s):"
|
||||
@@ -168,10 +155,7 @@ async def call_mcp(
|
||||
return invalid_arguments
|
||||
if arguments is not None and not isinstance(arguments, dict):
|
||||
return invalid_arguments
|
||||
try:
|
||||
available = await entry.session.list_tools()
|
||||
except McpConnectionUnavailableError:
|
||||
return _errored_tool_output(_unavailable_connection(connection))
|
||||
available = await entry.server.list_tools()
|
||||
valid_names = {mcp_tool.name for mcp_tool in available}
|
||||
if tool not in valid_names:
|
||||
offered = ", ".join(sorted(valid_names)) or "(none)"
|
||||
@@ -180,7 +164,8 @@ async def call_mcp(
|
||||
f"Tools this connection offers: {offered}. "
|
||||
"Call describe_mcp for their input schemas."
|
||||
)
|
||||
return await entry.session.dispatch(
|
||||
return await dispatch_mcp_call(
|
||||
entry.server,
|
||||
tool,
|
||||
arguments or {},
|
||||
label=namespaced_tool_name(connection, tool),
|
||||
|
||||
@@ -31,8 +31,6 @@ from agents.mcp import (
|
||||
)
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcpSession
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
@@ -54,18 +52,17 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConnectedMcpServer(NamedTuple):
|
||||
"""One successfully connected MCP connection and how many tools it offers.
|
||||
"""One successfully connected MCP server and how many tools it offers.
|
||||
|
||||
``session`` is the :class:`~strix.tools.mcp.session.SupervisedMcpSession` that
|
||||
owns the live connection on its own task, so the caller cleans it up when the
|
||||
run ends (``await session.aclose()``) and hands it to the run's
|
||||
``server`` is kept so the caller can clean it up when the run ends, and so
|
||||
the caller can hand the live session to the run's
|
||||
:class:`~strix.tools.mcp.registry.McpRegistry`; ``name`` and ``tool_count``
|
||||
let the caller show the user a startup summary and fill the prompt inventory;
|
||||
``notes`` carries the connection's optional free-text description so the
|
||||
caller can surface it as the connection's purpose in the inventory.
|
||||
"""
|
||||
|
||||
session: SupervisedMcpSession
|
||||
server: MCPServer
|
||||
name: str
|
||||
tool_count: int
|
||||
notes: str | None = None
|
||||
@@ -230,75 +227,66 @@ def _errored_tool_output(tool_output: Any) -> dict[str, Any]:
|
||||
return {"success": False, "content": tool_output}
|
||||
|
||||
|
||||
async def _count_session_tools(config: McpConnectionConfig, session: SupervisedMcpSession) -> int:
|
||||
"""Count a connected session's reachable tools for the startup summary.
|
||||
async def _count_server_tools(config: McpConnectionConfig, server: MCPServer) -> int:
|
||||
"""Count a connected server's reachable tools for the startup summary.
|
||||
|
||||
``allowed_tools`` of ``None`` counts every listed tool; a list counts only
|
||||
those names. The count matches what ``describe_mcp`` will show, because the
|
||||
static tool filter built in :func:`_build_server` restricts the server's own
|
||||
``list_tools`` to the same allowlist. The listing goes through the session's
|
||||
owning task like every other call.
|
||||
``list_tools`` to the same allowlist.
|
||||
"""
|
||||
allowed = config.allowed_tools
|
||||
mcp_tools = await session.list_tools()
|
||||
mcp_tools = await server.list_tools()
|
||||
return sum(1 for mcp_tool in mcp_tools if allowed is None or mcp_tool.name in allowed)
|
||||
|
||||
|
||||
async def connect_mcp_servers(
|
||||
configs: list[McpConnectionConfig],
|
||||
) -> list[ConnectedMcpServer]:
|
||||
"""Connect each MCP config on its own supervising task and return the sessions.
|
||||
"""Connect to each MCP server and return its live session.
|
||||
|
||||
Each connection becomes a :class:`~strix.tools.mcp.session.SupervisedMcpSession`
|
||||
that owns ``connect()``, the held-open session, and ``cleanup()`` on one
|
||||
dedicated task, so a later background failure in one session is contained to
|
||||
that task and never cancels the run. Returns one :class:`ConnectedMcpServer`
|
||||
per session that connected, carrying the session (the caller closes it with
|
||||
``await session.aclose()`` when the run ends and hands it to the run's
|
||||
registry) plus the connection name, tool count, and notes. A connection whose
|
||||
initial connect fails is skipped rather than raised (fail-open).
|
||||
|
||||
If this coroutine is itself cancelled mid-attach (the run going down), every
|
||||
session started so far is closed on its own task before the cancellation is
|
||||
re-raised, so nothing is orphaned.
|
||||
Returns one :class:`ConnectedMcpServer` per server that connected, carrying
|
||||
the SDK server (so the caller can clean it up when the run ends and hand it to
|
||||
the run's registry) plus the server name, how many tools it offers, and the
|
||||
connection's notes. Connections that fail are skipped rather than raised.
|
||||
|
||||
Nothing is registered as an agent tool: the caller builds a per-run
|
||||
:class:`~strix.tools.mcp.registry.McpRegistry` from these sessions, and the
|
||||
agent reaches each tool on demand through ``describe_mcp`` / ``call_mcp``.
|
||||
"""
|
||||
connected: list[ConnectedMcpServer] = []
|
||||
sessions: list[SupervisedMcpSession] = []
|
||||
try:
|
||||
for config in configs:
|
||||
session = SupervisedMcpSession(config)
|
||||
sessions.append(session)
|
||||
if not await session.start():
|
||||
# Initial connect failed; already logged inside the session. Drop it.
|
||||
await session.aclose()
|
||||
sessions.remove(session)
|
||||
continue
|
||||
try:
|
||||
tool_count = await _count_session_tools(config, session)
|
||||
except McpConnectionUnavailableError:
|
||||
# The session died between connecting and its first listing; skip it.
|
||||
logger.warning("MCP connection %r died before its first listing", config.name)
|
||||
await session.aclose()
|
||||
sessions.remove(session)
|
||||
continue
|
||||
logger.info("Connected MCP server %r (%d tools)", config.name, tool_count)
|
||||
connected.append(
|
||||
ConnectedMcpServer(
|
||||
session=session, name=config.name, tool_count=tool_count, notes=config.notes
|
||||
)
|
||||
for config in configs:
|
||||
server: MCPServer | None = None
|
||||
try:
|
||||
server = _build_server(config)
|
||||
await server.connect() # type: ignore[no-untyped-call]
|
||||
tool_count = await _count_server_tools(config, server)
|
||||
except Exception:
|
||||
logger.exception("Skipping MCP connection %r", config.name)
|
||||
if server is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
continue
|
||||
except BaseException:
|
||||
# A cancellation (or other non-Exception failure) mid-connect must not
|
||||
# orphan MCP subprocesses or HTTP sessions. Clean up the server being
|
||||
# connected and every server already connected, then re-raise so the
|
||||
# caller still stops. The runner only receives the list on a clean
|
||||
# return, so on an abnormal exit this function owns the cleanup.
|
||||
if server is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
for established in connected:
|
||||
with contextlib.suppress(Exception):
|
||||
await established.server.cleanup() # type: ignore[no-untyped-call]
|
||||
raise
|
||||
|
||||
logger.info("Connected MCP server %r (%d tools)", config.name, tool_count)
|
||||
connected.append(
|
||||
ConnectedMcpServer(
|
||||
server=server, name=config.name, tool_count=tool_count, notes=config.notes
|
||||
)
|
||||
except BaseException:
|
||||
# Cancelled or errored mid-attach: close every session started so far,
|
||||
# each on its own task, then re-raise. The runner only receives the list
|
||||
# on a clean return, so on an abnormal exit this function owns the cleanup.
|
||||
for session in sessions:
|
||||
with contextlib.suppress(BaseException):
|
||||
await session.aclose()
|
||||
raise
|
||||
)
|
||||
|
||||
return connected
|
||||
|
||||
@@ -330,7 +318,7 @@ async def attach_mcp_requests(
|
||||
request = request_by_name[connection.name]
|
||||
registry.add(
|
||||
name=connection.name,
|
||||
session=connection.session,
|
||||
server=connection.server,
|
||||
tool_count=connection.tool_count,
|
||||
purpose=request.purpose or connection.notes,
|
||||
provider=request.provider,
|
||||
|
||||
@@ -25,8 +25,6 @@ from __future__ import annotations
|
||||
import dataclasses
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple
|
||||
|
||||
from strix.tools.mcp.session import SupervisedMcpSession
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.mcp import MCPServer
|
||||
@@ -56,45 +54,24 @@ MCP_DISPATCH_TOOLS = frozenset({CALL_MCP_TOOL, DESCRIBE_MCP_TOOL})
|
||||
class McpConnectionEntry:
|
||||
"""One live MCP connection a scan may reach, keyed by ``name``.
|
||||
|
||||
``session`` is the :class:`~strix.tools.mcp.session.SupervisedMcpSession` that
|
||||
owns the connection on its own task; the dispatch tools list tools and call
|
||||
tools through it (``session.list_tools`` / ``session.dispatch``) so a session
|
||||
failure is contained and can reconnect. ``purpose`` is the human label
|
||||
``list_mcps`` reports as the connection's description (the user's connection
|
||||
notes, or whatever the caller supplies). ``tool_count`` is how many tools the
|
||||
connection offers, also reported by ``list_mcps``. ``result_transform``, when
|
||||
set, runs on each call's structured result at the single dispatch point
|
||||
(strix-pro's sanitizer uses it). ``provider`` is an optional source label
|
||||
(e.g. ``"supabase"``) the caller tags the connection with; the command-line
|
||||
path leaves it ``None``, and event tagging surfaces it when set.
|
||||
|
||||
The connection config the session reconnects with (and its bearer token) lives
|
||||
on ``session`` in memory only. It is reached via :attr:`config` for the
|
||||
reconnect path and is never logged, serialized into the event stream, or
|
||||
written to disk.
|
||||
``server`` is the connected SDK session the dispatch tools list tools on and
|
||||
call tools through. ``purpose`` is the human label ``list_mcps`` reports as the
|
||||
connection's description (the user's connection notes, or whatever the caller
|
||||
supplies). ``tool_count`` is how many tools the connection offers, also
|
||||
reported by ``list_mcps``. ``result_transform``, when set, runs on each call's structured result
|
||||
at the single dispatch point (strix-pro's sanitizer uses it). ``provider`` is
|
||||
an optional source label (e.g. ``"supabase"``) the caller tags the connection
|
||||
with; the command-line path leaves it ``None``, and event tagging surfaces it
|
||||
when set.
|
||||
"""
|
||||
|
||||
session: SupervisedMcpSession
|
||||
server: MCPServer
|
||||
name: str
|
||||
purpose: str | None = None
|
||||
tool_count: int = 0
|
||||
result_transform: ResultTransform | None = None
|
||||
provider: str | None = None
|
||||
|
||||
@property
|
||||
def server(self) -> MCPServer | None:
|
||||
"""The current live server behind the session (swapped on reconnect).
|
||||
|
||||
Kept so existing callers that read ``entry.server`` keep working; new code
|
||||
should call through ``entry.session`` so reconnect and containment apply.
|
||||
"""
|
||||
return self.session.server
|
||||
|
||||
@property
|
||||
def config(self) -> McpConnectionConfig | None:
|
||||
"""The session's reconnect config. Carries the bearer token; never log it."""
|
||||
return self.session.config
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class McpConnectionSummary:
|
||||
@@ -107,24 +84,6 @@ class McpConnectionSummary:
|
||||
provider: str | None = None
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class McpConnectionStatus:
|
||||
"""One connection's live status for the interfaces (the TUI panel, the app
|
||||
strip, and the roster signal the app consumes).
|
||||
|
||||
Non-secret by construction: only the connection ``name``, its ``provider``
|
||||
label, its ``tool_count``, and whether its live session is currently ``dead``
|
||||
(its reconnect-retry gave up). No config, token, url, or purpose rides here.
|
||||
``dead`` is read live off the connection's session at the moment this is
|
||||
built, so a fresh :meth:`McpRegistry.statuses` reflects the current health.
|
||||
"""
|
||||
|
||||
name: str
|
||||
provider: str | None
|
||||
tool_count: int
|
||||
dead: bool
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class McpConnectionRequest:
|
||||
"""A source-agnostic request to attach one MCP connection to a run.
|
||||
@@ -170,28 +129,15 @@ class McpRegistry:
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
session: SupervisedMcpSession | None = None,
|
||||
server: MCPServer | None = None,
|
||||
config: McpConnectionConfig | None = None,
|
||||
server: MCPServer,
|
||||
purpose: str | None = None,
|
||||
tool_count: int = 0,
|
||||
result_transform: ResultTransform | None = None,
|
||||
provider: str | None = None,
|
||||
) -> McpConnectionEntry:
|
||||
"""Register one connection under ``name`` (last write wins).
|
||||
|
||||
Pass ``session`` for a session the engine already supervises (the attach
|
||||
path does this). Pass ``server`` for an already-connected server the caller
|
||||
owns (strix-pro's cloud sessions): it is adopted into a session that runs
|
||||
calls inline against it, and reconnects only when a ``config`` is also
|
||||
given. Exactly one of ``session`` or ``server`` is required.
|
||||
"""
|
||||
if session is None:
|
||||
if server is None:
|
||||
raise ValueError("McpRegistry.add requires either 'session' or 'server'")
|
||||
session = SupervisedMcpSession.adopt(server, name=name, config=config)
|
||||
"""Register one connection under ``name`` (last write wins)."""
|
||||
entry = McpConnectionEntry(
|
||||
session=session,
|
||||
server=server,
|
||||
name=name,
|
||||
purpose=purpose,
|
||||
tool_count=tool_count,
|
||||
@@ -221,23 +167,6 @@ class McpRegistry:
|
||||
for entry in self._entries.values()
|
||||
]
|
||||
|
||||
def statuses(self) -> list[McpConnectionStatus]:
|
||||
"""One live status per connection, in insertion order.
|
||||
|
||||
Reads each connection's ``dead`` flag off its session at call time, so the
|
||||
interfaces (the TUI panel via the Python backend projection, and the
|
||||
roster signal the app consumes) get the current health each time they
|
||||
rebuild. Non-secret: name, provider, tool_count, dead only."""
|
||||
return [
|
||||
McpConnectionStatus(
|
||||
name=entry.name,
|
||||
provider=entry.provider,
|
||||
tool_count=entry.tool_count,
|
||||
dead=entry.session.is_dead,
|
||||
)
|
||||
for entry in self._entries.values()
|
||||
]
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Drop every connection (the sessions themselves are closed by the
|
||||
runner)."""
|
||||
|
||||
@@ -1,536 +0,0 @@
|
||||
"""Own each MCP connection's live session on its own supervising task.
|
||||
|
||||
The bug this fixes: the streamable-HTTP transport (the ``mcp`` SDK) opens an
|
||||
internal anyio task group when ``server.connect()`` runs, and that task group's
|
||||
cancel scope is entered on whatever task called ``connect()`` and stays open for
|
||||
the session's whole life. In the old code that task was the run's main task, the
|
||||
one the agent loop runs on. So when a provider returned an HTTP error on one of
|
||||
the transport's background tasks (for example a ``403`` on a background POST),
|
||||
the task group cancelled its scope, the cancellation
|
||||
propagated to the main task, and the whole scan died with a bare
|
||||
``CancelledError`` (mislabeled as a user interrupt). Teardown then raised
|
||||
"Attempted to exit cancel scope in a different task than it was entered in"
|
||||
because cleanup ran on a different task than connect.
|
||||
|
||||
The fix, mirroring how child agents run on their own ``asyncio.create_task``
|
||||
(see :func:`strix.core.execution.spawn_child_agent`): give each connection its
|
||||
own dedicated supervising task that owns ``connect()``, the session's held-open
|
||||
lifetime, and ``cleanup()``. Three consequences:
|
||||
|
||||
- **Containment.** The transport's cancel scope is now entered on the supervising
|
||||
task, so a background failure cancels only that task. The run and every other
|
||||
connection keep going.
|
||||
- **Co-located teardown.** ``connect()`` and ``cleanup()`` run on the same task,
|
||||
so the "exit cancel scope in a different task" error cannot happen.
|
||||
- **A value, not a cancellation, reaches the caller.** The agent never touches the
|
||||
live session directly. It hands a call to the supervising task over a queue and
|
||||
awaits the result as a value; if the session task dies, the caller gets a
|
||||
"connection unavailable" value instead of a cancellation propagating into the
|
||||
agent loop.
|
||||
|
||||
When a call fails the supervisor rebuilds and reconnects the session once (reusing
|
||||
the same config, so the same bearer token, never re-fetching credentials) and
|
||||
re-runs the one failed call once. If that still fails, the connection is marked
|
||||
dead: every later call returns the standard failed-tool output.
|
||||
|
||||
Security: the connection's :class:`~strix.tools.mcp.config.McpConnectionConfig`
|
||||
holds a live bearer credential and is kept here in memory only, on the same
|
||||
in-process object that already holds the live session. It is never logged,
|
||||
serialized into the run's event stream, or written to disk; :meth:`__repr__`
|
||||
omits it and the token field's own ``repr`` is already suppressed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from agents.mcp import MCPServer
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from strix.tools.mcp.client import ResultTransform
|
||||
from strix.tools.mcp.config import McpConnectionConfig
|
||||
|
||||
# One operation to run against the live session, e.g. ``list_tools`` or a tool
|
||||
# call. Runs on the supervising task (supervised sessions) or inline (adopted
|
||||
# sessions), and its return value becomes the caller's result.
|
||||
Job = Callable[[MCPServer], Awaitable[Any]]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# How long a graceful (sentinel) shutdown waits for the serve loop to drain
|
||||
# before the supervising task is cancelled instead. Bounds teardown so a slow or
|
||||
# hung in-flight call cannot stall it forever.
|
||||
_SHUTDOWN_TIMEOUT = 10.0
|
||||
|
||||
|
||||
class McpConnectionUnavailableError(RuntimeError):
|
||||
"""A dead MCP connection could not be reached and did not come back.
|
||||
|
||||
Raised by :meth:`SupervisedMcpSession.list_tools` when the connection is dead
|
||||
so the read-only dispatch tools (``describe_mcp``) can report it cleanly.
|
||||
:meth:`SupervisedMcpSession.dispatch` does not raise it: a call to a dead
|
||||
connection returns the standard failed-tool output instead.
|
||||
"""
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _Outcome:
|
||||
"""What running one job resolved to: a value, or the connection being dead."""
|
||||
|
||||
value: Any = None
|
||||
dead: bool = False
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _Request:
|
||||
"""One job handed to the supervising task, with the future its result lands in."""
|
||||
|
||||
job: Job
|
||||
future: asyncio.Future[_Outcome]
|
||||
|
||||
|
||||
class SupervisedMcpSession:
|
||||
"""One MCP connection whose live session is owned by a dedicated task.
|
||||
|
||||
Built two ways:
|
||||
|
||||
- :meth:`__init__` + :meth:`start` for a *supervised* session: the engine owns
|
||||
connecting. ``start`` spawns the supervising task, which builds and connects
|
||||
the server on itself and then serves calls handed to it over a queue. This is
|
||||
the path that contains a background session failure to one task.
|
||||
- :meth:`adopt` for an *adopted* session: the caller already holds a connected
|
||||
server (strix-pro's cloud sessions, and the test fakes). There is no
|
||||
supervising task; calls run inline against the given server. Reconnect works
|
||||
only when a config was supplied.
|
||||
|
||||
Public async API used by the dispatch tools: :meth:`list_tools` and
|
||||
:meth:`dispatch`. Lifecycle: :meth:`start`, :meth:`aclose`. Read-only:
|
||||
:attr:`name`, :attr:`server`, :attr:`config`, :attr:`is_dead`.
|
||||
"""
|
||||
|
||||
def __init__(self, config: McpConnectionConfig) -> None:
|
||||
self._name = config.name
|
||||
self._config: McpConnectionConfig | None = config
|
||||
self._server: MCPServer | None = None
|
||||
self._supervised = True
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
self._queue: asyncio.Queue[_Request | None] | None = None
|
||||
self._ready: asyncio.Future[bool] | None = None
|
||||
self._pending: set[asyncio.Future[_Outcome]] = set()
|
||||
self._dead = False
|
||||
self._closing = False
|
||||
self._on_dead: Callable[[], None] | None = None
|
||||
# Guards the idle-death self-heal against a flapping server: set after an
|
||||
# idle reconnect, cleared once a real call runs. If the session dies idle
|
||||
# again before serving anything, we give up instead of reconnecting in a
|
||||
# tight loop.
|
||||
self._healed_without_progress = False
|
||||
|
||||
@classmethod
|
||||
def adopt(
|
||||
cls,
|
||||
server: MCPServer,
|
||||
*,
|
||||
name: str,
|
||||
config: McpConnectionConfig | None = None,
|
||||
) -> SupervisedMcpSession:
|
||||
"""Wrap an already-connected server without a supervising task.
|
||||
|
||||
Calls run inline against ``server`` on the caller's task, matching the old
|
||||
direct-dispatch behavior. Reconnect is available only when ``config`` is
|
||||
given; otherwise a failed call marks the connection dead.
|
||||
"""
|
||||
self = cls.__new__(cls)
|
||||
self._name = name
|
||||
self._config = config
|
||||
self._server = server
|
||||
self._supervised = False
|
||||
self._task = None
|
||||
self._queue = None
|
||||
self._ready = None
|
||||
self._pending = set()
|
||||
self._dead = False
|
||||
self._closing = False
|
||||
self._on_dead = None
|
||||
self._healed_without_progress = False
|
||||
return self
|
||||
|
||||
# -- read-only accessors --------------------------------------------------
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def server(self) -> MCPServer | None:
|
||||
"""The current live server, or ``None`` once dead. Swapped on reconnect."""
|
||||
return self._server
|
||||
|
||||
@property
|
||||
def config(self) -> McpConnectionConfig | None:
|
||||
"""The connection config kept for reconnect. Carries the bearer token, so
|
||||
never log or serialize this."""
|
||||
return self._config
|
||||
|
||||
@property
|
||||
def is_dead(self) -> bool:
|
||||
return self._dead
|
||||
|
||||
def set_on_dead(self, callback: Callable[[], None] | None) -> None:
|
||||
"""Register a one-shot callback fired when the connection transitions to dead.
|
||||
|
||||
The callback runs on whatever task marks the connection dead (the
|
||||
supervising task for a supervised session, the caller's task for an
|
||||
adopted one), so it must not block. It fires at most once, on the
|
||||
healthy->dead edge, and never for a connection that only ever shut down
|
||||
cleanly. The interfaces use it to push a live "offline" status without
|
||||
polling. Exceptions from the callback are swallowed (logged) so a status
|
||||
push can never take down the session task.
|
||||
"""
|
||||
self._on_dead = callback
|
||||
|
||||
def _mark_dead(self) -> None:
|
||||
"""Flip the connection to dead and fire ``on_dead`` once on the transition."""
|
||||
if self._dead:
|
||||
return
|
||||
self._dead = True
|
||||
callback = self._on_dead
|
||||
if callback is None:
|
||||
return
|
||||
try:
|
||||
callback()
|
||||
except Exception:
|
||||
logger.exception("MCP on_dead callback for %r failed", self._name)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
# Deliberately omits the config so the bearer token can never reach a log
|
||||
# line through an accidental repr of this object.
|
||||
return f"SupervisedMcpSession(name={self._name!r}, dead={self._dead})"
|
||||
|
||||
# -- lifecycle ------------------------------------------------------------
|
||||
|
||||
async def start(self) -> bool:
|
||||
"""Spawn the supervising task, connect on it, and wait until it is ready.
|
||||
|
||||
Returns ``True`` when the session connected, ``False`` when the initial
|
||||
connect failed (the caller then skips this connection, fail-open). Only
|
||||
valid for a supervised session.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
self._queue = asyncio.Queue()
|
||||
self._ready = loop.create_future()
|
||||
self._task = asyncio.create_task(self._supervise(), name=f"mcp-session-{self._name}")
|
||||
return await self._ready
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Shut the connection down and clean up its session on its owning task.
|
||||
|
||||
For a connected supervised session this signals the supervising task with a
|
||||
sentinel so ``cleanup()`` runs on the same task that ran ``connect()``,
|
||||
giving an orderly shutdown the supervisor tells apart from a session death.
|
||||
Teardown is always bounded: if the serve loop cannot drain the sentinel in
|
||||
time (a slow or hung in-flight call), or the session never finished
|
||||
connecting (including a connect cancelled mid-await), the task is cancelled
|
||||
instead. ``_closing`` is set first, so the supervisor treats that
|
||||
cancellation as shutdown and still cleans up on its own task.
|
||||
"""
|
||||
self._closing = True
|
||||
if self._supervised and self._task is not None:
|
||||
if not self._task.done():
|
||||
# A cancelled readiness future (the connect was cancelled mid-await)
|
||||
# counts as "not connected": never call ``.result()`` on it, which
|
||||
# would raise here and skip the cleanup below.
|
||||
connected = (
|
||||
self._ready is not None
|
||||
and self._ready.done()
|
||||
and not self._ready.cancelled()
|
||||
and self._ready.result()
|
||||
)
|
||||
if connected and self._queue is not None:
|
||||
# Reached the serve loop: a sentinel gives a clean, cancel-free
|
||||
# teardown, with cleanup() running on the supervising task. Bound
|
||||
# it, though: a hung in-flight call would otherwise leave the
|
||||
# sentinel queued behind it forever, so cancel the task if the
|
||||
# drain does not finish in time (wait_for cancels it on timeout).
|
||||
with contextlib.suppress(Exception):
|
||||
await self._queue.put(None)
|
||||
with contextlib.suppress(
|
||||
asyncio.TimeoutError, asyncio.CancelledError, Exception
|
||||
):
|
||||
await asyncio.wait_for(self._task, _SHUTDOWN_TIMEOUT)
|
||||
else:
|
||||
# Still stuck in connect(), never connected, or connect
|
||||
# cancelled: cancel to unstick it.
|
||||
self._task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||||
await self._task
|
||||
else:
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
|
||||
# -- caller-facing operations --------------------------------------------
|
||||
|
||||
async def list_tools(self) -> list[MCPTool]:
|
||||
"""List the connection's tools, reconnecting once if the session died.
|
||||
|
||||
Raises :class:`McpConnectionUnavailableError` when the connection is dead.
|
||||
"""
|
||||
outcome = await self._run_job(lambda server: server.list_tools())
|
||||
if outcome.dead:
|
||||
raise McpConnectionUnavailableError(self._unavailable_message())
|
||||
return cast("list[MCPTool]", outcome.value)
|
||||
|
||||
async def dispatch(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
*,
|
||||
label: str,
|
||||
result_transform: ResultTransform | None = None,
|
||||
) -> Any:
|
||||
"""Run one tool call, reconnecting once and retrying once on session death.
|
||||
|
||||
Returns the tool output on success, or the standard failed-tool output
|
||||
(``success: False``) with a "connection unavailable" message when the
|
||||
connection is dead.
|
||||
"""
|
||||
from strix.tools.mcp.client import dispatch_mcp_call
|
||||
|
||||
async def job(server: MCPServer) -> Any:
|
||||
return await dispatch_mcp_call(
|
||||
server,
|
||||
tool_name,
|
||||
arguments,
|
||||
label=label,
|
||||
result_transform=result_transform,
|
||||
)
|
||||
|
||||
outcome = await self._run_job(job)
|
||||
if outcome.dead:
|
||||
from strix.tools.mcp.client import _errored_tool_output
|
||||
|
||||
return _errored_tool_output(self._unavailable_message())
|
||||
return outcome.value
|
||||
|
||||
# -- job routing ----------------------------------------------------------
|
||||
|
||||
async def _run_job(self, job: Job) -> _Outcome:
|
||||
"""Route one job to the owning task (supervised) or run it inline (adopted)."""
|
||||
if self._supervised:
|
||||
return await self._submit(job)
|
||||
return await self._execute(job)
|
||||
|
||||
async def _submit(self, job: Job) -> _Outcome:
|
||||
"""Hand a job to the supervising task and await its result as a value."""
|
||||
if self._dead or self._closing or self._task is None or self._task.done():
|
||||
return _Outcome(dead=True)
|
||||
loop = asyncio.get_running_loop()
|
||||
future: asyncio.Future[_Outcome] = loop.create_future()
|
||||
self._pending.add(future)
|
||||
if self._queue is None:
|
||||
self._pending.discard(future)
|
||||
return _Outcome(dead=True)
|
||||
await self._queue.put(_Request(job=job, future=future))
|
||||
# The task may have ended between the guard above and the put; ``_fail_pending``
|
||||
# would then never see this future, so resolve it here.
|
||||
if self._task.done() and not future.done():
|
||||
self._pending.discard(future)
|
||||
return _Outcome(dead=True)
|
||||
return await future
|
||||
|
||||
# -- the supervising task -------------------------------------------------
|
||||
|
||||
async def _supervise(self) -> None:
|
||||
"""Own the session for its whole life on one task: connect, serve, clean up."""
|
||||
try:
|
||||
self._server = await self._open()
|
||||
except asyncio.CancelledError:
|
||||
# The connect was cancelled (the run is going down, or the transport
|
||||
# scope cancelled mid-connect). Report not-ready so the attach path
|
||||
# treats it as a skipped connection; do not propagate.
|
||||
self._report_ready(value=False)
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("Skipping MCP connection %r", self._name)
|
||||
self._report_ready(value=False)
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
return
|
||||
|
||||
self._report_ready(value=True)
|
||||
try:
|
||||
await self._serve_loop()
|
||||
finally:
|
||||
await self._safe_cleanup()
|
||||
self._fail_pending()
|
||||
|
||||
async def _serve_loop(self) -> None:
|
||||
assert self._queue is not None
|
||||
while True:
|
||||
try:
|
||||
request = await self._queue.get()
|
||||
except asyncio.CancelledError:
|
||||
# A cancellation while idle is the transport's task group cancelling
|
||||
# this supervising task because a background session task failed.
|
||||
# Contained here. If we are closing, this is an ordinary shutdown,
|
||||
# so let it propagate. Otherwise try to self-heal once: reconnect a
|
||||
# fresh session and keep serving. The flag stops a flapping server
|
||||
# (one that dies again before serving any call) from reconnecting in
|
||||
# a tight loop; there we give up and mark the connection dead. Later
|
||||
# calls then short-circuit to the dead output without this task.
|
||||
if self._closing:
|
||||
raise
|
||||
if not self._healed_without_progress:
|
||||
logger.warning(
|
||||
"MCP connection %r session died while idle; reconnecting once",
|
||||
self._name,
|
||||
)
|
||||
if await self._reconnect():
|
||||
logger.info(
|
||||
"MCP connection %r reconnected after an idle death", self._name
|
||||
)
|
||||
self._healed_without_progress = True
|
||||
continue
|
||||
else:
|
||||
logger.warning(
|
||||
"MCP connection %r died again before serving a call; "
|
||||
"marking it unavailable",
|
||||
self._name,
|
||||
)
|
||||
self._mark_dead()
|
||||
await self._safe_cleanup()
|
||||
return
|
||||
if request is None: # shutdown sentinel
|
||||
return
|
||||
outcome = await self._execute(request.job)
|
||||
# A served call is real progress: clear the idle-heal guard so a future
|
||||
# idle death is again allowed one reconnect.
|
||||
self._healed_without_progress = False
|
||||
if not request.future.done():
|
||||
request.future.set_result(outcome)
|
||||
self._pending.discard(request.future)
|
||||
|
||||
# -- run one job with reconnect-once + retry-once -------------------------
|
||||
|
||||
async def _execute(self, job: Job) -> _Outcome:
|
||||
"""Run one job; on a session failure reconnect once and retry it once."""
|
||||
if self._dead or self._server is None:
|
||||
return _Outcome(dead=True)
|
||||
try:
|
||||
return _Outcome(value=await job(self._server))
|
||||
except asyncio.CancelledError:
|
||||
# For a supervised session a cancellation here is the transport scope
|
||||
# dying under an in-flight call: a session death, not a real cancel
|
||||
# (shutdown never cancels the task, it uses the sentinel). For an
|
||||
# adopted session there is no such scope, so a cancel is real.
|
||||
if not self._supervised or self._closing:
|
||||
raise
|
||||
logger.warning(
|
||||
"MCP connection %r was cancelled mid-call (session died); reconnecting once",
|
||||
self._name,
|
||||
)
|
||||
except Exception: # noqa: BLE001 - any call failure is treated as a session death
|
||||
logger.warning(
|
||||
"MCP connection %r failed mid-call; reconnecting once", self._name
|
||||
)
|
||||
|
||||
if not await self._reconnect():
|
||||
self._mark_dead()
|
||||
return _Outcome(dead=True)
|
||||
|
||||
try:
|
||||
return _Outcome(value=await job(self._server))
|
||||
except asyncio.CancelledError:
|
||||
if not self._supervised or self._closing:
|
||||
raise
|
||||
logger.warning(
|
||||
"MCP connection %r was cancelled again after reconnect; marking it unavailable",
|
||||
self._name,
|
||||
)
|
||||
self._mark_dead()
|
||||
return _Outcome(dead=True)
|
||||
except Exception: # noqa: BLE001 - any retry failure means the connection is dead
|
||||
logger.warning(
|
||||
"MCP connection %r failed again after reconnect; marking it unavailable",
|
||||
self._name,
|
||||
)
|
||||
self._mark_dead()
|
||||
return _Outcome(dead=True)
|
||||
|
||||
async def _reconnect(self) -> bool:
|
||||
"""Rebuild and reconnect the session once, reusing the stored config/token."""
|
||||
await self._safe_cleanup()
|
||||
if self._config is None:
|
||||
return False
|
||||
try:
|
||||
self._server = await self._open()
|
||||
except asyncio.CancelledError:
|
||||
if self._closing:
|
||||
raise
|
||||
logger.warning("MCP reconnect for %r was cancelled; giving up", self._name)
|
||||
self._server = None
|
||||
return False
|
||||
except Exception:
|
||||
logger.exception("MCP reconnect for %r failed", self._name)
|
||||
self._server = None
|
||||
return False
|
||||
logger.info("MCP connection %r reconnected", self._name)
|
||||
return True
|
||||
|
||||
async def _open(self) -> MCPServer:
|
||||
"""Build and connect the SDK server, reusing the existing setup steps.
|
||||
|
||||
If ``connect()`` fails, the just-built server is cleaned up here on this
|
||||
same task before the error propagates, so a failed connect never orphans
|
||||
an MCP subprocess or half-open HTTP session.
|
||||
"""
|
||||
from strix.tools.mcp.client import _build_server
|
||||
|
||||
if self._config is None:
|
||||
raise RuntimeError(f"MCP connection {self._name!r} has no config to connect")
|
||||
server = _build_server(self._config)
|
||||
try:
|
||||
await server.connect() # type: ignore[no-untyped-call]
|
||||
except BaseException:
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
raise
|
||||
return server
|
||||
|
||||
# -- helpers --------------------------------------------------------------
|
||||
|
||||
async def _safe_cleanup(self) -> None:
|
||||
server = self._server
|
||||
self._server = None
|
||||
if server is None:
|
||||
return
|
||||
with contextlib.suppress(Exception):
|
||||
await server.cleanup() # type: ignore[no-untyped-call]
|
||||
|
||||
def _report_ready(self, value: bool) -> None:
|
||||
if self._ready is not None and not self._ready.done():
|
||||
self._ready.set_result(value)
|
||||
|
||||
def _fail_pending(self) -> None:
|
||||
for future in self._pending:
|
||||
if not future.done():
|
||||
future.set_result(_Outcome(dead=True))
|
||||
self._pending.clear()
|
||||
|
||||
def _unavailable_message(self) -> str:
|
||||
return (
|
||||
f"MCP connection {self._name!r} is unavailable: its live session could "
|
||||
"not be reached and a reconnect attempt failed. It is marked unavailable "
|
||||
"for the rest of this run."
|
||||
)
|
||||
@@ -1,29 +1,28 @@
|
||||
"""Run-scoped threat models — mirrored to ``{state_dir}/threat_models.json``.
|
||||
"""Target-scoped threat models — cached under ``~/.strix/threat-models``.
|
||||
|
||||
A threat model is the scan's shared answer to who the attacker is, where the
|
||||
trust boundaries sit, and what counts as critical for the target. One agent
|
||||
derives it and every other agent on the same run reads it back instead of
|
||||
A threat model describes the target, not the scan: a host, an application, an
|
||||
API, a repository, or whatever else the engagement is pointed at. It stays
|
||||
valid across unrelated runs against the same target, so it is keyed by target
|
||||
identity rather than by run id — one agent derives it, every later agent in
|
||||
this run and in future runs against the same target reads it back instead of
|
||||
re-deriving trust boundaries from scratch.
|
||||
|
||||
It does not outlive the scan. The mirror lives in the run's own state directory
|
||||
and exists only so a resumed scan keeps the baseline its earlier agents agreed
|
||||
on; a new scan against the same host or checkout starts with no model and
|
||||
derives its own. Agents do spell one target several ways within a run — the URL
|
||||
they were handed, the page they happen to be testing, a checkout path — so a
|
||||
model is keyed by a normalized target identity to keep them converging on one
|
||||
document instead of each starting a fresh one.
|
||||
Where the target is a checkout, the model is additionally pinned to the git
|
||||
revision, so a moved ``HEAD`` marks it stale. Black-box targets have no
|
||||
revision to pin to; those age out instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
from datetime import UTC, datetime
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
@@ -36,19 +35,16 @@ from strix.core.agents import AgentCoordinator
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_CACHE_DIR = Path.home() / ".strix" / "threat-models"
|
||||
_MAX_MODEL_BYTES = 512 * 1024
|
||||
_MIN_MODEL_CHARS = 400
|
||||
_MIN_AMENDMENT_CHARS = 80
|
||||
_MAX_AMENDMENTS = 40
|
||||
_GIT_TIMEOUT_SECONDS = 10
|
||||
_UNVERSIONED = "unversioned"
|
||||
_MAX_AGE_DAYS = 14
|
||||
_DEFAULT_PORTS = {"http": "80", "https": "443"}
|
||||
|
||||
_store_lock = threading.RLock()
|
||||
|
||||
# The whole store: target identity -> model. It holds exactly the models this
|
||||
# scan derived, and is mirrored to the run's state directory for resume.
|
||||
_MODELS: dict[str, dict[str, Any]] = {}
|
||||
_store_path: Path | None = None
|
||||
_cache_lock = threading.RLock()
|
||||
|
||||
_REQUIRED_SECTIONS = (
|
||||
"overview",
|
||||
@@ -99,7 +95,7 @@ def _remote_authority(target: str) -> str:
|
||||
|
||||
|
||||
def _normalize_remote_target(target: str) -> str:
|
||||
"""Collapse the spellings of one remote target onto a single key."""
|
||||
"""Collapse the spellings of one remote target onto a single cache key."""
|
||||
authority = _remote_authority(target)
|
||||
if not authority:
|
||||
return re.sub(r"\s+", " ", target.lower()).strip()
|
||||
@@ -136,23 +132,30 @@ def _normalize_git_remote(remote: str) -> str:
|
||||
return normalized.removesuffix(".git")
|
||||
|
||||
|
||||
def _target_identity(target: str) -> str:
|
||||
"""Return the stable identity a model is stored under.
|
||||
def _target_identity(target: str) -> tuple[str, str]:
|
||||
"""Return the (stable identity, revision) pair a cached model is keyed on.
|
||||
|
||||
A checkout is keyed on its remote, so the same repository checked out at
|
||||
two paths shares one model and a subdirectory resolves to the whole tree.
|
||||
Everything else — a host, a URL, an API base, a named scope — is keyed on
|
||||
its normalized form. Both routes run through the same normalization, so a
|
||||
checkout and the URL it was cloned from land on one key.
|
||||
A checkout is keyed on its remote (so the same repository cloned to two
|
||||
paths shares one model, and a subdirectory resolves to the whole tree) and
|
||||
pinned to ``HEAD``. Everything else — a host, a URL, an API base, a named
|
||||
scope — is keyed on its normalized form and carries no revision. Both
|
||||
routes run through the same normalization, so a checkout and the URL it
|
||||
was cloned from land on one key.
|
||||
"""
|
||||
directory = _local_directory(target)
|
||||
if directory is None:
|
||||
return _normalize_remote_target(target).removesuffix(".git")
|
||||
return _normalize_remote_target(target).removesuffix(".git"), _UNVERSIONED
|
||||
remote = _git(directory, ["config", "--get", "remote.origin.url"])
|
||||
revision = _git(directory, ["rev-parse", "HEAD"]) or _UNVERSIONED
|
||||
if remote:
|
||||
return _normalize_git_remote(remote)
|
||||
return _normalize_git_remote(remote), revision
|
||||
toplevel = _git(directory, ["rev-parse", "--show-toplevel"])
|
||||
return toplevel or str(directory)
|
||||
return toplevel or str(directory), revision
|
||||
|
||||
|
||||
def _cache_path(identity: str) -> Path:
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:16]
|
||||
return _CACHE_DIR / f"{digest}.json"
|
||||
|
||||
|
||||
def _snap_to_scan_target(raw: str, scan_targets: list[str]) -> str:
|
||||
@@ -160,13 +163,13 @@ def _snap_to_scan_target(raw: str, scan_targets: list[str]) -> str:
|
||||
|
||||
Agents name the same target differently — one passes the URL it was given,
|
||||
the next the page it happens to be testing, a third the checkout path. Left
|
||||
alone those become separate keys, every lookup misses, and each agent
|
||||
alone those become separate cache keys, every lookup misses, and each agent
|
||||
quietly derives its own model, which is the exact failure the shared model
|
||||
exists to prevent. So a target that is recognisably one of the scan's own
|
||||
targets is resolved to that target instead.
|
||||
"""
|
||||
identity = _target_identity(raw)
|
||||
scoped = [(target, _target_identity(target)) for target in scan_targets]
|
||||
identity, _ = _target_identity(raw)
|
||||
scoped = [(target, _target_identity(target)[0]) for target in scan_targets]
|
||||
if any(known == identity for _, known in scoped):
|
||||
return raw
|
||||
|
||||
@@ -205,51 +208,38 @@ def _resolve_target(
|
||||
return (_snap_to_scan_target(raw, known) if known else raw), None
|
||||
|
||||
|
||||
def hydrate_threat_models_from_disk(state_dir: Path) -> None:
|
||||
"""Point the store at this run's mirror and load whatever it already holds.
|
||||
|
||||
A resumed scan is the same scan, so its agents have to keep the baseline
|
||||
the earlier ones agreed on. The mirror lives under the run directory, so a
|
||||
different scan never reads it.
|
||||
"""
|
||||
global _store_path # noqa: PLW0603
|
||||
_store_path = state_dir / "threat_models.json"
|
||||
with _store_lock:
|
||||
_MODELS.clear()
|
||||
if not _store_path.is_file():
|
||||
return
|
||||
try:
|
||||
data = json.loads(_store_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.exception(
|
||||
"threat_models.json at %s is unreadable; starting with no models",
|
||||
_store_path,
|
||||
)
|
||||
return
|
||||
if not isinstance(data, dict):
|
||||
return
|
||||
_MODELS.update(
|
||||
{
|
||||
identity: model
|
||||
for identity, model in data.items()
|
||||
if isinstance(identity, str) and isinstance(model, dict)
|
||||
}
|
||||
)
|
||||
logger.info("threat models hydrated from %s (%d)", _store_path, len(_MODELS))
|
||||
|
||||
|
||||
def _persist_locked() -> None:
|
||||
"""Mirror the store to disk. Callers must already hold ``_store_lock``.
|
||||
|
||||
Serializing and renaming in one critical section keeps a writer holding an
|
||||
older serialization from winning the rename and dropping a concurrent
|
||||
agent's model or amendment.
|
||||
"""
|
||||
path = _store_path
|
||||
if path is None:
|
||||
return
|
||||
def _is_expired(created_at: str | None) -> bool:
|
||||
if not created_at:
|
||||
return True
|
||||
try:
|
||||
created = datetime.fromisoformat(created_at)
|
||||
except ValueError:
|
||||
return True
|
||||
if created.tzinfo is None:
|
||||
created = created.replace(tzinfo=UTC)
|
||||
return datetime.now(UTC) - created > timedelta(days=_MAX_AGE_DAYS)
|
||||
|
||||
|
||||
def _missing_sections(content: str) -> list[str]:
|
||||
lowered = content.lower()
|
||||
return [section for section in _REQUIRED_SECTIONS if section not in lowered]
|
||||
|
||||
|
||||
def _read_cache(path: Path) -> dict[str, Any] | None:
|
||||
"""Load a cached model. Callers must already hold ``_cache_lock``."""
|
||||
if not path.is_file():
|
||||
return None
|
||||
try:
|
||||
cached = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.exception("threat model cache at %s is unreadable", path)
|
||||
return None
|
||||
return cached if isinstance(cached, dict) else None
|
||||
|
||||
|
||||
def _write_cache(path: Path, payload: dict[str, Any]) -> str | None:
|
||||
"""Atomically persist a model. Callers must already hold ``_cache_lock``."""
|
||||
try:
|
||||
payload = json.dumps(_MODELS, ensure_ascii=False, default=str)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
@@ -259,61 +249,85 @@ def _persist_locked() -> None:
|
||||
suffix=".tmp",
|
||||
delete=False,
|
||||
) as tmp:
|
||||
tmp.write(payload)
|
||||
tmp.write(json.dumps(payload, ensure_ascii=False))
|
||||
tmp_path = Path(tmp.name)
|
||||
tmp_path.replace(path)
|
||||
except OSError:
|
||||
logger.exception("threat model mirror to %s failed", path)
|
||||
except OSError as exc:
|
||||
logger.exception("threat model persist to %s failed", path)
|
||||
return f"Failed to persist threat model: {exc}"
|
||||
return None
|
||||
|
||||
|
||||
def _missing_sections(content: str) -> list[str]:
|
||||
lowered = content.lower()
|
||||
return [section for section in _REQUIRED_SECTIONS if section not in lowered]
|
||||
|
||||
|
||||
def _amendments_of(model: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
raw = model.get("amendments")
|
||||
def _amendments_of(cached: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
raw = cached.get("amendments")
|
||||
if not isinstance(raw, list):
|
||||
return []
|
||||
return [item for item in raw if isinstance(item, dict)]
|
||||
|
||||
|
||||
def _not_found(identity: str) -> dict[str, Any]:
|
||||
def _not_found(identity: str, revision: str) -> dict[str, Any]:
|
||||
return {
|
||||
"success": True,
|
||||
"found": False,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"message": (
|
||||
"No threat model for this target on this scan. Nothing carries over "
|
||||
"from other scans, so derive one — from the code if you have it, from "
|
||||
"recon output if you do not — and share it with save_threat_model, so "
|
||||
"every agent on this scan works from one view of the trust boundaries "
|
||||
"instead of each inventing their own."
|
||||
"No threat model cached for this target. Derive one — from the code if "
|
||||
"you have it, from recon output if you do not — and persist it with "
|
||||
"save_threat_model, so every agent on this scan shares one view of the "
|
||||
"trust boundaries instead of each inventing their own."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _staleness(cached: dict[str, Any], revision: str) -> tuple[bool, str | None]:
|
||||
"""Decide whether a cached model can still be trusted, and why not."""
|
||||
if revision != _UNVERSIONED:
|
||||
if cached.get("revision") == revision:
|
||||
return False, None
|
||||
return True, (
|
||||
"This model was derived against a different revision. Use it as a "
|
||||
"starting point, re-check the boundaries it names against the current "
|
||||
"tree, and save the corrected version."
|
||||
)
|
||||
created_at = cached.get("created_at")
|
||||
if not _is_expired(created_at if isinstance(created_at, str) else None):
|
||||
return False, None
|
||||
return True, (
|
||||
f"This model is more than {_MAX_AGE_DAYS} days old and there is no revision "
|
||||
"to pin it to, so the target may have moved under it. Treat its surface "
|
||||
"inventory as a lead list to re-confirm during recon, not as fact, and save "
|
||||
"the corrected version."
|
||||
)
|
||||
|
||||
|
||||
def _get_impl(target: str, scan_targets: list[str] | None = None) -> dict[str, Any]:
|
||||
resolved, error = _resolve_target(target, scan_targets)
|
||||
if resolved is None:
|
||||
return {"success": False, "error": error}
|
||||
|
||||
identity = _target_identity(resolved)
|
||||
with _store_lock:
|
||||
model = _MODELS.get(identity)
|
||||
if model is None:
|
||||
return _not_found(identity)
|
||||
content = model.get("content")
|
||||
amendments = list(_amendments_of(model))
|
||||
identity, revision = _target_identity(resolved)
|
||||
path = _cache_path(identity)
|
||||
with _cache_lock:
|
||||
cached = _read_cache(path)
|
||||
if cached is None:
|
||||
return _not_found(identity, revision)
|
||||
content = cached.get("content")
|
||||
if not isinstance(content, str) or not content.strip():
|
||||
return _not_found(identity)
|
||||
return _not_found(identity, revision)
|
||||
|
||||
stale, stale_message = _staleness(cached, revision)
|
||||
result: dict[str, Any] = {
|
||||
"success": True,
|
||||
"found": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"cached_revision": cached.get("revision"),
|
||||
"created_at": cached.get("created_at"),
|
||||
"stale": stale,
|
||||
"content": content,
|
||||
}
|
||||
amendments = _amendments_of(cached)
|
||||
if amendments:
|
||||
result["amendments"] = amendments
|
||||
result["amendments_note"] = (
|
||||
@@ -321,6 +335,8 @@ def _get_impl(target: str, scan_targets: list[str] | None = None) -> dict[str, A
|
||||
"correct or extend it and have not been folded in yet - read them as "
|
||||
"part of the model, and prefer the later one where they conflict."
|
||||
)
|
||||
if stale_message:
|
||||
result["message"] = stale_message
|
||||
return result
|
||||
|
||||
|
||||
@@ -360,21 +376,25 @@ def _save_impl(
|
||||
),
|
||||
}
|
||||
|
||||
identity = _target_identity(resolved)
|
||||
with _store_lock:
|
||||
existing = _MODELS.get(identity)
|
||||
identity, revision = _target_identity(resolved)
|
||||
path = _cache_path(identity)
|
||||
payload: dict[str, Any] = {
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"created_at": datetime.now(UTC).isoformat(),
|
||||
"created_by": agent_name,
|
||||
"content": body,
|
||||
}
|
||||
with _cache_lock:
|
||||
existing = _read_cache(path)
|
||||
folded = len(_amendments_of(existing)) if existing else 0
|
||||
_MODELS[identity] = {
|
||||
"target": identity,
|
||||
"written_at": datetime.now(UTC).isoformat(),
|
||||
"written_by": agent_name,
|
||||
"content": body,
|
||||
}
|
||||
_persist_locked()
|
||||
error = _write_cache(path, payload)
|
||||
if error:
|
||||
return {"success": False, "error": error}
|
||||
|
||||
message = (
|
||||
"Threat model shared with this scan. Subagents should call get_threat_model "
|
||||
"before they start, and treat its trust boundaries as the shared baseline."
|
||||
"Threat model saved. Subagents should call get_threat_model before they "
|
||||
"start, and treat its trust boundaries as the shared baseline."
|
||||
)
|
||||
if folded:
|
||||
message += (
|
||||
@@ -384,35 +404,34 @@ def _save_impl(
|
||||
return {
|
||||
"success": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"amendments_cleared": folded,
|
||||
"message": message,
|
||||
}
|
||||
|
||||
|
||||
def _append_amendment(
|
||||
identity: str, amendment: dict[str, Any]
|
||||
path: Path, amendment: dict[str, Any]
|
||||
) -> tuple[list[dict[str, Any]] | None, str | None]:
|
||||
"""Add an amendment to the stored model. Returns (amendments, error)."""
|
||||
with _store_lock:
|
||||
model = _MODELS.get(identity)
|
||||
if model is None or not str(model.get("content", "")).strip():
|
||||
"""Add an amendment to the cached model. Returns (amendments, error)."""
|
||||
with _cache_lock:
|
||||
cached = _read_cache(path)
|
||||
if cached is None or not str(cached.get("content", "")).strip():
|
||||
return None, (
|
||||
"No threat model exists for this target yet, so there is nothing to "
|
||||
"amend. Derive the base model and call save_threat_model instead."
|
||||
)
|
||||
amendments = _amendments_of(model)
|
||||
amendments = _amendments_of(cached)
|
||||
if len(amendments) >= _MAX_AMENDMENTS:
|
||||
return None, (
|
||||
f"This model already carries {len(amendments)} amendments. Fold them "
|
||||
"into the base model with save_threat_model before adding more."
|
||||
)
|
||||
candidate = [*amendments, amendment]
|
||||
sized = {**model, "amendments": candidate}
|
||||
if len(json.dumps(sized, ensure_ascii=False).encode("utf-8")) > _MAX_MODEL_BYTES:
|
||||
amendments.append(amendment)
|
||||
cached["amendments"] = amendments
|
||||
if len(json.dumps(cached, ensure_ascii=False).encode("utf-8")) > _MAX_MODEL_BYTES:
|
||||
return None, "Threat model with this amendment exceeds 512KB; tighten it."
|
||||
model["amendments"] = candidate
|
||||
_persist_locked()
|
||||
return candidate, None
|
||||
return amendments, _write_cache(path, cached)
|
||||
|
||||
|
||||
def _amend_impl(
|
||||
@@ -436,12 +455,13 @@ def _amend_impl(
|
||||
),
|
||||
}
|
||||
|
||||
identity = _target_identity(resolved)
|
||||
identity, revision = _target_identity(resolved)
|
||||
amendments, amend_error = _append_amendment(
|
||||
identity,
|
||||
_cache_path(identity),
|
||||
{
|
||||
"at": datetime.now(UTC).isoformat(),
|
||||
"by": agent_name,
|
||||
"revision": revision,
|
||||
"content": body,
|
||||
},
|
||||
)
|
||||
@@ -451,6 +471,7 @@ def _amend_impl(
|
||||
return {
|
||||
"success": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"amendment_count": len(amendments),
|
||||
"message": (
|
||||
"Amendment recorded. Agents calling get_threat_model will now see it "
|
||||
@@ -479,25 +500,25 @@ def _scan_targets(ctx: RunContextWrapper) -> list[str]:
|
||||
|
||||
@function_tool(timeout=30)
|
||||
async def get_threat_model(ctx: RunContextWrapper, target: str) -> str:
|
||||
"""Read this scan's threat model for a target, if an agent has derived one.
|
||||
"""Read the cached threat model for a target, if one exists.
|
||||
|
||||
The threat model is this run's shared answer to who the attacker
|
||||
is, where the trust boundaries sit, and what counts as critical
|
||||
here. Call it before you start hunting so you inherit the shared
|
||||
view instead of re-deriving it, and so every agent on this run
|
||||
agrees on what "attacker-controlled" means.
|
||||
|
||||
It is scoped to this scan and nothing is carried over from an
|
||||
earlier run, so an empty result means no agent has derived one yet.
|
||||
A threat model belongs to the target, not to this scan — the same
|
||||
trust boundaries hold across unrelated runs against the same host
|
||||
or application. Call this before you start hunting so you inherit
|
||||
the shared view instead of re-deriving it, and so every agent on
|
||||
this run agrees on what "attacker-controlled" means here.
|
||||
|
||||
Works black-box or white-box. The target can be a host, a URL, an
|
||||
API base, or a repository path; equivalent spellings of the same
|
||||
host resolve to the same model, and a checkout resolves to its
|
||||
remote, so a model derived white-box by one agent is read back by
|
||||
another testing the deployment.
|
||||
remote, so a model derived white-box is read back by a black-box
|
||||
agent testing the deployment.
|
||||
|
||||
Returns ``found: false`` when nothing has been derived yet — derive
|
||||
one and share it with ``save_threat_model``.
|
||||
Returns ``found: false`` when nothing is cached — derive one and
|
||||
persist it with ``save_threat_model``. ``stale: true`` means the
|
||||
checkout moved to a different revision, or that a model with no
|
||||
revision to pin to has aged out: use it as a starting point,
|
||||
re-confirm what it claims, and save the corrected version.
|
||||
|
||||
Any ``amendments`` in the response are corrections other agents
|
||||
recorded after the base model was written. They are part of the
|
||||
@@ -519,10 +540,10 @@ async def get_threat_model(ctx: RunContextWrapper, target: str) -> str:
|
||||
|
||||
@function_tool(timeout=30)
|
||||
async def save_threat_model(ctx: RunContextWrapper, target: str, content: str) -> str:
|
||||
"""Share a target-scoped threat model with the other agents on this scan.
|
||||
"""Persist a target-scoped threat model for reuse by other agents.
|
||||
|
||||
The model lives for this run only — it is not written to disk and a
|
||||
later scan of the same host or tree starts without it.
|
||||
Keyed by target identity, so a later scan of the same host or tree
|
||||
reads it back instead of paying to derive it again.
|
||||
|
||||
**This replaces the whole document, and clears any amendments** —
|
||||
it is for the agent establishing the baseline (normally root,
|
||||
@@ -540,9 +561,9 @@ async def save_threat_model(ctx: RunContextWrapper, target: str, content: str) -
|
||||
necessarily provisional — say which parts are inferred rather than
|
||||
observed, and let later agents amend it as the picture fills in.
|
||||
|
||||
**Scope it to the target, not to your slice of it.** Do not centre
|
||||
it on the diff you were handed, the subsystem you were assigned, or
|
||||
the one host that happened to answer first. With source, distinguish
|
||||
**Scope it to the target, not to this scan.** Do not centre it on
|
||||
the diff you were handed, the subsystem you were assigned, or the
|
||||
one host that happened to answer first. With source, distinguish
|
||||
real product and runtime surfaces from test, docs, example, and
|
||||
developer-tooling paths — in a monorepo, do not let ``tests/`` or
|
||||
one-off scripts become the centre of gravity unless the code shows
|
||||
|
||||
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
from strix.config import loader
|
||||
from strix.config.settings import DedupeSettings
|
||||
from strix.report.dedupe import _dedupe_model_settings, resolve_dedupe_model
|
||||
from strix.report.dedupe import _dedupe_model_settings
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -16,49 +16,32 @@ if TYPE_CHECKING:
|
||||
import pytest
|
||||
|
||||
|
||||
def _unwrap(model: object) -> object:
|
||||
while hasattr(model, "_inner"):
|
||||
model = model._inner
|
||||
return model
|
||||
|
||||
|
||||
def test_dedupe_key_bound_to_model_client_not_global_env() -> None:
|
||||
def test_dedupe_key_sent_per_call_not_via_global_env() -> None:
|
||||
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap", DEDUPE_LLM_API_KEY="dedupe-key")
|
||||
model = _unwrap(resolve_dedupe_model(dedupe, "deepseek/cheap"))
|
||||
# The key is bound to the dedupe model's own client, so a shared-provider
|
||||
# main key can't clobber it (and vice versa) through the process globals —
|
||||
# and it never rides on the request, where every model implementation's own
|
||||
# api_key kwarg would collide with it.
|
||||
assert model.api_key == "dedupe-key" # type: ignore[attr-defined]
|
||||
settings = _dedupe_model_settings(dedupe, "deepseek/cheap", 300)
|
||||
# The key rides on the request, so a shared-provider main key can't clobber
|
||||
# it (and vice versa) through the global provider env var.
|
||||
assert (settings.extra_args or {})["api_key"] == "dedupe-key"
|
||||
|
||||
|
||||
def test_dedupe_settings_carry_no_request_credentials() -> None:
|
||||
dedupe = DedupeSettings(
|
||||
STRIX_DEDUPE_MODEL="deepseek/cheap",
|
||||
DEDUPE_LLM_API_KEY="dedupe-key",
|
||||
DEDUPE_LLM_API_BASE="https://dedupe.example/v1",
|
||||
)
|
||||
def test_dedupe_settings_omit_api_key_when_unset() -> None:
|
||||
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap")
|
||||
settings = _dedupe_model_settings(dedupe, "deepseek/cheap", 300)
|
||||
assert "api_key" not in (settings.extra_args or {})
|
||||
assert "api_base" not in (settings.extra_args or {})
|
||||
|
||||
|
||||
def test_dedupe_endpoint_bound_to_model_client() -> None:
|
||||
def test_dedupe_endpoint_sent_per_call() -> None:
|
||||
dedupe = DedupeSettings(
|
||||
STRIX_DEDUPE_MODEL="openai/cheap",
|
||||
DEDUPE_LLM_API_KEY="dedupe-key",
|
||||
DEDUPE_LLM_API_BASE="https://dedupe.example/v1",
|
||||
)
|
||||
model = _unwrap(resolve_dedupe_model(dedupe, "openai/cheap"))
|
||||
client = model._client # type: ignore[attr-defined]
|
||||
assert client.api_key == "dedupe-key"
|
||||
assert str(client.base_url).startswith("https://dedupe.example/v1")
|
||||
|
||||
|
||||
def test_dedupe_without_credentials_uses_default_provider() -> None:
|
||||
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap")
|
||||
model = _unwrap(resolve_dedupe_model(dedupe, "deepseek/cheap"))
|
||||
assert model.api_key is None # type: ignore[attr-defined]
|
||||
settings = _dedupe_model_settings(dedupe, "openai/cheap", 300)
|
||||
# A distinct dedupe endpoint rides on the request instead of the
|
||||
# process-wide base URL, so it can't clobber the main model's endpoint.
|
||||
assert (settings.extra_args or {})["api_base"] == "https://dedupe.example/v1"
|
||||
assert (settings.extra_args or {})["api_key"] == "dedupe-key"
|
||||
|
||||
|
||||
def test_dedicated_dedupe_model_uses_own_headers_not_main() -> None:
|
||||
|
||||
@@ -458,49 +458,6 @@ async def test_non_user_send_does_not_resume_budget_pause(tmp_path: Any) -> None
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_send_starts_fresh_resume_attempt_after_failure() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
await coordinator.park_waiting("child", wait_kind="stalled")
|
||||
await coordinator.record_recovery("child")
|
||||
await coordinator.record_idle_resume("child")
|
||||
await coordinator.set_status("child", "failed", error="provider rejected request")
|
||||
assert await coordinator.claim_parent_notice("child") is True
|
||||
|
||||
delivered = await coordinator.send("child", {"from": "user", "content": "try again"})
|
||||
|
||||
assert delivered is True
|
||||
assert coordinator.statuses["child"] == "waiting"
|
||||
assert coordinator.pending_counts["child"] == 1
|
||||
assert "child" not in coordinator.errors
|
||||
assert "child" not in coordinator.wait_kinds
|
||||
assert "child" not in coordinator.recovery_counts
|
||||
assert "child" not in coordinator.idle_resume_counts
|
||||
assert await coordinator.claim_parent_notice("child") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_user_send_preserves_failed_resume_state() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
await coordinator.park_waiting("child", wait_kind="stalled")
|
||||
await coordinator.record_recovery("child")
|
||||
await coordinator.record_idle_resume("child")
|
||||
await coordinator.set_status("child", "failed", error="provider rejected request")
|
||||
|
||||
delivered = await coordinator.send("child", {"from": "root", "content": "status"})
|
||||
|
||||
assert delivered is True
|
||||
assert coordinator.statuses["child"] == "failed"
|
||||
assert coordinator.errors["child"] == "provider rejected request"
|
||||
assert coordinator.wait_kinds["child"] == "stalled"
|
||||
assert coordinator.recovery_counts["child"] == 1
|
||||
assert coordinator.idle_resume_counts["child"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_stops_clears_pause_and_normalizes_statuses() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
|
||||
@@ -857,37 +857,6 @@ async def test_agent_state_sync_does_not_mask_root_failure_with_completed_report
|
||||
assert runtime.controller.error == "finalization failed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_state_sync_clears_root_failure_after_user_resume() -> None:
|
||||
runtime = GoTuiRuntime(args())
|
||||
await runtime.coordinator.register("root", "Strix", parent_id=None)
|
||||
await runtime.coordinator.set_status("root", "failed", error="provider rejected request")
|
||||
|
||||
await runtime._sync_agent_state()
|
||||
assert runtime.controller.scan_state == "failed"
|
||||
assert runtime.live_view.agents["root"]["error_message"] == "provider rejected request"
|
||||
|
||||
await runtime.coordinator.send("root", {"from": "user", "content": "try again"})
|
||||
await runtime._sync_agent_state()
|
||||
|
||||
assert runtime.controller.scan_state == "running"
|
||||
assert runtime.controller.error is None
|
||||
root = runtime.live_view.agents["root"]
|
||||
assert root["status"] == "waiting"
|
||||
assert "error_message" not in root
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_state_sync_does_not_reopen_stopped_scan_with_active_root() -> None:
|
||||
runtime = GoTuiRuntime(args())
|
||||
runtime.controller.scan_state = "stopped"
|
||||
await runtime.coordinator.register("root", "Strix", parent_id=None)
|
||||
|
||||
await runtime._sync_agent_state()
|
||||
|
||||
assert runtime.controller.scan_state == "stopped"
|
||||
|
||||
|
||||
def _direct_launch_args() -> argparse.Namespace:
|
||||
launch_args = args()
|
||||
launch_args.needs_setup = False
|
||||
|
||||
@@ -1,99 +0,0 @@
|
||||
"""The import warm-up thread must never leave the import system poisoned.
|
||||
|
||||
Field failure: the warm-up thread's ``strix.core.runner`` import and the main
|
||||
thread's ``strix.report`` import both walked the agents SDK graph, and the two
|
||||
held each other's import locks (report -> dedupe -> agents while runner ->
|
||||
hooks -> report.state). CPython's deadlock avoidance breaks such a cycle by
|
||||
failing one import, which strands finished submodules in ``sys.modules`` with
|
||||
their parent package gone — and the next import of one of those submodules
|
||||
crashes with "partially initialized module".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
|
||||
from strix.llm import warmup
|
||||
|
||||
|
||||
def _run(code: str) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run( # noqa: S603
|
||||
[sys.executable, "-c", textwrap.dedent(code)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
timeout=300,
|
||||
)
|
||||
|
||||
|
||||
def test_strix_report_does_not_import_the_agents_graph() -> None:
|
||||
result = _run(
|
||||
"""
|
||||
import sys
|
||||
|
||||
import strix.report
|
||||
|
||||
agents_modules = [m for m in sys.modules if m == "agents" or m.startswith("agents.")]
|
||||
assert not agents_modules, agents_modules
|
||||
assert "strix.report.dedupe" not in sys.modules
|
||||
"""
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_check_duplicate_resolves_lazily() -> None:
|
||||
result = _run(
|
||||
"""
|
||||
import strix.report
|
||||
from strix.report import check_duplicate
|
||||
from strix.report.dedupe import check_duplicate as direct
|
||||
|
||||
assert strix.report.check_duplicate is direct is check_duplicate
|
||||
"""
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_failed_warm_import_purges_orphaned_submodules() -> None:
|
||||
result = _run(
|
||||
"""
|
||||
import sys
|
||||
|
||||
from strix.llm.warmup import _warm
|
||||
|
||||
# A package whose import fails after a submodule already completed:
|
||||
# CPython removes the package but leaves the submodule stranded.
|
||||
import pathlib
|
||||
import tempfile
|
||||
|
||||
root = pathlib.Path(tempfile.mkdtemp())
|
||||
pkg = root / "stranded_pkg"
|
||||
pkg.mkdir()
|
||||
(pkg / "ok.py").write_text("VALUE = 1")
|
||||
(pkg / "__init__.py").write_text("from . import ok\\nraise RuntimeError('boom')")
|
||||
sys.path.insert(0, str(root))
|
||||
|
||||
_warm(("stranded_pkg",))
|
||||
|
||||
assert "stranded_pkg" not in sys.modules
|
||||
assert "stranded_pkg.ok" not in sys.modules, "orphan survived the purge"
|
||||
|
||||
# And the subtree imports cleanly afterwards up to the real error.
|
||||
try:
|
||||
import stranded_pkg # noqa: F401
|
||||
except RuntimeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("expected the package's own error")
|
||||
"""
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_purge_does_not_touch_preexisting_or_healthy_modules() -> None:
|
||||
before = frozenset(sys.modules) - {"strix.llm.warmup"}
|
||||
warmup._purge_orphaned_modules(before)
|
||||
assert "strix.llm.warmup" in sys.modules # parent chain intact -> kept
|
||||
assert "strix" in sys.modules
|
||||
@@ -90,16 +90,6 @@ def test_make_model_settings_enables_prompt_cache_for_non_bedrock_claude(model_n
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
["claude-sonnet-4-5", "openai/claude-sonnet-4-5", "any-llm/anthropic/claude-sonnet-4-5"],
|
||||
)
|
||||
def test_no_prompt_cache_for_claude_off_the_litellm_route(model_name: str) -> None:
|
||||
# These names are served by SDK clients that raise TypeError on LiteLLM-only
|
||||
# request kwargs — e.g. a gateway in front of Claude reached with a bare name.
|
||||
assert _cache_points(model_name) is None
|
||||
|
||||
|
||||
def test_tool_config_point_not_leaked_to_non_bedrock_claude() -> None:
|
||||
# LiteLLM only consumes tool_config on Bedrock; elsewhere it leaks onto the
|
||||
# wire and native Anthropic 400s.
|
||||
@@ -153,8 +143,7 @@ def test_max_reasoning_effort_sent_as_raw_body_field() -> None:
|
||||
"max", model_name="deepseek/deepseek-v4-flash", request_timeout=30
|
||||
)
|
||||
assert settings.reasoning is None
|
||||
assert settings.extra_args == {"timeout": 30}
|
||||
assert settings.extra_body == {"reasoning_effort": "max"}
|
||||
assert settings.extra_args == {"timeout": 30, "extra_body": {"reasoning_effort": "max"}}
|
||||
|
||||
|
||||
def test_conversation_tail_breakpoint_moves_with_appended_transcript() -> None:
|
||||
|
||||
@@ -8,7 +8,6 @@ two dispatch tools ``describe_mcp`` and ``call_mcp``.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any
|
||||
@@ -30,7 +29,6 @@ from strix.tools.mcp import (
|
||||
McpConnectionConfig,
|
||||
McpConnectionRequest,
|
||||
McpRegistry,
|
||||
SupervisedMcpSession,
|
||||
attach_mcp_requests,
|
||||
call_mcp,
|
||||
describe_mcp,
|
||||
@@ -40,11 +38,9 @@ from strix.tools.mcp import (
|
||||
resolve_mcp_call,
|
||||
)
|
||||
from strix.tools.mcp import client as mcp_client
|
||||
from strix.tools.mcp import session as mcp_session_mod
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@@ -175,13 +171,6 @@ def _ctx(registry: McpRegistry | None) -> ToolContext[dict[str, Any]]:
|
||||
)
|
||||
|
||||
|
||||
async def _aclose_all(connections: list[Any]) -> None:
|
||||
"""Close every supervised session a connect/attach test opened, so no
|
||||
supervising task leaks into the event loop's teardown."""
|
||||
for connection in connections:
|
||||
await connection.session.aclose()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_mcp_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Hide any MCP settings the developer has exported in their own shell."""
|
||||
@@ -309,8 +298,6 @@ async def test_connect_returns_sessions_without_registering_agent_tools(
|
||||
assert [(c.name, c.tool_count) for c in connections] == [("fs", 2), ("db", 1)]
|
||||
assert list(factory.registered_agent_tools()) == before
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_count_honors_the_allowlist(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@@ -321,8 +308,6 @@ async def test_tool_count_honors_the_allowlist(monkeypatch: pytest.MonkeyPatch)
|
||||
|
||||
assert connections[0].tool_count == 1
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_notes_ride_on_the_connection(
|
||||
@@ -341,8 +326,6 @@ async def test_connection_notes_ride_on_the_connection(
|
||||
|
||||
assert connections[0].notes == "Staging analytics DB; read-only."
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
# --- server build branch -----------------------------------------------------
|
||||
|
||||
@@ -416,18 +399,11 @@ async def test_list_mcps_returns_connections_with_ids_and_descriptions() -> None
|
||||
out = await list_mcps.on_invoke_tool(_ctx(registry), "{}")
|
||||
|
||||
# ``id`` is the exact connection name describe_mcp/call_mcp accept;
|
||||
# ``description`` is the summary's purpose; ``dead`` is the connection's live
|
||||
# health (both healthy here); no tool schemas are included.
|
||||
# ``description`` is the summary's purpose; no tool schemas are included.
|
||||
assert out == {
|
||||
"connections": [
|
||||
{
|
||||
"id": "fs",
|
||||
"name": "fs",
|
||||
"description": "local files",
|
||||
"tool_count": 2,
|
||||
"dead": False,
|
||||
},
|
||||
{"id": "db", "name": "db", "description": None, "tool_count": 1, "dead": False},
|
||||
{"id": "fs", "name": "fs", "description": "local files", "tool_count": 2},
|
||||
{"id": "db", "name": "db", "description": None, "tool_count": 1},
|
||||
]
|
||||
}
|
||||
|
||||
@@ -825,82 +801,37 @@ def test_loader_exclude_selection_drops_named(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_skips_a_connection_whose_connect_is_cancelled(
|
||||
async def test_connect_cleans_up_when_cancelled_mid_connect(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Each connection now connects on its own supervising task. A cancellation of
|
||||
# one session's connect (the transport scope dying mid-connect) is contained
|
||||
# to that task: the connection is skipped and cleaned up, and the run's attach
|
||||
# keeps going rather than being cancelled.
|
||||
cleaned: list[str] = []
|
||||
|
||||
class _Tracking(FakeMCPServer):
|
||||
def __init__(self, name: str, *, cancel_connect: bool = False) -> None:
|
||||
def __init__(self, name: str, *, fail_connect: bool = False) -> None:
|
||||
super().__init__(name, [_mcp_tool("t")])
|
||||
self._cancel_connect = cancel_connect
|
||||
self._fail_connect = fail_connect
|
||||
|
||||
async def connect(self) -> None:
|
||||
if self._cancel_connect:
|
||||
if self._fail_connect:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
cleaned.append(self._name)
|
||||
|
||||
servers = {"good": _Tracking("good"), "bad": _Tracking("bad", cancel_connect=True)}
|
||||
servers = {"good": _Tracking("good"), "bad": _Tracking("bad", fail_connect=True)}
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda config: servers[config.name])
|
||||
|
||||
configs = [_config("good", ["t"]), _config("bad", ["t"])]
|
||||
configs = [
|
||||
McpConnectionConfig(name="good", url="https://mcp.example.com", allowed_tools=["t"]),
|
||||
McpConnectionConfig(name="bad", url="https://mcp.example.com", allowed_tools=["t"]),
|
||||
]
|
||||
|
||||
connections = await mcp_client.connect_mcp_servers(configs)
|
||||
|
||||
# The cancelled connect is skipped and cleaned up; the good one is returned.
|
||||
assert [c.name for c in connections] == ["good"]
|
||||
assert "bad" in cleaned
|
||||
|
||||
await _aclose_all(connections)
|
||||
assert "good" in cleaned
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_cleans_up_started_sessions_when_attach_is_cancelled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# If the attach coroutine itself is cancelled (the run going down) while a
|
||||
# later connection is still connecting, every session started so far is closed
|
||||
# on its own task before the cancellation is re-raised, so nothing is orphaned.
|
||||
cleaned: list[str] = []
|
||||
|
||||
class _Tracking(FakeMCPServer):
|
||||
def __init__(self, name: str, *, block_connect: bool = False) -> None:
|
||||
super().__init__(name, [_mcp_tool("t")])
|
||||
self._block_connect = block_connect
|
||||
|
||||
async def connect(self) -> None:
|
||||
if self._block_connect:
|
||||
await asyncio.Event().wait() # never completes
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
cleaned.append(self._name)
|
||||
|
||||
servers = {"good": _Tracking("good"), "slow": _Tracking("slow", block_connect=True)}
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda config: servers[config.name])
|
||||
|
||||
async def _attach() -> list[Any]:
|
||||
# Connect "good" first, then hang forever connecting "slow".
|
||||
return await mcp_client.connect_mcp_servers(
|
||||
[_config("good", ["t"]), _config("slow", ["t"])]
|
||||
)
|
||||
|
||||
task = asyncio.create_task(_attach())
|
||||
# Give the loop time to connect good and reach slow's hanging connect.
|
||||
for _ in range(100):
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
await mcp_client.connect_mcp_servers(configs)
|
||||
|
||||
# The already-connected "good" session was cleaned up, not orphaned.
|
||||
assert "good" in cleaned
|
||||
# The server being connected when cancelled, and the one already connected,
|
||||
# are both cleaned up rather than orphaned.
|
||||
assert cleaned == ["bad", "good"]
|
||||
|
||||
|
||||
# --- reading a tool call back to the server it went out to -------------------
|
||||
@@ -985,8 +916,6 @@ async def test_attach_populates_registry_with_provider_and_transform(
|
||||
assert entry.purpose == "Customer DB"
|
||||
assert entry.result_transform is transform
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attach_bare_request_matches_the_command_line_shape(
|
||||
@@ -1004,7 +933,7 @@ async def test_attach_bare_request_matches_the_command_line_shape(
|
||||
)
|
||||
|
||||
registry = McpRegistry()
|
||||
connections = await attach_mcp_requests([McpConnectionRequest(config=config)], registry)
|
||||
await attach_mcp_requests([McpConnectionRequest(config=config)], registry)
|
||||
|
||||
entry = registry.get("db")
|
||||
assert entry is not None
|
||||
@@ -1012,8 +941,6 @@ async def test_attach_bare_request_matches_the_command_line_shape(
|
||||
assert entry.result_transform is None
|
||||
assert entry.purpose == "Staging analytics DB; read-only."
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attach_is_fail_open_and_skips_a_failed_connection(
|
||||
@@ -1043,8 +970,6 @@ async def test_attach_is_fail_open_and_skips_a_failed_connection(
|
||||
assert registry.get("good") is not None
|
||||
assert registry.get("bad") is None
|
||||
|
||||
await _aclose_all(connections)
|
||||
|
||||
|
||||
# --- provider on the registry ------------------------------------------------
|
||||
|
||||
@@ -1167,435 +1092,3 @@ async def test_errored_structured_output_is_wrapped_with_success_false() -> None
|
||||
# Structured content serializes to a JSON string; it too is wrapped under
|
||||
# ``content`` so the failure flag has a top-level dict to ride on.
|
||||
assert out == {"success": False, "content": json.dumps({"error": "boom"})}
|
||||
|
||||
|
||||
# --- per-session isolation: containment, reconnect-retry, mark-dead ----------
|
||||
# Each MCP connection's live session is owned by its own supervising task. A
|
||||
# background failure in one session is contained to that task: the agent's call
|
||||
# comes back as a value, the run keeps going, and the session reconnects once and
|
||||
# retries the failed call once before it is marked unavailable.
|
||||
|
||||
|
||||
class _DyingHttpServer(FakeMCPServer):
|
||||
"""A connected server whose ``call_tool`` fails to model a session death.
|
||||
|
||||
``death`` is the exception raised on a call: a plain ``Exception`` models an
|
||||
HTTP/transport error, and ``asyncio.CancelledError`` models the streamable-HTTP
|
||||
transport's task group cancelling the supervising task from a background POST
|
||||
error (for example a provider 403). ``alive`` flips to stop dying, so a
|
||||
reconnected replacement can succeed.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
tools: list[MCPTool],
|
||||
*,
|
||||
death: BaseException,
|
||||
alive: bool = False,
|
||||
) -> None:
|
||||
super().__init__(name, tools)
|
||||
self._death = death
|
||||
self.alive = alive
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any] | None,
|
||||
meta: dict[str, Any] | None = None,
|
||||
) -> CallToolResult:
|
||||
if not self.alive:
|
||||
raise self._death
|
||||
return await super().call_tool(tool_name, arguments)
|
||||
|
||||
|
||||
def _secret_config(name: str) -> McpConnectionConfig:
|
||||
return McpConnectionConfig(
|
||||
name=name,
|
||||
url="https://mcp.example.com",
|
||||
auth=BearerAuth(token="super-secret-bearer-token-42"),
|
||||
allowed_tools=["read_file"],
|
||||
)
|
||||
|
||||
|
||||
async def _started_session(config: McpConnectionConfig) -> SupervisedMcpSession:
|
||||
session = SupervisedMcpSession(config)
|
||||
assert await session.start()
|
||||
return session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_reconnects_and_retries_after_a_session_death(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# The first session dies on its call; the supervisor rebuilds the connection
|
||||
# once (reusing the existing _build_server + connect), retries the one call
|
||||
# once, and the retry lands on the healthy replacement.
|
||||
first = _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403"))
|
||||
second = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
built = iter([first, second])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(built))
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
|
||||
# The caller gets the tool output as a value, and the retried call ran on the
|
||||
# reconnected server.
|
||||
assert out == {"type": "text", "text": "routed:read_file"}
|
||||
assert second.calls == [("read_file", {})]
|
||||
assert session.is_dead is False
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_marks_connection_dead_when_reconnect_keeps_failing(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# The session dies and the reconnect attempt also fails: the connection is
|
||||
# marked dead and the call returns the standard failed-tool output.
|
||||
first = _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403"))
|
||||
built = {"n": 0}
|
||||
|
||||
def _build(_config: McpConnectionConfig) -> MCPServer:
|
||||
built["n"] += 1
|
||||
if built["n"] == 1:
|
||||
return first
|
||||
raise ConnectionError("cannot reconnect")
|
||||
|
||||
monkeypatch.setattr(mcp_client, "_build_server", _build)
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
|
||||
# A dead connection surfaces as an ordinary failed tool call, not an exception.
|
||||
assert isinstance(out, dict)
|
||||
assert out["success"] is False
|
||||
assert "unavailable" in out["content"]
|
||||
assert session.is_dead is True
|
||||
|
||||
# A later call short-circuits to the same failed output without a new attempt.
|
||||
again = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
assert again["success"] is False
|
||||
|
||||
# describe_mcp reports the connection unavailable, and list_mcps still lists it.
|
||||
described = await describe_mcp.on_invoke_tool(_ctx(registry), json.dumps({"connection": "fs"}))
|
||||
assert "unavailable" in described
|
||||
listed = await list_mcps.on_invoke_tool(_ctx(registry), "{}")
|
||||
assert [c["id"] for c in listed["connections"]] == ["fs"]
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
async def _pump_until(predicate: Callable[[], bool], *, limit: int = 100) -> None:
|
||||
"""Yield to the event loop until ``predicate`` holds, so a background
|
||||
supervising task can advance its reconnect without a real timer."""
|
||||
for _ in range(limit):
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(0)
|
||||
raise AssertionError("condition not reached")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_idle_session_death_self_heals_on_reconnect(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A session that dies while idle (its supervising task cancelled between calls,
|
||||
# modeling the transport scope dying with no call in flight) reconnects once on
|
||||
# its own and keeps serving, rather than staying dead until a later call would
|
||||
# have triggered a reconnect.
|
||||
first = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
second = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
built = iter([first, second])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(built))
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
assert session._task is not None
|
||||
session._task.cancel() # idle transport death: no call in flight
|
||||
await _pump_until(lambda: session.server is second)
|
||||
assert session.is_dead is False
|
||||
|
||||
# The reconnected session serves calls normally.
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
assert out == {"type": "text", "text": "routed:read_file"}
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flapping_idle_session_is_marked_dead_without_looping(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# If a session reconnects after an idle death but dies again before serving any
|
||||
# call, the supervisor stops reconnecting and marks the connection dead, so a
|
||||
# server that instantly drops on connect cannot spin in a reconnect loop.
|
||||
first = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
second = FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
built = iter([first, second])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(built))
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
assert session._task is not None
|
||||
# First idle death heals onto the second server (only two builds ever happen).
|
||||
session._task.cancel()
|
||||
await _pump_until(lambda: session.server is second)
|
||||
assert session.is_dead is False
|
||||
|
||||
# Second idle death before any call is served: give up rather than reconnect.
|
||||
session._task.cancel()
|
||||
await _pump_until(lambda: session._task is not None and session._task.done())
|
||||
assert session.is_dead is True
|
||||
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
assert out["success"] is False
|
||||
assert "unavailable" in out["content"]
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
class _HangingCallServer(FakeMCPServer):
|
||||
"""A connected server whose ``call_tool`` never returns, modeling a hung
|
||||
in-flight call so teardown can be tested for boundedness."""
|
||||
|
||||
def __init__(self, name: str, tools: list[MCPTool]) -> None:
|
||||
super().__init__(name, tools)
|
||||
self.cleaned = False
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any] | None,
|
||||
meta: dict[str, Any] | None = None,
|
||||
) -> CallToolResult:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
self.cleaned = True
|
||||
|
||||
|
||||
class _HangingConnectServer(FakeMCPServer):
|
||||
"""A server whose ``connect`` never finishes, so ``start`` blocks on readiness
|
||||
and can be cancelled mid-connect."""
|
||||
|
||||
def __init__(self, name: str, tools: list[MCPTool]) -> None:
|
||||
super().__init__(name, tools)
|
||||
self.cleaned = False
|
||||
|
||||
async def connect(self) -> None:
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
self.cleaned = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_is_bounded_when_an_in_flight_call_hangs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A hung call must not queue the shutdown sentinel behind itself forever:
|
||||
# aclose falls back to cancelling the supervising task, and cleanup still runs.
|
||||
monkeypatch.setattr(mcp_session_mod, "_SHUTDOWN_TIMEOUT", 0.2)
|
||||
server = _HangingCallServer("fs", [_mcp_tool("read_file")])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server)
|
||||
|
||||
session = await _started_session(_secret_config("fs"))
|
||||
call = asyncio.create_task(session.dispatch("read_file", {}, label="fs_read_file"))
|
||||
await asyncio.sleep(0.05) # let the serve loop pick up the request and hang
|
||||
|
||||
# Must return promptly rather than block on the hung call.
|
||||
await asyncio.wait_for(session.aclose(), timeout=3.0)
|
||||
assert session._task is not None and session._task.done()
|
||||
assert server.cleaned is True
|
||||
|
||||
# The abandoned caller gets a value (dead), not a hang.
|
||||
out = await asyncio.wait_for(call, timeout=3.0)
|
||||
assert isinstance(out, dict) and out["success"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_cleans_up_when_connect_is_cancelled_mid_await(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# If the scan is cancelled while start() awaits readiness, the readiness future
|
||||
# is cancelled; aclose must not raise on it and must still cancel + clean up the
|
||||
# partially connected supervisor.
|
||||
server = _HangingConnectServer("fs", [_mcp_tool("read_file")])
|
||||
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server)
|
||||
|
||||
session = SupervisedMcpSession(_secret_config("fs"))
|
||||
start = asyncio.create_task(session.start())
|
||||
await asyncio.sleep(0.05) # let the supervisor reach the hanging connect()
|
||||
start.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await start
|
||||
|
||||
await asyncio.wait_for(session.aclose(), timeout=3.0)
|
||||
assert session._task is not None and session._task.done()
|
||||
assert server.cleaned is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_session_death_is_contained_and_other_connections_survive(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# A background failure that surfaces as a cancellation (the transport scope
|
||||
# dying) is contained to that one session: the caller gets a value, not a
|
||||
# raised CancelledError, and a second healthy connection keeps working.
|
||||
dying = _DyingHttpServer("dying", [_mcp_tool("read_file")], death=asyncio.CancelledError())
|
||||
healthy = FakeMCPServer("healthy", [_mcp_tool("read_file")])
|
||||
dying_builds = {"n": 0}
|
||||
|
||||
def _build(config: McpConnectionConfig) -> MCPServer:
|
||||
if config.name == "healthy":
|
||||
return healthy
|
||||
# The dying connection connects once, then its rebuild raises, so it ends
|
||||
# up marked dead rather than recovering.
|
||||
dying_builds["n"] += 1
|
||||
if dying_builds["n"] == 1:
|
||||
return dying
|
||||
raise ConnectionError("cannot reconnect")
|
||||
|
||||
monkeypatch.setattr(mcp_client, "_build_server", _build)
|
||||
|
||||
dying_session = await _started_session(_secret_config("dying"))
|
||||
healthy_session = await _started_session(_secret_config("healthy"))
|
||||
registry = McpRegistry()
|
||||
registry.add(name="dying", session=dying_session, tool_count=1)
|
||||
registry.add(name="healthy", session=healthy_session, tool_count=1)
|
||||
|
||||
dead_out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "dying", "tool": "read_file"})
|
||||
)
|
||||
# Contained: a value came back rather than a CancelledError tearing down the run.
|
||||
assert isinstance(dead_out, dict)
|
||||
assert dead_out["success"] is False
|
||||
|
||||
good_out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "healthy", "tool": "read_file"})
|
||||
)
|
||||
assert good_out == {"type": "text", "text": "routed:read_file"}
|
||||
|
||||
await dying_session.aclose()
|
||||
await healthy_session.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconnect_reuses_the_stored_config_and_never_logs_the_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
# The reconnect path rebuilds from the config held on the session, reusing the
|
||||
# same bearer token, and that token never reaches a log line, a repr, or the
|
||||
# inventory list_mcps emits.
|
||||
seen_tokens: list[str | None] = []
|
||||
|
||||
def _build(config: McpConnectionConfig) -> MCPServer:
|
||||
seen_tokens.append(config.auth.token if config.auth else None)
|
||||
if len(seen_tokens) == 1:
|
||||
return _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403"))
|
||||
return FakeMCPServer("fs", [_mcp_tool("read_file")])
|
||||
|
||||
monkeypatch.setattr(mcp_client, "_build_server", _build)
|
||||
|
||||
config = _secret_config("fs")
|
||||
token = config.auth.token if config.auth else ""
|
||||
session = await _started_session(config)
|
||||
registry = McpRegistry()
|
||||
entry = registry.add(name="fs", session=session, tool_count=1)
|
||||
|
||||
with caplog.at_level("DEBUG", logger="strix.tools.mcp.session"):
|
||||
out = await call_mcp.on_invoke_tool(
|
||||
_ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"})
|
||||
)
|
||||
|
||||
# The retry succeeded, and both the initial connect and the reconnect used the
|
||||
# same token from the stored config (never re-fetched).
|
||||
assert out == {"type": "text", "text": "routed:read_file"}
|
||||
assert seen_tokens == [token, token]
|
||||
|
||||
# The token appears in no log line, no repr of the session or entry, and not in
|
||||
# the inventory the agent sees.
|
||||
assert token not in caplog.text
|
||||
assert token not in repr(session)
|
||||
assert token not in repr(entry)
|
||||
listed = await list_mcps.on_invoke_tool(_ctx(registry), "{}")
|
||||
assert token not in json.dumps(listed)
|
||||
# The config is still reachable in memory for the reconnect path.
|
||||
assert entry.config is config
|
||||
|
||||
await session.aclose()
|
||||
|
||||
|
||||
# --- connection status signal ------------------------------------------------
|
||||
|
||||
|
||||
class _RaisingMCPServer(FakeMCPServer):
|
||||
"""A connected server whose every call raises, so the session dies."""
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any] | None,
|
||||
meta: dict[str, Any] | None = None,
|
||||
) -> CallToolResult:
|
||||
raise RuntimeError("connection lost")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_on_dead_fires_once_on_the_death_transition() -> None:
|
||||
# An adopted session with no config cannot reconnect, so the first failed
|
||||
# call marks it dead; the on-dead callback fires exactly once, on the edge.
|
||||
server = _RaisingMCPServer("db", [_mcp_tool("read")])
|
||||
session = SupervisedMcpSession.adopt(server, name="db")
|
||||
fires: list[int] = []
|
||||
session.set_on_dead(lambda: fires.append(1))
|
||||
|
||||
out = await session.dispatch("read", {}, label="db_read")
|
||||
|
||||
assert session.is_dead is True
|
||||
assert isinstance(out, dict) and out.get("success") is False
|
||||
assert fires == [1]
|
||||
|
||||
# A later call to the already-dead session must not fire the callback again.
|
||||
await session.dispatch("read", {}, label="db_read")
|
||||
assert fires == [1]
|
||||
|
||||
|
||||
def test_registry_statuses_report_the_live_dead_flag_and_provider() -> None:
|
||||
registry = McpRegistry()
|
||||
alive = SupervisedMcpSession.adopt(FakeMCPServer("a", []), name="a")
|
||||
gone = SupervisedMcpSession.adopt(FakeMCPServer("b", []), name="b")
|
||||
registry.add(name="a", session=alive, tool_count=2, provider="supabase")
|
||||
registry.add(name="b", session=gone, tool_count=1, provider=None)
|
||||
gone._mark_dead()
|
||||
|
||||
statuses = {status.name: status for status in registry.statuses()}
|
||||
assert statuses["a"].dead is False
|
||||
assert statuses["a"].tool_count == 2
|
||||
assert statuses["a"].provider == "supabase"
|
||||
assert statuses["b"].dead is True
|
||||
assert statuses["b"].provider is None
|
||||
|
||||
@@ -3,17 +3,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
from agents.model_settings import ModelSettings
|
||||
|
||||
from strix.config.models import (
|
||||
RECOMMENDED_MODEL_NAMES,
|
||||
StrixProvider,
|
||||
_NonStreamingModel,
|
||||
_TurnGuardModel,
|
||||
is_recommended_or_frontier_model,
|
||||
request_timeout_extra_args,
|
||||
routes_through_litellm,
|
||||
supports_strict_tool_schemas,
|
||||
)
|
||||
|
||||
@@ -117,38 +112,3 @@ def test_claude_routes_reject_strict_tool_schemas(model_name: str) -> None:
|
||||
)
|
||||
def test_other_routes_keep_strict_tool_schemas(model_name: str) -> None:
|
||||
assert supports_strict_tool_schemas(model_name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "litellm"),
|
||||
[
|
||||
("claude-sonnet-4-5", False),
|
||||
("openai/claude-sonnet-4-5", False),
|
||||
("any-llm/anthropic/claude-sonnet-4-5", False),
|
||||
("anthropic/claude-sonnet-4-5", True),
|
||||
("litellm/anthropic/claude-sonnet-4-5", True),
|
||||
("bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", True),
|
||||
("ollama/llama3", True),
|
||||
],
|
||||
)
|
||||
def test_routes_through_litellm_matches_the_provider(
|
||||
monkeypatch: pytest.MonkeyPatch, model_name: str, litellm: bool
|
||||
) -> None:
|
||||
"""The helper must agree with what StrixProvider actually builds.
|
||||
|
||||
Callers use it to decide whether a LiteLLM-only request field is safe to
|
||||
attach; on the SDK's own clients such a field raises TypeError mid-turn, so
|
||||
drift here breaks every request on that route.
|
||||
"""
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
||||
assert routes_through_litellm(model_name) is litellm
|
||||
try:
|
||||
model = StrixProvider().get_model(model_name)
|
||||
except ImportError:
|
||||
# any-llm's client is an optional dependency; reaching it at all already
|
||||
# proves the route is not LiteLLM's.
|
||||
assert not litellm
|
||||
return
|
||||
while isinstance(model, _NonStreamingModel | _TurnGuardModel):
|
||||
model = model._inner
|
||||
assert isinstance(model, LitellmModel) is litellm
|
||||
|
||||
@@ -63,34 +63,6 @@ def report_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> ReportState
|
||||
return state
|
||||
|
||||
|
||||
def test_record_mcp_connection_status_persists_and_dedupes(
|
||||
report_state: ReportState, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The roster lands on the run record so run.json carries it for the viewer,
|
||||
and an unchanged re-write is a no-op (it does not re-save)."""
|
||||
roster = [{"name": "local_fs", "provider": None, "tool_count": 3, "dead": False}]
|
||||
report_state.record_mcp_connection_status(roster)
|
||||
assert report_state.run_record["mcp_connection_status"] == roster
|
||||
|
||||
saves = 0
|
||||
original_save = report_state.save_run_data
|
||||
|
||||
def _counting_save(*args: Any, **kwargs: Any) -> None:
|
||||
nonlocal saves
|
||||
saves += 1
|
||||
original_save(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(report_state, "save_run_data", _counting_save)
|
||||
report_state.record_mcp_connection_status(roster)
|
||||
assert saves == 0, "an identical roster must not trigger another save"
|
||||
|
||||
report_state.record_mcp_connection_status(
|
||||
[{"name": "local_fs", "provider": None, "tool_count": 3, "dead": True}]
|
||||
)
|
||||
assert saves == 1
|
||||
assert report_state.run_record["mcp_connection_status"][0]["dead"] is True
|
||||
|
||||
|
||||
async def test_create_report_persists_new_fields(report_state: ReportState) -> None:
|
||||
result = await _do_create(
|
||||
title="Reflected XSS in search",
|
||||
|
||||
@@ -140,51 +140,3 @@ async def test_supplied_requests_are_attached_and_the_user_file_is_not_read(
|
||||
)
|
||||
|
||||
assert captured == [supplied]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_roster_is_persisted_even_without_a_status_sink(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Any
|
||||
) -> None:
|
||||
"""The viewer reads the roster off disk, so persistence must not depend on the
|
||||
interface status sink: with ``mcp_status_sink=None`` the connect-time roster is
|
||||
still written, carrying only the non-secret name/provider/tool_count/dead."""
|
||||
_wire_runner(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(
|
||||
mcp_pkg,
|
||||
"load_user_mcp_configs",
|
||||
lambda: [McpConnectionConfig(name="local_fs", transport="stdio", command="npx")],
|
||||
)
|
||||
|
||||
class _FakeSession:
|
||||
is_dead = False
|
||||
|
||||
def set_on_dead(self, _callback: Any) -> None:
|
||||
return None
|
||||
|
||||
async def _attach(_requests: list[McpConnectionRequest], registry: Any) -> list[Any]:
|
||||
registry.add(name="local_fs", session=_FakeSession(), tool_count=3, provider=None)
|
||||
entry = registry.get("local_fs")
|
||||
return [types.SimpleNamespace(name="local_fs", tool_count=3, session=entry.session)]
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _attach)
|
||||
|
||||
persisted: list[list[dict[str, Any]]] = []
|
||||
|
||||
def _capture_persist(roster: list[dict[str, Any]]) -> None:
|
||||
persisted.append(roster)
|
||||
|
||||
monkeypatch.setattr(runner, "_persist_mcp_status", _capture_persist)
|
||||
|
||||
await runner.run_strix_scan(
|
||||
scan_config={"targets": [], "scan_mode": "deep"},
|
||||
scan_id="scan-persist",
|
||||
image="img",
|
||||
coordinator=AgentCoordinator(),
|
||||
mcp_status_sink=None,
|
||||
)
|
||||
|
||||
assert persisted, "roster must persist even when no status sink is attached"
|
||||
assert persisted[-1] == [
|
||||
{"name": "local_fs", "provider": None, "tool_count": 3, "dead": False}
|
||||
]
|
||||
|
||||
@@ -194,13 +194,9 @@ async def test_mcp_available_flag_set_when_a_connection_attaches(
|
||||
scope_context: dict[str, Any] = {"scope": "built-in"}
|
||||
captured = _patch_engine_scaffold(monkeypatch, tmp_path, scope_context)
|
||||
|
||||
async def _aclose() -> None:
|
||||
return None
|
||||
|
||||
async def _attach(_requests: Any, registry: Any) -> list[Any]:
|
||||
registry.add(name="fs", server=object(), purpose="local files", tool_count=2)
|
||||
session = types.SimpleNamespace(aclose=_aclose)
|
||||
return [types.SimpleNamespace(name="fs", tool_count=2, session=session)]
|
||||
return [types.SimpleNamespace(name="fs", tool_count=2, server=object())]
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _attach)
|
||||
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
"""Tests for the run-scoped threat model store."""
|
||||
"""Tests for the target-scoped threat model cache."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
@@ -15,7 +17,6 @@ from strix.tools.threat_model.tools import (
|
||||
_save_impl,
|
||||
amend_threat_model,
|
||||
get_threat_model,
|
||||
hydrate_threat_models_from_disk,
|
||||
save_threat_model,
|
||||
)
|
||||
|
||||
@@ -63,10 +64,8 @@ def _make_repo(tmp_path: Path, name: str = "repo") -> Path:
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _empty_store() -> None:
|
||||
"""Each test is its own run, so it starts with an empty, unmirrored store."""
|
||||
threat_model_tools._MODELS.clear()
|
||||
threat_model_tools._store_path = None
|
||||
def _isolated_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(threat_model_tools, "_CACHE_DIR", tmp_path / "cache")
|
||||
|
||||
|
||||
def test_missing_model_reports_not_found(tmp_path: Path) -> None:
|
||||
@@ -86,53 +85,11 @@ def test_saved_model_round_trips(tmp_path: Path) -> None:
|
||||
result = _get_impl(str(repo))
|
||||
|
||||
assert result["found"] is True
|
||||
assert result["stale"] is False
|
||||
assert "multi-tenant billing API" in result["content"]
|
||||
|
||||
|
||||
def test_nothing_is_written_outside_the_run(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The model must not outlive the scan, so nothing may land in the home dir."""
|
||||
home = tmp_path / "home"
|
||||
home.mkdir()
|
||||
monkeypatch.setenv("HOME", str(home))
|
||||
repo = _make_repo(tmp_path)
|
||||
|
||||
_save_impl(str(repo), _MODEL, "root")
|
||||
_amend_impl(str(repo), _ADDENDUM, "agent-a")
|
||||
|
||||
assert list(home.rglob("*")) == []
|
||||
|
||||
|
||||
def test_a_new_run_starts_without_the_model(tmp_path: Path) -> None:
|
||||
"""A later scan of the same target inherits nothing from this one."""
|
||||
repo = _make_repo(tmp_path)
|
||||
hydrate_threat_models_from_disk(tmp_path / "first-run")
|
||||
_save_impl(str(repo), _MODEL, "root")
|
||||
|
||||
hydrate_threat_models_from_disk(tmp_path / "second-run") # a different scan
|
||||
|
||||
assert _get_impl(str(repo))["found"] is False
|
||||
|
||||
|
||||
def test_resuming_the_same_run_keeps_the_model(tmp_path: Path) -> None:
|
||||
"""A resumed scan is the same scan, so its agents keep the shared baseline."""
|
||||
state_dir = tmp_path / "state"
|
||||
repo = _make_repo(tmp_path)
|
||||
hydrate_threat_models_from_disk(state_dir)
|
||||
_save_impl(str(repo), _MODEL, "root")
|
||||
_amend_impl(str(repo), _ADDENDUM, "agent-a")
|
||||
|
||||
threat_model_tools._MODELS.clear() # what the resuming process starts from
|
||||
hydrate_threat_models_from_disk(state_dir)
|
||||
|
||||
result = _get_impl(str(repo))
|
||||
assert result["found"] is True
|
||||
assert [a["content"] for a in result["amendments"]] == [_ADDENDUM]
|
||||
|
||||
|
||||
def test_model_survives_a_new_revision_within_the_run(tmp_path: Path) -> None:
|
||||
"""The model is not pinned to a revision; a commit mid-run does not drop it."""
|
||||
def test_model_is_stale_after_new_revision(tmp_path: Path) -> None:
|
||||
repo = _make_repo(tmp_path)
|
||||
_save_impl(str(repo), _MODEL, None)
|
||||
|
||||
@@ -143,10 +100,11 @@ def test_model_survives_a_new_revision_within_the_run(tmp_path: Path) -> None:
|
||||
result = _get_impl(str(repo))
|
||||
|
||||
assert result["found"] is True
|
||||
assert "multi-tenant billing API" in result["content"]
|
||||
assert result["stale"] is True
|
||||
assert result["content"]
|
||||
|
||||
|
||||
def test_store_is_keyed_per_repository(tmp_path: Path) -> None:
|
||||
def test_cache_is_keyed_per_repository(tmp_path: Path) -> None:
|
||||
first = _make_repo(tmp_path, "first")
|
||||
second = _make_repo(tmp_path, "second")
|
||||
_save_impl(str(first), _MODEL, None)
|
||||
@@ -260,6 +218,8 @@ def test_blackbox_target_round_trips() -> None:
|
||||
result = _get_impl(target)
|
||||
|
||||
assert result["found"] is True
|
||||
assert result["stale"] is False, "a fresh model with no revision is not stale"
|
||||
assert result["revision"] == "unversioned"
|
||||
assert "Inferred from recon" in result["content"]
|
||||
|
||||
|
||||
@@ -272,6 +232,22 @@ def test_blackbox_target_spellings_share_one_model() -> None:
|
||||
assert _get_impl("https://other.example.com")["found"] is False
|
||||
|
||||
|
||||
def test_blackbox_model_goes_stale_with_age() -> None:
|
||||
target = "https://app.example.com"
|
||||
_save_impl(target, _BLACKBOX_MODEL, "recon")
|
||||
|
||||
aged = (datetime.now(UTC) - timedelta(days=threat_model_tools._MAX_AGE_DAYS + 1)).isoformat()
|
||||
path = threat_model_tools._cache_path("app.example.com:443")
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
payload["created_at"] = aged
|
||||
path.write_text(json.dumps(payload), encoding="utf-8")
|
||||
|
||||
result = _get_impl(target)
|
||||
|
||||
assert result["stale"] is True
|
||||
assert "re-confirm" in result["message"]
|
||||
|
||||
|
||||
def test_blackbox_target_can_be_amended() -> None:
|
||||
target = "https://app.example.com"
|
||||
_save_impl(target, _BLACKBOX_MODEL, "recon")
|
||||
|
||||
@@ -12,16 +12,6 @@ from strix.config.settings import DEFAULT_MAX_TURNS
|
||||
from strix.interface.tui.backend.controller import TuiController
|
||||
|
||||
|
||||
class _SendingCoordinator:
|
||||
def __init__(self, delivered: bool = True) -> None:
|
||||
self.delivered = delivered
|
||||
self.messages: list[tuple[str, dict[str, object]]] = []
|
||||
|
||||
async def send(self, agent_id: str, message: dict[str, object]) -> bool:
|
||||
self.messages.append((agent_id, message))
|
||||
return self.delivered
|
||||
|
||||
|
||||
def args() -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
needs_setup=True,
|
||||
@@ -71,25 +61,6 @@ async def test_setup_state_is_serializable() -> None:
|
||||
assert snapshot["diff_base"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connections_snapshot_reflects_the_pushed_mcp_roster() -> None:
|
||||
controller = TuiController(args())
|
||||
# A run with no MCP connections carries an empty roster, so the sidebar
|
||||
# omits the panel entirely.
|
||||
assert controller.snapshot()["connections"] == []
|
||||
|
||||
controller.set_mcp_connections(
|
||||
[
|
||||
{"name": "supabase", "tool_count": 3, "dead": False},
|
||||
{"name": "vercel", "tool_count": 1, "dead": True},
|
||||
]
|
||||
)
|
||||
assert controller.snapshot()["connections"] == [
|
||||
{"name": "supabase", "tool_count": 3, "dead": False},
|
||||
{"name": "vercel", "tool_count": 1, "dead": True},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_setup_instruction_starts_from_cli_and_can_be_cleared() -> None:
|
||||
setup_args = args()
|
||||
@@ -323,34 +294,6 @@ def test_snapshot_exposes_working_directory() -> None:
|
||||
assert controller.snapshot()["pending_mount"] == ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_message_updates_live_agent_projection_immediately() -> None:
|
||||
coordinator = _SendingCoordinator()
|
||||
controller = TuiController(args(), coordinator=coordinator)
|
||||
controller.setup_mode = False
|
||||
controller.scan_started = True
|
||||
controller.scan_loop = asyncio.get_running_loop()
|
||||
controller.live_view.upsert_agent(
|
||||
"root",
|
||||
name="Strix",
|
||||
status="failed",
|
||||
error_message="provider rejected request",
|
||||
)
|
||||
|
||||
result = await controller.handle(
|
||||
"agent.send_message",
|
||||
{"agent_id": "root", "message": "try again"},
|
||||
)
|
||||
|
||||
assert result == {"sent": True}
|
||||
assert coordinator.messages == [
|
||||
("root", {"from": "user", "content": "try again", "type": "instruction"})
|
||||
]
|
||||
agent = controller.live_view.agents["root"]
|
||||
assert agent["status"] == "waiting"
|
||||
assert "error_message" not in agent
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_forwards_verify_flag_by_default() -> None:
|
||||
seen_verify: bool | None = None
|
||||
|
||||
@@ -243,15 +243,13 @@ def test_defensive_state_projection_preserves_usage_summary() -> None:
|
||||
),
|
||||
)
|
||||
state = controller.snapshot()
|
||||
state["pending_mount"] = "current-project"
|
||||
state["provider"] = None
|
||||
state["future_oversized_field"] = "x" * 100_000
|
||||
|
||||
snapshot = bounded_state_projection(state)
|
||||
|
||||
assert snapshot["projection_truncated"] is True
|
||||
assert snapshot["usage"] == {"total_tokens": 720_400, "cost": 20.0}
|
||||
assert snapshot["working_dir"] == state["working_dir"]
|
||||
assert snapshot["pending_mount"] == "current-project"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -104,29 +104,6 @@ def test_resume_hides_system_guidance_injected_as_user_turns(tmp_path: Path) ->
|
||||
] == ["starting", "continuing"]
|
||||
|
||||
|
||||
def test_resume_hydrates_saved_agent_errors(tmp_path: Path) -> None:
|
||||
run_dir = tmp_path / "run"
|
||||
state_dir = runtime_state_dir(run_dir)
|
||||
state_dir.mkdir(parents=True, exist_ok=True)
|
||||
(state_dir / "agents.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"statuses": {"root": "failed"},
|
||||
"names": {"root": "Strix"},
|
||||
"parent_of": {"root": None},
|
||||
"errors": {"root": "provider rejected request"},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
view = GoTuiLiveView()
|
||||
view.hydrate_from_run_dir(run_dir)
|
||||
|
||||
assert view.agents["root"]["status"] == "failed"
|
||||
assert view.agents["root"]["error_message"] == "provider rejected request"
|
||||
|
||||
|
||||
def test_resume_keeps_messages_the_user_actually_typed(tmp_path: Path) -> None:
|
||||
run_dir = tmp_path / "run"
|
||||
_write_run(
|
||||
|
||||
@@ -96,25 +96,6 @@ def test_read_run_summary_finished_flag(tmp_path: Path) -> None:
|
||||
assert read_run_summary(partial)["finished"] is False
|
||||
|
||||
|
||||
def test_read_run_summary_surfaces_mcp_connection_status(tmp_path: Path) -> None:
|
||||
"""The engine persists the non-secret MCP roster under mcp_connection_status;
|
||||
read_run_summary spreads the whole record, so /api/run carries it to the
|
||||
viewer verbatim."""
|
||||
run_dir = _make_run(tmp_path, "mcp", status="running", end_time=None)
|
||||
roster = [
|
||||
{"name": "local_fs", "provider": None, "tool_count": 3, "dead": False},
|
||||
{"name": "db", "provider": "supabase", "tool_count": 7, "dead": True},
|
||||
]
|
||||
record = {
|
||||
"run_name": "mcp",
|
||||
"status": "running",
|
||||
"end_time": None,
|
||||
"mcp_connection_status": roster,
|
||||
}
|
||||
(run_dir / "run.json").write_text(json.dumps(record), encoding="utf-8")
|
||||
assert read_run_summary(run_dir)["mcp_connection_status"] == roster
|
||||
|
||||
|
||||
def test_read_missing_artifacts_return_defaults(tmp_path: Path) -> None:
|
||||
run_dir = _make_run(tmp_path, "empty", status="running", end_time=None)
|
||||
assert read_vulnerabilities(run_dir) == []
|
||||
|
||||
Reference in New Issue
Block a user