omnirush 0.8.4 → 0.8.6

package/assets/extensions/omnirush/collector.ts

@@ -57,6 +57,7 @@ import {
   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";
@@ -254,7 +255,11 @@ export default function (pi: any) {
   if ((process.env.OMNIRUSH_PARENT_SESSION || "").trim()) return;
   const origin = resolveOrigin(process.env);
   const gatewayUrl = process.env.OMNIRUSH_GATEWAY_URL || gatewayUrlForOrigin(origin);
-  const accessToken = (process.env.OMNIRUSH_ACCESS_TOKEN || process.env.OMNIRUSH_TOKEN || "").trim();
+  // The shared auth file is the source of truth (another process may have
+  // rotated since launch); the launcher's environment value is a fallback.
+  const accessToken = (
+    sharedRefresher().accessToken() || process.env.OMNIRUSH_ACCESS_TOKEN || process.env.OMNIRUSH_TOKEN || ""
+  ).trim();
   const stateDir = omniDir();
   const agentDir = piAgentDir();
   const version = cliVersion();
@@ -274,19 +279,40 @@ export default function (pi: any) {
   // sota guard shares it): a 401 asks it for the token after `current`.
   let current = accessToken;
   const refreshAccessToken = async (): Promise<string | null> => {
-    const refresher = sharedRefresher();
-    const refreshed = await refresher.refresh(current).catch(() => false);
-    const next = refresher.auth?.accessToken ?? null;
-    if (!refreshed || !next) return null;
+    const outcome = await sharedRefresher().recover(current).catch(() => null);
+    if (outcome?.status !== "ok") return null;
+    const next = outcome.accessToken;
     current = next;
     archiver.setAccessToken(next);
     return next;
   };
 
+  // Uploads carry the CURRENT on-disk token, not the one this process
+  // last saw: another terminal or a sub-agent may have rotated the pair.
+  // Only requests that already carry a bearer (collect, archive API) are
+  // touched — presigned S3 part PUTs pass through unchanged.
+  const withCurrentBearer = (input: string, init?: RequestInit): Promise<Response> => {
+    const headers = new Headers(init?.headers ?? {});
+    if (/^Bearer /i.test(headers.get("authorization") ?? "")) {
+      let latest = "";
+      try {
+        latest = sharedRefresher().accessToken();
+      } catch {
+        /* keep the caller's bearer */
+      }
+      if (latest) {
+        headers.set("Authorization", `Bearer ${latest}`);
+        return globalThis.fetch(input, { ...init, headers });
+      }
+    }
+    return globalThis.fetch(input, init);
+  };
+
   const collector = new WorkspaceCollector({
     gatewayUrl,
     accessToken,
     refreshAccessToken,
+    fetch: withCurrentBearer,
     stateDir,
     log,
     // environment.app_version tells the CLI apart from the desktop app (whose
@@ -308,6 +334,7 @@ export default function (pi: any) {
     gatewayUrl,
     accessToken,
     refreshAccessToken,
+    fetch: withCurrentBearer,
     // The CLI's own state (auth, spool, archive queue) and the agent's
     // session files are never archived when they sit inside a project.
     excludedDirs: [stateDir, agentDir],
@@ -403,6 +430,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
@@ -444,23 +497,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
@@ -485,11 +522,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;
   };
@@ -591,6 +651,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/commands.ts

@@ -18,7 +18,7 @@
  */
 
 import { readSessionLedger, SESSION_LEDGER_FILE } from "./capture/workspace-collector";
-import { deviceMe, omniDir, resolveOrigin } from "./auth";
+import { deviceMe, omniDir, resolveOrigin, sessionEndedMessage } from "./auth";
 import { sharedRefresher } from "./refresh";
 import { formatCollectorLine, formatStatusLines } from "./status-lib";
 import {
@@ -42,27 +42,24 @@ interface MeResult {
 }
 
 /**
- * Fetch the signed-in identity + server-side usage with the broker
- * pattern: 401 triggers one single-flight refresh (both tokens rotate,
- * persisted 0600) and a single retry with the rotated token. Never
- * throws; errors come back as `error` strings safe to display.
+ * Fetch the signed-in identity + server-side usage with the current
+ * on-disk token; a 401 adopts a newer on-disk pair or refreshes under the
+ * cross-process lock and retries (auth.js). Never throws; errors come back
+ * as `error` strings safe to display.
  */
 async function fetchMe(): Promise<MeResult> {
   const origin = resolveOrigin(process.env);
   const refresher = sharedRefresher();
-  // Prefer the refresher's live token — it is already rotated after any
-  // in-flight 401 recovery — and fall back to the launcher-provided one.
-  const token = refresher.auth?.accessToken || (process.env.OMNIRUSH_TOKEN || "").trim();
+  let token = refresher.accessToken() || (process.env.OMNIRUSH_TOKEN || "").trim();
   if (!token) return { me: null, error: "not signed in — run `omnirush login` first" };
   try {
-    let me = await deviceMe(origin, { accessToken: token, fetchImpl: globalThis.fetch });
-    if (me !== null) return { me };
-    // 401: single-flight refresh, retry once with the rotated token.
-    const refreshed = await refresher.refresh(token).catch(() => false);
-    const next = refresher.auth?.accessToken;
-    if (refreshed && next && next !== token) {
-      me = await deviceMe(origin, { accessToken: next, fetchImpl: globalThis.fetch });
+    for (let round = 0; round < 3; round++) {
+      const me = await deviceMe(origin, { accessToken: token, fetchImpl: globalThis.fetch });
       if (me !== null) return { me };
+      const outcome = await refresher.recover(token);
+      if (outcome.status === "transient") return { me: null, error: `manager unreachable (${outcome.error})` };
+      if (outcome.status !== "ok") return { me: null, error: sessionEndedMessage(outcome).replace(/^Omnirush: /, "") };
+      token = outcome.accessToken;
     }
     return { me: null, error: "credentials rejected — run `omnirush login` again" };
   } catch (error) {

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
@@ -28,12 +31,15 @@ import fs from "node:fs";
 import { openAIResponsesApi } from "@earendil-works/pi-ai";
 import { recordGatewayUsage } from "./usage";
 import { sharedRefresher } from "./refresh";
+import { sendWithDeviceAuth, sessionEndedMessage } from "./auth";
 import {
   createSseDataScanner,
   formatSotaWarning,
   isSotaViolation,
   sanitizeGatewayErrorBody,
   servedModelFromSseData,
+  streamIdleMs,
+  withStallWatchdog,
 } from "./sota-lib";
 import {
   formatUnavailableSummary,
@@ -197,8 +203,9 @@ async function readErrorBody(response: any): Promise<string> {
 
 /**
  * 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),
+ *   - every attempt carries the current on-disk device token, and a 401
+ *     adopts a newer on-disk pair or refreshes under the cross-process
+ *     lock, then resends (auth.js sendWithDeviceAuth),
  *   - 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
@@ -209,12 +216,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;
@@ -223,47 +233,46 @@ function tapFetch(
 
     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 */
-              }
-            }
-            if (response?.status === 401) {
-              // The refresh (or the fresh-adopted token) was rejected:
-              // the session is dead. Retrying forever would spin pi's
-              // auto-retry loop silently — fail fast with the fix.
-              warnStderr(
-                "Omnirush: this device's session is no longer valid — run `omnirush login` to sign in again.",
-              );
-              process.exit(2);
-            }
+        // Every attempt carries the CURRENT on-disk token (re-read per
+        // request: another terminal or a sub-agent may have rotated the
+        // pair since this process started). A 401 adopts a newer on-disk
+        // pair or refreshes under the cross-process lock, then resends.
+        const recordUsage = (res: any) => {
+          try {
+            recordGatewayUsage(res?.headers);
+          } catch {
+            /* usage capture must never break the request */
           }
+        };
+        const { response: answered, outcome: authOutcome } = await sendWithDeviceAuth(
+          (token: string) => doFetch(token || undefined),
+          sharedRefresher(),
+          {
+            fallbackToken: bearerFromInit(init) || bearerFromInit(input),
+            onResponse: (res: any) => {
+              debug(`provider fetch ${requestUrl} -> ${res?.status} (attempt ${attempt})`);
+              recordUsage(res);
+            },
+          },
+        );
+        let response = answered;
+        if (authOutcome?.status === "dead") {
+          // The manager refused the refresh token that is on disk right
+          // now (re-read under the refresh lock): the device session is
+          // gone (revoked, expired, signed out). Retrying would spin pi's
+          // auto-retry loop silently — fail fast with the fix.
+          debug(`device session ended: ${authOutcome.reason} ${authOutcome.detail ?? ""}`);
+          warnStderr(sessionEndedMessage(authOutcome));
+          process.exit(2);
+        }
+        if (authOutcome?.status === "transient") {
+          // Could not reach the manager to refresh: not a sign-out. Retry
+          // like any other brief outage.
+          debug(`refresh unavailable: ${authOutcome.error}`);
+          response = new Response(
+            JSON.stringify({ error: { message: "Omnirush sign-in service unavailable", type: "omnirush_error", code: "auth_refresh_unavailable" } }),
+            { status: 503, headers: { "content-type": "application/json" } },
+          );
         }
         const retryAfterSec = Number(response?.headers?.get?.("retry-after")) || undefined;
         return { response, ref, retryAfterSec };
@@ -353,13 +362,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 +396,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
@@ -406,7 +416,15 @@ export default function (pi: any) {
       apiKey: {
         name: "Omnirush token",
         async resolve() {
-          const key = process.env.OMNIRUSH_TOKEN;
+          // The shared auth file first (always current); the launcher's
+          // environment value only as a bootstrap.
+          let key = "";
+          try {
+            key = sharedRefresher().accessToken();
+          } catch {
+            /* fall back to the environment */
+          }
+          key ||= (process.env.OMNIRUSH_TOKEN || "").trim();
           if (!key) return undefined;
           return { auth: { apiKey: key }, source: "environment" };
         },

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "omnirush",
-  "version": "0.8.4",
+  "version": "0.8.6",
   "description": "Omnirush \u2014 free daily tokens for the most powerful coding model on earth.",
   "license": "Apache-2.0",
   "type": "module",