From 66166d5d68245a3f711ec73b8641126485b36aa8 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 19 Mar 2026 10:08:48 +0100 Subject: [PATCH] Fix resume: auto-detect checkpoint by target hash + save on Ctrl+C Root cause of resume not working: - generate_run_name() adds a random suffix every time, so without --run-name the checkpoint from a previous session was never found. - Ctrl+C during the first iteration (before any checkpoint was saved) left no checkpoint to resume from. Fixes: 1. _find_checkpoint_by_target_hash(): scans strix_runs/ for the most recent checkpoint whose target_hash matches the current targets. Now running `strix --target example.com` again automatically resumes the last interrupted scan without needing --run-name. 2. _save_checkpoint_on_interrupt(): saves current agent state in both the signal handler and atexit in cli.py and tui.py, so a Ctrl+C mid-first-iteration still produces a valid checkpoint. 3. _setup_checkpoint_on_args() restructured: handles run_name=None, --force-new, explicit --run-name, and auto-detect in one place. Co-Authored-By: Claude Sonnet 4.6 --- strix/interface/cli.py | 23 ++++++++ strix/interface/main.py | 125 +++++++++++++++++++++++++++++++--------- strix/interface/tui.py | 21 +++++++ 3 files changed, 141 insertions(+), 28 deletions(-) diff --git a/strix/interface/cli.py b/strix/interface/cli.py index ad7b0a34..edb73b3d 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -273,13 +273,35 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 tracer.vulnerability_found_callback = display_vulnerability + # Added for Resume Feature — mutable container so the nested closures can + # update the reference once the agent is created. + _agent_ref: list[Any] = [] + + def _save_checkpoint_on_interrupt() -> None: + """Persist current agent state before exit so the scan can be resumed.""" + if not checkpoint_manager or not _agent_ref: + return + try: + agent_instance = _agent_ref[0] + checkpoint_manager.save( + agent_instance.state, + tracer, + scan_config, + target_hash, + agent_instance.max_iterations, + ) + except Exception: # noqa: BLE001 + pass # non-fatal + def cleanup_on_exit() -> None: from strix.runtime import cleanup_runtime + _save_checkpoint_on_interrupt() tracer.cleanup() cleanup_runtime() def signal_handler(_signum: int, _frame: Any) -> None: + _save_checkpoint_on_interrupt() tracer.cleanup() sys.exit(1) @@ -329,6 +351,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 try: agent = StrixAgent(agent_config) + _agent_ref.append(agent) # expose to interrupt handler result = await agent.execute_scan(scan_config) if isinstance(result, dict) and not result.get("success", True): diff --git a/strix/interface/main.py b/strix/interface/main.py index 68e23a8a..9b5ad343 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -546,6 +546,48 @@ def persist_config() -> None: save_current_config() +def _find_checkpoint_by_target_hash( + strix_runs_dir: "Path", target_hash: str +) -> "tuple[str, Any] | None": + """Scan strix_runs/ for the most recent checkpoint matching *target_hash*. + + Returns ``(run_name, CheckpointModel)`` or ``None``. + Added for Resume Feature — enables auto-resume without --run-name. + """ + import json + + from strix.telemetry.checkpoint import CheckpointModel + + if not strix_runs_dir.exists(): + return None + + best_run_name: str | None = None + best_cp: Any = None + best_time: str = "" + + for run_dir in strix_runs_dir.iterdir(): + if not run_dir.is_dir(): + continue + cp_path = run_dir / "checkpoint.json" + if not cp_path.exists(): + continue + try: + raw = json.loads(cp_path.read_text(encoding="utf-8")) + if raw.get("target_hash") != target_hash: + continue + saved_at = raw.get("saved_at", "") + if best_run_name is None or saved_at > best_time: + best_cp = CheckpointModel.model_validate(raw) + best_run_name = run_dir.name + best_time = saved_at + except Exception: # noqa: BLE001 + continue + + if best_run_name and best_cp: + return (best_run_name, best_cp) + return None + + def _setup_checkpoint_on_args(args: argparse.Namespace) -> None: """Resolve checkpoint / resume state and attach it to ``args``. @@ -559,42 +601,71 @@ def _setup_checkpoint_on_args(args: argparse.Namespace) -> None: from strix.telemetry.checkpoint import CheckpointManager, compute_target_hash - run_dir = Path("strix_runs") / args.run_name - mgr = CheckpointManager(args.run_name, run_dir) target_hash = compute_target_hash(args.targets_info) - - args._checkpoint_manager = mgr args._target_hash = target_hash args._checkpoint_data = None args.resume_from_checkpoint = False if args.force_new: - mgr.delete() + # Delete any checkpoint for the current run name (if explicit) and start fresh + if args.run_name: + run_dir = Path("strix_runs") / args.run_name + CheckpointManager(args.run_name, run_dir).delete() + if not args.run_name: + args.run_name = generate_run_name(args.targets_info) + run_dir = Path("strix_runs") / args.run_name + args._checkpoint_manager = CheckpointManager(args.run_name, run_dir) return - if mgr.exists(): - checkpoint = mgr.load() - if checkpoint is None: - # Corrupted checkpoint — warn and start fresh + # If no explicit run name, try auto-detect by target hash + if not args.run_name: + found = _find_checkpoint_by_target_hash(Path("strix_runs"), target_hash) + if found: + run_name, checkpoint = found console = Console() console.print( - "[bold yellow]Warning:[/] Checkpoint file is corrupted or unreadable. " - "Starting a fresh scan." + f"[bold #22c55e]Auto-resuming previous scan[/] " + f"[dim](run: {run_name})[/] " + f"[dim]Use --new to start fresh.[/]" ) + args.run_name = run_name + run_dir = Path("strix_runs") / run_name + args._checkpoint_manager = CheckpointManager(run_name, run_dir) + args._checkpoint_data = checkpoint + args.resume_from_checkpoint = True return + # No checkpoint found — generate a fresh run name + args.run_name = generate_run_name(args.targets_info) - if checkpoint.target_hash != target_hash: - console = Console() - console.print( - "[bold yellow]Warning:[/] Checkpoint target mismatch " - f"(run '{args.run_name}' was for a different target). " - "Starting a fresh scan." - ) - return + # Explicit run name (or freshly generated) — look for its checkpoint + run_dir = Path("strix_runs") / args.run_name + mgr = CheckpointManager(args.run_name, run_dir) + args._checkpoint_manager = mgr - # Valid checkpoint found — auto-resume (or explicit --resume) - args._checkpoint_data = checkpoint - args.resume_from_checkpoint = True + if not mgr.exists(): + return + + checkpoint = mgr.load() + if checkpoint is None: + console = Console() + console.print( + "[bold yellow]Warning:[/] Checkpoint file is corrupted or unreadable. " + "Starting a fresh scan." + ) + return + + if checkpoint.target_hash != target_hash: + console = Console() + console.print( + "[bold yellow]Warning:[/] Checkpoint target mismatch " + f"(run '{args.run_name}' was for a different target). " + "Starting a fresh scan." + ) + return + + # Valid checkpoint found — resume + args._checkpoint_data = checkpoint + args.resume_from_checkpoint = True def main() -> None: @@ -614,12 +685,10 @@ def main() -> None: persist_config() - # Added for Resume Feature — determine run_name and whether to resume - if args.run_name_override: - args.run_name = args.run_name_override - else: - args.run_name = generate_run_name(args.targets_info) - + # Added for Resume Feature — determine run_name and whether to resume. + # _setup_checkpoint_on_args sets args.run_name (using override if given, + # auto-detecting by target hash, or generating a fresh name as fallback). + args.run_name = getattr(args, "run_name_override", None) or None _setup_checkpoint_on_args(args) for target_info in args.targets_info: diff --git a/strix/interface/tui.py b/strix/interface/tui.py index 4580feac..94a52daa 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -753,6 +753,7 @@ class StrixTUIApp(App): # type: ignore[misc] self._scan_thread: threading.Thread | None = None self._scan_stop_event = threading.Event() self._scan_completed = threading.Event() + self._current_agent: Any | None = None # set in _start_scan_thread for checkpointing self._spinner_frame_index: int = 0 # Current animation frame index self._sweep_num_squares: int = 6 # Number of squares in sweep animation @@ -823,13 +824,32 @@ class StrixTUIApp(App): # type: ignore[misc] return config def _setup_cleanup_handlers(self) -> None: + # Added for Resume Feature — save checkpoint on interrupt + def _save_checkpoint_on_interrupt() -> None: + mgr = self.agent_config.get("checkpoint_manager") + agent = getattr(self, "_current_agent", None) + if not mgr or not agent: + return + try: + mgr.save( + agent.state, + self.tracer, + self.scan_config, + self.agent_config.get("target_hash", ""), + agent.max_iterations, + ) + except Exception: # noqa: BLE001 + pass # non-fatal + def cleanup_on_exit() -> None: from strix.runtime import cleanup_runtime + _save_checkpoint_on_interrupt() self.tracer.cleanup() cleanup_runtime() def signal_handler(_signum: int, _frame: Any) -> None: + _save_checkpoint_on_interrupt() self.tracer.cleanup() sys.exit(0) @@ -1538,6 +1558,7 @@ class StrixTUIApp(App): # type: ignore[misc] try: agent = StrixAgent(self.agent_config) + self._current_agent = agent # expose for checkpoint on interrupt if not self._scan_stop_event.is_set(): loop.run_until_complete(agent.execute_scan(self.scan_config))