From d538acf66bb5c5e671f7ea76d572b7cf82d5e0aa Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sun, 26 Apr 2026 00:29:37 -0700 Subject: [PATCH] feat(orchestration): always-on resume across the agent graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A scan that crashes or is stopped can now be resumed by re-invoking ``strix`` with the same ``--run-name``. Resume is implicit — presence of ``{run_dir}/bus.json`` triggers it. To force a fresh start, delete the run dir. What survives a process restart with the same scan_id: * Root agent's LLM history — already worked (root SDK SQLiteSession). * Every non-terminal subagent's LLM history — new. ``create_agent`` now opens SQLiteSession(session_id=child_id, db_path={run_dir}/sessions/{child_id}.db) per child and passes it to ``run_with_continuation``. * Bus topology — new. ``AgentMessageBus`` gains snapshot/restore/ _maybe_snapshot async methods plus a ``metadata`` field that holds per-agent {task, skills, is_whitebox, scan_mode, diff_scope}. ``register``, ``finalize``, ``park``, and ``mark_llm_failed`` each call ``_maybe_snapshot`` to atomically persist the bus to {run_dir}/bus.json (tempfile + Path.replace). * Vulnerability reports — new. ``ScanArtifactWriter._write_ vulnerabilities`` now also writes ``vulnerabilities.json`` (atomic). ``Tracer.hydrate_from_run_dir`` reads it on resume so new vuln-NNNN ids don't collide with prior on-disk files. What does not survive: the sandbox container itself (fresh per process), so ``/workspace/scratch`` and Caido state are lost. ``/workspace/sources`` re-mounts from the host so source code is unchanged. ``orchestration/scan.py:run_strix_scan`` does the actual resume: 1. Resolve run_dir up front; if bus.json exists it's a resume. 2. Acquire {run_dir}/.lock (fcntl.flock) so a second strix process can't run concurrently on the same scan_id. 3. ``bus.set_snapshot_path(...)``, ``tracer.hydrate_from_run_dir()``. 4. On resume: load + bus.restore, find root_id from snapshot (the agent with parent_of[id] is None), spawn the sandbox, skip the root's bus.register (already in snapshot). 5. ``_respawn_subagents`` walks every agent with status in running/waiting/llm_failed: reopens its SQLiteSession, rebuilds the child agent via the captured factory, builds run config / context, asyncio.create_task the run with initial_input=[] so the SDK replays from session. Per-child failure (missing/corrupt DB, factory raises) finalizes that child as crashed and continues. 6. Open root SQLiteSession at the same path, run the root with initial_input=[] on resume (or the formatted root task on a fresh run), and let SDK replay drive the next turn. 7. ``finally``: close every per-agent session, take a final snapshot, tear down sandbox, release the lock. HARNESS_WIKI.md updated with the new run-dir layout (sessions/, bus.json, vulnerabilities.json, .lock) and the resume contract. Net: +500 LoC across 7 files. No new deps. Co-Authored-By: Claude Opus 4.7 (1M context) --- HARNESS_WIKI.md | 30 +++- pyproject.toml | 2 +- strix/orchestration/bus.py | 108 ++++++++++- strix/orchestration/scan.py | 290 +++++++++++++++++++++++++++--- strix/telemetry/scan_artifacts.py | 39 +++- strix/telemetry/tracer.py | 33 ++++ strix/tools/agents_graph/tools.py | 36 +++- 7 files changed, 496 insertions(+), 42 deletions(-) diff --git a/HARNESS_WIKI.md b/HARNESS_WIKI.md index a00341f..f12e8e4 100644 --- a/HARNESS_WIKI.md +++ b/HARNESS_WIKI.md @@ -530,13 +530,35 @@ Async run loop. Rich panels for vuln-found events, live stats panel updated ever ### 9.6 Run Directory Layout (`strix_runs//`) -Created and managed by `telemetry/tracer.py`. Contents: +Created and managed by `telemetry/tracer.py` + `orchestration/scan.py`. Contents: - `events.jsonl` — every span/event in append-only JSONL (thread-safe writes). -- `vulnerabilities/vuln_{id}.json` — one file per finding, sorted by severity, dedup-checked. -- `penetration_test_report.md` — final markdown report (executive summary + methodology + technical analysis + recommendations). +- `vulnerabilities/vuln_{id}.md` — one file per finding, sorted by severity, dedup-checked. +- `vulnerabilities.csv` — index of every finding for spreadsheet consumption. +- `vulnerabilities.json` — machine-readable mirror, used by `Tracer.hydrate_from_run_dir` to repopulate vuln state on resume so id allocation doesn't collide. +- `penetration_test_report.md` — final markdown report. +- `session.db` — SDK `SQLiteSession` for the **root** agent's conversation history. +- `sessions/{child_id}.db` — per-subagent `SQLiteSession`, one file per spawned child. +- `bus.json` — atomic snapshot of the orchestration bus (topology, statuses, inboxes, per-agent stats, per-agent metadata `{task, skills, is_whitebox, scan_mode, diff_scope}`). Written after every `bus.register` / `finalize` / `park` / `mark_llm_failed`, plus a final dump at scan teardown. +- `.lock` — advisory `flock`-style file lock; prevents two `strix` processes from running on the same `scan_id` concurrently. +- `strix.log` — per-scan log file (DEBUG to file, ERROR to stderr). - `/` — local source clones, per-target. -There is **no execution checkpointing** — if the process crashes mid-run, the agent restarts from scratch on retry. Resumability is limited to the interactive-mode wait/resume on inbound messages. +### 9.7 Resume + +Resume is **always on**: presence of `bus.json` triggers it. Fresh runs simply have no `bus.json` to begin with. To force a fresh start with the same `--run-name`, delete the run dir. + +What survives a process restart with the same `scan_id`: +- Root agent's full LLM conversation (replayed by SDK from `session.db`). +- Every non-terminal subagent's full LLM conversation (replayed from `sessions/{child_id}.db`). +- Bus topology: `parent_of`, `statuses`, `names`, `stats_live`, `stats_completed`, pending `inboxes`, `metadata` (task + skills per agent). +- Vulnerability reports (hydrated from `vulnerabilities.json`). +- Run log (appended to existing `strix.log`). + +What does **not** survive: +- The sandbox container itself — fresh container per process. Files agents wrote under `/workspace/scratch/` or scanner outputs to `/workspace/.strix-source-aware/` are lost. `/workspace/sources` re-mounts from the host so source code is unchanged. +- Caido proxy state (request log, scopes, replay sessions). + +On resume, every subagent in `running` / `waiting` / `llm_failed` status is respawned via `_respawn_subagents` with `initial_input=[]`; the SDK's session replay drives them from where they stopped. Terminal-status agents (`completed` / `crashed` / `stopped`) are left alone but their stats remain in `stats_completed` for the TUI footer. Per-child failure (corrupt session DB, missing skill module) finalizes that child as `crashed` and continues with the rest. --- diff --git a/pyproject.toml b/pyproject.toml index 5b69b33..87cb0ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -237,7 +237,7 @@ ignore = [ # resolution past where mypy needs it. ``_build_root_task`` legitimately # walks every supported target type — splitting it into per-type # helpers would add indirection without simplifying anything. -"strix/orchestration/scan.py" = ["TC003", "PLR0912"] +"strix/orchestration/scan.py" = ["TC003", "PLR0912", "PLR0915", "PLC0415"] # Tracer carries a long event surface and a runtime ``Callable`` # annotation on ``vulnerability_found_callback``. "strix/telemetry/tracer.py" = ["TC003", "PLR0912", "PLR0915", "E501"] diff --git a/strix/orchestration/bus.py b/strix/orchestration/bus.py index 50ea156..9c18d86 100644 --- a/strix/orchestration/bus.py +++ b/strix/orchestration/bus.py @@ -9,8 +9,11 @@ from __future__ import annotations import asyncio import contextlib +import json import logging +import tempfile from dataclasses import dataclass, field +from pathlib import Path from typing import TYPE_CHECKING, Any @@ -23,6 +26,9 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +_SNAPSHOT_VERSION = 1 + + @dataclass class AgentMessageBus: """Shared state for multi-agent orchestration. @@ -58,16 +64,33 @@ class AgentMessageBus: stats_live: dict[str, dict[str, Any]] = field(default_factory=dict) stats_completed: dict[str, dict[str, Any]] = field(default_factory=dict) stopping: set[str] = field(default_factory=set) + metadata: dict[str, dict[str, Any]] = field(default_factory=dict) _events: dict[str, asyncio.Event] = field(default_factory=dict) _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + _snapshot_path: Path | None = None + + def set_snapshot_path(self, path: Path) -> None: + """Wire the on-disk snapshot path. Triggers persist after each lifecycle event.""" + self._snapshot_path = path async def register( self, agent_id: str, name: str, parent_id: str | None, + *, + task: str | None = None, + skills: list[str] | None = None, + is_whitebox: bool = False, + scan_mode: str = "deep", + diff_scope: dict[str, Any] | None = None, ) -> None: - """Add a new agent to the bus before its Runner.run task starts.""" + """Add a new agent to the bus before its Runner.run task starts. + + ``task`` / ``skills`` / ``is_whitebox`` / ``scan_mode`` / + ``diff_scope`` are persisted on the bus's ``metadata`` map so a + process restart can rebuild the same agent from the snapshot. + """ async with self._lock: self.inboxes[agent_id] = [] self.statuses[agent_id] = "running" @@ -82,7 +105,15 @@ class AgentMessageBus: "warned_85": False, "warned_final": False, } + self.metadata[agent_id] = { + "task": task or "", + "skills": list(skills or []), + "is_whitebox": bool(is_whitebox), + "scan_mode": scan_mode, + "diff_scope": diff_scope, + } logger.info("bus.register %s (%s) parent=%s", agent_id, name, parent_id or "-") + await self._maybe_snapshot() async def send(self, target: str, msg: dict[str, Any]) -> None: """Append a message to ``target``'s inbox. @@ -207,7 +238,9 @@ class AgentMessageBus: self.streams.pop(agent_id, None) self.stopping.discard(agent_id) self._events.pop(agent_id, None) + self.metadata.pop(agent_id, None) logger.info("bus.finalize %s status=%s", agent_id, status) + await self._maybe_snapshot() async def park(self, agent_id: str) -> None: """Mark an agent as ``waiting`` without finalizing. @@ -222,6 +255,7 @@ class AgentMessageBus: if agent_id in self.statuses: self.statuses[agent_id] = "waiting" logger.debug("bus.park %s", agent_id) + await self._maybe_snapshot() async def mark_llm_failed(self, agent_id: str) -> None: """Mark an agent as ``llm_failed`` after retries exhausted. @@ -236,6 +270,7 @@ class AgentMessageBus: if agent_id in self.statuses: self.statuses[agent_id] = "llm_failed" logger.warning("bus.mark_llm_failed %s — awaiting user resume", agent_id) + await self._maybe_snapshot() @contextlib.asynccontextmanager async def attach_stream( @@ -345,3 +380,74 @@ class AgentMessageBus: ) for _aid, streamed in streams_to_cancel: streamed.cancel(mode="after_turn") + + # ------------------------------------------------------------------ + # Snapshot / restore — persist serializable state to ``bus.json`` so + # a process restart can rebuild the topology + per-agent metadata + # and respawn each non-terminal agent. + # ------------------------------------------------------------------ + async def snapshot(self) -> dict[str, Any]: + """Return a JSON-ready deep copy of every persistable bus field. + + Held under ``_lock`` for the copy so the caller can't observe a + partial mutation. The actual JSON serialisation happens outside + the lock — that's fine for ``json.dumps``: pure-Python primitives, + no I/O. + """ + async with self._lock: + return { + "version": _SNAPSHOT_VERSION, + "inboxes": {aid: list(msgs) for aid, msgs in self.inboxes.items()}, + "statuses": dict(self.statuses), + "parent_of": dict(self.parent_of), + "names": dict(self.names), + "stats_live": {aid: dict(s) for aid, s in self.stats_live.items()}, + "stats_completed": {aid: dict(s) for aid, s in self.stats_completed.items()}, + "stopping": list(self.stopping), + "metadata": {aid: dict(m) for aid, m in self.metadata.items()}, + } + + async def restore(self, snap: dict[str, Any]) -> None: + """Repopulate from a prior :meth:`snapshot`. Tasks/streams/events + stay empty — they're rebuilt by the resume path as agents respawn. + """ + async with self._lock: + self.inboxes = {aid: list(msgs) for aid, msgs in snap.get("inboxes", {}).items()} + self.statuses = dict(snap.get("statuses", {})) + self.parent_of = dict(snap.get("parent_of", {})) + self.names = dict(snap.get("names", {})) + self.stats_live = {aid: dict(s) for aid, s in snap.get("stats_live", {}).items()} + self.stats_completed = { + aid: dict(s) for aid, s in snap.get("stats_completed", {}).items() + } + self.stopping = set(snap.get("stopping", [])) + self.metadata = {aid: dict(m) for aid, m in snap.get("metadata", {}).items()} + + async def _maybe_snapshot(self) -> None: + """Persist the current state to ``_snapshot_path`` if one is set. + + No-op when ``_snapshot_path`` is unset (e.g. in tests). Atomic + ``tempfile`` + ``os.replace`` so a crash mid-write leaves the + previous valid file intact. Errors are logged + swallowed; a + snapshot failure should never tear down the run. + """ + path = self._snapshot_path + if path is None: + return + try: + data = await self.snapshot() + payload = json.dumps(data, ensure_ascii=False, default=str) + path.parent.mkdir(parents=True, exist_ok=True) + with 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("bus snapshot to %s failed", path) diff --git a/strix/orchestration/scan.py b/strix/orchestration/scan.py index 8e21c2a..441cb38 100644 --- a/strix/orchestration/scan.py +++ b/strix/orchestration/scan.py @@ -1,20 +1,30 @@ -"""Top-level scan entry point. +"""Top-level scan entry point with auto-resume. -1. Build the per-scan ``AgentMessageBus``. -2. Bring up (or reuse) a sandbox session for ``scan_id`` via the - :mod:`strix.runtime.session_manager`. -3. Build the root ``Agent`` via :func:`build_strix_agent` and a - matching child factory via :func:`make_child_factory`. -4. Build the root context dict (bus + sandbox bundle + agent_factory). -5. Register the root in the bus. -6. Build the ``RunConfig`` via the factory. -7. Call ``Runner.run(...)`` and surface the result. -8. ``finally`` cleanup the sandbox session — even on cancel, the bus - propagates ``cancel_descendants`` to every spawned child task. +1. Build (or take from caller) the per-scan ``AgentMessageBus``. +2. Wire a snapshot path so every lifecycle event auto-persists ``bus.json``. +3. Acquire an advisory file lock so a second ``strix`` process can't run + on the same ``scan_id`` concurrently. +4. **Resume detection**: if ``{run_dir}/bus.json`` already exists, restore + the bus, hydrate the tracer, reuse the persisted ``root_id`` instead + of generating a fresh one, and respawn every non-terminal subagent + from its per-child ``SQLiteSession`` before starting the root. +5. Bring up (or reuse) a sandbox session for ``scan_id``. +6. Build the root ``Agent`` + child factory. +7. Open root ``SQLiteSession`` at the same path so the SDK replays prior + turns on resume. +8. Call ``Runner.run`` (via ``run_with_continuation``). +9. ``finally``: close every per-agent session, take a final snapshot, + tear down the sandbox, release the lock. + +Resume is **always on**: there is no flag — presence of ``bus.json`` is +the trigger. Fresh runs simply have no ``bus.json`` to begin with. """ from __future__ import annotations +import asyncio +import contextlib +import json import logging import uuid from pathlib import Path @@ -209,10 +219,20 @@ async def run_strix_scan( if tracer is not None and hasattr(tracer, "get_run_dir") else Path.cwd() / "strix_runs" / scan_id ) + run_dir.mkdir(parents=True, exist_ok=True) teardown_logging = setup_scan_logging(run_dir) set_scan_id(scan_id) + + bus_path = run_dir / "bus.json" + is_resume = bus_path.exists() + sessions_dir = run_dir / "sessions" + sessions_dir.mkdir(parents=True, exist_ok=True) + + lock_handle = _acquire_run_lock(run_dir) + logger.info( - "Starting Strix scan %s (image=%s, max_turns=%d, interactive=%s, run_dir=%s)", + "%s Strix scan %s (image=%s, max_turns=%d, interactive=%s, run_dir=%s)", + "Resuming" if is_resume else "Starting", scan_id, image, max_turns, @@ -222,6 +242,7 @@ async def run_strix_scan( resolved_model = model or load_settings().llm.model if not resolved_model: + _release_run_lock(lock_handle) raise RuntimeError( "No LLM model configured. Set STRIX_LLM env or pass model= to run_strix_scan().", ) @@ -232,9 +253,41 @@ async def run_strix_scan( # own the bus internally for the scan's lifetime. if bus is None: bus = AgentMessageBus() - root_id = uuid.uuid4().hex[:8] - logger.info("Bringing up sandbox session for scan %s", scan_id) + bus.set_snapshot_path(bus_path) + if tracer is not None and hasattr(tracer, "hydrate_from_run_dir"): + tracer.hydrate_from_run_dir() + + root_id: str | None = None + if is_resume: + try: + snap = json.loads(bus_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + _release_run_lock(lock_handle) + raise RuntimeError( + f"Cannot resume scan {scan_id}: bus.json is unreadable: {exc}", + ) from exc + await bus.restore(snap) + for aid, parent in bus.parent_of.items(): + if parent is None: + root_id = aid + break + if root_id is None: + _release_run_lock(lock_handle) + raise RuntimeError( + f"Cannot resume scan {scan_id}: bus.json has no root agent (parent=None)", + ) + logger.info( + "Resume: restored bus with %d agent(s); root=%s; %d non-terminal to respawn", + len(bus.statuses), + root_id, + sum(1 for s in bus.statuses.values() if s in {"running", "waiting", "llm_failed"}) + - 1, # subtract root + ) + else: + root_id = uuid.uuid4().hex[:8] + + logger.info("Bringing up sandbox session for scan %s", scan_id) bundle = await session_manager.create_or_reuse( scan_id, image=image, @@ -242,6 +295,8 @@ async def run_strix_scan( ) logger.info("Sandbox ready for scan %s", scan_id) + sessions_to_close: list[SQLiteSession] = [] + try: # Lazy: ``strix.interface`` pulls cli→tui→scan which would cycle. from strix.interface.utils import is_whitebox_scan @@ -264,7 +319,17 @@ async def run_strix_scan( system_prompt_context=scope_context, ) - await bus.register(root_id, "strix", parent_id=None) + if not is_resume: + await bus.register( + root_id, + "strix", + parent_id=None, + task=_build_root_task(scan_config), + skills=skills, + is_whitebox=is_whitebox, + scan_mode=scan_mode, + diff_scope=diff_scope, + ) agent_factory = make_child_factory( scan_mode=scan_mode, @@ -287,9 +352,11 @@ async def run_strix_scan( "agent_finish_called": False, "is_whitebox": is_whitebox, "interactive": interactive, + "scan_mode": scan_mode, "diff_scope": diff_scope, "run_id": run_id, "agent_factory": agent_factory, + "_sessions_to_close": sessions_to_close, } reasoning_effort: Literal["low", "medium", "high"] | None = ( @@ -314,20 +381,30 @@ async def run_strix_scan( trace_include_sensitive_data=False, ) - # Native SDK session: persists conversation history to - # ``strix_runs//session.db`` so a second invocation - # with the same ``scan_id`` resumes from where we left off. - session_db = ( - (tracer.get_run_dir() / "session.db") - if tracer is not None and hasattr(tracer, "get_run_dir") - else Path.cwd() / "strix_runs" / scan_id / "session.db" - ) - session_db.parent.mkdir(parents=True, exist_ok=True) - session = SQLiteSession(session_id=scan_id, db_path=session_db) + if is_resume: + await _respawn_subagents( + bus=bus, + sessions_dir=sessions_dir, + factory=agent_factory, + parent_ctx=context, + resolved_model=resolved_model, + reasoning_effort=reasoning_effort, + root_id=root_id, + sessions_to_close=sessions_to_close, + ) + + # Root SDK session — same path on fresh + resume so SDK replay + # picks up prior turns automatically when ``initial_input`` is + # an empty list. + session_db = run_dir / "session.db" + root_session = SQLiteSession(session_id=scan_id, db_path=session_db) + sessions_to_close.append(root_session) + + initial_input: Any = [] if is_resume else _build_root_task(scan_config) return await run_with_continuation( agent=root_agent, - initial_input=_build_root_task(scan_config), + initial_input=initial_input, run_config=run_config, context=context, hooks=StrixOrchestrationHooks(), @@ -335,17 +412,172 @@ async def run_strix_scan( bus=bus, agent_id=root_id, interactive=interactive, - session=session, + session=root_session, ) except BaseException: logger.exception("Strix scan %s failed", scan_id) # Cancel any descendant tasks the root spawned before unwinding. # cancel_descendants is idempotent and handles the empty-tree case. - await bus.cancel_descendants(root_id) + if root_id is not None: + await bus.cancel_descendants(root_id) raise finally: + for s in sessions_to_close: + with contextlib.suppress(Exception): + s.close() + with contextlib.suppress(Exception): + await bus._maybe_snapshot() if cleanup_on_exit: logger.info("Tearing down sandbox session for scan %s", scan_id) await session_manager.cleanup(scan_id) + _release_run_lock(lock_handle) logger.info("Strix scan %s done", scan_id) teardown_logging() + + +async def _respawn_subagents( + *, + bus: AgentMessageBus, + sessions_dir: Path, + factory: Any, + parent_ctx: dict[str, Any], + resolved_model: str, + reasoning_effort: Literal["low", "medium", "high"] | None, + root_id: str, + sessions_to_close: list[SQLiteSession], +) -> None: + """Re-spawn every non-terminal subagent from a restored bus snapshot. + + Each child gets its own :class:`SQLiteSession` reopened at + ``sessions_dir/.db`` so the SDK replays its prior + conversation. Per-child failure (missing/corrupt session DB, + factory raising) finalizes that child as ``crashed`` and continues. + Terminal-status agents (``completed`` / ``crashed`` / ``stopped``) + are left alone — their stats stay in ``stats_completed`` for the + TUI, but no task respawns. + """ + async with bus._lock: + candidates = [ + ( + aid, + bus.names.get(aid, aid), + bus.parent_of.get(aid), + dict(bus.metadata.get(aid, {})), + ) + for aid, status in bus.statuses.items() + if status in {"running", "waiting", "llm_failed"} + and bus.parent_of.get(aid) is not None + and aid != root_id + ] + + for child_id, name, parent_id, md in candidates: + try: + session_path = sessions_dir / f"{child_id}.db" + if not session_path.exists(): + logger.warning( + "respawn %s (%s): session db missing at %s — finalizing as crashed", + child_id, + name, + session_path, + ) + await bus.finalize(child_id, "crashed") + continue + + child_session = SQLiteSession(session_id=child_id, db_path=session_path) + sessions_to_close.append(child_session) + + child_skills = list(md.get("skills") or []) + child_agent = factory(name=name, skills=child_skills) + + child_ctx: dict[str, Any] = dict(parent_ctx) + child_ctx["agent_id"] = child_id + child_ctx["parent_id"] = parent_id + child_ctx["agent_finish_called"] = False + child_ctx["task"] = md.get("task", "") + + child_model_settings = ModelSettings( + parallel_tool_calls=False, + tool_choice="required", + retry=DEFAULT_RETRY, + ) + if reasoning_effort is not None: + child_model_settings = child_model_settings.resolve( + ModelSettings(reasoning=Reasoning(effort=reasoning_effort)), + ) + child_run_config = RunConfig( + model=resolved_model, + model_provider=build_multi_provider(), + model_settings=child_model_settings, + sandbox=SandboxRunConfig( + client=parent_ctx["sandbox_client"], + session=parent_ctx["sandbox_session"], + ), + call_model_input_filter=inject_messages_filter, + tracing_disabled=False, + trace_include_sensitive_data=False, + ) + + task_handle = asyncio.create_task( + run_with_continuation( + agent=child_agent, + initial_input=[], + run_config=child_run_config, + context=child_ctx, + hooks=StrixOrchestrationHooks(), + max_turns=int(parent_ctx.get("max_turns", 300)), + bus=bus, + agent_id=child_id, + interactive=bool(parent_ctx.get("interactive", False)), + session=child_session, + ), + name=f"agent-{name}-{child_id}", + ) + async with bus._lock: + bus.tasks[child_id] = task_handle + logger.info( + "respawned %s (%s) parent=%s task_len=%d", + child_id, + name, + parent_id or "-", + len(md.get("task", "")), + ) + except Exception: + logger.exception("respawn %s failed; marking crashed", child_id) + with contextlib.suppress(Exception): + await bus.finalize(child_id, "crashed") + + +def _acquire_run_lock(run_dir: Path) -> Any: + """Take an exclusive flock on ``{run_dir}/.lock`` so two strix processes + can't run on the same scan_id concurrently. Raises ``RuntimeError`` if + another holder is detected. Best-effort on platforms without ``fcntl``. + """ + lock_path = run_dir / ".lock" + try: + import fcntl + except ImportError: + return None + handle = lock_path.open("a+", encoding="utf-8") + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exc: + handle.close() + raise RuntimeError( + f"Another strix process appears to be running on this scan " + f"(could not acquire lock at {lock_path}). Aborting.", + ) from exc + return handle + + +def _release_run_lock(handle: Any) -> None: + if handle is None: + return + try: + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + except (ImportError, OSError): + pass + finally: + with contextlib.suppress(Exception): + handle.close() diff --git a/strix/telemetry/scan_artifacts.py b/strix/telemetry/scan_artifacts.py index 4c25fe7..7409ad0 100644 --- a/strix/telemetry/scan_artifacts.py +++ b/strix/telemetry/scan_artifacts.py @@ -1,19 +1,21 @@ """Per-scan artifact writer. Writes the customer-facing penetration-test report and per-vulnerability -markdown + a ``vulnerabilities.csv`` index under ``strix_runs//``. +markdown + a ``vulnerabilities.csv`` index under ``strix_runs//``, +plus a machine-readable ``vulnerabilities.json`` so a resumed scan's +:class:`~strix.telemetry.tracer.Tracer` can hydrate its in-memory list +back from disk (otherwise vuln-id allocation collides post-restart). """ from __future__ import annotations import csv +import json import logging +import tempfile from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any - - -if TYPE_CHECKING: - from pathlib import Path +from pathlib import Path +from typing import Any logger = logging.getLogger(__name__) @@ -107,6 +109,15 @@ class ScanArtifactWriter: }, ) + # JSON index: machine-readable mirror used by ``Tracer.hydrate_from_run_dir`` + # so a process restart (resume path in ``orchestration/scan.py``) can + # rebuild ``vulnerability_reports`` and re-establish the next id slot + # before any new ``add_vulnerability_report`` call collides on disk. + _atomic_write_text( + self._run_dir / "vulnerabilities.json", + json.dumps(reports, ensure_ascii=False, indent=2, default=str), + ) + if new_reports: logger.info( "Saved %d new vulnerability report(s) to: %s", @@ -116,6 +127,22 @@ class ScanArtifactWriter: logger.info("Updated vulnerability index: %s", csv_path) +def _atomic_write_text(path: Path, payload: str) -> None: + """``tempfile`` + atomic rename so a crash mid-write leaves the prior file.""" + path.parent.mkdir(parents=True, exist_ok=True) + with 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) + + def _render_vulnerability_md(report: dict[str, Any]) -> str: lines: list[str] = [ f"# {report.get('title', 'Untitled Vulnerability')}\n", diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index ac220be..c809212 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -1,3 +1,4 @@ +import json import logging from collections.abc import Callable from datetime import UTC, datetime @@ -93,6 +94,38 @@ class Tracer: return self._run_dir + def hydrate_from_run_dir(self) -> None: + """Reload ``vulnerability_reports`` from ``{run_dir}/vulnerabilities.json``. + + 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. + """ + 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")) + if not isinstance(data, list): + return + 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") + if isinstance(rid, str): + writer._saved_vuln_ids.add(rid) + logger.info( + "tracer hydrated %d vulnerability report(s) from %s", + len(self.vulnerability_reports), + json_path, + ) + except Exception: + logger.exception("tracer hydrate_from_run_dir failed; starting fresh") + def _get_writer(self) -> ScanArtifactWriter: if self._writer is None: self._writer = ScanArtifactWriter(self.get_run_dir()) diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index ba30652..1194725 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -20,10 +20,12 @@ import json import logging import uuid from datetime import UTC, datetime +from pathlib import Path from typing import TYPE_CHECKING, Any, Literal from agents import RunConfig, RunContextWrapper, function_tool from agents.items import TResponseInputItem +from agents.memory import SQLiteSession from agents.model_settings import ModelSettings from agents.sandbox import SandboxRunConfig @@ -454,7 +456,16 @@ async def create_agent( default=str, ) - await bus.register(child_id, name, parent_id) + await bus.register( + child_id, + name, + parent_id, + task=task, + skills=list(skills or []), + is_whitebox=bool(inner.get("is_whitebox", False)), + scan_mode=str(inner.get("scan_mode", "deep")), + diff_scope=inner.get("diff_scope"), + ) # ``ctx.turn_input`` carries the parent's full conversation up to and # including the call that's currently invoking ``create_agent`` @@ -503,13 +514,35 @@ async def create_agent( "agent_finish_called": False, "is_whitebox": bool(inner.get("is_whitebox", False)), "interactive": bool(inner.get("interactive", False)), + "scan_mode": str(inner.get("scan_mode", "deep")), "diff_scope": inner.get("diff_scope"), "run_id": inner.get("run_id"), "agent_factory": factory, # Stashed for ``agent_finish`` to echo back in its completion report. "task": task, + # Inherited so ``create_agent`` calls from this child can also + # add their per-child sessions to the same teardown list. + "_sessions_to_close": inner.get("_sessions_to_close", []), } + # Per-child SQLiteSession at ``{run_dir}/sessions/{child_id}.db`` so + # this subagent's full conversation survives a process restart and + # can be replayed by the SDK on resume. Path is derived from the + # tracer's run_dir (root-side construction in + # :func:`run_strix_scan`); fall back to ``./strix_runs/{run_id}/`` if + # the tracer is absent (unit-test path). + tracer = inner.get("tracer") + if tracer is not None and hasattr(tracer, "get_run_dir"): + child_session_path = tracer.get_run_dir() / "sessions" / f"{child_id}.db" + else: + run_id = inner.get("run_id") or "default" + child_session_path = Path.cwd() / "strix_runs" / str(run_id) / "sessions" / f"{child_id}.db" + child_session_path.parent.mkdir(parents=True, exist_ok=True) + child_session = SQLiteSession(session_id=child_id, db_path=child_session_path) + sessions_list = child_ctx.get("_sessions_to_close") + if isinstance(sessions_list, list): + sessions_list.append(child_session) + child_model_settings = ModelSettings( parallel_tool_calls=False, tool_choice="required", @@ -544,6 +577,7 @@ async def create_agent( bus=bus, agent_id=child_id, interactive=bool(inner.get("interactive", False)), + session=child_session, ), name=f"agent-{name}-{child_id}", )