tools/local-harness/server.py

#!/usr/bin/env python3
"""Run a local browser UI for scanning AI agent skills with Vigil."""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import mimetypes
import os
import platform
import shutil
import subprocess
import tempfile
import time
import tarfile
import webbrowser
import zipfile
from dataclasses import dataclass
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path, PurePosixPath
from typing import Any
from urllib.parse import urlparse
from urllib.request import Request, urlopen


ROOT = Path(__file__).resolve().parent
REPOSITORY_URL = os.environ.get("VIGIL_REPOSITORY_URL", "https://huggingface.co/turenlabs/Vigil")
MODEL_URL = f"{REPOSITORY_URL}/resolve/main/compact-model.onnx"
METADATA_URL = MODEL_URL + ".json"
MODEL_SHA256 = "a56667baba56811b35dd7dffc75270f8c0d3f42de88f35dd198666a663e17f1e"
METADATA_SHA256 = "9c90f3bc1869f77452ed4f1cec1cdb17ffac5e5a20b229060432d6c9596f26db"
RUNTIME_VERSION = "v0.9.0-beta.3"
RUNTIMES = {
    ("darwin", "arm64"): (
        "vigil-compact-darwin-arm64.tar.gz",
        "b47164c9e7db7cdc199f5212e90c9865202f0c54fd05fe9225bd9dda4d170436",
    ),
    ("linux", "amd64"): (
        "vigil-compact-linux-amd64.tar.gz",
        "d4479903615788ebd1c1217318070cbb202639d5993099d9f9eece3d7a256af6",
    ),
    ("linux", "arm64"): (
        "vigil-compact-linux-arm64.tar.gz",
        "a344dac51a6a2691061494ae5ea24e5f8edab7a468d3ee9fad5efd14986b95d2",
    ),
    ("windows", "amd64"): (
        "vigil-compact-windows-amd64.zip",
        "2c5e65696b6704416b9fdbbd63d1079a6af5c3da49e7242541da17c5d412164f",
    ),
    ("windows", "arm64"): (
        "vigil-compact-windows-arm64.zip",
        "064d7243b71e735a1e9169a25bdda1e08039bb1080e343d19dec500c1e9ceff5",
    ),
}
MAX_REQUEST = 32 * 1024 * 1024
MAX_FILE = 4 * 1024 * 1024
MAX_FILES = 512


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def download(url: str, destination: Path, expected_hash: str) -> None:
    if destination.is_file() and sha256(destination) == expected_hash:
        return
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(destination.suffix + ".download")
    request = Request(url, headers={"User-Agent": "vigil-local-harness/1"})
    try:
        with urlopen(request, timeout=60) as response, temporary.open("wb") as output:
            shutil.copyfileobj(response, output)
        actual = sha256(temporary)
        if actual != expected_hash:
            raise RuntimeError(f"download hash mismatch: {actual}")
        temporary.replace(destination)
    finally:
        temporary.unlink(missing_ok=True)


def runtime_library_name(system: str | None = None) -> str:
    system = system or platform.system().lower()
    if system == "darwin":
        return "libonnxruntime.dylib"
    if system == "windows":
        return "onnxruntime.dll"
    return "libonnxruntime.so"


def device() -> tuple[str, str]:
    system = platform.system().lower()
    machine = platform.machine().lower()
    architectures = {
        "x86_64": "amd64",
        "amd64": "amd64",
        "aarch64": "arm64",
        "arm64": "arm64",
    }
    architecture = architectures.get(machine, machine)
    if (system, architecture) not in RUNTIMES:
        raise RuntimeError(
            f"This device is not supported yet: {platform.system()} {platform.machine()}"
        )
    return system, architecture


def safe_archive_path(value: str) -> Path:
    path = PurePosixPath(value.replace("\\", "/"))
    if path.is_absolute() or not path.parts or ".." in path.parts:
        raise RuntimeError(f"unsafe runtime archive path: {value}")
    return Path(*path.parts)


def extract_runtime(archive: Path, destination: Path) -> None:
    temporary = destination.with_name(destination.name + ".installing")
    if temporary.exists():
        shutil.rmtree(temporary)
    temporary.mkdir(parents=True)
    try:
        if archive.suffix == ".zip":
            with zipfile.ZipFile(archive) as bundle:
                for member in bundle.infolist():
                    relative = safe_archive_path(member.filename)
                    if member.is_dir():
                        (temporary / relative).mkdir(parents=True, exist_ok=True)
                        continue
                    target = temporary / relative
                    target.parent.mkdir(parents=True, exist_ok=True)
                    with bundle.open(member) as source, target.open("wb") as output:
                        shutil.copyfileobj(source, output)
        else:
            with tarfile.open(archive, "r:gz") as bundle:
                for member in bundle.getmembers():
                    relative = safe_archive_path(member.name)
                    if member.issym() or member.islnk() or member.isdev():
                        raise RuntimeError(f"unsupported runtime archive entry: {member.name}")
                    if member.isdir():
                        (temporary / relative).mkdir(parents=True, exist_ok=True)
                        continue
                    if not member.isfile():
                        continue
                    source = bundle.extractfile(member)
                    if source is None:
                        raise RuntimeError(f"could not extract runtime file: {member.name}")
                    target = temporary / relative
                    target.parent.mkdir(parents=True, exist_ok=True)
                    with source, target.open("wb") as output:
                        shutil.copyfileobj(source, output)
                    target.chmod(member.mode & 0o777)
        if destination.exists():
            shutil.rmtree(destination)
        temporary.replace(destination)
    finally:
        if temporary.exists():
            shutil.rmtree(temporary)


def install_runtime(runtime_dir: Path) -> tuple[Path, Path]:
    system, architecture = device()
    archive_name, archive_hash = RUNTIMES[(system, architecture)]
    package_name = archive_name.removesuffix(".tar.gz").removesuffix(".zip")
    install_root = runtime_dir / RUNTIME_VERSION / package_name
    binary_name = "vigil-compact.exe" if system == "windows" else "vigil-compact"
    binary = install_root / package_name / binary_name
    library = install_root / package_name / runtime_library_name(system)
    if binary.is_file() and library.is_file():
        return binary, library

    runtime_dir.mkdir(parents=True, exist_ok=True)
    archive = runtime_dir / RUNTIME_VERSION / archive_name
    url = f"{REPOSITORY_URL}/resolve/main/runtime/{RUNTIME_VERSION}/{archive_name}"
    print(f"Installing the Vigil runtime for {system} {architecture}...", flush=True)
    download(url, archive, archive_hash)
    extract_runtime(archive, install_root)
    if not binary.is_file() or not library.is_file():
        raise RuntimeError("the installed runtime package is incomplete")
    if system != "windows":
        binary.chmod(binary.stat().st_mode | 0o111)
    return binary, library


@dataclass(frozen=True)
class Config:
    binary: Path
    runtime_library: Path
    model: Path
    metadata: Path
    timeout: float

    def missing(self) -> list[str]:
        pairs = (
            ("vigil-compact binary", self.binary),
            ("ONNX Runtime library", self.runtime_library),
            ("ONNX model", self.model),
            ("model metadata", self.metadata),
        )
        return [label for label, path in pairs if not path.is_file()]


def resolve_config(args: argparse.Namespace) -> Config:
    model_dir = Path(args.model_dir).expanduser().resolve()
    model = model_dir / "compact-model.onnx"
    metadata = model_dir / "compact-model.onnx.json"
    download(MODEL_URL, model, MODEL_SHA256)
    download(METADATA_URL, metadata, METADATA_SHA256)

    binary_value = args.binary or os.environ.get("VIGIL_COMPACT_BIN")
    library_value = (
        args.runtime_lib
        or os.environ.get("VIGIL_COMPACT_RUNTIME_LIB")
        or os.environ.get("ONNXRUNTIME_LIB")
    )
    if binary_value and library_value:
        binary = Path(binary_value).expanduser().resolve()
        library = Path(library_value).expanduser().resolve()
    else:
        binary, library = install_runtime(Path(args.runtime_dir).expanduser().resolve())

    metadata_value = json.loads(metadata.read_text(encoding="utf-8"))
    if metadata_value["model"]["sha256"] != MODEL_SHA256:
        raise RuntimeError("metadata does not bind the downloaded model")
    return Config(binary, library, model, metadata, max(1.0, args.timeout))


def safe_relative_path(value: str) -> Path:
    path = PurePosixPath(value.replace("\\", "/"))
    if path.is_absolute() or not path.parts or ".." in path.parts:
        raise ValueError("invalid uploaded path")
    clean = [part for part in path.parts if part not in ("", ".")]
    if not clean:
        raise ValueError("invalid uploaded path")
    return Path(*clean)


def normalize_files(files: list[tuple[Path, bytes]]) -> list[tuple[Path, bytes]]:
    if len(files) == 1 and files[0][0].name.casefold() == "skill.md":
        return [(Path("SKILL.md"), files[0][1])]
    while not any(path.as_posix().casefold() == "skill.md" for path, _ in files):
        parts = [path.parts for path, _ in files]
        if any(len(value) < 2 for value in parts):
            break
        prefix = parts[0][0]
        if any(value[0] != prefix for value in parts):
            break
        files = [(Path(*path.parts[1:]), data) for path, data in files]
    if not any(path.as_posix().casefold() == "skill.md" for path, _ in files):
        raise ValueError("the package must contain a top-level SKILL.md")
    return files


def decode_files(body: bytes) -> list[tuple[Path, bytes]]:
    payload = json.loads(body.decode("utf-8"))
    items = payload.get("files") if isinstance(payload, dict) else None
    if not isinstance(items, list) or not items or len(items) > MAX_FILES:
        raise ValueError("choose between 1 and 512 files")
    decoded: list[tuple[Path, bytes]] = []
    seen: set[str] = set()
    for item in items:
        if not isinstance(item, dict) or not isinstance(item.get("path"), str):
            raise ValueError("invalid uploaded file")
        path = safe_relative_path(item["path"])
        key = path.as_posix().casefold()
        if key in seen:
            raise ValueError("duplicate uploaded path")
        seen.add(key)
        try:
            data = base64.b64decode(item.get("data", ""), validate=True)
        except Exception as exc:
            raise ValueError("invalid uploaded file data") from exc
        if len(data) > MAX_FILE:
            raise ValueError("an uploaded file exceeds 4 MiB")
        decoded.append((path, data))
    return normalize_files(decoded)


def scan(config: Config, files: list[tuple[Path, bytes]]) -> tuple[dict[str, Any], float]:
    if config.missing():
        raise RuntimeError("runtime_not_ready")
    started = time.perf_counter()
    with tempfile.TemporaryDirectory(prefix="vigil-harness-") as temporary:
        root = Path(temporary)
        for relative, data in files:
            destination = root / relative
            destination.parent.mkdir(parents=True, exist_ok=True)
            destination.write_bytes(data)
        command = [
            str(config.binary), "--require-model",
            "--model", str(config.model),
            "--metadata", str(config.metadata),
            "--runtime-lib", str(config.runtime_library),
            "--format", "json", str(root),
        ]
        completed = subprocess.run(
            command, stdin=subprocess.DEVNULL, capture_output=True, text=True,
            timeout=config.timeout, check=False,
        )
    output = completed.stdout if completed.returncode == 0 else completed.stderr
    lines = [line for line in output.splitlines() if line.strip()]
    if not lines:
        raise RuntimeError("scanner_returned_no_result")
    result = json.loads(lines[-1])
    if completed.returncode != 0 or result.get("schema_version") != "vigil.compact-score.v1":
        raise RuntimeError(str(result.get("error_code", "scan_failed")))
    return result, (time.perf_counter() - started) * 1000


class Handler(BaseHTTPRequestHandler):
    server_version = "VigilLocalHarness/1"

    @property
    def config(self) -> Config:
        return self.server.config  # type: ignore[attr-defined]

    def send_json(self, value: Any, status: HTTPStatus = HTTPStatus.OK) -> None:
        data = json.dumps(value, separators=(",", ":")).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self) -> None:  # noqa: N802
        route = urlparse(self.path).path
        if route == "/api/config":
            self.send_json({
                "ready": not self.config.missing(),
                "missing": self.config.missing(),
                "model_sha256": sha256(self.config.model),
            })
            return
        filename = "index.html" if route in ("/", "/index.html") else route.lstrip("/")
        if filename not in {"index.html", "app.js", "styles.css"}:
            self.send_json({"error": "not_found"}, HTTPStatus.NOT_FOUND)
            return
        data = (ROOT / filename).read_bytes()
        content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(data)

    def do_POST(self) -> None:  # noqa: N802
        if urlparse(self.path).path != "/api/scan":
            self.send_json({"error": "not_found"}, HTTPStatus.NOT_FOUND)
            return
        try:
            length = int(self.headers.get("Content-Length", "-1"))
            if length < 0 or length > MAX_REQUEST:
                raise ValueError("request exceeds 32 MiB")
            result, elapsed = scan(self.config, decode_files(self.rfile.read(length)))
            result["elapsed_ms"] = round(elapsed, 1)
            self.send_json({"ok": True, "result": result})
        except RuntimeError as exc:
            self.send_json({"error": str(exc)}, HTTPStatus.SERVICE_UNAVAILABLE)
        except (ValueError, OSError, json.JSONDecodeError) as exc:
            self.send_json({"error": str(exc)}, HTTPStatus.BAD_REQUEST)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8787)
    parser.add_argument("--bin", dest="binary", help="path to vigil-compact")
    parser.add_argument("--runtime-lib", help="path to the ONNX Runtime shared library")
    parser.add_argument("--model-dir", default=str(ROOT / ".model"))
    parser.add_argument("--runtime-dir", default=str(ROOT / ".runtime"))
    parser.add_argument("--timeout", type=float, default=30.0)
    parser.add_argument("--no-open", action="store_true", help="do not open a browser")
    args = parser.parse_args()
    config = resolve_config(args)
    server = ThreadingHTTPServer((args.host, args.port), Handler)
    server.config = config  # type: ignore[attr-defined]
    print(f"Vigil local harness: http://{args.host}:{args.port}", flush=True)
    if config.missing():
        print("Runtime setup needed: " + ", ".join(config.missing()), flush=True)
    if not args.no_open:
        webbrowser.open(f"http://{args.host}:{args.port}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())