omnirush 0.8.5 → 0.9.0

package/assets/extensions/omnirush/bgshell.ts

@@ -0,0 +1,392 @@
+// omnirush-bgshell — long shell commands move to the background instead of
+// being killed.
+//
+// Overrides the built-in bash tool (same execution, output handling and
+// rendering: pi's own bash tool definition runs each call, with an
+// execution backend that keeps the process in a job — bgshell-lib.ts):
+//   - a command that finishes within its foreground window (the call's
+//     `timeout` in seconds, default 300, OMNIRUSH_BASH_FOREGROUND_SECONDS)
+//     answers the call as before;
+//   - one still running then is NOT killed: it moves to the background and
+//     the call returns its id and the output so far; `run_in_background:
+//     true` does that at once (servers, watchers, long builds);
+//   - bash_output (new output and status), bash_wait and bash_kill manage
+//     background commands; when one finishes, its result comes back to the
+//     main agent as a message, the way background sub-agents' results do
+//     (deliveries.ts): steered into a running turn, starting a turn when the
+//     agent is idle, and holding a one-shot run open until it arrives;
+//   - background commands still running when the session ends (exit, /new,
+//     /resume) are stopped with a message saying so; their output is kept
+//     in the session (and so in the trace) as a bash tool part.
+
+import { createBashToolDefinition, createLocalBashOperations } from "@earendil-works/pi-coding-agent";
+import { writeSync } from "node:fs";
+import { Type } from "typebox";
+
+import {
+  BGSHELL_HOOK,
+  type BgShellHook,
+  DEFAULT_FOREGROUND_SECONDS,
+  jobDetails,
+  jobOutcome,
+  MAX_FOREGROUND_SECONDS,
+  MAX_WAIT_SECONDS,
+  renderJobDelivery,
+  renderJobResult,
+  renderJobStatus,
+  shortCommand,
+  type ShellExec,
+  type ShellJob,
+  ShellJobManager,
+  tailText,
+} from "./bgshell-lib";
+import { deliveryHub } from "./deliveries";
+import { BACKGROUND_BASH_TYPE } from "./pi-engine";
+
+/** customType of the message that brings finished background commands back (the trace maps it to bash tool parts). */
+export const BASH_DELIVERY_TYPE = BACKGROUND_BASH_TYPE;
+/** customType of the one-shot reminder about background commands still running. */
+export const BASH_REMINDER_TYPE = "omnirush-bash-reminder";
+
+class MovedToBackground extends Error {
+  constructor() {
+    super("moved to the background");
+  }
+}
+
+/** The default foreground window (seconds), OMNIRUSH_BASH_FOREGROUND_SECONDS overriding it. */
+export function defaultForegroundSeconds(env: NodeJS.ProcessEnv = process.env): number {
+  const raw = Number(String(env.OMNIRUSH_BASH_FOREGROUND_SECONDS ?? "").trim());
+  return Number.isFinite(raw) && raw > 0 ? Math.min(raw, MAX_FOREGROUND_SECONDS) : DEFAULT_FOREGROUND_SECONDS;
+}
+
+/** The foreground window of one call: its `timeout` (seconds), bounded; placeholders (<= 0) take the default. */
+export function foregroundSeconds(timeout: unknown, fallback: number = defaultForegroundSeconds()): number {
+  if (typeof timeout !== "number" || !Number.isFinite(timeout) || timeout <= 0) return fallback;
+  return Math.min(Math.max(timeout, 1), MAX_FOREGROUND_SECONDS);
+}
+
+function plural(count: number, word: string): string {
+  return `${count} ${word}${count === 1 ? "" : "s"}`;
+}
+
+export default function (pi: any, options: { exec?: ShellExec; now?: () => number } = {}) {
+  const hub = deliveryHub(pi);
+  const exec: ShellExec = options.exec ?? (createLocalBashOperations().exec as ShellExec);
+  const manager = new ShellJobManager({ exec, ...(options.now ? { now: options.now } : {}), onDeliverable: () => hub.notify() });
+  const fgDefault = defaultForegroundSeconds();
+
+  /** New output of each job, marked read. */
+  const collect = (jobs: ShellJob[]) => jobs.map((job) => ({ job, ...manager.read(job) }));
+
+  hub.addSource({
+    hasPending: (session) => manager.hasPending(session),
+    take: (session) => {
+      const jobs = manager.takeDeliverable(session);
+      if (jobs.length === 0) return [];
+      const parts = collect(jobs);
+      return [{
+        customType: BASH_DELIVERY_TYPE,
+        content: renderJobDelivery(parts),
+        details: {
+          jobs: parts.map(({ job, text, dropped }) => {
+            const shown = tailText(text);
+            return jobDetails(job, shown.text, shown.truncated || dropped > 0);
+          }),
+        },
+      }];
+    },
+    changed: () => manager.changed(),
+    beforeEnd: (session) => {
+      const jobs = manager.remindable(session);
+      if (jobs.length === 0) return null;
+      return {
+        customType: BASH_REMINDER_TYPE,
+        content: [
+          `This run is about to end while ${plural(jobs.length, "background command")} you started still run${jobs.length === 1 ? "s" : ""}:`,
+          renderJobStatus(jobs),
+          "If you need their results, wait for them with bash_wait; stop the ones you no longer need (servers, watchers) with bash_kill. Background commands still running when you finish are stopped.",
+        ].join("\n"),
+        details: { jobs: jobs.map((job) => job.id) },
+      };
+    },
+  });
+
+  // --- stopping at the end of the session -------------------------------------
+
+  /** What the session's end stopped (the collector's hook may do it before our own handler runs). */
+  let stoppedAtEnd: ShellJob[] = [];
+  const stopAll = async (session?: string): Promise<ShellJob[]> => {
+    const live = manager.running(session);
+    if (live.length === 0) return [];
+    const stopped = await manager.kill(live, { interrupted: true });
+    stoppedAtEnd.push(...stopped);
+    // Kept in the session (and so in the trace) like any tool output; the
+    // model never sees it — nothing runs after this.
+    const background = stopped.filter((job) => job.background);
+    if (background.length > 0) {
+      try {
+        const parts = collect(background);
+        pi.sendMessage(
+          {
+            customType: BASH_DELIVERY_TYPE,
+            content: [
+              `The session ended: ${plural(background.length, "background command")} stopped. Output since it was last read:`,
+              "",
+              parts.map(({ job, text, dropped }) => renderJobResult(job, text, { dropped }).text).join("\n\n---\n\n"),
+            ].join("\n"),
+            display: false,
+            details: {
+              jobs: parts.map(({ job, text, dropped }) => {
+                const shown = tailText(text);
+                return jobDetails(job, shown.text, shown.truncated || dropped > 0);
+              }),
+            },
+          },
+          { triggerTurn: false },
+        );
+      } catch {
+        /* the session may already be gone */
+      }
+    }
+    return stopped;
+  };
+  const hook: BgShellHook = {
+    async stopAll(sessionId?: string) {
+      const stopped = await stopAll(sessionId);
+      return stopped.map((job) => ({ id: job.id, command: job.command }));
+    },
+  };
+  (globalThis as any)[BGSHELL_HOOK] = hook;
+
+  pi.on("session_shutdown", async (event: any, ctx: any) => {
+    await stopAll();
+    const stopped = stoppedAtEnd.filter((job) => job.background || job.status === "killed");
+    stoppedAtEnd = [];
+    if (stopped.length === 0) return;
+    const text = `omnirush: stopped ${plural(stopped.length, "background command")} as the session ended: ${stopped
+      .map((job) => `${job.id} (${shortCommand(job.command, 60)})`)
+      .join(", ")}`;
+    if (!event?.reason || event.reason === "quit") {
+      // After the terminal is restored, so the line stays readable.
+      process.once("exit", () => {
+        try {
+          writeSync(2, `${text}\n`);
+        } catch {
+          /* stderr may be closed */
+        }
+      });
+    } else if (ctx?.ui?.notify) {
+      ctx.ui.notify(text, "warning");
+    }
+  });
+
+  // --- the bash tool ---------------------------------------------------------------
+
+  const base = createBashToolDefinition(process.cwd());
+  const BashParams = Type.Object({
+    command: Type.String({ description: "Shell command to execute" }),
+    timeout: Type.Optional(Type.Number({
+      description: `Seconds to wait for the command in the foreground (default ${fgDefault}, at most ${MAX_FOREGROUND_SECONDS}). A command still running then is NOT killed: it moves to the background, you get its id and the output so far, and its result is delivered to you when it finishes`,
+    })),
+    run_in_background: Type.Optional(Type.Boolean({
+      description: "true: start the command in the background and return at once with its id (dev servers, watchers, long builds or test suites while you do other work). Its result is delivered to you as a message when it finishes",
+    })),
+  });
+
+  pi.registerTool({
+    ...base,
+    name: "bash",
+    label: "bash",
+    description: [
+      "Execute a bash command in the current working directory. Returns stdout and stderr (output is truncated to the last lines; full output is saved to a temp file).",
+      `A command that runs longer than its foreground window (timeout, default ${fgDefault} s) is not killed: it continues in the background and its result arrives later as a message.`,
+      "run_in_background: true starts it in the background right away.",
+      "bash_output reads a background command's new output and status, bash_wait waits for it, bash_kill stops it.",
+    ].join(" "),
+    promptSnippet: "Execute bash commands (ls, grep, find, etc.); long-running commands continue in the background",
+    promptGuidelines: [
+      ...(Array.isArray(base.promptGuidelines) ? base.promptGuidelines : []),
+      "Long-running commands (dev servers, watchers, long builds or test suites) you do not need to wait for: bash with run_in_background: true, then keep working; their result arrives as a message when they finish, so do not poll bash_output in a loop.",
+      "Stop background commands you no longer need (servers, watchers) with bash_kill before you finish.",
+    ],
+    parameters: BashParams,
+
+    async execute(toolCallId: string, params: any, signal: AbortSignal | undefined, onUpdate: any, ctx: any) {
+      const session = hub.sessionOf(ctx) || "session";
+      const explicit = params?.run_in_background === true;
+      const windowMs = explicit ? 0 : foregroundSeconds(params?.timeout, fgDefault) * 1000;
+      let job: ShellJob | null = null;
+
+      const operations = {
+        exec: (command: string, cwd: string, execOptions: any) =>
+          new Promise<{ exitCode: number | null }>((resolve, reject) => {
+            const current = manager.start({ session, command, cwd, env: execOptions?.env, toolCallId, onData: execOptions?.onData });
+            job = current;
+            let timer: ReturnType<typeof setTimeout> | null = null;
+            const onAbort = () => { void manager.kill([current]); };
+            const cleanup = () => {
+              if (timer) clearTimeout(timer);
+              execOptions?.signal?.removeEventListener?.("abort", onAbort);
+            };
+            const moveToBackground = () => {
+              if (current.status !== "running") return;
+              cleanup();
+              manager.toBackground(current, explicit);
+              reject(new MovedToBackground());
+            };
+            if (execOptions?.signal) {
+              if (execOptions.signal.aborted) onAbort();
+              else execOptions.signal.addEventListener("abort", onAbort, { once: true });
+            }
+            if (windowMs <= 0) moveToBackground();
+            else timer = setTimeout(moveToBackground, windowMs);
+            void current.done.then(() => {
+              if (current.background) return;
+              cleanup();
+              // Its output reaches the model through this tool result.
+              current.delivered = true;
+              if (current.status === "killed") reject(new Error("aborted"));
+              else if (current.status === "failed") reject(new Error(current.error ?? "command failed"));
+              else resolve({ exitCode: current.exitCode });
+            });
+          }),
+      };
+
+      const definition = createBashToolDefinition(ctx?.cwd ? String(ctx.cwd) : process.cwd(), { operations: operations as any });
+      try {
+        return await definition.execute(toolCallId, { command: params?.command }, signal, onUpdate, ctx);
+      } catch (error) {
+        const moved = job as ShellJob | null;
+        if (!(error instanceof MovedToBackground) || !moved) throw error;
+        const hint = "Its result arrives as a message in this conversation when it finishes (no need to poll). bash_output reads its new output, bash_wait waits for it, bash_kill stops it.";
+        let text: string;
+        if (explicit) {
+          text = `Started in the background as ${moved.id}: ${shortCommand(moved.command, 200)}\n${hint}`;
+        } else {
+          const soFar = tailText(moved.output.since(0).text);
+          text = [
+            `Still running after ${Math.round(windowMs / 1000)} s — moved to the background as ${moved.id} (not killed).`,
+            hint,
+            "",
+            "Output so far:",
+            `${soFar.truncated ? "[earlier output omitted]\n" : ""}${soFar.text.trim() ? soFar.text.replace(/\s+$/, "") : "(no output yet)"}`,
+          ].join("\n");
+        }
+        return {
+          content: [{ type: "text", text }],
+          details: { background: { id: moved.id, command: moved.command, explicit, started_at: moved.startedAt } },
+        };
+      }
+    },
+  });
+
+  // --- companion tools -----------------------------------------------------------
+
+  const IdsParam = Type.Optional(Type.Array(Type.String(), {
+    description: 'Background command ids (e.g. "sh2"); omit (or "all") for every background command of this session',
+  }));
+
+  const readResult = (jobs: ShellJob[], unknown: string[]) => {
+    const parts = collect(jobs);
+    const sections = parts.map(({ job, text, dropped }) => renderJobResult(job, text, { dropped }).text);
+    if (unknown.length) sections.push(`Unknown ids: ${unknown.join(", ")}`);
+    return {
+      content: [{ type: "text", text: sections.join("\n\n---\n\n") || "No background commands in this session." }],
+      details: {
+        jobs: parts.map(({ job, text, dropped }) => {
+          const shown = tailText(text);
+          return jobDetails(job, shown.text, shown.truncated || dropped > 0);
+        }),
+      },
+    };
+  };
+
+  pi.registerTool({
+    name: "bash_output",
+    label: "bash output",
+    description: "A background command's output since you last read it, and its status (running, or its exit code). Without an id: the status of every background command of this session. Results are delivered automatically when a command finishes — use this to check progress, not to poll in a loop.",
+    promptSnippet: "bash_output: new output and status of a background command",
+    parameters: Type.Object({
+      id: Type.Optional(Type.String({ description: 'Background command id (e.g. "sh2"); omit for the status of all' })),
+    }),
+    async execute(_toolCallId: string, params: any, _signal: AbortSignal | undefined, _onUpdate: any, ctx: any) {
+      const session = hub.sessionOf(ctx) || "session";
+      const id = typeof params?.id === "string" ? params.id.trim() : "";
+      if (!id || id === "all") {
+        const jobs = manager.resolve(session, []).jobs;
+        return {
+          content: [{ type: "text", text: renderJobStatus(jobs) }],
+          details: { jobs: jobs.map((job) => ({ id: job.id, command: job.command, status: job.status, exit_code: job.exitCode })) },
+        };
+      }
+      const { jobs, unknown } = manager.resolve(session, [id]);
+      return readResult(jobs, unknown);
+    },
+  });
+
+  pi.registerTool({
+    name: "bash_wait",
+    label: "bash wait",
+    description: "Wait until background commands finish (default: every running one) and return their new output and exit codes. timeout_seconds bounds the WAIT only — the commands keep running past it.",
+    promptSnippet: "bash_wait: wait for background commands to finish",
+    parameters: Type.Object({
+      ids: IdsParam,
+      timeout_seconds: Type.Optional(Type.Number({ description: `Leave out (or 0) to wait until they finish (at most ${MAX_WAIT_SECONDS} s per call); otherwise stop WAITING after this many seconds` })),
+    }),
+    async execute(_toolCallId: string, params: any, signal: AbortSignal | undefined, onUpdate: any, ctx: any) {
+      const session = hub.sessionOf(ctx) || "session";
+      const explicit = Array.isArray(params?.ids) && params.ids.length > 0 && !params.ids.includes("all");
+      const resolved = manager.resolve(session, params?.ids);
+      const jobs = explicit ? resolved.jobs : resolved.jobs.filter((job) => job.status === "running" || !job.delivered);
+      if (jobs.length === 0 && resolved.unknown.length === 0) {
+        return { content: [{ type: "text", text: "No background commands to wait for." }], details: { jobs: [] } };
+      }
+      const seconds = typeof params?.timeout_seconds === "number" && params.timeout_seconds > 0
+        ? Math.min(params.timeout_seconds, MAX_WAIT_SECONDS)
+        : MAX_WAIT_SECONDS;
+      onUpdate?.({ content: [{ type: "text", text: `bash_wait: waiting for ${jobs.filter((job) => job.status === "running").map((job) => job.id).join(", ") || "nothing"}` }], details: {} });
+      await manager.wait(jobs, { timeoutMs: seconds * 1000, ...(signal ? { signal } : {}) });
+      return readResult(jobs, resolved.unknown);
+    },
+  });
+
+  pi.registerTool({
+    name: "bash_kill",
+    label: "bash kill",
+    description: 'Stop background commands (the whole process tree) by id, or ["all"]; returns their last output.',
+    promptSnippet: "bash_kill: stop background commands",
+    parameters: Type.Object({
+      ids: Type.Array(Type.String(), { description: 'Background command ids, or ["all"]', minItems: 1 }),
+    }),
+    async execute(_toolCallId: string, params: any, _signal: AbortSignal | undefined, _onUpdate: any, ctx: any) {
+      const session = hub.sessionOf(ctx) || "session";
+      const { jobs, unknown } = manager.resolve(session, params?.ids);
+      const live = jobs.filter((job) => job.status === "running");
+      await manager.kill(live);
+      if (live.length === 0) {
+        const done = jobs.map((job) => `${job.id} already ${jobOutcome(job)}`);
+        return {
+          content: [{ type: "text", text: `Nothing to stop${done.length ? ` (${done.join(", ")})` : ""}${unknown.length ? `; unknown ids: ${unknown.join(", ")}` : ""}.` }],
+          details: { jobs: [] },
+        };
+      }
+      return readResult(live, unknown);
+    },
+  });
+
+  pi.registerCommand("jobs", {
+    description: "Background shell commands of this session; /jobs kill <id|all> stops them",
+    handler: async (args: string, ctx: any) => {
+      const session = hub.sessionOf(ctx) || "session";
+      const words = String(args ?? "").trim().split(/\s+/).filter(Boolean);
+      if (words[0] === "kill") {
+        const { jobs } = manager.resolve(session, words.slice(1).length ? words.slice(1) : ["all"]);
+        const stopped = await manager.kill(jobs.filter((job) => job.status === "running"), { deliver: true });
+        ctx?.ui?.notify?.(stopped.length ? `Stopped ${stopped.map((job) => job.id).join(", ")}.` : "No running background commands.");
+        return;
+      }
+      ctx?.ui?.notify?.(renderJobStatus(manager.resolve(session, []).jobs));
+    },
+  });
+}

package/assets/extensions/omnirush/collector.ts

@@ -58,6 +58,7 @@ import {
 } from "./pi-engine";
 import { sharedRefresher } from "./refresh";
 import { AGENTS_HOOK, type AgentsHook } from "./agents-lib";
+import { BGSHELL_HOOK, type BgShellHook } from "./bgshell-lib";
 
 /** info.agent of the root session's messages (a sub-agent's is its role). */
 export const ROOT_AGENT = "omnirush-cli";
@@ -83,11 +84,12 @@ const SWITCH_BUDGET_MS = 30_000;
  */
 const PRINT_EXIT_GRACE_MS = 2_000;
 
-/** The pi agent config dir: session files live under <dir>/sessions. */
+/** The agent data dir: session files live under <dir>/sessions (the launcher always sets it). */
 function piAgentDir(): string {
-  const override = process.env.PI_CODING_AGENT_DIR || process.env.OMNIRUSH_CODING_AGENT_DIR || "";
+  const override = process.env.OMNIRUSH_CODING_AGENT_DIR || process.env.PI_CODING_AGENT_DIR || "";
   if (override.trim()) return path.resolve(override.trim());
-  return path.join(os.homedir(), ".pi", "agent");
+  const stateDir = (process.env.OMNIRUSH_DIR || "").trim();
+  return path.join(stateDir ? path.resolve(stateDir) : path.join(os.homedir(), ".omnirush"), "agent");
 }
 
 /** pi session ids are UUIDs; the collector refuses anything outside its id alphabet. */
@@ -124,26 +126,68 @@ function removeLegacyJournals(stateDir: string): void {
   }
 }
 
-/** Sub-agents a session started: its spawn_agents results (session id, role, task). */
-export function spawnedChildren(entries: readonly PiEntry[]): Array<{ id: string; role: string | null; task: string | null }> {
-  const children: Array<{ id: string; role: string | null; task: string | null }> = [];
+/** A sub-agent that ran on the main model instead of the picked one (spawn_agents' model_fallback). */
+export type ChildModelFallback = { requested: string; used: string; effort: string | null; reason: string; kind: string; note: string };
+
+/** Tools and messages that report sub-agent results after spawn_agents named them. */
+const CHILD_RESULT_TOOLS = new Set(["agents_wait", "agents_result", "agents_cancel"]);
+const CHILD_RESULT_MESSAGE = "omnirush-agents-result";
+
+function modelFallbackOf(result: any): ChildModelFallback | null {
+  const fallback = result && typeof result.model_fallback === "object" ? result.model_fallback : null;
+  if (!fallback || typeof fallback.requested !== "string" || typeof fallback.used !== "string") return null;
+  return {
+    requested: fallback.requested,
+    used: fallback.used,
+    effort: typeof fallback.effort === "string" ? fallback.effort : null,
+    reason: typeof fallback.reason === "string" ? fallback.reason : "refused",
+    kind: typeof fallback.kind === "string" ? fallback.kind : "selection",
+    note: typeof fallback.note === "string" ? fallback.note : `ran on ${fallback.used}`,
+  };
+}
+
+/**
+ * Sub-agents a session started: its spawn_agents results (session id, role,
+ * task), with the model fallback a later result reported (a background
+ * child's delivery, agents_wait / agents_result) when there is one.
+ */
+export function spawnedChildren(entries: readonly PiEntry[]): Array<{ id: string; role: string | null; task: string | null; modelFallback?: ChildModelFallback }> {
+  const children: Array<{ id: string; role: string | null; task: string | null; modelFallback?: ChildModelFallback }> = [];
   for (const entry of entries) {
     const message = entry.type === "message" ? entry.message : null;
-    if (!message || message.role !== "toolResult" || message.toolName !== "spawn_agents") continue;
-    const results = message.details && Array.isArray(message.details.results) ? message.details.results : [];
+    const spawned = Boolean(message && message.role === "toolResult" && message.toolName === "spawn_agents");
+    const details = spawned
+      ? message!.details
+      : message && message.role === "toolResult" && CHILD_RESULT_TOOLS.has(message.toolName)
+        ? message.details
+        : entry.type === "custom_message" && (entry as any).customType === CHILD_RESULT_MESSAGE
+          ? (entry as any).details
+          : null;
+    const results = details && Array.isArray(details.results) ? details.results : [];
     for (const result of results) {
       const id = result && typeof result.session_id === "string" ? result.session_id : "";
-      if (!sanitizeSessionId(id) || children.some((child) => child.id === id)) continue;
+      if (!sanitizeSessionId(id)) continue;
+      const known = children.find((child) => child.id === id);
+      const modelFallback = modelFallbackOf(result);
+      if (known) {
+        if (modelFallback) known.modelFallback = modelFallback;
+        continue;
+      }
+      if (!spawned) continue;
       children.push({
         id,
         role: typeof result.role === "string" ? result.role : null,
         task: typeof result.task === "string" ? result.task : null,
+        ...(modelFallback ? { modelFallback } : {}),
       });
     }
   }
   return children;
 }
 
+/** Fallbacks already recorded as trace events (root session : child session). */
+const recordedFallbacks = new Set<string>();
+
 /** The session file of `sessionId`: `<ts>_<id>.jsonl` in one of the session dirs. */
 async function findSessionFile(sessionId: string, dirs: readonly string[]): Promise<string | null> {
   for (const dir of dirs) {
@@ -233,13 +277,26 @@ export async function captureChildSessions(input: {
     if (history.omitted > 0) {
       input.collector.recordTrace(input.rootSessionId, "session.messages_omitted", { child_session_id: child.id, count: history.omitted });
     }
+    const fallbackKey = `${input.rootSessionId}:${child.id}`;
+    if (child.modelFallback && !recordedFallbacks.has(fallbackKey)) {
+      // As the desktop: the sub-agent ran on the main model, and why.
+      recordedFallbacks.add(fallbackKey);
+      input.collector.recordTrace(input.rootSessionId, "subagent.model_fallback", {
+        child_session_id: child.id,
+        requested: child.modelFallback.requested,
+        used: child.modelFallback.used,
+        effort: child.modelFallback.effort,
+        reason: child.modelFallback.reason,
+        kind: child.modelFallback.kind,
+      });
+    }
     if (history.delta.length > 0 || !input.known.has(child.id)) {
       input.beforeRecord(history.sizes.reduce((total, bytes) => total + bytes, 0));
       input.collector.recordChildSession(input.rootSessionId, {
         childSessionId: child.id,
         parentSessionId: input.parentSessionId,
         depth: input.depth,
-        title: child.task,
+        title: child.task && child.modelFallback ? `${child.task} (${child.modelFallback.note})` : child.task,
         agent: turnModelFromMessages(messages)?.agent ?? agent,
         messages: history.delta,
         lastMessageId: typeof messages.at(-1)?.info.id === "string" ? (messages.at(-1)!.info.id as string) : null,
@@ -255,7 +312,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();
@@ -275,26 +336,47 @@ 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
     // value is a bare release number); client_version / client_platform are
     // the backend's optional envelope metadata.
     appVersion: version ? `omnirush-cli/${version}` : "omnirush-cli",
-    ...(engineVersion ? { engineVersion: `pi/${engineVersion}` } : {}),
+    ...(engineVersion ? { engineVersion: `omnirush-core/${engineVersion}` } : {}),
     envelopeMetadata: {
       ...(version ? { client_version: version } : {}),
       client_platform: `${process.platform}/${process.arch}`,
@@ -309,6 +391,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],
@@ -501,6 +584,10 @@ export default function (pi: any) {
     const hook = (globalThis as any)[AGENTS_HOOK] as AgentsHook | undefined;
     const followed = session;
     const interrupted = followed?.started && hook ? await hook.interruptAll(followed.id).catch(() => []) : [];
+    // Background shell commands still running are stopped too; their output
+    // is added to the session as a message, so the last capture keeps it.
+    const shell = (globalThis as any)[BGSHELL_HOOK] as BgShellHook | undefined;
+    if (followed?.started && shell) await shell.stopAll(followed.id).catch(() => []);
     await settleRunning("interrupted");
     const state = session;
     if (!state?.started) return;

package/assets/extensions/omnirush/commands.ts

@@ -18,9 +18,10 @@
  */
 
 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 { formatModeStatus } from "./yolo-lib";
 import {
   formatDeviceLines,
   formatGrantLines,
@@ -42,27 +43,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) {
@@ -156,6 +154,7 @@ export default function (pi: ExtensionAPI) {
         collector: { sessionId, record },
         origin: resolveOrigin(process.env),
         version: CLI_VERSION,
+        yolo: formatModeStatus(),
       });
       notify(ctx, ["Omnirush status:", ...statusLines.map((line) => `  ${line}`)].join("\n"));
     },