omnirush 0.8.5 → 0.9.0

package/assets/extensions/omnirush/pi-engine.ts

@@ -33,6 +33,10 @@
 //                              a command the user ran)
 //   custom / branch summary -> user message with synthetic text parts (what
 //                              the model was shown; never a turn opener)
+//   background command      -> the delivery message as above, answered by an
+//     results (bgshell.ts)     assistant message holding one bash tool part
+//                              per finished command (its output and exit
+//                              code; no model step, no tokens)
 //   compaction              -> user message with a compaction part and the
 //                              summary as an assistant `summary` message
 //   system prompt           -> not a message (the desktop records it in
@@ -49,6 +53,9 @@
 import { createHash } from "node:crypto";
 import { isAbsolute, resolve } from "node:path";
 
+/** customType of bgshell.ts's finished-background-command message. */
+export const BACKGROUND_BASH_TYPE = "omnirush-bash-result";
+
 export type EngineMessage = { info: Record<string, unknown>; parts: Array<Record<string, unknown>> };
 
 /** One pi session entry (loosely typed: the files are parsed without validation). */
@@ -295,6 +302,60 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
     return messageVariant ?? null;
   };
 
+  /**
+   * Finished background shell commands (bgshell.ts delivery): their output
+   * as bash tool parts, the way the engine records a command the user ran.
+   */
+  const backgroundBash = (parentId: string, entryId: string, details: unknown, at: number | null, index: number): void => {
+    const jobs = isRecord(details) && Array.isArray(details.jobs) ? details.jobs.filter(isRecord) : [];
+    if (jobs.length === 0) return;
+    const shellId = engineMessageId(entryId, "_bash");
+    const shellModel = model ?? nextAssistantModel.get(index) ?? null;
+    const time = at ?? 0;
+    const parts: Array<Record<string, unknown>> = jobs.map((job, jobIndex) => {
+      const start = typeof job.started_at === "number" ? job.started_at : time;
+      const end = typeof job.finished_at === "number" ? job.finished_at : time;
+      const failed = job.status === "failed";
+      return {
+        type: "tool",
+        callID: `bg_${typeof job.id === "string" ? job.id : jobIndex}_${entryId}`,
+        tool: "bash",
+        state: {
+          status: failed ? "error" : "completed",
+          input: { command: typeof job.command === "string" ? job.command : "", background_id: job.id ?? null },
+          ...(failed ? { error: typeof job.error === "string" ? job.error : "failed" } : {}),
+          output: typeof job.output === "string" ? job.output : "",
+          metadata: {
+            background: true,
+            ...(typeof job.exit_code === "number" ? { exit: job.exit_code } : {}),
+            ...(typeof job.status === "string" ? { status: job.status } : {}),
+            ...(job.interrupted === true ? { interrupted: true } : {}),
+            ...(job.truncated === true ? { truncated: true } : {}),
+            ...(typeof job.tool_call_id === "string" ? { startedBy: job.tool_call_id } : {}),
+          },
+          time: { start, end },
+        },
+      };
+    });
+    out.push({
+      info: {
+        id: shellId,
+        sessionID: sessionId,
+        role: "assistant",
+        parentID: parentId,
+        time: { created: time, completed: time },
+        ...(shellModel ? { modelID: shellModel.modelID, providerID: shellModel.providerID } : {}),
+        agent,
+        mode: agent,
+        path: { cwd, root: cwd },
+        cost: 0,
+        tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+        finish: "stop",
+      },
+      parts: parts.map((part, partIndex) => ({ id: partId(shellId, partIndex), sessionID: sessionId, messageID: shellId, ...part })),
+    });
+  };
+
   entries.forEach((entry, index) => {
     const entryId = typeof entry.id === "string" && entry.id ? entry.id : null;
     if (!entryId) return;
@@ -340,7 +401,9 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
     }
     if (entry.type === "custom_message") {
       const parts = contentParts(entry.content, { synthetic: true, metadata: { customType: entry.customType ?? null } });
-      if (parts.length > 0) userMessage(engineMessageId(entryId), at, parts, index);
+      const id = engineMessageId(entryId);
+      if (parts.length > 0) userMessage(id, at, parts, index);
+      if (entry.customType === BACKGROUND_BASH_TYPE) backgroundBash(id, entryId, (entry as any).details, at, index);
       return;
     }
     if (entry.type !== "message" || !isRecord(entry.message)) return;
@@ -356,6 +419,7 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
       case "custom": {
         const parts = contentParts(message.content, { synthetic: true, metadata: { customType: message.customType ?? null } });
         if (parts.length > 0) userMessage(id, created, parts, index);
+        if (message.customType === BACKGROUND_BASH_TYPE) backgroundBash(id, entryId, message.details, created, index);
         return;
       }
       case "branchSummary":

package/assets/extensions/omnirush/sota.ts

@@ -31,6 +31,7 @@ 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,
@@ -40,6 +41,14 @@ import {
   streamIdleMs,
   withStallWatchdog,
 } from "./sota-lib";
+import { omniDir } from "./auth";
+import {
+  ENV_FALLBACK_EFFORT,
+  ENV_FALLBACK_MODEL,
+  FALLBACK_MARKER,
+  subagentFallbackFetch,
+  type FallbackRoute,
+} from "./subagents-lib";
 import {
   formatUnavailableSummary,
   isRetryableStatus,
@@ -202,8 +211,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
@@ -220,7 +230,7 @@ async function readErrorBody(response: any): Promise<string> {
  */
 function tapFetch(
   baseFetch: any,
-  expectedModelId: string,
+  expectedModel: string | (() => string),
   onViolation: (expected: string, served: string) => void,
   idleMs: number,
 ): any {
@@ -231,47 +241,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 };
@@ -355,6 +364,8 @@ function tapFetch(
       const onData = (data: string) => {
         if (reported) return;
         const served = servedModelFromSseData(data);
+        // A sub-agent moved to the main model expects that one.
+        const expectedModelId = typeof expectedModel === "function" ? expectedModel() : expectedModel;
         if (served === undefined || !isSotaViolation(expectedModelId, served)) return;
         reported = true;
         onViolation(expectedModelId, served);
@@ -389,22 +400,60 @@ export default function (pi: any) {
   const responses = openAIResponsesApi();
   const strictSota = (): boolean => pi.getFlag("strict-sota") === true;
 
+  // A sub-agent on a picked model: the main model it falls back to when the
+  // gateway refuses the picked one (set by the parent's spawn_agents).
+  const fallbackModel = (process.env[ENV_FALLBACK_MODEL] || "").trim();
+  const fallbackEffort = (process.env[ENV_FALLBACK_EFFORT] || "").trim() || null;
+
   const streamWith =
     (simple: boolean) =>
     (model: any, context: any, options: any) => {
-      const baseFetch = options?.fetch ?? globalThis.fetch;
+      let baseFetch = options?.fetch ?? globalThis.fetch;
+      const route: FallbackRoute = { expected: model.id, movedTo: null };
+      if (fallbackModel && fallbackModel !== model.id) {
+        baseFetch = subagentFallbackFetch(baseFetch, {
+          model: fallbackModel,
+          effort: fallbackEffort,
+          dir: omniDir(),
+          route,
+          report: (event) => {
+            // The parent's spawn_agents reads this line from the child's stderr.
+            try {
+              process.stderr.write(`\n${FALLBACK_MARKER}${JSON.stringify(event)}\n`);
+            } catch {
+              /* stderr may be closed */
+            }
+          },
+        });
+      }
       const tapped = {
         ...(options ?? {}),
         fetch: tapFetch(
           baseFetch,
-          model.id,
+          () => route.expected,
           (expected, served) => reportViolation(strictSota(), expected, served),
           streamIdleMs(process.env.OMNIRUSH_STREAM_IDLE_MS, options?.timeoutMs),
         ),
       };
-      return simple
+      const stream = simple
         ? responses.streamSimple(model, context, tapped)
         : responses.stream(model, context, tapped);
+      if (fallbackModel && stream && typeof stream.push === "function") {
+        // The message records the model (and effort) that really answered.
+        const push = stream.push.bind(stream);
+        stream.push = (event: any) => {
+          if (route.movedTo) {
+            for (const message of [event?.partial, event?.message, event?.error]) {
+              if (message && typeof message === "object" && message.role === "assistant") {
+                message.model = route.movedTo.model;
+                if (route.movedTo.effort) message.providerThinkingLevel = route.movedTo.effort;
+              }
+            }
+          }
+          return push(event);
+        };
+      }
+      return stream;
     };
 
   pi.registerProvider({
@@ -415,7 +464,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/assets/extensions/omnirush/status-lib.ts

@@ -52,6 +52,8 @@ export interface StatusInput {
   origin: string;
   /** omnirush CLI version ("unknown" when not provided by the launcher). */
   version: string;
+  /** The approval mode line (yolo-lib formatModeStatus), when known. */
+  yolo?: string | null;
 }
 
 /**
@@ -79,6 +81,7 @@ export function formatStatusLines(input: StatusInput): string[] {
   }
 
   lines.push(formatCollectorLine(input.collector));
+  if (input.yolo) lines.push(input.yolo);
   lines.push(`manager: ${input.origin}`);
   lines.push(`version: ${input.version}`);
   return lines;