diff --git a/clearwing/llm/budget.py b/clearwing/llm/budget.py index 2d48b6e..7e5df33 100644 --- a/clearwing/llm/budget.py +++ b/clearwing/llm/budget.py @@ -130,9 +130,14 @@ def __init__( default_max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, manifest_filename: str = "manifest.json", endpoint: LLMEndpoint | None = None, + initial_spent_usd: float = 0.0, ) -> None: if not math.isfinite(limit_usd) or limit_usd < 0: raise BudgetConfigurationError("LLM budget must be a finite value >= 0") + if not math.isfinite(initial_spent_usd) or initial_spent_usd < 0: + raise BudgetConfigurationError( + "initial_spent_usd must be a finite value >= 0" + ) if (input_price_per_million is None) != (output_price_per_million is None): raise BudgetConfigurationError( "input and output token prices must be provided together" @@ -173,7 +178,9 @@ def __init__( self._endpoint_pricing = _endpoint_pricing(endpoint) self._lock = threading.RLock() - self._spent_usd = 0.0 + # Carried-forward spend from a resumed session's prior spend-ledger, so + # `remaining`/`exhausted` honor the original cap from the first call. + self._spent_usd = float(initial_spent_usd) self._reserved_usd = 0.0 self._input_tokens = 0 self._output_tokens = 0 @@ -193,6 +200,7 @@ def __init__( "event": "run_started", "session_id": self.session_id, "budget_usd": self.limit_usd, + "carried_forward_usd": self._spent_usd, "timestamp": self._timestamp(), } ) diff --git a/clearwing/sourcehunt/checkpoints.py b/clearwing/sourcehunt/checkpoints.py new file mode 100644 index 0000000..89abce3 --- /dev/null +++ b/clearwing/sourcehunt/checkpoints.py @@ -0,0 +1,191 @@ +"""Per-stage checkpoints for a single SourceHuntRunner run. + +Writes a full, human-readable snapshot of the findings at each natural pipeline +phase boundary — end of hunt, end of verify, end of exploit — under +``/checkpoints/.json``. Each file is a superset of +``findings.json`` (``json.load``-simple) so it doubles as an eval artifact +showing how findings evolve across phases. + +Serialization is LOSSLESS: findings are dataclasses (``clearwing.findings.types.Finding``) +serialized via ``dataclasses.asdict`` — unlike ``FindingsPool``'s deliberately +lossy discovery-time subset, this keeps every verify/exploit/patch field. Reload +uses the same tolerant ``__dataclass_fields__`` filter the pool uses. + +These checkpoints power ``SourceHuntRunner(resume_session=...)``: a crashed run +resumes from the last completed phase instead of re-hunting from scratch. +""" +from __future__ import annotations + +import json +import logging +import os +import tempfile +from dataclasses import asdict +from pathlib import Path +from typing import Any + +from clearwing.findings.types import Finding + +logger = logging.getLogger(__name__) + +# Ordered most- to least-complete. `latest_checkpoint` returns the first present. +STAGES = ("exploit", "verify", "hunt") + +_CHECKPOINT_DIRNAME = "checkpoints" +_LEDGER_FILENAME = "spend-ledger.jsonl" + + +def _checkpoint_dir(session_dir: Path) -> Path: + return Path(session_dir) / _CHECKPOINT_DIRNAME + + +def _finding_to_dict(finding: Any) -> dict[str, Any]: + """Losslessly serialize a Finding dataclass (or pass through a plain dict).""" + if isinstance(finding, Finding): + return asdict(finding) + if isinstance(finding, dict): + return dict(finding) + # Best effort for anything dataclass-like. + if hasattr(finding, "__dict__"): + return dict(finding.__dict__) + raise TypeError(f"Cannot serialize finding of type {type(finding)!r}") + + +def _finding_from_dict(data: dict[str, Any]) -> Finding: + """Rehydrate a Finding, ignoring unknown keys (tolerant, matches FindingsPool).""" + return Finding(**{ + k: v for k, v in data.items() + if k in Finding.__dataclass_fields__ + }) + + +def write_stage_checkpoint( + session_dir: Path | str, + stage: str, + *, + findings: list[Any], + budget_spent_usd: float, + verified: list[Any] | None = None, + rejected: list[Any] | None = None, + exploited: list[Any] | None = None, + session_id: str | None = None, +) -> Path: + """Atomically write a stage checkpoint to ``checkpoints/.json``. + + Only the finding collections relevant to the stage are included: + - hunt: findings + - verify: findings + verified + rejected + - exploit: findings + verified + rejected + exploited + + Returns the path written. + """ + if stage not in STAGES: + raise ValueError(f"unknown stage {stage!r}; expected one of {STAGES}") + + ckpt_dir = _checkpoint_dir(Path(session_dir)) + ckpt_dir.mkdir(parents=True, exist_ok=True) + target = ckpt_dir / f"{stage}.json" + + payload: dict[str, Any] = { + "stage": stage, + "session_id": session_id or Path(session_dir).name, + "budget_spent_usd": float(budget_spent_usd), + "findings": [_finding_to_dict(f) for f in findings], + } + if verified is not None: + payload["verified"] = [_finding_to_dict(f) for f in verified] + if rejected is not None: + payload["rejected"] = [_finding_to_dict(f) for f in rejected] + if exploited is not None: + payload["exploited"] = [_finding_to_dict(f) for f in exploited] + + # Atomic write: temp file in the same dir + os.replace (mirrors campaign.py). + fd, tmp_path = tempfile.mkstemp(dir=str(ckpt_dir), suffix=".tmp") + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, sort_keys=True) + os.replace(tmp_path, str(target)) + except Exception: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + return target + + +def load_stage_checkpoint( + session_dir: Path | str, + stage: str, +) -> dict[str, Any] | None: + """Load a stage checkpoint, rehydrating finding lists into Finding objects. + + Returns a dict with keys ``stage``, ``session_id``, ``budget_spent_usd``, + ``findings`` (list[Finding]) and, for verify/exploit stages, ``verified`` / + ``rejected`` / ``exploited`` (list[Finding]). Returns None if absent. + """ + target = _checkpoint_dir(Path(session_dir)) / f"{stage}.json" + if not target.exists(): + return None + data = json.loads(target.read_text(encoding="utf-8")) + + out: dict[str, Any] = { + "stage": data.get("stage", stage), + "session_id": data.get("session_id", ""), + "budget_spent_usd": float(data.get("budget_spent_usd", 0.0)), + "findings": [_finding_from_dict(d) for d in data.get("findings", [])], + } + for key in ("verified", "rejected", "exploited"): + if key in data: + out[key] = [_finding_from_dict(d) for d in data[key]] + return out + + +def latest_checkpoint(session_dir: Path | str) -> str | None: + """Return the most-complete checkpoint stage present, or None. + + Priority: exploit > verify > hunt (a later phase implies the earlier ones ran). + """ + ckpt_dir = _checkpoint_dir(Path(session_dir)) + for stage in STAGES: + if (ckpt_dir / f"{stage}.json").exists(): + return stage + return None + + +def sum_prior_spend(session_dir: Path | str) -> float: + """Reconstruct total settled spend from an existing spend-ledger.jsonl. + + Sums ``cost_usd`` over every ``call_settled`` event. Tolerant of malformed + lines. Returns 0.0 if the ledger is absent. + """ + ledger = Path(session_dir) / _LEDGER_FILENAME + if not ledger.exists(): + return 0.0 + total = 0.0 + for line in ledger.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + if event.get("event") == "call_settled": + try: + total += float(event.get("cost_usd", 0.0) or 0.0) + except (TypeError, ValueError): + continue + return total + + +def resolve_session_dir(resume_session: str, default_output_dir: str | Path) -> Path: + """Resolve a ``resume_session`` argument to a concrete session directory. + + Accepts either a bare session name (``sh-`` → ``/``) + or a path (absolute or relative) that points directly at a session dir. + """ + candidate = Path(resume_session) + # A path-like argument (contains a separator or exists as given) is used directly. + if candidate.parent != Path(".") or candidate.exists(): + return candidate + return Path(default_output_dir) / resume_session diff --git a/clearwing/sourcehunt/runner.py b/clearwing/sourcehunt/runner.py index 92659f7..dc72e30 100644 --- a/clearwing/sourcehunt/runner.py +++ b/clearwing/sourcehunt/runner.py @@ -33,6 +33,13 @@ ) from ..sandbox.hunter_sandbox import HunterSandbox +from .checkpoints import ( + latest_checkpoint, + load_stage_checkpoint, + resolve_session_dir, + sum_prior_spend, + write_stage_checkpoint, +) from .config import SourceHuntConfig from .disclosure import ( DisclosureGenerator, @@ -275,6 +282,7 @@ def __init__( emit_rejection_certificates: bool = True, falsify: bool = True, on_progress: SourceHuntProgressCallback | None = None, + resume_session: str | None = None, ): # --- Resolve from SourceHuntConfig when provided ---------------------- if config is not None: @@ -511,7 +519,30 @@ def __init__( self.sandbox_factory = sandbox_factory self._sandbox_manager: HunterSandbox | None = None self._preprocessor: Preprocessor | None = None - self._session_id = parent_session_id or f"sh-{uuid.uuid4().hex[:8]}" + # --- Resume wiring --------------------------------------------------- + # When resuming, adopt the prior run's session id BEFORE the + # instrumentation (and every output_dir/session_id/... path) is built, + # so findings_pool, trajectories, spend-ledger, and checkpoints all + # resolve into the resumed session directory. Resume is deliberate: + # if no checkpoint is present, we raise rather than silently start over. + self._resume_session = resume_session + self._resume_from: str | None = None + if resume_session: + resumed_dir = resolve_session_dir(resume_session, self.output_dir) + self._resume_from = latest_checkpoint(resumed_dir) + if self._resume_from is None: + raise ValueError( + f"resume_session={resume_session!r} resolved to {resumed_dir} " + "but no stage checkpoint was found there " + "(expected checkpoints/hunt.json, verify.json, or exploit.json). " + "Resume is only for continuing an existing run." + ) + # Anchor all session paths on the resumed dir: parent = output_dir, + # name = session id. This works for both bare-name and path forms. + self.output_dir = str(resumed_dir.parent) + self._session_id = resumed_dir.name + else: + self._session_id = parent_session_id or f"sh-{uuid.uuid4().hex[:8]}" self._agent_mode_override = agent_mode self._prompt_mode = prompt_mode self._campaign_hint = campaign_hint @@ -643,6 +674,36 @@ def _run_spent_usd(self) -> float: return 0.0 return self._spend_ledger.spent_usd + def _write_stage_checkpoint_safe( + self, + stage: str, + *, + findings: list[Finding], + verified: list[Finding] | None = None, + rejected: list[Finding] | None = None, + exploited: list[Finding] | None = None, + ) -> None: + """Write a stage checkpoint, swallowing any error. + + A checkpoint is an optional durability aid — a failure to write one + (disk full, permissions, serialization edge case) must never crash an + otherwise-successful run. + """ + try: + path = write_stage_checkpoint( + Path(self.output_dir) / self._session_id, + stage, + findings=findings, + budget_spent_usd=self._run_spent_usd(), + verified=verified, + rejected=rejected, + exploited=exploited, + session_id=self._session_id, + ) + logger.info("Wrote %s checkpoint: %s", stage, path) + except Exception: + logger.warning("Failed to write %s checkpoint", stage, exc_info=True) + @property def _shard_entry_points(self) -> bool: if self._shard_entry_points_override is not None: @@ -720,6 +781,14 @@ def _finalize_instrumentation(self, status: str) -> None: def _ensure_spend_ledger(self) -> SpendLedger: if self._spend_ledger is None: + # On resume, carry forward the prior run's settled spend so the + # original cap is honored across the resumed continuation. The + # existing ledger is appended to (not truncated) by SpendLedger. + initial_spent = 0.0 + if self._resume_session: + initial_spent = sum_prior_spend( + Path(self.output_dir) / self._session_id + ) self._spend_ledger = SpendLedger( limit_usd=self.budget_usd, session_id=self._session_id, @@ -730,6 +799,7 @@ def _ensure_spend_ledger(self) -> SpendLedger: manifest_filename=( "spend-summary.json" if self._flow == "proof" else "manifest.json" ), + initial_spent_usd=initial_spent, ) return self._spend_ledger @@ -737,18 +807,30 @@ def _budget_exhausted(self) -> bool: return self._spend_ledger is not None and self._spend_ledger.exhausted def _preflight_budget_clients(self) -> None: - """Validate every model this run may use before the first paid call.""" + """Validate every model this run may use before the first paid call. + + On resume, phases already completed in the prior run are skipped, so we + don't validate their models — a resume-from-exploit run must not fail + just because the hunter/verifier model is no longer configured. + """ if self._spend_ledger is None or not self._spend_ledger.enforcing: return + skip_rank = self._no_rank or bool(self._resume_from) + skip_hunt = self._resume_from in ("hunt", "verify", "exploit") + skip_verify = self._resume_from in ("verify", "exploit") + skip_exploit = self._resume_from == "exploit" roles: list[tuple[str, AsyncLLMClient | None, str]] = [] - if not self._no_rank: + if not skip_rank: roles.append(("ranker", self.ranker_llm, "rank")) if self.depth != "quick": - roles.append(("hunter", self.hunter_llm, "hunt")) - if not self.no_verify: + if not skip_hunt: + roles.append(("hunter", self.hunter_llm, "hunt")) + if not self.no_verify and not skip_verify: roles.append(("verifier", self.verifier_llm, "verify")) - if not self.no_exploit or self.enable_auto_patch or self.enable_elaboration: + if ( + not self.no_exploit or self.enable_auto_patch or self.enable_elaboration + ) and not skip_exploit: roles.append(("sourcehunt_exploit", self.exploiter_llm, "exploit")) for task, override, stage in roles: self._get_native_client(task, override, budget_stage=stage) @@ -1012,6 +1094,35 @@ async def arun(self) -> SourceHuntResult: self._ensure_output_dir_layout() self._ensure_spend_ledger() pipeline_status = PipelineStatus() + + # --- Resume: load the highest checkpoint and derive phase skip flags. + # Preprocess always re-runs (rebuilds repo checkout / callgraph / sandbox + # that verify+exploit need); rank is skipped (its priority fields are + # only consumed by the hunt phase). Completed phases are skipped and + # their findings are seeded from the checkpoint. + resume_ckpt: dict | None = None + skip_hunt = skip_verify = skip_exploit = False + if self._resume_from: + resume_ckpt = load_stage_checkpoint( + Path(self.output_dir) / self._session_id, self._resume_from + ) + if resume_ckpt is None: + raise RuntimeError( + f"resume checkpoint {self._resume_from!r} vanished for session " + f"{self._session_id}" + ) + skip_hunt = self._resume_from in ("hunt", "verify", "exploit") + skip_verify = self._resume_from in ("verify", "exploit") + skip_exploit = self._resume_from == "exploit" + logger.info( + "Resuming session %s from %r checkpoint (skip hunt=%s verify=%s exploit=%s)", + self._session_id, + self._resume_from, + skip_hunt, + skip_verify, + skip_exploit, + ) + logger.info("Sourcehunt session %s starting on %s", self._session_id, self.repo_url) self._instrumentation.record( "run", @@ -1037,10 +1148,15 @@ async def arun(self) -> SourceHuntResult: ) self._ensure_sandbox_factory(repo_path, files) - # 2. Rank — unless depth=quick AND no LLM available, or --no-rank + # 2. Rank — unless depth=quick AND no LLM available, or --no-rank. + # On resume rank is skipped entirely: its priority fields are only + # consumed by the (now-skipped) hunt phase, so re-ranking would just + # burn budget. Force ranker_llm to None so the --no-rank heuristic + # path assigns default scores without any LLM call. + resume_skips_rank = bool(self._resume_from) ranker_llm = ( None - if self._no_rank + if (self._no_rank or resume_skips_rank) else self._get_native_client( "ranker", self.ranker_llm, @@ -1165,10 +1281,13 @@ async def arun(self) -> SourceHuntResult: # 2.5. Harness Generator (crash-first ordering) — at depth=deep or # when seed_harness_crashes is explicitly enabled (spec 018). + # Feeds the hunt only; skipped on resume (hunt already done). seeded_crashes: list[SeededCrash] = [] if ( - self.depth == "deep" or self._seed_harness_crashes - ) and not self._budget_exhausted(): + (self.depth == "deep" or self._seed_harness_crashes) + and not self._budget_exhausted() + and not skip_hunt + ): harness_llm = self._get_native_client( "hunter", self.hunter_llm, @@ -1294,13 +1413,20 @@ async def arun(self) -> SourceHuntResult: logger.warning("Historical findings DB load failed", exc_info=True) historical_db = None - # 3. Tiered hunt - hunter_llm = self._get_native_client( - "hunter", - self.hunter_llm, - budget_stage="hunt", + # 3. Tiered hunt. On resume the hunt+subsystem phases are skipped, so + # we don't resolve (and can't require) the hunter model. + hunter_llm = ( + None + if skip_hunt + else self._get_native_client( + "hunter", + self.hunter_llm, + budget_stage="hunt", + ) + ) + all_findings: list[Finding] = ( + list(resume_ckpt["findings"]) if skip_hunt else [] ) - all_findings: list[Finding] = [] files_hunted = 0 spent_per_tier: dict[str, float] = {"A": 0.0, "B": 0.0, "C": 0.0} band_stats: dict | None = None @@ -1313,7 +1439,22 @@ async def arun(self) -> SourceHuntResult: } ) - if self._no_per_file_hunt: + if skip_hunt: + logger.info( + "Hunt skipped (resumed from %r checkpoint): %d findings loaded", + self._resume_from, + len(all_findings), + ) + self._emit_stage( + "hunt", + "completed", + findings_so_far=len(all_findings), + detail=f"Resumed: {len(all_findings)} findings from checkpoint", + files=[str(finding.file or "") for finding in all_findings], + symbols=self._finding_symbols(all_findings), + finding_ids=[finding.id for finding in all_findings], + ) + elif self._no_per_file_hunt: logger.info("Per-file hunt skipped (--no-per-file-hunt)") self._emit_stage( "hunt", @@ -1452,7 +1593,7 @@ async def arun(self) -> SourceHuntResult: ) # 3.5. v0.6: Behavioral monitoring of findings text (spec 013). - if self._enable_behavior_monitor and all_findings: + if self._enable_behavior_monitor and all_findings and not skip_hunt: try: from .behavior_monitor import BehaviorMonitor @@ -1473,12 +1614,20 @@ async def arun(self) -> SourceHuntResult: logger.debug("Behavior monitor failed", exc_info=True) # Promote static findings into the all_findings list so depth=quick - # output is still useful when no hunter llm is available - all_findings = self._merge_static_findings(all_findings, preprocess_result) + # output is still useful when no hunter llm is available. On resume the + # checkpoint already contains the merged static findings — re-merging + # would duplicate them (each gets a fresh random id), so skip it. + if not skip_hunt: + all_findings = self._merge_static_findings(all_findings, preprocess_result) # 3.5. Persist findings to historical DB (spec 005) # Skip when running under campaign — campaign handles bulk ingestion. - if historical_db is not None and all_findings and self._injected_findings_pool is None: + if ( + historical_db is not None + and all_findings + and self._injected_findings_pool is None + and not skip_hunt + ): try: count = historical_db.ingest_campaign( all_findings, @@ -1494,13 +1643,18 @@ async def arun(self) -> SourceHuntResult: # 3.7. Subsystem hunt (spec 006) subsystems_hunted = 0 subsystem_spent = 0.0 - if self._enable_subsystem_hunt and hunter_llm is not None and self._budget_exhausted(): + if ( + self._enable_subsystem_hunt + and hunter_llm is not None + and self._budget_exhausted() + and not skip_hunt + ): logger.info( "Subsystem hunt skipped: budget $%.2f exhausted ($%.2f spent)", self.budget_usd, self._run_spent_usd(), ) - elif self._enable_subsystem_hunt and hunter_llm is not None: + elif self._enable_subsystem_hunt and hunter_llm is not None and not skip_hunt: from .subsystem import ( SubsystemHuntConfig, identify_subsystems_auto, @@ -1638,9 +1792,20 @@ async def arun(self) -> SourceHuntResult: error={"type": type(exc).__name__, "message": str(exc)}, ) + # Checkpoint the end of the hunt phase (preprocess + rank + per-file + # hunt + static-merge + subsystem hunt all resolved). A crashed run + # resumes verify from here. Never on resume (already present); never + # fatal (a checkpoint failure must not sink the run). + if not skip_hunt: + self._write_stage_checkpoint_safe("hunt", findings=all_findings) + # 4. Verify (unless --no-verify) - verified: list[Finding] = [] - rejected: list[Finding] = [] + verified: list[Finding] = ( + list(resume_ckpt.get("verified", [])) if skip_verify else [] + ) + rejected: list[Finding] = ( + list(resume_ckpt.get("rejected", [])) if skip_verify else [] + ) verify_status = "completed" self._emit_stage( "verify", @@ -1651,7 +1816,15 @@ async def arun(self) -> SourceHuntResult: symbols=self._finding_symbols(all_findings), finding_ids=[finding.id for finding in all_findings], ) - if not self.no_verify: + if skip_verify: + logger.info( + "Verify skipped (resumed from %r checkpoint): " + "%d verified, %d rejected loaded", + self._resume_from, + len(verified), + len(rejected), + ) + elif not self.no_verify: if self._budget_exhausted(): verify_status = "budget_exhausted" pipeline_status.record( @@ -1716,7 +1889,12 @@ async def arun(self) -> SourceHuntResult: # 4.5. v0.3: Extract mechanisms from verified findings and persist them # to the cross-run store. Cheap LLM pass; failures are non-fatal. - if self._mechanism_store is not None and verified and not self._budget_exhausted(): + if ( + self._mechanism_store is not None + and verified + and not self._budget_exhausted() + and not skip_verify + ): verifier_llm_for_extract = self._get_native_client( "verifier", self.verifier_llm, @@ -1744,7 +1922,12 @@ async def arun(self) -> SourceHuntResult: # match as a new suspicion-level finding linked back to the # original. v0.3 scope: we surface the matches in the report; # we don't re-spawn hunters on each match (that's a v1.0 pass). - if self.enable_variant_loop and verified and not self._budget_exhausted(): + if ( + self.enable_variant_loop + and verified + and not self._budget_exhausted() + and not skip_verify + ): variant_llm = self._get_native_client( "verifier", self.verifier_llm, @@ -1823,6 +2006,7 @@ async def arun(self) -> SourceHuntResult: and verified and self._sandbox_manager is not None and not self._budget_exhausted() + and not skip_verify ): from .stability import StabilityVerifier, apply_stability_result @@ -1871,6 +2055,17 @@ async def arun(self) -> SourceHuntResult: non_poc = [f for f in verified if f not in stability_eligible] verified = stable_verified + non_poc + # Checkpoint the end of the verify phase (verify + mechanism extraction + # + variant loop + stability all resolved). A crashed run resumes + # exploit from here. + if not skip_verify: + self._write_stage_checkpoint_safe( + "verify", + findings=all_findings, + verified=verified, + rejected=rejected, + ) + # 5. Exploit-triage (unless --no-exploit) — gated on evidence_level self._emit_stage( "exploit", @@ -1880,11 +2075,13 @@ async def arun(self) -> SourceHuntResult: symbols=self._finding_symbols(verified), finding_ids=[finding.id for finding in verified], ) - exploited: list[Finding] = [] + exploited: list[Finding] = ( + list(resume_ckpt.get("exploited", [])) if skip_exploit else [] + ) # 5.5 v0.3: Auto-patch (opt-in) — runs after exploiter on verified # critical/high findings with root_cause_explained evidence. patched: list[Finding] = [] - if not self.no_exploit and not self._budget_exhausted(): + if not self.no_exploit and not self._budget_exhausted() and not skip_exploit: exploiter_llm = self._get_native_client( "sourcehunt_exploit", self.exploiter_llm, @@ -1951,7 +2148,12 @@ async def arun(self) -> SourceHuntResult: # 5.25. Stage 1.5: Exploit elaboration (autonomous, opt-in). elaborated: list[Finding] = [] - if self.enable_elaboration and exploited and not self._budget_exhausted(): + if ( + self.enable_elaboration + and exploited + and not self._budget_exhausted() + and not skip_exploit + ): from .elaboration import ( ElaborationAgent, prioritize_for_elaboration, @@ -2009,7 +2211,12 @@ async def arun(self) -> SourceHuntResult: # 5.5. v0.3: Auto-patch mode (opt-in). # The verify-by-recompile gate is MANDATORY — a patch is only marked # `validated` if we actually applied it, rebuilt, and re-ran the PoC. - if self.enable_auto_patch and verified and not self._budget_exhausted(): + if ( + self.enable_auto_patch + and verified + and not self._budget_exhausted() + and not skip_exploit + ): patcher_llm = self._get_native_client( "sourcehunt_exploit", self.exploiter_llm, @@ -2058,15 +2265,17 @@ async def arun(self) -> SourceHuntResult: logger.warning("Auto-patcher failed", exc_info=True) # 5.75. v0.3: Populate the cross-run knowledge graph with source - # findings. Best-effort — never blocks the run. + # findings. Best-effort — never blocks the run. On resume from + # the exploit checkpoint these cross-run side effects already + # fired in the original run; skip them and only re-run report. try: - if self.enable_knowledge_graph and all_findings: + if self.enable_knowledge_graph and all_findings and not skip_exploit: self._populate_knowledge_graph_source(repo_path, all_findings) except Exception: logger.warning("Knowledge graph population failed", exc_info=True) # 5.85. v0.4: Coordinated-disclosure templates (opt-in). - if self.export_disclosures and verified: + if self.export_disclosures and verified and not skip_exploit: try: self._export_disclosure_bundle(verified) except Exception: @@ -2089,7 +2298,7 @@ async def arun(self) -> SourceHuntResult: logger.warning("Disclosure DB queue failed", exc_info=True) # 5.87. v0.6: Store exploits in encrypted artifact store (spec 013). - if self._enable_artifact_store and exploited: + if self._enable_artifact_store and exploited and not skip_exploit: try: from .artifact_store import ArtifactStore @@ -2111,7 +2320,11 @@ async def arun(self) -> SourceHuntResult: try: from .commitment import CommitmentLog - committable = filter_by_evidence(verified, "root_cause_explained") + committable = ( + filter_by_evidence(verified, "root_cause_explained") + if not skip_exploit + else [] + ) if committable: commitment_log = CommitmentLog() for f in committable: @@ -2139,6 +2352,17 @@ async def arun(self) -> SourceHuntResult: finding_ids=[finding.id for finding in all_findings], ) + # Checkpoint the end of the exploit phase. A crashed run resumes at + # report from here (report is cheap and always re-runs on resume). + if not skip_exploit: + self._write_stage_checkpoint_safe( + "exploit", + findings=all_findings, + verified=verified, + rejected=rejected, + exploited=exploited, + ) + # 6. Report self._emit_stage( "report", diff --git a/tests/test_sourcehunt_checkpoints.py b/tests/test_sourcehunt_checkpoints.py new file mode 100644 index 0000000..a7234de --- /dev/null +++ b/tests/test_sourcehunt_checkpoints.py @@ -0,0 +1,495 @@ +"""Tests for SourceHuntRunner stage checkpoints and budget carry-forward. + +Covers the standalone checkpoint module (lossless finding round-trip, latest +checkpoint priority, prior-spend reconstruction, session-dir resolution) and the +``SpendLedger`` seeding that lets a resumed run honor the original dollar cap. +""" +from __future__ import annotations + +import json +from dataclasses import asdict +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest +from genai_pyo3 import ChatResponse + +from clearwing.findings.types import Finding +from clearwing.llm.budget import BudgetConfigurationError, SpendLedger +from clearwing.sourcehunt.checkpoints import ( + latest_checkpoint, + load_stage_checkpoint, + resolve_session_dir, + sum_prior_spend, + write_stage_checkpoint, +) +from clearwing.sourcehunt.runner import SourceHuntRunner + +FIXTURE_C_PROPAGATION = ( + Path(__file__).parent / "fixtures" / "vuln_samples" / "c_propagation" +) + + +def _make_ranker_llm() -> AsyncMock: + llm = AsyncMock() + llm.aask_json.return_value = ({"results": []}, ChatResponse()) + return llm + + +def _make_hunter_llm() -> MagicMock: + llm = MagicMock() + bound = MagicMock() + response = MagicMock() + response.content = "No vulnerabilities found." + response.tool_calls = [] + bound.invoke.return_value = response + llm.bind_tools.return_value = bound + return llm + + +def _make_verifier_llm() -> AsyncMock: + llm = AsyncMock() + llm.aask_text.return_value = ChatResponse( + content=[{"text": json.dumps({"is_real": True, "severity": "high"})}] + ) + return llm + + +def _fully_populated_finding() -> Finding: + """A finding carrying hunter + verify + exploit + nested-dict fields.""" + return Finding( + id="find-001", + finding_type="sql-injection", + cwe="CWE-89", + file="src/db.py", + line_number=42, + end_line=48, + code_snippet="cursor.execute(q)", + severity="high", + severity_verified="critical", + confidence="high", + description="Unsanitized query", + crash_evidence="stack trace here", + poc="' OR 1=1 --", + discovered_by="hunter", + evidence_level="crash_reproduced", + related_cve="CVE-2026-9999", + primitive_type="oob_write", + verified=True, + verifier_pro_argument="reaches sink", + verifier_counter_argument="input is validated upstream", + verifier_tie_breaker="not on this path", + patch_oracle_passed=True, + exploit="payload", + exploit_success=True, + auto_patch="--- a/db.py", + auto_patch_validated=True, + hunter_session_id="sh-abc", + verifier_session_id="sh-def", + crypto_evidence={"nested": {"a": 1}, "list": [1, 2, 3]}, + extra={"exploit_cost_usd": 0.12, "stable_finding_id": "stable-1"}, + ) + + +# --- checkpoint round-trip -------------------------------------------------- + + +def test_write_then_load_is_lossless(tmp_path): + finding = _fully_populated_finding() + write_stage_checkpoint( + tmp_path, + "hunt", + findings=[finding], + budget_spent_usd=1.23, + ) + + loaded = load_stage_checkpoint(tmp_path, "hunt") + assert loaded is not None + assert loaded["stage"] == "hunt" + assert loaded["session_id"] == tmp_path.name + assert loaded["budget_spent_usd"] == pytest.approx(1.23) + assert len(loaded["findings"]) == 1 + + # No field loss: the round-tripped dataclass equals the original. + assert asdict(loaded["findings"][0]) == asdict(finding) + + +def test_checkpoint_file_is_plain_json(tmp_path): + write_stage_checkpoint( + tmp_path, + "hunt", + findings=[_fully_populated_finding()], + budget_spent_usd=0.5, + ) + path = tmp_path / "checkpoints" / "hunt.json" + assert path.exists() + # json.load-simple, as required for eval transparency. + data = json.loads(path.read_text()) + assert data["stage"] == "hunt" + assert isinstance(data["findings"], list) + assert data["findings"][0]["id"] == "find-001" + + +def test_verify_stage_carries_verified_and_rejected(tmp_path): + good = Finding(id="ok", verified=True) + bad = Finding(id="no", verified=False) + write_stage_checkpoint( + tmp_path, + "verify", + findings=[good, bad], + verified=[good], + rejected=[bad], + budget_spent_usd=2.0, + ) + loaded = load_stage_checkpoint(tmp_path, "verify") + assert [f.id for f in loaded["verified"]] == ["ok"] + assert [f.id for f in loaded["rejected"]] == ["no"] + + +def test_exploit_stage_carries_exploited(tmp_path): + f = Finding(id="e", exploit_success=True) + write_stage_checkpoint( + tmp_path, + "exploit", + findings=[f], + verified=[f], + rejected=[], + exploited=[f], + budget_spent_usd=3.0, + ) + loaded = load_stage_checkpoint(tmp_path, "exploit") + assert [f.id for f in loaded["exploited"]] == ["e"] + assert loaded["rejected"] == [] + + +def test_unknown_stage_rejected(tmp_path): + with pytest.raises(ValueError): + write_stage_checkpoint( + tmp_path, "bogus", findings=[], budget_spent_usd=0.0 + ) + + +def test_load_absent_checkpoint_returns_none(tmp_path): + assert load_stage_checkpoint(tmp_path, "hunt") is None + + +def test_reload_ignores_unknown_keys(tmp_path): + ckpt_dir = tmp_path / "checkpoints" + ckpt_dir.mkdir() + (ckpt_dir / "hunt.json").write_text( + json.dumps( + { + "stage": "hunt", + "session_id": "sh-x", + "budget_spent_usd": 0.0, + "findings": [{"id": "f1", "not_a_real_field": "ignore me"}], + } + ) + ) + loaded = load_stage_checkpoint(tmp_path, "hunt") + assert loaded["findings"][0].id == "f1" + assert not hasattr(loaded["findings"][0], "not_a_real_field") + + +def test_write_is_atomic_no_tmp_left_behind(tmp_path): + write_stage_checkpoint( + tmp_path, "hunt", findings=[Finding(id="a")], budget_spent_usd=0.0 + ) + ckpt_dir = tmp_path / "checkpoints" + leftovers = [p for p in ckpt_dir.iterdir() if p.suffix == ".tmp"] + assert leftovers == [] + + +# --- latest_checkpoint priority --------------------------------------------- + + +def test_latest_checkpoint_priority_order(tmp_path): + assert latest_checkpoint(tmp_path) is None + + write_stage_checkpoint(tmp_path, "hunt", findings=[], budget_spent_usd=0.0) + assert latest_checkpoint(tmp_path) == "hunt" + + write_stage_checkpoint( + tmp_path, "verify", findings=[], verified=[], rejected=[], budget_spent_usd=0.0 + ) + assert latest_checkpoint(tmp_path) == "verify" + + write_stage_checkpoint( + tmp_path, + "exploit", + findings=[], + verified=[], + rejected=[], + exploited=[], + budget_spent_usd=0.0, + ) + assert latest_checkpoint(tmp_path) == "exploit" + + +# --- sum_prior_spend -------------------------------------------------------- + + +def test_sum_prior_spend_sums_settled_only(tmp_path): + ledger = tmp_path / "spend-ledger.jsonl" + ledger.write_text( + "\n".join( + [ + json.dumps({"event": "run_started"}), + json.dumps({"event": "call_reserved", "cost_usd": 99.0}), + json.dumps({"event": "call_settled", "cost_usd": 0.5}), + json.dumps({"event": "call_settled", "cost_usd": 1.25}), + json.dumps({"event": "budget_snapshot", "cost_usd": 0.0}), + ] + ) + ) + assert sum_prior_spend(tmp_path) == pytest.approx(1.75) + + +def test_sum_prior_spend_absent_ledger(tmp_path): + assert sum_prior_spend(tmp_path) == 0.0 + + +def test_sum_prior_spend_tolerates_malformed_lines(tmp_path): + ledger = tmp_path / "spend-ledger.jsonl" + ledger.write_text( + "\n".join( + [ + "not json at all", + json.dumps({"event": "call_settled", "cost_usd": 2.0}), + json.dumps({"event": "call_settled", "cost_usd": None}), + json.dumps({"event": "call_settled"}), # missing cost + "", + ] + ) + ) + assert sum_prior_spend(tmp_path) == pytest.approx(2.0) + + +# --- resolve_session_dir ---------------------------------------------------- + + +def test_resolve_session_dir_bare_name(tmp_path): + resolved = resolve_session_dir("sh-1234abcd", tmp_path) + assert resolved == tmp_path / "sh-1234abcd" + + +def test_resolve_session_dir_explicit_path(tmp_path): + session = tmp_path / "runs" / "sh-9999" + session.mkdir(parents=True) + resolved = resolve_session_dir(str(session), "/some/other/output") + assert resolved == session + + +# --- budget carry-forward --------------------------------------------------- + + +def _seeded_ledger(tmp_path, *, initial: float, budget: float) -> SpendLedger: + return SpendLedger( + limit_usd=budget, + session_id="resume-test", + repo_url="/tmp/repo", + output_dir=tmp_path, + input_price_per_million=0.0, + output_price_per_million=1_000_000.0, + initial_spent_usd=initial, + ) + + +def test_initial_spent_seeds_spent_and_remaining(tmp_path): + ledger = _seeded_ledger(tmp_path, initial=5.0, budget=10.0) + assert ledger.spent_usd == pytest.approx(5.0) + assert ledger.remaining_usd == pytest.approx(5.0) + + +def test_initial_spent_recorded_in_run_started_event(tmp_path): + _seeded_ledger(tmp_path, initial=3.0, budget=10.0) + ledger_file = tmp_path / "resume-test" / "spend-ledger.jsonl" + events = [json.loads(line) for line in ledger_file.read_text().splitlines() if line] + run_started = next(e for e in events if e["event"] == "run_started") + assert run_started["carried_forward_usd"] == pytest.approx(3.0) + + +def test_resume_seeded_past_cap_reports_exhausted_on_reserve(tmp_path): + # Prior run already spent the whole cap; the next reservation must fail. + ledger = _seeded_ledger(tmp_path, initial=10.0, budget=10.0) + assert ledger.remaining_usd == pytest.approx(0.0) + from clearwing.llm.budget import BudgetExceeded + + with pytest.raises(BudgetExceeded): + ledger.reserve_call( + model="test-model", + provider="test", + stage="verify", + input_token_upper_bound=0, + requested_max_output_tokens=1, + supports_output_limit=True, + ) + assert ledger.exhausted is True + + +def test_negative_initial_spent_rejected(tmp_path): + with pytest.raises(BudgetConfigurationError): + _seeded_ledger(tmp_path, initial=-1.0, budget=10.0) + + +def test_non_finite_initial_spent_rejected(tmp_path): + with pytest.raises(BudgetConfigurationError): + _seeded_ledger(tmp_path, initial=float("inf"), budget=10.0) + + +# --- runner resume wiring (deliberate exit) --------------------------------- + + +def test_runner_resume_without_checkpoint_raises(tmp_path): + """Resume is deliberate: a session dir with no checkpoint exits, not restarts.""" + from clearwing.sourcehunt.runner import SourceHuntRunner + + empty_session = tmp_path / "sh-empty" + empty_session.mkdir() + with pytest.raises(ValueError, match="no stage checkpoint"): + SourceHuntRunner( + repo_url="/tmp/repo", + output_dir=str(tmp_path), + resume_session=str(empty_session), + ) + + +def test_runner_resume_nonexistent_session_raises(tmp_path): + from clearwing.sourcehunt.runner import SourceHuntRunner + + with pytest.raises(ValueError, match="no stage checkpoint"): + SourceHuntRunner( + repo_url="/tmp/repo", + output_dir=str(tmp_path), + resume_session="sh-does-not-exist", + ) + + +def test_runner_resume_adopts_session_id_and_detects_stage(tmp_path): + """A present checkpoint sets the resumed session id and the resume stage.""" + from clearwing.sourcehunt.runner import SourceHuntRunner + + session = tmp_path / "sh-resume01" + write_stage_checkpoint( + session, + "verify", + findings=[Finding(id="f")], + verified=[Finding(id="f")], + rejected=[], + budget_spent_usd=1.0, + ) + runner = SourceHuntRunner( + repo_url="/tmp/repo", + output_dir=str(tmp_path), + resume_session=str(session), + ) + assert runner._session_id == "sh-resume01" + assert runner._resume_from == "verify" + + +def test_runner_resume_by_bare_name(tmp_path): + from clearwing.sourcehunt.runner import SourceHuntRunner + + session = tmp_path / "sh-bare" + write_stage_checkpoint( + session, "hunt", findings=[Finding(id="f")], budget_spent_usd=0.0 + ) + runner = SourceHuntRunner( + repo_url="/tmp/repo", + output_dir=str(tmp_path), + resume_session="sh-bare", + ) + assert runner._session_id == "sh-bare" + assert runner._resume_from == "hunt" + + +def test_runner_no_resume_leaves_state_clean(tmp_path): + from clearwing.sourcehunt.runner import SourceHuntRunner + + runner = SourceHuntRunner( + repo_url="/tmp/repo", + output_dir=str(tmp_path), + ) + assert runner._resume_session is None + assert runner._resume_from is None + assert runner._session_id.startswith("sh-") + + +# --- end-to-end: checkpoints written, then resumed -------------------------- + + +def _standard_runner(tmp_path, **overrides): + kwargs = dict( + repo_url=str(FIXTURE_C_PROPAGATION), + local_path=str(FIXTURE_C_PROPAGATION), + depth="standard", + budget_usd=1.0, + max_parallel=2, + output_dir=str(tmp_path), + ranker_llm=_make_ranker_llm(), + hunter_llm=_make_hunter_llm(), + verifier_llm=_make_verifier_llm(), + no_exploit=True, + ) + kwargs.update(overrides) + return SourceHuntRunner(**kwargs) + + +def test_full_run_writes_all_stage_checkpoints(tmp_path): + runner = _standard_runner(tmp_path) + result = runner.run() + session_dir = tmp_path / result.session_id + # hunt + verify are always reached; exploit is reached even with no_exploit + # (the stage runs and simply exploits nothing). + assert (session_dir / "checkpoints" / "hunt.json").exists() + assert (session_dir / "checkpoints" / "verify.json").exists() + assert (session_dir / "checkpoints" / "exploit.json").exists() + # Each is plain, loadable JSON with the documented shape. + hunt = json.loads((session_dir / "checkpoints" / "hunt.json").read_text()) + assert hunt["stage"] == "hunt" + assert "findings" in hunt + verify = json.loads((session_dir / "checkpoints" / "verify.json").read_text()) + assert "verified" in verify and "rejected" in verify + + +def test_resume_from_verify_skips_hunt_and_verify(tmp_path): + # First: a full run to produce real checkpoints. + first = _standard_runner(tmp_path) + result = first.run() + session_dir = tmp_path / result.session_id + assert (session_dir / "checkpoints" / "verify.json").exists() + + # Delete the exploit checkpoint so the highest present is verify. + (session_dir / "checkpoints" / "exploit.json").unlink() + + # Resume: hunter + verifier must never be invoked; report still runs. + hunter = _make_hunter_llm() + verifier = _make_verifier_llm() + resumed = _standard_runner( + tmp_path, + hunter_llm=hunter, + verifier_llm=verifier, + resume_session=str(session_dir), + ) + assert resumed._resume_from == "verify" + resumed_result = resumed.run() + + assert resumed_result.session_id == result.session_id + # Skipped phases never touched their models. + hunter.bind_tools.assert_not_called() + verifier.aask_text.assert_not_called() + # Report re-ran: outputs exist. + for path in resumed_result.output_paths.values(): + assert Path(path).exists() + + +def test_resume_budget_carries_forward(tmp_path): + first = _standard_runner(tmp_path) + result = first.run() + session_dir = tmp_path / result.session_id + prior_spend = sum_prior_spend(session_dir) + + resumed = _standard_runner(tmp_path, resume_session=str(session_dir)) + ledger = resumed._ensure_spend_ledger() + # The resumed ledger starts already having "spent" the prior run's cost. + assert ledger.spent_usd == pytest.approx(prior_spend)