diff --git a/benchmarks/measurements/2026-09-08-module-audit.json b/benchmarks/measurements/2026-09-08-module-audit.json index 17b1136..7820271 100644 --- a/benchmarks/measurements/2026-09-08-module-audit.json +++ b/benchmarks/measurements/2026-09-08-module-audit.json @@ -1,7 +1,7 @@ { - "modules": 123, + "modules": 125, "counts": { - "load_bearing": 114, + "load_bearing": 116, "standalone_capability": 9 }, "orphaned": [], @@ -56,6 +56,7 @@ "complexity.ledger", "complexity.map", "map_gate", + "plane.tasks.verify", "regress.candidate", "semantic.prove_filter", "tool_hunter" @@ -90,6 +91,7 @@ "complexity.hints", "complexity.ledger", "map_gate", + "plane.tasks.verify", "regress.candidate", "semantic.prove_filter", "tool_hunter" @@ -655,13 +657,17 @@ "module": "plane", "classification": "load_bearing", "reason": "ai-service-plane Req 1-4: the service plane package; runner and reconciler integration in tasks 3-4", - "importers": [] + "importers": [ + "plane.tasks.verify" + ] }, { "module": "plane.budget", "classification": "load_bearing", "reason": "ai-service-plane Req 2.3/2.4: metered chat client with usd and calls ceilings; task integration in tasks 3, 6", - "importers": [] + "importers": [ + "plane.tasks.verify" + ] }, { "module": "plane.manifests", @@ -672,7 +678,15 @@ { "module": "plane.tasks", "classification": "load_bearing", - "reason": "ai-service-plane Req 4.1: the task entrypoint package the runner imports by name; runner integration pending (task 3)", + "reason": "ai-service-plane Req 2.1: the tasks package, each a main() under the OUSAST_* env contract; runner wiring in task 3", + "importers": [ + "plane.tasks.agree" + ] + }, + { + "module": "plane.tasks.agree", + "classification": "load_bearing", + "reason": "ai-service-plane Req 5.3/2.1: agreement across two verify passes with per-candidate cost, turns and site match", "importers": [] }, { @@ -681,6 +695,14 @@ "reason": "ai-service-plane Req 5.1/5.2: source-only functions and callers per file; runner integration pending", "importers": [] }, + { + "module": "plane.tasks.verify", + "classification": "load_bearing", + "reason": "ai-service-plane Req 5.3/2.1-2.3/7.2: the batched hunt with known callers, metered and resumable per file hunt", + "importers": [ + "plane.tasks.agree" + ] + }, { "module": "policy", "classification": "load_bearing", @@ -1335,7 +1357,8 @@ "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ "model.endpoint", - "plane.budget" + "plane.budget", + "plane.tasks.verify" ] }, { diff --git a/src/openultrasast/model/audit.py b/src/openultrasast/model/audit.py index 9cc8f22..75c628e 100644 --- a/src/openultrasast/model/audit.py +++ b/src/openultrasast/model/audit.py @@ -59,6 +59,9 @@ SPEC_OWNED = { "plane.manifests": "ai-service-plane Req 1.1-1.4: ax-shaped Task/Workspace/Model and the Run kind, validated before anything runs", "plane.tasks": "ai-service-plane Req 4.1: the task entrypoint package the runner imports by name; runner integration pending (task 3)", "plane.tasks.repo_facts": "ai-service-plane Req 5.1/5.2: source-only functions and callers per file; runner integration pending", + "plane.tasks": "ai-service-plane Req 2.1: the tasks package, each a main() under the OUSAST_* env contract; runner wiring in task 3", + "plane.tasks.verify": "ai-service-plane Req 5.3/2.1-2.3/7.2: the batched hunt with known callers, metered and resumable per file hunt", + "plane.tasks.agree": "ai-service-plane Req 5.3/2.1: agreement across two verify passes with per-candidate cost, turns and site match", } STANDALONE = { diff --git a/src/openultrasast/plane/tasks/__init__.py b/src/openultrasast/plane/tasks/__init__.py index e69de29..db95bbd 100644 --- a/src/openultrasast/plane/tasks/__init__.py +++ b/src/openultrasast/plane/tasks/__init__.py @@ -0,0 +1 @@ +"""Tasks of the service plane: each a `main()` driven by the `OUSAST_*` env contract of `design.md`.""" diff --git a/src/openultrasast/plane/tasks/agree.py b/src/openultrasast/plane/tasks/agree.py new file mode 100644 index 0000000..9abc873 --- /dev/null +++ b/src/openultrasast/plane/tasks/agree.py @@ -0,0 +1,195 @@ +"""`agree` task (ai-service-plane task 6): agreement across two `verify` passes, no model calls. + +The rule is `verify_batched.agreement`: a candidate is reported when every pass flagged it; a candidate flagged +by one pass is `disputed`. Beside the verdicts go the measurement's numbers: per-candidate cost (a hunt's usd +split evenly over the candidates it asked about, summed over the passes), tool turns per hunt, and -- when the +case's declared sites are given -- the site match by `benchmarks/independent/evaluate.matches_v2`, the rule +protocol v2 scores by, imported from the benchmark tree rather than copied. + +Env contract: `OUSAST_INPUT_PASS_A`, `OUSAST_INPUT_PASS_B` (a verify output directory or its `units.jsonl`), +`OUSAST_INPUT_CANDIDATES`, `OUSAST_OUTPUT_DIR`, optional `OUSAST_INPUT_SITES` (a JSON case record with `family` +and `sites` in the `population-v2.toml` shape, or that TOML itself with `OUSAST_CASE_ID` naming the case) and +`OUSAST_BENCHMARKS_DIR` (where `independent/evaluate.py` lives when the package is not run from the repository). +Writes `agreed.json`, `disputed.json`, `summary.json`. Exit 0 done, 2 failed, 3 unfinished. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +import tomllib +import traceback +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import Any + +from .verify import EXIT_CODES, USAGE_FIELDS, load_candidates, read_units + +Matcher = Callable[[dict[str, Any], dict[str, Any], dict[str, list[tuple[int, int]]]], bool] + + +def agreement(passes: Sequence[Sequence[Mapping[str, Any]]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """(agreed, disputed): flagged in every pass vs in some; the rule of `verify_batched.agreement`.""" + if not passes: + return [], [] + by_pass = [{str(r["candidate"]): dict(r) for r in flagged} for flagged in passes] + every = set.intersection(*(set(b) for b in by_pass)) + some = set.union(*(set(b) for b in by_pass)) + first = by_pass[0] + pick = lambda c: next(b[c] for b in by_pass if c in b) # noqa: E731 + return [dict(first[c], passes=len(by_pass)) for c in sorted(every)], [ + dict(pick(c), passes=sum(c in b for b in by_pass)) for c in sorted(some - every) + ] + + +def load_pass(source: Path) -> list[dict[str, Any]]: + """The finished unit rows of one verify pass (a directory holding `units.jsonl`, or the file itself).""" + log = source / "units.jsonl" if source.is_dir() else source + return [r for r in read_units(log) if "error" not in r and "flagged" in r] + + +def load_sites(path: Path | None, case_id: str | None) -> dict[str, Any] | None: + """A case record with `family`, `sites` and optional `ranges` (`{path: [[first, last], ...]}`).""" + if path is None: + return None + if path.suffix == ".toml": + cases = tomllib.loads(path.read_text()).get("case", []) + found = [c for c in cases if not case_id or c.get("id") == case_id] + if len(found) != 1: + raise ValueError(f"{path}: {len(found)} cases match id {case_id!r}; set OUSAST_CASE_ID") + return dict(found[0]) + record = json.loads(path.read_text()) + if not isinstance(record, dict): + raise ValueError(f"{path}: a case record (object with family and sites) was expected") + return record + + +def load_matcher(benchmarks_dir: Path | None = None) -> Matcher | None: + """`evaluate.matches_v2` from the benchmark tree, loaded by file so the package never imports it by name; + None when the tree is not present (site matching is then reported as not assessed, never as a miss).""" + roots = [benchmarks_dir] if benchmarks_dir is not None else [] + roots += [p / "benchmarks" for p in (Path.cwd(), *Path(__file__).resolve().parents)] + for root in roots: + script = root / "independent" / "evaluate.py" + if script.is_file(): + spec = importlib.util.spec_from_file_location("ousast_benchmark_evaluate", script) + if spec is None or spec.loader is None: + continue + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + matcher: Matcher = module.matches_v2 + return matcher + return None + + +def _per_candidate_cost(rows: Sequence[Mapping[str, Any]]) -> tuple[dict[str, float | None], dict[str, int]]: + """A hunt's usd split evenly over its candidates; its turns and calls counted once per candidate.""" + usd: dict[str, float | None] = {} + turns: dict[str, int] = {} + for row in rows: + names = [f"{row['path']}::{c[1]}" for c in row.get("candidates", [])] + share = None if row.get("usd") is None or not names else float(row["usd"]) / len(names) + for name in names: + usd[name] = None if share is None or usd.get(name, 0.0) is None else (usd.get(name) or 0.0) + share + turns[name] = turns.get(name, 0) + int(row.get("turns") or 0) + return usd, turns + + +def _mean(values: Sequence[float]) -> float | None: + return round(sum(values) / len(values), 3) if values else None + + +def run( + output_dir: Path, + candidates: object, + pass_a: Sequence[Mapping[str, Any]], + pass_b: Sequence[Mapping[str, Any]], + *, + sites: Mapping[str, Any] | None = None, + matcher: Matcher | None = None, +) -> dict[str, Any]: + output_dir.mkdir(parents=True, exist_ok=True) + family, unique = load_candidates(candidates) + names = [f"{p}::{fn}" for p, fn, _ in unique] + lines = {f"{p}::{fn}": ln for p, fn, ln in unique} + asked = [{f"{r['path']}::{c[1]}" for r in rows for c in r.get("candidates", [])} for rows in (pass_a, pass_b)] + flagged = [{str(f["candidate"]): dict(f) for r in rows for f in r.get("flagged", [])} for rows in (pass_a, pass_b)] + agreed, disputed = agreement([list(flagged[0].values()), list(flagged[1].values())]) + usd_by, turns_by = _per_candidate_cost([*pass_a, *pass_b]) + case = dict(sites) if sites else None + if case is not None and "family" not in case and family: + case["family"] = family + ranges = {p: [(int(a), int(b)) for a, b in spans] for p, spans in (case or {}).get("ranges", {}).items()} + match_state = ( + "not assessed: no sites given" if case is None else ("not assessed: evaluate.py not found" if matcher is None else "assessed") + ) + rows: list[dict[str, Any]] = [] + for name in names: + in_a, in_b = (name in flagged[0] if name in asked[0] else None), (name in flagged[1] if name in asked[1] else None) + finding = flagged[0].get(name) or flagged[1].get(name) + site = finding["site"] if finding else f"{name.partition('::')[0]}:{lines[name]}:{name.partition('::')[2]}" + site_match = None + if case is not None and matcher is not None: + site_match = matcher({"site": site, "family": (finding or {}).get("family", family)}, case, ranges) + cost = usd_by.get(name) + rows.append({ + "candidate": name, "site": site, "a": in_a, "b": in_b, "agreed": bool(in_a and in_b), "disputed": bool(in_a) != bool(in_b), + "site_match": site_match, "usd": (None if cost is None else round(cost, 6)), "turns": turns_by.get(name, 0), "finding": finding, + }) # fmt: skip + unasked = [r["candidate"] for r in rows if r["a"] is None or r["b"] is None] + declared = [r for r in rows if r["site_match"]] + hunts = [*pass_a, *pass_b] + usd_total = None if any(r.get("usd") is None for r in hunts) else round(sum(float(r["usd"]) for r in hunts), 6) + metrics = { + "candidates": len(names), "asked_in_both": len(names) - len(unasked), "flagged_a": len(flagged[0]), "flagged_b": len(flagged[1]), + "agreed": len(agreed), "disputed": len(disputed), "hunts": len(hunts), + "declared_sites": len(set((case or {}).get("sites", []))), "declared_sites_matched": len(declared), + "declared_sites_matched_agreed": sum(1 for r in declared if r["agreed"]), "site_matching": match_state, + "usd_total": usd_total, "cost_per_candidate": (None if usd_total is None or not names else round(usd_total / len(names), 6)), + "turns_per_hunt": _mean([float(r.get("turns") or 0) for r in hunts]), + "calls_per_hunt": _mean([float((r.get("usage") or {}).get("calls") or 0) for r in hunts]), + } # fmt: skip + (output_dir / "agreed.json").write_text( + json.dumps({"family": family, "candidates": rows, "agreed": agreed, "disputed": disputed, "metrics": metrics}, indent=1) + "\n" + ) + (output_dir / "disputed.json").write_text(json.dumps(disputed, indent=1) + "\n") + summary: dict[str, Any] = { + "status": "done" if not unasked else "unfinished", "units_done": len(names) - len(unasked), "units_total": len(names), + "usd": 0, "calls": 0, "usage": dict.fromkeys(USAGE_FIELDS, 0), "priced": True, "model": None, "pass": None, "metrics": metrics, + } # fmt: skip + if unasked: + summary["reason"] = f"{len(unasked)} candidates not answered by both passes: {unasked[:5]}" + (output_dir / "summary.json").write_text(json.dumps(summary, indent=1) + "\n") + return summary + + +def main() -> int: + output_dir = Path(os.environ["OUSAST_OUTPUT_DIR"]) if os.environ.get("OUSAST_OUTPUT_DIR") else None + try: + if output_dir is None: + raise ValueError("OUSAST_OUTPUT_DIR is not set") + pass_a, pass_b = (load_pass(Path(os.environ[name])) for name in ("OUSAST_INPUT_PASS_A", "OUSAST_INPUT_PASS_B")) + candidates = json.loads(Path(os.environ["OUSAST_INPUT_CANDIDATES"]).read_text()) + sites_path = Path(os.environ["OUSAST_INPUT_SITES"]) if os.environ.get("OUSAST_INPUT_SITES") else None + sites = load_sites(sites_path, os.environ.get("OUSAST_CASE_ID")) + benchmarks = Path(os.environ["OUSAST_BENCHMARKS_DIR"]) if os.environ.get("OUSAST_BENCHMARKS_DIR") else None + summary = run(output_dir, candidates, pass_a, pass_b, sites=sites, matcher=load_matcher(benchmarks) if sites else None) + except Exception: # noqa: BLE001 -- a crash is `failed` with its traceback + summary = { + "status": "failed", "units_done": 0, "units_total": 0, "usd": 0, "calls": 0, "usage": dict.fromkeys(USAGE_FIELDS, 0), + "priced": True, "model": None, "pass": None, "reason": traceback.format_exc()[-2000:], + } # fmt: skip + if output_dir is not None: + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "summary.json").write_text(json.dumps(summary, indent=1) + "\n") + print(json.dumps({k: summary.get(k) for k in ("status", "units_done", "units_total", "metrics")}), file=sys.stderr) + return EXIT_CODES.get(str(summary["status"]), 2) + + +__all__ = ["agreement", "load_matcher", "load_pass", "load_sites", "main", "run"] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/openultrasast/plane/tasks/verify.py b/src/openultrasast/plane/tasks/verify.py new file mode 100644 index 0000000..494f9b8 --- /dev/null +++ b/src/openultrasast/plane/tasks/verify.py @@ -0,0 +1,373 @@ +"""`verify` task (ai-service-plane task 6): the batched sink hunt of `benchmarks/independent/verify_batched.py` +moved into the package, metered and resumable. + +One tool hunt per file for up to `PER_HUNT` candidates and `MAX_STEPS` steps; a verdict is anchored to the +candidate asked about (its file, function and line), the model's own location kept beside it. Two things are +new against the reference: each candidate's known call sites from `repo-facts` (Requirement 5.3) go into the +prompt, and every hunt records its tool turns and usage in `units.jsonl` (one row per file hunt, appended as it +completes, read back on start: Requirement 2.2). The chat client is wrapped by `MeteredClient` (Requirement +2.3/2.4): a ceiling ends the task `unfinished`, an account refusal ends it `failed`, and `summary.json` -- written +last -- names the bound Model (Requirement 7.2) and is `done` only when every unit finished. + +Env contract: `OUSAST_WORKSPACE_DIR`, `OUSAST_OUTPUT_DIR`, `OUSAST_INPUT_CANDIDATES`, `OUSAST_INPUT_FACTS` +(optional), `OUSAST_BUDGET_USD`, `OUSAST_BUDGET_CALLS`, `OUSAST_MODEL`, `OUSAST_MODEL_PARAMS`, `OUSAST_PASS`. +Exit 0 done, 2 failed, 3 unfinished. +""" + +from __future__ import annotations + +import json +import os +import sys +import traceback +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from ...tool_hunter import CLIENT_ENV, ChatClient, ChatResponse, resolve_hunter_client, run_tool_hunter, scripted_hunter_client +from ..budget import AccountError, BudgetExhausted, MeteredClient + +PER_HUNT = 6 +MAX_STEPS = 6 +CALLERS_CAP = 8 +EXIT_CODES = {"done": 0, "failed": 2, "unfinished": 3} +USAGE_FIELDS = ("prompt_tokens", "prompt_cache_hit_tokens", "completion_tokens") + +# What each family's candidate operation does, in the words of `model_sinks.py` (the stage-1 classifier). +OPERATIONS = { + "injection": ( + "executes a database query (SQL, NoSQL, JCR/graph query), an operating-system command or shell, code or " + "expression evaluation (eval, exec, dynamic parsing into executable form), server-side template rendering " + "from a template string, or deserialization" + ), + "untrusted_destination": "makes an outbound HTTP/network request, downloads a URL, or redirects the user to a URL", + "config_secrets": ( + "sets security-relevant configuration: CORS policy, debug mode, TLS/certificate verification, cookie or " + "session security, or embeds a secret, password or API key" + ), +} +SYSTEM = ( + "You are a security reviewer with repository tools (read_file, grep_repo, find_refs). You are given several " + "operations in ONE file, each named by its function. For EACH operation decide whether data an attacker " + "controls -- HTTP request parameters, body, headers, cookies, uploaded or imported content, webhook or API " + "payloads, values stored earlier from such input -- can reach it without a sufficient guard (validation or " + "allowlisting, escaping that fits the exact context, parameterization, or a fixed destination). Use the tools " + "to follow callers and data. Report an operation only if you can state the concrete path from the source to " + "it. Answer with a JSON array: one object per vulnerable operation with path, line, function_name (exactly as " + "listed), title and rationale (the source, each hop, and why no guard applies); [] if none is vulnerable or " + "you cannot establish a path." +) + +Candidate = tuple[str, str, int] + + +class _Recording: + """Counts a hunt's tool turns on the way through to the metered client.""" + + def __init__(self, client: ChatClient) -> None: + self.client = client + self.turns = 0 + self.calls = 0 + + def complete(self, **kw: Any) -> ChatResponse: + response = self.client.complete(**kw) + self.calls += 1 + if response.tool_calls: + self.turns += 1 + return response + + +def load_candidates(payload: object) -> tuple[str | None, list[Candidate]]: + """(family, unique candidates) from a recorded candidate list: the `independent-v2-modelsinks` case record + (`{"family", "candidates": [[path, function, line], ...]}`) or a bare list of such triples.""" + family: str | None = None + rows: object = payload + if isinstance(payload, Mapping): + family = str(payload["family"]) if payload.get("family") else None + rows = payload.get("candidates", []) + if not isinstance(rows, list): + raise ValueError("candidates must be a list of [path, function, line]") + unique: dict[tuple[str, str], Candidate] = {} + for row in sorted(rows, key=lambda r: (str(r[0]), str(r[1]), int(r[2] or 0) if len(r) > 2 else 0)): + path, function = str(row[0]), str(row[1]) + line = int(row[2]) if len(row) > 2 and isinstance(row[2], int | float) else 0 + unique.setdefault((path, function), (path, function, line)) + return family, list(unique.values()) + + +def callers_of(facts: Mapping[str, Any] | None, path: str, function: str) -> list[dict[str, Any]]: + """Known call sites of `path::function` from `facts.json`: the task-5 shape (`callers` keyed by + `path::function`, rows `{"path", "line", "enclosing"}`) or the design's per-file shape (`{path: {"callers": + {function: [{"path", "line", "in"}]}}`).""" + if not facts: + return [] + rows: object = None + callers = facts.get("callers") + if isinstance(callers, Mapping): + rows = callers.get(f"{path}::{function}") + if rows is None: + per_file = facts.get("files", facts) + entry = per_file.get(path) if isinstance(per_file, Mapping) else None + if isinstance(entry, Mapping) and isinstance(entry.get("callers"), Mapping): + rows = entry["callers"].get(function) + return [dict(r) for r in rows if isinstance(r, Mapping)] if isinstance(rows, list) else [] + + +def callers_line(rows: Sequence[Mapping[str, Any]], cap: int = CALLERS_CAP) -> str: + """`Known callers: a.py:12 in handler(), ...` (at most `cap`); empty when nothing is known, so the prompt + only ever gains facts and never asserts an absence the extractor may have missed.""" + parts = [] + for row in rows[:cap]: + enclosing = row.get("enclosing") or row.get("in") or "" + where = f"{row.get('path', '?')}:{row.get('line', 0)}" + parts.append(f"{where} in {enclosing}()" if enclosing else where) + if not parts: + return "" + more = f", and {len(rows) - cap} more" if len(rows) > cap else "" + return "Known callers: " + ", ".join(parts) + more + + +def hunt_prompt(path: str, group: Sequence[Candidate], family: str, facts: Mapping[str, Any] | None) -> str: + lines = [] + for _, function, line in group: + lines.append(f"- function `{function}` around line {line}") + known = callers_line(callers_of(facts, path, function)) + if known: + lines.append(f" {known}") + listing = "\n".join(lines) + return ( + f"Operations to judge, all in `{path}`, each of which {OPERATIONS[family]}:\n{listing}\n" + "The file is above. For each, trace where the operation's data comes from and whether it is guarded." + ) + + +def hunt( + root: Path, path: str, group: Sequence[Candidate], family: str, *, client: ChatClient, model: str, facts: Mapping[str, Any] | None +) -> tuple[list[dict[str, Any]], int, int]: + """One batched hunt: (flagged rows anchored to the candidates, tool turns, calls).""" + from ...complexity.map import Hotspot + + recording = _Recording(client) + spots = [ + Hotspot( + path=path, function_name=fn, score=1.0, band="candidate", signals={}, rationale="model-classified sink", + test_hint=None, inventory_finding_ids=(), + ) + for _, fn, _ in group + ] # fmt: skip + found = run_tool_hunter( + root, spots, client=recording, model=model, max_steps=MAX_STEPS, system_prompt=SYSTEM, + user_prompt=hunt_prompt(path, group, family, facts), context_files=[path], tags=("sink-verifier",), + ) # fmt: skip + names = {fn for _, fn, _ in group} + rows = [] + for f in found: + fn = f.function_name or "" + if fn not in names: # attribute by the nearest listed line when the model renamed the function + fn = min(group, key=lambda c: abs((c[2] or 0) - (f.line or 0)))[1] + # The verdict is about the candidate operation, so the site is the candidate's file and function; the + # model's own location is kept beside it (it often names the file where the input enters instead). + line = f.line if f.path == path else next(c[2] for c in group if c[1] == fn) + rows.append({ + "candidate": f"{path}::{fn}", "site": f"{path}:{line or 0}:{fn}", "reported_at": f"{f.path}:{f.line or 0}", + "family": family, "title": f.title, "witness": f.rationale[:600], + }) # fmt: skip + return rows, recording.turns, recording.calls + + +def units_of(candidates: Sequence[Candidate], per_hunt: int = PER_HUNT) -> list[tuple[str, list[Candidate]]]: + by_file: dict[str, list[Candidate]] = {} + for c in candidates: + by_file.setdefault(c[0], []).append(c) + return [(path, group[i : i + per_hunt]) for path, group in by_file.items() for i in range(0, len(group), per_hunt)] + + +def unit_key(path: str, group: Sequence[Sequence[object]]) -> tuple[str, tuple[str, ...]]: + return path, tuple(str(c[1]) for c in group) + + +def read_units(log: Path) -> list[dict[str, Any]]: + if not log.is_file(): + return [] + rows = [] + for text in log.read_text().splitlines(): + if text.strip(): + row = json.loads(text) + if isinstance(row, dict): + rows.append(row) + return rows + + +def _sum_usage(rows: Sequence[Mapping[str, Any]]) -> tuple[dict[str, int], int, float | None]: + usage = dict.fromkeys(USAGE_FIELDS, 0) + calls, usd, priced = 0, 0.0, True + for row in rows: + used = row.get("usage") or {} + for field in USAGE_FIELDS: + usage[field] += int(used.get(field) or 0) + calls += int(used.get("calls") or 0) + if row.get("usd") is None: + priced = False + else: + usd += float(row["usd"]) + return usage, calls, (round(usd, 6) if priced else None) + + +def run( + workspace: Path, + output_dir: Path, + candidates: object, + facts: Mapping[str, Any] | None, + *, + client: ChatClient, + budget: MeteredClient | None = None, + model: str = "", + pass_label: str = "a", + family: str | None = None, +) -> dict[str, Any]: + """Hunt every unfinished unit, append `units.jsonl` per hunt, write `summary.json` last; returns the summary. + + `budget` is the meter to run under (built from `client` when absent); `client` is what the meter wraps. + """ + output_dir.mkdir(parents=True, exist_ok=True) + log, summary_path = output_dir / "units.jsonl", output_dir / "summary.json" + metered = budget if budget is not None else MeteredClient(client) + model = model or getattr(metered, "model", "") or "unbound" + recorded_family, unique = load_candidates(candidates) + family = family or recorded_family + units = units_of(unique) + rows = read_units(log) + done = {unit_key(r["path"], r["candidates"]): r for r in rows if "error" not in r} + status, reason = "done", "" + if family is None or family not in OPERATIONS: + status, reason = "failed", f"unknown family {family!r}: expected one of {sorted(OPERATIONS)}" + family = "" + with log.open("a") as sink: + for path, group in units if status == "done" else []: + if unit_key(path, group) in done: + continue + before = (dict(metered.usage), metered.usd, metered.calls) + row: dict[str, Any] = {"path": path, "candidates": [list(c) for c in group], "pass": pass_label} + try: + flagged, turns, _ = hunt(workspace, path, group, family, client=metered, model=model, facts=facts) + row.update(flagged=flagged, turns=turns) + except BudgetExhausted as stop: + status, reason = "unfinished", str(stop) + row["error"] = f"BudgetExhausted: {stop}" # a hunt cut short is re-asked on resume + except AccountError as refusal: + status, reason = "failed", refusal.provider_message + row["error"] = f"AccountError: {refusal.provider_message[:200]}" + except Exception as error: # noqa: BLE001 -- recorded, re-asked on resume; never a clean verdict + row["error"] = f"{type(error).__name__}: {str(error)[:200]}" + usage = {k: metered.usage[k] - before[0][k] for k in USAGE_FIELDS} | {"calls": metered.calls - before[2]} + row.update(usage=usage, usd=(round(metered.usd - before[1], 8) if metered.usd is not None and before[1] is not None else None)) + if status != "done" and not usage["calls"]: + break # nothing was asked; no row to keep + sink.write(json.dumps(row) + "\n") + sink.flush() + rows.append(row) + if "error" not in row: + done[unit_key(path, group)] = row + if status != "done": + break + if status == "done" and len(done) < len(units): + errors = [r["error"] for r in rows if "error" in r] + status, reason = "unfinished", f"{len(units) - len(done)} of {len(units)} units not finished; errors: {errors[:5]}" + usage, calls, usd = _sum_usage(rows) + summary: dict[str, Any] = { + "status": status, "units_done": len(done), "units_total": len(units), "usd": usd, "calls": calls, "usage": usage, + "priced": metered.priced, "model": model, "pass": pass_label, "family": family, + } # fmt: skip + if reason: + summary["reason"] = reason + summary_path.write_text(json.dumps(summary, indent=1) + "\n") + return summary + + +def _env_path(name: str, *, required: bool = True) -> Path | None: + value = os.environ.get(name, "").strip() + if not value: + if required: + raise ValueError(f"{name} is not set") + return None + return Path(value) + + +def _env_number(name: str, cast: Any) -> Any: + value = os.environ.get(name, "").strip() + return cast(value) if value else None + + +def build_client() -> ChatClient: + flag = os.environ.get("OUSAST_CLIENT", "").strip().lower() + if flag == "scripted": + return scripted_hunter_client("scripted") + client = resolve_hunter_client() + if client is None: + raise ValueError(f"no model client: set DEEPSEEK_API_KEY, OPENROUTER_API_KEY or {CLIENT_ENV}") + return client + + +def _failed(output_dir: Path | None, reason: str, model: str, pass_label: str) -> dict[str, Any]: + summary = { + "status": "failed", "units_done": 0, "units_total": 0, "usd": None, "calls": 0, "usage": dict.fromkeys(USAGE_FIELDS, 0), + "priced": False, "model": model, "pass": pass_label, "reason": reason, + } # fmt: skip + if output_dir is not None: + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "summary.json").write_text(json.dumps(summary, indent=1) + "\n") + return summary + + +def main() -> int: + output_dir, model, pass_label = ( + _env_path("OUSAST_OUTPUT_DIR", required=False), + os.environ.get("OUSAST_MODEL", ""), + os.environ.get("OUSAST_PASS", "a"), + ) + try: + workspace = _env_path("OUSAST_WORKSPACE_DIR") + candidates_path = _env_path("OUSAST_INPUT_CANDIDATES") + facts_path = _env_path("OUSAST_INPUT_FACTS", required=False) + if output_dir is None or workspace is None or candidates_path is None: + raise ValueError("OUSAST_OUTPUT_DIR is not set") + if not model: + raise ValueError("OUSAST_MODEL is not set") + params = json.loads(os.environ.get("OUSAST_MODEL_PARAMS", "") or "{}") + candidates = json.loads(candidates_path.read_text()) + facts = json.loads(facts_path.read_text()) if facts_path is not None else None + client = build_client() + metered = MeteredClient( + client, prices=params, budget_usd=_env_number("OUSAST_BUDGET_USD", float), budget_calls=_env_number("OUSAST_BUDGET_CALLS", int) + ) + summary = run( + workspace, output_dir, candidates, facts, client=client, budget=metered, model=model, pass_label=pass_label, + family=os.environ.get("OUSAST_FAMILY") or None, + ) # fmt: skip + except Exception: # noqa: BLE001 -- a crash is `failed` with its traceback, never a silent zero + summary = _failed(output_dir, traceback.format_exc()[-2000:], model, pass_label) + print( + json.dumps({k: summary.get(k) for k in ("status", "units_done", "units_total", "usd", "calls", "model", "pass")}), file=sys.stderr + ) + return EXIT_CODES.get(str(summary["status"]), 2) + + +__all__ = [ + "MAX_STEPS", + "OPERATIONS", + "PER_HUNT", + "SYSTEM", + "build_client", + "callers_line", + "callers_of", + "hunt", + "hunt_prompt", + "load_candidates", + "main", + "run", + "units_of", +] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_plane_verify.py b/tests/test_plane_verify.py new file mode 100644 index 0000000..ff7e724 --- /dev/null +++ b/tests/test_plane_verify.py @@ -0,0 +1,285 @@ +"""verify and agree tasks (ai-service-plane task 6): callers in the prompt, per-unit resume, 402 -> failed, +budget -> unfinished, the bound model and summed usage in summary.json, agreement and metrics over two passes.""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import pytest + +from openultrasast.plane.budget import MeteredClient +from openultrasast.plane.tasks import agree, verify +from openultrasast.tool_hunter import ChatResponse, ToolCall + +FLASH = {"cache_hit_per_m": 0.014, "input_per_m": 0.44, "output_per_m": 1.32} +ROW = {"prompt_tokens": 1000, "prompt_cache_hit_tokens": 200, "completion_tokens": 50} + + +class HuntClient: + """A scripted hunter that answers from the conversation: a read_file turn first, then a JSON verdict naming + the candidates in `flag` for the file the prompt asked about. Records provider usage rows like + `DeepSeekChatClient`, so `MeteredClient` can price it; a step index in `raise_at` raises instead.""" + + def __init__(self, flag: set[str], *, raise_at: dict[int, Exception] | None = None) -> None: + self.flag = flag + self.raise_at = raise_at or {} + self.usage: list[dict[str, object]] = [] + self.calls: list[list[dict[str, object]]] = [] + + def complete(self, *, model: str, messages: list[dict[str, object]], tools: list[dict[str, object]], **_: object) -> ChatResponse: + index = len(self.calls) + self.calls.append(list(messages)) + if index in self.raise_at: + raise self.raise_at[index] + self.usage.append(dict(ROW)) + prompt = str(messages[1]["content"]) + path = re.search(r"all in `([^`]+)`", prompt).group(1) # type: ignore[union-attr] + if not any(m.get("role") == "tool" for m in messages): + return ChatResponse(tool_calls=(ToolCall(id="c1", name="read_file", arguments={"path": path}),)) + names = re.findall(r"- function `([^`]+)` around line (\d+)", prompt) + found = [ + { + "path": path, + "line": int(line), + "function_name": fn, + "title": f"{fn} sink", + "rationale": "request -> sink", + "family": "injection", + } + for fn, line in names + if f"{path}::{fn}" in self.flag + ] + return ChatResponse(content=json.dumps(found)) + + def prompts(self) -> list[str]: + return [str(call[1]["content"]) for call in self.calls] + + +def _workspace(tmp_path: Path) -> Path: + root = tmp_path / "repo" + root.mkdir() + (root / "app.py").write_text("def run(q):\n return eval(q)\n\ndef other(x):\n return exec(x)\n") + (root / "db.py").write_text("def query(sql):\n return cursor.execute(sql)\n") + return root + + +CANDIDATES = { + "id": "case", + "family": "injection", + "candidates": [["app.py", "run", 2], ["app.py", "other", 5], ["db.py", "query", 2], ["app.py", "run", 2]], +} +FACTS = { + "files": {"app.py": {"functions": ["run", "other"]}, "db.py": {"functions": ["query"]}}, + "callers": { + "app.py::run": [{"path": "web.py", "line": 12, "enclosing": "handler"}, {"path": "cli.py", "line": 40, "enclosing": ""}], + "db.py::query": [{"path": f"m{i}.py", "line": i, "enclosing": f"f{i}"} for i in range(12)], + }, + "counts": {}, +} + + +def _run(root: Path, out: Path, client: HuntClient, *, facts: dict | None = FACTS, **budget: object) -> dict: + metered = MeteredClient(client, prices=FLASH, **budget) # type: ignore[arg-type] + return verify.run(root, out, CANDIDATES, facts, client=client, budget=metered, model="deepseek-chat", pass_label="a") + + +def test_prompt_carries_the_known_callers_line_capped_at_eight(tmp_path: Path) -> None: + client = HuntClient({"app.py::run"}) + summary = _run(_workspace(tmp_path), tmp_path / "out", client) + assert summary["status"] == "done", summary + app = next(p for p in client.prompts() if "all in `app.py`" in p) + # candidates in the reference's sorted order; the callers line follows the candidate it belongs to + assert ( + "- function `other` around line 5\n- function `run` around line 2\n" + " Known callers: web.py:12 in handler(), cli.py:40\nThe file is above." in app + ) + assert ( + "Known callers" not in app.split("- function `other`")[1].split("- function `run`")[0] + ) # nothing known: no line, no claim of absence + db = next(p for p in client.prompts() if "all in `db.py`" in p) + assert db.count("m") >= 8 and "and 4 more" in db and "m8.py" not in db + assert verify.SYSTEM in str(client.calls[0][0]["content"]) # the reference's system prompt, unchanged + + +def test_without_facts_the_prompt_is_the_reference_prompt(tmp_path: Path) -> None: + client = HuntClient(set()) + _run(_workspace(tmp_path), tmp_path / "out", client, facts=None) + assert all("Known callers" not in p for p in client.prompts()) + assert any( + p.endswith("The file is above. For each, trace where the operation's data comes from and whether it is guarded.") + for p in client.prompts() + ) + + +def test_units_are_one_row_per_file_hunt_with_anchored_sites(tmp_path: Path) -> None: + client = HuntClient({"app.py::other", "db.py::query"}) + out = tmp_path / "out" + summary = _run(_workspace(tmp_path), out, client) + rows = [json.loads(t) for t in (out / "units.jsonl").read_text().splitlines()] + assert [(r["path"], [c[1] for c in r["candidates"]], r["turns"], r["pass"]) for r in rows] == [ + ("app.py", ["other", "run"], 1, "a"), + ("db.py", ["query"], 1, "a"), + ] + assert [f["site"] for r in rows for f in r["flagged"]] == ["app.py:5:other", "db.py:2:query"] + assert rows[0]["usage"] == {**ROW, "prompt_tokens": 2000, "prompt_cache_hit_tokens": 400, "completion_tokens": 100, "calls": 2} + assert summary["units_done"] == summary["units_total"] == 2 + + +def test_resume_asks_only_the_unfinished_file(tmp_path: Path) -> None: + root, out = _workspace(tmp_path), tmp_path / "out" + out.mkdir() + done_row = { + "path": "app.py", + "candidates": [["app.py", "other", 5], ["app.py", "run", 2]], + "flagged": [], + "turns": 3, + "usage": {**ROW, "calls": 4}, + "usd": 0.01, + "pass": "a", + } + errored = { + "path": "db.py", + "candidates": [["db.py", "query", 2]], + "error": "TimeoutError: x", + "usage": {**ROW, "calls": 1}, + "usd": 0.001, + "pass": "a", + } + (out / "units.jsonl").write_text(json.dumps(done_row) + "\n" + json.dumps(errored) + "\n") + client = HuntClient({"db.py::query"}) + summary = _run(root, out, client) + assert [p for p in client.prompts() if "all in `" in p] and all("all in `db.py`" in p for p in client.prompts()) + assert summary["status"] == "done" and summary["units_done"] == 2 and summary["units_total"] == 2 + assert summary["calls"] == 4 + 1 + 2 # the finished row, the errored attempt and this run's hunt all cost money + rows = [json.loads(t) for t in (out / "units.jsonl").read_text().splitlines()] + assert len(rows) == 3 and rows[-1]["path"] == "db.py" and "error" not in rows[-1] + + +def test_a_scripted_402_yields_failed_with_the_provider_message(tmp_path: Path) -> None: + client = HuntClient(set(), raise_at={1: RuntimeError("HTTP Error 402: Insufficient Balance")}) + out = tmp_path / "out" + summary = _run(_workspace(tmp_path), out, client) + assert summary["status"] == "failed" and "402" in summary["reason"] and summary["units_done"] == 0 + assert len(client.calls) == 2 # nothing further was asked + rows = [json.loads(t) for t in (out / "units.jsonl").read_text().splitlines()] + assert len(rows) == 1 and rows[0]["error"].startswith("AccountError") and rows[0]["usage"]["calls"] == 1 + assert json.loads((out / "summary.json").read_text())["status"] == "failed" + + +def test_a_budget_stop_yields_unfinished_and_resumes_under_a_larger_one(tmp_path: Path) -> None: + root, out = _workspace(tmp_path), tmp_path / "out" + client = HuntClient({"db.py::query"}) + summary = _run(root, out, client, budget_calls=3) + assert summary["status"] == "unfinished" and "calls budget exhausted" in summary["reason"] + assert summary["units_done"] == 1 and summary["units_total"] == 2 and summary["calls"] == 3 + rows = [json.loads(t) for t in (out / "units.jsonl").read_text().splitlines()] + assert [("error" in r) for r in rows] == [False, True] # the cut-short hunt is on record, marked for a re-ask + again = HuntClient({"db.py::query"}) + summary = _run(root, out, again, budget_calls=10) + assert summary["status"] == "done" and all("all in `db.py`" in p for p in again.prompts()) + assert summary["units_done"] == 2 + + +def test_summary_records_the_bound_model_and_summed_usage(tmp_path: Path) -> None: + client = HuntClient(set()) + out = tmp_path / "out" + summary = _run(_workspace(tmp_path), out, client) + assert summary["model"] == "deepseek-chat" and summary["pass"] == "a" and summary["priced"] is True + assert summary["usage"] == {"prompt_tokens": 4000, "prompt_cache_hit_tokens": 800, "completion_tokens": 200} and summary["calls"] == 4 + per_call = (200 * 0.014 + 800 * 0.44 + 50 * 1.32) / 1_000_000 + assert summary["usd"] == pytest.approx(4 * per_call, abs=1e-6) + on_disk = json.loads((out / "summary.json").read_text()) + assert on_disk == summary + + +def test_an_unpriced_model_reports_usd_none_not_zero(tmp_path: Path) -> None: + client = HuntClient(set()) + summary = verify.run(_workspace(tmp_path), tmp_path / "out", CANDIDATES, None, client=client, budget=MeteredClient(client), model="m") + assert summary["usd"] is None and summary["priced"] is False and summary["calls"] == 4 + + +def test_main_runs_the_env_contract_with_the_scripted_seam(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + root, out = _workspace(tmp_path), tmp_path / "out" + (tmp_path / "candidates.json").write_text(json.dumps(CANDIDATES)) + (tmp_path / "facts.json").write_text(json.dumps(FACTS)) + env = { + "OUSAST_WORKSPACE_DIR": str(root), "OUSAST_OUTPUT_DIR": str(out), "OUSAST_INPUT_CANDIDATES": str(tmp_path / "candidates.json"), + "OUSAST_INPUT_FACTS": str(tmp_path / "facts.json"), "OUSAST_MODEL": "deepseek-chat", "OUSAST_MODEL_PARAMS": json.dumps(FLASH), + "OUSAST_BUDGET_USD": "1.0", "OUSAST_BUDGET_CALLS": "100", "OUSAST_PASS": "b", "OUSAST_CLIENT": "scripted", + } # fmt: skip + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert verify.main() == 0 + summary = json.loads((out / "summary.json").read_text()) + assert summary["status"] == "done" and summary["model"] == "deepseek-chat" and summary["pass"] == "b" and summary["units_total"] == 2 + monkeypatch.delenv("OUSAST_MODEL") + assert verify.main() == 2 + assert "OUSAST_MODEL" in json.loads((out / "summary.json").read_text())["reason"] + + +def _two_passes(tmp_path: Path) -> tuple[Path, Path, Path]: + root = _workspace(tmp_path) + a, b = tmp_path / "pass-a", tmp_path / "pass-b" + _run(root, a, HuntClient({"app.py::run", "db.py::query"})) + _run(root, b, HuntClient({"app.py::run", "app.py::other"})) + return root, a, b + + +def test_agree_over_two_scripted_passes_computes_agreement_and_metrics(tmp_path: Path) -> None: + _, a, b = _two_passes(tmp_path) + out = tmp_path / "agree" + sites = {"id": "case", "family": "injection", "sites": ["app.py::run", "db.py::query"]} + summary = agree.run(out, CANDIDATES, agree.load_pass(a), agree.load_pass(b), sites=sites, matcher=agree.load_matcher()) + assert summary["status"] == "done" and summary["usd"] == 0 and summary["calls"] == 0 and summary["model"] is None + report = json.loads((out / "agreed.json").read_text()) + by_name = {r["candidate"]: r for r in report["candidates"]} + assert {n: (r["a"], r["b"], r["agreed"], r["site_match"]) for n, r in by_name.items()} == { + "app.py::run": (True, True, True, True), + "app.py::other": (False, True, False, False), + "db.py::query": (True, False, False, True), + } + assert [f["candidate"] for f in report["agreed"]] == ["app.py::run"] and report["agreed"][0]["passes"] == 2 + assert sorted(f["candidate"] for f in report["disputed"]) == ["app.py::other", "db.py::query"] + assert json.loads((out / "disputed.json").read_text()) == report["disputed"] + metrics = summary["metrics"] + per_call = (200 * 0.014 + 800 * 0.44 + 50 * 1.32) / 1_000_000 + assert metrics["candidates"] == 3 and metrics["agreed"] == 1 and metrics["disputed"] == 2 and metrics["hunts"] == 4 + assert metrics["declared_sites"] == 2 and metrics["declared_sites_matched"] == 2 and metrics["declared_sites_matched_agreed"] == 1 + assert metrics["usd_total"] == pytest.approx(8 * per_call, abs=1e-6) + assert metrics["cost_per_candidate"] == pytest.approx(8 * per_call / 3, abs=1e-6) + assert metrics["turns_per_hunt"] == 1.0 and metrics["calls_per_hunt"] == 2.0 and metrics["site_matching"] == "assessed" + # a hunt's usd is split over its candidates: app.py's two share one hunt per pass, db.py's one has a hunt alone + assert by_name["db.py::query"]["usd"] == pytest.approx(4 * per_call, abs=1e-6) + assert by_name["app.py::run"]["usd"] == pytest.approx(2 * per_call, abs=1e-6) + + +def test_agree_without_sites_leaves_site_matching_unassessed_and_flags_unasked_candidates(tmp_path: Path) -> None: + _, a, b = _two_passes(tmp_path) + rows_b = [r for r in agree.load_pass(b) if r["path"] != "db.py"] + summary = agree.run(tmp_path / "agree", CANDIDATES, agree.load_pass(a), rows_b) + assert summary["status"] == "unfinished" and "db.py::query" in summary["reason"] + report = json.loads((tmp_path / "agree" / "agreed.json").read_text()) + assert {r["candidate"]: (r["b"], r["site_match"]) for r in report["candidates"]}["db.py::query"] == (None, None) + assert summary["metrics"]["site_matching"] == "not assessed: no sites given" + + +def test_agree_main_reads_the_env_contract(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _, a, b = _two_passes(tmp_path) + (tmp_path / "candidates.json").write_text(json.dumps(CANDIDATES)) + (tmp_path / "case.toml").write_text( + '[[case]]\nid = "case"\nfamily = "injection"\nsites = ["app.py::run"]\n\n[[case]]\nid = "x"\nfamily = "injection"\nsites = []\n' + ) + env = { + "OUSAST_OUTPUT_DIR": str(tmp_path / "agree"), "OUSAST_INPUT_PASS_A": str(a), "OUSAST_INPUT_PASS_B": str(b / "units.jsonl"), + "OUSAST_INPUT_CANDIDATES": str(tmp_path / "candidates.json"), "OUSAST_INPUT_SITES": str(tmp_path / "case.toml"), + "OUSAST_CASE_ID": "case", + } # fmt: skip + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert agree.main() == 0 + summary = json.loads((tmp_path / "agree" / "summary.json").read_text()) + assert summary["metrics"]["declared_sites_matched"] == 1 and summary["metrics"]["agreed"] == 1 + monkeypatch.delenv("OUSAST_CASE_ID") + assert agree.main() == 2 # two cases match: the record is ambiguous, not silently the first