feat: initial public release (MAESTRO)

2026-06-03 05:08:00 +00:00
commit f5c7666f6b
823 changed files with 184150 additions and 0 deletions
@@ -0,0 +1,530 @@
+/**
+ * Probe implementations for BackendStatusRegistry.
+ *
+ * Two upstream shapes are supported in Phase B:
+ *
+ * 1. **Direct workers** (llama-server compatible). We hit `<endpoint>/slots`
+ *    and derive busySlots / totalSlots / loadedModel from the slot array.
+ *    Recent llama-server builds let operators disable /slots via
+ *    `--no-slots`; in that case the endpoint returns 404/405/501 and we
+ *    fall back to `/health` to determine `online` only.
+ *
+ * 2. **Proxy workers** (LiteLLM Proxy). We hit `<endpoint>/health` and
+ *    convert each healthy_endpoints / unhealthy_endpoints entry into a
+ *    per-deployment NodeStatus. Slot / model info isn't currently
+ *    surfaced by LiteLLM's /health; v1 reports `online` + cache hits and
+ *    leaves the rest null. Phase C can switch to `/metrics` or
+ *    `/model/info` once they're wired through the proxy auth surface.
+ *
+ * Both probes:
+ *  - apply a 3s per-request timeout via AbortController (cluster-wide
+ *    hangs would otherwise wedge the entire registry tick),
+ *  - send the worker's apiKey as a Bearer header so team-scoped tokens
+ *    work,
+ *  - normalize the endpoint by stripping trailing slashes and stripping
+ *    a single trailing `/v1` segment (the worker config typically points
+ *    at the OpenAI-shaped `…/v1` base; llama-server's `/slots` lives at
+ *    the server root, not under /v1).
+ *
+ * Long-form rationale lives in
+ * docs/superpowers/specs/2026-05-18-multi-team-gpu-pool-and-node-status-design.md (Phase B).
+ */
+
+import type { WorkerDef } from '../config.js';
+import type { NodeStatus, ProbeContext } from './backend-status-registry.js';
+
+/**
+ * Per-probe HTTP timeout. Default was 3s historically but llama-server
+ * stalls its HTTP loop while a long-running generation is in flight on
+ * its event loop, so `/slots` regularly comes back >3s on a busy GPU
+ * and the registry would flip the node to offline mid-inference. 10s
+ * gives the server-side enough headroom to interleave the probe with
+ * an active generation without losing real liveness sensitivity (a
+ * truly dead node still trips ECONNREFUSED / DNS failure inside the
+ * window). 2026-05-21 dogfooding feedback.
+ */
+export const DEFAULT_PROBE_TIMEOUT_MS = 10_000;
+
+/** Trim trailing slashes and one optional `/v1` suffix so we can build sibling paths. */
+export function normalizeWorkerBase(endpoint: string): string {
+  const trimmed = endpoint.replace(/\/+$/, '');
+  return trimmed.endsWith('/v1') ? trimmed.slice(0, -3) : trimmed;
+}
+
+function authHeaders(apiKey: string | undefined): Record<string, string> {
+  const h: Record<string, string> = { Accept: 'application/json' };
+  if (apiKey) h['Authorization'] = `Bearer ${apiKey}`;
+  return h;
+}
+
+async function fetchWithTimeout(
+  url: string,
+  init: RequestInit,
+  timeoutMs: number,
+  fetchImpl: typeof fetch,
+  externalSignal?: AbortSignal,
+): Promise<Response> {
+  const ctrl = new AbortController();
+  const t = setTimeout(() => ctrl.abort(), timeoutMs);
+  // Chain the registry-wide cycle signal so shutdown can cancel pending
+  // probes immediately instead of waiting for timeoutMs.
+  let onExternalAbort: (() => void) | undefined;
+  if (externalSignal) {
+    if (externalSignal.aborted) {
+      ctrl.abort();
+    } else {
+      onExternalAbort = () => ctrl.abort();
+      externalSignal.addEventListener('abort', onExternalAbort, { once: true });
+    }
+  }
+  try {
+    return await fetchImpl(url, { ...init, signal: ctrl.signal });
+  } finally {
+    clearTimeout(t);
+    if (externalSignal && onExternalAbort) {
+      externalSignal.removeEventListener('abort', onExternalAbort);
+    }
+  }
+}
+
+export interface ProbeDeps {
+  fetchImpl?: typeof fetch;
+  timeoutMs?: number;
+  now?: () => string;
+}
+
+// ── llama-server (direct worker) ───────────────────────────────────────────
+
+/**
+ * Parse a Prometheus text-format /metrics payload from llama-server and
+ * extract a current decode throughput in tokens/sec.
+ *
+ * llama-server exposes several Prometheus metrics relevant to throughput:
+ *  - `llamacpp:n_decode_total` (counter)
+ *  - `llamacpp:n_prompt_tokens_total` (counter)
+ *  - `llamacpp:requests_processing` (gauge)
+ *  - `llamacpp:tokens_per_second` (gauge) — newer builds
+ *  - `llamacpp:prompt_tokens_seconds` (gauge) — older builds
+ *
+ * We prefer a directly-named tokens/sec gauge when present; otherwise we
+ * fall back to `prompt_tokens_seconds` (also a gauge in /metrics).
+ * Counter-based rate computation requires two scrapes and a per-scrape
+ * delta — that's a Phase C+ enhancement; Phase C v1 only surfaces gauges.
+ *
+ * Exported for unit testing. Returns null when no recognised gauge is
+ * present in the body.
+ */
+export function parseLlamaMetricsThroughput(text: string): number | null {
+  if (typeof text !== 'string' || text.length === 0) return null;
+  // Match `metric_name <number>` ignoring optional `{labels}` and
+  // trailing timestamp. Anchored per-line, comment lines (`# HELP`/
+  // `# TYPE`) skipped.
+  //
+  // Order across candidates matters: the first recognised metric family
+  // wins. tokens_per_second is the cleanest signal; prompt_tokens_seconds
+  // is a fallback. Within a family we use the `g` flag so multi-label
+  // exports (one llama-server process serving multiple slots/models with
+  // distinct `{model="..."}` labels) are all collected — first-match
+  // would silently drop sibling series and misreport the cluster's
+  // throughput. We surface the max across labels so the widget shows
+  // the fastest currently-decoding model instead of an arbitrary one.
+  const candidates = [
+    /^llamacpp:tokens_per_second(?:\{[^}]*\})?\s+([0-9eE.+-]+)/gm,
+    /^llamacpp:prompt_tokens_seconds(?:\{[^}]*\})?\s+([0-9eE.+-]+)/gm,
+    /^llamacpp:n_decode_tokens_per_second(?:\{[^}]*\})?\s+([0-9eE.+-]+)/gm,
+  ];
+  for (const re of candidates) {
+    const values: number[] = [];
+    for (const m of text.matchAll(re)) {
+      const v = Number(m[1]);
+      if (Number.isFinite(v) && v >= 0) values.push(v);
+    }
+    if (values.length > 0) return Math.max(...values);
+  }
+  return null;
+}
+
+interface LlamaSlot {
+  id?: number;
+  // `is_processing` is the canonical field on recent llama-server builds.
+  // Older builds called it `state` (0 = idle, 1 = processing).
+  is_processing?: boolean;
+  state?: number;
+  model?: string;
+  task_id?: number | string;
+}
+
+/**
+ * Parse a llama-server `/slots` JSON payload into busy / total / model.
+ *
+ * Tolerates both the modern object-array shape and the older shape where
+ * the response was already an envelope (`{ slots: [...] }`). Returns
+ * conservative zeros if the payload doesn't look like a slot array.
+ */
+export function parseLlamaSlots(payload: unknown): {
+  busySlots: number;
+  totalSlots: number;
+  loadedModel: string | null;
+} {
+  let slots: LlamaSlot[] = [];
+  if (Array.isArray(payload)) {
+    slots = payload as LlamaSlot[];
+  } else if (payload && typeof payload === 'object') {
+    const inner = (payload as { slots?: unknown }).slots;
+    if (Array.isArray(inner)) slots = inner as LlamaSlot[];
+  }
+  let busy = 0;
+  let model: string | null = null;
+  for (const s of slots) {
+    const isBusy = s.is_processing === true || (typeof s.state === 'number' && s.state !== 0);
+    if (isBusy) busy++;
+    if (!model && typeof s.model === 'string' && s.model.length > 0) model = s.model;
+  }
+  return { busySlots: busy, totalSlots: slots.length, loadedModel: model };
+}
+
+/**
+ * Probe a direct (llama-server) worker.
+ *
+ * Tries `/slots` first; if the server rejects it (--no-slots or older
+ * builds), falls back to `/health` so we can still report online/offline.
+ */
+export function buildDirectProbe(deps: ProbeDeps = {}): (worker: WorkerDef, ctx?: ProbeContext) => Promise<NodeStatus> {
+  const fetchImpl = deps.fetchImpl ?? fetch;
+  const timeoutMs = deps.timeoutMs ?? DEFAULT_PROBE_TIMEOUT_MS;
+  const now = deps.now ?? (() => new Date().toISOString());
+
+  return async function probeDirect(worker: WorkerDef, ctx?: ProbeContext): Promise<NodeStatus> {
+    const base = normalizeWorkerBase(worker.endpoint);
+    // Security: direct probes (llama-server compatible) must NOT receive the
+    // worker.apiKey. That key is a LiteLLM virtual key intended for the
+    // proxy auth surface — forwarding it to a direct llama-server would
+    // leak the secret to upstream access logs in plaintext. Only the proxy
+    // probe path (buildProxyProbe) sends Authorization.
+    const headers = authHeaders(undefined);
+    const slotsUrl = `${base}/slots`;
+
+    let slotsRes: Response;
+    try {
+      slotsRes = await fetchWithTimeout(slotsUrl, { method: 'GET', headers }, timeoutMs, fetchImpl, ctx?.signal);
+    } catch (err) {
+      // Network-level failure (DNS, refused, timeout). The server may
+      // still be alive but we can't talk to it; record online=false.
+      return errorStatus(worker, 'direct', err, now());
+    }
+
+    if (slotsRes.ok) {
+      let body: unknown;
+      try {
+        body = await slotsRes.json();
+      } catch (err) {
+        return errorStatus(worker, 'direct', err, now());
+      }
+      const parsed = parseLlamaSlots(body);
+      // /metrics is a best-effort enrichment: a failure here must NOT
+      // demote the status to offline. Some llama-server builds disable
+      // /metrics by default (`--metrics` opt-in) or run behind a proxy
+      // that strips Prometheus endpoints.
+      const throughputTps = await tryFetchThroughput(
+        base, headers, timeoutMs, fetchImpl, ctx?.signal,
+      );
+      return {
+        nodeId: worker.id,
+        workerId: worker.id,
+        source: 'direct',
+        online: true,
+        busy: parsed.busySlots > 0,
+        busySlots: parsed.busySlots,
+        totalSlots: parsed.totalSlots,
+        loadedModel: parsed.loadedModel ?? worker.model ?? null,
+        throughputTps,
+        lastSeen: now(),
+      };
+    }
+
+    // /slots disabled (--no-slots) returns 404 on most builds and 501 on
+    // some forks. /health is the documented liveness endpoint and is
+    // always available.
+    if (slotsRes.status === 404 || slotsRes.status === 405 || slotsRes.status === 501) {
+      return probeHealthFallback(base, headers, worker, timeoutMs, fetchImpl, now, ctx?.signal);
+    }
+
+    return errorStatus(
+      worker,
+      'direct',
+      new Error(`/slots returned HTTP ${slotsRes.status}`),
+      now(),
+    );
+  };
+}
+
+/**
+ * Best-effort throughput fetch. Returns null on any failure (404 from
+ * builds without --metrics, network error, parse failure) so the
+ * primary /slots-derived status is unaffected.
+ *
+ * Uses the same per-request timeout cap as /slots so a hung /metrics
+ * endpoint can't double the worst-case probe latency.
+ */
+async function tryFetchThroughput(
+  base: string,
+  headers: Record<string, string>,
+  timeoutMs: number,
+  fetchImpl: typeof fetch,
+  externalSignal?: AbortSignal,
+): Promise<number | null> {
+  try {
+    const res = await fetchWithTimeout(
+      `${base}/metrics`,
+      { method: 'GET', headers: { ...headers, Accept: 'text/plain' } },
+      timeoutMs,
+      fetchImpl,
+      externalSignal,
+    );
+    if (!res.ok) return null;
+    const text = await res.text();
+    return parseLlamaMetricsThroughput(text);
+  } catch {
+    return null;
+  }
+}
+
+async function probeHealthFallback(
+  base: string,
+  headers: Record<string, string>,
+  worker: WorkerDef,
+  timeoutMs: number,
+  fetchImpl: typeof fetch,
+  now: () => string,
+  externalSignal?: AbortSignal,
+): Promise<NodeStatus> {
+  try {
+    const res = await fetchWithTimeout(`${base}/health`, { method: 'GET', headers }, timeoutMs, fetchImpl, externalSignal);
+    if (!res.ok) {
+      return errorStatus(worker, 'direct', new Error(`/health returned HTTP ${res.status}`), now());
+    }
+    return {
+      nodeId: worker.id,
+      workerId: worker.id,
+      source: 'direct',
+      online: true,
+      busy: false,
+      busySlots: 0,
+      totalSlots: 0,
+      loadedModel: worker.model ?? null,
+      throughputTps: null,
+      lastSeen: now(),
+    };
+  } catch (err) {
+    return errorStatus(worker, 'direct', err, now());
+  }
+}
+
+// ── LiteLLM (proxy worker) ─────────────────────────────────────────────────
+
+/**
+ * Single LiteLLM `/health` entry. The exact field names depend on the
+ * LiteLLM version; we treat all fields as optional and only consume the
+ * ones we recognise. `model` is the deployment alias; some versions
+ * nest it under `litellm_params.model`.
+ */
+interface LiteLLMHealthEntry {
+  model?: unknown;
+  litellm_params?: { model?: unknown } | null;
+  cache?: unknown;
+  error?: unknown;
+  /**
+   * AAO Gateway extension fields (additive — vanilla LiteLLM omits them
+   * and we treat absence as "unknown busy state", same as before this
+   * extension existed). When present, parseLiteLLMHealth populates the
+   * NodeStatus with the gateway-aggregated busy view so every AAO
+   * client pointed at the same gateway sees consistent backend usage.
+   */
+  aao_busy_slots?: unknown;
+  aao_total_slots?: unknown;
+  aao_saturated?: unknown;
+  aao_last_seen?: unknown;
+}
+
+interface LiteLLMHealthBody {
+  healthy_endpoints?: LiteLLMHealthEntry[];
+  unhealthy_endpoints?: LiteLLMHealthEntry[];
+  // Some LiteLLM versions return `healthy_count` etc.; we don't need them
+  // since we walk the arrays directly.
+  [key: string]: unknown;
+}
+
+/**
+ * Coerce an arbitrary `aao_*_slots` value into a non-negative integer.
+ * Floors fractional inputs, drops anything non-finite or negative to
+ * zero so a malformed entry never corrupts the registry.
+ */
+function coerceNonNegInt(v: unknown): number {
+  if (typeof v !== 'number' || !Number.isFinite(v) || v < 0) return 0;
+  return Math.floor(v);
+}
+
+function extractDeploymentId(entry: LiteLLMHealthEntry): string | null {
+  if (typeof entry.model === 'string' && entry.model.length > 0) return entry.model.trim();
+  const nested = entry.litellm_params?.model;
+  if (typeof nested === 'string' && nested.length > 0) return nested.trim();
+  return null;
+}
+
+/**
+ * Coerce a LiteLLM `error` field to a human-readable string.
+ *
+ * LiteLLM pre-1.40 surfaced errors as plain strings, but recent
+ * versions (post-1.40) wrap them in `{ message: string, type: string }`
+ * objects. The previous string-only guard silently dropped the object
+ * form, leaving the widget with "unhealthy but reason unknown" — a
+ * red icon with nothing to action.
+ *
+ * Strategy:
+ *  - string → return as-is
+ *  - object with `.message` → return that
+ *  - other object → JSON.stringify so something useful surfaces in
+ *    logs / UI (better than the cryptic "[object Object]")
+ *  - null/undefined → undefined (caller skips the field)
+ */
+function extractErrorMessage(err: unknown): string | undefined {
+  if (typeof err === 'string') return err;
+  if (err && typeof err === 'object') {
+    const obj = err as { message?: unknown };
+    if (typeof obj.message === 'string') return obj.message;
+    try {
+      return JSON.stringify(err);
+    } catch {
+      return String(err);
+    }
+  }
+  return undefined;
+}
+
+/**
+ * Convert a LiteLLM /health response into one NodeStatus per deployment.
+ * Exported for unit-testing without a live proxy.
+ *
+ * Dedup policy (Phase C): **unhealthy wins precedence**. When the same
+ * deployment id appears in both `healthy_endpoints` and
+ * `unhealthy_endpoints` (a flapping backend), we surface the unhealthy
+ * row so the widget doesn't mask a degraded state behind a green icon.
+ * The error message from the unhealthy entry is preserved in
+ * `lastProbeError`. Phase B used "first appearance wins" which silently
+ * hid flap conditions from operators.
+ */
+export function parseLiteLLMHealth(
+  body: unknown,
+  workerId: string,
+  ts: string,
+): NodeStatus[] {
+  const obj = (body && typeof body === 'object' ? body : {}) as LiteLLMHealthBody;
+  const out: NodeStatus[] = [];
+  const seen = new Set<string>();
+  const push = (entry: LiteLLMHealthEntry, online: boolean): void => {
+    const id = extractDeploymentId(entry);
+    if (!id || seen.has(id)) return;
+    seen.add(id);
+    // AAO Gateway extension: when the gateway annotates each entry
+    // with aao_busy_slots / aao_total_slots, we inherit its
+    // multi-client-aggregated view. Without these fields (vanilla
+    // LiteLLM or older gateway) busy stays 0 — same as before the
+    // extension existed.
+    const busySlots = coerceNonNegInt(entry.aao_busy_slots);
+    const totalSlots = coerceNonNegInt(entry.aao_total_slots);
+    out.push({
+      nodeId: id,
+      workerId,
+      source: 'proxy',
+      online,
+      busy: busySlots > 0,
+      busySlots,
+      totalSlots,
+      loadedModel: id,
+      throughputTps: null,
+      lastSeen: ts,
+      lastProbeError: online ? undefined : extractErrorMessage(entry.error),
+    });
+  };
+  // Iterate unhealthy first so a flapping deployment registers as
+  // unhealthy and the healthy entry's subsequent push is skipped by the
+  // `seen` guard.
+  for (const entry of obj.unhealthy_endpoints ?? []) push(entry, false);
+  for (const entry of obj.healthy_endpoints ?? []) push(entry, true);
+  return out;
+}
+
+/**
+ * Probe a LiteLLM proxy worker.
+ *
+ * Returns one NodeStatus per deployment the proxy reports. If the proxy
+ * itself is unreachable we synthesize a single offline NodeStatus keyed
+ * to the workerId so the widget can render a "proxy down" row.
+ */
+export function buildProxyProbe(deps: ProbeDeps = {}): (worker: WorkerDef, ctx?: ProbeContext) => Promise<NodeStatus[]> {
+  const fetchImpl = deps.fetchImpl ?? fetch;
+  const timeoutMs = deps.timeoutMs ?? DEFAULT_PROBE_TIMEOUT_MS;
+  const now = deps.now ?? (() => new Date().toISOString());
+
+  return async function probeProxy(worker: WorkerDef, ctx?: ProbeContext): Promise<NodeStatus[]> {
+    const base = normalizeWorkerBase(worker.endpoint);
+    const headers = authHeaders(worker.apiKey);
+    const url = `${base}/health`;
+    let res: Response;
+    try {
+      res = await fetchWithTimeout(url, { method: 'GET', headers }, timeoutMs, fetchImpl, ctx?.signal);
+    } catch (err) {
+      return [errorStatus(worker, 'proxy', err, now())];
+    }
+    if (!res.ok) {
+      return [errorStatus(worker, 'proxy', new Error(`/health returned HTTP ${res.status}`), now())];
+    }
+    let body: unknown;
+    try {
+      body = await res.json();
+    } catch (err) {
+      return [errorStatus(worker, 'proxy', err, now())];
+    }
+    const ts = now();
+    const parsed = parseLiteLLMHealth(body, worker.id, ts);
+    if (parsed.length === 0) {
+      // Proxy is alive but reports zero deployments: surface a single
+      // synthetic row so the widget shows the proxy itself rather than
+      // a blank panel.
+      return [{
+        nodeId: worker.id,
+        workerId: worker.id,
+        source: 'proxy',
+        online: true,
+        busy: false,
+        busySlots: 0,
+        totalSlots: 0,
+        loadedModel: null,
+        throughputTps: null,
+        lastSeen: ts,
+      }];
+    }
+    return parsed;
+  };
+}
+
+function errorStatus(
+  worker: WorkerDef,
+  source: 'direct' | 'proxy',
+  err: unknown,
+  ts: string,
+): NodeStatus {
+  return {
+    nodeId: worker.id,
+    workerId: worker.id,
+    source,
+    online: false,
+    busy: false,
+    busySlots: 0,
+    totalSlots: 0,
+    loadedModel: null,
+    throughputTps: null,
+    lastSeen: ts,
+    lastProbeError: err instanceof Error ? err.message : String(err),
+  };
+}