"""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")