omnirush 0.9.0 → 0.10.0

package/assets/extensions/omnirush/collector.ts

@@ -32,10 +32,20 @@
 // session reaches the parent's trace as session.child events, like the
 // desktop's task sub-agents.
 //
+// A turn's prompt is only what the user submitted. A sub-agent process
+// marks its own session (an `omnirush-subagent` entry naming its parent),
+// so when that session is later opened as a main session (`--continue`,
+// `--resume`, `--session`) its first user message, the task the parent
+// agent wrote, is recorded as synthetic; for sub-agents that wrote no
+// marker, the parent session file naming it in a spawn result confirms
+// it. A prompt an extension sends (pi's input event, source "extension")
+// is synthetic too.
+//
 // Silent no-op without device credentials: collection only runs for
 // signed-in omnirush identities.
 
 import { createHash } from "node:crypto";
+import markSubagentSession from "./subagent-marker";
 import { createReadStream, readdirSync, readFileSync, rmSync } from "node:fs";
 import { readdir, stat } from "node:fs/promises";
 import os from "node:os";
@@ -51,13 +61,16 @@ import {
   engineMessagesFromEntries,
   MAX_TURN_MESSAGE_BYTES,
   messageParts,
+  subagentTask,
+  SUBAGENT_ENTRY_TYPE,
   turnHistory,
   turnModelFromMessages,
+  userText,
   type EngineMessage,
   type PiEntry,
 } from "./pi-engine";
 import { sharedRefresher } from "./refresh";
-import { AGENTS_HOOK, type AgentsHook } from "./agents-lib";
+import { AGENTS_HOOK, SUBAGENT_SESSIONS_DIR, 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). */
@@ -203,6 +216,71 @@ async function findSessionFile(sessionId: string, dirs: readonly string[]): Prom
   return null;
 }
 
+/** Session files read at most when looking for a sub-agent's parent, and for how long. */
+const MAX_PARENT_SCAN_FILES = 64;
+const PARENT_SCAN_BUDGET_MS = 3_000;
+
+/** Whether `file` contains `needle` (streamed in chunks, never held whole). */
+async function fileContains(file: string, needle: string, deadline: number): Promise<boolean> {
+  const input = createReadStream(file, { encoding: "utf8", highWaterMark: 256 * 1024 });
+  let tail = "";
+  try {
+    for await (const chunk of input) {
+      const text = tail + (chunk as string);
+      if (text.includes(needle)) return true;
+      if (Date.now() > deadline) return false;
+      tail = text.slice(-needle.length);
+    }
+  } catch {
+    /* unreadable: not the parent */
+  } finally {
+    input.destroy();
+  }
+  return false;
+}
+
+/**
+ * The session that started sub-agent `childId`: the session file in `dir`
+ * whose spawn_agents (or agents_*) results name it, among the files
+ * written since the sub-agent started (`since`, epoch ms; the newest
+ * first, at most MAX_PARENT_SCAN_FILES, within PARENT_SCAN_BUDGET_MS).
+ */
+export async function findSpawningSession(childId: string, dir: string, since: number): Promise<string | null> {
+  const deadline = Date.now() + PARENT_SCAN_BUDGET_MS;
+  const found = await findSpawningSessionIn(childId, dir, since, deadline);
+  if (found) return found;
+  // A sub-agent session moved into <session dir>/subagents/ (src/sessions.js,
+  // or written there by a newer parent): its parent sits one level up.
+  if (path.basename(dir) === SUBAGENT_SESSIONS_DIR) {
+    return findSpawningSessionIn(childId, path.dirname(dir), since, deadline);
+  }
+  return null;
+}
+
+async function findSpawningSessionIn(childId: string, dir: string, since: number, deadline: number): Promise<string | null> {
+  let names: string[];
+  try {
+    names = await readdir(dir);
+  } catch {
+    return null;
+  }
+  const candidates: Array<{ file: string; id: string; mtime: number }> = [];
+  for (const name of names) {
+    const match = /_([A-Za-z0-9._:-]{8,128})\.jsonl$/.exec(name);
+    if (!match || match[1] === childId) continue;
+    const file = path.join(dir, name);
+    const stats = await stat(file).catch(() => null);
+    if (stats && stats.mtimeMs >= since) candidates.push({ file, id: match[1]!, mtime: stats.mtimeMs });
+  }
+  candidates.sort((a, b) => b.mtime - a.mtime);
+  const needle = `"session_id":${JSON.stringify(childId)}`;
+  for (const candidate of candidates.slice(0, MAX_PARENT_SCAN_FILES)) {
+    if (Date.now() > deadline) break;
+    if (await fileContains(candidate.file, needle, deadline)) return candidate.id;
+  }
+  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;
@@ -253,7 +331,8 @@ export async function captureChildSessions(input: {
   visited: Set<string>;
   /** Unchanged session files are not read again in this process. */
   seenFiles: Map<string, string>;
-  beforeRecord: (bytes: number) => void;
+  /** Before each event is recorded (may wait: the caller flushes and drains the upload queue). */
+  beforeRecord: (bytes: number) => void | Promise<void>;
 }): Promise<void> {
   if (input.depth > MAX_CHILD_DEPTH) return;
   for (const child of spawnedChildren(input.parentEntries).slice(0, MAX_CHILDREN)) {
@@ -291,7 +370,7 @@ export async function captureChildSessions(input: {
       });
     }
     if (history.delta.length > 0 || !input.known.has(child.id)) {
-      input.beforeRecord(history.sizes.reduce((total, bytes) => total + bytes, 0));
+      await input.beforeRecord(history.sizes.reduce((total, bytes) => total + bytes, 0));
       input.collector.recordChildSession(input.rootSessionId, {
         childSessionId: child.id,
         parentSessionId: input.parentSessionId,
@@ -309,7 +388,13 @@ export async function captureChildSessions(input: {
 
 export default function (pi: any) {
   // A sub-agent: the parent session captures it (session.child events).
-  if ((process.env.OMNIRUSH_PARENT_SESSION || "").trim()) return;
+  const parentSession = (process.env.OMNIRUSH_PARENT_SESSION || "").trim();
+  if (parentSession) {
+    // Its session says so (subagent-marker.ts; index.ts loads it for
+    // sub-agent processes, which do not load the collector).
+    markSubagentSession(pi, parentSession);
+    return;
+  }
   const origin = resolveOrigin(process.env);
   const gatewayUrl = process.env.OMNIRUSH_GATEWAY_URL || gatewayUrlForOrigin(origin);
   // The shared auth file is the source of truth (another process may have
@@ -429,21 +514,34 @@ export default function (pi: any) {
     pendingPrompts: Array<{ at: number; model: { providerID: string; modelID: string } | null; variant: string | null }>;
     /** User entry id -> what its prompt was sent with (pi-engine promptSettings). */
     promptSettings: Map<string, { model: { providerID: string; modelID: string } | null; variant: string | null }>;
+    /** Prompts an extension sent (pi's input event, source "extension"), until their user entry is found. */
+    extensionInputs: Array<{ at: number; text: string }>;
+    /** User entries the user did not submit -> their parts' metadata (pi-engine injectedPrompts). */
+    injectedPrompts: Map<string, Record<string, unknown>>;
+    /** Whether the session is a sub-agent's (its first prompt the parent agent's task), once looked up. */
+    origin: Promise<void> | null;
   } | null = null;
   let settling: Promise<void> = Promise.resolve();
 
+  // Where sub-agent session files are looked up: the session's own
+  // subagents/ folder first (where spawn_agents puts them, and where the
+  // launcher moves older releases' children), then the session dirs
+  // themselves (children of releases up to 0.9.0 not moved yet), then
+  // every project's.
   const sessionDirs = (ctx: any): string[] => {
     const dirs: string[] = [];
+    const add = (dir: unknown) => {
+      if (typeof dir !== "string" || !dir) return;
+      for (const candidate of [path.join(dir, SUBAGENT_SESSIONS_DIR), dir]) {
+        if (!dirs.includes(candidate)) dirs.push(candidate);
+      }
+    };
     const file = ctx?.sessionManager?.getSessionFile?.();
-    if (typeof file === "string" && file) dirs.push(path.dirname(file));
-    const sessionDir = ctx?.sessionManager?.getSessionDir?.();
-    if (typeof sessionDir === "string" && sessionDir && !dirs.includes(sessionDir)) dirs.push(sessionDir);
+    if (typeof file === "string" && file) add(path.dirname(file));
+    add(ctx?.sessionManager?.getSessionDir?.());
     try {
       const root = path.join(agentDir, "sessions");
-      for (const name of readdirSync(root)) {
-        const dir = path.join(root, name);
-        if (!dirs.includes(dir)) dirs.push(dir);
-      }
+      for (const name of readdirSync(root)) add(path.join(root, name));
     } catch {
       /* no sessions yet */
     }
@@ -463,7 +561,23 @@ export default function (pi: any) {
       // never wrote a user entry (handled by an extension) is let go.
       else if (pending.at > Date.now() - 60_000) state.pendingPrompts.push(pending);
     }
-    return engineMessagesFromEntries(entries, { sessionId: state.id, agent: ROOT_AGENT, cwd: state.root, toolStarts: state.toolStarts, promptSettings: state.promptSettings });
+    // An extension's prompt is the first user entry with its text written at or after it went out.
+    for (const input of state.extensionInputs.splice(0)) {
+      const entry = entries.find((candidate: any) =>
+        candidate?.type === "message" && candidate.message?.role === "user" && typeof candidate.id === "string"
+        && !state.injectedPrompts.has(candidate.id) && userText(candidate.message) === input.text
+        && Number(candidate.message.timestamp ?? Date.parse(candidate.timestamp)) >= input.at - 1_000);
+      if (entry) state.injectedPrompts.set(entry.id, { source: "extension" });
+      else if (input.at > Date.now() - 60_000) state.extensionInputs.push(input);
+    }
+    return engineMessagesFromEntries(entries, {
+      sessionId: state.id,
+      agent: ROOT_AGENT,
+      cwd: state.root,
+      toolStarts: state.toolStarts,
+      promptSettings: state.promptSettings,
+      injectedPrompts: state.injectedPrompts,
+    });
   };
 
   /**
@@ -476,7 +590,21 @@ export default function (pi: any) {
     if (!id) return null;
     if (session?.id !== id) {
       const root = ctx?.cwd || process.cwd();
-      session = { id, root, ctx, started: false, running: false, lastMessageId: undefined, toolStarts: new Map(), seenFiles: new Map(), pendingPrompts: [], promptSettings: new Map() };
+      session = {
+        id,
+        root,
+        ctx,
+        started: false,
+        running: false,
+        lastMessageId: undefined,
+        toolStarts: new Map(),
+        seenFiles: new Map(),
+        pendingPrompts: [],
+        promptSettings: new Map(),
+        extensionInputs: [],
+        injectedPrompts: new Map(),
+        origin: null,
+      };
     }
     session.ctx = ctx;
     // Idempotent: a session already started in this process keeps its segment.
@@ -492,7 +620,7 @@ export default function (pi: any) {
    * 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 captureChildren = async (state: NonNullable<typeof session>, room: (bytes: number) => void | Promise<void>): Promise<void> => {
     const { id } = state;
     try {
       await captureChildSessions({
@@ -513,6 +641,26 @@ export default function (pi: any) {
     }
   };
 
+  /**
+   * Room for the next sub-agent event: once the recorded messages would pass
+   * one flush (MAX_TURN_MESSAGE_BYTES), they flush, and the next child is
+   * read only after that upload went out. A turn with many sub-agents (20
+   * children, MiBs each) used to read and queue all of them at once, holding
+   * every child's messages, their trace documents and envelopes in memory
+   * together.
+   */
+  const childRoom = (id: string) => {
+    let pending = 0;
+    return async (bytes: number): Promise<void> => {
+      if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
+        collector.flushTrace(id);
+        pending = 0;
+        await collector.drained(id);
+      }
+      pending += bytes;
+    };
+  };
+
   /**
    * Everything a settled run records, in the desktop observer's order:
    * messages, sub-agents, the idle event, the snapshot, the trace and the
@@ -540,22 +688,17 @@ export default function (pi: any) {
       for (const [index, part] of parts.slice(0, -1).entries()) {
         collector.recordTrace(id, "turn.messages", { messages: part.messages, part: index + 1, parts: parts.length });
         collector.flushTrace(id);
+        // One part waits in memory at a time (a resumed session's backlog is up to 32 MiB).
+        await collector.drained(id);
       }
     } catch (error) {
       collector.recordTrace(id, "session.messages_failed", { error: error instanceof Error ? error.message : "unknown" });
     }
     const newest = parts.at(-1)!;
     // Sub-agent messages flush on their own before they would crowd the turn's out.
-    let pending = 0;
-    const room = (bytes: number) => {
-      if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
-        collector.flushTrace(id);
-        pending = 0;
-      }
-      pending += bytes;
-    };
+    const room = childRoom(id);
     await captureChildren(state, room);
-    room(newest.bytes);
+    await room(newest.bytes);
     collector.recordTrace(id, "session.idle", { status });
     // The turn snapshot runs first so the artifacts it discovers are part of
     // the trace flushed right behind it.
@@ -598,14 +741,7 @@ export default function (pi: any) {
     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;
-        });
+        await captureChildren(state, childRoom(state.id));
         collector.flushTrace(state.id);
       }).catch(() => undefined);
       await settling;
@@ -649,11 +785,45 @@ export default function (pi: any) {
     clearTimeout(budget);
   };
 
+  /**
+   * A session opened as a main session that is a sub-agent's own (its file
+   * sits with the main sessions, so `--continue` or the resume picker can
+   * reopen it): its first prompt is the task the parent agent wrote, not
+   * something the user typed. Looked up once, before the first prompt of
+   * this process is written (every user entry then on the branch came
+   * before this process).
+   */
+  const lookUpOrigin = (state: NonNullable<typeof session>): Promise<void> =>
+    (state.origin ??= (async () => {
+      const entries: PiEntry[] = state.ctx?.sessionManager?.getBranch?.() ?? [];
+      const task = subagentTask(entries);
+      if (!task) return;
+      let parent = task.parentSessionId;
+      if (!task.marked) {
+        // Sub-agents that wrote no marker: the parent session names it.
+        const file = state.ctx?.sessionManager?.getSessionFile?.();
+        const started = Date.parse(String(state.ctx?.sessionManager?.getHeader?.()?.timestamp ?? entries[0]?.timestamp ?? ""));
+        parent = typeof file === "string" && file
+          ? await findSpawningSession(state.id, path.dirname(file), Number.isFinite(started) ? started - 10 * 60_000 : 0).catch(() => null)
+          : null;
+        if (!parent) return;
+      }
+      state.injectedPrompts.set(task.entryId, { subagentTask: true, ...(parent ? { parentSessionID: parent } : {}) });
+    })().catch(() => undefined));
+
   pi.on("session_start", (_event: any, ctx: any) => {
     follow(ctx, false);
   });
 
-  pi.on("before_agent_start", (event: any, ctx: any) => {
+  // What the user submits comes as input; an extension's prompt
+  // (pi.sendUserMessage) is marked so it never reads as the user's.
+  pi.on("input", (event: any, ctx: any) => {
+    if (event?.source !== "extension" || typeof event?.text !== "string") return;
+    const state = follow(ctx, false);
+    state?.extensionInputs.push({ at: Date.now(), text: event.text });
+  });
+
+  pi.on("before_agent_start", async (event: any, ctx: any) => {
     const state = follow(ctx, true);
     if (!state) return;
     state.running = true;
@@ -710,6 +880,7 @@ export default function (pi: any) {
       },
     });
     archive.turnFollowed(state.id);
+    await lookUpOrigin(state);
   });
 
   // A run that no prompt started (a background sub-agent's results

package/assets/extensions/omnirush/index.ts

@@ -5,29 +5,43 @@
 // extension (without it, none of these files were ever registered).
 // EVERY module that registers commands/tools/handlers must be wired here:
 // a module left out of this list is silently dead in sessions.
+//
+// A sub-agent (a headless spawn_agents child: OMNIRUSH_PARENT_SESSION is
+// set) does not load what only the interactive top agent uses — the slash
+// commands (commands), the session capture (collector: the parent records
+// its children) and voice input (TUI only) — so each of the many child
+// processes a fan-out starts stays smaller. They are imported on demand,
+// in the same order as before, so the top agent registers exactly what it
+// did.
 
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import collector from "./collector";
-import commands from "./commands";
-import mcp from "./mcp";
 import recorder from "./recorder";
 import sota from "./sota";
+import mcp from "./mcp";
 import agents from "./agents";
 import bgshell from "./bgshell";
 import subagents from "./subagents";
 import swarm from "./swarm";
 import plan from "./plan";
-import voice from "./voice";
 import webfetch from "./webfetch";
 import websearch from "./websearch";
 import guard from "./guard";
 import yolo from "./yolo";
 
-export default function (pi: ExtensionAPI) {
-  commands(pi as any);
+/** A spawn_agents child process (the parent session captures it; no TUI). */
+export function isSubagentProcess(env: NodeJS.ProcessEnv = process.env): boolean {
+  return Boolean(String(env.OMNIRUSH_PARENT_SESSION ?? "").trim());
+}
+
+export default async function (pi: ExtensionAPI) {
+  const subagent = isSubagentProcess();
+  if (!subagent) (await import("./commands")).default(pi as any);
   recorder(pi as any);
   sota(pi as any);
-  collector(pi as any);
+  if (!subagent) (await import("./collector")).default(pi as any);
+  // Children still mark their own session as a sub-agent's (the collector
+  // does that for the top agent's sessions; children do not load it).
+  else (await import("./subagent-marker")).default(pi as any);
   mcp(pi as any);
   webfetch(pi as any);
   websearch(pi as any);
@@ -35,7 +49,7 @@ export default function (pi: ExtensionAPI) {
   agents(pi as any);
   subagents(pi as any);
   bgshell(pi as any);
-  voice(pi as any);
+  if (!subagent) (await import("./voice")).default(pi as any);
   yolo(pi as any);
   guard(pi as any);
   swarm(pi as any);

package/assets/extensions/omnirush/mcp.ts

@@ -18,10 +18,21 @@
 // The MCP SDK is imported from the committed sibling bundle
 // (mcp-client.bundle.js): pi copies this directory into
 // ~/.pi/agent/extensions/omnirush/ and jiti there resolves only sibling
-// files, never npm packages.
+// files, never npm packages. It is loaded on first connect, so a process
+// without MCP servers never pays for it.
+//
+// Sub-agents (spawn_agents children) do not start every server up front:
+// each child used to spawn its own copy of every configured stdio server
+// at session start (20 children x 3 servers = 60 extra processes). A child
+// registers the tools from the list its top agent saved
+// (<omnirush dir>/mcp-tools.json, per server config) and connects a server
+// only when it first calls one of its tools; without a saved list it
+// connects at start as before.
 
+import { createHash, randomUUID } from "node:crypto";
+import { mkdirSync, readFileSync, renameSync, writeFileSync } from "node:fs";
 import path from "node:path";
-import { Client, McpError, StdioClientTransport, StreamableHTTPClientTransport } from "./mcp-client.bundle.js";
+import type { Client } from "./mcp-client.bundle.js";
 import { omniDir } from "./auth";
 import {
   dedupeToolNames,
@@ -32,6 +43,7 @@ import {
   type McpServerConfig,
 } from "./mcp-lib";
 import { truncateTail } from "@earendil-works/pi-coding-agent";
+import { sanitizeToolEnvironment } from "./secret-env";
 
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 
@@ -53,7 +65,62 @@ const SHUTDOWN_BUDGET_MS = 3_000;
 // Stderr lines kept per stdio server for failure diagnostics.
 const STDERR_TAIL_LINES = 5;
 
-type ServerStatus = "connecting" | "connected" | "failed";
+/** "deferred": tools registered from the saved list, connects on first use (sub-agents). */
+type ServerStatus = "connecting" | "connected" | "failed" | "deferred";
+
+type McpSdk = typeof import("./mcp-client.bundle.js");
+let sdkPromise: Promise<McpSdk> | null = null;
+let loadedSdk: McpSdk | null = null;
+/** The MCP SDK bundle, imported on first use. */
+function loadSdk(): Promise<McpSdk> {
+  sdkPromise ??= import("./mcp-client.bundle.js").then((sdk) => {
+    loadedSdk = sdk;
+    return sdk;
+  });
+  return sdkPromise;
+}
+
+/** Saved tool lists (by the process that connected), keyed by server name. */
+export const MCP_TOOLS_CACHE_FILE = "mcp-tools.json";
+/** A saved list older than this is not trusted (the server may have changed). */
+const TOOLS_CACHE_MAX_AGE_MS = 7 * 24 * 3600_000;
+
+/** Identity of a server config: a saved tool list only counts for the same config. */
+export function serverConfigKey(config: McpServerConfig): string {
+  return createHash("sha256").update(JSON.stringify(config)).digest("hex").slice(0, 32);
+}
+
+function readToolsCache(dir: string): Record<string, { key: string; at: number; tools: McpToolInfo[] }> {
+  try {
+    const parsed = JSON.parse(readFileSync(path.join(dir, MCP_TOOLS_CACHE_FILE), "utf8"));
+    return parsed && typeof parsed === "object" && !Array.isArray(parsed) ? parsed : {};
+  } catch {
+    return {};
+  }
+}
+
+/** Save one server's tool list (atomic; other servers' entries kept). */
+function saveToolsCache(dir: string, state: ServerState): void {
+  try {
+    const cache = readToolsCache(dir);
+    cache[state.name] = { key: serverConfigKey(state.config), at: Date.now(), tools: state.tools };
+    mkdirSync(dir, { recursive: true, mode: 0o700 });
+    const file = path.join(dir, MCP_TOOLS_CACHE_FILE);
+    const tmp = `${file}.${randomUUID()}.tmp`;
+    writeFileSync(tmp, JSON.stringify(cache), { mode: 0o600 });
+    renameSync(tmp, file);
+  } catch {
+    /* only an optimization for sub-agents */
+  }
+}
+
+/** A saved tool list for this server config, or null. */
+function cachedTools(dir: string, state: ServerState, now = Date.now()): McpToolInfo[] | null {
+  const entry = readToolsCache(dir)[state.name];
+  if (!entry || entry.key !== serverConfigKey(state.config) || !Array.isArray(entry.tools)) return null;
+  if (!(now - Number(entry.at) < TOOLS_CACHE_MAX_AGE_MS)) return null;
+  return entry.tools.filter((tool) => tool && typeof tool.name === "string");
+}
 
 /** OMNIRUSH_DEBUG=1 diagnostics (same convention as the sota guard). */
 function debug(line: string): void {
@@ -79,6 +146,8 @@ interface ServerState {
   tools: McpToolInfo[];
   client?: Client;
   stderrTail: string[];
+  /** A deferred server's connect in flight (single-flight for parallel tool calls). */
+  connecting?: Promise<void>;
 }
 
 function withTimeout<T>(promise: Promise<T>, ms: number, label: string): Promise<T> {
@@ -138,6 +207,7 @@ function contentText(content: Array<Record<string, any>>): string {
 
 /** Connect one server (bounded). Throws on failure — caller isolates. */
 async function connectServer(state: ServerState): Promise<void> {
+  const { Client, StdioClientTransport, StreamableHTTPClientTransport } = await loadSdk();
   const client = new Client({
     name: "omnirush",
     version: (process.env.OMNIRUSH_VERSION || "").trim() || "0.0.0",
@@ -149,7 +219,7 @@ async function connectServer(state: ServerState): Promise<void> {
     transport = new StdioClientTransport({
       command: state.config.command!,
       args: state.config.args ?? [],
-      ...(state.config.env ? { env: state.config.env } : {}),
+      env: sanitizeToolEnvironment({ ...process.env, ...(state.config.env ?? {}) }),
       stderr: "pipe",
     });
   }
@@ -194,6 +264,11 @@ async function connectServer(state: ServerState): Promise<void> {
   }
 }
 
+/** A spawn_agents child process. */
+function isSubagent(): boolean {
+  return Boolean(String(process.env.OMNIRUSH_PARENT_SESSION ?? "").trim());
+}
+
 function describeServer(state: ServerState): string {
   if (state.config.url) return `${state.name} (${state.config.url})`;
   const args = (state.config.args ?? []).join(" ");
@@ -231,6 +306,11 @@ export default function (pi: any) {
           void toolCallId;
           void onUpdate;
           const current = servers.get(state.name);
+          if (current?.status === "deferred") {
+            // A sub-agent's first call to this server: connect now.
+            current.connecting ??= connectServer(current).finally(() => { current.connecting = undefined; });
+            await current.connecting.catch(() => undefined);
+          }
           const client = current?.client;
           if (current?.status !== "connected" || !client) {
             throw new Error(
@@ -248,7 +328,7 @@ export default function (pi: any) {
               { timeout: CALL_TIMEOUT_MS, signal },
             );
           } catch (error: any) {
-            const detail = error instanceof McpError ? error.message : error?.message ?? String(error);
+            const detail = loadedSdk && error instanceof loadedSdk.McpError ? error.message : error?.message ?? String(error);
             throw new Error(`MCP call failed on "${state.name}": ${detail}`);
           }
           const content = mcpContentToPiContent(result?.content);
@@ -284,6 +364,7 @@ export default function (pi: any) {
         try {
           await connectServer(state);
           registerServerTools(state);
+          saveToolsCache(omniDir(), state);
         } catch {
           /* connectServer recorded the failure; isolation keeps going */
         }
@@ -316,7 +397,9 @@ export default function (pi: any) {
       const stateLabel =
         state.status === "connected"
           ? `connected, ${state.tools.length} tool${state.tools.length === 1 ? "" : "s"}`
-          : state.status === "connecting"
+          : state.status === "deferred"
+            ? `connects on first use, ${state.tools.length} tool${state.tools.length === 1 ? "" : "s"}`
+            : state.status === "connecting"
             ? "connecting…"
             : `failed: ${state.error ?? "unknown error"}`;
       lines.push(`  ${describeServer(state)} — ${stateLabel}`);
@@ -367,10 +450,23 @@ export default function (pi: any) {
     for (const error of errors) {
       ctx?.ui?.notify?.(`mcp config: ${error}`, "warning");
     }
-    if (servers.size > 0) {
-      // Isolated + bounded per server; failures surface via /mcp only.
-      await connectAll();
+    if (servers.size === 0) return;
+    if (isSubagent()) {
+      // Tools from the saved list; the server starts on first use.
+      const dir = omniDir();
+      for (const state of servers.values()) {
+        const tools = cachedTools(dir, state);
+        if (!tools) continue;
+        state.tools = tools;
+        state.status = "deferred";
+        registerServerTools(state);
+      }
+      const rest = [...servers.values()].filter((state) => state.status !== "deferred");
+      await Promise.allSettled(rest.map((state) => connectAll(state.name)));
+      return;
     }
+    // Isolated + bounded per server; failures surface via /mcp only.
+    await connectAll();
   });
 
   pi.on("session_shutdown", async () => {