mirror of
https://github.com/usestrix/strix.git
synced 2026-08-23 03:12:37 +02:00
feat(tui): replace Textual with a Go/Bubble Tea interface (#941)
This commit is contained in:
oyasumi
@@ -0,0 +1,497 @@
|
||||
"""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 enter_setup(self) -> None:
|
||||
"""Return a session to the start screen, e.g. on a declined mount."""
|
||||
self.setup_mode = True
|
||||
self.scan_started = False
|
||||
self.scan_state = "setup"
|
||||
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
|
||||
if not approved:
|
||||
# Nothing was prepared, so return to the start screen untouched.
|
||||
self.workspace_mount = None
|
||||
self.enter_setup()
|
||||
return {"approved": False}
|
||||
self.workspace_mount = mount
|
||||
await self._begin_scan(self._pending_verify)
|
||||
return {"approved": True}
|
||||
|
||||
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")
|
||||
Reference in New Issue
Block a user