mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 19:59:26 +00:00
Merge branch 'worktree-agent-ad03a357c98ed488d'
# Conflicts: # benchmarks/measurements/2026-09-08-module-audit.json # src/openultrasast/model/audit.py
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"modules": 125,
|
||||
"modules": 127,
|
||||
"counts": {
|
||||
"load_bearing": 116,
|
||||
"load_bearing": 118,
|
||||
"standalone_capability": 9
|
||||
},
|
||||
"orphaned": [],
|
||||
@@ -658,7 +658,10 @@
|
||||
"classification": "load_bearing",
|
||||
"reason": "ai-service-plane Req 1-4: the service plane package; runner and reconciler integration in tasks 3-4",
|
||||
"importers": [
|
||||
"plane.tasks.verify"
|
||||
"cli",
|
||||
"plane.reconciler",
|
||||
"plane.tasks.verify",
|
||||
"plane.workspaces"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -673,7 +676,17 @@
|
||||
"module": "plane.manifests",
|
||||
"classification": "load_bearing",
|
||||
"reason": "ai-service-plane Req 1.1-1.4: ax-shaped Task/Workspace/Model and the Run kind, validated before anything runs",
|
||||
"importers": []
|
||||
"importers": [
|
||||
"plane.reconciler"
|
||||
]
|
||||
},
|
||||
{
|
||||
"module": "plane.reconciler",
|
||||
"classification": "load_bearing",
|
||||
"reason": "ai-service-plane Req 3.1-3.5, 7.3: ax-backed Run execution, status attribution and doctor; `ousast plane`",
|
||||
"importers": [
|
||||
"plane.workspaces"
|
||||
]
|
||||
},
|
||||
{
|
||||
"module": "plane.tasks",
|
||||
@@ -703,6 +716,14 @@
|
||||
"plane.tasks.agree"
|
||||
]
|
||||
},
|
||||
{
|
||||
"module": "plane.workspaces",
|
||||
"classification": "load_bearing",
|
||||
"reason": "ai-service-plane task 4: Workspace manifests per case pin of a population; `ousast plane workspaces`",
|
||||
"importers": [
|
||||
"cli"
|
||||
]
|
||||
},
|
||||
{
|
||||
"module": "policy",
|
||||
"classification": "load_bearing",
|
||||
|
||||
@@ -41,6 +41,7 @@ pythonpath = ["src"]
|
||||
markers = [
|
||||
"docker: requires a docker daemon",
|
||||
"semantic: requires openultrasast[semantic] tree-sitter extra",
|
||||
"ax: requires the live ax cluster of ops/ax/up.sh (skipped unless `ousast plane doctor` passes)",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
|
||||
@@ -238,6 +238,20 @@ def main(argv: list[str] | None = None) -> int:
|
||||
candidates.add_argument("--generator", choices=("ir", "ruleset"), default="ir", help="which pass enumerates the candidate sites")
|
||||
candidates.add_argument("--json", action="store_true")
|
||||
|
||||
plane = subparsers.add_parser("plane", help="maintainer: the ax-backed service plane on this host's kind cluster")
|
||||
plane_sub = plane.add_subparsers(dest="plane_command", required=True)
|
||||
plane_run = plane_sub.add_parser("run", help="execute a Run manifest on ax; a rerun skips done tasks")
|
||||
plane_run.add_argument("run_manifest", type=Path)
|
||||
plane_run.add_argument("--workers", type=int, default=1, help="tasks in flight at once (serialize labels never overlap)")
|
||||
plane_run.add_argument("--ax", default="ax", help="path of the ax CLI")
|
||||
plane_status = plane_sub.add_parser("status", help="per-task status and the token attribution table of a run")
|
||||
plane_status.add_argument("run")
|
||||
plane_status.add_argument("--units", action="store_true", help="per-unit rows from each task's units.jsonl")
|
||||
plane_sub.add_parser("doctor", help="check kind, Agent Substrate, the ax controller and the runner image")
|
||||
plane_ws = plane_sub.add_parser("workspaces", help="one Workspace manifest per case pin of a population file")
|
||||
plane_ws.add_argument("population", type=Path)
|
||||
plane_ws.add_argument("--out", type=Path, default=Path("plane/workspaces"))
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
# Diagnosis needs the stage costs, and nothing configures logging, so the default root level of
|
||||
# WARNING silently dropped every informational line. A run that reports nothing is
|
||||
@@ -350,6 +364,8 @@ def main(argv: list[str] | None = None) -> int:
|
||||
return _repos(args.dir, args.repo, fetch=args.fetch, json_out=args.json)
|
||||
if args.command == "model":
|
||||
return _model_candidates(args)
|
||||
if args.command == "plane":
|
||||
return _plane(args)
|
||||
if args.command == "mcp":
|
||||
from .mcp import serve # lazy: keeps the import cycle (mcp -> cli) one-directional
|
||||
|
||||
@@ -1272,6 +1288,28 @@ def _load_scan_findings(run_dir: Path) -> list[StaticFinding]:
|
||||
return load_findings(run_dir / "findings.json")
|
||||
|
||||
|
||||
def _plane(args: argparse.Namespace) -> int:
|
||||
from .plane import reconciler # lazy: the plane is a maintainer surface, not the scan path
|
||||
|
||||
if args.plane_command == "run":
|
||||
result = reconciler.run(args.run_manifest, workers=args.workers, ax=args.ax)
|
||||
print(f"run finished: {result}")
|
||||
return 0 if result == "done" else 1
|
||||
if args.plane_command == "status":
|
||||
print(reconciler.status(args.run, units=args.units))
|
||||
return 0
|
||||
if args.plane_command == "doctor":
|
||||
checks = reconciler.doctor()
|
||||
for name, ok, text in checks:
|
||||
print(f"[{'ok' if ok else 'FAIL'}] {name}: {text}")
|
||||
return 0 if all(ok for _, ok, _ in checks) else 1
|
||||
from .plane.workspaces import workspaces
|
||||
|
||||
written = workspaces(args.population, args.out)
|
||||
print(f"{len(written)} Workspace manifests written to {args.out}")
|
||||
return 0
|
||||
|
||||
|
||||
def _model_candidates(args: argparse.Namespace) -> int:
|
||||
"""What the candidate enumerator can reach, per family and per slice. No model is called anywhere here.
|
||||
|
||||
|
||||
@@ -62,6 +62,8 @@ SPEC_OWNED = {
|
||||
"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",
|
||||
"plane.reconciler": "ai-service-plane Req 3.1-3.5, 7.3: ax-backed Run execution, status attribution and doctor; `ousast plane`",
|
||||
"plane.workspaces": "ai-service-plane task 4: Workspace manifests per case pin of a population; `ousast plane workspaces`",
|
||||
}
|
||||
|
||||
STANDALONE = {
|
||||
|
||||
@@ -0,0 +1,492 @@
|
||||
"""ax-backed reconciler of a ``Run``: submission, status, resume, artifact receipt (ai-service-plane Req 3, 7).
|
||||
|
||||
Every task of a Run is executed by ax on this host's kind cluster; nothing here runs a task as a local subprocess
|
||||
and nothing here knows a prompt, a score or a model call. For each ready task the reconciler renders one YAML
|
||||
file (the Task with its env extended, its Workspaces, one generated Workspace carrying the consumed artifacts as
|
||||
``files`` entries, and the bound Model), runs ``ax apply -f``, polls ``ax get task`` until a terminal phase and
|
||||
``ax delete task`` afterwards. The outcome is read from the delivered ``summary.json``, never from the phase alone.
|
||||
|
||||
Model binding: a Task names its Model in ``metadata.annotations["openultrasast.io/model"]`` (ax's ObjectMeta has no
|
||||
annotations, so the rendered Task carries only ``name``/``atespace``); the task receives ``OUSAST_MODEL`` and the
|
||||
Model's ``spec.parameters`` as ``OUSAST_MODEL_PARAMS`` JSON (Req 7.1). Inputs: ``OUSAST_INPUT_<NAME>`` points into a
|
||||
generated Workspace bound at ``/workspace/.ousast-in/<task>`` (ax mounts every workspace under ``/workspace``).
|
||||
Artifacts: the runner POSTs an uncompressed tar of its output directory to ``OUSAST_ARTIFACT_URL`` with
|
||||
``X-Ousast-Run`` and ``X-Ousast-Task`` headers; the receiver listens on ``0.0.0.0`` at ``OUSAST_ARTIFACT_PORT``
|
||||
(default: a free port) and advertises ``OUSAST_ARTIFACT_HOST`` (default ``172.17.0.1``, the docker bridge address
|
||||
of this host, which kind nodes reach; ``127.0.0.1`` for a fake ax on the host). Run state lives in
|
||||
``~/ousast-results/plane/<run>/`` (``OUSAST_RESULTS`` overrides the root): ``state.json``, a PID ``lock``, one
|
||||
directory per task, ``attribution.json`` written by ``status``; ``OUSAST_POLL_SECONDS`` (default 3) paces polling.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import subprocess
|
||||
import tarfile
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
from .manifests import AX_API_VERSION, Manifests, Model, Run, RunTask, Task, Workspace, load_manifests
|
||||
|
||||
MODEL_ANNOTATION = "openultrasast.io/model"
|
||||
OUTPUT_ROOT = "/workspace/.ousast-out"
|
||||
INPUTS_ROOT = "/workspace/.ousast-in"
|
||||
DEFAULT_ARTIFACT_HOST = "172.17.0.1"
|
||||
FAILURE_PHASES = ("failed", "error")
|
||||
TERMINAL_PHASES = ("succeeded", "completed", *FAILURE_PHASES)
|
||||
UP = "ops/ax/up.sh brings it up"
|
||||
_ROW = "{:<28} {:<20} {:>7} {:>10} {:>10} {:>9} {:>10}"
|
||||
_COLUMNS = ("calls", "prompt", "cache_hit", "output", "usd")
|
||||
_USAGE_KEYS = {"prompt": "prompt_tokens", "cache_hit": "prompt_cache_hit_tokens", "output": "completion_tokens"}
|
||||
|
||||
|
||||
def results_root() -> Path:
|
||||
return Path(os.environ.get("OUSAST_RESULTS") or Path.home() / "ousast-results") / "plane"
|
||||
|
||||
|
||||
def run_dir(run_name: str, root: Path | None = None) -> Path:
|
||||
return (root or results_root()) / run_name
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(UTC).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
def _write_json(path: Path, payload: object) -> None:
|
||||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||||
tmp.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
||||
tmp.replace(path)
|
||||
|
||||
|
||||
def _read_json(path: Path) -> Any:
|
||||
try:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _ax_name(*parts: str) -> str:
|
||||
return re.sub(r"[^a-z0-9-]+", "-", "-".join(parts).lower()).strip("-")[:63]
|
||||
|
||||
|
||||
def _is(phase: str, words: tuple[str, ...]) -> bool:
|
||||
return any(word in phase.lower() for word in words)
|
||||
|
||||
|
||||
def load_run(run_manifest: Path) -> tuple[Run, Manifests]:
|
||||
"""Load the Run file; a file under ``.../runs/`` also loads the sibling ``tasks``, ``models`` and ``workspaces``."""
|
||||
paths = [run_manifest]
|
||||
if run_manifest.parent.name == "runs":
|
||||
for kind in ("tasks", "models", "workspaces"):
|
||||
paths.extend(sorted((run_manifest.parent.parent / kind).glob("*.y*ml")))
|
||||
manifests = load_manifests(paths)
|
||||
if len(manifests.runs) != 1:
|
||||
raise ValueError(f"{run_manifest}: expected exactly one Run, found {len(manifests.runs)}")
|
||||
return next(iter(manifests.runs.values())), manifests
|
||||
|
||||
|
||||
def _doc(kind: str, name: str, spec: dict[str, Any], namespace: str | None = None) -> dict[str, Any]:
|
||||
meta = {"name": name, **({"atespace": namespace} if namespace else {})}
|
||||
return {"apiVersion": AX_API_VERSION, "kind": kind, "metadata": meta, "spec": spec}
|
||||
|
||||
|
||||
def _fields(value: object) -> dict[str, Any]:
|
||||
return {k: v for k, v in vars(value).items() if v is not None}
|
||||
|
||||
|
||||
def _workspace_doc(ws: Workspace) -> dict[str, Any]:
|
||||
spec: dict[str, Any] = {}
|
||||
if ws.git:
|
||||
spec["git"] = [_fields(g) for g in ws.git]
|
||||
if ws.files:
|
||||
spec["files"] = [_fields(f) for f in ws.files]
|
||||
if ws.mcp:
|
||||
spec["mcp"] = {"registries": list(ws.mcp.registries), "servers": list(ws.mcp.servers)}
|
||||
if ws.skills:
|
||||
spec["skills"] = {"registries": list(ws.skills.registries), **({"path": ws.skills.path} if ws.skills.path else {})}
|
||||
return _doc("Workspace", ws.metadata.name, spec, ws.metadata.namespace)
|
||||
|
||||
|
||||
def _model_doc(model: Model) -> dict[str, Any]:
|
||||
spec: dict[str, Any] = {"provider": model.provider, "model": model.model}
|
||||
if model.secret_key:
|
||||
spec["secretKey"] = _fields(model.secret_key)
|
||||
if model.parameters:
|
||||
spec["parameters"] = dict(model.parameters)
|
||||
return _doc("Model", model.metadata.name, spec, model.metadata.namespace)
|
||||
|
||||
|
||||
def _input_files(entry: RunTask, base: Path) -> list[dict[str, str]]:
|
||||
files: list[dict[str, str]] = []
|
||||
for ref in entry.inputs.values():
|
||||
source = base / ref
|
||||
for path in sorted(source.rglob("*")) if source.is_dir() else [source]:
|
||||
if path.is_file():
|
||||
files.append({"path": path.relative_to(base).as_posix(), "content": path.read_text(encoding="utf-8")})
|
||||
return files
|
||||
|
||||
|
||||
def render_task(run: Run, entry: RunTask, manifests: Manifests, base: Path, artifact_url: str) -> list[dict[str, Any]]:
|
||||
"""The documents ``ax apply`` receives for one Run task: Task, its Model, its Workspaces, the inputs Workspace."""
|
||||
task: Task = manifests.tasks[entry.task]
|
||||
ax_name = _ax_name(run.metadata.name, entry.name)
|
||||
env = {e.name: e.value for e in task.env}
|
||||
env.update(OUSAST_RUN=run.metadata.name, OUSAST_TASK=entry.name, OUSAST_OUTPUT_DIR=f"{OUTPUT_ROOT}/{entry.name}")
|
||||
env["OUSAST_ARTIFACT_URL"] = artifact_url
|
||||
budget = _fields(entry.budget) if entry.budget else {}
|
||||
env.update({f"OUSAST_BUDGET_{k.upper()}": str(v) for k, v in budget.items()})
|
||||
for input_name, ref in entry.inputs.items():
|
||||
env[f"OUSAST_INPUT_{re.sub(r'[^A-Z0-9]+', '_', input_name.upper())}"] = f"{INPUTS_ROOT}/{entry.name}/{ref}"
|
||||
docs: list[dict[str, Any]] = []
|
||||
model_name = task.metadata.annotations.get(MODEL_ANNOTATION)
|
||||
if model_name:
|
||||
model = manifests.models[model_name]
|
||||
env.update(OUSAST_MODEL=model_name, OUSAST_MODEL_PARAMS=json.dumps(dict(model.parameters), sort_keys=True))
|
||||
docs.append(_model_doc(model))
|
||||
bindings = [_fields(b) for b in task.workspaces]
|
||||
docs.extend(_workspace_doc(manifests.workspaces[b.name]) for b in task.workspaces)
|
||||
if entry.inputs:
|
||||
docs.append(_doc("Workspace", f"{ax_name}-inputs", {"files": _input_files(entry, base)}, task.metadata.namespace))
|
||||
bindings.append({"name": f"{ax_name}-inputs", "path": f"{INPUTS_ROOT}/{entry.name}"})
|
||||
spec: dict[str, Any] = {"command": list(task.command), "env": [{"name": k, "value": v} for k, v in env.items()]}
|
||||
if task.image:
|
||||
spec["image"] = task.image
|
||||
if task.resources:
|
||||
spec["resources"] = {k: _fields(v) for k, v in _fields(task.resources).items()}
|
||||
if bindings:
|
||||
spec["workspaces"] = bindings
|
||||
if task.debug:
|
||||
spec["debug"] = True
|
||||
return [_doc("Task", ax_name, spec, task.metadata.namespace), *docs]
|
||||
|
||||
|
||||
class Ax:
|
||||
"""The three ``ax`` verbs the reconciler uses; the executable path is injectable so tests use a fake."""
|
||||
|
||||
def __init__(self, executable: str = "ax") -> None:
|
||||
self.executable = executable
|
||||
|
||||
def _run(self, *args: str) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run([self.executable, *args], capture_output=True, text=True, check=False)
|
||||
|
||||
def apply(self, path: Path) -> None:
|
||||
proc = self._run("apply", "-f", str(path))
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(f"ax apply failed ({proc.returncode}): {proc.stderr.strip() or proc.stdout.strip()}")
|
||||
|
||||
def phase(self, name: str) -> str:
|
||||
"""``status.phase`` from ``ax get task <name>`` (YAML; Pending when unset; ``Unknown: ...`` when ax fails)."""
|
||||
proc = self._run("get", "task", name)
|
||||
if proc.returncode != 0:
|
||||
return "Unknown: " + (proc.stderr.strip() or proc.stdout.strip())
|
||||
try:
|
||||
doc = yaml.safe_load(proc.stdout)
|
||||
except yaml.YAMLError:
|
||||
return "Unknown: unreadable ax output"
|
||||
status = doc.get("status") if isinstance(doc, dict) else None
|
||||
return str(status.get("phase") or "Pending") if isinstance(status, dict) else "Pending"
|
||||
|
||||
def delete(self, name: str) -> None:
|
||||
self._run("delete", "task", name)
|
||||
|
||||
|
||||
class Receiver(ThreadingHTTPServer):
|
||||
"""Accepts one tar per running task and extracts it under ``<run dir>/<task>/``."""
|
||||
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, run_name: str, base: Path, port: int) -> None:
|
||||
super().__init__(("0.0.0.0", port), _Handler)
|
||||
self.run_name, self.base, self.running, self.delivered = run_name, base, set[str](), set[str]()
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def url(self) -> str:
|
||||
return f"http://{os.environ.get('OUSAST_ARTIFACT_HOST') or DEFAULT_ARTIFACT_HOST}:{self.server_address[1]}/"
|
||||
|
||||
def accept(self, headers: Any, body: bytes) -> tuple[int, str]:
|
||||
task = headers.get("X-Ousast-Task", "")
|
||||
if headers.get("X-Ousast-Run") != self.run_name or task not in self.running:
|
||||
return 403, "unknown run or task"
|
||||
if not str(headers.get("Content-Type", "")).startswith("application/x-tar"):
|
||||
return 415, "expected application/x-tar"
|
||||
dest = (self.base / task).resolve()
|
||||
try:
|
||||
with tarfile.open(fileobj=io.BytesIO(body), mode="r:") as tar:
|
||||
members = tar.getmembers()
|
||||
for m in members:
|
||||
if m.issym() or m.islnk() or not (dest / m.name).resolve().is_relative_to(dest):
|
||||
return 400, f"rejected member {m.name}"
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
tar.extractall(dest, members) # every member checked above
|
||||
except tarfile.TarError as exc:
|
||||
return 400, f"bad tar: {exc}"
|
||||
with self.lock:
|
||||
self.delivered.add(task)
|
||||
return 200, "ok"
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
server: Receiver = self.server # type: ignore[assignment]
|
||||
body = self.rfile.read(int(self.headers.get("Content-Length") or 0))
|
||||
code, text = server.accept(self.headers, body)
|
||||
self.send_response(code)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(text.encode())
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return
|
||||
|
||||
|
||||
# --- the run -------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _State:
|
||||
path: Path
|
||||
data: dict[str, Any]
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def set(self, task: str, **fields: object) -> None:
|
||||
with self.lock:
|
||||
self.data["tasks"].setdefault(task, {"status": "pending", "started": None, "finished": None, "usd": None, "model": None})
|
||||
self.data["tasks"][task].update(fields)
|
||||
_write_json(self.path, self.data)
|
||||
|
||||
def status(self, task: str) -> str:
|
||||
return str(self.data["tasks"].get(task, {}).get("status", "pending"))
|
||||
|
||||
|
||||
def outcome_of(summary: Mapping[str, Any] | None, phase: str) -> tuple[str, str]:
|
||||
"""(status, reason) of a finished task from its delivered summary and ax's terminal phase."""
|
||||
if _is(phase, FAILURE_PHASES):
|
||||
return "failed", f"ax phase {phase}"
|
||||
if summary is None:
|
||||
return "failed", "no summary.json delivered"
|
||||
if summary.get("status") == "done" and summary.get("units_done") == summary.get("units_total"):
|
||||
return "done", ""
|
||||
if summary.get("status") == "unfinished":
|
||||
return "unfinished", str(summary.get("reason") or "budget stop")
|
||||
return "failed", str(summary.get("reason") or summary.get("status") or "task did not report done")
|
||||
|
||||
|
||||
def _acquire_lock(base: Path) -> Path:
|
||||
lock = base / "lock"
|
||||
pid = _read_json(lock)
|
||||
if isinstance(pid, int) and pid != os.getpid():
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except OSError:
|
||||
pass
|
||||
else:
|
||||
raise RuntimeError(f"run is held by pid {pid} ({lock})")
|
||||
lock.write_text(json.dumps(os.getpid()), encoding="utf-8")
|
||||
return lock
|
||||
|
||||
|
||||
def _execute(ax: Ax, run: Run, entry: RunTask, manifests: Manifests, base: Path, receiver: Receiver, poll: float) -> tuple[str, str]:
|
||||
ax_name = _ax_name(run.metadata.name, entry.name)
|
||||
manifest = base / entry.name / "task.yaml"
|
||||
manifest.parent.mkdir(parents=True, exist_ok=True)
|
||||
docs = render_task(run, entry, manifests, base, receiver.url())
|
||||
manifest.write_text(yaml.safe_dump_all(docs, sort_keys=False), encoding="utf-8")
|
||||
try:
|
||||
ax.apply(manifest)
|
||||
except RuntimeError as exc:
|
||||
return "failed", str(exc)
|
||||
phase, unknown = "Pending", 0
|
||||
try:
|
||||
while not _is(phase, TERMINAL_PHASES):
|
||||
time.sleep(poll)
|
||||
phase = ax.phase(ax_name)
|
||||
unknown = unknown + 1 if phase.startswith("Unknown") else 0
|
||||
if unknown >= 10: # ten consecutive unreadable polls: ax itself is broken, not the task
|
||||
phase = "Error: " + phase
|
||||
for _ in range(5): # delivery may trail the phase change by a few seconds
|
||||
if entry.name in receiver.delivered or _is(phase, FAILURE_PHASES):
|
||||
break
|
||||
time.sleep(poll)
|
||||
finally:
|
||||
ax.delete(ax_name)
|
||||
return outcome_of(_read_json(base / entry.name / "summary.json"), phase)
|
||||
|
||||
|
||||
def run(run_manifest: Path, *, workers: int = 1, ax: str = "ax", results_root: Path | None = None) -> str:
|
||||
"""Execute a Run on ax and return its status: ``done``, ``unfinished`` or ``failed`` (Req 3.1-3.3, 2.2)."""
|
||||
run_spec, manifests = load_run(Path(run_manifest))
|
||||
base = run_dir(run_spec.metadata.name, results_root)
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
lock = _acquire_lock(base)
|
||||
poll = float(os.environ.get("OUSAST_POLL_SECONDS") or 3)
|
||||
fresh = {"run": run_spec.metadata.name, "started": _now(), "tasks": {}}
|
||||
state = _State(base / "state.json", _read_json(base / "state.json") or fresh)
|
||||
receiver = Receiver(run_spec.metadata.name, base, int(os.environ.get("OUSAST_ARTIFACT_PORT") or 0))
|
||||
threading.Thread(target=receiver.serve_forever, daemon=True).start()
|
||||
cli, finished = Ax(ax), queue.Queue[tuple[str, str, str]]()
|
||||
pending = [t for t in run_spec.tasks if state.status(t.name) != "done"]
|
||||
for entry in run_spec.tasks:
|
||||
state.set(entry.name, **({"status": "pending"} if entry in pending else {}))
|
||||
running: dict[str, RunTask] = {}
|
||||
halted = False
|
||||
|
||||
def worker(entry: RunTask) -> None:
|
||||
if state.status(entry.name) in ("running", "suspended"):
|
||||
cli.delete(_ax_name(run_spec.metadata.name, entry.name)) # an interrupted attempt; ax must not keep it
|
||||
state.set(entry.name, status="running", started=_now(), finished=None)
|
||||
finished.put((entry.name, *_execute(cli, run_spec, entry, manifests, base, receiver, poll)))
|
||||
|
||||
try:
|
||||
while pending or running:
|
||||
blocked = {t.serialize for t in running.values() if t.serialize}
|
||||
for entry in list(pending):
|
||||
if halted or len(running) >= max(workers, 1) or (entry.serialize and entry.serialize in blocked):
|
||||
continue
|
||||
if not all(state.status(p) == "done" for p in entry.producers):
|
||||
continue
|
||||
pending.remove(entry)
|
||||
running[entry.name] = entry
|
||||
blocked.add(entry.serialize or "")
|
||||
with receiver.lock:
|
||||
receiver.running.add(entry.name)
|
||||
threading.Thread(target=worker, args=(entry,), daemon=True).start()
|
||||
if not running:
|
||||
break # nothing can start: a producer stopped short
|
||||
name, status, reason = finished.get()
|
||||
del running[name]
|
||||
summary = _read_json(base / name / "summary.json") or {}
|
||||
state.set(name, status=status, finished=_now(), usd=summary.get("usd"), model=summary.get("model"), reason=reason or None)
|
||||
halted = halted or status == "failed"
|
||||
if status == "unfinished":
|
||||
pending = [t for t in pending if name not in t.producers]
|
||||
finally:
|
||||
receiver.shutdown()
|
||||
lock.unlink(missing_ok=True)
|
||||
statuses = [state.status(t.name) for t in run_spec.tasks]
|
||||
result = "failed" if "failed" in statuses else "done" if all(s == "done" for s in statuses) else "unfinished"
|
||||
state.data.update(status=result, finished=_now())
|
||||
_write_json(state.path, state.data)
|
||||
return result
|
||||
|
||||
|
||||
# --- status and attribution (Req 7.3) ------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _number(value: object) -> float | None:
|
||||
return value if isinstance(value, int | float) and not isinstance(value, bool) else None
|
||||
|
||||
|
||||
def _row(name: str, model: str | None, payload: Mapping[str, Any]) -> dict[str, Any]:
|
||||
raw = payload.get("usage")
|
||||
usage: Mapping[str, Any] = raw if isinstance(raw, Mapping) else {}
|
||||
row: dict[str, Any] = {"task": name, "model": model, "calls": _number(payload.get("calls", payload.get("turns")))}
|
||||
row.update({column: _number(usage.get(key)) for column, key in _USAGE_KEYS.items()})
|
||||
row["usd"] = _number(payload.get("usd"))
|
||||
return row
|
||||
|
||||
|
||||
def _sum(rows: Iterable[Mapping[str, Any]], label: str, model: str | None) -> dict[str, Any]:
|
||||
"""Column sums over rows bound to a Model; a missing value makes the sum ``n/a``, never a smaller number."""
|
||||
priced = [r for r in rows if r.get("model")]
|
||||
total: dict[str, Any] = {"task": label, "model": model}
|
||||
for column in _COLUMNS:
|
||||
values = [r.get(column) for r in priced]
|
||||
total[column] = None if not values or any(v is None for v in values) else round(sum(v or 0 for v in values), 6)
|
||||
return total
|
||||
|
||||
|
||||
def _fmt(row: Mapping[str, Any]) -> str:
|
||||
cells = [str(row["task"])[:28], str(row.get("model") or "-")[:20]]
|
||||
for column in _COLUMNS:
|
||||
value = row.get(column)
|
||||
cells.append("n/a" if value is None else f"{value:.4f}" if column == "usd" else str(int(value)))
|
||||
return _ROW.format(*cells)
|
||||
|
||||
|
||||
def attribution(run_name: str, *, units: bool = False, results_root: Path | None = None) -> dict[str, Any]:
|
||||
"""Token attribution from the tasks' ``summary.json`` usage fields; ``units`` adds rows from ``units.jsonl``."""
|
||||
base = run_dir(run_name, results_root)
|
||||
state = _read_json(base / "state.json") or {}
|
||||
tasks: dict[str, dict[str, Any]] = state.get("tasks") or {}
|
||||
rows: list[dict[str, Any]] = []
|
||||
unit_rows: dict[str, list[dict[str, Any]]] = {}
|
||||
for name, info in tasks.items():
|
||||
summary = _read_json(base / name / "summary.json") or {}
|
||||
rows.append({**_row(name, summary.get("model") or info.get("model"), summary), "status": info.get("status")})
|
||||
if units and (base / name / "units.jsonl").is_file():
|
||||
lines = (base / name / "units.jsonl").read_text(encoding="utf-8").splitlines()
|
||||
unit_rows[name] = [
|
||||
_row(str(u.get("unit") or u.get("path") or u.get("candidate") or f"unit {i}"), rows[-1]["model"], u)
|
||||
for i, u in enumerate(json.loads(line) for line in lines if line.strip())
|
||||
]
|
||||
names = sorted({r["model"] for r in rows if r["model"]})
|
||||
models = [_sum([r for r in rows if r["model"] == m], f"subtotal {m}", m) for m in names]
|
||||
return {"run": run_name, "tasks": rows, "models": models, "total": _sum(rows, "total", None), "units": unit_rows}
|
||||
|
||||
|
||||
def status(run_name: str, *, units: bool = False, results_root: Path | None = None) -> str:
|
||||
"""Per-task status and the attribution table (Req 3.2, 7.3); written to ``<run dir>/attribution.json``."""
|
||||
report = attribution(run_name, units=units, results_root=results_root)
|
||||
base = run_dir(run_name, results_root)
|
||||
state = _read_json(base / "state.json") or {}
|
||||
header = _ROW.format("TASK", "MODEL", "CALLS", "PROMPT", "CACHE-HIT", "OUTPUT", "USD")
|
||||
lines = [f"run {run_name}: {state.get('status') or 'not started'} ({base})", "", header]
|
||||
for row in report["tasks"]:
|
||||
lines.append(_fmt(row) + f" [{row.get('status') or 'pending'}]")
|
||||
lines.extend(" " + _fmt(u) for u in report["units"].get(row["task"], []))
|
||||
lines.append("-" * len(header))
|
||||
lines.extend(_fmt(m) for m in report["models"])
|
||||
lines.append(_fmt(report["total"]))
|
||||
if base.is_dir():
|
||||
_write_json(base / "attribution.json", report)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# --- doctor (Req 3.5) ----------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sh(*args: str) -> tuple[bool, str]:
|
||||
try:
|
||||
proc = subprocess.run(args, capture_output=True, text=True, timeout=60, check=False)
|
||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||
return False, str(exc)
|
||||
return proc.returncode == 0, (proc.stdout + proc.stderr).strip()
|
||||
|
||||
|
||||
def _pods_ready(namespace: str, context: str) -> tuple[bool, str]:
|
||||
ok, out = _sh("kubectl", "--context", context, "-n", namespace, "get", "pods", "--no-headers")
|
||||
if not ok or not out:
|
||||
return False, out or "no pods"
|
||||
rows = [line.split() for line in out.splitlines()]
|
||||
bad = [r[0] for r in rows if len(r) > 2 and r[2] not in ("Completed", "Succeeded") and r[1].split("/")[0] != r[1].split("/")[-1]]
|
||||
return not bad, f"not ready: {', '.join(bad)}" if bad else f"{len(rows)} pods ready"
|
||||
|
||||
|
||||
def doctor() -> list[tuple[str, bool, str]]:
|
||||
"""Four checks of the host's ax: kind, Agent Substrate, the ax controller, the runner image (Req 3.5)."""
|
||||
context = "kind-" + (os.environ.get("KIND_CLUSTER_NAME") or "ousast")
|
||||
ok, out = _sh("kubectl", "--context", context, "cluster-info")
|
||||
checks = [("kind cluster", ok, "reachable" if ok else f"context {context} unreachable ({out[:80]}); {UP}")]
|
||||
for label, ns in (("agent substrate (ate-system)", "ate-system"), ("ax controller (ax-system)", "ax-system")):
|
||||
ok, out = _pods_ready(ns, context)
|
||||
checks.append((label, ok, out if ok else f"{out}; {UP}"))
|
||||
ok, out = _sh("curl", "-s", "localhost:5001/v2/_catalog")
|
||||
present = ok and "ousast-runner" in out
|
||||
text = "ousast-runner present" if present else f"missing from localhost:5001 ({out[:80]}); {UP} and loads it"
|
||||
checks.append(("runner image in kind registry", present, text))
|
||||
return checks
|
||||
|
||||
|
||||
__all__ = ["Ax", "Receiver", "attribution", "doctor", "load_run", "outcome_of", "render_task", "run", "run_dir", "status"]
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Workspace manifests for the cases of an independent population file (ai-service-plane task 4).
|
||||
|
||||
``ousast plane workspaces <population.toml> [--out DIR]`` writes one ax ``Workspace`` per case pin (vulnerable,
|
||||
fixed, benign base, benign tip) with the case's repository URL and the pinned commit, so a Run binds a case at a
|
||||
pin by name. The layout of a case is that of ``benchmarks/independent/evaluate.py``'s ``pins()``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from .reconciler import _ax_name, _doc
|
||||
|
||||
__all__ = ["workspaces"]
|
||||
|
||||
|
||||
def workspaces(population: Path, out: Path) -> list[Path]:
|
||||
"""One Workspace manifest per case pin of a population file; returns the paths written."""
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
written: list[Path] = []
|
||||
for case in tomllib.loads(population.read_text(encoding="utf-8"))["case"]:
|
||||
benign = case["benign"]
|
||||
pins = {"vulnerable": case["vulnerable"], "fixed": case["fixed"], "benign-base": benign["base"], "benign-tip": benign["tip"]}
|
||||
for label, commit in pins.items():
|
||||
name = _ax_name(case["id"], label)
|
||||
doc = _doc("Workspace", name, {"git": [{"name": "repo", "repo": case["repo"], "commit": commit}]})
|
||||
text = f"# generated by `ousast plane workspaces` from {population.name}\n" + yaml.safe_dump(doc, sort_keys=False)
|
||||
(out / f"{name}.yaml").write_text(text, encoding="utf-8")
|
||||
written.append(out / f"{name}.yaml")
|
||||
return written
|
||||
@@ -0,0 +1,43 @@
|
||||
"""ai-service-plane Req 3.5/3.6: one trivial Task end to end on the real ax cluster of ``ops/ax/up.sh``.
|
||||
|
||||
Marker ``ax``; skipped unless every check of ``ousast plane doctor`` passes, so the host suite never depends on
|
||||
the cluster. Run it on purpose: ``pytest -m ax tests/test_plane_ax_live.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from openultrasast.plane.reconciler import Ax, doctor
|
||||
|
||||
pytestmark = pytest.mark.ax
|
||||
SMOKE = Path("ops/ax/smoke-task.yaml")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def live_ax() -> Ax:
|
||||
checks = doctor()
|
||||
failing = [f"{name}: {text}" for name, ok, text in checks if not ok]
|
||||
if failing:
|
||||
pytest.skip("no live ax cluster: " + "; ".join(failing))
|
||||
if not SMOKE.is_file():
|
||||
pytest.skip(f"{SMOKE} is not in this checkout")
|
||||
return Ax()
|
||||
|
||||
|
||||
def test_smoke_task_completes_on_the_cluster(live_ax: Ax) -> None:
|
||||
live_ax.apply(SMOKE)
|
||||
deadline = time.monotonic() + 600
|
||||
phase = "Pending"
|
||||
try:
|
||||
while time.monotonic() < deadline:
|
||||
phase = live_ax.phase("ousast-smoke")
|
||||
if any(word in phase.lower() for word in ("completed", "succeeded", "failed", "error")):
|
||||
break
|
||||
time.sleep(5)
|
||||
finally:
|
||||
live_ax.delete("ousast-smoke")
|
||||
assert phase.lower() in ("completed", "succeeded"), f"smoke task ended in phase {phase!r}"
|
||||
@@ -0,0 +1,485 @@
|
||||
"""ai-service-plane Req 3.1-3.5, 2.2, 7.3: a fake ``ax`` drives a Run through the reconciler.
|
||||
|
||||
The fake is a Python script the tests write into ``tmp_path``: it logs every argv line, parses what ``apply``
|
||||
receives, answers ``get task`` with Running until a scripted poll count and then a terminal phase, and on reaching
|
||||
it POSTs a scripted tar to the ``OUSAST_ARTIFACT_URL`` it found in the applied Task's env. Nothing here runs a
|
||||
task; every outcome comes from a scripted ``summary.json`` delivery.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import tarfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from openultrasast.plane import reconciler
|
||||
from openultrasast.plane.workspaces import workspaces
|
||||
|
||||
RECONCILER = Path("src/openultrasast/plane/reconciler.py")
|
||||
|
||||
FAKE_AX = '''#!{python}
|
||||
"""Fake ax CLI for the reconciler tests: argv log, scripted phases, scripted deliveries."""
|
||||
import io, json, sys, tarfile, time, urllib.error, urllib.request
|
||||
from pathlib import Path
|
||||
import yaml
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
SCRIPT = json.loads((HERE / "script.json").read_text())
|
||||
argv = sys.argv[1:]
|
||||
(HERE / "ax.log").open("a").write(f"{{time.monotonic():.4f}} " + " ".join(argv) + "\\n")
|
||||
(HERE / "applied").mkdir(exist_ok=True)
|
||||
if argv[:2] == ["apply", "-f"]:
|
||||
docs = list(yaml.safe_load_all(Path(argv[2]).read_text()))
|
||||
for doc in docs:
|
||||
if doc["kind"] == "Task":
|
||||
env = {{e["name"]: e["value"] for e in doc["spec"].get("env", [])}}
|
||||
(HERE / "applied" / (doc["metadata"]["name"] + ".json")).write_text(json.dumps({{"docs": docs, "env": env, "polls": 0}}))
|
||||
print("task created")
|
||||
elif argv[:2] == ["get", "task"]:
|
||||
record = HERE / "applied" / (argv[2] + ".json")
|
||||
if not record.exists():
|
||||
print("Error: not found", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
rec = json.loads(record.read_text())
|
||||
rec["polls"] += 1
|
||||
env, task = rec["env"], rec["env"]["OUSAST_TASK"]
|
||||
script = SCRIPT.get(task, {{}})
|
||||
needed = script.get("polls", 1)
|
||||
phase = "Running" if rec["polls"] < needed else script.get("phase", "Completed")
|
||||
if rec["polls"] == needed and script.get("deliver", True):
|
||||
buf = io.BytesIO()
|
||||
with tarfile.open(fileobj=buf, mode="w") as tar:
|
||||
for path, content in script.get("files", {{}}).items():
|
||||
data = (content if isinstance(content, str) else json.dumps(content)).encode()
|
||||
info = tarfile.TarInfo(path)
|
||||
info.size = len(data)
|
||||
tar.addfile(info, io.BytesIO(data))
|
||||
headers = {{"Content-Type": "application/x-tar", "X-Ousast-Task": script.get("as_task", task), "X-Ousast-Run": env["OUSAST_RUN"]}}
|
||||
req = urllib.request.Request(env["OUSAST_ARTIFACT_URL"], data=buf.getvalue(), method="POST", headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req) as resp:
|
||||
(HERE / "ax.log").open("a").write(f"{{time.monotonic():.4f}} delivered {{task}} {{resp.status}}\\n")
|
||||
except urllib.error.HTTPError as exc:
|
||||
(HERE / "ax.log").open("a").write(f"{{time.monotonic():.4f}} rejected {{task}} {{exc.code}}\\n")
|
||||
record.write_text(json.dumps(rec))
|
||||
print(yaml.safe_dump({{"apiVersion": "ax.io/v1alpha1", "kind": "Task", "metadata": {{"name": argv[2]}}, "status": {{"phase": phase}}}}))
|
||||
elif argv[:2] == ["delete", "task"]:
|
||||
print("deleted")
|
||||
'''
|
||||
|
||||
MANIFESTS = """
|
||||
apiVersion: ax.io/v1alpha1
|
||||
kind: Model
|
||||
metadata:
|
||||
name: deepseek-flash
|
||||
annotations: {openultrasast.io/provider-extension: "true"}
|
||||
spec:
|
||||
provider: deepseek
|
||||
model: deepseek-chat
|
||||
parameters: {cache_hit_per_m: 0.07, input_per_m: 0.27, output_per_m: 1.1}
|
||||
---
|
||||
apiVersion: ax.io/v1alpha1
|
||||
kind: Workspace
|
||||
metadata:
|
||||
name: case-pin
|
||||
spec:
|
||||
git: [{name: repo, repo: https://example.invalid/repo.git, commit: abc123}]
|
||||
---
|
||||
apiVersion: ax.io/v1alpha1
|
||||
kind: Task
|
||||
metadata:
|
||||
name: repo-facts
|
||||
spec:
|
||||
image: ousast-runner:dev
|
||||
command: [repo_facts]
|
||||
workspaces: [{name: case-pin, path: /workspace}]
|
||||
---
|
||||
apiVersion: ax.io/v1alpha1
|
||||
kind: Task
|
||||
metadata:
|
||||
name: verify
|
||||
annotations: {openultrasast.io/model: deepseek-flash}
|
||||
spec:
|
||||
image: ousast-runner:dev
|
||||
command: [verify]
|
||||
env: [{name: OUSAST_PASS, value: a}]
|
||||
workspaces: [{name: case-pin, path: /workspace}]
|
||||
---
|
||||
apiVersion: ax.io/v1alpha1
|
||||
kind: Task
|
||||
metadata:
|
||||
name: agree
|
||||
spec:
|
||||
command: [agree]
|
||||
---
|
||||
apiVersion: openultrasast.io/v1alpha1
|
||||
kind: Run
|
||||
metadata:
|
||||
name: {run}
|
||||
spec:
|
||||
tasks:
|
||||
{tasks}
|
||||
"""
|
||||
|
||||
CHAIN = """
|
||||
- name: repo-facts
|
||||
task: repo-facts
|
||||
outputs: [facts.json]
|
||||
serialize: engine
|
||||
- name: verify
|
||||
task: verify
|
||||
inputs: {facts: repo-facts/facts.json}
|
||||
outputs: [agreed.json]
|
||||
budget: {usd: 2.5, calls: 40}
|
||||
- name: agree
|
||||
task: agree
|
||||
dependsOn: [verify]
|
||||
inputs: {pass: verify/agreed.json}
|
||||
"""
|
||||
|
||||
FACTS_SUMMARY = {"status": "done", "units_done": 3, "units_total": 3, "usd": None, "calls": 0, "usage": {}}
|
||||
VERIFY_SUMMARY = {
|
||||
"status": "done",
|
||||
"units_done": 2,
|
||||
"units_total": 2,
|
||||
"usd": 0.5,
|
||||
"calls": 4,
|
||||
"model": "deepseek-flash",
|
||||
"usage": {"prompt_tokens": 1000, "prompt_cache_hit_tokens": 300, "completion_tokens": 200},
|
||||
}
|
||||
AGREE_SUMMARY = {"status": "done", "units_done": 1, "units_total": 1, "usd": None, "calls": 0}
|
||||
SCRIPT = {
|
||||
"repo-facts": {"polls": 2, "files": {"summary.json": FACTS_SUMMARY, "facts.json": {"a.py": {"functions": ["f"]}}}},
|
||||
"verify": {"polls": 2, "files": {"summary.json": VERIFY_SUMMARY, "agreed.json": {"agreed": ["a.py:f"]}}},
|
||||
"agree": {"polls": 1, "files": {"summary.json": AGREE_SUMMARY}},
|
||||
}
|
||||
|
||||
|
||||
class Fake:
|
||||
def __init__(self, root: Path) -> None:
|
||||
self.root = root
|
||||
self.ax = root / "ax"
|
||||
self.ax.write_text(FAKE_AX.format(python=sys.executable), encoding="utf-8")
|
||||
self.ax.chmod(self.ax.stat().st_mode | stat.S_IXUSR)
|
||||
|
||||
def script(self, script: dict) -> None:
|
||||
(self.root / "script.json").write_text(json.dumps(script), encoding="utf-8")
|
||||
|
||||
def log(self) -> list[tuple[float, list[str]]]:
|
||||
lines = (self.root / "ax.log").read_text(encoding="utf-8").splitlines() if (self.root / "ax.log").exists() else []
|
||||
return [(float(line.split()[0]), line.split()[1:]) for line in lines]
|
||||
|
||||
def applied(self, run: str, task: str) -> dict:
|
||||
return json.loads((self.root / "applied" / f"{run}-{task}.json").read_text(encoding="utf-8"))
|
||||
|
||||
def events(self, verb: str) -> list[tuple[float, str]]:
|
||||
"""(time, name) of every ``apply``/``delete``/``delivered``/``rejected`` line, name = run task or ax name."""
|
||||
out = []
|
||||
for stamp, words in self.log():
|
||||
if verb == "apply" and words[:2] == ["apply", "-f"]:
|
||||
out.append((stamp, Path(words[2]).parent.name))
|
||||
elif verb == "delete" and words[:2] == ["delete", "task"]:
|
||||
out.append((stamp, words[2]))
|
||||
elif verb in ("delivered", "rejected") and words[0] == verb:
|
||||
out.append((stamp, words[1]))
|
||||
return out
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Fake:
|
||||
monkeypatch.setenv("OUSAST_RESULTS", str(tmp_path / "results"))
|
||||
monkeypatch.setenv("OUSAST_ARTIFACT_HOST", "127.0.0.1")
|
||||
monkeypatch.setenv("OUSAST_POLL_SECONDS", "0.02")
|
||||
fake = Fake(tmp_path)
|
||||
fake.script(SCRIPT)
|
||||
return fake
|
||||
|
||||
|
||||
def write_run(tmp_path: Path, name: str, tasks: str = CHAIN) -> Path:
|
||||
path = tmp_path / f"{name}.yaml"
|
||||
path.write_text(MANIFESTS.replace("{run}", name).replace("{tasks}", tasks.rstrip("\n")), encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def state_of(name: str) -> dict:
|
||||
return json.loads((reconciler.run_dir(name) / "state.json").read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def statuses(name: str) -> dict[str, str]:
|
||||
return {task: info["status"] for task, info in state_of(name)["tasks"].items()}
|
||||
|
||||
|
||||
# --- the DAG through ax -------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_chain_runs_in_order_with_inputs_env_and_model(fake: Fake, tmp_path: Path) -> None:
|
||||
result = reconciler.run(write_run(tmp_path, "chain"), ax=str(fake.ax))
|
||||
assert result == "done"
|
||||
assert statuses("chain") == {"repo-facts": "done", "verify": "done", "agree": "done"}
|
||||
applies = [name for _, name in fake.events("apply")]
|
||||
assert applies == ["repo-facts", "verify", "agree"]
|
||||
assert [n for _, n in fake.events("delete")] == ["chain-repo-facts", "chain-verify", "chain-agree"]
|
||||
verify = fake.applied("chain", "verify")
|
||||
env = verify["env"]
|
||||
assert env["OUSAST_PASS"] == "a" and env["OUSAST_RUN"] == "chain" and env["OUSAST_TASK"] == "verify"
|
||||
assert env["OUSAST_OUTPUT_DIR"] == "/workspace/.ousast-out/verify"
|
||||
assert env["OUSAST_BUDGET_USD"] == "2.5" and env["OUSAST_BUDGET_CALLS"] == "40"
|
||||
assert env["OUSAST_INPUT_FACTS"] == "/workspace/.ousast-in/verify/repo-facts/facts.json"
|
||||
assert env["OUSAST_ARTIFACT_URL"].startswith("http://127.0.0.1:")
|
||||
assert env["OUSAST_MODEL"] == "deepseek-flash"
|
||||
assert json.loads(env["OUSAST_MODEL_PARAMS"]) == {"cache_hit_per_m": 0.07, "input_per_m": 0.27, "output_per_m": 1.1}
|
||||
kinds = [(d["kind"], d["metadata"]["name"]) for d in verify["docs"]]
|
||||
assert kinds == [("Task", "chain-verify"), ("Model", "deepseek-flash"), ("Workspace", "case-pin"), ("Workspace", "chain-verify-inputs")]
|
||||
inputs_ws = verify["docs"][3]
|
||||
assert inputs_ws["spec"]["files"] == [{"path": "repo-facts/facts.json", "content": json.dumps({"a.py": {"functions": ["f"]}})}]
|
||||
assert verify["docs"][0]["spec"]["workspaces"][-1] == {"name": "chain-verify-inputs", "path": "/workspace/.ousast-in/verify"}
|
||||
assert "annotations" not in verify["docs"][0]["metadata"], "ax's ObjectMeta has no annotations"
|
||||
facts = fake.applied("chain", "repo-facts")
|
||||
assert "OUSAST_MODEL" not in facts["env"] and "OUSAST_BUDGET_USD" not in facts["env"]
|
||||
assert facts["docs"][0]["spec"]["image"] == "ousast-runner:dev"
|
||||
delivered = reconciler.run_dir("chain") / "verify" / "agreed.json"
|
||||
assert json.loads(delivered.read_text()) == {"agreed": ["a.py:f"]}
|
||||
assert state_of("chain")["tasks"]["verify"]["usd"] == 0.5 and state_of("chain")["tasks"]["verify"]["model"] == "deepseek-flash"
|
||||
assert not (reconciler.run_dir("chain") / "lock").exists()
|
||||
|
||||
|
||||
PAIR = """
|
||||
- name: x
|
||||
task: repo-facts
|
||||
{serialize}
|
||||
- name: y
|
||||
task: repo-facts
|
||||
{serialize}
|
||||
"""
|
||||
|
||||
|
||||
def overlap(fake: Fake) -> bool:
|
||||
applies = dict((name, stamp) for stamp, name in fake.events("apply"))
|
||||
deletes = dict((name.split("-", 1)[1], stamp) for stamp, name in fake.events("delete"))
|
||||
return max(applies.values()) < min(deletes.values())
|
||||
|
||||
|
||||
def test_serialize_label_never_overlaps_but_others_do(fake: Fake, tmp_path: Path) -> None:
|
||||
fake.script({"x": {"polls": 4, "files": {"summary.json": FACTS_SUMMARY}}, "y": {"polls": 4, "files": {"summary.json": FACTS_SUMMARY}}})
|
||||
assert reconciler.run(write_run(tmp_path, "par", PAIR.replace("{serialize}", "")), ax=str(fake.ax), workers=2) == "done"
|
||||
assert overlap(fake), "two unlabeled tasks with --workers 2 run at once"
|
||||
(fake.root / "ax.log").unlink()
|
||||
serialized = write_run(tmp_path, "ser", PAIR.replace("{serialize}", "serialize: engine"))
|
||||
assert reconciler.run(serialized, ax=str(fake.ax), workers=2) == "done"
|
||||
assert not overlap(fake), "tasks sharing a serialize label never overlap"
|
||||
|
||||
|
||||
def test_unfinished_stops_dependants_and_rerun_skips_done(fake: Fake, tmp_path: Path) -> None:
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
script["verify"]["files"]["summary.json"] = {**VERIFY_SUMMARY, "status": "unfinished", "units_done": 1, "reason": "budget"}
|
||||
fake.script(script)
|
||||
manifest = write_run(tmp_path, "resume")
|
||||
assert reconciler.run(manifest, ax=str(fake.ax)) == "unfinished"
|
||||
assert statuses("resume") == {"repo-facts": "done", "verify": "unfinished", "agree": "pending"}
|
||||
assert state_of("resume")["tasks"]["verify"]["reason"] == "budget"
|
||||
fake.script(SCRIPT)
|
||||
(fake.root / "ax.log").unlink()
|
||||
assert reconciler.run(manifest, ax=str(fake.ax)) == "done"
|
||||
assert [name for _, name in fake.events("apply")] == ["verify", "agree"], "done tasks are skipped on rerun (Req 2.2, 3.3)"
|
||||
assert statuses("resume") == {"repo-facts": "done", "verify": "done", "agree": "done"}
|
||||
|
||||
|
||||
def test_failed_task_halts_the_run(fake: Fake, tmp_path: Path) -> None:
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
failed = {"status": "failed", "units_done": 0, "units_total": 3, "reason": "402 insufficient balance"}
|
||||
script["repo-facts"]["files"]["summary.json"] = failed
|
||||
fake.script(script)
|
||||
assert reconciler.run(write_run(tmp_path, "halt"), ax=str(fake.ax)) == "failed"
|
||||
assert statuses("halt") == {"repo-facts": "failed", "verify": "pending", "agree": "pending"}
|
||||
assert state_of("halt")["tasks"]["repo-facts"]["reason"] == "402 insufficient balance"
|
||||
assert [name for _, name in fake.events("apply")] == ["repo-facts"]
|
||||
|
||||
|
||||
def test_ax_failure_phase_and_missing_delivery_are_failed(fake: Fake, tmp_path: Path) -> None:
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
script["repo-facts"]["phase"] = "Failed"
|
||||
fake.script(script)
|
||||
assert reconciler.run(write_run(tmp_path, "phase"), ax=str(fake.ax)) == "failed"
|
||||
assert state_of("phase")["tasks"]["repo-facts"]["reason"] == "ax phase Failed"
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
script["repo-facts"]["deliver"] = False
|
||||
fake.script(script)
|
||||
assert reconciler.run(write_run(tmp_path, "silent"), ax=str(fake.ax)) == "failed"
|
||||
assert state_of("silent")["tasks"]["repo-facts"]["reason"] == "no summary.json delivered"
|
||||
assert [n for _, n in fake.events("delete")] == ["phase-repo-facts", "silent-repo-facts"], "ax delete follows every terminal task"
|
||||
|
||||
|
||||
def test_summary_done_requires_every_unit(fake: Fake, tmp_path: Path) -> None:
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
script["repo-facts"]["files"]["summary.json"] = {**FACTS_SUMMARY, "units_done": 2}
|
||||
fake.script(script)
|
||||
assert reconciler.run(write_run(tmp_path, "short"), ax=str(fake.ax)) == "failed"
|
||||
assert statuses("short")["repo-facts"] == "failed"
|
||||
|
||||
|
||||
# --- receiver and lock ---------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def tar_of(files: dict[str, bytes]) -> bytes:
|
||||
buf = io.BytesIO()
|
||||
with tarfile.open(fileobj=buf, mode="w") as tar:
|
||||
for name, data in files.items():
|
||||
info = tarfile.TarInfo(name)
|
||||
info.size = len(data)
|
||||
tar.addfile(info, io.BytesIO(data))
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_receiver_rejects_traversal_unknown_tasks_and_wrong_types(tmp_path: Path) -> None:
|
||||
receiver = reconciler.Receiver("r", tmp_path / "r", 0)
|
||||
receiver.running.add("t")
|
||||
headers = {"X-Ousast-Run": "r", "X-Ousast-Task": "t", "Content-Type": "application/x-tar"}
|
||||
code, text = receiver.accept(headers, tar_of({"../escape.json": b"{}"}))
|
||||
assert code == 400 and "escape" in text and not (tmp_path / "escape.json").exists()
|
||||
assert receiver.accept({**headers, "X-Ousast-Task": "other"}, tar_of({"summary.json": b"{}"}))[0] == 403
|
||||
assert receiver.accept({**headers, "X-Ousast-Run": "someone"}, tar_of({"summary.json": b"{}"}))[0] == 403
|
||||
assert receiver.accept({**headers, "Content-Type": "application/json"}, b"{}")[0] == 415
|
||||
assert receiver.accept(headers, b"not a tar")[0] == 400
|
||||
assert receiver.accept(headers, tar_of({"summary.json": b"{}", "sub/x.json": b"1"}))[0] == 200
|
||||
assert (tmp_path / "r" / "t" / "sub" / "x.json").read_text() == "1" and "t" in receiver.delivered
|
||||
receiver.server_close()
|
||||
|
||||
|
||||
def test_delivery_for_a_task_that_is_not_running_is_rejected(fake: Fake, tmp_path: Path) -> None:
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
script["repo-facts"]["as_task"] = "agree" # the runner claims to be a task that has not started
|
||||
fake.script(script)
|
||||
assert reconciler.run(write_run(tmp_path, "spoof"), ax=str(fake.ax)) == "failed"
|
||||
assert [n for _, n in fake.events("rejected")] == ["repo-facts"]
|
||||
assert not (reconciler.run_dir("spoof") / "agree" / "summary.json").exists()
|
||||
|
||||
|
||||
def test_lock_refuses_a_second_runner_and_ignores_a_dead_one(fake: Fake, tmp_path: Path) -> None:
|
||||
manifest = write_run(tmp_path, "locked")
|
||||
base = reconciler.run_dir("locked")
|
||||
base.mkdir(parents=True)
|
||||
holder = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"])
|
||||
try:
|
||||
(base / "lock").write_text(str(holder.pid))
|
||||
with pytest.raises(RuntimeError, match=f"held by pid {holder.pid}"):
|
||||
reconciler.run(manifest, ax=str(fake.ax))
|
||||
assert not fake.events("apply")
|
||||
finally:
|
||||
holder.kill()
|
||||
holder.wait()
|
||||
(base / "lock").write_text(str(holder.pid)) # the process is gone: a stale lock
|
||||
assert reconciler.run(manifest, ax=str(fake.ax)) == "done"
|
||||
|
||||
|
||||
# --- status and attribution (Req 7.3) ------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def usage(prompt: int, cache_hit: int, output: int) -> dict[str, int]:
|
||||
return {"prompt_tokens": prompt, "prompt_cache_hit_tokens": cache_hit, "completion_tokens": output}
|
||||
|
||||
|
||||
def test_status_prints_attribution_from_usage_fields_only(fake: Fake, tmp_path: Path) -> None:
|
||||
script = json.loads(json.dumps(SCRIPT))
|
||||
script["agree"]["files"]["summary.json"] = {**AGREE_SUMMARY, "model": "deepseek-flash", "usd": 0.1, "calls": 1} # no usage
|
||||
script["verify"]["files"]["units.jsonl"] = "\n".join(
|
||||
json.dumps(u)
|
||||
for u in (
|
||||
{"path": "a.py", "turns": 3, "usd": 0.3, "usage": usage(700, 300, 150)},
|
||||
{"path": "b.py", "turns": 1, "usd": 0.2, "usage": usage(300, 0, 50)},
|
||||
)
|
||||
)
|
||||
fake.script(script)
|
||||
assert reconciler.run(write_run(tmp_path, "attr"), ax=str(fake.ax)) == "done"
|
||||
text = reconciler.status("attr", units=True)
|
||||
lines = {line.split()[0]: line.split() for line in text.splitlines() if line.strip() and not line.startswith("run ")}
|
||||
assert lines["verify"][1:7] == ["deepseek-flash", "4", "1000", "300", "200", "0.5000"]
|
||||
assert lines["agree"][1:7] == ["deepseek-flash", "1", "n/a", "n/a", "n/a", "0.1000"], "missing usage is n/a, never estimated"
|
||||
assert lines["repo-facts"][1:7] == ["-", "0", "n/a", "n/a", "n/a", "n/a"]
|
||||
assert lines["a.py"][1:7] == ["deepseek-flash", "3", "700", "300", "150", "0.3000"]
|
||||
assert lines["b.py"][1:7] == ["deepseek-flash", "1", "300", "0", "50", "0.2000"]
|
||||
assert lines["subtotal"][1:8] == ["deepseek-flash", "deepseek-flash", "5", "n/a", "n/a", "n/a", "0.6000"]
|
||||
assert lines["total"][1:7] == ["-", "5", "n/a", "n/a", "n/a", "0.6000"]
|
||||
report = json.loads((reconciler.run_dir("attr") / "attribution.json").read_text())
|
||||
assert report["total"] == {"task": "total", "model": None, "calls": 5, "prompt": None, "cache_hit": None, "output": None, "usd": 0.6}
|
||||
assert [u["task"] for u in report["units"]["verify"]] == ["a.py", "b.py"]
|
||||
assert [row["status"] for row in report["tasks"]] == ["done", "done", "done"]
|
||||
fake.script(SCRIPT)
|
||||
reconciler.run(write_run(tmp_path, "attr2"), ax=str(fake.ax))
|
||||
report = reconciler.attribution("attr2")
|
||||
assert report["total"]["prompt"] == 1000 and report["total"]["cache_hit"] == 300 and report["total"]["output"] == 200
|
||||
assert report["models"][0]["task"] == "subtotal deepseek-flash" and report["models"][0]["usd"] == 0.5
|
||||
|
||||
|
||||
def test_status_of_an_unknown_run_says_so(fake: Fake) -> None:
|
||||
assert "not started" in reconciler.status("nothing")
|
||||
|
||||
|
||||
# --- doctor, workspaces, size ---------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_doctor_names_the_bring_up_script_when_nothing_is_reachable(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setenv("PATH", str(tmp_path)) # no kubectl, no curl
|
||||
checks = reconciler.doctor()
|
||||
assert [name for name, _, _ in checks] == [
|
||||
"kind cluster",
|
||||
"agent substrate (ate-system)",
|
||||
"ax controller (ax-system)",
|
||||
"runner image in kind registry",
|
||||
]
|
||||
assert all(not ok and "ops/ax/up.sh" in text for _, ok, text in checks)
|
||||
|
||||
|
||||
def test_workspaces_one_manifest_per_case_pin(tmp_path: Path) -> None:
|
||||
population = tmp_path / "population.toml"
|
||||
population.write_text(
|
||||
'[[case]]\nid = "demo-sqli"\nrepo = "https://github.com/x/y"\nvulnerable = "aaa"\nfixed = "bbb"\n'
|
||||
'benign = { base = "ccc", tip = "ddd" }\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
written = workspaces(population, tmp_path / "out")
|
||||
labels = ("vulnerable", "fixed", "benign-base", "benign-tip")
|
||||
assert [p.name for p in written] == [f"demo-sqli-{label}.yaml" for label in labels]
|
||||
doc = yaml.safe_load(written[0].read_text(encoding="utf-8"))
|
||||
assert doc == {
|
||||
"apiVersion": "ax.io/v1alpha1",
|
||||
"kind": "Workspace",
|
||||
"metadata": {"name": "demo-sqli-vulnerable"},
|
||||
"spec": {"git": [{"name": "repo", "repo": "https://github.com/x/y", "commit": "aaa"}]},
|
||||
}
|
||||
from openultrasast.plane.manifests import load_manifests
|
||||
|
||||
assert set(load_manifests(written).workspaces) == {p.stem for p in written}
|
||||
|
||||
|
||||
def test_reconciler_is_under_500_lines_and_holds_no_pipeline_logic() -> None:
|
||||
text = RECONCILER.read_text(encoding="utf-8")
|
||||
assert len(text.splitlines()) < 500, "Req 3.4: the reconciler stays under 500 lines"
|
||||
for word in ("prompt(", "openai", "ChatClient", "score(", "system_prompt"):
|
||||
assert word not in text, f"Req 3.4: no pipeline logic in the reconciler ({word})"
|
||||
assert "Popen" not in text and "openultrasast.plane.tasks" not in text, "ax is the only executor"
|
||||
|
||||
|
||||
def test_cli_wires_plane_subcommands() -> None:
|
||||
from openultrasast.cli import main
|
||||
|
||||
with pytest.raises(SystemExit) as exc:
|
||||
main(["plane", "--help"])
|
||||
assert exc.value.code == 0
|
||||
with pytest.raises(SystemExit):
|
||||
main(["plane"])
|
||||
|
||||
|
||||
def test_elapsed_polling_is_paced(fake: Fake, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("OUSAST_POLL_SECONDS", "0.1")
|
||||
fake.script({"x": {"polls": 3, "files": {"summary.json": FACTS_SUMMARY}}, "y": {"polls": 1, "files": {"summary.json": FACTS_SUMMARY}}})
|
||||
started = time.monotonic()
|
||||
assert reconciler.run(write_run(tmp_path, "paced", PAIR.replace("{serialize}", "")), ax=str(fake.ax)) == "done"
|
||||
assert time.monotonic() - started >= 0.4, "three polls of x plus one of y at 0.1s each"
|
||||
Reference in New Issue
Block a user