From 87b82564b4d40b7c64fdf74af456357eac9e37c3 Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Mon, 5 Oct 2026 18:58:24 +0200 Subject: [PATCH] search-with-proof: pilot driver (6.2) One side per run over selected pilot pairs: board, brain on the VM with the project DeepSeek client, executor tasks on kube-ax, verification on the VM lane (AX verify lane switchable). Hard run ceiling with reservation before every call, per-search and step caps; fixed side never sees the vulnerable revision; records carry digests only. Dry run with the stub model reaches demonstrated. Co-Authored-By: Claude Opus 5.5 --- benchmarks/search/README.md | 78 ++++- benchmarks/search/pilot_run.py | 456 +++++++++++++++++++++++++++++ src/openultrasast/search/worker.py | 26 ++ tests/test_search_pilot_run.py | 322 ++++++++++++++++++++ 4 files changed, 881 insertions(+), 1 deletion(-) create mode 100644 benchmarks/search/pilot_run.py create mode 100644 tests/test_search_pilot_run.py diff --git a/benchmarks/search/README.md b/benchmarks/search/README.md index faf42d6..12271f8 100644 --- a/benchmarks/search/README.md +++ b/benchmarks/search/README.md @@ -1,6 +1,8 @@ # Search feasibility tools -These host-side tools make no model calls. Run them from the repository root. +Run these host-side tools from the repository root. Selection and executor smoke +make no model calls; the pilot driver calls the configured model unless `--dry-run` +is specified. ## Executor smoke @@ -55,3 +57,77 @@ The default identity-bearing manifest is counts, selection biases, input instrumentation, and digests only. Neither output is overwritten. A successful selection is feasibility input, not evidence of buildability, oracle operation, unseen performance, or BLOCK eligibility. + + +## One-sided pilot driver (task 6.2) + +The shared virtualenv may be installed against another worktree; `PYTHONPATH` +below explicitly selects this worktree's source. + +Offline probe, using a scripted model and in-process repository executor with real +host verification (no network or provider calls): + +```sh +PYTHONPATH="$PWD/src" /home/mc/Source/OpenUltraSAST/.venv/bin/python -m benchmarks.search.pilot_run \ + --dry-run --side vulnerable --pairs 0 --ceiling-usd 0.50 \ + --out /tmp/search-pilot-dry.jsonl +``` + +For three **live** searches, after the oracle infrastructure probe and budget +approval, set `SEARCH_TASK_IMAGE` to the tested search image's complete +`REGISTRY/IMAGE@sha256:DIGEST` reference, then run: + +```sh +PYTHONPATH="$PWD/src" /home/mc/Source/OpenUltraSAST/.venv/bin/python -m benchmarks.search.pilot_run \ + --manifest benchmarks/search/private/pilot.json \ + --side vulnerable --pairs 0,1,2 --ceiling-usd 1.50 \ + --verify-lane vm --image "${SEARCH_TASK_IMAGE:?Set the tested digest-pinned search image}" \ + --out benchmarks/measurements/2026-10-05-search-pilot-vulnerable.jsonl +``` + +This command spends money and contacts AX and the store. The ceiling shown is an +example for the three $0.50 searches, not evidence that a live run was approved or +performed. Use a fresh output path; output files are never overwritten. + +The existing DeepSeek key/base URL environment variables and provider model/prices +are reused (`deepseek-flash` by default). `.env` does not override exported values. +Retries are disabled. Unknown models require both `--usd-per-mtok-in` and +`--usd-per-mtok-out`. Every call reserves its conservative input-token bound plus +`--max-tokens` output allowance against the task allowance and shared run ceiling; +failed calls and calls with missing usage settle the maximum. Reservations in the +record are cumulative admitted maxima, not outstanding spend. Their sum can exceed +the ceiling when usage settlements release unused allowances. Settled spend remains +bounded by the run ceiling and $0.50/search. Search defaults also cap reason rounds +at 4, intents at 3 per round, explore at 15 minutes including acquisition, and demos +at 3. Physical AX submissions and coordinator steps have separate 32-task caps. + +The board defaults to a private local FileStore under a fresh directory in +`benchmarks/search/private/runs/`; `--board-memory` selects another board store URI. +The AX mailbox still requires the existing `OUSAST_MEMORY=s3://bucket/prefix`, public +HTTPS `S3_ENDPOINT`, and host store credentials. Board storage does not replace the +mailbox. Executors receive only presigned object URLs. The brain runs on the VM; +repository tools and build/export run in key-free AX executors. VM verification +requires working user namespaces and consumes executor-built archives. `--verify-lane +ax` switches to the existing dispatcher, with three fresh verifier tasks per side +and no executor holding a worker slot during those repetitions. + +Each selected index produces a JSON line followed by one run summary: Req 3.4 +outcome, reserved/settled spend, calls, tokens, physical tasks, coordinator steps, +phase timings, failure codes, input bytes, board bytes, end reason, and hashed +repository/revision identities. The private board and checkouts contain identity-bearing +data. The public JSONL omits raw model/executor diagnostics. Selection advantages +are explicit in the summary; this is feasibility evidence, not BLOCK qualification. + +A fixed invocation creates a fresh board/demo and retains only the fixed revision. +Its verifier runs two independent sets of repetitions on that same fixed revision; +no vulnerable checkout or previously generated vulnerable demo is passed in. +`fixed_effect_observed` reports a consistent owned effect on the first three fixed +runs (`null` when unavailable or inconsistent). Such an effect requires adjudication; +it is not promoted to a differential `demonstrated` outcome. The normal safe probe +therefore ends `inconclusive`, never “not vulnerable”. + +Generic `injection` needs an `oracle` value of `sql` or `command` in the private +manifest; without a subtype verification returns `no_oracle`. No subtype is guessed +from repository names. The selected manifest's pair 1 currently lacks this metadata. +Live transport, provider authentication/cost, real project buildability, and AX +oracle operation must still be measured; an offline probe cannot establish them. diff --git a/benchmarks/search/pilot_run.py b/benchmarks/search/pilot_run.py new file mode 100644 index 0000000..be11278 --- /dev/null +++ b/benchmarks/search/pilot_run.py @@ -0,0 +1,456 @@ +"""One-sided feasibility searches. Live mode is opt-in by omitting --dry-run.""" + +from __future__ import annotations + +import argparse +import hashlib +import io +import json +import math +import os +import re +import tarfile +import time +import uuid +from collections import Counter +from contextlib import contextmanager +from dataclasses import replace +from pathlib import Path + +from benchmarks.ax.batch import AXLane, SearchExecutorTask, Workload, validate_image +from benchmarks.ax.search_verify import AXSideDispatcher, validate_result +from benchmarks.search.executor_smoke import CLONE +from openultrasast.config import load_dotenv +from openultrasast.model.endpoint import ( + DEEPSEEK_BASE_ENV, + DEEPSEEK_BASE_URL, + DEEPSEEK_KEY_ENV, + DEFAULT_DETECTOR_MODEL, + Prices, + price_of, +) +from openultrasast.plane.budget import BudgetExhausted, SpendBudget +from openultrasast.plane.memory import FileStore, open_store +from openultrasast.provider.openrouter import OpenRouterChatClient +from openultrasast.search import _sandbox +from openultrasast.search.board import Board +from openultrasast.search.budget import SearchBudget +from openultrasast.search.coordinator import Coordinator +from openultrasast.search.demo import load_demo +from openultrasast.search.executor import InProcessExecutor +from openultrasast.search.oracles import BrowserExecutor, oracle_for +from openultrasast.search.probe import materialise +from openultrasast.search.verify import Side, VerificationRecord, verify +from openultrasast.search.verify_task import extract_checkout, validate_spec +from openultrasast.search.worker import StubModel, Worker + + +def digest(value): + return hashlib.sha256(value.encode()).hexdigest() + + +def tool(name, **args): + return {"tool_calls": [{"id": "call", "type": "function", "function": {"name": name, "arguments": json.dumps(args)}}]} + + +def scripted(schema): + return StubModel( + [ + {"content": json.dumps({"intents": [{"description": "Prove the candidate via the external interface", "from_facts": []}]})}, + tool("read_file", path="app.py"), + tool("write_demo", schema=schema), + tool("finish", fact={"text": "Candidate demo", "evidence_refs": ["tool:1"], "demo": "demo/"}), + *[{"content": "{}"} for _ in range(4)], + ] + ) + + +class Pilot: + def __init__(self, args, row, index, root, store, run_budget, client, prices, lane=None): + self.args, self.root, self.lane = args, root, lane + # Fixed searches retain no vulnerable pin, even in trusted verification state. + self.row = {k: v for k, v in row.items() if args.side != "fixed" or k != "vulnerable"} + self.phase = "setup" + self.deadline = float("inf") + self.record = dict( + type="search", + pair_index=index, + side=args.side, + repo_digest=digest(row["repo"]), + revision_digest=digest(row[args.side]), + family=row["family"], + language=row["language"], + dry_run=args.dry_run, + fixed_effect_observed=None, + outcome="no_evidence", + end_reason="exhausted", + tasks_submitted=0, + executor_tasks=0, + verify_tasks=0, + input_bytes=0, + wall_seconds_by_phase={}, + failures=[], + ) + self.board = Board(store, "pilot-" + uuid.uuid4().hex, coordinator="pilot") + self.board.commit( + writer="pilot", + hints=[ + dict( + id="candidate", step="reason", task_id="manifest", text=json.dumps({k: row[k] for k in ("family", "file", "function")}) + ) + ], + ) + self.worker = Worker( + client, + repo=root / "unmounted", + demo=root / "demo", + model=args.model, + prices=prices, + run_budget=run_budget, + max_output_tokens=args.max_tokens, + ) + self.family = {"output_encoding": "xss", "untrusted_destination": "ssrf"}.get(row["family"], row["family"]) + # Injection is ambiguous without a manifest oracle subtype; never guess SQL vs command. + if self.family == "injection": + self.family = row.get("oracle", "injection") + self.sides = None + self.verifications = [] + + def failure(self, phase, reason): + # No untrusted diagnostic text in public records: it can contain repository identities or secrets. + self.record["failures"].append(dict(phase=phase, reason=reason)) + + @contextmanager + def executor(self, revision, deadline=900): + deadline = min(deadline, self.deadline - time.monotonic()) + if deadline < 1: + raise BudgetExhausted("task deadline", usd=0, calls=0) + acquired_at = time.monotonic() + empty = io.BytesIO() + with tarfile.open(fileobj=empty, mode="w"): + pass + self.admit_task() + self.record["executor_tasks"] += 1 + with SearchExecutorTask(self.lane, empty.getvalue(), deadline=deadline) as task: + remaining = deadline - (time.monotonic() - acquired_at) + if remaining < 1: + raise BudgetExhausted("acquisition deadline", usd=0, calls=0) + timeout = min(180, remaining) + result = task.client.submit( + "run", + {"command": ["sh", "-ec", CLONE, "clone", self.row["repo"], revision], "timeout_seconds": timeout}, + timeout_seconds=timeout, + ) + if result.get("exit_code") != 0 or result.get("error"): + self.failure("acquisition", "checkout_command_failed") + raise ValueError("checkout failed") + remaining = deadline - (time.monotonic() - acquired_at) + if remaining < 1: + raise BudgetExhausted("acquisition deadline", usd=0, calls=0) + result = task.client.submit("read_file", {"path": self.row["file"]}, timeout_seconds=min(30, remaining)) + if not isinstance(result.get("input_bytes"), int) or result["input_bytes"] <= 0: + self.failure("acquisition", "candidate_input_unreadable") + raise ValueError("unreadable input") + self.record["input_bytes"] += result["input_bytes"] + yield task + + def admit_task(self): + if self.record["tasks_submitted"] >= 32: + raise BudgetExhausted("executor task cap", usd=0, calls=0) + self.record["tasks_submitted"] += 1 + + def prepare(self, revision, spec): + before = time.monotonic() + try: + with self.executor(revision) as task: + return task.prepare(spec) + except BudgetExhausted: + raise + except Exception: + self.failure("build_export", "executor_preparation_failed") + raise + finally: + wall = self.record["wall_seconds_by_phase"] + wall["build_export"] = wall.get("build_export", 0) + time.monotonic() - before + + def verification(self): + demo = load_demo(self.root / "demo") + browser = BrowserExecutor(task_boundary=self.args.verify_lane == "ax") if self.family == "xss" else None + if oracle_for(self.family, browser) is None: + return VerificationRecord("no_oracle", (), 0, "no configured owned oracle") + if self.args.dry_run: + sides = self.sides if self.args.side == "vulnerable" else (self.sides[1], self.sides[1]) + return verify(*sides, self.root / "demo", self.family) + revisions = [self.row[self.args.side], self.row["fixed"]] + if self.args.verify_lane == "ax": + # Side paths are opaque lookup keys; no checkout or opposite revision reaches the brain. + sides = [Side(self.root / str(i)) for i in range(2)] + pins = {s.checkout: pin for s, pin in zip(sides, revisions, strict=True)} + parent = self + + class CountingLane: + def __call__(self, *a, **kw): + parent.admit_task() + parent.record["verify_tasks"] += 1 + return parent.lane(*a, **kw) + + dispatcher = AXSideDispatcher( + self.args, + self.args.image, + self.root / "verify", + lane=CountingLane(), + prepare=lambda side, spec: self.prepare(pins[side.checkout], spec), + ) + return verify( + *sides, + self.root / "demo", + self.family, + task_dispatcher=dispatcher, + browser=BrowserExecutor(task_boundary=True) if self.family == "xss" else None, + ) + if _sandbox.isolation_mode() != "userns": + raise _sandbox.IsolationUnavailable("VM verification requires userns") + sides = [] + spec = dict(demo=demo, family=self.family, timeout_seconds=5) + for revision in revisions: + archive = self.prepare(revision, spec) + bundle = self.root / ("built-" + uuid.uuid4().hex) + extract_checkout(archive, bundle) + sides.append(Side(bundle / "checkout", products=bundle / "products")) + # Export contains the executor's build-free schema. + (self.root / "verify-demo").mkdir(exist_ok=True) + built_spec = validate_spec(json.loads((bundle / "spec.json").read_text())) + expected_demo = {**demo, "build": {"recipe": "none", "arguments": []}} + if built_spec != {**spec, "demo": expected_demo}: + raise ValueError("executor altered verification contract") + (self.root / "verify-demo/demo.json").write_text(json.dumps(built_spec["demo"])) + return verify(*sides, self.root / "verify-demo", self.family, browser=BrowserExecutor() if self.family == "xss" else None) + + def dispatch(self, task): + self.phase = task.step + before = time.monotonic() + self.deadline = before + task.limits.get("task_wall_seconds", 900) + try: + if task.step == "verify": + result = self.verification() + self.verifications.append(result.outcome) + self.record["outcome"] = result.outcome + if self.args.side == "fixed" and result.sides: + observations = result.sides[0] + if observations.outcome == "observed" and len(observations.runs) == 3: + effects = {r.observed for r in observations.runs} + if len(effects) == 1: + self.record["fixed_effect_observed"] = effects.pop() + wall = self.record["wall_seconds_by_phase"] + for phase in ("build", "ready", "run"): + wall["verify_" + phase] = wall.get("verify_" + phase, 0) + sum(getattr(s, phase + "_seconds") for s in result.sides) + + for side in result.sides: + if side.outcome != "observed": + self.failure("verify/" + side.outcome, "side_verification_failed") + if result.outcome != "demonstrated": + self.failure("verify", "no_consistent_differential" if result.outcome == "inconclusive" else result.outcome) + return dict(status="ok", outcome=result.outcome, evidence_refs=["verify:owned"], cost_usd=0) + if task.step == "explore": + if self.args.dry_run: + self.worker.executor = InProcessExecutor(self.sides[0 if self.args.side == "vulnerable" else 1].checkout) + result = self.worker(task) + else: + with self.executor(self.row[self.args.side], deadline=task.limits["task_wall_seconds"]) as remote: + self.worker.executor = remote.client + remaining = task.limits["task_wall_seconds"] - (time.monotonic() - before) + if remaining < 1: + raise BudgetExhausted("explore deadline", usd=0, calls=0) + result = self.worker(replace(task, limits={**task.limits, "task_wall_seconds": remaining})) + else: + result = self.worker(task) + if self.worker.budget_exhausted: + raise BudgetExhausted("run ceiling", usd=0, calls=0) + if result.get("status") != "ok": + self.failure(task.step, "worker_execution_failed") + if not self.verifications: + self.record["outcome"] = "inconclusive" + return result + finally: + wall = self.record["wall_seconds_by_phase"] + wall[task.step] = wall.get(task.step, 0) + time.monotonic() - before + + def run(self): + started = time.monotonic() + coordinator = Coordinator(self.board, self.dispatch, budget=SearchBudget(retries=0, memory_bytes=2 * 1024**3)) + try: + state = coordinator.run() + self.record["end_reason"] = state["end_reason"] + if not self.verifications and any(f.get("step") == "explore" for f in state["facts"]): + self.record["outcome"] = "plausible_unproven" + except Exception as exc: + self.failure(self.phase, type(exc).__name__) + self.record["outcome"] = ( + "could_not_build" + if self.record["failures"] and any(f["phase"] in ("build_export", "acquisition") for f in self.record["failures"]) + else "inconclusive" + ) + self.record["end_reason"] = "instrument_failure" + metrics = self.worker.metrics + self.record.update( + spend={k: metrics[k] for k in ("reserved", "settled")}, + **{k: metrics[k] for k in ("model_calls", "tokens_in", "tokens_out")}, + coordinator_tasks=coordinator.progress["tasks"], + board_bytes=len(json.dumps(self.board.state).encode()), + wall_seconds=time.monotonic() - started, + ) + return self.record + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--side", choices=("vulnerable", "fixed"), required=True) + parser.add_argument("--pairs", required=True) + parser.add_argument("--ceiling-usd", required=True, type=float) + parser.add_argument("--manifest", type=Path, default=Path("benchmarks/search/private/pilot.json")) + parser.add_argument("--private-root", type=Path, default=Path("benchmarks/search/private/runs")) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--board-memory", help="Memory store URI; defaults to a private local FileStore") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--verify-lane", choices=("vm", "ax"), default="vm") + parser.add_argument("--model", default=DEFAULT_DETECTOR_MODEL) + parser.add_argument("--max-tokens", type=int, default=1024) + parser.add_argument("--usd-per-mtok-in", type=float) + parser.add_argument("--usd-per-mtok-out", type=float) + parser.add_argument("--image") + parser.add_argument("--ax-bin", default="ax") + parser.add_argument("--atespace", default="default") + parser.add_argument("--kubeconfig") + args = parser.parse_args(argv) + if not math.isfinite(args.ceiling_usd) or args.ceiling_usd < 0: + parser.error("ceiling must be finite and nonnegative") + run_budget = SpendBudget(args.ceiling_usd) + if not re.fullmatch(r"\d+(,\d+)*", args.pairs) or args.max_tokens < 1: + parser.error("pairs must be comma-separated nonnegative indices; max tokens must be positive") + indices = [int(i) for i in args.pairs.split(",")] + if len(indices) != len(set(indices)): + parser.error("duplicate pair indices") + prices = price_of(args.model) + if args.usd_per_mtok_in is not None or args.usd_per_mtok_out is not None: + if any(v is None or not math.isfinite(v) or v < 0 for v in (args.usd_per_mtok_in, args.usd_per_mtok_out)): + parser.error("both finite nonnegative prices required") + prices = Prices(args.usd_per_mtok_in, args.usd_per_mtok_in, args.usd_per_mtok_out) + if prices is None: + parser.error("unknown model: supply both price flags") + lane = None + if args.dry_run: + if indices != [0] or args.verify_lane != "vm" or args.board_memory: + parser.error("dry run uses probe index 0, VM verifier and private FileStore only") + rows = [ + dict( + repo="probe", + vulnerable="probe-before", + fixed="probe-after", + family="path", + language="python", + file="app.py", + function="main", + ) + ] + else: + try: + validate_image(args.image or "") + except ValueError: + parser.error("image must be pinned by sha256 digest") + if not args.private_root.resolve().is_relative_to(Path("benchmarks/search/private").resolve()): + parser.error("live working files must stay in benchmarks/search/private") + if args.verify_lane == "vm": + if _sandbox.isolation_mode() != "userns": + parser.error("VM verification requires userns") + try: + _sandbox.isolation_check(timeout_seconds=5) + except _sandbox.IsolationUnavailable: + parser.error("VM userns isolation probe failed") + # Validate private manifest before opening any remote store or client. + if not args.manifest.resolve().is_relative_to(Path("benchmarks/search/private").resolve()): + parser.error("manifest must stay in benchmarks/search/private") + try: + rows = json.loads(args.manifest.read_bytes())["pairs"] + if not isinstance(rows, list) or any(index >= len(rows) for index in indices): + parser.error("pair index outside manifest") + except (OSError, ValueError, KeyError, TypeError): + parser.error("private manifest unreadable or invalid") + for index in indices: + row = rows[index] + if not isinstance(row, dict) or not all( + isinstance(row.get(k), str) and row[k] for k in ("repo", args.side, "fixed", "family", "language", "file", "function") + ): + parser.error("invalid manifest pair") + if row["family"] not in ( + "injection", + "path", + "output_encoding", + "untrusted_destination", + "sql", + "command", + "xss", + "ssrf", + "access_control", + "deserialization", + "configuration", + "secrets", + ): + parser.error("unknown manifest family") + if row["language"] not in ("python", "javascript", "typescript", "php", "java"): + parser.error("unknown manifest language") + path = Path(row["file"]) + if path.is_absolute() or ".." in path.parts or "\\" in row["file"]: + parser.error("invalid candidate path") + + if not re.fullmatch(r"https://github\.com/[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", row["repo"]): + parser.error("invalid repository URL") + if not all(re.fullmatch("[0-9a-f]{40}", row[s]) for s in (args.side, "fixed")): + parser.error("invalid revision") + load_dotenv() + if not os.environ.get(DEEPSEEK_KEY_ENV): + parser.error("DeepSeek key is not configured") + client = OpenRouterChatClient( + api_key=os.environ[DEEPSEEK_KEY_ENV], base_url=os.environ.get(DEEPSEEK_BASE_ENV, DEEPSEEK_BASE_URL), max_attempts=1 + ) + lane = AXLane(args, Workload(args.image, frozenset({"VERIFY_INPUT_URL", "RESULT_URL"}), 900, validate_result)) + root = args.private_root / uuid.uuid4().hex + root.mkdir(parents=True, mode=0o700) + store = open_store(args.board_memory) if args.board_memory else FileStore(root / "boards") + records = [] + started = time.monotonic() + args.out.parent.mkdir(parents=True, exist_ok=True) + with args.out.open("x") as output: + for index in indices: + local = root / str(index) + local.mkdir() + if args.dry_run: + sides, demos = materialise(local / "probe", "path") + client = scripted(load_demo(demos["real"])) + pilot = Pilot(args, rows[index], index, local, store, run_budget, client, prices, lane) + if args.dry_run: + pilot.sides = sides + pilot.record["input_bytes"] = len((sides[0 if args.side == "vulnerable" else 1].checkout / "app.py").read_bytes()) + record = pilot.run() + records.append(record) + output.write(json.dumps(record, allow_nan=False) + "\n") + output.flush() + summary = dict( + type="summary", + searches=len(records), + dry_run=args.dry_run, + side=args.side, + ceiling_usd=args.ceiling_usd, + purpose="feasibility_only", + advantages=["known_location", "labelled_family", "runnable_project_selection_bias"], + fixed_effects_observed=sum(r["fixed_effect_observed"] is True for r in records), + outcomes=dict(Counter(r["outcome"] for r in records)), + spend={k: sum(r["spend"][k] for r in records) for k in ("reserved", "settled")}, + **{k: sum(r[k] for r in records) for k in ("model_calls", "tokens_in", "tokens_out", "tasks_submitted")}, + wall_seconds=time.monotonic() - started, + ) + output.write(json.dumps(summary, allow_nan=False) + "\n") + return int(any(r["end_reason"] == "instrument_failure" for r in records)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/openultrasast/search/worker.py b/src/openultrasast/search/worker.py index 7042f27..6d1043c 100644 --- a/src/openultrasast/search/worker.py +++ b/src/openultrasast/search/worker.py @@ -109,10 +109,14 @@ class Worker: max_output_tokens: int = 1024, max_steps: int = 32, executor: RepositoryExecutor | None = None, + run_budget: SpendBudget | None = None, ) -> None: if getattr(client, "max_attempts", None) != 1: raise ValueError("model client must disable unreserved retries (max_attempts=1)") self.client, self.model = client, model + self.run_budget = run_budget + self.metrics = dict(reserved=0.0, settled=0.0, model_calls=0, tokens_in=0, tokens_out=0) + self.budget_exhausted = False selected_prices = prices if prices is not None else getattr(client, "prices", None) if not isinstance(selected_prices, Prices) or any(not math.isfinite(v) or v < 0 for v in vars(selected_prices).values()): raise ValueError("nonnegative finite model prices required before admission") @@ -144,6 +148,17 @@ class Worker: input_bound * max(self.prices.input_per_m, self.prices.cache_hit_per_m) + self.max_output_tokens * self.prices.output_per_m ) / 1_000_000 reservation = meter.reserve(maximum) + run_reservation = None + try: + if self.run_budget is not None: + run_reservation = self.run_budget.reserve(maximum) + except BudgetExhausted: + reservation.release() + self.budget_exhausted = True + raise + self.metrics["reserved"] += maximum + self.metrics["model_calls"] += 1 + actual = maximum try: raw = self.client.complete_chat_raw( model=self.model, @@ -162,12 +177,21 @@ class Worker: isinstance(usage.get(k), int) and usage[k] >= 0 for k in ("prompt_tokens", "completion_tokens") ): actual = cost_of(usage, self.prices) + self.metrics["tokens_in"] += usage["prompt_tokens"] + self.metrics["tokens_out"] += usage["completion_tokens"] reservation.settle(actual) + if run_reservation is not None: + run_reservation.settle(actual) return cast(dict[str, Any], raw["choices"][0]["message"]) except Exception: if reservation._active: reservation.settle(maximum) + actual = maximum + if run_reservation is not None and run_reservation._active: + run_reservation.settle(maximum) raise + finally: + self.metrics["settled"] += actual def tool(self, name: str, args: dict[str, Any]) -> dict[str, Any]: timeout = self._remaining() @@ -292,6 +316,7 @@ class Worker: return {"status": "execution_failure", "failure": "tool step limit exceeded"} def __call__(self, task: SearchTask) -> dict[str, Any]: + self.budget_exhausted = False self.limits = task.limits self.deadline = time.monotonic() + task.limits.get("task_wall_seconds", 900) meter = SpendBudget(task.limits["max_call_usd"]) @@ -316,6 +341,7 @@ class Worker: "isolation_mode": getattr(self.executor, "isolation_mode", "brain-only"), } except Exception as exc: + self.budget_exhausted = isinstance(exc, BudgetExhausted) # Preserve the diagnostic after removing credentials and URLs. return { "status": "execution_failure", diff --git a/tests/test_search_pilot_run.py b/tests/test_search_pilot_run.py new file mode 100644 index 0000000..2276419 --- /dev/null +++ b/tests/test_search_pilot_run.py @@ -0,0 +1,322 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import json + +import pytest + +from benchmarks.search.pilot_run import main + + +@pytest.mark.parametrize("side,outcome", [("vulnerable", "demonstrated"), ("fixed", "inconclusive")]) +def test_probe_pilot_end_to_end(tmp_path, capsys, side, outcome): + output = tmp_path / "records.jsonl" + assert ( + main( + [ + "--dry-run", + "--side", + side, + "--pairs", + "0", + "--ceiling-usd", + "0.5", + "--out", + str(output), + "--private-root", + str(tmp_path / "private"), + ] + ) + == 0 + ) + records = [json.loads(line) for line in output.read_text().splitlines()] + search, summary = records + assert search["outcome"] == outcome + assert search["model_calls"] == (4 if side == "vulnerable" else 5) + assert search["input_bytes"] > 0 + assert search["board_bytes"] > 0 + assert search["spend"]["settled"] == 0 + assert summary["searches"] == 1 + assert summary["outcomes"] == {outcome: 1} + assert "app.py" not in output.read_text() + + +def test_ceiling_refuses_before_stub_call(tmp_path): + output = tmp_path / "records.jsonl" + main( + [ + "--dry-run", + "--side", + "vulnerable", + "--pairs", + "0", + "--ceiling-usd", + "0", + "--out", + str(output), + "--private-root", + str(tmp_path / "private"), + ] + ) + record = json.loads(output.read_text().splitlines()[0]) + assert record["model_calls"] == 0 + assert record["end_reason"] == "budget_spent" + + +def make_pilot(tmp_path, *, side="fixed", dry_run=False): + from argparse import Namespace + + from benchmarks.search.pilot_run import Pilot + from openultrasast.model.endpoint import Prices + from openultrasast.plane.budget import SpendBudget + from openultrasast.plane.memory import FileStore + from openultrasast.search.worker import StubModel + + args = Namespace(side=side, dry_run=dry_run, model="scripted", max_tokens=128, verify_lane="vm", image="unused") + row = dict( + repo="https://github.com/private/project", + vulnerable="a" * 40, + fixed="b" * 40, + family="path", + language="python", + file="app.py", + function="main", + ) + return Pilot(args, row, 0, tmp_path, FileStore(tmp_path / "board"), SpendBudget(1), StubModel([{"content": "{}"}]), Prices(1, 1, 1)) + + +def test_fixed_verification_never_acquires_vulnerable_revision(tmp_path, monkeypatch): + from benchmarks.search import pilot_run + from openultrasast.search.probe import materialise + + pilot = make_pilot(tmp_path) + sides, demos = materialise(tmp_path / "probe", "path") + schema = json.loads((demos["real"] / "demo.json").read_text()) + (tmp_path / "demo/demo.json").write_text(json.dumps(schema)) + prepared = [] + + def prepare(revision, spec): + prepared.append(revision) + return b"executor-built-bundle" + + monkeypatch.setattr(pilot, "prepare", prepare) + monkeypatch.setattr(pilot_run._sandbox, "isolation_mode", lambda: "userns") + + # Inspect requested pins without relying on archive mechanics. + class StopVerification(Exception): + pass + + def stop(archive, bundle): + bundle.mkdir() + (bundle / "spec.json").write_text(json.dumps({"demo": schema, "family": "path", "timeout_seconds": 5})) + + monkeypatch.setattr(pilot_run, "extract_checkout", stop) + monkeypatch.setattr(pilot_run, "verify", lambda *args, **kwargs: (_ for _ in ()).throw(StopVerification())) + with pytest.raises(StopVerification): + pilot.verification() + assert prepared == ["b" * 40, "b" * 40] + assert "a" * 40 not in json.dumps(pilot.row) + + +def test_worker_failure_is_not_no_evidence(tmp_path): + pilot = make_pilot(tmp_path, dry_run=True) + pilot.worker.client.responses = iter([{"content": "invalid json"}]) + record = pilot.run() + assert record["outcome"] == "inconclusive" + assert record["failures"] == [{"phase": "reason", "reason": "worker_execution_failed"}] + + +def test_shared_budget_reserved_and_settled_on_error(tmp_path): + from openultrasast.model.endpoint import Prices + from openultrasast.plane.budget import SpendBudget + from openultrasast.search.worker import Worker + + run = SpendBudget(0.5) + task = SpendBudget(0.05) + + class Client: + max_attempts = 1 + + def complete_chat_raw(self, **kwargs): + assert run.reserved == task.reserved > 0 + raise RuntimeError("secret-provider-diagnostic") + + worker = Worker(Client(), repo=tmp_path, demo=tmp_path / "demo", prices=Prices(1, 1, 1), run_budget=run) + with pytest.raises(RuntimeError): + worker.reason("facts: []", task) + assert run.reserved == task.reserved == 0 + assert run.spent == task.spent == worker.metrics["settled"] > 0 + assert worker.metrics["model_calls"] == 1 + + +@pytest.mark.parametrize("verify_lane", ["vm", "ax"]) +def test_live_wiring_with_fake_mailbox_no_network(tmp_path, monkeypatch, verify_lane): + import shutil + from dataclasses import asdict + + from benchmarks.search import pilot_run + from openultrasast.search.demo import load_demo + from openultrasast.search.executor import InProcessExecutor + from openultrasast.search.probe import materialise + from openultrasast.search.verify import RunObservation, SideRecord + from openultrasast.search.verify_task import pack_checkout + + pilot = make_pilot(tmp_path, side="vulnerable") + pilot.args.verify_lane = verify_lane + pilot.args.image = "registry.example/search@sha256:" + "0" * 64 + sides, demos = materialise(tmp_path / "probe", "path") + pilot.worker.client = pilot_run.scripted(load_demo(demos["real"])) + events = [] + active = [] + + class Mailbox: + def submit(self, name, args, **kwargs): + if name == "run": + assert args["command"][:3] == ["sh", "-ec", pilot_run.CLONE] + self.side = 0 if args["command"][-1] == "a" * 40 else 1 + events.append(("clone", self.side)) + return {"exit_code": 0} + return InProcessExecutor(sides[self.side].checkout).submit(name, args, **kwargs) + + class Executor: + def __init__(self, lane, archive, *, deadline): + assert 0 < deadline <= 900 + self.client = Mailbox() + + def __enter__(self): + assert not active + active.append(self) + return self + + def __exit__(self, *exc): + active.remove(self) + events.append(("cleanup", self.client.side)) + + def prepare(self, spec): + bundle = tmp_path / ("export-" + str(len(events))) + shutil.copytree(sides[self.client.side].checkout, bundle / "checkout") + (bundle / "products").mkdir() + (bundle / "spec.json").write_text(json.dumps(spec)) + return pack_checkout(bundle) + + def verify_lane_call(item, output, attempt): + assert not active # No executor consumes a slot during verification. + assert item.command == ("python3", "-m", "openultrasast.search.verify_task") + assert set(item.inputs) == {"VERIFY_INPUT_URL"} + events.append(("verify", None)) + # The dispatcher's six fresh repetitions get owned observations. + observed = sum(kind == "verify" for kind, _ in events) <= 3 + return asdict(SideRecord("observed", (RunObservation(observed, "owned", 0.1),), 0.1, "completed", isolation_mode="task-boundary")) + + monkeypatch.setattr(pilot_run, "SearchExecutorTask", Executor) + pilot.lane = verify_lane_call + record = pilot.run() + assert record["outcome"] == "demonstrated", record + assert record["end_reason"] == "goal_met" + assert record["model_calls"] == 4 + assert record["executor_tasks"] == 3 + assert record["verify_tasks"] == (6 if verify_lane == "ax" else 0) + assert record["tasks_submitted"] == (9 if verify_lane == "ax" else 3) + assert record["input_bytes"] > 0 + assert [event for event in events if event[0] == "clone"] == [("clone", 0), ("clone", 0), ("clone", 1)] + assert not active + assert "private/project" not in json.dumps(record) + calls = json.dumps(pilot.worker.client.calls) + assert "b" * 40 not in calls + assert "family" in calls and "main" in calls and "app.py" in calls + + +def test_no_oracle_does_not_acquire_or_build(tmp_path, monkeypatch): + from openultrasast.search.probe import materialise + + pilot = make_pilot(tmp_path) + _, demos = materialise(tmp_path / "probe", "path") + (tmp_path / "demo/demo.json").write_bytes((demos["real"] / "demo.json").read_bytes()) + pilot.family = "injection" + monkeypatch.setattr(pilot, "prepare", lambda *args: pytest.fail("no-oracle must not build")) + assert pilot.verification().outcome == "no_oracle" + assert pilot.record["tasks_submitted"] == 0 + + +@pytest.mark.parametrize("ceiling", ["-1", "nan", "inf"]) +def test_invalid_ceiling_rejected_before_work(tmp_path, ceiling): + with pytest.raises(SystemExit, match="2"): + main(["--dry-run", "--side", "fixed", "--pairs", "0", "--ceiling-usd", ceiling, "--out", str(tmp_path / "record")]) + assert not (tmp_path / "record").exists() + + +@pytest.mark.parametrize("observed,expected", [([True] * 3, True), ([False] * 3, False), ([True, False, True], None)]) +def test_fixed_effect_reporting_is_independent_of_differential(tmp_path, monkeypatch, observed, expected): + from dataclasses import asdict + + from openultrasast.search.budget import SearchBudget + from openultrasast.search.coordinator import SearchTask + from openultrasast.search.verify import RunObservation, SideRecord, VerificationRecord + + pilot = make_pilot(tmp_path) + side = SideRecord("observed", tuple(RunObservation(value, "owned", 0) for value in observed), 0, "completed") + result = VerificationRecord("inconclusive", (side, side), 0, "no differential") + monkeypatch.setattr(pilot, "verification", lambda: result) + task = SearchTask("task-1", "search-1", "verify", "", {}, asdict(SearchBudget())) + assert pilot.dispatch(task)["outcome"] == "inconclusive" + assert pilot.record["fixed_effect_observed"] is expected + + +def test_vm_preflight_refuses_before_client_or_store(tmp_path, monkeypatch): + from benchmarks.search import pilot_run + + monkeypatch.setattr(pilot_run._sandbox, "isolation_mode", lambda: "task-boundary") + monkeypatch.setattr(pilot_run, "OpenRouterChatClient", lambda **kwargs: pytest.fail("client must not open")) + monkeypatch.setattr(pilot_run, "AXLane", lambda *args: pytest.fail("store must not open")) + with pytest.raises(SystemExit, match="2"): + main( + [ + "--side", + "vulnerable", + "--pairs", + "0,1,2", + "--ceiling-usd", + "1.5", + "--image", + "registry.example/search@sha256:" + "0" * 64, + "--out", + str(tmp_path / "records"), + ] + ) + + +def test_actual_usage_settles_shared_budget_and_counts_tokens(tmp_path): + from openultrasast.model.endpoint import Prices + from openultrasast.plane.budget import SpendBudget + from openultrasast.search.worker import Worker + + run, task = SpendBudget(0.5), SpendBudget(0.05) + + class Client: + max_attempts = 1 + + def complete_chat_raw(self, **kwargs): + assert run.reserved == task.reserved > 0.0003 + return {"choices": [{"message": {"content": "{}"}}], "usage": {"prompt_tokens": 100, "completion_tokens": 100}} + + worker = Worker(Client(), repo=tmp_path, demo=tmp_path / "demo", prices=Prices(1, 1, 2), run_budget=run) + assert worker.reason("facts: []", task) == {} + assert worker.metrics["tokens_in"] == worker.metrics["tokens_out"] == 100 + assert worker.metrics["settled"] == run.spent == task.spent == 0.0003 + assert run.reserved == task.reserved == 0 + + +def test_refused_shared_reservation_releases_task_allowance(tmp_path): + from openultrasast.model.endpoint import Prices + from openultrasast.plane.budget import BudgetExhausted, SpendBudget + from openultrasast.search.worker import StubModel, Worker + + run, task = SpendBudget(0), SpendBudget(0.05) + client = StubModel([]) + worker = Worker(client, repo=tmp_path, demo=tmp_path / "demo", prices=Prices(1, 1, 2), run_budget=run) + with pytest.raises(BudgetExhausted): + worker.reason("facts: []", task) + assert not client.calls + assert task.reserved == task.spent == run.reserved == run.spent == 0