refactor: consolidate run state layout

This commit is contained in:
0xallam
2026-04-26 15:01:35 -07:00
parent 49623ee625
commit 61a4c18a02
14 changed files with 223 additions and 133 deletions
+7 -2
View File
@@ -89,6 +89,10 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
"run_name": args.run_name,
"diff_scope": getattr(args, "diff_scope", {"active": False}),
"scan_mode": scan_mode,
"non_interactive": bool(getattr(args, "non_interactive", False)),
"local_sources": getattr(args, "local_sources", None) or [],
"scope_mode": getattr(args, "scope_mode", "auto"),
"diff_base": getattr(args, "diff_base", None),
# Forward the new --instruction (if any) to the resume path so it
# can deliver it as a fresh user message after SDK session replay.
# Empty string when the user didn't pass one on resume — no-op.
@@ -96,8 +100,9 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
}
report_state = ReportState(args.run_name)
report_state.set_scan_config(scan_config)
report_state.hydrate_from_run_dir()
report_state.set_scan_config(scan_config)
report_state.save_run_data()
def display_vulnerability(report: dict[str, Any]) -> None:
report_id = report.get("id", "unknown")
@@ -121,7 +126,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
report_state.cleanup()
def signal_handler(_signum: int, _frame: Any) -> None:
report_state.cleanup()
report_state.cleanup(status="interrupted")
sys.exit(1)
atexit.register(cleanup_on_exit)
+31 -35
View File
@@ -5,9 +5,9 @@ Strix Agent Interface
import argparse
import asyncio
import json
import shutil
import sys
from datetime import UTC, datetime
from pathlib import Path
from agents.model_settings import ModelSettings
@@ -24,6 +24,7 @@ from strix.config import (
persist_current,
)
from strix.config.models import configure_sdk_model_defaults, normalize_model_name
from strix.core.paths import run_dir_for, runtime_state_dir
from strix.interface.cli import run_cli
from strix.interface.tui import run_tui
from strix.interface.utils import (
@@ -42,6 +43,7 @@ from strix.interface.utils import (
validate_config_file,
)
from strix.report.state import get_global_report_state
from strix.report.writer import read_run_record, write_run_record
from strix.telemetry import posthog
from strix.telemetry.logging import configure_dependency_logging
@@ -433,7 +435,7 @@ Examples:
"the prior run left off, including the original target list."
)
_load_resume_state(args, parser)
agents_path = Path("strix_runs") / args.resume / "agents.json"
agents_path = runtime_state_dir(run_dir_for(args.resume)) / "agents.json"
if not agents_path.exists():
parser.error(
f"--resume {args.resume}: missing {agents_path}. The run was "
@@ -469,17 +471,16 @@ Examples:
return args
def _persist_scan_state(args: argparse.Namespace) -> None:
"""Dump the resolved scan inputs to ``{run_dir}/scan_state.json``.
Called once at the end of fresh-run setup. ``--resume <run_name>`` on
a future invocation reads this file to repopulate targets, scan_mode,
instruction, local_sources, and diff_scope without the user having to
retype them.
"""
run_dir = Path("strix_runs") / args.run_name
def _persist_run_record(args: argparse.Namespace) -> None:
"""Write the single public run descriptor used by resume and reporting."""
run_dir = run_dir_for(args.run_name)
run_dir.mkdir(parents=True, exist_ok=True)
state = {
run_record = {
"run_id": args.run_name,
"run_name": args.run_name,
"status": "running",
"start_time": datetime.now(UTC).isoformat(),
"end_time": None,
"targets_info": args.targets_info,
"scan_mode": args.scan_mode,
"instruction": args.instruction,
@@ -489,34 +490,26 @@ def _persist_scan_state(args: argparse.Namespace) -> None:
"scope_mode": args.scope_mode,
"diff_base": args.diff_base,
}
(run_dir / "scan_state.json").write_text(
json.dumps(state, ensure_ascii=False, indent=2, default=str),
encoding="utf-8",
)
write_run_record(run_dir, run_record)
def _load_resume_state(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
"""Populate ``args.targets_info`` and friends from a prior run's scan state.
Reads ``strix_runs/<run_name>/scan_state.json`` written at the end of the
fresh-run setup in ``main()``. Only fields the user did not explicitly
set on this invocation are restored — passing ``--instruction`` on
resume, for example, overrides the persisted instruction.
"""
state_path = Path("strix_runs") / args.resume / "scan_state.json"
"""Populate ``args.targets_info`` and friends from a prior run's run.json."""
run_dir = run_dir_for(args.resume)
state_path = run_dir / "run.json"
if not state_path.exists():
parser.error(
f"--resume {args.resume}: no such run "
f"(missing {state_path}; remove --resume for a fresh start)"
)
try:
state = json.loads(state_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
parser.error(f"--resume {args.resume}: scan_state.json unreadable: {exc}")
state = read_run_record(run_dir)
except RuntimeError as exc:
parser.error(f"--resume {args.resume}: run.json unreadable: {exc}")
args.targets_info = state.get("targets_info") or []
if not args.targets_info:
parser.error(f"--resume {args.resume}: scan_state.json has no targets_info")
parser.error(f"--resume {args.resume}: run.json has no targets_info")
# Validate any persisted ``cloned_repo_path`` still exists on disk.
# The resume path skips re-cloning, so a missing dir would mean the
@@ -540,8 +533,6 @@ def _load_resume_state(args: argparse.Namespace, parser: argparse.ArgumentParser
if args.instruction is None:
args.instruction = state.get("instruction")
if state.get("instruction_file") and args.instruction_file is None:
args.instruction_file = state.get("instruction_file")
if state.get("local_sources"):
args.local_sources = state.get("local_sources")
if state.get("diff_scope"):
@@ -561,8 +552,8 @@ def display_completion_message(args: argparse.Namespace, results_path: Path) ->
report_state = get_global_report_state()
scan_completed = False
if report_state and report_state.scan_results:
scan_completed = report_state.scan_results.get("scan_completed", False)
if report_state:
scan_completed = report_state.run_record.get("status") == "completed"
completion_text = Text()
if scan_completed:
@@ -739,10 +730,10 @@ def main() -> None:
else:
args.instruction = diff_scope.instruction_block
# Persist the fully-resolved scan state so a future
# Persist the fully-resolved run descriptor so a future
# ``--resume <run_name>`` invocation can pick up without
# re-supplying targets / instructions / scope.
_persist_scan_state(args)
_persist_run_record(args)
posthog.start(
model=load_settings().llm.model,
@@ -767,9 +758,14 @@ def main() -> None:
finally:
report_state = get_global_report_state()
if report_state:
status = {"interrupted": "interrupted", "error": "failed"}.get(
exit_reason,
"stopped",
)
report_state.cleanup(status=status)
posthog.end(report_state, exit_reason=exit_reason)
results_path = Path("strix_runs") / args.run_name
results_path = run_dir_for(args.run_name)
display_completion_message(args, results_path)
if args.non_interactive:
+11 -3
View File
@@ -708,8 +708,9 @@ class StrixTUIApp(App): # type: ignore[misc]
self.scan_config = self._build_scan_config(args)
self.report_state = ReportState(self.scan_config["run_name"])
self.report_state.set_scan_config(self.scan_config)
self.report_state.hydrate_from_run_dir()
self.report_state.set_scan_config(self.scan_config)
self.report_state.save_run_data()
set_global_report_state(self.report_state)
self.live_view = TuiLiveView()
self.live_view.hydrate_from_run_dir(self.report_state.get_run_dir())
@@ -759,6 +760,10 @@ class StrixTUIApp(App): # type: ignore[misc]
"run_name": args.run_name,
"diff_scope": getattr(args, "diff_scope", {"active": False}),
"scan_mode": getattr(args, "scan_mode", "deep"),
"non_interactive": bool(getattr(args, "non_interactive", False)),
"local_sources": getattr(args, "local_sources", None) or [],
"scope_mode": getattr(args, "scope_mode", "auto"),
"diff_base": getattr(args, "diff_base", None),
# Forward the new --instruction (if any) so the resume path
# can deliver it as a fresh user message after session replay.
"resume_instruction": getattr(args, "user_explicit_instruction", None) or "",
@@ -769,7 +774,7 @@ class StrixTUIApp(App): # type: ignore[misc]
self.report_state.cleanup()
def signal_handler(_signum: int, _frame: Any) -> None:
self.report_state.cleanup()
self.report_state.cleanup(status="interrupted")
sys.exit(0)
atexit.register(cleanup_on_exit)
@@ -1611,13 +1616,16 @@ class StrixTUIApp(App): # type: ignore[misc]
)
target_agent_id = self.selected_agent_id
send_user_message_to_agent(
submitted = send_user_message_to_agent(
coordinator=self.coordinator,
loop=self._scan_loop,
live_view=self.live_view,
target_agent_id=target_agent_id,
message=message,
)
if not submitted:
self.notify("Scan loop is not ready; message was not sent", severity="warning")
return
self._displayed_events.clear()
self._update_chat_view()
+3 -1
View File
@@ -8,6 +8,8 @@ import sqlite3
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
from strix.core.paths import runtime_state_dir
if TYPE_CHECKING:
from pathlib import Path
@@ -17,7 +19,7 @@ logger = logging.getLogger(__name__)
def load_session_history(run_dir: Path, agent_ids: Any) -> list[tuple[str, dict[str, Any], str]]:
agents_db = run_dir / "agents.db"
agents_db = runtime_state_dir(run_dir) / "agents.db"
session_ids = [aid for aid in agent_ids if isinstance(aid, str)]
if not agents_db.exists() or not session_ids:
return []
+3 -1
View File
@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from pathlib import Path
from strix.core.paths import runtime_state_dir
from strix.interface.tui.history import load_session_history
@@ -24,7 +25,8 @@ class TuiLiveView:
self._tool_event_by_call_id: dict[str, dict[str, Any]] = {}
def hydrate_from_run_dir(self, run_dir: Path) -> None:
agents_path = run_dir / "agents.json"
state_dir = runtime_state_dir(run_dir)
agents_path = state_dir / "agents.json"
if not agents_path.exists():
return
try:
+17 -3
View File
@@ -3,9 +3,13 @@
from __future__ import annotations
import asyncio
import logging
from typing import Any
logger = logging.getLogger(__name__)
def send_user_message_to_agent(
*,
coordinator: Any,
@@ -15,16 +19,26 @@ def send_user_message_to_agent(
message: str,
) -> bool:
"""Record a local user message and enqueue it into the target SDK session."""
live_view.record_user_message(target_agent_id, message)
if loop is None or loop.is_closed():
return False
asyncio.run_coroutine_threadsafe(
live_view.record_user_message(target_agent_id, message)
future = asyncio.run_coroutine_threadsafe(
coordinator.send(
target_agent_id,
{"from": "user", "content": message, "type": "instruction"},
),
loop,
)
future.add_done_callback(_log_delivery_failure)
return True
def _log_delivery_failure(future: Any) -> None:
try:
delivered = bool(future.result())
except Exception:
logger.exception("TUI user message delivery failed")
return
if not delivered:
logger.warning("TUI user message was not persisted to the SDK session")