omnirush 0.8.6 → 0.9.1

package/assets/extensions/omnirush/agents.ts

@@ -4,7 +4,11 @@
 // the SAME workspace:
 //   - role presets: code-searcher (read-only exploration),
 //     researcher-web (web_fetch research), general-worker (full tools)
-//   - parallel execution, every task at once unless max_parallel throttles
+//   - parallel execution, every task at once unless max_parallel throttles;
+//     no count cap, but memory-aware admission (memory-lib.ts): a child
+//     starts while the machine keeps enough RAM free and the agents' process
+//     tree stays under its budget, the rest wait ("queued: waiting for
+//     memory") and start as memory frees up
 //   - no wall-clock limit: a child runs until it is done. Only an explicit
 //     timeout_minutes, the inactivity watchdog (no output at all for
 //     OMNIRUSH_AGENT_INACTIVITY_MINUTES, default 30, longer while one of
@@ -20,7 +24,8 @@
 //     running children
 //
 // Tracing is free: every child writes its own session JSONL under
-// <agentDir>/sessions/<cwd-slug>/ (children run without --no-session)
+// <session dir>/subagents/ (children run without --no-session; that
+// folder keeps them out of /resume and --continue)
 // under a session id the parent picks, and the parent session's
 // collector records each child's conversation as a sub-agent of the
 // turn (session.child events with parent and depth, as the desktop
@@ -29,6 +34,11 @@
 // stopped and recorded as interrupted. Children get
 // OMNIRUSH_PARENT_SESSION, so a child process uploads nothing of its own.
 //
+// Nesting: a sub-agent may start sub-agents of its own, up to
+// MAX_SUBAGENT_DEPTH (3) layers below the main session (the desktop's
+// OMNIRUSH_SUBAGENT_DEPTH). Each child gets OMNIRUSH_AGENT_DEPTH (its layer)
+// and a note saying it; layer 3 has no spawn_agents.
+//
 // `/agents` lists this session's subagents (running + finished);
 // `/agents cancel <id|all>` stops running ones.
 
@@ -36,6 +46,8 @@ import { StringEnum } from "@earendil-works/pi-ai";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { Type } from "typebox";
 
+import { deliveryHub } from "./deliveries";
+import { subagentTasks } from "./subagents";
 import {
   AGENT_ROLES,
   AGENTS_HOOK,
@@ -44,15 +56,25 @@ import {
   type AgentRunner,
   type AgentsHook,
   type AgentRecord,
+  agentDepth,
+  canDelegate,
   childInactivityMs,
   renderAgentStatus,
+  subagentLayerNote,
   renderChildResults,
   renderDelivery,
-  modelSelectionFromContext,
   runChildAgent,
+  childModelFallback,
+  subagentSessionDir,
   type ChildResult,
   type ChildTask,
 } from "./agents-lib";
+import { SWARM_GUIDELINE } from "./swarm-lib";
+import { omniDir } from "./auth";
+import { MemoryAdmission, readMemorySettings } from "./memory-lib";
+
+/** Blocking spawn_agents: progress updates at most this often (ms). */
+const PROGRESS_UPDATE_MS = 500;
 
 /** A timeout_minutes below this is taken for a placeholder, not a limit. */
 const MIN_TIMEOUT_MINUTES = 1;
@@ -68,7 +90,7 @@ const SpawnAgentsParams = Type.Object({
       }),
       task: Type.String({ description: "Self-contained task description for the child agent (it sees ONLY this, not your conversation)" }),
       model: Type.Optional(Type.String({
-        description: 'Optional model[:effort] for THIS task (e.g. "muse-spark-1.3", "gpt-6-sol:high"). Omit (or leave empty) to inherit the parent session\'s current model and effort',
+        description: 'Optional model[:effort] for THIS task (e.g. "muse-spark-1.3", "muse-spark-1.2-contributor" for cheap swarm workers, "meta-muse-spark", "muse-spark-1.1:low", "gpt-6-astra", "gpt-6-sol:high", "gpt-5.6-sol"). Omit (or leave empty) for the sub-agent model the user picked with /subagents, else the parent session\'s current model and effort',
       })),
     }),
     { description: "Tasks to delegate; they all run in parallel", minItems: 1 },
@@ -106,21 +128,28 @@ function renderSome(records: AgentRecord[], unknown: string[] = []): string {
   return parts.join("\n\n") || "No matching sub-agents.";
 }
 
-export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {}) {
+export default function (pi: ExtensionAPI, options: { run?: AgentRunner; admission?: MemoryAdmission | null; env?: NodeJS.ProcessEnv } = {}) {
+  // Sub-agents nest up to MAX_SUBAGENT_DEPTH layers below the main session,
+  // as in the desktop app: every sub-agent is told its layer, and the last
+  // layer gets no spawn_agents (nor the tools that manage sub-agents).
+  const depth = agentDepth(options.env ?? process.env);
+  if (depth > 0) {
+    pi.on("before_agent_start", async (event: any) => {
+      const base = typeof event?.systemPrompt === "string" ? event.systemPrompt : "";
+      return { systemPrompt: `${base}\n\n${subagentLayerNote(depth)}` };
+    });
+  }
+  if (!canDelegate(depth)) return;
+
   const watchdogMinutes = Math.round(childInactivityMs() / 60_000);
-  /** The session this runtime is on (deliveries go to it). */
-  let currentSession = "";
-  let cwd = process.cwd();
-  let inheritedModel = "";
+  const hub = deliveryHub(pi);
+  /** The last context seen (the TUI status bar for running agents). */
   let lastContext: any = null;
   let statusTimer: ReturnType<typeof setTimeout> | null = null;
-  /** An agent run is under way (one-shot runs deliver through the settle loop otherwise). */
-  let busy = false;
-  let oneShot = false;
 
   function publishStatus(): void {
     statusTimer = null;
-    const session = currentSession;
+    const session = hub.currentSession;
     if (!session || !lastContext?.ui?.setStatus) return;
     const records = manager.list(session);
     if (records.length > 0) lastContext.ui.setStatus("omnirush-agents", agentStatusSummary(records));
@@ -132,40 +161,60 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
     statusTimer.unref?.();
   }
 
-  const send = (records: AgentRecord[]) => {
-    pi.sendMessage(
-      {
-        customType: DELIVERY_TYPE,
-        content: renderDelivery(records),
-        display: true,
-        details: { results: records.map(resultOf) },
-      },
-      // Busy: into the running turn at its next step. Idle: a new turn.
-      { triggerTurn: true, deliverAs: "steer" },
-    );
+  // Memory-aware admission of children; thresholds from the environment or
+  // settings.json "subagentMemory" (memory-lib.ts). Built on first use.
+  let admission: MemoryAdmission | null = null;
+  const admissionOf = (): MemoryAdmission => {
+    admission ??= options.admission ?? new MemoryAdmission({ settings: readMemorySettings(omniDir()) });
+    return admission;
   };
-
-  const deliver = () => {
-    if (!currentSession) return;
-    // One-shot and between runs: the agent_settled loop delivers.
-    if (oneShot && !busy) {
-      return;
-    }
-    for (const group of manager.takeDeliverable(currentSession)) send(group.records);
+  const lazyAdmission = {
+    check: () => admissionOf().check(),
+    started: (key: object, pid?: number | null) => admissionOf().started(key, pid ?? null),
+    spawned: (key: object, pid: number | null | undefined) => admissionOf().spawned(key, pid),
+    finished: (key: object) => admissionOf().finished(key),
   };
 
+  /** Tell the user when children wait for memory (a warning when it is critically low). */
+  function memoryWait(info: { critical: boolean; detail: string; queued: number } | null): void {
+    scheduleStatus();
+    if (!info) return;
+    const text = info.critical
+      ? `omnirush: ${info.detail}. Not starting more sub-agents until memory frees up (${info.queued} queued); running ones continue.`
+      : `omnirush: ${info.queued} sub-agent${info.queued === 1 ? "" : "s"} queued: waiting for memory (${info.detail}).`;
+    if (lastContext?.hasUI && lastContext?.ui?.notify) lastContext.ui.notify(text, info.critical ? "warning" : "info");
+    // Headless: the top agent says it on stderr (a sub-agent's stderr only reaches its parent).
+    else if (info.critical || !String(process.env.OMNIRUSH_PARENT_SESSION ?? "").trim()) process.stderr.write(`${text}\n`);
+  }
+
   const manager = new AgentManager({
     run: options.run ?? ((task, options) => runChildAgent(task, {
-      cwd: options.cwd ?? cwd,
+      cwd: options.cwd ?? hub.cwd,
       parentSessionId: options.parentSessionId,
       sessionId: options.sessionId,
+      ...(options.sessionDir ? { sessionDir: options.sessionDir } : {}),
       signal: options.signal,
       onActivity: options.onActivity,
+      ...(options.onSpawn ? { onSpawn: options.onSpawn } : {}),
+      extraEnv: admissionOf().childEnv(),
       ...(options.timeoutMs !== undefined ? { timeoutMs: options.timeoutMs } : {}),
     })),
     // After the settle's own bookkeeping (waiters took theirs).
-    onDeliverable: () => queueMicrotask(deliver),
+    onDeliverable: () => hub.notify(),
     onChanged: scheduleStatus,
+    admission: options.admission === null ? null : lazyAdmission,
+    onMemoryWait: memoryWait,
+  });
+
+  hub.addSource({
+    hasPending: (session) => manager.hasPending(session),
+    // One message per finished batch (or per child with notify "each").
+    take: (session) => manager.takeDeliverable(session).map(({ records }) => ({
+      customType: DELIVERY_TYPE,
+      content: renderDelivery(records),
+      details: { results: records.map(resultOf) },
+    })),
+    changed: () => manager.changed(),
   });
 
   const hook: AgentsHook = {
@@ -177,17 +226,8 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
   (globalThis as any)[AGENTS_HOOK] = hook;
 
   const sessionOf = (ctx: any): string => {
-    const id = String(ctx?.sessionManager?.getSessionId?.() ?? "");
-    if (id && id !== currentSession) {
-      currentSession = id;
-      inheritedModel = "";
-    }
     lastContext = ctx ?? lastContext;
-    const selected = modelSelectionFromContext(ctx);
-    if (selected) inheritedModel = selected;
-    if (ctx?.cwd) cwd = String(ctx.cwd);
-    if (ctx?.mode === "print" || ctx?.mode === "json") oneShot = true;
-    return id;
+    return hub.sessionOf(ctx);
   };
 
   pi.registerTool({
@@ -212,6 +252,7 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
       "Short fan-outs whose results you need right away (quick searches, lookups): the default blocking call.",
       "Long work (implementation chunks, builds, test suites, long research — anything that may take many minutes) while you have other things to do: spawn_agents with wait:false, then keep working; the results arrive on their own as a message, so do not poll agents_status in a loop. Call agents_wait only when you have nothing else to do and need the results before going on.",
       "Do not use spawn_agents for a single quick action — doing it yourself is cheaper.",
+      SWARM_GUIDELINE,
     ],
     parameters: SpawnAgentsParams,
 
@@ -222,16 +263,22 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
         if (!raw || typeof raw.task !== "string" || !raw.task.trim()) {
           throw new Error("invalid tasks: every entry needs a non-empty task string");
         }
-        // An omitted model inherits the live parent selection. Explicit task
-        // models remain available for deliberate mixed-model swarms.
+        // A task's own model[:effort] wins; an omitted (or empty) one takes
+        // the /subagents pick, else the parent session's current model and
+        // effort (resolved below by subagentTasks).
         const requestedModel = typeof raw.model === "string" ? raw.model.trim() : "";
         tasks.push({
           role: raw.role,
           task: raw.task.trim(),
-          ...((requestedModel || inheritedModel) ? { model: requestedModel || inheritedModel } : {}),
+          ...(requestedModel ? { model: requestedModel } : {}),
         });
       }
       if (tasks.length === 0) throw new Error("no tasks given");
+      // The model and effort each child runs on (subagents.ts): the task's
+      // own model[:effort], else the /subagents pick, else the parent
+      // session's current model and effort; the main model when a pick is
+      // unavailable.
+      const resolvedTasks = subagentTasks(pi, ctx, tasks);
       if (!parentSessionId) throw new Error("no parent session id — subagents cannot be traced");
       const workdir = ctx?.cwd ? String(ctx.cwd) : process.cwd();
       // Only a deliberate limit counts: models fill optional numbers with
@@ -246,8 +293,10 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
         ? Math.floor(params.max_parallel)
         : undefined;
       const background = params.wait === false;
+      let progressTimer: ReturnType<typeof setTimeout> | null = null;
+      let finished = false;
 
-      const dispatched = manager.dispatch(parentSessionId, tasks, {
+      const dispatched = manager.dispatch(parentSessionId, resolvedTasks, {
         background,
         concurrency,
         notify: params.notify === "each" ? "each" : "batch",
@@ -260,18 +309,31 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
           });
         },
         onChanged: () => {
-          if (background) return;
-          onUpdate?.({
-            content: [{ type: "text", text: agentStatusSummary(manager.list(parentSessionId)) }],
-            details: { results: manager.list(parentSessionId).filter((candidate) => candidate.result).map(resultOf) },
-          });
+          // Every sign of life of every child lands here: the progress update
+          // (which carries the finished results) goes out at most every
+          // PROGRESS_UPDATE_MS, not once per output chunk.
+          if (background || progressTimer) return;
+          progressTimer = setTimeout(() => {
+            progressTimer = null;
+            if (finished) return;
+            onUpdate?.({
+              content: [{ type: "text", text: agentStatusSummary(manager.list(parentSessionId)) }],
+              details: { results: manager.list(parentSessionId).filter((candidate) => candidate.result).map(resultOf) },
+            });
+          }, PROGRESS_UPDATE_MS);
+          progressTimer.unref?.();
         },
         ...(timeoutMs !== undefined ? { timeoutMs } : {}),
         cwd: workdir,
+        // Children's sessions stay out of the user's /resume list.
+        sessionDir: subagentSessionDir(ctx?.sessionManager?.getSessionDir?.()),
       });
 
       if (background) {
-        const lines = dispatched.records.map((record) => `- ${record.id}: ${record.role}${record.model ? ` [${record.model}]` : ""} — ${record.task.length > 100 ? `${record.task.slice(0, 99)}…` : record.task}`);
+        const lines = dispatched.records.map((record, index) => {
+          const note = childModelFallback(resolvedTasks[index])?.note;
+          return `- ${record.id}: ${record.role}${record.model ? ` [${record.model}${record.effort ? ` · ${record.effort}` : ""}]` : ""} — ${record.task.length > 100 ? `${record.task.slice(0, 99)}…` : record.task}${note ? ` (${note})` : ""}`;
+        });
         return {
           content: [{
             type: "text",
@@ -288,20 +350,27 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
           details: {
             background: true,
             batch: dispatched.batch,
-            results: dispatched.records.map((record) => ({
-              agent_id: record.id,
-              role: record.role,
-              task: record.task,
-              session_id: record.sessionId,
-              ...(record.model ? { model: record.model } : {}),
-              parallel_limit: record.parallelLimit,
-              status: "running",
-            })),
+            results: dispatched.records.map((record, index) => {
+              const modelFallback = childModelFallback(resolvedTasks[index]);
+              return {
+                agent_id: record.id,
+                role: record.role,
+                task: record.task,
+                session_id: record.sessionId,
+                ...(record.model ? { model: record.model } : {}),
+                ...(record.effort ? { effort: record.effort } : {}),
+                ...(modelFallback ? { model_fallback: modelFallback } : {}),
+                parallel_limit: record.parallelLimit,
+                status: "running",
+              };
+            }),
           },
         };
       }
 
       await dispatched.done;
+      finished = true;
+      if (progressTimer) clearTimeout(progressTimer);
       const records = dispatched.records;
       return {
         content: [{ type: "text", text: renderChildResults(records.map((record) => record.result as ChildResult), records.map((record) => record.id)) }],
@@ -404,55 +473,6 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
     },
   });
 
-  pi.on("session_start", (_event: any, ctx: any) => {
-    sessionOf(ctx);
-  });
-
-  pi.on("agent_start", (_event: any, ctx: any) => {
-    sessionOf(ctx);
-    busy = true;
-  });
-
-  /**
-   * One-shot runs (print/json mode) end when the agent settles: background
-   * children still out would be lost. The run waits for them here and
-   * continues with their results (a new run nested in this settle, whose
-   * own settle resumes this one when everything is done).
-   */
-  let resumeOuter: (() => void) | null = null;
-  pi.on("agent_settled", async (_event: any, ctx: any) => {
-    busy = false;
-    const session = sessionOf(ctx);
-    if (!oneShot) {
-      // Anything that finished right as the run ended.
-      deliver();
-      return;
-    }
-    const outer = resumeOuter;
-    resumeOuter = null;
-    try {
-      for (;;) {
-        if (typeof ctx?.isIdle === "function" && !ctx.isIdle()) {
-          // A run started (a delivery): its settle resumes this one.
-          await new Promise<void>((resolve) => { resumeOuter = resolve; });
-          break;
-        }
-        if (!session || !manager.hasPending(session)) break;
-        const groups = manager.takeDeliverable(session);
-        if (groups.length === 0) {
-          await manager.changed();
-          continue;
-        }
-        busy = true;
-        const records = groups.flatMap((group) => group.records);
-        send(records);
-        if (typeof ctx?.isIdle === "function" && ctx.isIdle()) busy = false;
-      }
-    } finally {
-      outer?.();
-    }
-  });
-
   pi.on("session_shutdown", async () => {
     // Children the session ends under are stopped (the collector, when it
     // runs, already did this before its last capture).
@@ -460,7 +480,7 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
   });
 
   pi.registerCommand("agents", {
-    description: "List this session's subagents (running and finished); /agents cancel <id|all> stops running ones",
+    description: "List this session's subagents (running and finished) and the memory they may use; /agents cancel <id|all> stops running ones; /agents memory shows the memory admission",
     handler: async (args: string, ctx: any) => {
       const session = sessionOf(ctx);
       const words = String(args ?? "").trim().split(/\s+/).filter(Boolean);
@@ -472,11 +492,16 @@ export default function (pi: ExtensionAPI, options: { run?: AgentRunner } = {})
         return;
       }
       const records = manager.list(session);
+      const memory = options.admission === null ? "" : `\n${admissionOf().describe()}`;
+      if (words[0] === "memory") {
+        ctx?.ui?.notify?.(memory.trim() || "Memory admission is off.");
+        return;
+      }
       if (records.length === 0) {
-        ctx?.ui?.notify?.("No subagents have been spawned in this session yet.");
+        ctx?.ui?.notify?.(`No subagents have been spawned in this session yet.${memory}`);
         return;
       }
-      ctx?.ui?.notify?.(renderAgentStatus(records));
+      ctx?.ui?.notify?.(`${renderAgentStatus(records)}${memory}`);
     },
   });
 }