strix/strix/telemetry/checkpoint.py

"""Checkpoint system for Strix scan resume feature.

Added for Resume Feature - Original behavior is 100% unchanged when
checkpoint_manager is not injected into the agent config.
"""

import hashlib
import json
import logging
import os
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

CHECKPOINT_VERSION = "1.0"


class CheckpointModel(BaseModel):
    """Pydantic model for a full scan checkpoint snapshot.

    Added for Resume Feature.
    """

    version: str = CHECKPOINT_VERSION
    run_name: str
    target_hash: str  # Short SHA-256 of sorted target strings — used for validation
    saved_at: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())

    # Agent loop progress
    iteration: int
    original_max_iterations: int  # The max that was set when the scan started

    # Full AgentState dump (messages, sandbox_id, sandbox_token, etc.)
    agent_state: dict[str, Any]

    # Tracer state needed to restore stats and replay findings
    tracer_chat_messages: list[dict[str, Any]] = Field(default_factory=list)
    tracer_vulnerability_reports: list[dict[str, Any]] = Field(default_factory=list)

    # Added for sub-agent persistence: full agent registry and tool execution log
    # tracer.agents  → every agent (root + sub) with status, name, task, parent_id
    # tracer.tool_executions → every tool call result across all agents
    tracer_agents: dict[str, Any] = Field(default_factory=dict)
    tracer_tool_executions: dict[str, Any] = Field(default_factory=dict)  # key is str(int)
    tracer_next_execution_id: int = 1  # restore the ID counter so new IDs don't collide

    # Original scan configuration (passed to execute_scan)
    scan_config: dict[str, Any] = Field(default_factory=dict)


def compute_target_hash(targets_info: list[dict[str, Any]]) -> str:
    """Return a short stable hash of the target list for checkpoint validation.

    Added for Resume Feature.
    """
    target_strings = sorted(t.get("original", "") for t in (targets_info or []))
    combined = "|".join(target_strings)
    return hashlib.sha256(combined.encode()).hexdigest()[:16]


class CheckpointManager:
    """Saves and loads scan checkpoints to ``strix_runs/<run_name>/checkpoint.json``.

    All operations are *non-fatal*: any I/O error is logged as a warning and
    the scan continues normally.

    Added for Resume Feature.
    """

    def __init__(self, run_name: str, run_dir: Path) -> None:
        self.run_name = run_name
        self.run_dir = run_dir
        self.checkpoint_path = run_dir / "checkpoint.json"
        self._tmp_path = run_dir / "checkpoint.json.tmp"

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def exists(self) -> bool:
        """Return True if a checkpoint file is present."""
        return self.checkpoint_path.exists()

    def save(
        self,
        agent_state: Any,
        tracer: Any | None,
        scan_config: dict[str, Any],
        target_hash: str,
        original_max_iterations: int,
    ) -> None:
        """Atomically persist the current scan state.

        Writes to a ``.tmp`` file then renames to prevent corruption during a
        crash mid-write.  All errors are non-fatal (warning only).
        """
        try:
            self.run_dir.mkdir(parents=True, exist_ok=True)

            state_dict: dict[str, Any] = (
                agent_state.model_dump() if hasattr(agent_state, "model_dump") else {}
            )

            tracer_chat_messages: list[dict[str, Any]] = []
            tracer_vulnerability_reports: list[dict[str, Any]] = []
            tracer_agents: dict[str, Any] = {}
            tracer_tool_executions: dict[str, Any] = {}
            tracer_next_execution_id: int = 1
            if tracer:
                tracer_chat_messages = list(getattr(tracer, "chat_messages", []))
                tracer_vulnerability_reports = list(
                    getattr(tracer, "vulnerability_reports", [])
                )
                # Added for sub-agent persistence — capture full agent registry
                # and all tool execution records so sub-agents are fully restored
                tracer_agents = dict(getattr(tracer, "agents", {}))
                # tool_executions keys are ints; serialise as strings for JSON
                raw_execs = getattr(tracer, "tool_executions", {})
                tracer_tool_executions = {str(k): v for k, v in raw_execs.items()}
                tracer_next_execution_id = getattr(tracer, "_next_execution_id", 1)

            checkpoint = CheckpointModel(
                run_name=self.run_name,
                target_hash=target_hash,
                iteration=agent_state.iteration,
                original_max_iterations=original_max_iterations,
                agent_state=state_dict,
                tracer_chat_messages=tracer_chat_messages,
                tracer_vulnerability_reports=tracer_vulnerability_reports,
                tracer_agents=tracer_agents,
                tracer_tool_executions=tracer_tool_executions,
                tracer_next_execution_id=tracer_next_execution_id,
                scan_config=scan_config,
            )

            # Atomic write: .tmp → rename
            self._tmp_path.write_text(checkpoint.model_dump_json(indent=2), encoding="utf-8")
            os.rename(self._tmp_path, self.checkpoint_path)

        except Exception as e:  # noqa: BLE001
            logger.warning("[Resume] Checkpoint save failed (non-fatal): %s", e)

    def load(self) -> "CheckpointModel | None":
        """Load and parse the checkpoint file.

        Returns ``None`` and logs a warning on any error (corruption, missing
        file, version mismatch).
        """
        if not self.checkpoint_path.exists():
            return None
        try:
            raw = json.loads(self.checkpoint_path.read_text(encoding="utf-8"))
            return CheckpointModel.model_validate(raw)
        except Exception as e:  # noqa: BLE001
            logger.warning("[Resume] Checkpoint load failed: %s", e)
            return None

    def delete(self) -> None:
        """Remove the checkpoint file (called when the scan finishes cleanly)."""
        try:
            if self.checkpoint_path.exists():
                self.checkpoint_path.unlink()
            if self._tmp_path.exists():
                self._tmp_path.unlink()
        except Exception as e:  # noqa: BLE001
            logger.warning("[Resume] Checkpoint delete failed (non-fatal): %s", e)