oira666_pi-subagent 0.1.0

package/runner.ts

@@ -0,0 +1,419 @@
+/**
+ * Subagent process runner.
+ *
+ * Spawns isolated `pi` processes and streams results back via callbacks.
+ */
+
+import { spawn } from "node:child_process";
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import type { AgentToolResult } from "@mariozechner/pi-agent-core";
+import type { Message } from "@mariozechner/pi-ai";
+import type { AgentConfig } from "./agents.js";
+import {
+  type DelegationMode,
+  type SingleResult,
+  type SubagentDetails,
+  emptyUsage,
+  getFinalOutput,
+} from "./types.js";
+
+const SIGKILL_TIMEOUT_MS = 5000;
+const SUBAGENT_DEPTH_ENV = "PI_SUBAGENT_DEPTH";
+const SUBAGENT_MAX_DEPTH_ENV = "PI_SUBAGENT_MAX_DEPTH";
+const SUBAGENT_STACK_ENV = "PI_SUBAGENT_STACK";
+const SUBAGENT_PREVENT_CYCLES_ENV = "PI_SUBAGENT_PREVENT_CYCLES";
+const PI_OFFLINE_ENV = "PI_OFFLINE";
+
+type OnUpdateCallback = (partial: AgentToolResult<SubagentDetails>) => void;
+
+// ---------------------------------------------------------------------------
+// Temp file helpers
+// ---------------------------------------------------------------------------
+
+function writePromptToTempFile(
+  agentName: string,
+  prompt: string,
+): { dir: string; filePath: string } {
+  const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-subagent-"));
+  const safeName = agentName.replace(/[^\w.-]+/g, "_");
+  const filePath = path.join(tmpDir, `prompt-${safeName}.md`);
+  fs.writeFileSync(filePath, prompt, { encoding: "utf-8", mode: 0o600 });
+  return { dir: tmpDir, filePath };
+}
+
+function writeForkSessionToTempFile(
+  agentName: string,
+  sessionJsonl: string,
+): { dir: string; filePath: string } {
+  const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-subagent-"));
+  const safeName = agentName.replace(/[^\w.-]+/g, "_");
+  const filePath = path.join(tmpDir, `fork-${safeName}.jsonl`);
+  fs.writeFileSync(filePath, sessionJsonl, { encoding: "utf-8", mode: 0o600 });
+  return { dir: tmpDir, filePath };
+}
+
+function cleanupTempDir(dir: string | null): void {
+  if (!dir) return;
+  try {
+    fs.rmSync(dir, { recursive: true, force: true });
+  } catch {
+    /* ignore */
+  }
+}
+
+function resolveExtensionArg(value: string): string {
+  if (!value) return value;
+  if (value.startsWith("npm:") || value.startsWith("git:")) return value;
+  if (value.startsWith("~/")) return path.join(os.homedir(), value.slice(2));
+  if (path.isAbsolute(value)) return value;
+
+  const resolved = path.resolve(process.cwd(), value);
+  return fs.existsSync(resolved) ? resolved : value;
+}
+
+function getInheritedExtensionArgs(argv: string[]): string[] {
+  const args: string[] = [];
+  for (let i = 2; i < argv.length; i++) {
+    const arg = argv[i];
+
+    if (arg === "--no-extensions" || arg === "-ne") {
+      args.push("--no-extensions");
+      continue;
+    }
+
+    if (arg === "--extension" || arg === "-e") {
+      const value = argv[i + 1];
+      if (value !== undefined) {
+        args.push("--extension", resolveExtensionArg(value));
+        i++;
+      }
+      continue;
+    }
+
+    if (arg.startsWith("--extension=")) {
+      args.push("--extension", resolveExtensionArg(arg.slice("--extension=".length)));
+    }
+  }
+  return args;
+}
+
+// ---------------------------------------------------------------------------
+// JSON-line stream processing
+// ---------------------------------------------------------------------------
+
+function processJsonLine(line: string, result: SingleResult): boolean {
+  if (!line.trim()) return false;
+
+  let event: any;
+  try {
+    event = JSON.parse(line);
+  } catch {
+    return false;
+  }
+
+  if (event.type === "message_end" && event.message) {
+    const msg = event.message as Message;
+    result.messages.push(msg);
+
+    if (msg.role === "assistant") {
+      result.usage.turns++;
+      const usage = msg.usage;
+      if (usage) {
+        result.usage.input += usage.input || 0;
+        result.usage.output += usage.output || 0;
+        result.usage.cacheRead += usage.cacheRead || 0;
+        result.usage.cacheWrite += usage.cacheWrite || 0;
+        result.usage.cost += usage.cost?.total || 0;
+        result.usage.contextTokens = usage.totalTokens || 0;
+      }
+      if (!result.model && msg.model) result.model = msg.model;
+      if (msg.stopReason) result.stopReason = msg.stopReason;
+      if (msg.errorMessage) result.errorMessage = msg.errorMessage;
+    }
+    return true;
+  }
+
+  if (event.type === "tool_result_end" && event.message) {
+    result.messages.push(event.message as Message);
+    return true;
+  }
+
+  return false;
+}
+
+// ---------------------------------------------------------------------------
+// Build pi CLI arguments
+// ---------------------------------------------------------------------------
+
+function buildPiArgs(
+  agent: AgentConfig,
+  systemPromptPath: string | null,
+  task: string,
+  delegationMode: DelegationMode,
+  forkSessionPath: string | null,
+): string[] {
+  const args: string[] = [
+    "--mode",
+    "json",
+    ...getInheritedExtensionArgs(process.argv),
+    "-p",
+  ];
+
+  if (delegationMode === "spawn") {
+    args.push("--no-session");
+  } else if (forkSessionPath) {
+    args.push("--session", forkSessionPath);
+  }
+
+  if (agent.model) args.push("--model", agent.model);
+  if (agent.thinking) args.push("--thinking", agent.thinking);
+  if (agent.tools && agent.tools.length > 0)
+    args.push("--tools", agent.tools.join(","));
+  if (systemPromptPath) args.push("--append-system-prompt", systemPromptPath);
+  args.push(`Task: ${task}`);
+  return args;
+}
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+export interface RunAgentOptions {
+  /** Fallback working directory when the task doesn't specify one. */
+  cwd: string;
+  /** All available agent configs. */
+  agents: AgentConfig[];
+  /** Name of the agent to run. */
+  agentName: string;
+  /** Task description. */
+  task: string;
+  /** Optional override working directory. */
+  taskCwd?: string;
+  /** Context mode: spawn (fresh) or fork (session snapshot + task). */
+  delegationMode: DelegationMode;
+  /** Serialized parent session snapshot used when delegationMode is "fork". */
+  forkSessionSnapshotJsonl?: string;
+  /** Current delegation depth of the caller process. */
+  parentDepth: number;
+  /** Delegation stack from the caller process (ancestor agent names). */
+  parentAgentStack: string[];
+  /** Maximum allowed delegation depth to propagate to child processes. */
+  maxDepth: number;
+  /** Whether cycle prevention should be enforced in child processes. */
+  preventCycles: boolean;
+  /** Abort signal for cancellation. */
+  signal?: AbortSignal;
+  /** Streaming update callback. */
+  onUpdate?: OnUpdateCallback;
+  /** Factory to wrap results into SubagentDetails. */
+  makeDetails: (results: SingleResult[]) => SubagentDetails;
+}
+
+/**
+ * Spawn a single subagent process and collect its results.
+ *
+ * Returns a SingleResult even on failure (exitCode > 0, stderr populated).
+ */
+export async function runAgent(opts: RunAgentOptions): Promise<SingleResult> {
+  const {
+    cwd,
+    agents,
+    agentName,
+    task,
+    taskCwd,
+    delegationMode,
+    forkSessionSnapshotJsonl,
+    parentDepth,
+    parentAgentStack,
+    maxDepth,
+    preventCycles,
+    signal,
+    onUpdate,
+    makeDetails,
+  } = opts;
+
+  const agent = agents.find((a) => a.name === agentName);
+  if (!agent) {
+    const available = agents.map((a) => `"${a.name}"`).join(", ") || "none";
+    return {
+      agent: agentName,
+      agentSource: "unknown",
+      task,
+      exitCode: 1,
+      messages: [],
+      stderr: `Unknown agent: "${agentName}". Available agents: ${available}.`,
+      usage: emptyUsage(),
+    };
+  }
+
+  if (
+    delegationMode === "fork" &&
+    (!forkSessionSnapshotJsonl || !forkSessionSnapshotJsonl.trim())
+  ) {
+    return {
+      agent: agentName,
+      agentSource: agent.source,
+      task,
+      exitCode: 1,
+      messages: [],
+      stderr:
+        "Cannot run in fork mode: missing parent session snapshot context.",
+      usage: emptyUsage(),
+      model: agent.model,
+      stopReason: "error",
+      errorMessage:
+        "Cannot run in fork mode: missing parent session snapshot context.",
+    };
+  }
+
+  const result: SingleResult = {
+    agent: agentName,
+    agentSource: agent.source,
+    task,
+    exitCode: -1,
+    messages: [],
+    stderr: "",
+    usage: emptyUsage(),
+    model: agent.model,
+  };
+
+  const emitUpdate = () => {
+    onUpdate?.({
+      content: [
+        {
+          type: "text",
+          text: getFinalOutput(result.messages) || "(running...)",
+        },
+      ],
+      details: makeDetails([result]),
+    });
+  };
+
+  // Write system prompt to temp file if needed
+  let promptTmpDir: string | null = null;
+  let promptTmpPath: string | null = null;
+  if (agent.systemPrompt.trim()) {
+    const tmp = writePromptToTempFile(agent.name, agent.systemPrompt);
+    promptTmpDir = tmp.dir;
+    promptTmpPath = tmp.filePath;
+  }
+
+  // Write forked session snapshot if needed
+  let forkSessionTmpDir: string | null = null;
+  let forkSessionTmpPath: string | null = null;
+  if (delegationMode === "fork" && forkSessionSnapshotJsonl) {
+    const tmp = writeForkSessionToTempFile(agent.name, forkSessionSnapshotJsonl);
+    forkSessionTmpDir = tmp.dir;
+    forkSessionTmpPath = tmp.filePath;
+  }
+
+  try {
+    const piArgs = buildPiArgs(
+      agent,
+      promptTmpPath,
+      task,
+      delegationMode,
+      forkSessionTmpPath,
+    );
+    let wasAborted = false;
+
+    const exitCode = await new Promise<number>((resolve) => {
+      const nextDepth = Math.max(0, Math.floor(parentDepth)) + 1;
+      const propagatedMaxDepth = Math.max(0, Math.floor(maxDepth));
+      const propagatedStack = [...parentAgentStack, agentName];
+      const proc = spawn("pi", piArgs, {
+        cwd: taskCwd ?? cwd,
+        shell: false,
+        stdio: ["ignore", "pipe", "pipe"],
+        env: {
+          ...process.env,
+          [SUBAGENT_DEPTH_ENV]: String(nextDepth),
+          [SUBAGENT_MAX_DEPTH_ENV]: String(propagatedMaxDepth),
+          [SUBAGENT_STACK_ENV]: JSON.stringify(propagatedStack),
+          [SUBAGENT_PREVENT_CYCLES_ENV]: preventCycles ? "1" : "0",
+          [PI_OFFLINE_ENV]: "1",
+        },
+      });
+
+      let buffer = "";
+
+      const flushLine = (line: string) => {
+        if (processJsonLine(line, result)) emitUpdate();
+      };
+
+      proc.stdout.on("data", (chunk: Buffer) => {
+        buffer += chunk.toString();
+        const lines = buffer.split("\n");
+        buffer = lines.pop() || "";
+        for (const line of lines) flushLine(line);
+      });
+
+      proc.stderr.on("data", (chunk: Buffer) => {
+        result.stderr += chunk.toString();
+      });
+
+      proc.on("close", (code) => {
+        if (buffer.trim()) flushLine(buffer);
+        resolve(code ?? 0);
+      });
+
+      proc.on("error", () => resolve(1));
+
+      // Abort handling
+      if (signal) {
+        const kill = () => {
+          wasAborted = true;
+          proc.kill("SIGTERM");
+          setTimeout(() => {
+            if (!proc.killed) proc.kill("SIGKILL");
+          }, SIGKILL_TIMEOUT_MS);
+        };
+        if (signal.aborted) kill();
+        else signal.addEventListener("abort", kill, { once: true });
+      }
+    });
+
+    result.exitCode = exitCode;
+    if (wasAborted) {
+      result.exitCode = 130;
+      result.stopReason = "aborted";
+      result.errorMessage = "Subagent was aborted.";
+      if (!result.stderr.trim()) result.stderr = "Subagent was aborted.";
+    }
+    return result;
+  } finally {
+    cleanupTempDir(promptTmpDir);
+    cleanupTempDir(forkSessionTmpDir);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Concurrency helper
+// ---------------------------------------------------------------------------
+
+/**
+ * Map over items with a bounded number of concurrent workers.
+ */
+export async function mapConcurrent<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 worker = async () => {
+    while (true) {
+      const i = nextIndex++;
+      if (i >= items.length) return;
+      results[i] = await fn(items[i], i);
+    }
+  };
+
+  await Promise.all(Array.from({ length: limit }, () => worker()));
+  return results;
+}

package/types.ts

@@ -0,0 +1,135 @@
+/**
+ * Shared type definitions for the subagent extension.
+ */
+
+import type { Message } from "@mariozechner/pi-ai";
+
+/** Context mode for delegated runs. */
+export type DelegationMode = "spawn" | "fork";
+
+/** Default context mode for delegated runs. */
+export const DEFAULT_DELEGATION_MODE: DelegationMode = "spawn";
+
+/** Aggregated token usage from a subagent run. */
+export interface UsageStats {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	cost: number;
+	contextTokens: number;
+	turns: number;
+}
+
+/** Result of a single subagent invocation. */
+export interface SingleResult {
+	agent: string;
+	agentSource: "user" | "project" | "builtin" | "unknown";
+	task: string;
+	exitCode: number;
+	messages: Message[];
+	stderr: string;
+	usage: UsageStats;
+	model?: string;
+	stopReason?: string;
+	errorMessage?: string;
+}
+
+/** Metadata attached to every tool result for rendering. */
+export interface SubagentDetails {
+	mode: "single" | "parallel";
+	delegationMode: DelegationMode;
+	projectAgentsDir: string | null;
+	results: SingleResult[];
+}
+
+/** Nested subagent tool result captured from a delegated run. */
+export interface NestedSubagentResult {
+	details: SubagentDetails;
+	isError: boolean;
+	toolCallId: string;
+}
+
+/** A display-friendly representation of a message part. */
+export type DisplayItem =
+	| { type: "text"; text: string }
+	| { type: "toolCall"; name: string; args: Record<string, unknown> };
+
+/** Create an empty UsageStats object. */
+export function emptyUsage(): UsageStats {
+	return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 };
+}
+
+/** Sum usage across multiple results. */
+export function aggregateUsage(results: SingleResult[]): UsageStats {
+	const total = emptyUsage();
+	for (const r of results) {
+		total.input += r.usage.input;
+		total.output += r.usage.output;
+		total.cacheRead += r.usage.cacheRead;
+		total.cacheWrite += r.usage.cacheWrite;
+		total.cost += r.usage.cost;
+		total.turns += r.usage.turns;
+	}
+	return total;
+}
+
+/** Whether a result represents an error. */
+export function isResultError(r: SingleResult): boolean {
+	return r.exitCode > 0 || r.stopReason === "error" || r.stopReason === "aborted";
+}
+
+/** Check whether a value looks like SubagentDetails. */
+export function isSubagentDetails(value: unknown): value is SubagentDetails {
+	if (!value || typeof value !== "object") return false;
+	const maybe = value as Partial<SubagentDetails>;
+	return (maybe.mode === "single" || maybe.mode === "parallel") &&
+		(maybe.delegationMode === "spawn" || maybe.delegationMode === "fork") &&
+		Array.isArray(maybe.results);
+}
+
+/** Extract the last assistant text from a message history. */
+export function getFinalOutput(messages: Message[]): string {
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const msg = messages[i];
+		if (msg.role === "assistant") {
+			for (const part of msg.content) {
+				if (part.type === "text") return part.text;
+			}
+		}
+	}
+	return "";
+}
+
+/** Extract all display-worthy items from a message history. */
+export function getDisplayItems(messages: Message[]): DisplayItem[] {
+	const items: DisplayItem[] = [];
+	for (const msg of messages) {
+		if (msg.role === "assistant") {
+			for (const part of msg.content) {
+				if (part.type === "text") {
+					items.push({ type: "text", text: part.text });
+				} else if (part.type === "toolCall") {
+					items.push({ type: "toolCall", name: part.name, args: part.arguments });
+				}
+			}
+		}
+	}
+	return items;
+}
+
+/** Extract nested subagent tool results from a message history. */
+export function getNestedSubagentResults(messages: Message[]): NestedSubagentResult[] {
+	const results: NestedSubagentResult[] = [];
+	for (const msg of messages) {
+		if (msg.role !== "toolResult") continue;
+		if (msg.toolName !== "subagent") continue;
+		if (!isSubagentDetails(msg.details)) continue;
+		results.push({
+			details: msg.details,
+			isError: msg.isError,
+			toolCallId: msg.toolCallId,
+		});
+	}
+	return results;
+}