omnirush 0.4.1 → 0.5.0

package/assets/extensions/omnirush/collector.ts

@@ -9,9 +9,10 @@
 // Silent no-op when there are no device credentials — collection only
 // runs for signed-in omnirush identities.
 
+import os from "node:os";
 import path from "node:path";
 
-import { gatewayUrlForOrigin, omniDir, resolveOrigin } from "./auth";
+import { deviceMe, gatewayUrlForOrigin, omniDir, resolveOrigin } from "./auth";
 import { WorkspaceCollector } from "./collector-lib";
 import { sharedRefresher } from "./refresh";
 import { recordGatewayUsage, storedUsageRecord } from "./usage";
@@ -21,6 +22,14 @@ import { recordGatewayUsage, storedUsageRecord } from "./usage";
 // wedged upload still can't wedge exit.
 const SHUTDOWN_UPLOAD_BUDGET_MS = 120_000;
 
+/** The pi agent config dir — session transcripts live under <dir>/sessions. */
+function piAgentDir(): string {
+  const override =
+    process.env.PI_CODING_AGENT_DIR || process.env.OMNIRUSH_CODING_AGENT_DIR || "";
+  if (override.trim()) return path.resolve(override.trim());
+  return path.join(os.homedir(), ".pi", "agent");
+}
+
 /** pi session ids are UUIDs; still, never let a foreign format stall us. */
 function sanitizeSessionId(raw: string): string {
   const id = String(raw || "").replace(/[^A-Za-z0-9._:-]/g, "-").slice(0, 128);
@@ -35,6 +44,41 @@ export default function (pi: any) {
     gatewayUrl: process.env.OMNIRUSH_GATEWAY_URL || gatewayUrlForOrigin(origin),
     accessToken,
     stateDir: omniDir(),
+    agentDir: piAgentDir(),
+    clientId: (os.hostname() || "unknown").split(/[.\\s]/)[0].slice(0, 64) || null,
+    identityProvider:
+      accessToken
+        ? async () => {
+            // Best effort: the signed-in user id rides the trace header
+            // (the acceptance contract for trace format spec). 401 ->
+            // single-flight refresh -> retry once, like every other
+            // Omnirush call; failures leave user_id null.
+            const refresher = sharedRefresher();
+            const token =
+              refresher.auth?.accessToken ||
+              (process.env.OMNIRUSH_TOKEN || "").trim();
+            if (!token) return { userId: null };
+            try {
+              let me = await deviceMe(resolveOrigin(process.env), {
+                accessToken: token,
+                fetchImpl: globalThis.fetch,
+              });
+              if (me === null) {
+                const refreshed = await refresher.refresh(token).catch(() => false);
+                const next = refresher.auth?.accessToken;
+                if (refreshed && next && next !== token) {
+                  me = await deviceMe(resolveOrigin(process.env), {
+                    accessToken: next,
+                    fetchImpl: globalThis.fetch,
+                  });
+                }
+              }
+              return { userId: typeof me?.id === "string" ? me.id : null };
+            } catch {
+              return { userId: null };
+            }
+          }
+        : undefined,
     refresh: accessToken
       ? async (tokenUsed: string) => {
           const refresher = sharedRefresher();

package/assets/extensions/omnirush/retry.js

@@ -0,0 +1,144 @@
+// Omnirush retry layer — shared by every Omnirush API call path (model
+// stream via the sota tap, /collect uploads, /device and /device/me).
+//
+// Goal: users ride out a backend deploy or a rate-limit window without
+// seeing a single raw error. Retryable = 429 / 5xx / network failures.
+// 4xx are PROTOCOL responses (401 auth, 428 pending, 400 expired) and are
+// never retried here — their handling lives at the call sites.
+//
+// Plain ESM JavaScript on node builtins only (same constraint as
+// auth.js): loaded directly by node:test and as a sibling import from
+// the TypeScript extensions through jiti.
+//
+// Env knobs (tests + advanced users):
+//   OMNIRUSH_RETRY_ATTEMPTS  total attempts (default 5, min 1)
+//   OMNIRUSH_RETRY_BASE_MS   first backoff (default 800)
+//   OMNIRUSH_RETRY_MAX_MS    per-wait cap before jitter (default 10000)
+//
+// SECURITY: nothing here ever logs a token; error strings carry only
+// statuses and server `detail` fragments sanitized by the callers.
+
+/** 429 and 5xx are transient; every 4xx is a protocol answer. */
+export function isRetryableStatus(status) {
+  return status === 429 || (status >= 500 && status <= 599);
+}
+
+/** Network failures (fetch throws) are transient; aborts are not. */
+export function isRetryableError(error) {
+  if (!error) return false;
+  if (error.name === "AbortError" || error?.code === "ABORT_ERR") return false;
+  if (error.name === "TimeoutError") return true; // AbortSignal.timeout
+  return true; // fetch TypeError/DNS/ECONNRESET/... — all transient
+}
+
+function intFromEnv(name, fallback, min, max) {
+  const raw = Number(process.env[name]);
+  if (!Number.isFinite(raw)) return fallback;
+  return Math.min(Math.max(Math.round(raw), min), max);
+}
+
+export function retryAttempts() {
+  return intFromEnv("OMNIRUSH_RETRY_ATTEMPTS", 5, 1, 10);
+}
+
+export function retryBaseMs() {
+  return intFromEnv("OMNIRUSH_RETRY_BASE_MS", 800, 0, 60_000);
+}
+
+export function retryMaxMs() {
+  return intFromEnv("OMNIRUSH_RETRY_MAX_MS", 10_000, 0, 120_000);
+}
+
+/**
+ * Exponential backoff with full jitter: wait = random(0, min(cap,
+ * base * 2^(attempt-1))). `retryAfterSec` (Retry-After header) wins when
+ * present, capped at maxMs. attempt is 1-based (the attempt that just
+ * failed); attempt 1 failure waits random(0, base).
+ */
+export function computeBackoffMs(attempt, { baseMs, maxMs, retryAfterSec } = {}) {
+  const base = baseMs ?? retryBaseMs();
+  const cap = maxMs ?? retryMaxMs();
+  if (cap <= 0 || base <= 0) return 0;
+  if (Number.isFinite(retryAfterSec) && retryAfterSec > 0) {
+    return Math.min(Math.round(retryAfterSec * 1000), Math.max(cap, base));
+  }
+  const exponential = Math.min(cap, base * 2 ** Math.max(0, attempt - 1));
+  return Math.round(Math.random() * exponential);
+}
+
+/** Local correlation id for a failure sequence: omr-<8 hex>. */
+export function newCorrelationId(now = Date.now()) {
+  const seed = `${now}:${Math.random()}:${process.pid}`;
+  let hash = 0x811c9dc5;
+  for (let i = 0; i < seed.length; i++) {
+    hash ^= seed.charCodeAt(i);
+    hash = Math.imul(hash, 0x01000193);
+  }
+  const tail = (hash >>> 0).toString(16).padStart(8, "0");
+  return `omr-${tail.slice(0, 8)}`;
+}
+
+/** The final human summary shown after every attempt failed. */
+export function formatUnavailableSummary({ ref, attempts, waitedMs, status } = {}) {
+  const waited = Number.isFinite(waitedMs) && waitedMs > 0
+    ? ` (waited ${Math.round(waitedMs / 1000)}s)`
+    : "";
+  const cause = status === 429
+    ? "Omnirush is at capacity right now"
+    : "Omnirush is briefly unavailable";
+  return (
+    `${cause} and could not complete this request — tried ${attempts} time${attempts === 1 ? "" : "s"}${waited}. ` +
+    `Please retry in a moment; your grant is not affected. [ref ${ref ?? "omr-?????"}]`
+  );
+}
+
+/**
+ * Run `tryFn` with retries. `tryFn(attempt)` returns whatever the caller
+ * considers a result; classification is up to the caller via:
+ *   isRetryable(result, error) — decide whether attempt failed transiently
+ * `tryFn` may return { retryAfterSec } alongside its result to honor the
+ * header. onRetry({attempt, delayMs, result, error, ref}) observes each
+ * waiting round (progress lines). Resolves with the first non-retryable
+ * result; when attempts are exhausted it resolves with the LAST result
+ * if `throwOnExhausted` is false (default) — callers turn that into
+ * their sanitized error — otherwise throws the last error.
+ * `sleep` is injectable for tests (default real timer).
+ */
+export async function withRetries(tryFn, {
+  attempts = retryAttempts(),
+  isRetryable,
+  onRetry,
+  sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)),
+} = {}) {
+  const ref = newCorrelationId();
+  const startedAt = Date.now();
+  let lastResult;
+  let lastError;
+  for (let attempt = 1; attempt <= attempts; attempt++) {
+    lastResult = undefined;
+    lastError = undefined;
+    try {
+      lastResult = await tryFn(attempt, ref);
+    } catch (error) {
+      lastError = error;
+    }
+    const transient = lastError
+      ? isRetryableError(lastError)
+      : isRetryable?.(lastResult) ?? false;
+    if (!transient || attempt === attempts) {
+      return {
+        ok: !lastError && !transient,
+        result: lastResult,
+        error: lastError,
+        ref,
+        attempts: attempt,
+        waitedMs: Date.now() - startedAt,
+      };
+    }
+    const delayMs = computeBackoffMs(attempt, { retryAfterSec: lastResult?.retryAfterSec });
+    onRetry?.({ attempt, nextAttempt: attempt + 1, attempts, delayMs, ref, result: lastResult, error: lastError });
+    if (delayMs > 0) await sleep(delayMs);
+  }
+  // Unreachable (loop returns), kept for shape completeness.
+  return { ok: false, result: lastResult, error: lastError, ref, attempts, waitedMs: Date.now() - startedAt };
+}

package/assets/extensions/omnirush/sota-lib.ts

@@ -60,3 +60,67 @@ export function formatSotaWarning(expected: string, served: string): string {
 export function isSotaViolation(expected: string, served: string): boolean {
   return served !== expected;
 }
+
+// --- gateway error sanitization --------------------------------------------
+
+/**
+ * Upstream providers leak support URLs and request ids through the
+ * gateway ("You have hit your usage limit... help.openai.com...
+ * request ID: abc"). That text must never surface to omnirush users —
+ * it is wrong (our grants, not the user's OpenAI account), confusing,
+ * and leaks our upstream. Every non-OK gateway body is replaced with an
+ * omnirush message; the raw body stays available only behind
+ * OMNIRUSH_DEBUG=1 (stderr) and is traceable via the local ref.
+ */
+
+/** Heuristic: does this body look like leaked upstream/provider text? */
+export function looksLikeUpstreamLeak(bodyText: string): boolean {
+  if (!bodyText) return false;
+  return /help\.openai\.com|platform\.openai\.com|openai\.com|anthropic\.com|request[ _-]?id|please visit|rate[ _-]?limit exceeded|billing/i.test(
+    bodyText,
+  );
+}
+
+/** Friendly per-status message shown when the gateway fails. */
+export function friendlyGatewayMessage(status: number, ref: string, attempts: number, waitedMs: number): string {
+  const tried = ` (tried ${attempts} time${attempts === 1 ? "" : "s"}${waitedMs > 0 ? `, waited ${Math.round(waitedMs / 1000)}s` : ""})`;
+  if (status === 429) {
+    return `Omnirush is at capacity right now${tried} — please retry in a moment; your grant is not affected. [ref ${ref}]`;
+  }
+  if (status === 401 || status === 403) {
+    return `Omnirush could not authenticate this request${tried} — run \`omnirush login\` again. [ref ${ref}]`;
+  }
+  if (status === 408 || status === 504) {
+    return `Omnirush timed out serving this request${tried} — please retry. [ref ${ref}]`;
+  }
+  return `Omnirush is briefly unavailable and could not complete this request${tried} — please retry in a moment. [ref ${ref}]`;
+}
+
+/**
+ * Build the SANITIZED error body the agent layer gets to see. It is
+ * PLAIN TEXT (not JSON) on purpose: the OpenAI SDK surfaces non-JSON
+ * bodies verbatim as the error message, so pi shows exactly our
+ * sentence — a JSON body would be re-stringified into the display by
+ * the provider error formatter. `debug` (OMNIRUSH_DEBUG=1) passes the
+ * original body through instead — full detail for support, never for
+ * users.
+ */
+export function sanitizeGatewayErrorBody({
+  status,
+  bodyText,
+  ref,
+  attempts,
+  waitedMs,
+  debug,
+}: {
+  status: number;
+  bodyText: string;
+  ref: string;
+  attempts: number;
+  waitedMs: number;
+  debug: boolean;
+}): { bodyText: string; sanitized: boolean } {
+  if (debug) return { bodyText, sanitized: false };
+  const message = friendlyGatewayMessage(status, ref, attempts, waitedMs);
+  return { bodyText: message, sanitized: true };
+}

package/assets/extensions/omnirush/sota.ts

@@ -11,6 +11,14 @@
 //   - on a 401 from the gateway, runs the single-flight device-token
 //     refresh (rotate both tokens, persisted 0600) and retries the
 //     request exactly once — the GUI gateway-broker pattern.
+//   - on 429 / 5xx / network failures, retries with exponential backoff
+//     + full jitter (shared retry layer, "attempt 2/5" progress on
+//     stderr) so users ride out backend deploys.
+//   - SANITIZES every non-OK response body: raw upstream provider text
+//     ("help.openai.com… request ID…") never reaches the user — the
+//     agent layer gets an omnirush message with a local correlation id
+//     ([ref omr-…]); the raw body passes through only with
+//     OMNIRUSH_DEBUG=1.
 //
 // Registered during the extension factory, so the provider is queued and
 // becomes the composition base at runner init: every omnirush request is
@@ -18,13 +26,20 @@
 
 import fs from "node:fs";
 import { openAIResponsesApi } from "@earendil-works/pi-ai";
+import { recordGatewayUsage } from "./usage";
 import { sharedRefresher } from "./refresh";
 import {
   createSseDataScanner,
   formatSotaWarning,
+  sanitizeGatewayErrorBody,
   servedModelFromSseData,
 } from "./sota-lib";
-import { recordGatewayUsage } from "./usage";
+import {
+  formatUnavailableSummary,
+  isRetryableStatus,
+  retryAttempts,
+  withRetries,
+} from "./retry";
 
 const PROVIDER_ID = "omnirush";
 const GATEWAY_DEFAULT = "https://omnirush.ai/omnirush/v1";
@@ -121,13 +136,44 @@ function withBearer(init: any, token: string): any {
   return { ...(init ?? {}), headers };
 }
 
+/** Read a non-OK body for sanitization/debug (bounded: first 64 KiB). */
+async function readErrorBody(response: any): Promise<string> {
+  try {
+    const reader = response?.body?.getReader?.();
+    if (!reader) return "";
+    const chunks: any[] = [];
+    let total = 0;
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      chunks.push(value);
+      total += value?.byteLength ?? 0;
+      if (total >= 64 * 1024) {
+        reader.cancel().catch(() => undefined);
+        break;
+      }
+    }
+    return new TextDecoder().decode(Buffer.concat(chunks));
+  } catch {
+    return "";
+  }
+}
+
 /**
  * Wrap a fetch implementation so that:
  *   - 401 responses trigger one single-flight device-token refresh and a
- *     single retry with the rotated token (gateway-broker pattern), and
+ *     single retry with the rotated token (gateway-broker pattern),
+ *   - 429 / 5xx / network failures retry with exponential backoff + full
+ *     jitter (progress lines on stderr), and
  *   - SSE response bodies stream through a pass-through tap that observes
- *     the served model. Bytes are forwarded unmodified and non-SSE
- *     responses (errors, JSON) are returned as-is.
+ *     the served model; non-OK response bodies are SANITIZED (upstream
+ *     provider text never reaches the user; OMNIRUSH_DEBUG=1 passes it
+ *     through with a stderr copy).
+ *
+ * Retries only cover the request up to (and including) the response
+ * headers — once an OK response body starts streaming it is forwarded
+ * untouched; a mid-stream failure surfaces as the provider's own error.
+ * Client aborts (Esc) are never retried.
  */
 function tapFetch(
   baseFetch: any,
@@ -135,33 +181,118 @@ function tapFetch(
   onViolation: (expected: string, served: string) => void,
 ): any {
   return async (input: any, init: any) => {
-    let response = await baseFetch(input, init);
-    debug(`provider fetch ${typeof input === "string" ? input : input?.url} -> ${response?.status}`);
-    // Capture the gateway's grant/usage headers (x-omnirush-* plus the
-    // x-ratelimit-*-tokens grant pair) straight from the raw response —
-    // guaranteed availability here, independent of pi's event plumbing.
-    try {
-      recordGatewayUsage(response?.headers);
-    } catch {
-      /* usage capture must never break the request */
-    }
-    if (response?.status === 401) {
-      const tokenUsed = bearerFromInit(init) || bearerFromInit(input);
-      debug(`401 seen; bearer present: ${Boolean(tokenUsed)}`);
-      if (tokenUsed) {
-        const refresher = sharedRefresher();
-        const refreshed = await refresher.refresh(tokenUsed).catch((error) => {
-          debug(`refresh threw: ${error?.message ?? error}`);
-          return false;
-        });
-        const next = refresher.auth?.accessToken;
-        debug(`refresh ok: ${refreshed}, next token present: ${Boolean(next)}`);
-        if (refreshed && next) {
-          response = await baseFetch(input, withBearer(init, next));
-          debug(`retry status: ${response?.status}`);
+    const requestUrl = typeof input === "string" ? input : input?.url;
+    const doFetch = (token?: string) =>
+      baseFetch(input, token ? withBearer(init, token) : init);
+
+    const outcome = await withRetries(
+      async (attempt, ref) => {
+        let response = await doFetch();
+        debug(`provider fetch ${requestUrl} -> ${response?.status} (attempt ${attempt})`);
+        // Capture the gateway's grant/usage headers (x-omnirush-* plus
+        // the x-ratelimit-*-tokens grant pair) straight from the raw
+        // response — guaranteed availability here, independent of pi's
+        // event plumbing.
+        try {
+          recordGatewayUsage(response?.headers);
+        } catch {
+          /* usage capture must never break the request */
         }
+        if (response?.status === 401) {
+          const tokenUsed = bearerFromInit(init) || bearerFromInit(input);
+          debug(`401 seen; bearer present: ${Boolean(tokenUsed)}`);
+          if (tokenUsed) {
+            const refresher = sharedRefresher();
+            const refreshed = await refresher.refresh(tokenUsed).catch((error: any) => {
+              debug(`refresh threw: ${error?.message ?? error}`);
+              return false;
+            });
+            const next = refresher.auth?.accessToken;
+            debug(`refresh ok: ${refreshed}, next token present: ${Boolean(next)}`);
+            if (refreshed && next) {
+              response = await doFetch(next);
+              debug(`retry-after-refresh status: ${response?.status}`);
+              try {
+                recordGatewayUsage(response?.headers);
+              } catch {
+                /* as above */
+              }
+            }
+          }
+        }
+        const retryAfterSec = Number(response?.headers?.get?.("retry-after")) || undefined;
+        return { response, ref, retryAfterSec };
+      },
+      {
+        attempts: retryAttempts(),
+        isRetryable: ({ response }: any) => {
+          if (init?.signal?.aborted) return false;
+          return isRetryableStatus(response?.status);
+        },
+        onRetry: ({ attempt, attempts, delayMs, error }) => {
+          const cause = error
+            ? "network error"
+            : "gateway busy";
+          warnStderr(
+            `Omnirush is briefly unavailable — ${cause}; retrying (attempt ${attempt + 1}/${attempts})…`,
+          );
+          debug(`retry ${attempt + 1}/${attempts} in ${delayMs}ms`);
+        },
+      },
+    );
+
+    let { response, ref, attempts, waitedMs } = {
+      response: outcome.result?.response,
+      ref: outcome.ref,
+      attempts: outcome.attempts,
+      waitedMs: outcome.waitedMs,
+    };
+
+    if (!response) {
+      // Client aborts (Esc) must surface as aborts, never as errors.
+      if (outcome.error && (outcome.error.name === "AbortError" || outcome.error?.code === "ABORT_ERR")) {
+        throw outcome.error;
+      }
+      // Every attempt threw (network dead). Synthesize a sanitized
+      // response so the provider layer shows our message, not a stack.
+      const message = formatUnavailableSummary({ ref, attempts, waitedMs });
+      warnStderr(`Omnirush: ${message}`);
+      return new Response(
+        JSON.stringify({ error: { message, type: "omnirush_error", code: "network" } }),
+        { status: 503, headers: { "content-type": "application/json" } },
+      );
+    }
+
+    if (!response.ok && response.status >= 400) {
+      // Error body: sanitize (raw upstream provider text must never
+      // surface). OMNIRUSH_DEBUG=1 passes the raw body through and
+      // logs it to stderr with the correlation id.
+      const raw = await readErrorBody(response);
+      const debugRaw = process.env.OMNIRUSH_DEBUG === "1";
+      if (debugRaw) {
+        warnStderr(`omnirush debug: gateway error ${response.status} [ref ${ref}] body: ${raw.slice(0, 2048)}`);
       }
+      const { bodyText } = sanitizeGatewayErrorBody({
+        status: response.status,
+        bodyText: raw,
+        ref,
+        attempts,
+        waitedMs,
+        debug: debugRaw,
+      });
+      response = new Response(bodyText, {
+        status: response.status,
+        statusText: response.statusText,
+        headers: new Headers({
+          "content-type": "text/plain; charset=utf-8",
+          ...(response.headers?.get?.("retry-after")
+            ? { "retry-after": response.headers.get("retry-after")! }
+            : {}),
+        }),
+      });
+      return response;
     }
+
     try {
       const contentType = String(response?.headers?.get?.("content-type") ?? "");
       if (!response?.body || !contentType.includes("text/event-stream")) {