From 5fd2a64562febac8df22b9ef2593f66617708989 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sun, 26 Apr 2026 00:57:52 -0700 Subject: [PATCH] fix(persistence): close all 9 gaps from the resume audit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three critical correctness fixes + six TUI/audit/UX fixes from the parallel-agent audit. All changes verified by an end-to-end smoke that builds, persists, and re-hydrates state across two simulated process boundaries. Critical (resume integrity): 1. ``bus.cancel_descendants_graceful`` now calls ``_maybe_snapshot`` after mutating the ``stopping`` set. Previously, a process crash between user-initiated graceful-stop and the next finalize lost the stop signal — respawned agents would run forever instead of exiting. ``_respawn_subagents`` also gains a guard that skips agents in ``stopping`` so a previously-cancelled agent is not resurrected on resume. 2. ``Tracer.hydrate_from_run_dir`` now **raises** on corrupt ``vulnerabilities.json`` instead of swallowing the exception. The prior behaviour silently reset ``vulnerability_reports`` to empty, so the next ``add_vulnerability_report`` would allocate ``vuln-0001`` and overwrite the prior MD on disk — silent data loss. 3. ``--instruction`` passed on resume now reaches the model. The CLI captures whether the user explicitly passed an instruction (``args.user_explicit_instruction``) before ``_load_resume_state`` loads the persisted one. ``run_strix_scan`` reads ``scan_config["resume_instruction"]`` and, on resume, sends the new instruction to root's bus inbox before calling ``run_with_continuation`` (which uses ``initial_input=[]`` for SDK replay). The inject filter surfaces it on the next turn. 4. ``--resume X`` errors loudly when ``scan_state.json`` exists but ``bus.json`` doesn't. Previously this silently fresh-started in the same dir, confusing the user who explicitly asked to resume. TUI / audit / UX: 5. ``Tracer.hydrate_from_run_dir`` now reads ``bus.json`` too and pre-populates ``tracer.agents`` from the snapshot's ``statuses`` / ``names`` / ``parent_of``. Before this, the TUI tree on resume showed only currently-running agents; completed/crashed children from the prior run were invisible. 6. ``Tracer.hydrate_from_run_dir`` also seeds ``self._llm_stats`` from ``bus.stats_live + bus.stats_completed`` so the resume's footer shows cumulative tokens / requests across the prior run plus the resume segment, instead of resetting to zero. 7. ``Tracer.save_run_data`` now also writes ``run_metadata.json`` (start_time, run_id, run_name, targets, status), and ``hydrate_from_run_dir`` restores ``start_time`` from it. Prior behaviour reset start_time to ``now()`` on every Tracer init, breaking the final report's duration calc on resumed scans. 8. Per-agent todos persist to ``{run_dir}/todos.json`` (atomic write on every CRUD). ``hydrate_todos_from_disk`` (called from ``run_strix_scan``) reloads them so respawned subagents find their lists intact. Previously, the module-level ``_todos_storage`` was lost on every process restart. 9. ``_load_resume_state`` validates each ``cloned_repo_path`` from the persisted ``scan_state.json`` still exists on disk. Previously a deleted clone dir would let the resume proceed with an empty source tree, with agents silently scanning nothing. Bonus: ``bus.finalize`` no longer pops ``parent_of`` and ``names`` for finalized agents. Routing protection (don't accept ``send`` to finalized agents) comes from the ``statuses[id]`` terminal-state check in ``send`` itself, so dropping those keys was overzealous and made completed children invisible in ``view_agent_graph`` and the TUI tree. Co-Authored-By: Claude Opus 4.7 (1M context) --- strix/interface/cli.py | 4 + strix/interface/main.py | 33 +++++++ strix/interface/tui.py | 3 + strix/orchestration/bus.py | 16 +++- strix/orchestration/scan.py | 31 +++++++ strix/telemetry/scan_artifacts.py | 10 ++- strix/telemetry/tracer.py | 142 ++++++++++++++++++++++++++---- strix/tools/todo/tools.py | 106 +++++++++++++++++++++- 8 files changed, 320 insertions(+), 25 deletions(-) diff --git a/strix/interface/cli.py b/strix/interface/cli.py index d382a80..5369e5f 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -111,6 +111,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, + # 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. + "resume_instruction": getattr(args, "user_explicit_instruction", None) or "", } tracer = Tracer(args.run_name) diff --git a/strix/interface/main.py b/strix/interface/main.py index 0b597c1..63fea43 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -415,6 +415,11 @@ Examples: except Exception as e: parser.error(f"Failed to read instruction file '{instruction_path}': {e}") + # Capture before ``_load_resume_state`` overrides — used by the resume + # path in ``run_strix_scan`` to decide whether to inject the new + # instruction into the root's bus inbox after session replay. + args.user_explicit_instruction = args.instruction if args.resume else None + if args.resume: if args.target: parser.error( @@ -422,6 +427,14 @@ Examples: "the prior run left off, including the original target list." ) _load_resume_state(args, parser) + bus_path = Path("strix_runs") / args.resume / "bus.json" + if not bus_path.exists(): + parser.error( + f"--resume {args.resume}: missing {bus_path}. The run was " + f"persisted but never reached its first bus snapshot — " + f"there's nothing to resume from. Pick a fresh --run-name " + f"or remove --resume to start over with the same targets." + ) else: if not args.target: parser.error( @@ -499,6 +512,26 @@ def _load_resume_state(args: argparse.Namespace, parser: argparse.ArgumentParser if not args.targets_info: parser.error(f"--resume {args.resume}: scan_state.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 + # container mounts an empty source tree and agents silently scan + # nothing. + for target in args.targets_info: + if not isinstance(target, dict): + continue + details = target.get("details") or {} + if target.get("type") != "repository": + continue + cloned = details.get("cloned_repo_path") + if not cloned: + continue + if not Path(cloned).expanduser().exists(): + parser.error( + f"--resume {args.resume}: cloned repo at {cloned} is missing. " + f"It was deleted between runs. Pick a fresh --run-name to " + f"re-clone, or restore the directory before resuming." + ) + if args.instruction is None: args.instruction = state.get("instruction") if state.get("instruction_file") and args.instruction_file is None: diff --git a/strix/interface/tui.py b/strix/interface/tui.py index 1e1d9af..2c4a3c2 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -763,6 +763,9 @@ 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"), + # 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 "", } def _setup_cleanup_handlers(self) -> None: diff --git a/strix/orchestration/bus.py b/strix/orchestration/bus.py index 9c18d86..846bc14 100644 --- a/strix/orchestration/bus.py +++ b/strix/orchestration/bus.py @@ -226,15 +226,19 @@ class AgentMessageBus: async def finalize(self, agent_id: str, status: str) -> None: """Move an agent from live to completed; clean up routing state. - Also clears ``inboxes``, ``parent_of``, ``names`` so siblings - that send to a finished agent can't accumulate orphan messages. + Clears live routing state (``inboxes``, ``streams``, + ``_events``, ``stopping``, ``metadata`` for the spawn args) and + moves stats from ``stats_live`` to ``stats_completed``. Keeps + ``parent_of`` and ``names`` so the agent remains visible in + ``view_agent_graph`` and the TUI tree post-finalize. Routing + protection against late ``send`` calls comes from the + ``statuses[id]`` terminal-state check in :meth:`send`, not from + removing the parent/name keys. """ async with self._lock: self.statuses[agent_id] = status self.stats_completed[agent_id] = self.stats_live.pop(agent_id, {}) self.inboxes.pop(agent_id, None) - self.parent_of.pop(agent_id, None) - self.names.pop(agent_id, None) self.streams.pop(agent_id, None) self.stopping.discard(agent_id) self._events.pop(agent_id, None) @@ -380,6 +384,10 @@ class AgentMessageBus: ) for _aid, streamed in streams_to_cancel: streamed.cancel(mode="after_turn") + # Persist ``stopping`` so a process crash before any child finalizes + # doesn't lose the user's stop signal — otherwise resume would + # respawn the cancelled agents and they'd run forever. + await self._maybe_snapshot() # ------------------------------------------------------------------ # Snapshot / restore — persist serializable state to ``bus.json`` so diff --git a/strix/orchestration/scan.py b/strix/orchestration/scan.py index 441cb38..4ef4098 100644 --- a/strix/orchestration/scan.py +++ b/strix/orchestration/scan.py @@ -258,6 +258,13 @@ async def run_strix_scan( if tracer is not None and hasattr(tracer, "hydrate_from_run_dir"): tracer.hydrate_from_run_dir() + # Wire the per-agent todo store to ``{run_dir}/todos.json`` (mirrored + # on every CRUD) and reload any prior todos so respawned subagents + # find their lists intact. + from strix.tools.todo.tools import hydrate_todos_from_disk + + hydrate_todos_from_disk(run_dir) + root_id: str | None = None if is_resume: try: @@ -402,6 +409,27 @@ async def run_strix_scan( initial_input: Any = [] if is_resume else _build_root_task(scan_config) + # Resume + new ``--instruction``: SDK replay drives root from + # session.db with ``initial_input=[]``, so a brand-new instruction + # passed on the resume CLI would otherwise be silently ignored. + # Inject it as a fresh user message in root's inbox; the + # ``inject_messages_filter`` will surface it on the very next turn. + resume_instruction = str(scan_config.get("resume_instruction") or "").strip() + if is_resume and resume_instruction: + await bus.send( + root_id, + { + "from": "user", + "type": "instruction", + "priority": "high", + "content": resume_instruction, + }, + ) + logger.info( + "Resume: injected new instruction into root inbox (len=%d)", + len(resume_instruction), + ) + return await run_with_continuation( agent=root_agent, initial_input=initial_input, @@ -468,6 +496,9 @@ async def _respawn_subagents( if status in {"running", "waiting", "llm_failed"} and bus.parent_of.get(aid) is not None and aid != root_id + # Honour a previously-issued graceful stop — the user clicked + # "stop" before the crash; don't respawn what they cancelled. + and aid not in bus.stopping ] for child_id, name, parent_id, md in candidates: diff --git a/strix/telemetry/scan_artifacts.py b/strix/telemetry/scan_artifacts.py index 7409ad0..dee9b8e 100644 --- a/strix/telemetry/scan_artifacts.py +++ b/strix/telemetry/scan_artifacts.py @@ -46,9 +46,11 @@ class ScanArtifactWriter: *, vulnerability_reports: list[dict[str, Any]], final_scan_result: str | None, + run_metadata: dict[str, Any] | None = None, ) -> None: """Write any new vulnerability MDs + rewrite the CSV index + - write the executive penetration-test report if available. + write the executive penetration-test report if available + dump + ``run_metadata.json`` for resume hydration. Tolerant of OSError / RuntimeError — logs and swallows so a cleanup failure can't prevent the next scan from finishing. @@ -62,6 +64,12 @@ class ScanArtifactWriter: if vulnerability_reports: self._write_vulnerabilities(vulnerability_reports) + if run_metadata is not None: + _atomic_write_text( + self._run_dir / "run_metadata.json", + json.dumps(run_metadata, ensure_ascii=False, indent=2, default=str), + ) + logger.info("📊 Essential scan data saved to: %s", self._run_dir) except (OSError, RuntimeError): logger.exception("Failed to save scan data") diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index c809212..72eb589 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -95,24 +95,67 @@ class Tracer: return self._run_dir def hydrate_from_run_dir(self) -> None: - """Reload ``vulnerability_reports`` from ``{run_dir}/vulnerabilities.json``. + """Reload prior-scan state from ``{run_dir}/`` for resume. - Called by the resume path in :func:`run_strix_scan` before any - new agent runs. Ensures id allocation in - :meth:`add_vulnerability_report` does not collide on disk - (``vuln-0001`` re-used would otherwise overwrite the prior MD). - Idempotent — calling without a JSON file is a no-op. + Called by :func:`run_strix_scan` before any new agent runs. + Restores: + + - ``vulnerability_reports`` from ``vulnerabilities.json`` so + :meth:`add_vulnerability_report` doesn't allocate a colliding + ``vuln-0001`` and overwrite the prior on-disk MD. + - ``run_metadata`` (start_time, run_id, targets, status) from + ``run_metadata.json`` so audit-trail timestamps + the final + report's duration calc reflect the original scan, not just + this resume segment. + + Idempotent on missing files (fresh runs land here too via the + same code path). **Raises on corruption** — silently swallowing + a corrupt ``vulnerabilities.json`` would let the next vuln + allocate ``vuln-0001`` and overwrite the prior MD on disk + (data loss). Caller is expected to fail the run loud and let + the user inspect ``{run_dir}`` or pick a fresh ``--run-name``. """ - try: - json_path = self.get_run_dir() / "vulnerabilities.json" - if not json_path.exists(): - return - data = json.loads(json_path.read_text(encoding="utf-8")) + run_dir = self.get_run_dir() + + meta_path = run_dir / "run_metadata.json" + if meta_path.exists(): + try: + data = json.loads(meta_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"run_metadata.json at {meta_path} is unreadable: {exc}", + ) from exc + if isinstance(data, dict): + if isinstance(data.get("start_time"), str): + self.start_time = data["start_time"] + self.run_metadata.update( + { + k: v + for k, v in data.items() + if k in {"run_id", "run_name", "start_time", "targets", "status"} + }, + ) + logger.info( + "tracer hydrated run_metadata from %s (start_time=%s)", + meta_path, + self.start_time, + ) + + json_path = run_dir / "vulnerabilities.json" + if json_path.exists(): + try: + data = json.loads(json_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"vulnerabilities.json at {json_path} is corrupt ({exc}); " + f"refusing to start fresh — that would overwrite prior " + f"vulnerability MDs on disk. Inspect or delete the run dir.", + ) from exc if not isinstance(data, list): - return + raise RuntimeError( + f"vulnerabilities.json at {json_path} is not a list", + ) self.vulnerability_reports = [r for r in data if isinstance(r, dict)] - # Pre-mark these ids as already-saved so the writer doesn't - # re-emit per-vuln markdown on the next save() call. writer = self._get_writer() for r in self.vulnerability_reports: rid = r.get("id") @@ -123,8 +166,74 @@ class Tracer: len(self.vulnerability_reports), json_path, ) - except Exception: - logger.exception("tracer hydrate_from_run_dir failed; starting fresh") + + bus_path = run_dir / "bus.json" + if bus_path.exists(): + try: + bus_data = json.loads(bus_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + # Caller will surface this same corruption via ``bus.restore``; + # no need to fail twice. Skip the agents/stats hydrate path. + bus_data = None + if isinstance(bus_data, dict): + self._hydrate_agents_tree(bus_data) + self._hydrate_llm_stats(bus_data) + + def _hydrate_agents_tree(self, bus_data: dict[str, Any]) -> None: + """Populate ``self.agents`` from a bus snapshot. + + Without this, the TUI tree on resume would only show agents + currently running (mirrored by ``on_agent_start``); completed / + crashed / stopped children from the prior run would be invisible + even though the bus knows about them. + """ + statuses = bus_data.get("statuses") or {} + names = bus_data.get("names") or {} + parent_of = bus_data.get("parent_of") or {} + if not isinstance(statuses, dict): + return + timestamp = self.start_time + for agent_id, status in statuses.items(): + if not isinstance(agent_id, str): + continue + self.agents[agent_id] = { + "id": agent_id, + "name": names.get(agent_id, agent_id) if isinstance(names, dict) else agent_id, + "parent_id": parent_of.get(agent_id) if isinstance(parent_of, dict) else None, + "status": status, + "created_at": timestamp, + "updated_at": timestamp, + } + logger.info("tracer hydrated %d agent(s) into tree", len(self.agents)) + + def _hydrate_llm_stats(self, bus_data: dict[str, Any]) -> None: + """Seed ``self._llm_stats`` from the bus snapshot's per-agent counters. + + Aggregates ``stats_live + stats_completed`` so the resumed scan's + TUI footer shows cumulative tokens / requests across the prior + run plus whatever the resume adds, instead of resetting to zero. + """ + totals = {"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0, "requests": 0} + for bucket_key in ("stats_live", "stats_completed"): + bucket = bus_data.get(bucket_key) or {} + if not isinstance(bucket, dict): + continue + for entry in bucket.values(): + if not isinstance(entry, dict): + continue + totals["input_tokens"] += int(entry.get("in", 0) or 0) + totals["output_tokens"] += int(entry.get("out", 0) or 0) + totals["cached_tokens"] += int(entry.get("cached", 0) or 0) + totals["requests"] += int(entry.get("calls", 0) or 0) + for k, v in totals.items(): + self._llm_stats[k] = v + logger.info( + "tracer hydrated llm stats from bus (in=%d out=%d cached=%d requests=%d)", + totals["input_tokens"], + totals["output_tokens"], + totals["cached_tokens"], + totals["requests"], + ) def _get_writer(self) -> ScanArtifactWriter: if self._writer is None: @@ -280,6 +389,7 @@ class Tracer: self._get_writer().save( vulnerability_reports=self.vulnerability_reports, final_scan_result=self.final_scan_result, + run_metadata=dict(self.run_metadata), ) def log_tool_start( diff --git a/strix/tools/todo/tools.py b/strix/tools/todo/tools.py index 6ba5ddf..73d0623 100644 --- a/strix/tools/todo/tools.py +++ b/strix/tools/todo/tools.py @@ -1,21 +1,32 @@ """Per-agent todo tools. -In-memory only — todos live for the lifetime of one scan, scoped per -agent via ``ctx.context['agent_id']``. Bulk forms are preserved so the -prompt-template documentation still works (``todos`` / ``updates`` / -``todo_ids`` accept JSON strings or comma-separated strings). +Per-agent in-memory dict, scoped via ``ctx.context['agent_id']``. The +table is mirrored to ``{run_dir}/todos.json`` after every mutation so a +process restart can ``hydrate_todos_from_disk`` and each respawned +agent finds its prior list intact. The persistence is best-effort — +errors are logged and swallowed so a disk failure can't kill the agent +mid-call. Bulk forms are preserved so the prompt-template documentation +still works (``todos`` / ``updates`` / ``todo_ids`` accept JSON strings +or comma-separated strings). """ from __future__ import annotations import json +import logging +import tempfile +import threading import uuid from datetime import UTC, datetime +from pathlib import Path from typing import Any from agents import RunContextWrapper, function_tool +logger = logging.getLogger(__name__) + + VALID_PRIORITIES = ["low", "normal", "high", "critical"] VALID_STATUSES = ["pending", "in_progress", "done"] @@ -36,6 +47,86 @@ def _todo_sort_key(todo: dict[str, Any]) -> tuple[int, int, str]: # don't see each other's lists. _todos_storage: dict[str, dict[str, dict[str, Any]]] = {} +# On-disk mirror path. Set by ``hydrate_todos_from_disk`` once per scan; +# unset means "no persistence" (e.g. unit tests). All writes go through +# ``_persist`` which is a no-op until the path is set. +_todos_path: Path | None = None +_todos_io_lock = threading.RLock() + + +def hydrate_todos_from_disk(run_dir: Path) -> None: + """Wire the on-disk mirror at ``{run_dir}/todos.json`` and reload it. + + Called by :func:`run_strix_scan` once at scan setup. Subsequent CRUD + calls auto-persist after every mutation. Idempotent on missing file. + Tolerant of corruption — logs and starts empty rather than failing + the scan over a broken sidecar artifact. + """ + global _todos_path # noqa: PLW0603 + _todos_path = run_dir / "todos.json" + with _todos_io_lock: + _todos_storage.clear() + if not _todos_path.exists(): + return + try: + data = json.loads(_todos_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + logger.exception( + "todos.json at %s is unreadable; starting with empty todos", + _todos_path, + ) + return + if not isinstance(data, dict): + return + loaded = 0 + for aid, by_id in data.items(): + if not isinstance(aid, str) or not isinstance(by_id, dict): + continue + cleaned = { + str(tid): t + for tid, t in by_id.items() + if isinstance(tid, str) and isinstance(t, dict) + } + if cleaned: + _todos_storage[aid] = cleaned + loaded += len(cleaned) + logger.info( + "todos hydrated from %s (%d agent(s), %d todo(s))", + _todos_path, + len(_todos_storage), + loaded, + ) + + +def _persist() -> None: + """Atomic-rename mirror of ``_todos_storage`` → ``{run_dir}/todos.json``. + + No-op when ``_todos_path`` isn't wired (tests). Errors are logged + and swallowed. + """ + path = _todos_path + if path is None: + return + try: + payload = json.dumps(_todos_storage, ensure_ascii=False, default=str) + path.parent.mkdir(parents=True, exist_ok=True) + with ( + _todos_io_lock, + tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=str(path.parent), + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as tmp, + ): + tmp.write(payload) + tmp_path = Path(tmp.name) + tmp_path.replace(path) + except Exception: + logger.exception("todos persist to %s failed", path) + def _agent_id_from(ctx: RunContextWrapper) -> str: inner = ctx.context if isinstance(ctx.context, dict) else {} @@ -279,6 +370,7 @@ async def create_todo( default=str, ) + _persist() return json.dumps( { "success": True, @@ -419,6 +511,8 @@ async def update_todo( except (ValueError, TypeError) as e: return json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, default=str) + if updated: + _persist() response: dict[str, Any] = { "success": len(errors) == 0, "updated": updated, @@ -464,6 +558,8 @@ def _mark( except (ValueError, TypeError) as e: return json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, default=str) + if marked: + _persist() key = "marked_done" if new_status == "done" else "marked_pending" response: dict[str, Any] = { "success": len(errors) == 0, @@ -556,6 +652,8 @@ async def delete_todo( except (ValueError, TypeError) as e: return json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, default=str) + if deleted: + _persist() response: dict[str, Any] = { "success": len(errors) == 0, "deleted": deleted,