diff --git a/pyproject.toml b/pyproject.toml index ebe7bbe..992d635 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,8 @@ requires-python = ">=3.11" dependencies = ["pyyaml>=6.0"] [project.optional-dependencies] +# Hunter learning research only; the core install stays PyYAML. +hunt = ["dspy>=3.4,<3.5", "pandas>=2.2,<3"] # The documentation site (mkdocs.yml, docs/). Mermaid renders client-side through Material's superfences fence. docs = [ "mkdocs>=1.6,<2", diff --git a/src/openultrasast/hunt/__init__.py b/src/openultrasast/hunt/__init__.py new file mode 100644 index 0000000..10ed0c3 --- /dev/null +++ b/src/openultrasast/hunt/__init__.py @@ -0,0 +1 @@ +"""Repository annotations and offline contracts for hunter learning research.""" diff --git a/src/openultrasast/hunt/artifact.py b/src/openultrasast/hunt/artifact.py new file mode 100644 index 0000000..96ae253 --- /dev/null +++ b/src/openultrasast/hunt/artifact.py @@ -0,0 +1,99 @@ +"""Content-addressed run metadata and budget accounting without a model runtime.""" + +from __future__ import annotations + +import hashlib +import json +import math +from collections.abc import Mapping +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +from .table import canonical_json + + +def _exhausted(usage: Mapping[str, float], budget: Mapping[str, float]) -> bool: + for name in ("usd", "calls"): + if name in budget: + bound = budget[name] + if not math.isfinite(bound) or bound < 0: + raise ValueError("budget bounds must be finite and non-negative") + if usage.get(name, 0) >= bound: + return True + return False + + +@dataclass +class RunArtifact: + program_id: str + schema_version: int + mask: tuple[str, ...] + data_digests: dict[str, str] + configurations: list[dict[str, Any]] + usage: dict[str, float] + versions: dict[str, str] + created: str + + @property + def artifact_id(self) -> str: + return hashlib.sha256(canonical_json(asdict(self)).encode("utf-8")).hexdigest() + + def finish(self, budget: Mapping[str, float]) -> RunArtifact | None: + """Return no completed artifact when either registered budget is exhausted.""" + return None if _exhausted(self.usage, budget) else self + + def write(self, directory: Path | str, *, budget: Mapping[str, float] | None = None) -> Path | None: + """Write metadata under its content id, after an optional final budget check.""" + if budget is not None and self.finish(budget) is None: + return None + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + path = directory / f"{self.artifact_id}.json" + path.write_text(canonical_json({**asdict(self), "artifact_id": self.artifact_id}) + "\n", encoding="utf-8") + return path + + @classmethod + def read(cls, path: Path | str) -> RunArtifact: + payload = json.loads(Path(path).read_text(encoding="utf-8")) + identifier = payload.pop("artifact_id", None) + payload["mask"] = tuple(payload["mask"]) + artifact = cls(**payload) + if identifier != artifact.artifact_id: + raise ValueError("run artifact id does not match its contents") + return artifact + + +@dataclass +class UsageLedger: + """Count paid calls, replay hits, spend and elapsed call time separately.""" + + artifact: RunArtifact | None = None + rows: list[dict[str, float]] = field(default_factory=list) + + @property + def usage(self) -> dict[str, float]: + return {name: sum(row[name] for row in self.rows) for name in ("usd", "calls", "replay_hits", "wall_seconds")} + + def add(self, call_row: Mapping[str, Any]) -> None: + replay_hits = call_row.get("replay_hits", int(bool(call_row.get("replay_hit", False)))) + values = {"usd": call_row.get("usd", 0), "calls": call_row.get("calls", 0 if replay_hits else 1), + "replay_hits": replay_hits, "wall_seconds": call_row.get("wall_seconds", 0)} + if any(not isinstance(v, int | float) or not math.isfinite(v) or v < 0 for v in values.values()): + raise ValueError("usage values must be finite and non-negative") + if any(int(values[k]) != values[k] for k in ("calls", "replay_hits")): + raise ValueError("call and replay counts must be integers") + self.rows.append(values) + + def exhausted(self, budget: Mapping[str, float]) -> bool: + return _exhausted(self.usage, budget) + + def finish(self, budget: Mapping[str, float], artifact: RunArtifact | None = None) -> RunArtifact | None: + """Attach measured usage to a completed run, or retain only this partial ledger.""" + if self.exhausted(budget): + return None + result = artifact if artifact is not None else self.artifact + if result is None: + raise ValueError("finish needs a run artifact") + result.usage = self.usage + return result diff --git a/src/openultrasast/hunt/cli.py b/src/openultrasast/hunt/cli.py new file mode 100644 index 0000000..ac25565 --- /dev/null +++ b/src/openultrasast/hunt/cli.py @@ -0,0 +1,80 @@ +"""Command line for the hunt program: run a program over a frozen manifest and write scorable records. + + python -m openultrasast.hunt.cli run MANIFEST --repos REPOS.json --out DIR [--mask producer,...] + [--only vulnerability|ordinary|all] [--slice 1] [--budget-usd 2.0] [--budget-calls 200] + [--no-tools] [--max-iters 8] [--pilot-report] + +The model key comes from the environment only. Records and the DSPy cache go under DIR; the pilot manifest is +refused unless --pilot-report is given (the held-out set is a report surface, not a tuning surface). +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import dspy + +from benchmarks.unseen.replay import load_slice + +from . import program as hp +from . import runner + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + commands = parser.add_subparsers(dest="command", required=True) + run = commands.add_parser("run") + run.add_argument("manifest", type=Path) + run.add_argument("--repos", type=Path, required=True, help="JSON map repo -> local clone with history") + run.add_argument("--out", type=Path, required=True) + run.add_argument("--mask", default="", help="comma-separated producers or inputs to drop (an ablation arm)") + run.add_argument("--only", choices=("vulnerability", "ordinary", "all"), default="all") + run.add_argument("--slice", default="1") + run.add_argument("--budget-usd", type=float, default=None) + run.add_argument("--budget-calls", type=float, default=None) + run.add_argument("--no-tools", action="store_true", help="one prediction per change instead of the tool loop") + run.add_argument("--max-iters", type=int, default=8) + run.add_argument("--pilot-report", action="store_true") + args = parser.parse_args(argv) + + mask = tuple(part.strip() for part in args.mask.split(",") if part.strip()) + args.out.mkdir(parents=True, exist_ok=True) + lm = hp.build_lm(cache_dir=args.out / "dspy-cache") + dspy.configure(lm=lm, adapter=dspy.ChatAdapter()) + tools_for = None if args.no_tools else hp.repository_tools + # The program is built with the tools of a placeholder root; the runner rebinds them per checkout. + program = hp.HuntProgram(mask=mask, tools=() if args.no_tools else hp.repository_tools(Path(".")), max_iters=args.max_iters) + pid = hp.program_id(program, mask=mask) + budget = {key: value for key, value in (("usd", args.budget_usd), ("calls", args.budget_calls)) if value is not None} + repos = json.loads(args.repos.read_text()) + versions = {"dspy": dspy.__version__, "python": sys.version.split()[0]} + with dspy.track_usage() as usage: + artifact = runner.run( + args.manifest, + repos, + args.out, + program=program, + program_id=pid, + mask=mask, + only=args.only, + slice_id=args.slice, + budget=budget or None, + pilot_report=args.pilot_report, + tools_for=tools_for, + load_changes=load_slice, + versions=versions, + ) + totals = usage.get_total_tokens() if hasattr(usage, "get_total_tokens") else {} + print( + json.dumps( + {"program_id": pid, "mask": list(mask), "artifact": artifact.artifact_id if artifact else None, "tokens": totals}, default=str + ) + ) + return 0 if artifact is not None else 3 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/openultrasast/hunt/labels.py b/src/openultrasast/hunt/labels.py new file mode 100644 index 0000000..371f036 --- /dev/null +++ b/src/openultrasast/hunt/labels.py @@ -0,0 +1,33 @@ +"""Attach benchmark fix-site labels without exposing them to feature producers.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from typing import Any + + +def label_rows(rows: Sequence[Mapping[str, Any]], change: Mapping[str, Any]) -> list[dict[str, Any]]: + """Label verified sites, leave other positive rows unknown and label ordinary rows zero.""" + kind = change.get("kind") + if kind not in ("vulnerability", "ordinary"): + raise ValueError("change kind must be vulnerability or ordinary") + if not change.get("case_id"): + raise ValueError("a label needs a traceable case_id") + if kind == "vulnerability" and not all(change.get(k) for k in ("path", "function", "function_lines", "family")): + raise ValueError("a vulnerability label needs its path, function, span and family") + result = [] + for row in rows: + data = deepcopy(dict(row)) + span = change.get("function_lines") + function = row.get("function", row.get("feature.source.function")) + site = kind == "vulnerability" and row.get("path") == change["path"] and ( + function == change["function"] or (span[0] <= row.get("line", 0) <= span[1]) + ) + data.update(label=1 if site else 0 if kind == "ordinary" else None, + family=change["family"] if site else None, + fix_span=list(span) if site else None, + label_ref=change["case_id"] if site or kind == "ordinary" else None, + verification=None) + result.append(data) + return result diff --git a/src/openultrasast/hunt/metric.py b/src/openultrasast/hunt/metric.py new file mode 100644 index 0000000..399cc20 --- /dev/null +++ b/src/openultrasast/hunt/metric.py @@ -0,0 +1,72 @@ +"""Search outcomes using the frozen benchmark's location and family rules.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class Outcome: + hit: bool + suspicions: int + on_diff: int + feedback: str + + +def _location_match(change: Mapping[str, Any], finding: Mapping[str, Any], tolerance: int) -> bool: + """Mirror unseen.score.location_match, including its config-secret exception.""" + for file in change.get("files", []): + if finding.get("path") != file["path"]: + continue + line = finding.get("line", 0) or 0 + if any(start - tolerance <= line <= end + tolerance for start, end in file.get("head_lines", [])): + return True + if change.get("family") != "config_secrets" and file["path"] == change.get("path", file["path"]): + function = change.get("function") + if function and function != "" and finding.get("function") == function: + return True + span = change.get("function_lines") or file.get("function_lines") + if span and span[0] <= line <= span[1]: + return True + return False + + +def _on_diff(change: Mapping[str, Any], finding: Mapping[str, Any], tolerance: int) -> bool: + return any(finding.get("path") == file["path"] and + any(start - tolerance <= (finding.get("line", 0) or 0) <= end + tolerance + for start, end in file.get("head_lines", [])) for file in change.get("files", [])) + + +def search_outcome(change: Mapping[str, Any], findings: Sequence[Mapping[str, Any]], *, tolerance: int = 10) -> Outcome: + """Count all suspicions and record whether one catches a positive site.""" + if change.get("kind") not in ("vulnerability", "ordinary"): + raise ValueError("change kind must be vulnerability or ordinary") + if tolerance < 0: + raise ValueError("location tolerance cannot be negative") + hit = change["kind"] == "vulnerability" and any( + finding.get("family") == change.get("family") and _location_match(change, finding, tolerance) for finding in findings + ) + on_diff = sum(_on_diff(change, finding, tolerance) for finding in findings) + feedback = [] + if change["kind"] == "vulnerability" and not hit: + span = change.get("function_lines") or [next((start for f in change.get("files", []) + for start, _ in f.get("head_lines", [])), 0)] + feedback.append(f"missed {change.get('path', '')}:{span[0]} {change.get('function', '')} family {change.get('family', '')}") + off_diff = [f"{f.get('path', '')}:{f.get('line', 0)} family {f.get('family', '')}" + for f in findings if not _on_diff(change, f, tolerance)] + if off_diff: + feedback.append("off-diff suspicions: " + ", ".join(off_diff)) + return Outcome(hit, len(findings), on_diff, "; ".join(feedback)) + + +def _field(value: object, name: str, default: Any = None) -> Any: + return value.get(name, default) if isinstance(value, Mapping) else getattr(value, name, default) + + +def search_metric(example: object, pred: object, trace: object = None) -> float: + """Provide DSPy's metric signature while keeping this module dependency free.""" + change = _field(example, "change", example) + outcome = search_outcome(change, _field(pred, "findings", [])) + return float(outcome.hit) if change["kind"] == "vulnerability" else 1 - min(1, outcome.suspicions / 5) diff --git a/src/openultrasast/hunt/producers.py b/src/openultrasast/hunt/producers.py new file mode 100644 index 0000000..289a078 --- /dev/null +++ b/src/openultrasast/hunt/producers.py @@ -0,0 +1,294 @@ +"""Four fixed annotations computed from head source and the tool's own facts.""" + +from __future__ import annotations + +import json +import re +import subprocess +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from functools import cached_property +from pathlib import Path +from typing import Any + +from .. import mapping, source +from ..model.taxonomy import DEFAULT_FAMILIES_PATH, load_families +from ..preprocess import build_file_target, detect_language +from ..semantic.facts import load_facts +from ..semantic.ir import FileIR, parse_file +from ..semantic.obligations import check_obligations, load_obligation_facts +from ..semantic.obligations.dominance import OrderDominance +from ..semantic.obligations.operations import covers, find_operations, valid_witness +from .labels import label_rows +from .table import FEATURE_NAMES + +Change = Mapping[str, Any] + + +class _Absent(ValueError): + def __init__(self, reason: str) -> None: + self.reason = reason + super().__init__(reason) + + +def _git(root: Path, *args: str) -> str: + result = subprocess.run(["git", "-C", str(root), *args], capture_output=True, check=True) + return result.stdout.decode("utf-8", errors="replace") + + +def _safe_path(root: Path, path: str) -> Path: + relative = Path(path) + resolved = (root / relative).resolve() + if relative.is_absolute() or any(part.startswith(".env") or part == ".git" for part in relative.parts): + raise ValueError("private or absolute source path") + if not resolved.is_relative_to(root.resolve()): + raise ValueError("source path is outside the repository") + return resolved + + +@dataclass +class Context: + """Cache work shared by every candidate of one checked-out change.""" + + root: Path + change: Change + texts: dict[str, str] = field(default_factory=dict) + spans: dict[str, tuple[source.FunctionRange, ...]] = field(default_factory=dict) + range_errors: dict[str, str] = field(default_factory=dict) + irs: dict[str, FileIR] = field(default_factory=dict) + diff_bytes: int = 0 + + @cached_property + def head_verified(self) -> bool: + """True once the checkout is proven to be at the change's head; absent otherwise.""" + if _git(self.root, "rev-parse", "HEAD").strip() != _git(self.root, "rev-parse", str(self.change["head"])).strip(): + raise _Absent("head_mismatch") + return True + + def verify_head(self) -> None: + if not self.head_verified: + raise _Absent("head_mismatch") + + def read(self, path: str) -> str: + safe = _safe_path(self.root, path) + if path not in self.texts: + self.verify_head() + data = safe.read_bytes() + if not data: + raise _Absent("empty_source") + self.texts[path] = data.decode("utf-8", errors="replace") + return self.texts[path] + + def ranges(self, path: str) -> tuple[source.FunctionRange, ...]: + if path not in self.spans: + self.spans[path] = source.ranges(path, self.read(path), source.canonical_language(detect_language(self.root / path))) + return self.spans[path] + + @cached_property + def product_files(self) -> list[str]: + return source.product_files(self.root) + + @cached_property + def entry_points(self) -> list[mapping.EntryPointRecord]: + for path in self.product_files: + self.read(path) + return mapping.analyze_entry_points(self.root, [build_file_target(self.root, self.root / path) for path in self.product_files]) + + @cached_property + def hunks(self) -> dict[str, list[list[int]]]: + self.verify_head() + # Exclude private files before git opens the patch; only hunk coordinates become features. + patch = _git(self.root, "diff", "--no-ext-diff", "--no-textconv", "-U0", + f"{self.change['base']}..{self.change['head']}", "--", ".", + ":(glob,exclude)**/.env*", ":(exclude).env*") + self.diff_bytes = len(patch.encode("utf-8")) + result: dict[str, list[list[int]]] = {} + path = None + for text in patch.splitlines(): + if text.startswith("+++ "): + token = text[4:] + token = json.loads(token) if token.startswith('"') else token + path = token[2:] if token.startswith("b/") else None + elif path is not None and (match := re.match(r"@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", text)): + start, count = int(match[1]), int(match[2] or "1") + if count: + result.setdefault(path, []).append([start, start + count - 1]) + return result + + def ir(self, path: str) -> FileIR: + if path not in self.irs: + self.irs[path] = parse_file(path, self.read(path), source.canonical_language(detect_language(self.root / path))) + return self.irs[path] + + +Compute = Callable[[Path, Change, str, int, Context], dict[str, object]] + + +@dataclass(frozen=True) +class Producer: + """A versioned producer with a closed output column list and private failure handling.""" + + name: str + version: str + columns: tuple[str, ...] + compute: Compute + + +def _producer(name: str, compute: Compute) -> Producer: + columns = tuple(f"feature.{name}.{column}" for column in FEATURE_NAMES[name]) + provenance = tuple(f"provenance.{name}.{column}" for column in ("ok", "reason", "version", "bytes_read")) + + def guarded(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: + result: dict[str, object] = dict.fromkeys(columns) + reason = None + before = set(ctx.texts) + try: + values = compute(root, change, path, line, ctx) + if set(values) != set(FEATURE_NAMES[name]): + raise ValueError("producer output does not match its columns") + result.update({f"feature.{name}.{key}": value for key, value in values.items()}) + except Exception as error: + reason = error.reason if isinstance(error, _Absent) else type(error).__name__ + bytes_read = sum(len(text.encode("utf-8")) for key, text in ctx.texts.items() if key not in before or key == path) + if name == "diff": + bytes_read = ctx.diff_bytes + elif name == "taxonomy" and reason is None: + bytes_read = len(DEFAULT_FAMILIES_PATH.read_bytes()) + result.update({f"provenance.{name}.ok": reason is None, f"provenance.{name}.reason": reason, + f"provenance.{name}.version": "1", f"provenance.{name}.bytes_read": bytes_read}) + return result + + return Producer(name, "1", columns + provenance, guarded) + + +def _diff(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: + hunks = ctx.hunks.get(path, []) + distance = min((max(start - line, 0, line - end) for start, end in hunks), default=None) + return {"hunks": hunks, "on_changed_line": distance == 0, "distance": distance} + + +def _source(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: + ctx.read(path) + if path in ctx.range_errors: + raise _Absent(ctx.range_errors[path]) + facts = source.facts_for(root, path, line, entry_points=ctx.entry_points) + # Select the occurrence at this line when several methods share a name. + spans = [r for r in ctx.ranges(path) if r.name == facts.function and r.start <= line <= r.end] + span = min(spans, key=lambda r: r.end - r.start) if spans else None + return {"function": facts.function, "span": [span.start, span.end] if span else None, + "callers": [list(caller) for caller in facts.callers], + "entry_points": [{"path": entry.path, "line": entry.line, "kind": entry.kind} for entry in facts.entry_points], + "language": facts.language} + + +def _roles(function: Any, facts: Any) -> dict[str, str]: + names = set(function.params) + names.update(call.name for call in function.calls) + for bind in function.binds: + names.update((bind.name, *bind.names)) + if bind.call_name: + names.add(bind.call_name) + for call in function.calls: + for arguments in call.arg_names: + names.update(arguments) + roles = {} + for role, entries, attribute in (("source", facts.sources, "patterns"), ("sink", facts.sinks, "calls"), + ("sanitizer", facts.sanitizers, "calls")): + for fact in entries: + for pattern in getattr(fact, attribute): + token = pattern.rstrip("(") + for name in sorted(names): + if name == token or name.endswith("." + token) or token.startswith(name + ".") and role == "source": + # A shared sink/sanitizer name remains a sink; parameterization needs more than a name. + roles.setdefault(name, role) + return roles + + +def _obligations(path: str, ctx: Context) -> list[dict[str, object]]: + ir, text = ctx.ir(path), ctx.read(path) + facts, flow = load_obligation_facts(), load_facts() + scoped = facts.for_language(ir.language) + if not scoped.operations and not scoped.dischargers: + raise _Absent("obligations_language_unsupported") + dominance = OrderDominance({path: text}) + result = check_obligations(irs={path: (ir, text)}, entries=ctx.entry_points, facts=facts, flow_facts=flow, + policy=None, paths=(), dominance=dominance, store_shapes=(), min_siblings=3) + # No finding does not imply discharge: the checker also gates on entry-point reachability. + return [{"operation": operation.kind, "line": operation.line, + "discharged": any(valid_witness(witness) and covers(witness, operation) and + dominance.dominates(witness, operation, None) for witness in result.discharges)} + for operation in find_operations(ir, scoped, text=text)] + + +def _semantic(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: + ir = ctx.ir(path) + if not ir.parse_ok: + raise _Absent("language_unsupported" if ir.reason == "language_unsupported" else "parse_failed") + functions = [fn for fn in ir.functions if fn.start_line <= line <= fn.end_line] + if not functions: + raise _Absent("function_unavailable") + function = min(functions, key=lambda fn: fn.end_line - fn.start_line) + return {"roles": _roles(function, load_facts().for_language(ir.language)), + "def_use": {"binds": [{"line": b.line, "name": b.name, "names": list(b.names), "call_name": b.call_name} + for b in function.binds], + "calls": [{"line": c.line, "name": c.name, "names": [list(names) for names in c.arg_names]} + for c in function.calls]}, + "obligations": _obligations(path, ctx)} + + +def _taxonomy(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: + return {"families": [{"id": f.id, "definition": " ".join(f.description.split())} for f in load_families().families]} + + +PRODUCERS = (_producer("diff", _diff), _producer("source", _source), + _producer("semantic", _semantic), _producer("taxonomy", _taxonomy)) + + +def _candidates(root: Path, change: Change, ctx: Context) -> list[dict[str, Any]]: + selected: set[tuple[str, int]] = set() + files = list(change.get("files", [])) + if change.get("kind") == "vulnerability" and not any(f["path"] == change["path"] for f in files): + files.append({"path": change["path"], "head_lines": []}) + for file in files: + path = file["path"] + _safe_path(root, path) + try: + ranges = ctx.ranges(path) + except Exception as error: + ctx.range_errors[path] = error.reason if isinstance(error, _Absent) else type(error).__name__ + line = next((start for start, _ in file.get("head_lines", [])), 1) + selected.add((path, max(1, line))) + continue + touched = [r for r in ranges if r.kind == "function" and + any(r.start <= end and start <= r.end for start, end in file.get("head_lines", []))] + selected.update((path, r.start) for r in touched) + if not touched and file.get("head_lines") and ranges: + selected.add((path, max(1, file["head_lines"][0][0]))) + if change.get("kind") == "vulnerability" and path == change["path"]: + labelled = [r for r in ranges if r.name == change["function"] and + r.start <= change["function_lines"][1] and change["function_lines"][0] <= r.end] + selected.update((path, r.start) for r in labelled) + if not labelled: + selected.add((path, change["function_lines"][0])) + return [{"change_id": change.get("id", change.get("case_id")), "repo": change.get("repo", change.get("repository")), + "head": change["head"], "path": path, "line": line, "kind": change["kind"], "fold": change.get("fold")} + for path, line in sorted(selected)] + + +def candidate_rows(root: Path, change: Change) -> list[dict[str, Any]]: + """Enumerate touched head functions and the positive fix site, retaining unreadable files.""" + root = Path(root) + return _candidates(root, change, Context(root, change)) + + +def compute_rows(root: Path, change: Change, producers: tuple[Producer, ...] = PRODUCERS) -> list[dict[str, Any]]: + """Fill candidate rows, then join labels after every producer has finished.""" + root = Path(root) + # Producers receive only revision coordinates; advisory and label fields never cross this boundary. + coordinates = {key: change[key] for key in ("base", "head")} + ctx = Context(root, coordinates) + rows = _candidates(root, change, ctx) + for row in rows: + for producer in producers: + row.update(producer.compute(root, coordinates, row["path"], row["line"], ctx)) + return label_rows(rows, change) diff --git a/src/openultrasast/hunt/program.py b/src/openultrasast/hunt/program.py new file mode 100644 index 0000000..14d20d7 --- /dev/null +++ b/src/openultrasast/hunt/program.py @@ -0,0 +1,214 @@ +"""The hunt as a DSPy program: a signature over the feature row, three repository tools, and the answer check. + +The program sees only the unmasked feature columns of a change, rendered in a fixed order, plus the changed +hunks. Its answer is the classifier's structure (family from the closed list, confidence, rationale, cited +lines). DSPy proposes instructions and demonstrations; the acceptance rule in `accept.py` decides, outside DSPy. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import dspy +import pydantic + +from .. import hunter_tools +from ..model.endpoint import _PRICES, DEEPSEEK_BASE_URL, DEFAULT_DETECTOR_MODEL +from ..model.taxonomy import FAMILY_IDS + +BASELINE_INSTRUCTION = ( + "You are a security reviewer looking at one pushed change. Use read_file, grep_repo and find_refs to check " + "whether the changed code introduces a vulnerability that an attacker could reach. Report only what the code " + "shows: a path, a line, the family from the list, your confidence between 0 and 1, a one-sentence rationale and " + "the lines you relied on. Report nothing when the change is ordinary." +) + +# Input name -> (feature column, renderer). Order is the order the model sees them in; masked columns are absent. +INPUTS: tuple[tuple[str, str, str], ...] = ( + ("change_summary", "", "The files and functions the push touched."), + ("hunks", "feature.diff.hunks", "Changed line ranges per file at head, as [start, end] pairs."), + ("function_text", "feature.source.span", "The enclosing function at head, numbered lines."), + ("callers", "feature.source.callers", "Where the enclosing function is called from: path and line."), + ("entry_points", "feature.source.entry_points", "In-repository entry points that reach the file."), + ("roles", "feature.semantic.roles", "Identifiers known as sources, sinks or sanitizers in this function."), + ("def_use", "feature.semantic.def_use", "Assignments and calls inside the function, with lines."), + ("obligations", "feature.semantic.obligations", "Checks the code is expected to perform before an operation."), + ("families", "feature.taxonomy.families", "The closed list of vulnerability families with one-line definitions."), +) + + +class Finding(pydantic.BaseModel): + """One suspicion in the classifier's answer structure.""" + + path: str + line: int + family: str + confidence: float = 0.5 + rationale: str = "" + cited_lines: list[int] = pydantic.Field(default_factory=list) + + +def signature_for(mask: Iterable[str] = ()) -> type[dspy.Signature]: + """The hunt signature with every unmasked input, in the fixed order, and the findings output.""" + dropped = set(mask) + fields: dict[str, Any] = {} + for name, column, desc in INPUTS: + producer = column.split(".")[1] if column else "" + if name in dropped or column in dropped or (producer and producer in dropped): + continue + fields[name] = (str, dspy.InputField(desc=desc)) + fields["findings"] = (list[Finding], dspy.OutputField(desc="JSON list; empty when the change is ordinary")) + return dspy.Signature(fields, BASELINE_INSTRUCTION) + + +def numbered(lines: Sequence[str], start: int, end: int) -> str: + """Lines start..end (1-based, inclusive) with their numbers, so cited lines can be checked.""" + return "\n".join(f"{number:>5} {lines[number - 1]}" for number in range(max(start, 1), min(end, len(lines)) + 1)) + + +def render_rows( + change: Mapping[str, Any], rows: Sequence[Mapping[str, Any]], texts: Mapping[str, Sequence[str]], *, mask: Iterable[str] = () +) -> dict[str, str]: + """Render one change's candidate rows into the signature's inputs. Masked inputs are not produced at all.""" + signature = signature_for(mask) + wanted = set(signature.input_fields) + out: dict[str, str] = {} + if "change_summary" in wanted: + touched = sorted({(row["path"], str(row.get("feature.source.function") or "")) for row in rows}) + out["change_summary"] = "\n".join(f"{path}::{function}" for path, function in touched) or "(nothing readable)" + if "hunks" in wanted: + hunks = {row["path"]: row.get("feature.diff.hunks") for row in rows if row.get("feature.diff.hunks")} + out["hunks"] = json.dumps(hunks, sort_keys=True) if hunks else "(no hunks)" + if "function_text" in wanted: + blocks = [] + for row in rows: + span = row.get("feature.source.span") + lines = texts.get(row["path"]) + if span and lines: + blocks.append(f"# {row['path']}::{row.get('feature.source.function')}\n{numbered(lines, span[0], span[1])}") + out["function_text"] = "\n\n".join(blocks) or "(no function text)" + for name, column, _ in INPUTS[3:]: + if name not in wanted: + continue + values = {f"{row['path']}::{row.get('feature.source.function') or ''}": row.get(column) for row in rows} + values = {key: value for key, value in values.items() if value not in (None, [], {})} + if name == "families": # one list for the whole change, not one per row + families = next((value for value in values.values()), None) + out[name] = "\n".join(f"{f['id']}: {f['definition']}" for f in families) if families else "(no families)" + else: + out[name] = json.dumps(values, sort_keys=True) if values else "(none)" + return out + + +def shown_lines(inputs: Mapping[str, str]) -> dict[str, set[int]]: + """The numbered lines the model was shown per path, parsed back from `function_text`.""" + shown: dict[str, set[int]] = {} + current = None + for line in inputs.get("function_text", "").splitlines(): + if line.startswith("# ") and "::" in line: + current = line[2:].split("::", 1)[0] + elif current and line[:5].strip().isdigit(): + shown.setdefault(current, set()).add(int(line[:5])) + return shown + + +def check(findings: Iterable[Finding | Mapping[str, Any]], inputs: Mapping[str, str]) -> list[dict[str, Any]]: + """Apply the answer rules: a family outside the closed list, or a cited line the model was not shown, scores + against the program by becoming `unknown` with a reason. Nothing is dropped, so the metric sees it.""" + shown = shown_lines(inputs) + allowed = {family for family in FAMILY_IDS if family != "unknown"} + out = [] + for item in findings: + data = item.model_dump() if isinstance(item, pydantic.BaseModel) else dict(item) + reasons = [] + if data.get("family") not in allowed: + reasons.append("family_not_in_list") + cited = [int(line) for line in data.get("cited_lines") or []] + if cited and not set(cited) <= shown.get(data.get("path", ""), set()): + reasons.append("cited_line_not_shown") + if reasons: + data["reported_family"] = data.get("family") + data["family"] = "unknown" + data["unknown_reason"] = ",".join(reasons) + out.append(data) + return out + + +def repository_tools(root: Path, *, max_chars: int = 4000, max_matches: int = 50) -> list[dspy.Tool]: + """The three tools the shipped hunter has, bound to one checkout and clamped to it.""" + root = Path(root) + + def read_file(path: str) -> str: + """Read a file of the repository by its relative path, numbered lines.""" + return hunter_tools.read_file(root, path, max_chars=max_chars) + + def grep_repo(pattern: str) -> str: + """Search the repository for a regular expression; returns path:line: text, one per line.""" + rows = hunter_tools.grep_repo(root, pattern, max_matches=max_matches) + return "\n".join(f"{row.get('path')}:{row.get('line')}: {row.get('text', '')}" for row in rows) or "(no matches)" + + def find_refs(symbol: str) -> str: + """Find where a symbol is referenced in the repository; returns path:line per reference.""" + rows = hunter_tools.find_refs(root, symbol) + return "\n".join(f"{row.get('path')}:{row.get('line')}" for row in rows) or "(no references)" + + return [dspy.Tool(read_file), dspy.Tool(grep_repo), dspy.Tool(find_refs)] + + +class HuntProgram(dspy.Module): + """One hunt over one change. With tools it is a bounded ReAct loop; without, one prediction (tests, dry runs).""" + + def __init__(self, *, mask: Iterable[str] = (), tools: Sequence[dspy.Tool] = (), max_iters: int = 8) -> None: + super().__init__() + self.mask = tuple(sorted(set(mask))) + signature = signature_for(self.mask) + self.hunt = dspy.ReAct(signature, tools=list(tools), max_iters=max_iters) if tools else dspy.Predict(signature) + + def forward(self, **inputs: str) -> dspy.Prediction: + prediction = self.hunt(**inputs) + findings = check(getattr(prediction, "findings", []) or [], inputs) + return dspy.Prediction(findings=findings, raw=getattr(prediction, "findings", [])) + + +def program_id(program: dspy.Module, *, mask: Iterable[str] = ()) -> str: + """Content address of a program: its saved state plus the mask, so two identical programs share an id.""" + state = program.dump_state() + payload = json.dumps({"state": state, "mask": sorted(set(mask))}, sort_keys=True, default=str) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def build_lm( + model: str = DEFAULT_DETECTOR_MODEL, *, cache_dir: Path | None = None, temperature: float = 0.0, max_tokens: int = 4096 +) -> dspy.LM: + """The research-side model client: DSPy's LM over the DeepSeek endpoint, prices registered from our table. + + The key is read from the environment only; it is never written anywhere. A cache directory per run keeps a + replay of every response next to the run's artifact.""" + key = os.environ.get("DEEPSEEK_API_KEY") + if not key: + raise RuntimeError("DEEPSEEK_API_KEY is not set") + if cache_dir is not None: + os.environ["DSPY_CACHEDIR"] = str(cache_dir) + try: + import litellm + + prices = _PRICES.get(model) + if prices is not None: + litellm.register_model( + { + f"openai/{model}": { + "input_cost_per_token": prices.input_per_m / 1_000_000, + "output_cost_per_token": prices.output_per_m / 1_000_000, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + except Exception as error: # noqa: BLE001 (a missing price only means cost is reported as unknown) + print(f"price registration skipped: {type(error).__name__}", flush=True) + return dspy.LM(f"openai/{model}", api_base=DEEPSEEK_BASE_URL, api_key=key, temperature=temperature, max_tokens=max_tokens, cache=True) diff --git a/src/openultrasast/hunt/runner.py b/src/openultrasast/hunt/runner.py new file mode 100644 index 0000000..c67618b --- /dev/null +++ b/src/openultrasast/hunt/runner.py @@ -0,0 +1,248 @@ +"""Run a hunt program over a frozen manifest and write records the benchmark scorer reads. + +One record per change, in the shape `benchmarks/valid/hunter_arm.py` writes: the instrument proof (base and head +verified, bytes read per changed file), the defects under `push.admission.defects`, the status. Resumable by +change id, input digest, program id and mask. The pilot manifest is refused by digest unless the run is a pilot +report run, so the loop cannot tune on the held-out set by accident. +""" + +from __future__ import annotations + +import hashlib +import json +import subprocess +import sys +import tempfile +import time +from collections.abc import Iterable, Mapping +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from .. import source +from ..preprocess import detect_language +from .artifact import RunArtifact, UsageLedger +from .labels import label_rows +from .producers import compute_rows +from .table import FeatureTable + +PILOT_FREEZE_DIGEST = "c866c0e812c0b54476c7168fdb0e4a20652de0d30485270afecbe51e5c35e424" + + +class InstrumentFailure(RuntimeError): + """A bounded, secret-free reason the record could not be measured.""" + + +def _git(root: Path, *args: str) -> str: + done = subprocess.run(["git", "-C", str(root), "-c", "core.hooksPath=/dev/null", *args], capture_output=True, check=True) + return done.stdout.decode("utf-8", errors="replace") + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".partial") + tmp.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n") + tmp.replace(path) + + +def record_path(out: Path, change_id: str) -> Path: + return Path(out) / (hashlib.sha256(change_id.encode("utf-8")).hexdigest() + ".json") + + +def input_digest(change: Mapping[str, Any], program_id: str, mask: Iterable[str]) -> str: + payload = json.dumps({"change": change, "program": program_id, "mask": sorted(set(mask))}, sort_keys=True) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def refuse_pilot(manifest: Path, *, pilot_report: bool) -> None: + """The held-out pilot is never a training or tuning input (requirement 2.4).""" + freeze = Path(manifest).with_name("freeze-" + Path(manifest).stem.removeprefix("pool-") + ".json") + if freeze.is_file(): + digest = json.loads(freeze.read_text()).get("freeze_digest") + if digest == PILOT_FREEZE_DIGEST and not pilot_report: + raise ValueError("the pilot manifest is held out; pass pilot_report=True only for a report run") + + +def defect_from(finding: Mapping[str, Any], *, program_id: str, mask: tuple[str, ...]) -> dict[str, Any]: + """The record's defect shape, with the program and mask that produced it.""" + path, line = str(finding.get("path", "")), int(finding.get("line", 0) or 0) + return { + "defect_id": f"hunt:{program_id[:12]}:{path}:{line}", + "family": finding.get("family", "unknown"), + "cwe": None, + "function": finding.get("function"), + "locations": [[path, line]], + "source": "hunt", + "title": f"{finding.get('family', 'unknown')} at {path}:{line}", + "rationale": finding.get("rationale", ""), + "confidence": finding.get("confidence"), + "cited_lines": list(finding.get("cited_lines") or []), + "reported_family": finding.get("reported_family"), + "unknown_reason": finding.get("unknown_reason"), + "program_id": program_id, + "mask": list(mask), + } + + +def hunt_change( + change: Mapping[str, Any], repo_root: Path, *, program: Any, program_id: str, mask: tuple[str, ...], render: Any, tools_for: Any +) -> dict[str, Any]: + """Check out the change's head, compute its feature rows, run the program once, return the record.""" + instrument: dict[str, Any] = { + "hook_exit": 0, + "hook_bytes_read": 0, + "changed_bytes": {}, + "head_verified": False, + "base_verified": False, + "failure": None, + } + record: dict[str, Any] = { + "id": change["id"], + "done": True, + "status": "hunt_covered", + "instrument": instrument, + "push": {"admission": {"defects": [], "dispositions": []}, "quick_tier": []}, + "rows": 0, + "program_id": program_id, + "mask": list(mask), + "inputs_present": [], + "seconds": 0.0, + } + started = time.monotonic() + try: + _git(repo_root, "cat-file", "-e", change["base"] + "^{commit}") + instrument["base_verified"] = True + with tempfile.TemporaryDirectory(prefix="ousast-hunt-") as temporary: + root = Path(temporary) / "head" + added = False + try: + _git(repo_root, "worktree", "add", "--detach", str(root), change["head"]) + added = True + if _git(root, "rev-parse", "HEAD").strip() != change["head"]: + raise InstrumentFailure("head_mismatch") + instrument["head_verified"] = True + texts: dict[str, list[str]] = {} + for relative in dict.fromkeys(file["path"] for file in change.get("files", [])): + target = root / relative + if not target.is_file(): + continue + data = target.read_bytes() + print(f"{change['id']} {relative}: {len(data)} bytes read", flush=True) + instrument["changed_bytes"][relative] = len(data) + instrument["hook_bytes_read"] += len(data) + texts[relative] = data.decode("utf-8", errors="replace").splitlines() + if not instrument["changed_bytes"]: + raise InstrumentFailure("no_readable_hotspots") + rows = label_rows(compute_rows(root, change), change) + FeatureTable(rows) # validates the schema before anything is rendered + record["rows"] = len(rows) + inputs = render(change, rows, texts, mask=mask) + record["inputs_present"] = sorted(inputs) + if tools_for is not None: + program.hunt.tools = ( + {tool.name: tool for tool in tools_for(root)} if hasattr(program.hunt, "tools") else program.hunt.tools + ) + prediction = program(**inputs) + for finding in prediction.findings: + if finding.get("path") in texts and not finding.get("function") and (span := _function_at(root, finding)): + finding["function"] = span + record["push"]["admission"]["defects"].append(defect_from(finding, program_id=program_id, mask=mask)) + finally: + if added: + _git(repo_root, "worktree", "remove", "--force", str(root)) + except Exception as error: # noqa: BLE001 (provider messages can carry a key: type only) + reason = str(error) if isinstance(error, InstrumentFailure) else type(error).__name__ + instrument.update(hook_exit=2, failure=reason) + record["status"] = "instrument_failure" + record["push"]["admission"]["defects"] = [] + print(f"{change['id']}: instrument_failure ({reason})", file=sys.stderr, flush=True) + record["seconds"] = round(time.monotonic() - started, 3) + return record + + +def _function_at(root: Path, finding: Mapping[str, Any]) -> str | None: + try: + path = str(finding["path"]) + text = (root / path).read_text(errors="replace") + language = source.canonical_language(detect_language(root / path)) + return source.enclosing(source.ranges(path, text, language), int(finding.get("line", 0) or 0)) + except Exception: # noqa: BLE001 + return None + + +def run( + manifest: Path, + repos: Mapping[str, str], + out: Path, + *, + program: Any, + program_id: str, + mask: Iterable[str] = (), + only: str = "all", + slice_id: str = "1", + budget: Mapping[str, float] | None = None, + pilot_report: bool = False, + render: Any = None, + tools_for: Any = None, + load_changes: Any = None, + versions: Mapping[str, str] | None = None, +) -> RunArtifact | None: + """Run the program over every change of the manifest's slice; resumable; one artifact at the end.""" + manifest, out = Path(manifest), Path(out) + refuse_pilot(manifest, pilot_report=pilot_report) + mask = tuple(sorted(set(mask))) + if render is None: + from .program import render_rows as render # noqa: PLC0415 (dspy import stays optional) + changes = load_changes(manifest, slice_id) + changes = [change for change in changes if only == "all" or change["kind"] == only] + out.mkdir(parents=True, exist_ok=True) + inputs_path = out / "inputs.json" + entries = json.loads(inputs_path.read_text()) if inputs_path.is_file() else [] + indexed = {entry["item_id"]: entry for entry in entries} + ledger = UsageLedger() + started = time.monotonic() + for change in changes: + digest = input_digest(change, program_id, mask) + entry = indexed.get(change["id"]) + if entry is None: + entry = {"item_id": change["id"], "change": change, "input_digest": digest} + entries.append(entry) + indexed[change["id"]] = entry + elif entry["input_digest"] != digest: + entry.update(change=change, input_digest=digest) + path = record_path(out, change["id"]) + if path.is_file(): + previous = json.loads(path.read_text()) + if previous.get("done") and previous.get("input_digest") == digest: + continue + if budget and ledger.exhausted(budget): + print("budget exhausted; stopping before the next change", flush=True) + break + record = hunt_change( + change, + Path(repos[change["repo"]]).expanduser(), + program=program, + program_id=program_id, + mask=mask, + render=render, + tools_for=tools_for, + ) + record["input_digest"] = digest + ledger.add({"calls": 1, "usd": 0.0}) + _write_json(path, record) + _write_json(inputs_path, entries) + _write_json(inputs_path, entries) + usage = dict(ledger.usage) + usage["wall_seconds"] = round(time.monotonic() - started, 1) + artifact = RunArtifact( + program_id=program_id, + schema_version=1, + mask=list(mask), + data_digests={"manifest": hashlib.sha256(manifest.read_bytes()).hexdigest()}, + configurations=[{"only": only, "slice": slice_id, "changes": len(changes)}], + usage=usage, + versions=dict(versions or {}), + created=datetime.now(UTC).isoformat(timespec="seconds"), + ) + written = artifact.write(out, budget=budget) if budget else artifact.write(out) + return artifact if written is not None else None diff --git a/src/openultrasast/hunt/table.py b/src/openultrasast/hunt/table.py new file mode 100644 index 0000000..92f85f7 --- /dev/null +++ b/src/openultrasast/hunt/table.py @@ -0,0 +1,97 @@ +"""Versioned rows of repository features, labels and producer provenance.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Iterable +from copy import deepcopy +from pathlib import Path +from typing import Any + +SCHEMA_VERSION = 1 +KEY_COLUMNS = ("change_id", "repo", "head", "path", "line", "kind", "fold") +LABEL_COLUMNS = ("label", "family", "fix_span", "label_ref", "verification") +FEATURE_NAMES = { + "diff": ("hunks", "on_changed_line", "distance"), + "source": ("function", "span", "callers", "entry_points", "language"), + "semantic": ("roles", "def_use", "obligations"), + "taxonomy": ("families",), +} +PROVENANCE_NAMES = ("ok", "reason", "version", "bytes_read") +FEATURE_COLUMNS = tuple(f"feature.{producer}.{name}" for producer, names in FEATURE_NAMES.items() for name in names) +PROVENANCE_COLUMNS = tuple(f"provenance.{producer}.{name}" for producer in FEATURE_NAMES for name in PROVENANCE_NAMES) + + +def canonical_json(value: object) -> str: + """Serialize data with stable keys and no ambiguous non-finite numbers.""" + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + + +class FeatureTable: + """Hold flat schema-v1 rows; JSONL starts with a version and mask header.""" + + schema_version = SCHEMA_VERSION + + def __init__(self, rows: Iterable[dict[str, Any]], *, masked: Iterable[str] = ()) -> None: + self.rows = deepcopy(list(rows)) + self.masked = tuple(sorted(set(masked))) + self._validate() + + def _validate(self) -> None: + allowed = set(KEY_COLUMNS + LABEL_COLUMNS + FEATURE_COLUMNS + PROVENANCE_COLUMNS) + for row in self.rows: + if set(row) - allowed: + raise ValueError("unknown feature table column") + label = row.get("label") + if label is not None and (type(label) is not int or label not in (0, 1)): + raise ValueError("label must be 1, 0 or None") + for column in row: + if column.startswith("feature."): + producer = column.split(".")[1] + if any(f"provenance.{producer}.{name}" not in row for name in PROVENANCE_NAMES): + raise ValueError("feature column is missing producer provenance") + canonical_json(self.rows) + + def write_jsonl(self, path: Path | str) -> None: + self._validate() + with Path(path).open("w", encoding="utf-8") as stream: + stream.write(canonical_json({"schema_version": self.schema_version, "masked": self.masked}) + "\n") + for row in self.rows: + stream.write(canonical_json(row) + "\n") + + @classmethod + def read_jsonl(cls, path: Path | str) -> FeatureTable: + with Path(path).open(encoding="utf-8") as stream: + header = json.loads(next(stream, "{}")) + if header.get("schema_version") != SCHEMA_VERSION: + raise ValueError("unsupported feature table schema version") + return cls((json.loads(line) for line in stream if line.strip()), masked=header.get("masked", ())) + + def to_pandas(self) -> Any: + """Create a dataframe only when the optional research dependency is installed.""" + try: + import pandas + except ImportError as error: + raise RuntimeError("pandas is required; install openultrasast[hunt]") from error + return pandas.DataFrame(deepcopy(self.rows)) + + def digest(self) -> str: + self._validate() + payload = {"schema_version": self.schema_version, "rows": self.rows} + return hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest() + + def mask(self, drop: Iterable[str]) -> FeatureTable: + names = set(drop) + unknown = names - set(FEATURE_NAMES) - set(FEATURE_COLUMNS) + if unknown: + raise ValueError("mask names an unknown producer or feature column") + columns = {c for c in FEATURE_COLUMNS if c in names or c.split(".")[1] in names} + return FeatureTable(({k: v for k, v in row.items() if k not in columns} for row in self.rows), + masked=(*self.masked, *names)) + + def feature_columns(self) -> list[str]: + return sorted({column for row in self.rows for column in row if column.startswith("feature.")}) + + def labelled(self) -> FeatureTable: + return FeatureTable((row for row in self.rows if row.get("label") is not None), masked=self.masked) diff --git a/tests/test_hunt_artifact.py b/tests/test_hunt_artifact.py new file mode 100644 index 0000000..a8d8cb4 --- /dev/null +++ b/tests/test_hunt_artifact.py @@ -0,0 +1,49 @@ +"""Checks for stable run records and budget accounting.""" + +import json + +import pytest + +from openultrasast.hunt.artifact import RunArtifact, UsageLedger + + +def artifact(): + return RunArtifact(program_id="program", schema_version=1, mask=("source",), data_digests={"train": "abc"}, + configurations=[{"seed": 0}], usage={"usd": 0.1, "calls": 1, "replay_hits": 0, "wall_seconds": 2}, + versions={"dspy": "3.4.0", "litellm": "1", "openultrasast": "2.0.1"}, created="2026-10-13T00:00:00Z") + + +def test_artifact_round_trip_and_tamper(tmp_path): + run = artifact() + assert run.artifact_id == artifact().artifact_id + path = run.write(tmp_path) + assert RunArtifact.read(path) == run + payload = json.loads(path.read_text()) + payload["program_id"] = "tampered" + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="id"): + RunArtifact.read(path) + + +def test_ledger_and_budget_prevent_artifact(tmp_path): + run = artifact() + ledger = UsageLedger(artifact=run) + ledger.add({"usd": 0.25, "calls": 1, "wall_seconds": 2}) + ledger.add({"usd": 0, "replay_hits": 1, "wall_seconds": 1}) + assert ledger.usage == {"usd": 0.25, "calls": 1, "replay_hits": 1, "wall_seconds": 3} + assert ledger.exhausted({"usd": 0.25, "calls": 10}) + assert ledger.exhausted({"usd": 10, "calls": 1}) + assert ledger.finish({"usd": 0.25}) is None + assert not list(tmp_path.iterdir()) + result = ledger.finish({"usd": 1, "calls": 2}) + assert result is run + assert result.usage == ledger.usage + assert run.finish({"usd": 0.25}) is None + assert run.write(tmp_path, budget={"usd": 0.25}) is None + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("value", [-1, float("nan"), float("inf")]) +def test_invalid_usage_rejected(value): + with pytest.raises(ValueError): + UsageLedger().add({"usd": value}) diff --git a/tests/test_hunt_labels.py b/tests/test_hunt_labels.py new file mode 100644 index 0000000..bdf6c93 --- /dev/null +++ b/tests/test_hunt_labels.py @@ -0,0 +1,47 @@ +"""Checks for fix-site and ordinary labels.""" + +import pytest + +from openultrasast.hunt.labels import label_rows +from openultrasast.hunt.table import FeatureTable + + +def positive(): + return {"kind": "vulnerability", "path": "app.py", "function": "handle", "function_lines": [10, 20], + "family": "injection", "case_id": "verified-case"} + + +def test_labels_match_name_or_span_on_same_path(): + rows = [{"path": "app.py", "function": "handle", "line": 1}, + {"path": "app.py", "function": "nested", "line": 12}, + {"path": "other.py", "function": "handle", "line": 12}, + {"path": "app.py", "function": "other", "line": 30}] + labelled = label_rows(rows, positive()) + assert [r["label"] for r in labelled] == [1, 1, None, None] + for data in labelled[:2]: + assert data["family"] == "injection" + assert data["fix_span"] == [10, 20] + assert data["label_ref"] == "verified-case" + assert data["verification"] is None + assert "label" not in rows[0] + + +def test_feature_function_and_ordinary_rows(): + rows = [{"path": "app.py", "line": 1, "feature.source.function": "handle"}] + assert label_rows(rows, positive())[0]["label"] == 1 + ordinary = label_rows([{"path": "app.py", "line": 1}], {"kind": "ordinary", "case_id": "draw"}) + assert ordinary[0]["label"] == 0 + assert len(FeatureTable(ordinary).labelled().rows) == 1 + + +@pytest.mark.parametrize("kind", [None, "fix", "unknown"]) +def test_unknown_change_kind_rejected(kind): + with pytest.raises(ValueError): + label_rows([], {"kind": kind}) + + +def test_positive_requires_traceable_reference(): + change = positive() + del change["case_id"] + with pytest.raises(ValueError, match="case_id"): + label_rows([], change) diff --git a/tests/test_hunt_metric.py b/tests/test_hunt_metric.py new file mode 100644 index 0000000..ba2a4c8 --- /dev/null +++ b/tests/test_hunt_metric.py @@ -0,0 +1,46 @@ +"""Checks that search scoring follows the benchmark's location rule.""" + +from types import SimpleNamespace + +import pytest + +from benchmarks.unseen.score import catches +from openultrasast.hunt.metric import search_metric, search_outcome + + +@pytest.mark.parametrize("change", [ + {"family": "injection", "path": "app.py", "function": "handler", "function_lines": [100, 140]}, + {"family": "path", "path": "other.py", "function": "read", "function_lines": [100, 140]}, + {"family": "config_secrets", "path": "app.py", "function": "handler", "function_lines": [100, 140]}, +]) +def test_matches_score_catches(change): + change = {**change, "kind": "vulnerability", "files": [{"path": "app.py", "head_lines": [[20, 22]]}]} + for path in ("app.py", "other.py"): + for line in (9, 10, 22, 32, 33, 110, 500): + for function in ("handler", "other"): + for family in (change["family"], "unknown"): + finding = {"path": path, "line": line, "function": function, "family": family} + record = {**change, "push": {"admission": {"defects": [ + {**finding, "locations": [[path, line]]} + ]}}} + assert search_outcome(change, [finding]).hit == bool(catches(record)) + + +def test_counts_scope_feedback_and_tolerance(): + change = {"kind": "vulnerability", "path": "app.py", "function": "handle", "family": "path", + "function_lines": [20, 22], "files": [{"path": "app.py", "head_lines": [[20, 22]]}]} + findings = [{"path": "app.py", "line": 10, "family": "path"}, + {"path": "other.py", "line": 30, "family": "path"}] + outcome = search_outcome(change, findings) + assert (outcome.hit, outcome.suspicions, outcome.on_diff) == (True, 2, 1) + assert "other.py:30" in outcome.feedback + missed = search_outcome(change, findings, tolerance=0) + assert not missed.hit + assert "missed app.py:20 handle family path" in missed.feedback + assert search_metric({"change": change}, {"findings": findings}) == 1.0 + assert search_metric(SimpleNamespace(change=change), SimpleNamespace(findings=[])) == 0.0 + + +@pytest.mark.parametrize("count,expected", [(0, 1.0), (1, 0.8), (5, 0.0), (7, 0.0)]) +def test_ordinary_metric(count, expected): + assert search_metric({"kind": "ordinary", "files": []}, {"findings": [{}] * count}) == expected diff --git a/tests/test_hunt_producers.py b/tests/test_hunt_producers.py new file mode 100644 index 0000000..611db3a --- /dev/null +++ b/tests/test_hunt_producers.py @@ -0,0 +1,158 @@ +"""Checks for repository-only annotations on real local git changes.""" + +import json +import subprocess + +import pytest + +from openultrasast import mapping, source +from openultrasast.hunt import producers as module +from openultrasast.hunt.producers import PRODUCERS, Context, candidate_rows, compute_rows +from openultrasast.hunt.table import FeatureTable +from openultrasast.model.taxonomy import FAMILY_IDS + + +@pytest.fixture +def repository(tmp_path): + def git(*args): + return subprocess.run(["git", "-C", str(tmp_path), *args], check=True, capture_output=True, text=True).stdout.strip() + + git("init", "-q") + git("config", "user.name", "Fixture") + git("config", "user.email", "fixture@example.invalid") + initial = ('def first():\n value = input()\n return value\n\n' + 'def second():\n value = first()\n return value\n') + (tmp_path / "app.py").write_text(initial) + (tmp_path / "notes.txt").write_text("GHSA-xxxx-xxxx-xxxx CVE-2024-0001\n") + (tmp_path / "unknown.zzz").write_text("opaque\n") + git("add", "app.py", "notes.txt", "unknown.zzz") + git("commit", "-qm", "initial") + base = git("rev-parse", "HEAD") + (tmp_path / "app.py").write_text(initial.replace("return value", "return eval(value)", 1)) + (tmp_path / "unknown.zzz").write_text("opaque changed\n") + git("add", "app.py", "unknown.zzz") + git("commit", "-qm", "fix XSS CVE-2024-0001") + head = git("rev-parse", "HEAD") + data = {"id": "change", "repo": "fixture/repo", "base": base, "head": head, "kind": "vulnerability", "fold": 1, + "path": "app.py", "function": "second", "function_lines": [5, 7], "family": "injection", "case_id": "verified", + "advisory": "GHSA-xxxx-xxxx-xxxx", "commit_message": "fix XSS CVE-2024-0001", + "files": [{"path": "app.py", "head_lines": [[3, 3]]}]} + return tmp_path, data + + +def test_every_column_candidates_and_cached_work(repository, monkeypatch): + root, change = repository + calls = [] + real = mapping.analyze_entry_points + + def entries(*args): + calls.append(1) + return real(*args) + + monkeypatch.setattr(mapping, "analyze_entry_points", entries) + candidates = candidate_rows(root, change) + assert [r["line"] for r in candidates] == [1, 5] + rows = compute_rows(root, change) + assert len(calls) == 1 + table = FeatureTable(rows) + assert len(table.feature_columns()) == 12 + assert [r["label"] for r in rows] == [None, 1] + first, second = rows + assert first["feature.diff.hunks"] == [[3, 3]] + assert first["feature.diff.distance"] == 2 + assert not first["feature.diff.on_changed_line"] + assert first["feature.source.function"] == "first" + assert first["feature.source.span"] == [1, 3] + assert first["feature.source.callers"] == [["app.py", 6]] + assert first["feature.source.language"] == "python" + assert isinstance(first["feature.source.entry_points"], list) + assert first["feature.semantic.roles"]["input"] == "source" + assert first["feature.semantic.roles"]["eval"] == "sink" + assert first["feature.semantic.def_use"]["binds"] == [{"line": 2, "name": "value", "names": ["input"], "call_name": "input"}] + assert {c["name"] for c in first["feature.semantic.def_use"]["calls"]} == {"input", "eval"} + assert isinstance(first["feature.semantic.obligations"], list) + assert [f["id"] for f in first["feature.taxonomy.families"]] == list(FAMILY_IDS) + assert second["feature.diff.distance"] == 2 + for row in rows: + for producer in PRODUCERS: + assert set(producer.columns) <= set(row) + assert row[f"provenance.{producer.name}.ok"] is True + assert row[f"provenance.{producer.name}.reason"] is None + assert row[f"provenance.{producer.name}.version"] + assert row[f"provenance.{producer.name}.bytes_read"] > 0 + ctx = Context(root, change) + diff = PRODUCERS[0] + assert diff.compute(root, change, "app.py", 3, ctx)["feature.diff.on_changed_line"] + assert diff.compute(root, change, "notes.txt", 1, ctx)["feature.diff.distance"] is None + + +def test_unsupported_language_is_absent_with_reason(repository): + root, change = repository + change = {**change, "kind": "ordinary", "files": [{"path": "unknown.zzz", "head_lines": [[1, 1]]}]} + rows = compute_rows(root, change) + assert len(rows) == 1 + assert rows[0]["provenance.semantic.ok"] is False + assert rows[0]["provenance.semantic.reason"] == "language_unsupported" + assert all(rows[0][column] is None for column in PRODUCERS[2].columns if column.startswith("feature.")) + assert rows[0]["label"] == 0 + FeatureTable(rows) + + +def test_planted_advisory_and_commit_text_never_enter_features(repository): + root, change = repository + rows = compute_rows(root, change) + values = json.dumps([{k: v for k, v in row.items() if k.startswith("feature.")} for row in rows]) + for planted in ("GHSA-xxxx-xxxx-xxxx", "CVE-2024-0001", "fix XSS"): + assert planted not in values + + +def test_producer_exceptions_are_type_only(repository, monkeypatch): + root, change = repository + + def failed(*args, **kwargs): + raise RuntimeError("secret credential") + + monkeypatch.setattr(source, "facts_for", failed) + row = compute_rows(root, change)[0] + assert row["provenance.source.reason"] == "RuntimeError" + assert row["provenance.source.ok"] is False + assert all(row[c] is None for c in PRODUCERS[1].columns if c.startswith("feature.")) + assert "secret credential" not in json.dumps(row) + + +def test_range_rejection_retains_row_and_reason(repository, monkeypatch): + root, change = repository + + def failed(*args): + raise ValueError("private text") + + monkeypatch.setattr(source, "ranges", failed) + rows = compute_rows(root, {**change, "kind": "ordinary"}) + assert len(rows) == 1 + assert rows[0]["provenance.source.reason"] == "ValueError" + assert rows[0]["provenance.source.ok"] is False + + +def test_private_or_outside_paths_are_not_read(repository): + root, change = repository + for path in (".env", "../outside.py"): + with pytest.raises(ValueError): + candidate_rows(root, {**change, "files": [{"path": path, "head_lines": [[1, 1]]}]}) + + +def test_head_mismatch_is_not_a_clean_measurement(repository): + root, change = repository + rows = compute_rows(root, {**change, "head": change["base"], "kind": "ordinary"}) + assert rows[0]["provenance.diff.ok"] is False + assert rows[0]["provenance.diff.reason"] == "head_mismatch" + + +def test_obligation_result_uses_dominance_not_absence_of_findings(repository, monkeypatch): + root, change = repository + text = ('from flask import request\n@app.route("/item")\ndef first():\n' + ' return User.query.get(request.args["id"])\n') + (root / "app.py").write_text(text) + monkeypatch.setattr(module, "_git", lambda *args: change["head"]) + row = compute_rows(root, {**change, "kind": "ordinary", "files": [{"path": "app.py", "head_lines": [[4, 4]]}]})[0] + assert row["feature.semantic.obligations"] + assert row["feature.semantic.obligations"][0]["discharged"] is False diff --git a/tests/test_hunt_program.py b/tests/test_hunt_program.py new file mode 100644 index 0000000..d4a4f51 --- /dev/null +++ b/tests/test_hunt_program.py @@ -0,0 +1,97 @@ +"""The hunt program with a scripted model: signature masks, rendering, the answer rules, the program id.""" + +from __future__ import annotations + +import json + +import pytest + +dspy = pytest.importorskip("dspy") +from dspy.utils import DummyLM # noqa: E402 + +from openultrasast.hunt import program as hp # noqa: E402 + +CHANGE = {"id": "c1", "kind": "vulnerability", "path": "app.py", "function": "second", "family": "injection"} +ROWS = [ + { + "change_id": "c1", + "path": "app.py", + "line": 5, + "feature.diff.hunks": [[3, 3]], + "feature.source.function": "second", + "feature.source.span": [5, 7], + "feature.source.callers": [["web.py", 2]], + "feature.source.entry_points": [], + "feature.semantic.roles": {"eval": "sink"}, + "feature.semantic.def_use": {"binds": [], "calls": [{"line": 6, "name": "eval"}]}, + "feature.semantic.obligations": [], + "feature.taxonomy.families": [{"id": "injection", "definition": "untrusted data reaches an interpreter"}], + } +] +TEXT = {"app.py": ["def first():", " v = input()", " return v", "", "def second():", " return eval(first())", ""]} + + +def test_signature_drops_masked_inputs_and_keeps_the_order(): + full = hp.signature_for() + assert list(full.input_fields) == [name for name, _, _ in hp.INPUTS] + masked = hp.signature_for(["semantic", "callers"]) + assert "roles" not in masked.input_fields and "def_use" not in masked.input_fields and "callers" not in masked.input_fields + assert list(masked.input_fields) == ["change_summary", "hunks", "function_text", "entry_points", "families"] + assert "findings" in masked.output_fields + + +def test_render_rows_numbers_the_function_and_omits_masked_inputs(): + inputs = hp.render_rows(CHANGE, ROWS, TEXT) + assert inputs["change_summary"] == "app.py::second" + assert json.loads(inputs["hunks"]) == {"app.py": [[3, 3]]} + assert " 5 def second():" in inputs["function_text"] and " 6 " in inputs["function_text"] + assert "input()" not in inputs["function_text"] # only the span is shown + assert inputs["families"].startswith("injection: ") + masked = hp.render_rows(CHANGE, ROWS, TEXT, mask=["diff", "function_text"]) + assert "hunks" not in masked and "function_text" not in masked and "roles" in masked + assert hp.shown_lines(inputs) == {"app.py": {5, 6, 7}} + + +def test_check_turns_off_list_families_and_unseen_citations_into_unknown(): + inputs = hp.render_rows(CHANGE, ROWS, TEXT) + good = {"path": "app.py", "line": 6, "family": "injection", "confidence": 0.9, "rationale": "eval", "cited_lines": [6]} + bad_family = {**good, "family": "sql-parser"} + bad_cite = {**good, "cited_lines": [2]} + out = hp.check([good, bad_family, bad_cite], inputs) + assert out[0]["family"] == "injection" and "unknown_reason" not in out[0] + assert out[1]["family"] == "unknown" and out[1]["reported_family"] == "sql-parser" and out[1]["unknown_reason"] == "family_not_in_list" + assert out[2]["family"] == "unknown" and out[2]["unknown_reason"] == "cited_line_not_shown" + + +def test_program_without_tools_runs_on_a_scripted_model_and_has_a_stable_id(): + answer = [{"path": "app.py", "line": 6, "family": "injection", "confidence": 0.8, "rationale": "eval of input", "cited_lines": [6]}] + lm = DummyLM([{"findings": json.dumps(answer)}, {"findings": json.dumps(answer)}]) + dspy.configure(lm=lm, adapter=dspy.ChatAdapter()) + program = hp.HuntProgram(mask=["obligations"]) + inputs = hp.render_rows(CHANGE, ROWS, TEXT, mask=program.mask) + prediction = program(**inputs) + assert prediction.findings == [{**answer[0]}] + assert prediction.findings[0]["family"] == "injection" + assert hp.program_id(program, mask=program.mask) == hp.program_id(hp.HuntProgram(mask=["obligations"]), mask=["obligations"]) + assert hp.program_id(program) != hp.program_id(program, mask=["obligations"]) + + +def test_tools_are_clamped_to_the_checkout(tmp_path): + (tmp_path / "a.py").write_text("x = 1\n") + tools = {tool.name: tool for tool in hp.repository_tools(tmp_path)} + assert set(tools) == {"read_file", "grep_repo", "find_refs"} + assert "x = 1" in tools["read_file"](path="a.py") + from openultrasast.hunter_tools import PathEscapesRepo + + with pytest.raises(PathEscapesRepo): + tools["read_file"](path="../../etc/passwd") + assert ( + "(no matches)" in tools["grep_repo"](pattern="nothing_here_\\d+") + or tools["grep_repo"](pattern="nothing_here_\\d+") == "(no matches)" + ) + + +def test_build_lm_needs_a_key_from_the_environment(monkeypatch): + monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False) + with pytest.raises(RuntimeError, match="DEEPSEEK_API_KEY"): + hp.build_lm() diff --git a/tests/test_hunt_runner.py b/tests/test_hunt_runner.py new file mode 100644 index 0000000..8dbd343 --- /dev/null +++ b/tests/test_hunt_runner.py @@ -0,0 +1,110 @@ +"""The hunt runner over a fixture repository with a scripted program: records, resumability, the pilot refusal.""" + +from __future__ import annotations + +import json +import subprocess + +import pytest + +dspy = pytest.importorskip("dspy") +from dspy.utils import DummyLM # noqa: E402 + +from benchmarks.unseen import score # noqa: E402 +from openultrasast.hunt import program as hp # noqa: E402 +from openultrasast.hunt import runner # noqa: E402 + + +@pytest.fixture +def repository(tmp_path): + root = tmp_path / "repo" + root.mkdir() + + def git(*args): + return subprocess.run(["git", "-C", str(root), *args], check=True, capture_output=True, text=True).stdout.strip() + + git("init", "-q") + git("config", "user.name", "Fixture") + git("config", "user.email", "fixture@example.invalid") + (root / "app.py").write_text( + "def first():\n value = input()\n return value\n\ndef second():\n value = first()\n return value\n" + ) + git("add", "app.py") + git("commit", "-qm", "initial") + base = git("rev-parse", "HEAD") + (root / "app.py").write_text( + "def first():\n value = input()\n return value\n\ndef second():\n value = first()\n return eval(value)\n" + ) + git("add", "app.py") + git("commit", "-qm", "change") + head = git("rev-parse", "HEAD") + positive = { + "id": "pos", "repo": "fixture/repo", "base": base, "head": head, "kind": "vulnerability", "case_id": "case-1", + "path": "app.py", "function": "second", "function_lines": [5, 7], "family": "injection", + "files": [{"path": "app.py", "head_lines": [[7, 7]]}], + } # fmt: skip + ordinary = { + "id": "ord", + "repo": "fixture/repo", + "base": base, + "head": head, + "kind": "ordinary", + "case_id": "ord-1", + "files": [{"path": "app.py", "head_lines": [[7, 7]]}], + } + manifest = tmp_path / "pool-fixture.toml" + manifest.write_text("status = 'fixture'\n") + return root, manifest, [positive, ordinary] + + +def _scripted_program(): + answer = [{"path": "app.py", "line": 7, "family": "injection", "confidence": 0.9, "rationale": "eval of input", "cited_lines": [7]}] + lm = DummyLM([{"findings": json.dumps(answer)}, {"findings": "[]"}, {"findings": json.dumps(answer)}, {"findings": "[]"}]) + dspy.configure(lm=lm, adapter=dspy.ChatAdapter()) + return hp.HuntProgram(mask=["obligations"]) + + +def test_runner_writes_scorable_records_and_resumes(repository, tmp_path): + root, manifest, changes = repository + program = _scripted_program() + pid = hp.program_id(program, mask=program.mask) + out = tmp_path / "out" + artifact = runner.run( + manifest, {"fixture/repo": str(root)}, out, program=program, program_id=pid, mask=program.mask, load_changes=lambda m, s: changes + ) + assert artifact is not None and artifact.program_id == pid + rows = score.load_records(out) + assert {r["id"]: r["status"] for r in rows} == {"pos": "hunt_covered", "ord": "hunt_covered"} + pos = next(r for r in rows if r["id"] == "pos") + assert pos["instrument"]["head_verified"] and pos["instrument"]["hook_bytes_read"] > 0 + defects = pos["push"]["admission"]["defects"] + assert defects and defects[0]["family"] == "injection" and defects[0]["locations"] == [["app.py", 7]] + assert defects[0]["program_id"] == pid and defects[0]["mask"] == ["obligations"] + assert defects[0]["function"] == "second" + assert score.catches(pos) + assert "obligations" not in pos["inputs_present"] and "function_text" in pos["inputs_present"] + # a second run with the same program and mask touches nothing + before = {p.name: p.read_bytes() for p in out.glob("*.json")} + runner.run( + manifest, {"fixture/repo": str(root)}, out, program=program, program_id=pid, mask=program.mask, load_changes=lambda m, s: changes + ) + assert {p.name: p.read_bytes() for p in out.glob("*.json") if p.name in before} == before + + +def test_runner_refuses_the_pilot_unless_it_is_a_report_run(tmp_path): + manifest = tmp_path / "pool-valid.toml" + manifest.write_text("status = 'frozen'\n") + (tmp_path / "freeze-valid.json").write_text(json.dumps({"freeze_digest": runner.PILOT_FREEZE_DIGEST})) + with pytest.raises(ValueError, match="held out"): + runner.refuse_pilot(manifest, pilot_report=False) + runner.refuse_pilot(manifest, pilot_report=True) + + +def test_instrument_failure_is_recorded_with_a_type_only(repository, tmp_path): + root, manifest, changes = repository + program = _scripted_program() + broken = {**changes[0], "id": "bad", "head": "0" * 40} + record = runner.hunt_change(broken, root, program=program, program_id="p", mask=(), render=hp.render_rows, tools_for=None) + assert record["status"] == "instrument_failure" + assert record["instrument"]["failure"] == "CalledProcessError" + assert record["push"]["admission"]["defects"] == [] diff --git a/tests/test_hunt_table.py b/tests/test_hunt_table.py new file mode 100644 index 0000000..0245784 --- /dev/null +++ b/tests/test_hunt_table.py @@ -0,0 +1,73 @@ +"""Checks for the versioned hunter table.""" + +import builtins +import json + +import pytest + +from openultrasast.hunt.table import SCHEMA_VERSION, FeatureTable + + +def row(label=1): + return { + "change_id": "c", "repo": "r", "head": "h", "path": "app.py", "line": 1, + "kind": "vulnerability", "fold": 0, "label": label, "family": "injection", + "fix_span": [1, 3], "label_ref": "case", "verification": None, + "feature.diff.distance": 0, "provenance.diff.ok": True, + "provenance.diff.reason": None, "provenance.diff.version": "1", "provenance.diff.bytes_read": 10, + } + + +def test_round_trip_digest_and_mask(tmp_path): + table = FeatureTable([row(), row(None), row(0)]) + path = tmp_path / "table.jsonl" + table.write_jsonl(path) + restored = FeatureTable.read_jsonl(path) + assert restored.rows == table.rows + assert restored.digest() == table.digest() + assert SCHEMA_VERSION == 1 + assert len(table.labelled().rows) == 2 + masked = table.mask(["diff"]) + assert masked.masked == ("diff",) + assert masked.feature_columns() == [] + assert table.feature_columns() == ["feature.diff.distance"] + assert masked.digest() != table.digest() + masked.write_jsonl(path) + assert FeatureTable.read_jsonl(path).masked == ("diff",) + assert table.mask(["feature.diff.distance"]).feature_columns() == [] + assert table.digest() == FeatureTable([dict(reversed(list(row().items()))), row(None), row(0)]).digest() + + +@pytest.mark.parametrize("extra", [{"unknown": 1}, {"feature.diff.typo": 1}, {"provenance.diff.secret": 1}]) +def test_unknown_columns_rejected(extra): + with pytest.raises(ValueError): + FeatureTable([{**row(), **extra}]) + + +def test_missing_provenance_rejected(): + data = row() + del data["provenance.diff.ok"] + with pytest.raises(ValueError, match="provenance"): + FeatureTable([data]) + + +def test_schema_version_and_empty_table(tmp_path): + path = tmp_path / "table.jsonl" + FeatureTable([]).write_jsonl(path) + assert FeatureTable.read_jsonl(path).rows == [] + path.write_text(json.dumps({"schema_version": 99, "masked": []}) + "\n") + with pytest.raises(ValueError, match="schema"): + FeatureTable.read_jsonl(path) + + +def test_pandas_is_optional(monkeypatch): + real_import = builtins.__import__ + + def importing(name, *args, **kwargs): + if name == "pandas": + raise ImportError("missing") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", importing) + with pytest.raises(RuntimeError, match="pandas.*hunt"): + FeatureTable([row()]).to_pandas()