@runuai/host 0.9.170 → 0.9.172

package/lib/openrouter-budget.ts

@@ -0,0 +1,247 @@
+/** ADR-164: OpenRouter key budgets, never account credit balance or CLI usage.
+ * A credential is read only while a task's current lease serves it. Several
+ * tasks sharing one credential/revision share one ten-minute poll. */
+import type { CustomConnectionReason, OpenRouterBudget, OpenRouterBudgetFrame, OpenRouterBudgetResult } from "../src/custom-connections";
+import { SafeHttpRefusal, readSafeBody, safeRequest } from "./safe-http";
+
+export const OPENROUTER_BUDGET_URL = "https://openrouter.ai/api/v1/key";
+export const OPENROUTER_BUDGET_POLL_MS = 10 * 60_000;
+export const OPENROUTER_BUDGET_TIMEOUT_MS = 20_000;
+export const OPENROUTER_BUDGET_MAX_BYTES = 64 * 1024;
+
+function object(value: unknown): Record<string, unknown> | null {
+  return value !== null && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : null;
+}
+function amount(value: unknown): value is number | null {
+  return value === null || (typeof value === "number" && Number.isFinite(value) && value >= 0);
+}
+/** Provider timestamps are ISO instants. An interval such as "monthly" is
+ * deliberately not converted into an invented reset date. */
+function timestamp(value: unknown): number | null | undefined {
+  if (value === null || value === undefined) return null;
+  if (typeof value !== "string" || !/^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d(?:\.\d+)?(?:Z|[+-]\d\d:\d\d)$/.test(value)) return undefined;
+  const ms = Date.parse(value);
+  return Number.isSafeInteger(ms) && ms > 0 ? ms : undefined;
+}
+
+/** Select only the bounded public values from an arbitrary /key response.
+ * Required nulls remain meaningful: a missing limit is not "unlimited". */
+export function parseOpenRouterKeyBudget(value: unknown): OpenRouterBudget | null {
+  const data = object(object(value)?.data);
+  if (!data || !Object.hasOwn(data, "usage") || !Object.hasOwn(data, "limit") || !Object.hasOwn(data, "limit_remaining") || !amount(data.usage) || !amount(data.limit) || !amount(data.limit_remaining)) return null;
+  const keyExpiresAt = timestamp(data.expires_at);
+  if (keyExpiresAt === undefined) return null;
+  let resetsAt: number | null = null;
+  // A future API may return a timestamp instead of the documented interval.
+  if (typeof data.limit_reset === "string" && data.limit_reset.includes("T")) {
+    const reset = timestamp(data.limit_reset);
+    if (reset === undefined) return null;
+    resetsAt = reset;
+  }
+  let effectiveBudget: number | null = null;
+  if (data.effective_budget !== undefined && data.effective_budget !== null) {
+    const budget = object(data.effective_budget);
+    if (!budget || !Object.hasOwn(budget, "remaining_usd") || !amount(budget.remaining_usd)) return null;
+    effectiveBudget = budget.remaining_usd;
+    const reset = timestamp(budget.resets_at);
+    if (reset === undefined) return null;
+    // This reset belongs to the effective, stricter budget when it binds.
+    if (effectiveBudget !== null && (data.limit_remaining === null || effectiveBudget < data.limit_remaining)) resetsAt = reset;
+  }
+  return { usage: data.usage, limit: data.limit, remaining: data.limit_remaining, effectiveBudget, resetsAt, keyExpiresAt };
+}
+
+function responseReason(status: number): CustomConnectionReason {
+  if (status === 401) return "unauthorized";
+  if (status === 402 || status === 403) return "forbidden";
+  if (status === 429) return "rate-limited";
+  return status >= 500 ? "upstream-error" : "wrong-protocol";
+}
+function failureReason(error: unknown): CustomConnectionReason {
+  if (error instanceof SafeHttpRefusal) {
+    if (error.reason === "non_public_address") return "private-address";
+    if (error.reason === "timeout") return "timeout";
+    if (error.reason === "response_too_large" || error.reason === "url") return "wrong-protocol";
+  }
+  return "unreachable";
+}
+export interface ReadOpenRouterBudgetOptions {
+  signal?: AbortSignal;
+  /** Tests inject transport; production always uses the DNS-pinned client. */
+  request?: typeof safeRequest;
+  now?: () => number;
+}
+
+/** Reads neither a caller-provided URL nor any redirect. No provider text,
+ * key, arbitrary object or exception escapes this boundary. */
+export async function readOpenRouterBudget(apiKey: string, options: ReadOpenRouterBudgetOptions = {}): Promise<OpenRouterBudgetResult> {
+  if (!apiKey || apiKey.length > 8192 || /[\s\u0000-\u001f\u007f]/.test(apiKey)) return { status: "error", reason: "unauthorized" };
+  if (options.signal?.aborted) return { status: "error", reason: "timeout" };
+  const controller = new AbortController();
+  const onAbort = (): void => controller.abort();
+  if (options.signal?.aborted) controller.abort();
+  else options.signal?.addEventListener("abort", onAbort, { once: true });
+  const timeout = setTimeout(() => controller.abort(), OPENROUTER_BUDGET_TIMEOUT_MS);
+  timeout.unref();
+  try {
+    const response = await (options.request ?? safeRequest)({
+      url: OPENROUTER_BUDGET_URL,
+      method: "GET",
+      headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" },
+      signal: controller.signal,
+    }, { connectMs: 10_000, headersMs: OPENROUTER_BUDGET_TIMEOUT_MS, idleMs: OPENROUTER_BUDGET_TIMEOUT_MS, maxRequestBodyBytes: 0 });
+    if (response.status !== 200) {
+      response.body.destroy();
+      return { status: "error", reason: responseReason(response.status) };
+    }
+    const raw = await readSafeBody(response, OPENROUTER_BUDGET_MAX_BYTES);
+    const budget = parseOpenRouterKeyBudget(JSON.parse(raw.toString("utf8")) as unknown);
+    return budget ? { status: "ok", ...budget } : { status: "error", reason: "wrong-protocol" };
+  } catch (error) {
+    return { status: "error", reason: controller.signal.aborted ? "timeout" : error instanceof SyntaxError ? "wrong-protocol" : failureReason(error) };
+  } finally {
+    clearTimeout(timeout);
+    options.signal?.removeEventListener("abort", onAbort);
+  }
+}
+
+export interface OpenRouterBudgetAccount {
+  credentialId: string;
+  revision: string;
+  apiKey: string;
+}
+interface Watcher {
+  identity: OpenRouterBudgetAccount;
+  tasks: Set<string>;
+  controller: AbortController | null;
+  timer: ReturnType<typeof setTimeout> | null;
+  last: { observedAt: number; result: OpenRouterBudgetResult } | null;
+}
+export interface OpenRouterBudgetPollerOptions {
+  emit: (frame: OpenRouterBudgetFrame) => void;
+  read?: (apiKey: string, options: ReadOpenRouterBudgetOptions) => Promise<OpenRouterBudgetResult>;
+  now?: () => number;
+}
+
+function identityKey(account: Pick<OpenRouterBudgetAccount, "credentialId" | "revision">): string {
+  return `${account.credentialId}\u0000${account.revision}`;
+}
+
+export class OpenRouterBudgetPoller {
+  private readonly watchers = new Map<string, Watcher>();
+  private readonly tasks = new Map<string, Set<string>>();
+  private readonly lastPoll = new Map<string, number>();
+  private readonly pending = new Set<Promise<void>>();
+  private readonly now: () => number;
+  private readonly read: NonNullable<OpenRouterBudgetPollerOptions["read"]>;
+
+  constructor(private readonly options: OpenRouterBudgetPollerOptions) {
+    this.now = options.now ?? Date.now;
+    this.read = options.read ?? readOpenRouterBudget;
+  }
+
+  /** The caller passes only OpenRouter accounts from the task's OPEN current
+   * lease. Passing [] ends that task's references and aborts unused reads. */
+  sync(taskId: string, accounts: readonly OpenRouterBudgetAccount[]): void {
+    this.pruneCooldowns();
+    const next = new Map(accounts.map(account => [identityKey(account), account]));
+    const previous = this.tasks.get(taskId) ?? new Set<string>();
+    for (const key of previous) if (!next.has(key)) this.removeTask(key, taskId);
+    const keys = new Set<string>();
+    for (const [key, identity] of next) {
+      keys.add(key);
+      let watcher = this.watchers.get(key);
+      if (!watcher) {
+        watcher = { identity: { ...identity }, tasks: new Set(), controller: null, timer: null, last: null };
+        this.watchers.set(key, watcher);
+      }
+      const added = !watcher.tasks.has(taskId);
+      watcher.tasks.add(taskId);
+      if (added && watcher.last) this.report(taskId, watcher, watcher.last);
+      if (!watcher.timer && !watcher.controller) this.schedule(key, watcher);
+    }
+    if (keys.size) this.tasks.set(taskId, keys);
+    else this.tasks.delete(taskId);
+  }
+
+  /** Removal releases the key, aborts its request when unused, and prevents a
+   * late response from publishing under a subsequent run or revision. */
+  stop(taskId?: string): void {
+    if (taskId !== undefined) {
+      for (const key of this.tasks.get(taskId) ?? []) this.removeTask(key, taskId);
+      this.tasks.delete(taskId);
+      return;
+    }
+    for (const watcher of this.watchers.values()) {
+      if (watcher.timer) clearTimeout(watcher.timer);
+      watcher.controller?.abort();
+    }
+    this.watchers.clear();
+    this.tasks.clear();
+  }
+
+  async drain(): Promise<void> {
+    await Promise.all([...this.pending]);
+  }
+
+  private removeTask(key: string, taskId: string): void {
+    const watcher = this.watchers.get(key);
+    if (!watcher) return;
+    watcher.tasks.delete(taskId);
+    if (watcher.tasks.size) return;
+    if (watcher.timer) clearTimeout(watcher.timer);
+    watcher.controller?.abort();
+    this.watchers.delete(key);
+  }
+
+  private schedule(key: string, watcher: Watcher): void {
+    const last = this.lastPoll.get(key);
+    const delay = last === undefined ? 0 : Math.max(0, last + OPENROUTER_BUDGET_POLL_MS - this.now());
+    if (delay === 0) {
+      this.start(key, watcher);
+      return;
+    }
+    watcher.timer = setTimeout(() => {
+      watcher.timer = null;
+      if (this.watchers.get(key) === watcher) this.start(key, watcher);
+    }, delay);
+    watcher.timer.unref();
+  }
+
+  private start(key: string, watcher: Watcher): void {
+    if (this.watchers.get(key) !== watcher || watcher.controller) return;
+    const controller = new AbortController();
+    watcher.controller = controller;
+    this.lastPoll.set(key, this.now());
+    const operation = (async (): Promise<void> => {
+      let result: OpenRouterBudgetResult;
+      try { result = await this.read(watcher.identity.apiKey, { signal: controller.signal, now: this.now }); }
+      catch { result = { status: "error", reason: "unreachable" }; }
+      if (this.watchers.get(key) !== watcher || controller.signal.aborted) return;
+      const observation = { observedAt: this.now(), result };
+      watcher.last = observation;
+      for (const taskId of watcher.tasks) this.report(taskId, watcher, observation);
+    })();
+    this.pending.add(operation);
+    void operation.finally(() => {
+      this.pending.delete(operation);
+      if (this.watchers.get(key) === watcher) {
+        watcher.controller = null;
+        this.schedule(key, watcher);
+      }
+    });
+  }
+
+  private report(taskId: string, watcher: Watcher, observation: NonNullable<Watcher["last"]>): void {
+    try {
+      this.options.emit({ kind: "openrouter.budget", taskId, credentialId: watcher.identity.credentialId, revision: watcher.identity.revision, observedAt: observation.observedAt, ...observation.result });
+    } catch {
+      // A disconnected transport is not an upstream response and never logs a key.
+    }
+  }
+
+  private pruneCooldowns(): void {
+    const before = this.now() - OPENROUTER_BUDGET_POLL_MS;
+    for (const [key, at] of this.lastPoll) if (at <= before && !this.watchers.has(key)) this.lastPoll.delete(key);
+  }
+}