omnirush 0.5.3 → 0.5.4

package/assets/extensions/omnirush/agents-lib.ts

@@ -0,0 +1,421 @@
+// agents-lib — pure orchestration logic for the spawn_agents tool
+// (headless pi child runs). Pure/injectable so node:test covers it
+// without ever spawning a real agent.
+//
+// Children are headless runs of THIS entry (process.argv[1], i.e. the
+// running pi entry — same runtime, same global extension/config dir)
+// in the parent's workspace, with `--mode json -p` so their output is
+// a parseable event stream. Each child writes its own session JSONL
+// under <agentDir>/sessions/<cwd-slug>/ — the workspace collector
+// ships every session file it finds there as a trace artifact linked
+// via parent_session_id, so subagent tracing is free.
+//
+// Linkage: children get OMNIRUSH_PARENT_SESSION=<parent session id> in
+// their environment (the collector stamps parent_session_id from it).
+
+import { spawn as nodeSpawn } from "node:child_process";
+import { existsSync } from "node:fs";
+import { mkdtemp, writeFile, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import path from "node:path";
+
+/** Default per-child wall clock budget. */
+export const CHILD_TIMEOUT_MS = 10 * 60_000;
+/** Grace between SIGTERM and SIGKILL on timeout. */
+export const CHILD_KILL_GRACE_MS = 5_000;
+/** Default number of children running at once. */
+export const MAX_CONCURRENT_CHILDREN = 4;
+/** Upper bound on tasks per spawn_agents call. */
+export const MAX_TASKS_PER_CALL = 8;
+/** Per-child output cap in the structured result (bytes, UTF-8). */
+export const CHILD_OUTPUT_CAP_BYTES = 50 * 1024;
+
+export const AGENT_ROLES = ["code-searcher", "researcher-web", "general-worker"] as const;
+export type AgentRole = (typeof AGENT_ROLES)[number];
+
+export interface RolePreset {
+  role: AgentRole;
+  label: string;
+  description: string;
+  /** Appended to the child's system prompt. */
+  systemPrompt: string;
+}
+
+/**
+ * Role presets. Guidance only — the child keeps every tool (hard tool
+ * filters would break extension tools like web_fetch); the system
+ * prompt sets expectations instead.
+ */
+export const ROLE_PRESETS: Record<AgentRole, RolePreset> = {
+  "code-searcher": {
+    role: "code-searcher",
+    label: "Code Searcher",
+    description: "Read-only codebase exploration: locate code, trace call paths, summarize findings.",
+    systemPrompt: [
+      "You are a code-searcher subagent.",
+      "Your job is READ-ONLY codebase exploration: find where things live, trace how calls flow, and report precise findings (file paths, symbol names, line references).",
+      "Do NOT modify files, run state-changing commands, or start long-running processes.",
+      "Report findings as a compact structured summary: what you found, where (paths), and any surprises — not a play-by-play of your searches.",
+    ].join(" "),
+  },
+  "researcher-web": {
+    role: "researcher-web",
+    label: "Researcher (Web)",
+    description: "Web research via web_fetch: read documentation pages, articles and API responses.",
+    systemPrompt: [
+      "You are a researcher subagent with web access.",
+      "Use the web_fetch tool to read documentation, articles and API responses (HTTPS URLs; the tool strips pages to readable text).",
+      "Cross-check claims across more than one page when it matters, and prefer primary sources (official docs, spec pages) over blog summaries.",
+      "Report a compact synthesis with the source URLs you actually used — not a list of everything you fetched.",
+    ].join(" "),
+  },
+  "general-worker": {
+    role: "general-worker",
+    label: "General Worker",
+    description: "Full-tool worker for self-contained implementation tasks (edits, commands, verification).",
+    systemPrompt: [
+      "You are a general-worker subagent with the full toolset.",
+      "Complete the delegated task end to end: implement, verify (run the relevant tests/commands), and keep changes tight and reviewable.",
+      "Stay inside the delegated task — do not refactor beyond it, do not touch unrelated files.",
+      "Report what you changed (paths), how you verified it, and anything you deliberately left undone.",
+    ].join(" "),
+  },
+};
+
+/** Resolve the child entry: this running script + the current runtime. */
+export function childInvocation(args: string[]): {
+  command: string;
+  args: string[];
+} {
+  const currentScript = process.argv[1];
+  const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
+  if (currentScript && !isBunVirtualScript && existsSync(currentScript)) {
+    return { command: process.execPath, args: [currentScript, ...args] };
+  }
+  const execName = path.basename(process.execPath).toLowerCase();
+  const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
+  if (!isGenericRuntime) {
+    return { command: process.execPath, args };
+  }
+  // Packaged bun builds: fall back to the pi on PATH.
+  return { command: "pi", args };
+}
+
+/**
+ * Build the child's argv for one task: JSON print mode with the role's
+ * system prompt appended via a temp file path (the caller writes and
+ * cleans up that file — see withRolePromptFile).
+ */
+export function buildChildArgs(task: string, promptFilePath: string | null): string[] {
+  const args: string[] = ["--mode", "json", "-p"];
+  if (promptFilePath) {
+    args.push("--append-system-prompt", promptFilePath);
+  }
+  args.push(`Task: ${task}`);
+  return args;
+}
+
+/**
+ * Write a role's system prompt to a private temp file and run `body`
+ * with its path; the file is cleaned up afterwards either way.
+ */
+export async function withRolePromptFile(
+  role: AgentRole,
+  body: (promptFilePath: string) => Promise<ChildResultInput>,
+): Promise<ChildResultInput> {
+  let dir: string | null = null;
+  try {
+    dir = await mkdtemp(path.join(tmpdir(), "omnirush-agent-"));
+    const filePath = path.join(dir, `prompt-${role.replace(/[^\w.-]+/g, "_")}.md`);
+    await writeFile(filePath, ROLE_PRESETS[role].systemPrompt, { encoding: "utf8", mode: 0o600 });
+    return await body(filePath);
+  } finally {
+    if (dir) await rm(dir, { recursive: true, force: true }).catch(() => undefined);
+  }
+}
+
+export interface ChildTask {
+  role: AgentRole;
+  task: string;
+}
+
+export type ChildStatus = "completed" | "failed" | "timeout" | "cancelled";
+
+export interface ChildResult {
+  role: AgentRole;
+  task: string;
+  status: ChildStatus;
+  /** Process exit code (null when killed by a signal or still unknown). */
+  exitCode: number | null;
+  /** Final assistant text (or the best available partial output). */
+  output: string;
+  /** True when `output` was capped (full stdout events were larger). */
+  outputTruncated: boolean;
+  durationMs: number;
+  /** Assistant-turn count seen in the child's event stream. */
+  turns: number;
+  error?: string;
+}
+
+export interface ChildResultInput {
+  exitCode: number | null;
+  finalText: string;
+  turns: number;
+  status?: ChildStatus;
+  error?: string;
+}
+
+/** Final structured result for one child. */
+export function buildChildResult(
+  task: ChildTask,
+  input: ChildResultInput,
+  durationMs: number,
+): ChildResult {
+  const output = input.finalText;
+  let capped = output;
+  const truncated = Buffer.byteLength(output, "utf8") > CHILD_OUTPUT_CAP_BYTES;
+  if (truncated) {
+    capped = output.slice(0, CHILD_OUTPUT_CAP_BYTES);
+    while (Buffer.byteLength(capped, "utf8") > CHILD_OUTPUT_CAP_BYTES) capped = capped.slice(0, -1);
+  }
+  return {
+    role: task.role,
+    task: task.task,
+    status: input.status ?? (input.exitCode === 0 ? "completed" : "failed"),
+    exitCode: input.exitCode,
+    output: capped,
+    outputTruncated: truncated,
+    durationMs,
+    turns: input.turns,
+    ...(input.error ? { error: input.error } : {}),
+  };
+}
+
+// --- child event-stream parsing -------------------------------------------
+
+export interface ParsedChildEvents {
+  /** Final assistant text (last assistant message with text parts). */
+  finalText: string;
+  /** Assistant messages seen (turn count). */
+  turns: number;
+}
+
+/**
+ * Parse the accumulated stdout of a `pi --mode json -p` run: one JSON
+ * event per line; `message_end` events carry the messages. Unparseable
+ * lines (banners, noise) are ignored.
+ */
+export function parseChildEvents(stdout: string): ParsedChildEvents {
+  let finalText = "";
+  let turns = 0;
+  for (const line of stdout.split("\n")) {
+    const trimmed = line.trim();
+    if (!trimmed) continue;
+    let event: any;
+    try {
+      event = JSON.parse(trimmed);
+    } catch {
+      continue;
+    }
+    if (event?.type !== "message_end" || !event.message) continue;
+    const message = event.message;
+    if (message.role !== "assistant") continue;
+    turns += 1;
+    const parts = Array.isArray(message.content) ? message.content : [];
+    const text = parts
+      .filter((part: any) => part?.type === "text" && typeof part.text === "string")
+      .map((part: any) => part.text)
+      .join("\n\n");
+    if (text.trim()) finalText = text;
+    else if (typeof message.content === "string" && message.content.trim()) {
+      finalText = message.content;
+    }
+  }
+  return { finalText, turns };
+}
+
+// --- the runner -------------------------------------------------------------
+
+export interface SpawnChildOptions {
+  cwd: string;
+  parentSessionId: string;
+  timeoutMs?: number;
+  /** Kill the child when this fires (the tool's cancellation signal). */
+  signal?: AbortSignal;
+  /** Grace between SIGTERM and SIGKILL when killing (default CHILD_KILL_GRACE_MS). */
+  killGraceMs?: number;
+  spawnImpl?: typeof nodeSpawn;
+  now?: () => number;
+  /** Progress callback: partial output lines as the child runs. */
+  onChildStdout?: (role: AgentRole, chunk: string) => void;
+}
+
+/**
+ * Run ONE child agent to completion (or timeout/cancellation). The
+ * child inherits the parent's environment plus OMNIRUSH_PARENT_SESSION.
+ * On timeout: SIGTERM, a short grace, then SIGKILL; whatever the child
+ * printed so far is kept as the partial result.
+ */
+export async function runChildAgent(
+  task: ChildTask,
+  options: SpawnChildOptions,
+): Promise<ChildResult> {
+  const preset = ROLE_PRESETS[task.role] ?? ROLE_PRESETS["general-worker"];
+  const startedAt = (options.now ?? Date.now)();
+  const timeoutMs = Math.max(1_000, options.timeoutMs ?? CHILD_TIMEOUT_MS);
+  const spawnImpl = options.spawnImpl ?? nodeSpawn;
+
+  return withRolePromptFile(task.role, async (promptFile) => {
+    const args = buildChildArgs(task.task, promptFile);
+    const invocation = childInvocation(args);
+
+    return await new Promise<ChildResult>((resolvePromise) => {
+      let stdout = "";
+      let stderr = "";
+      let settled = false;
+      let killedFor: "timeout" | "cancelled" | null = null;
+
+      const child = spawnImpl(invocation.command, invocation.args, {
+        cwd: options.cwd,
+        shell: false,
+        stdio: ["ignore", "pipe", "pipe"],
+        env: {
+          ...process.env,
+          OMNIRUSH_PARENT_SESSION: options.parentSessionId,
+        },
+      });
+
+      const finish = (input: ChildResultInput) => {
+        if (settled) return;
+        settled = true;
+        clearTimeout(timer);
+        if (options.signal && abortHandler) {
+          options.signal.removeEventListener("abort", abortHandler);
+        }
+        resolvePromise(
+          buildChildResult(task, input, (options.now ?? Date.now)() - startedAt),
+        );
+      };
+
+      const killTree = () => {
+        try {
+          child.kill("SIGTERM");
+        } catch {
+          /* already gone */
+        }
+        setTimeout(() => {
+          try {
+            child.kill("SIGKILL");
+          } catch {
+            /* already gone */
+          }
+        }, Math.max(50, options.killGraceMs ?? CHILD_KILL_GRACE_MS)).unref?.();
+      };
+
+      const timer = setTimeout(() => {
+        if (settled) return;
+        killedFor = "timeout";
+        killTree();
+      }, timeoutMs);
+      timer.unref?.();
+
+      const abortHandler = options.signal
+        ? () => {
+            if (settled) return;
+            killedFor = "cancelled";
+            killTree();
+          }
+        : null;
+      if (options.signal && abortHandler) {
+        if (options.signal.aborted) abortHandler();
+        else options.signal.addEventListener("abort", abortHandler, { once: true });
+      }
+
+      child.stdout?.on("data", (chunk: Buffer | string) => {
+        const text = String(chunk);
+        stdout += text;
+        options.onChildStdout?.(task.role, text);
+      });
+      child.stderr?.on("data", (chunk: Buffer | string) => {
+        stderr += String(chunk);
+      });
+      child.on("error", (error: Error) => {
+        finish({
+          exitCode: null,
+          finalText: "",
+          turns: 0,
+          status: "failed",
+          error: `failed to start: ${error.message}`,
+        });
+      });
+      child.on("close", (code: number | null, signalName: string | null) => {
+        const parsed = parseChildEvents(stdout);
+        if (killedFor === "timeout") {
+          finish({
+            exitCode: code,
+            finalText: parsed.finalText,
+            turns: parsed.turns,
+            status: "timeout",
+            error: `timed out after ${Math.round(timeoutMs / 60_000)} min${parsed.finalText ? " — partial result kept" : ""}`,
+          });
+          return;
+        }
+        if (killedFor === "cancelled") {
+          finish({
+            exitCode: code,
+            finalText: parsed.finalText,
+            turns: parsed.turns,
+            status: "cancelled",
+            error: "cancelled by the parent session",
+          });
+          return;
+        }
+        finish({
+          exitCode: code,
+          finalText: parsed.finalText || (code !== 0 ? stderr.slice(-2_000) : ""),
+          turns: parsed.turns,
+          status: code === 0 ? "completed" : "failed",
+          ...(code !== 0 && !parsed.finalText && stderr.trim() ? { error: stderr.trim().slice(-500) } : {}),
+        });
+        void signalName;
+      });
+    });
+  });
+}
+
+/**
+ * Run tasks with a concurrency cap, preserving input order in the
+ * results (ports the pi subagent example's mapWithConcurrencyLimit).
+ */
+export async function mapWithConcurrency<TIn, TOut>(
+  items: TIn[],
+  concurrency: number,
+  fn: (item: TIn, index: number) => Promise<TOut>,
+): Promise<TOut[]> {
+  if (items.length === 0) return [];
+  const limit = Math.max(1, Math.min(concurrency, items.length));
+  const results: TOut[] = new Array(items.length);
+  let nextIndex = 0;
+  const workers = Array.from({ length: limit }, async () => {
+    for (;;) {
+      const current = nextIndex++;
+      if (current >= items.length) return;
+      results[current] = await fn(items[current], current);
+    }
+  });
+  await Promise.all(workers);
+  return results;
+}
+
+/** Structured result text for the parent model (one section per child). */
+export function renderChildResults(results: ChildResult[]): string {
+  const succeeded = results.filter((result) => result.status === "completed").length;
+  const sections = results.map((result) => {
+    const minutes = Math.round((result.durationMs / 60_000) * 10) / 10;
+    const header = `### ${result.role} — ${result.status} (${minutes} min${result.outputTruncated ? ", output capped" : ""})`;
+    const meta: string[] = [`task: ${result.task}`];
+    if (result.error) meta.push(`error: ${result.error}`);
+    return `${header}\n${meta.join("\n")}\n\n${result.output || "(no output)"}`;
+  });
+  return `spawn_agents: ${succeeded}/${results.length} completed\n\n${sections.join("\n\n---\n\n")}`;
+}

package/assets/extensions/omnirush/agents.ts

@@ -0,0 +1,163 @@
+// omnirush-agents — subagent orchestration (the headline feature).
+//
+// `spawn_agents` delegates tasks to headless child agents running in
+// the SAME workspace:
+//   - role presets: code-searcher (read-only exploration),
+//     researcher-web (web_fetch research), general-worker (full tools)
+//   - parallel execution, default max 4 children at once
+//   - per-child wall-clock timeout (default 10 min): kill + keep the
+//     partial result
+//   - structured results returned to the parent model
+//
+// Tracing is free: every child writes its own session JSONL under
+// <agentDir>/sessions/<cwd-slug>/ (children run without --no-session),
+// and the workspace collector ships every session file it finds as a
+// trace artifact linked via parent_session_id. Children also get
+// OMNIRUSH_PARENT_SESSION so a collector running inside them knows the
+// parent session (linked mode: no second baseline, authoritative
+// parent stamping).
+//
+// `/agents` lists this session's subagents (running + finished).
+
+import { StringEnum } from "@earendil-works/pi-ai";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+
+import {
+  AGENT_ROLES,
+  CHILD_TIMEOUT_MS,
+  mapWithConcurrency,
+  MAX_CONCURRENT_CHILDREN,
+  MAX_TASKS_PER_CALL,
+  renderChildResults,
+  runChildAgent,
+  type ChildResult,
+  type ChildTask,
+} from "./agents-lib";
+
+const SpawnAgentsParams = Type.Object({
+  tasks: Type.Array(
+    Type.Object({
+      role: StringEnum(AGENT_ROLES, {
+        description: "code-searcher: read-only codebase exploration; researcher-web: web research via web_fetch; general-worker: full-tool implementation work",
+      }),
+      task: Type.String({ description: "Self-contained task description for the child agent (it sees ONLY this, not your conversation)" }),
+    }),
+    { description: "Tasks to delegate; they run in parallel", minItems: 1, maxItems: MAX_TASKS_PER_CALL },
+  ),
+  timeout_minutes: Type.Optional(Type.Number({
+    description: `Per-child wall-clock limit in minutes (default ${CHILD_TIMEOUT_MS / 60_000}; a timed-out child is killed and its partial result is kept)`,
+    minimum: 0.1,
+    maximum: 60,
+  })),
+});
+
+interface AgentRegistryEntry {
+  role: string;
+  task: string;
+  status: "running" | ChildResult["status"];
+  startedAt: number;
+  durationMs: number | null;
+}
+
+export default function (pi: ExtensionAPI) {
+  /** session id -> subagent records (this process only). */
+  const registry = new Map<string, AgentRegistryEntry[]>();
+
+  pi.registerTool({
+    name: "spawn_agents",
+    label: "Spawn Agents",
+    description: [
+      "Delegate tasks to parallel subagents with isolated context windows.",
+      "Each child is a headless agent run in this workspace with a role preset:",
+      "code-searcher (read-only exploration), researcher-web (web research via web_fetch), general-worker (full tools).",
+      `Children run in parallel (max ${MAX_CONCURRENT_CHILDREN} at once, ${MAX_TASKS_PER_CALL} tasks per call) and each gets its own timeout (default ${CHILD_TIMEOUT_MS / 60_000} min).`,
+      "Every task description must be SELF-CONTAINED: the child cannot see this conversation.",
+      "Use for parallelizable work: broad code surveys, independent research questions, independent implementation chunks.",
+    ].join(" "),
+    promptSnippet: "spawn_agents: delegate parallel tasks to role-preset subagents",
+    promptGuidelines: [
+      "Use spawn_agents when 2+ independent sub-tasks would each take several tool calls (surveys, research, independent fixes); write each task brief as a self-contained mini-prompt with the goal, relevant paths and the expected report format.",
+      "Do not use spawn_agents for a single quick action — doing it yourself is cheaper.",
+    ],
+    parameters: SpawnAgentsParams,
+
+    async execute(_toolCallId, params, signal, onUpdate, ctx) {
+      const tasks: ChildTask[] = [];
+      for (const raw of params.tasks ?? []) {
+        if (!raw || typeof raw.task !== "string" || !raw.task.trim()) {
+          throw new Error("invalid tasks: every entry needs a non-empty task string");
+        }
+        tasks.push({ role: raw.role, task: raw.task.trim() });
+      }
+      if (tasks.length === 0) throw new Error("no tasks given");
+      if (tasks.length > MAX_TASKS_PER_CALL) {
+        throw new Error(`too many tasks (${tasks.length}); the limit is ${MAX_TASKS_PER_CALL} per call`);
+      }
+      const parentSessionId = String(ctx?.sessionManager?.getSessionId?.() ?? "");
+      if (!parentSessionId) throw new Error("no parent session id — subagents cannot be traced");
+      const cwd = ctx?.cwd ? String(ctx.cwd) : process.cwd();
+      const timeoutMinutes = typeof params.timeout_minutes === "number" && params.timeout_minutes > 0
+        ? params.timeout_minutes
+        : undefined;
+
+      const entries = registry.get(parentSessionId) ?? [];
+      registry.set(parentSessionId, entries);
+      const marks = tasks.map((task) => {
+        const entry: AgentRegistryEntry = {
+          role: task.role,
+          task: task.task,
+          status: "running",
+          startedAt: Date.now(),
+          durationMs: null,
+        };
+        entries.push(entry);
+        return entry;
+      });
+
+      const results = await mapWithConcurrency<ChildTask, ChildResult>(
+        tasks,
+        MAX_CONCURRENT_CHILDREN,
+        async (task, index) => {
+          const result = await runChildAgent(task, {
+            cwd,
+            parentSessionId,
+            timeoutMs: timeoutMinutes === undefined ? undefined : timeoutMinutes * 60_000,
+            signal,
+          });
+          const mark = marks[index];
+          mark.status = result.status;
+          mark.durationMs = result.durationMs;
+          onUpdate?.({
+            content: [{ type: "text", text: `spawn_agents: ${results.filter(Boolean).length + 1}/${tasks.length} settled (latest: ${task.role})` }],
+            details: { results: [...results.filter(Boolean)] },
+          });
+          return result;
+        },
+      );
+
+      return {
+        content: [{ type: "text", text: renderChildResults(results) }],
+        details: { results },
+      };
+    },
+  });
+
+  pi.registerCommand("agents", {
+    description: "List this session's subagents (running and finished)",
+    handler: async (_args: string, ctx: any) => {
+      const sessionId = String(ctx?.sessionManager?.getSessionId?.() ?? "");
+      const entries = registry.get(sessionId) ?? [];
+      if (entries.length === 0) {
+        ctx?.ui?.notify?.("No subagents have been spawned in this session yet.");
+        return;
+      }
+      const lines = entries.map((entry, index) => {
+        const minutes = Math.round(((entry.durationMs ?? Date.now() - entry.startedAt) / 60_000) * 10) / 10;
+        const task = entry.task.length > 60 ? `${entry.task.slice(0, 59)}…` : entry.task;
+        return `  ${index + 1}. [${entry.status}] ${entry.role} (${minutes} min) — ${task}`;
+      });
+      ctx?.ui?.notify?.(`Subagents of this session:\n${lines.join("\n")}`);
+    },
+  });
+}