vignette/scripts/probe-engine-gateway.py

#!/usr/bin/env python3
"""Measure a running engine gateway without embedding credentials.

The gateway process owns Claude authentication. This script only sends local HTTP
requests to an already-running gateway and is not used by normal tests.
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any


class ProbeError(RuntimeError):
    pass


def _endpoint(base_url: str, path: str) -> str:
    return f"{base_url.rstrip('/')}/{path.lstrip('/')}"


def _json_request(
    base_url: str,
    method: str,
    path: str,
    payload: dict[str, Any] | None = None,
    timeout: float = 30.0,
) -> dict[str, Any]:
    body = None if payload is None else json.dumps(payload).encode("utf-8")
    headers = {"Accept": "application/json"}
    if body is not None:
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(
        _endpoint(base_url, path),
        data=body,
        headers=headers,
        method=method,
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise ProbeError(f"{method} {path} failed with HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise ProbeError(f"{method} {path} transport failed: {exc}") from exc
    if not raw:
        return {}
    try:
        return json.loads(raw)
    except json.JSONDecodeError as exc:
        raise ProbeError(f"{method} {path} returned non-JSON: {raw[:200]}") from exc


def _stream_request(
    base_url: str,
    payload: dict[str, Any],
    timeout: float,
) -> dict[str, Any]:
    req = urllib.request.Request(
        _endpoint(base_url, "/v1/stream"),
        data=json.dumps(payload).encode("utf-8"),
        headers={"Accept": "text/event-stream", "Content-Type": "application/json"},
        method="POST",
    )
    started = time.perf_counter()
    first_token_at: float | None = None
    chunks: list[str] = []
    done_meta: dict[str, Any] = {}
    done_seen = False
    event_name: str | None = None
    data_lines: list[str] = []

    def flush_event(now: float) -> None:
        nonlocal first_token_at, done_meta, done_seen, event_name, data_lines
        if not event_name:
            data_lines = []
            return
        data = "\n".join(data_lines)
        if event_name == "token":
            if first_token_at is None:
                first_token_at = now
            try:
                token_payload = json.loads(data)
                chunks.append(str(token_payload.get("text", "")))
            except json.JSONDecodeError:
                chunks.append(data)
        elif event_name == "done":
            done_meta = json.loads(data or "{}")
            done_seen = True
        elif event_name == "error":
            try:
                err = json.loads(data)
                detail = err.get("detail", data)
            except json.JSONDecodeError:
                detail = data
            raise ProbeError(f"stream returned error event: {detail}")
        event_name = None
        data_lines = []

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            for raw_line in resp:
                now = time.perf_counter()
                line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n")
                if not line:
                    flush_event(now)
                    if done_seen:
                        break
                    continue
                if line.startswith("event:"):
                    event_name = line[6:].strip()
                elif line.startswith("data:"):
                    data_lines.append(line[5:].lstrip())
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise ProbeError(f"POST /v1/stream failed with HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise ProbeError(f"POST /v1/stream transport failed: {exc}") from exc

    finished = time.perf_counter()
    return {
        "ttft_ms": None if first_token_at is None else round((first_token_at - started) * 1000, 1),
        "latency_ms": round((finished - started) * 1000, 1),
        "text_chars": len("".join(chunks)),
        "done": done_meta,
    }


def _generate_payload(
    *,
    prompt: str,
    system_prompt: str,
    session_id: str | None,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "ai_role": "client",
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.0,
        "max_tokens": 128,
        "metadata": {"probe": "engine-gateway-probe"},
    }
    if session_id:
        payload["session_id"] = session_id
    return payload


def _summarize(values: list[float | None]) -> dict[str, Any]:
    clean = [v for v in values if v is not None]
    if not clean:
        return {"count": 0}
    return {
        "count": len(clean),
        "min_ms": round(min(clean), 1),
        "median_ms": round(statistics.median(clean), 1),
        "mean_ms": round(statistics.fmean(clean), 1),
        "max_ms": round(max(clean), 1),
    }


def _summaries(measurements: list[dict[str, Any]]) -> dict[str, Any]:
    result: dict[str, Any] = {}
    for kind in sorted({item["kind"] for item in measurements}):
        group = [item for item in measurements if item["kind"] == kind]
        result[kind] = {
            "ttft": _summarize([item.get("ttft_ms") for item in group]),
            "latency": _summarize([item.get("latency_ms") for item in group]),
        }
    return result


def _print_text(result: dict[str, Any]) -> None:
    print(f"Gateway: {result['base_url']}")
    health = result["health"]
    print(f"Health: ok={health.get('ok')} engine={health.get('engine')} sessions={health.get('sessions')}")
    session_id = result.get("session_id")
    if session_id:
        print(f"Session: {session_id[:12]}... closed={result.get('closed')}")
    for kind, summary in result["summary"].items():
        print(
            f"{kind}: "
            f"ttft median={summary['ttft'].get('median_ms')} ms "
            f"latency median={summary['latency'].get('median_ms')} ms "
            f"n={summary['latency'].get('count', 0)}"
        )


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Probe a running Vignette engine gateway for streaming TTFT, total latency, "
            "and session reuse. Claude credentials are never passed to this script."
        )
    )
    parser.add_argument(
        "--base-url",
        default=os.environ.get("ENGINE_URL", "http://127.0.0.1:9099"),
        help="Gateway base URL. Defaults to ENGINE_URL or http://127.0.0.1:9099.",
    )
    parser.add_argument("--prompt", default="Reply with exactly OK.", help="Probe user prompt.")
    parser.add_argument(
        "--system-prompt",
        default="You are an engine gateway probe. Reply concisely.",
        help="System prompt used when creating the reusable session.",
    )
    parser.add_argument("--budget-usd", type=float, default=0.5, help="Budget for the created /session.")
    parser.add_argument("--reuse-runs", type=int, default=2, help="Number of /v1/stream calls using one session_id.")
    parser.add_argument(
        "--ephemeral-runs",
        type=int,
        default=0,
        help="Optional cold /v1/stream calls without session_id. Defaults to 0 to limit spend.",
    )
    parser.add_argument("--timeout-sec", type=float, default=120.0, help="HTTP timeout per generation.")
    parser.add_argument("--json", action="store_true", help="Print machine-readable JSON.")
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    if args.reuse_runs < 0 or args.ephemeral_runs < 0:
        raise ProbeError("run counts must be non-negative")

    measurements: list[dict[str, Any]] = []
    session_id: str | None = None
    closed = False
    health = _json_request(args.base_url, "GET", "/health", timeout=10.0)
    if not health.get("ok"):
        raise ProbeError(f"gateway health is not ok: {health}")

    try:
        if args.reuse_runs:
            created = _json_request(
                args.base_url,
                "POST",
                "/session",
                {"system_prompt": args.system_prompt, "budget_usd": args.budget_usd},
                timeout=args.timeout_sec,
            )
            session_id = str(created["session_id"])
            for index in range(args.reuse_runs):
                measured = _stream_request(
                    args.base_url,
                    _generate_payload(
                        prompt=args.prompt,
                        system_prompt=args.system_prompt,
                        session_id=session_id,
                    ),
                    timeout=args.timeout_sec,
                )
                measured.update({"kind": "reused_stream", "run": index + 1})
                measurements.append(measured)

        for index in range(args.ephemeral_runs):
            measured = _stream_request(
                args.base_url,
                _generate_payload(
                    prompt=args.prompt,
                    system_prompt=args.system_prompt,
                    session_id=None,
                ),
                timeout=args.timeout_sec,
            )
            measured.update({"kind": "ephemeral_stream", "run": index + 1})
            measurements.append(measured)
    finally:
        if session_id:
            quoted = urllib.parse.quote(session_id, safe="")
            closed_payload = _json_request(args.base_url, "DELETE", f"/session/{quoted}", timeout=10.0)
            closed = bool(closed_payload.get("closed"))

    result = {
        "base_url": args.base_url,
        "health": health,
        "session_id": session_id,
        "closed": closed,
        "measurements": measurements,
        "summary": _summaries(measurements),
    }
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        _print_text(result)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main(sys.argv[1:]))
    except ProbeError as exc:
        print(f"error: {exc}", file=sys.stderr)
        raise SystemExit(2)