omnirush 0.8.6 → 0.9.1

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,
@@ -53,7 +64,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 +145,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 +206,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",
@@ -194,6 +263,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 +305,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 +327,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 +363,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 +396,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 +449,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 () => {