omnirush 0.8.4 → 0.8.5

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

@@ -138,3 +138,132 @@ export function sanitizeGatewayErrorBody({
   const message = friendlyGatewayMessage(status, ref, attempts, waitedMs);
   return { bodyText: message, sanitized: true };
 }
+
+// --- stream stall watchdog -------------------------------------------------
+//
+// The agent core runs on Bun, and there its undici transport never
+// settles a pending body read when the connection dies mid-stream: the
+// HTTP idle timeout (or the gateway/edge dropping the socket) closes the
+// socket, but `reader.read()` stays pending forever. The turn then sits
+// on "Working…" until the user presses Esc — no error, no auto-retry
+// (field report: session 01a0dd1a, a 13-minute hang after a tool call).
+// On Node the same close surfaces as a `terminated` error that pi retries.
+// The watchdog gives Bun that behavior: it re-exposes the body through a
+// stream of our own and errors it when the upstream read has been
+// pending for `idleMs`, so pi's retry takes over.
+
+/** setTimeout's ceiling; pi passes it (2^31-1) for a disabled timeout. */
+const MAX_TIMER_MS = 2_147_483_647;
+export const DEFAULT_STREAM_IDLE_MS = 300_000;
+
+/**
+ * The idle limit for a model stream: OMNIRUSH_STREAM_IDLE_MS wins, then
+ * the request timeout pi hands the provider (its httpIdleTimeoutMs /
+ * retry.provider.timeoutMs), then 5 minutes. 0 (or a value at the timer
+ * ceiling, pi's "disabled") turns the watchdog off.
+ */
+export function streamIdleMs(envValue: string | undefined, requestTimeoutMs: unknown): number {
+  const pick = (value: unknown): number | undefined => {
+    if (value === undefined || value === null || value === "") return undefined;
+    const n = Number(value);
+    if (!Number.isFinite(n) || n < 0) return undefined;
+    return n >= MAX_TIMER_MS ? 0 : Math.floor(n);
+  };
+  return pick(envValue) ?? pick(requestTimeoutMs) ?? DEFAULT_STREAM_IDLE_MS;
+}
+
+/** The error a stalled stream fails with ("timed out" keeps it retryable for pi). */
+export function streamStallError(idleMs: number): Error {
+  const seconds = Math.max(1, Math.round(idleMs / 1000));
+  const error = new Error(
+    `Omnirush model stream timed out: no data from the gateway for ${seconds}s — retrying`,
+  );
+  error.name = "OmnirushStreamStall";
+  return error;
+}
+
+/**
+ * Re-expose `source` as a new ReadableStream that errors with
+ * `streamStallError` when one upstream read stays pending for `idleMs`
+ * (0 = never), and with the abort reason when `signal` aborts. Every
+ * chunk is handed to `onChunk` before it is enqueued. The upstream
+ * reader is cancelled (best effort, never awaited — it may be the hung
+ * one) whenever the wrapper errors or is cancelled.
+ */
+export function withStallWatchdog(
+  source: ReadableStream<any>,
+  {
+    idleMs,
+    signal,
+    onChunk,
+  }: { idleMs: number; signal?: AbortSignal | null; onChunk?: (chunk: any) => void },
+): ReadableStream<any> {
+  const reader = source.getReader();
+  let controller: ReadableStreamDefaultController<any>;
+  let finished = false;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+
+  const clearTimer = () => {
+    if (timer !== undefined) clearTimeout(timer);
+    timer = undefined;
+  };
+  const finish = () => {
+    finished = true;
+    clearTimer();
+    signal?.removeEventListener?.("abort", onAbort);
+  };
+  const fail = (error: unknown) => {
+    if (finished) return;
+    finish();
+    try {
+      controller.error(error);
+    } catch {
+      /* already closed */
+    }
+    try {
+      reader.cancel(error).catch(() => undefined);
+    } catch {
+      /* already released */
+    }
+  };
+  function onAbort() {
+    const reason = (signal as any)?.reason;
+    fail(reason ?? Object.assign(new Error("This operation was aborted"), { name: "AbortError" }));
+  }
+
+  return new ReadableStream<any>({
+    start(c) {
+      controller = c;
+      if (signal?.aborted) onAbort();
+      else signal?.addEventListener?.("abort", onAbort, { once: true });
+    },
+    async pull(c) {
+      if (finished) return;
+      if (idleMs > 0) timer = setTimeout(() => fail(streamStallError(idleMs)), Math.min(idleMs, MAX_TIMER_MS - 1));
+      let result: ReadableStreamReadResult<any>;
+      try {
+        result = await reader.read();
+      } catch (error) {
+        fail(error);
+        return;
+      }
+      clearTimer();
+      if (finished) return;
+      if (result.done) {
+        finish();
+        c.close();
+        return;
+      }
+      try {
+        onChunk?.(result.value);
+      } catch {
+        /* observers never break the stream */
+      }
+      c.enqueue(result.value);
+    },
+    cancel(reason) {
+      finish();
+      return reader.cancel(reason).catch(() => undefined);
+    },
+  });
+}

package/assets/extensions/omnirush/sota.ts

@@ -14,6 +14,9 @@
 //   - 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.
+//   - fails a model stream that stops delivering data (stall watchdog)
+//     with a retryable error, so a dropped connection never leaves the
+//     turn hanging on Bun (pi retries it like on Node).
 //   - 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
@@ -34,6 +37,8 @@ import {
   isSotaViolation,
   sanitizeGatewayErrorBody,
   servedModelFromSseData,
+  streamIdleMs,
+  withStallWatchdog,
 } from "./sota-lib";
 import {
   formatUnavailableSummary,
@@ -209,12 +214,15 @@ async function readErrorBody(response: any): Promise<string> {
  * 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.
+ * A body that stops delivering data for `idleMs` fails with a retryable
+ * "timed out" error (stall watchdog, see sota-lib) instead of hanging.
  * Client aborts (Esc) are never retried.
  */
 function tapFetch(
   baseFetch: any,
   expectedModelId: string,
   onViolation: (expected: string, served: string) => void,
+  idleMs: number,
 ): any {
   return async (input: any, init: any) => {
     const requestUrl = typeof input === "string" ? input : input?.url;
@@ -353,13 +361,11 @@ function tapFetch(
       };
       const scanner = createSseDataScanner(onData);
       const decoder = new TextDecoder();
-      const passthrough = new TransformStream({
-        transform(chunk: any, controller: any) {
-          controller.enqueue(chunk);
-          scanner(decoder.decode(chunk, { stream: true }));
-        },
+      const body = withStallWatchdog(response.body, {
+        idleMs,
+        signal: init?.signal,
+        onChunk: (chunk: any) => scanner(decoder.decode(chunk, { stream: true })),
       });
-      const body = response.body.pipeThrough(passthrough);
       return new Response(body, {
         status: response.status,
         statusText: response.statusText,
@@ -389,8 +395,11 @@ export default function (pi: any) {
       const baseFetch = options?.fetch ?? globalThis.fetch;
       const tapped = {
         ...(options ?? {}),
-        fetch: tapFetch(baseFetch, model.id, (expected, served) =>
-          reportViolation(strictSota(), expected, served),
+        fetch: tapFetch(
+          baseFetch,
+          model.id,
+          (expected, served) => reportViolation(strictSota(), expected, served),
+          streamIdleMs(process.env.OMNIRUSH_STREAM_IDLE_MS, options?.timeoutMs),
         ),
       };
       return simple

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "omnirush",
-  "version": "0.8.4",
+  "version": "0.8.5",
   "description": "Omnirush \u2014 free daily tokens for the most powerful coding model on earth.",
   "license": "Apache-2.0",
   "type": "module",
@@ -41,4 +41,4 @@
     "minimatch": "^10.2.6",
     "zod": "^4.6.5"
   }
-}
+}