mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
b260a4ee38
* make the working-directory prompt answer the mouse Its Confirm and Cancel were drawn as buttons and did nothing when clicked: the modal mouse handler had a case for every dialog except this one, so a click fell through and the scan sat waiting on an answer the user believed they had given. Only the keyboard could answer it. The prompt is docked in a corner rather than centered, so it also needs its own bounds; the centered ones every other dialog uses would have put the buttons in the wrong place. Those bounds now come from the same placement cornerOverlay draws with. Two returns that hand back the model alongside a call that mutates it are now sequenced explicitly. They work, but only because the compiler happens to evaluate the call first, and one of them is what puts the prompt back in the composer when the mount is declined. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * skip the mount instead of abandoning the scan Declining the working-directory prompt threw the whole launch away and dropped back to the start screen, which is a lot to lose for answering one question about one directory. The two answers are now about the directory alone: mount it, or run without it. The prompt is the whole of the input either way. The buttons say which is which - Mount and Skip rather than Confirm and Cancel - and the prompt says what skipping costs. A run with neither target nor directory is a real run, so two things follow it. It can be resumed: its instruction is what drives it, and that is in the run record. And it tells the agent plainly that it has neither, because an agent given no scope goes looking for the one it assumes it was meant to have. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
489 lines
20 KiB
Python
489 lines
20 KiB
Python
"""UI-independent state and command controller for interactive Strix clients."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import math
|
|
import webbrowser
|
|
from collections.abc import Awaitable, Callable
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from strix.config import load_settings
|
|
from strix.config.models import is_recommended_or_frontier_model
|
|
from strix.config.settings import DEFAULT_MAX_TURNS
|
|
from strix.interface.tui.backend.live_view import TuiLiveView
|
|
from strix.interface.tui.backend.projection import (
|
|
MAX_TERMINAL_EVENTS,
|
|
MAX_TERMINAL_VULNERABILITIES,
|
|
SCAN_MODES,
|
|
SCOPE_MODES,
|
|
bounded_state_projection,
|
|
collection_item_projection,
|
|
sanitize_terminal_text,
|
|
terminal_projection,
|
|
)
|
|
from strix.interface.utils import is_subscription_run
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
import argparse
|
|
|
|
from strix.report.state import ReportState
|
|
|
|
|
|
_STOPPABLE_AGENT_STATUSES = frozenset({"running", "waiting", "budget_paused"})
|
|
|
|
ChangeCallback = Callable[[], None]
|
|
StartCallback = Callable[[bool], Awaitable[None]]
|
|
QuitCallback = Callable[[], Awaitable[None]]
|
|
|
|
|
|
class TuiController:
|
|
"""Own setup state and expose serializable scan state to any TUI."""
|
|
|
|
def __init__(
|
|
self,
|
|
args: argparse.Namespace,
|
|
*,
|
|
live_view: TuiLiveView | None = None,
|
|
coordinator: Any = None,
|
|
report_state: ReportState | None = None,
|
|
on_start: StartCallback | None = None,
|
|
on_quit: QuitCallback | None = None,
|
|
on_change: ChangeCallback | None = None,
|
|
) -> None:
|
|
self.args = args
|
|
self.live_view = live_view or TuiLiveView()
|
|
self.coordinator = coordinator
|
|
self.report_state = report_state
|
|
self.scan_loop: asyncio.AbstractEventLoop | None = None
|
|
self.setup_mode = bool(args.needs_setup)
|
|
self.scan_started = not self.setup_mode
|
|
self._start_in_progress = False
|
|
self.scan_state = "setup" if self.setup_mode else "running"
|
|
self.targets = [
|
|
str(target["original"])
|
|
for target in args.targets_info
|
|
if isinstance(target, dict) and target.get("original")
|
|
]
|
|
instruction = args.instruction
|
|
self.instruction = instruction.strip() if isinstance(instruction, str) else ""
|
|
requested_scan_mode = str(args.scan_mode)
|
|
self.scan_mode = requested_scan_mode if requested_scan_mode in SCAN_MODES else "deep"
|
|
raw_budget = args.max_budget_usd
|
|
self.max_budget_usd = (
|
|
float(raw_budget)
|
|
if isinstance(raw_budget, int | float)
|
|
and not isinstance(raw_budget, bool)
|
|
and math.isfinite(float(raw_budget))
|
|
and raw_budget > 0
|
|
else None
|
|
)
|
|
raw_turns = args.max_turns
|
|
self.max_turns = (
|
|
raw_turns
|
|
if isinstance(raw_turns, int) and not isinstance(raw_turns, bool) and raw_turns > 0
|
|
else DEFAULT_MAX_TURNS
|
|
)
|
|
requested_scope = str(args.scope_mode)
|
|
self.scope_mode = requested_scope if requested_scope in SCOPE_MODES else "auto"
|
|
raw_diff_base = args.diff_base
|
|
self.diff_base = raw_diff_base.strip() if isinstance(raw_diff_base, str) else None
|
|
# Host directory mounted for the agent to work in when the scan has no
|
|
# target, set only once the user confirms it. It is a workspace, not a
|
|
# target: it carries no scan scope, and the instruction is the only
|
|
# source of truth for what to do.
|
|
self.workspace_mount: str | None = None
|
|
# A target-less launch enters the live view and asks there before
|
|
# anything is prepared; this holds the directory awaiting that answer.
|
|
self.pending_workspace_mount: str | None = None
|
|
self._pending_verify = True
|
|
self.messages: list[dict[str, str]] = []
|
|
self._next_message_id = 1
|
|
self.error: str | None = None
|
|
self.viewer_status = "idle"
|
|
self.viewer_url: str | None = None
|
|
self._viewer_httpd: Any = None
|
|
self._on_start = on_start
|
|
self._on_quit = on_quit
|
|
self._on_change = on_change
|
|
|
|
def set_change_callback(self, callback: ChangeCallback) -> None:
|
|
self._on_change = callback
|
|
|
|
def notify_changed(self) -> None:
|
|
if self._on_change is not None:
|
|
self._on_change()
|
|
|
|
def set_runtime(
|
|
self,
|
|
*,
|
|
report_state: ReportState | None = None,
|
|
scan_loop: asyncio.AbstractEventLoop | None = None,
|
|
) -> None:
|
|
if report_state is not None:
|
|
self.report_state = report_state
|
|
if scan_loop is not None:
|
|
self.scan_loop = scan_loop
|
|
|
|
def begin_preparation(self) -> None:
|
|
"""Mark a directly-launched run as preparing behind the live TUI."""
|
|
self.scan_state = "preparing"
|
|
self.notify_changed()
|
|
|
|
def fail_preparation(self, detail: str) -> None:
|
|
self.scan_state = "failed"
|
|
self.error = detail
|
|
self.notify_changed()
|
|
|
|
def add_message(self, text: str, level: str = "info") -> None:
|
|
self._append_message(text, level)
|
|
self.notify_changed()
|
|
|
|
def _append_message(self, text: str, level: str) -> None:
|
|
self.messages.append(
|
|
{
|
|
"id": f"message-{self._next_message_id}",
|
|
"text": sanitize_terminal_text(text),
|
|
"level": sanitize_terminal_text(level),
|
|
}
|
|
)
|
|
self._next_message_id += 1
|
|
self.messages = self.messages[-200:]
|
|
|
|
def snapshot(self) -> dict[str, Any]:
|
|
"""Return small mutable state; histories are streamed as collections."""
|
|
model = ""
|
|
with contextlib.suppress(Exception):
|
|
model = (load_settings().llm.model or "").strip()
|
|
usage: dict[str, Any] = {}
|
|
if self.report_state is not None:
|
|
usage = dict(self.report_state.get_total_llm_usage())
|
|
subscription = False
|
|
with contextlib.suppress(Exception):
|
|
subscription = is_subscription_run(self.report_state)
|
|
model_warning = ""
|
|
if model and not is_recommended_or_frontier_model(model):
|
|
model_warning = (
|
|
f"{model} is not a recommended frontier model; pentest quality could be degraded"
|
|
)
|
|
state = {
|
|
"setup_mode": self.setup_mode,
|
|
"scan_started": self.scan_started,
|
|
"scan_state": self.scan_state,
|
|
"targets": [
|
|
terminal_projection(target, max_string=128) for target in self.targets[:16]
|
|
],
|
|
"target_count": len(self.targets),
|
|
"working_dir": str(Path.cwd()),
|
|
"pending_mount": self.pending_workspace_mount or "",
|
|
"instruction": terminal_projection(self.instruction, max_string=2 * 1024),
|
|
"scan_mode": self.scan_mode,
|
|
"max_budget_usd": self.max_budget_usd,
|
|
"max_turns": self.max_turns,
|
|
"scope_mode": self.scope_mode,
|
|
"diff_base": terminal_projection(self.diff_base, max_string=256),
|
|
"model": terminal_projection(model, max_string=256),
|
|
"model_warning": terminal_projection(model_warning, max_string=512),
|
|
"caido_url": terminal_projection(
|
|
getattr(self.report_state, "caido_url", None), max_string=1024
|
|
),
|
|
"messages": [
|
|
{
|
|
"id": str(message.get("id", ""))[:64],
|
|
"text": terminal_projection(message.get("text", ""), max_string=256),
|
|
"level": str(message.get("level", "info"))[:32],
|
|
}
|
|
for message in self.messages[-10:]
|
|
],
|
|
"usage": terminal_projection(usage, max_string=256, max_items=20),
|
|
"subscription": subscription,
|
|
"viewer_status": self.viewer_status,
|
|
"viewer_url": terminal_projection(self.viewer_url, max_string=1024),
|
|
"error": terminal_projection(self.error, max_string=2 * 1024),
|
|
}
|
|
return bounded_state_projection(state)
|
|
|
|
def collection(self, name: str) -> list[dict[str, Any]]:
|
|
"""Return one bounded terminal projection with stable item identities."""
|
|
if name == "agents":
|
|
return [
|
|
{
|
|
key: terminal_projection(agent.get(key), max_string=256, max_items=5)
|
|
for key in (
|
|
"id",
|
|
"name",
|
|
"parent_id",
|
|
"status",
|
|
"error_message",
|
|
"created_at",
|
|
"updated_at",
|
|
)
|
|
if key in agent
|
|
}
|
|
for agent in self.live_view.agents.values()
|
|
]
|
|
if name == "events":
|
|
return [collection_item_projection(event) for event in self.live_view.events]
|
|
if name == "vulnerabilities":
|
|
reports = (
|
|
self.report_state.vulnerability_reports if self.report_state is not None else []
|
|
)[-MAX_TERMINAL_VULNERABILITIES:]
|
|
result: list[dict[str, Any]] = []
|
|
for index, report in enumerate(reports):
|
|
projected = collection_item_projection(report)
|
|
report_id = projected.get("id")
|
|
if not isinstance(report_id, str) or not report_id:
|
|
projected["id"] = f"vulnerability-{index}"
|
|
result.append(projected)
|
|
return result
|
|
raise ValueError(f"Unknown collection: {name}")
|
|
|
|
def collection_snapshot(self, name: str) -> tuple[int | None, list[dict[str, Any]]]:
|
|
"""Return a collection cursor and complete bounded projection."""
|
|
if name == "events":
|
|
cursor, events = self.live_view.event_snapshot(limit=MAX_TERMINAL_EVENTS)
|
|
return cursor, [collection_item_projection(event) for event in events]
|
|
return None, self.collection(name)
|
|
|
|
def collection_changes(
|
|
self,
|
|
name: str,
|
|
cursor: int,
|
|
) -> tuple[int, list[dict[str, Any]]]:
|
|
"""Return event upserts since a monotonic source cursor."""
|
|
if name != "events":
|
|
raise ValueError(f"Collection {name!r} does not expose incremental changes")
|
|
next_cursor, events = self.live_view.event_changes_since(cursor)
|
|
return next_cursor, [
|
|
collection_item_projection(event) for event in events[-MAX_TERMINAL_EVENTS:]
|
|
]
|
|
|
|
async def handle(self, command: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
handlers = {
|
|
"setup.add_target": self._add_target,
|
|
"setup.set_instruction": self._set_instruction,
|
|
"setup.start": self._start,
|
|
"setup.confirm_mount": self._confirm_mount,
|
|
"agent.send_message": self._send_message,
|
|
"agent.stop": self._stop_agent,
|
|
"viewer.open": self._open_viewer,
|
|
"app.quit": self._quit,
|
|
}
|
|
handler = handlers.get(command)
|
|
if handler is None:
|
|
raise ValueError(f"Unknown command: {command}")
|
|
result = await handler(payload)
|
|
self.notify_changed()
|
|
return result
|
|
|
|
async def _add_target(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
self._require_setup_mutable()
|
|
target = self._required_string(payload, "target")
|
|
if target not in self.targets:
|
|
self.targets.append(target)
|
|
return {"target": target, "total": len(self.targets)}
|
|
|
|
async def _set_instruction(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
self._require_setup_mutable()
|
|
instruction = payload.get("instruction", "")
|
|
if not isinstance(instruction, str):
|
|
raise TypeError("instruction must be a string")
|
|
self.instruction = instruction.strip()
|
|
return {"instruction": self.instruction}
|
|
|
|
async def _start(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
if self.scan_started or self._start_in_progress:
|
|
raise RuntimeError("Scan is already starting or running")
|
|
# A bare prompt launches optimistically, like a coding agent: it skips
|
|
# the network model preflight and surfaces any model error live. A named
|
|
# target keeps the preflight so a real scan does not commit blind.
|
|
verify = payload.get("verify", True)
|
|
if not isinstance(verify, bool):
|
|
raise TypeError("verify must be a boolean")
|
|
# Launching with no target mounts the working directory, so it requires
|
|
# the user's explicit confirmation rather than happening silently.
|
|
mount_working_dir = payload.get("mount_working_dir", False)
|
|
if not isinstance(mount_working_dir, bool):
|
|
raise TypeError("mount_working_dir must be a boolean")
|
|
model = (load_settings().llm.model or "").strip()
|
|
if not model:
|
|
raise ValueError("No model configured. Set STRIX_LLM first.")
|
|
if self._on_start is None:
|
|
raise RuntimeError("Scan start is unavailable")
|
|
if not self.targets:
|
|
if not mount_working_dir:
|
|
raise ValueError("No target set. Add a target first.")
|
|
# Mounting the working directory needs the user's confirmation, and
|
|
# that is asked in the live view. Enter it now and prepare nothing
|
|
# until the answer arrives, so declining leaves no run behind.
|
|
self.pending_workspace_mount = str(Path.cwd())
|
|
self._pending_verify = verify
|
|
self.setup_mode = False
|
|
self.scan_started = True
|
|
self.scan_state = "preparing"
|
|
return {"started": True}
|
|
await self._begin_scan(verify)
|
|
return {"started": True}
|
|
|
|
async def _begin_scan(self, verify: bool) -> None:
|
|
if self._on_start is None:
|
|
raise RuntimeError("Scan start is unavailable")
|
|
self._start_in_progress = True
|
|
try:
|
|
await self._on_start(verify)
|
|
finally:
|
|
self._start_in_progress = False
|
|
self.setup_mode = False
|
|
self.scan_started = True
|
|
self.scan_state = "running"
|
|
|
|
async def _confirm_mount(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Answer the pending working-directory mount asked for in the live view."""
|
|
mount = self.pending_workspace_mount
|
|
if mount is None:
|
|
raise RuntimeError("No mount confirmation is pending")
|
|
approved = payload.get("approved")
|
|
if not isinstance(approved, bool):
|
|
raise TypeError("approved must be a boolean")
|
|
self.pending_workspace_mount = None
|
|
# Declining skips the mount, it does not abandon the scan. The prompt is
|
|
# the whole of the input either way; the working directory is only an
|
|
# extra the agent may look at, so the run goes ahead without one.
|
|
self.workspace_mount = mount if approved else None
|
|
await self._begin_scan(self._pending_verify)
|
|
return {"approved": approved}
|
|
|
|
async def _send_message(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
agent_id = self._required_string(payload, "agent_id")
|
|
message = self._required_string(payload, "message")
|
|
if self.coordinator is None:
|
|
raise RuntimeError("Agent coordinator is unavailable")
|
|
if self.scan_loop is None or self.scan_loop.is_closed():
|
|
raise RuntimeError("Scan loop is not ready")
|
|
self.live_view.record_user_message(agent_id, message)
|
|
if self.scan_loop is asyncio.get_running_loop():
|
|
delivered = await self.coordinator.send(
|
|
agent_id,
|
|
{"from": "user", "content": message, "type": "instruction"},
|
|
)
|
|
else:
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
self.coordinator.send(
|
|
agent_id,
|
|
{"from": "user", "content": message, "type": "instruction"},
|
|
),
|
|
self.scan_loop,
|
|
)
|
|
delivered = await asyncio.wrap_future(future)
|
|
if not delivered:
|
|
raise RuntimeError("Message could not be delivered")
|
|
return {"sent": True}
|
|
|
|
async def _stop_agent(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
agent_id = self._required_string(payload, "agent_id")
|
|
agent = self.live_view.agents.get(agent_id)
|
|
if agent is None:
|
|
raise ValueError(f"Unknown agent: {agent_id}")
|
|
status = str(agent.get("status", ""))
|
|
if status not in _STOPPABLE_AGENT_STATUSES:
|
|
raise RuntimeError(f"Agent '{agent_id}' cannot be stopped while {status or 'unknown'}")
|
|
if self.coordinator is None or self.scan_loop is None or self.scan_loop.is_closed():
|
|
raise RuntimeError("Scan loop is not ready")
|
|
if self.scan_loop is asyncio.get_running_loop():
|
|
accepted = await self.coordinator.cancel_descendants_graceful(agent_id)
|
|
else:
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
self.coordinator.cancel_descendants_graceful(agent_id), self.scan_loop
|
|
)
|
|
accepted = await asyncio.wrap_future(future)
|
|
if not accepted:
|
|
raise RuntimeError(f"Agent '{agent_id}' is no longer active")
|
|
return {"stopped": True}
|
|
|
|
async def _open_viewer(self, _payload: dict[str, Any]) -> dict[str, Any]:
|
|
if self.viewer_url:
|
|
with contextlib.suppress(Exception):
|
|
webbrowser.open(self.viewer_url)
|
|
return {"status": "running", "url": self.viewer_url}
|
|
if self.report_state is None:
|
|
self.viewer_status = "failed"
|
|
return {"status": self.viewer_status, "error": "Scan output is not ready"}
|
|
try:
|
|
from strix.interface.tui.backend.messages import (
|
|
send_user_message_to_agent,
|
|
)
|
|
from strix.interface.viewer.server import (
|
|
authorized_url,
|
|
bundle_is_built,
|
|
serve,
|
|
)
|
|
|
|
if not bundle_is_built():
|
|
self.viewer_status = "unavailable"
|
|
return {"status": self.viewer_status, "error": "Viewer UI not built"}
|
|
|
|
def steer(agent_id: str, message: str) -> bool:
|
|
return send_user_message_to_agent(
|
|
coordinator=self.coordinator,
|
|
loop=self.scan_loop,
|
|
live_view=self.live_view,
|
|
target_agent_id=agent_id,
|
|
message=message,
|
|
notify_changed=self.notify_changed,
|
|
wait_for_delivery=True,
|
|
)
|
|
|
|
httpd, url, token = serve(
|
|
self.report_state.get_run_dir(),
|
|
open_browser=True,
|
|
steer_handler=steer,
|
|
)
|
|
self._viewer_httpd = httpd
|
|
self.viewer_url = authorized_url(url, token)
|
|
self.viewer_status = "running"
|
|
with contextlib.suppress(Exception):
|
|
from strix.telemetry import posthog
|
|
|
|
live = self.report_state.run_record.get("status") not in {
|
|
"completed",
|
|
"stopped",
|
|
"failed",
|
|
"interrupted",
|
|
}
|
|
posthog.viewer_opened(source="tui", live=live)
|
|
except Exception: # noqa: BLE001 - viewer startup failures must not crash the TUI
|
|
self.viewer_status = "failed"
|
|
return {"status": self.viewer_status, "error": "Viewer failed to start"}
|
|
else:
|
|
return {"status": self.viewer_status, "url": self.viewer_url}
|
|
|
|
def close_viewer(self) -> None:
|
|
httpd = self._viewer_httpd
|
|
if httpd is None:
|
|
return
|
|
self._viewer_httpd = None
|
|
with contextlib.suppress(Exception):
|
|
httpd.shutdown()
|
|
httpd.server_close()
|
|
|
|
async def _quit(self, _payload: dict[str, Any]) -> dict[str, Any]:
|
|
self.close_viewer()
|
|
if self._on_quit is not None:
|
|
await self._on_quit()
|
|
self.scan_state = "stopped"
|
|
return {"quitting": True}
|
|
|
|
@staticmethod
|
|
def _required_string(payload: dict[str, Any], name: str) -> str:
|
|
value = payload.get(name)
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise ValueError(f"{name} must be a non-empty string")
|
|
return value.strip()
|
|
|
|
def _require_setup_mutable(self) -> None:
|
|
if not self.setup_mode or self.scan_started or self._start_in_progress:
|
|
raise RuntimeError("Setup can no longer be changed after the scan starts")
|