omnirush 0.8.3 → 0.8.5

package/assets/extensions/omnirush/collector.ts

@@ -36,10 +36,11 @@
 // signed-in omnirush identities.
 
 import { createHash } from "node:crypto";
-import { readdirSync, readFileSync, rmSync } from "node:fs";
-import { readdir, readFile, stat } from "node:fs/promises";
+import { createReadStream, readdirSync, readFileSync, rmSync } from "node:fs";
+import { readdir, stat } from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
+import { createInterface } from "node:readline";
 
 import { gatewayUrlForOrigin, omniDir, resolveOrigin } from "./auth";
 import { WorkspaceCollector } from "./capture/workspace-collector";
@@ -50,13 +51,13 @@ import {
   engineMessagesFromEntries,
   MAX_TURN_MESSAGE_BYTES,
   messageParts,
-  parseSessionFile,
   turnHistory,
   turnModelFromMessages,
   type EngineMessage,
   type PiEntry,
 } from "./pi-engine";
 import { sharedRefresher } from "./refresh";
+import { AGENTS_HOOK, type AgentsHook } from "./agents-lib";
 
 /** info.agent of the root session's messages (a sub-agent's is its role). */
 export const ROOT_AGENT = "omnirush-cli";
@@ -158,6 +159,32 @@ async function findSessionFile(sessionId: string, dirs: readonly string[]): Prom
   return null;
 }
 
+/** Parse a session JSONL file without buffering the raw transcript or its line array. */
+export async function parseSessionFileStream(file: string): Promise<{ header: PiEntry | null; entries: PiEntry[] }> {
+  let header: PiEntry | null = null;
+  const entries: PiEntry[] = [];
+  const input = createReadStream(file, { encoding: "utf8" });
+  const lines = createInterface({ input, crlfDelay: Infinity });
+  try {
+    for await (const line of lines) {
+      if (!line.trim()) continue;
+      let value: unknown;
+      try {
+        value = JSON.parse(line);
+      } catch {
+        continue;
+      }
+      if (!value || typeof value !== "object" || Array.isArray(value)) continue;
+      if ((value as { type?: unknown }).type === "session") header ??= value as PiEntry;
+      else entries.push(value as PiEntry);
+    }
+  } finally {
+    lines.close();
+    input.destroy();
+  }
+  return { header, entries };
+}
+
 type Collector = Pick<
   WorkspaceCollector,
   "recordTrace" | "recordChildSession" | "childCheckpoints" | "childSessionIds"
@@ -195,7 +222,7 @@ export async function captureChildSessions(input: {
       const stats = await stat(file).catch(() => null);
       const signature = stats ? `${stats.size}:${stats.mtimeMs}` : "";
       if (signature && input.seenFiles.get(child.id) === signature && input.known.has(child.id)) continue;
-      const parsed = parseSessionFile(await readFile(file, "utf8").catch(() => ""));
+      const parsed = await parseSessionFileStream(file).catch(() => ({ header: null, entries: [] }));
       entries = branchOfEntries(parsed.entries);
       cwd = typeof parsed.header?.cwd === "string" ? parsed.header.cwd : "";
       if (signature) input.seenFiles.set(child.id, signature);
@@ -377,6 +404,32 @@ export default function (pi: any) {
     return session;
   };
 
+  /**
+   * The sub-agent tree's messages since their checkpoints. Background
+   * sub-agents keep running past the turn that started them: each later
+   * settle (and the session's end) records what they did since.
+   */
+  const captureChildren = async (state: NonNullable<typeof session>, room: (bytes: number) => void): Promise<void> => {
+    const { id } = state;
+    try {
+      await captureChildSessions({
+        collector,
+        rootSessionId: id,
+        parentSessionId: id,
+        parentEntries: state.ctx?.sessionManager?.getBranch?.() ?? [],
+        depth: 1,
+        sessionDirs: sessionDirs(state.ctx),
+        checkpoints: await collector.childCheckpoints(id),
+        known: new Set(await collector.childSessionIds(id)),
+        visited: new Set([id]),
+        seenFiles: state.seenFiles,
+        beforeRecord: room,
+      });
+    } catch (error) {
+      collector.recordTrace(id, "session.children_failed", { error: error instanceof Error ? error.message : "unknown" });
+    }
+  };
+
   /**
    * Everything a settled run records, in the desktop observer's order:
    * messages, sub-agents, the idle event, the snapshot, the trace and the
@@ -418,23 +471,7 @@ export default function (pi: any) {
       }
       pending += bytes;
     };
-    try {
-      await captureChildSessions({
-        collector,
-        rootSessionId: id,
-        parentSessionId: id,
-        parentEntries: state.ctx?.sessionManager?.getBranch?.() ?? [],
-        depth: 1,
-        sessionDirs: sessionDirs(state.ctx),
-        checkpoints: await collector.childCheckpoints(id),
-        known: new Set(await collector.childSessionIds(id)),
-        visited: new Set([id]),
-        seenFiles: state.seenFiles,
-        beforeRecord: room,
-      });
-    } catch (error) {
-      collector.recordTrace(id, "session.children_failed", { error: error instanceof Error ? error.message : "unknown" });
-    }
+    await captureChildren(state, room);
     room(newest.bytes);
     collector.recordTrace(id, "session.idle", { status });
     // The turn snapshot runs first so the artifacts it discovers are part of
@@ -459,11 +496,34 @@ export default function (pi: any) {
    * settled too, so they reach the trace.
    */
   const settleRemaining = async (): Promise<void> => {
+    // Background sub-agents still running are stopped first (their session
+    // files are complete once they exited) and marked interrupted.
+    const hook = (globalThis as any)[AGENTS_HOOK] as AgentsHook | undefined;
+    const followed = session;
+    const interrupted = followed?.started && hook ? await hook.interruptAll(followed.id).catch(() => []) : [];
     await settleRunning("interrupted");
     const state = session;
     if (!state?.started) return;
+    for (const child of interrupted) {
+      collector.recordTrace(state.id, "session.child_interrupted", { child_session_id: child.sessionId, agent: child.role });
+    }
     const lastId = messagesOf(state).at(-1)?.info.id;
-    if (typeof lastId !== "string" || lastId === state.lastMessageId) return;
+    if (typeof lastId !== "string" || lastId === state.lastMessageId) {
+      // No new message of its own: what its sub-agents did since the last settle.
+      settling = settling.then(async () => {
+        let pending = 0;
+        await captureChildren(state, (bytes) => {
+          if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
+            collector.flushTrace(state.id);
+            pending = 0;
+          }
+          pending += bytes;
+        });
+        collector.flushTrace(state.id);
+      }).catch(() => undefined);
+      await settling;
+      return;
+    }
     settling = settling.then(() => settle(state, "idle")).catch(() => undefined);
     await settling;
   };
@@ -565,6 +625,13 @@ export default function (pi: any) {
     archive.turnFollowed(state.id);
   });
 
+  // A run that no prompt started (a background sub-agent's results
+  // delivered as a message start one): it settles like any other.
+  pi.on("agent_start", (_event: any, ctx: any) => {
+    const state = follow(ctx, false);
+    if (state?.started) state.running = true;
+  });
+
   pi.on("tool_execution_start", (event: any) => {
     if (session && typeof event?.toolCallId === "string") session.toolStarts.set(event.toolCallId, Date.now());
   });

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.3",
+  "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"
   }
-}
+}