pi-dynamic-workflow 0.1.0

package/src/sandbox.ts

@@ -0,0 +1,128 @@
+/**
+ * Evaluates the LLM-authored workflow script body in a `node:vm` context with
+ * the orchestration hooks injected.
+ */
+
+import * as vm from "node:vm";
+import type { ScriptHooks } from "./types.ts";
+
+export class WorkflowScriptError extends Error {
+	constructor(kind: "syntax" | "runtime", cause: unknown) {
+		const err = cause instanceof Error ? cause : new Error(String(cause));
+		const stackLine = err.stack?.split("\n").find((line) => line.includes("<workflow>")) ?? "";
+		super(`Workflow script ${kind} error: ${err.message}${stackLine ? `\n${stackLine.trim()}` : ""}`);
+		this.name = "WorkflowScriptError";
+	}
+}
+
+/**
+ * Cap on the script's *initial synchronous slice* (code up to the first await).
+ * Catches the common LLM mistake of a top-level busy loop (`while (true) {}`)
+ * before it can wedge pi's event loop. Continuations after an await run as host
+ * microtasks and are NOT covered — that residual hang requires worker threads
+ * to fix and is accepted as a known limitation.
+ */
+const SYNC_SLICE_TIMEOUT_MS = 5000;
+
+/**
+ * Determinism guards: scripts must produce the same agent() call sequence on
+ * every run so a resumed run can replay cached results. Wall-clock and
+ * randomness would silently break that, so they throw with guidance instead.
+ */
+function blocked(name: string): () => never {
+	return () => {
+		throw new Error(
+			`${name} is blocked in workflow scripts: runs must be deterministic so they can be resumed. ` +
+				"Pass timestamps or random seeds in via `args` instead.",
+		);
+	};
+}
+
+const guardedMath = new Proxy(Math, {
+	get(target, prop, receiver) {
+		if (prop === "random") return blocked("Math.random()");
+		return Reflect.get(target, prop, receiver);
+	},
+});
+
+const guardedDate = new Proxy(Date, {
+	construct(target, argArray, newTarget) {
+		if (argArray.length === 0) blocked("new Date() without arguments")();
+		return Reflect.construct(target, argArray, newTarget);
+	},
+	get(target, prop, receiver) {
+		if (prop === "now") return blocked("Date.now()");
+		return Reflect.get(target, prop, receiver);
+	},
+	apply() {
+		return blocked("Date()")();
+	},
+});
+
+/** Run a plain async JS body with the hooks in scope. Returns the script's return value. */
+export async function runWorkflowScript(
+	source: string,
+	hooks: ScriptHooks,
+	options?: { syncTimeoutMs?: number },
+): Promise<unknown> {
+	const sandboxConsole = {
+		log: (...parts: unknown[]) => hooks.log(parts.map(stringify).join(" ")),
+		error: (...parts: unknown[]) => hooks.log(parts.map(stringify).join(" ")),
+		warn: (...parts: unknown[]) => hooks.log(parts.map(stringify).join(" ")),
+		info: (...parts: unknown[]) => hooks.log(parts.map(stringify).join(" ")),
+	};
+
+	// SECURITY NOTE: node:vm provides isolation hygiene, not a security boundary.
+	// The injected hooks are host-realm functions, so a determined script can reach
+	// host globals via e.g. `agent.constructor("return process")()`. This is
+	// acceptable here because the script author is the session LLM, which already
+	// has full tool access through pi itself. Do not treat this context as a
+	// jail for untrusted third-party code.
+	const context = vm.createContext({
+		agent: hooks.agent,
+		parallel: hooks.parallel,
+		pipeline: hooks.pipeline,
+		phase: hooks.phase,
+		log: hooks.log,
+		args: hooks.args,
+		budget: hooks.budget,
+		workflow: hooks.workflow,
+		console: sandboxConsole,
+		setTimeout,
+		clearTimeout,
+		// Shadow the context's realm intrinsics with the determinism guards.
+		Math: guardedMath,
+		Date: guardedDate,
+	});
+
+	let script: vm.Script;
+	try {
+		script = new vm.Script(`(async () => {\n${source}\n})()`, { filename: "<workflow>" });
+	} catch (error) {
+		throw new WorkflowScriptError("syntax", error);
+	}
+
+	try {
+		return await script.runInContext(context, { timeout: options?.syncTimeoutMs ?? SYNC_SLICE_TIMEOUT_MS });
+	} catch (error) {
+		if (error instanceof WorkflowAbortError) throw error;
+		throw new WorkflowScriptError("runtime", error);
+	}
+}
+
+/** Thrown by hooks when the parent tool call is aborted; passes through untouched. */
+export class WorkflowAbortError extends Error {
+	constructor() {
+		super("Workflow aborted");
+		this.name = "WorkflowAbortError";
+	}
+}
+
+function stringify(value: unknown): string {
+	if (typeof value === "string") return value;
+	try {
+		return JSON.stringify(value);
+	} catch {
+		return String(value);
+	}
+}

package/src/scheduler.ts

@@ -0,0 +1,80 @@
+/**
+ * Concurrency control and deterministic combinators for workflow scripts.
+ */
+
+import * as os from "node:os";
+
+export function defaultConcurrency(): number {
+	return Math.max(2, Math.min(8, os.cpus().length - 2));
+}
+
+/** Simple counting semaphore capping concurrent subagent subprocesses. */
+export class Semaphore {
+	private available: number;
+	private waiters: Array<() => void> = [];
+
+	constructor(capacity: number) {
+		this.available = Math.max(1, capacity);
+	}
+
+	async acquire(): Promise<void> {
+		if (this.available > 0) {
+			this.available--;
+			return;
+		}
+		await new Promise<void>((resolve) => this.waiters.push(resolve));
+	}
+
+	release(): void {
+		const next = this.waiters.shift();
+		if (next) next();
+		else this.available++;
+	}
+
+	async run<T>(fn: () => Promise<T>): Promise<T> {
+		await this.acquire();
+		try {
+			return await fn();
+		} finally {
+			this.release();
+		}
+	}
+}
+
+/**
+ * Run all thunks concurrently. A thrown/rejected thunk resolves to `null`
+ * (the batch never rejects), so scripts can `.filter(Boolean)`.
+ */
+export async function parallel(thunks: Array<() => Promise<unknown>>): Promise<unknown[]> {
+	return Promise.all(
+		thunks.map(async (thunk) => {
+			try {
+				return await thunk();
+			} catch {
+				return null;
+			}
+		}),
+	);
+}
+
+export type PipelineStage = (prev: unknown, item: unknown, index: number) => Promise<unknown> | unknown;
+
+/**
+ * Each item flows through all stages independently with no cross-item barrier.
+ * A throwing stage drops the item to `null` and skips its remaining stages.
+ */
+export async function pipeline(items: unknown[], ...stages: PipelineStage[]): Promise<unknown[]> {
+	return Promise.all(
+		items.map(async (item, index) => {
+			let value: unknown = item;
+			for (const stage of stages) {
+				try {
+					value = await stage(value, item, index);
+				} catch {
+					return null;
+				}
+			}
+			return value;
+		}),
+	);
+}

package/src/structured.ts

@@ -0,0 +1,73 @@
+/**
+ * Schema-enforced structured output for subagents.
+ *
+ * Generates a temporary pi extension that registers a terminating `emit_result`
+ * tool whose parameters are the caller-supplied JSON schema. pi validates tool
+ * arguments against the schema (and retries the model on validation errors), so
+ * the parent only needs to extract the final `emit_result` call arguments.
+ */
+
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+
+export const EMIT_RESULT_TOOL = "emit_result";
+
+export interface StructuredOutputExtension {
+	/** Path to the generated extension file (pass to `pi -e`). */
+	path: string;
+	/** Remove the temp file and directory. */
+	cleanup(): void;
+}
+
+export async function createStructuredOutputExtension(
+	schema: Record<string, unknown>,
+): Promise<StructuredOutputExtension> {
+	const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-workflow-schema-"));
+	const filePath = path.join(dir, "structured-output.ts");
+
+	const source = `import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+
+const schema = ${JSON.stringify(schema, null, "\t")} as const;
+
+export default function (pi: ExtensionAPI) {
+	pi.registerTool({
+		name: ${JSON.stringify(EMIT_RESULT_TOOL)},
+		label: "Emit Result",
+		description:
+			"Emit the final structured result for this task. Call this exactly once, as your last action, with the complete answer matching the required schema.",
+		promptSnippet: "Emit the final structured result (required last action)",
+		promptGuidelines: [
+			"You MUST finish by calling ${EMIT_RESULT_TOOL} exactly once with the final answer matching its parameter schema. Do not answer in plain text.",
+		],
+		parameters: Type.Unsafe(schema),
+		async execute(_toolCallId, params) {
+			return {
+				content: [{ type: "text", text: "Structured result recorded." }],
+				details: params,
+				terminate: true,
+			};
+		},
+	});
+}
+`;
+
+	await fs.promises.writeFile(filePath, source, { encoding: "utf-8", mode: 0o600 });
+
+	return {
+		path: filePath,
+		cleanup() {
+			try {
+				fs.unlinkSync(filePath);
+			} catch {
+				/* ignore */
+			}
+			try {
+				fs.rmdirSync(dir);
+			} catch {
+				/* ignore */
+			}
+		},
+	};
+}

package/src/subagent.ts

@@ -0,0 +1,304 @@
+/**
+ * Runs a single pi subagent as a subprocess (`pi --mode json -p --no-session`)
+ * and parses its newline-delimited JSON event stream.
+ */
+
+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 { AssistantMessage, Message } from "@earendil-works/pi-ai";
+import { createStructuredOutputExtension, EMIT_RESULT_TOOL } from "./structured.ts";
+import { addUsage, emptyUsage, type UsageStats } from "./types.ts";
+
+export interface SubagentRequest {
+	prompt: string;
+	model?: string;
+	tools?: string[];
+	cwd: string;
+	schema?: Record<string, unknown>;
+	signal?: AbortSignal;
+	/** Wall-clock cap in ms; on expiry the subprocess is killed and the result is a non-abort failure. */
+	timeoutMs?: number;
+	/** Replaces the subagent's system prompt (`pi --system-prompt`). */
+	systemPrompt?: string;
+	/** Appended to the system prompt in order (`pi --append-system-prompt`, via temp files). */
+	appendSystemPrompt?: string[];
+	/** Called on each parsed event so the parent can stream progress. */
+	onEvent?: (update: SubagentProgress) => void;
+}
+
+export interface SubagentProgress {
+	/** Short human-readable description of the latest activity. */
+	activity: string;
+	usage: UsageStats;
+}
+
+export interface SubagentResult {
+	ok: boolean;
+	/** Last assistant text output. */
+	outputText: string;
+	/** Parsed `emit_result` arguments when a schema was supplied. */
+	structured?: unknown;
+	usage: UsageStats;
+	model?: string;
+	stopReason?: string;
+	errorMessage?: string;
+	exitCode: number;
+	aborted: boolean;
+}
+
+/** Env var counting nested subagent depth to prevent runaway recursion (fork bombs). */
+const DEPTH_ENV = "PI_WORKFLOW_DEPTH";
+const MAX_SUBAGENT_DEPTH = 3;
+
+/**
+ * Resolve how to invoke pi.
+ *
+ * Unlike the official subagent example, this only reuses `process.argv[1]`
+ * when it plausibly IS pi's CLI entry point. Blindly re-invoking argv[1]
+ * would fork-bomb the machine if this module is ever loaded outside pi
+ * (e.g. from a test harness: argv[1] is the harness, which re-runs itself).
+ */
+function getPiInvocation(args: string[]): { command: string; args: string[] } {
+	const currentScript = process.argv[1];
+	const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
+	const looksLikePiEntry = currentScript !== undefined && /^(pi|cli\.(js|ts|mjs))$/.test(path.basename(currentScript));
+	if (currentScript && looksLikePiEntry && !isBunVirtualScript && fs.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 };
+	}
+
+	return { command: "pi", args };
+}
+
+function lastAssistantText(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" && part.text.trim()) return part.text;
+			}
+		}
+	}
+	return "";
+}
+
+function lastEmitResultArguments(messages: Message[]): unknown {
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const msg = messages[i];
+		if (msg?.role !== "assistant") continue;
+		for (let j = msg.content.length - 1; j >= 0; j--) {
+			const part = msg.content[j];
+			if (part?.type === "toolCall" && part.name === EMIT_RESULT_TOOL) {
+				return part.arguments;
+			}
+		}
+	}
+	return undefined;
+}
+
+export async function runSubagent(request: SubagentRequest): Promise<SubagentResult> {
+	const { prompt, model, tools, cwd, schema, signal, timeoutMs, systemPrompt, appendSystemPrompt, onEvent } = request;
+
+	// Hard recursion guard: a subagent that itself spawns workflows could
+	// otherwise multiply subprocesses without bound.
+	const depth = Number(process.env[DEPTH_ENV] ?? "0") || 0;
+	if (depth >= MAX_SUBAGENT_DEPTH) {
+		return {
+			ok: false,
+			outputText: "",
+			usage: emptyUsage(),
+			exitCode: 1,
+			aborted: false,
+			errorMessage: `Subagent nesting depth limit (${MAX_SUBAGENT_DEPTH}) reached; refusing to spawn`,
+		};
+	}
+
+	const args: string[] = ["--mode", "json", "-p", "--no-session"];
+	if (model) args.push("--model", model);
+	if (tools && tools.length > 0) {
+		const toolList = schema ? [...new Set([...tools, EMIT_RESULT_TOOL])] : tools;
+		args.push("--tools", toolList.join(","));
+	}
+	if (systemPrompt) args.push("--system-prompt", systemPrompt);
+
+	// --append-system-prompt accepts text or a file path; prompt text that
+	// happens to look like a path would be misread, so always pass a temp file.
+	let promptDir: string | null = null;
+	if (appendSystemPrompt && appendSystemPrompt.length > 0) {
+		promptDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-workflow-sysprompt-"));
+		appendSystemPrompt.forEach((text, i) => {
+			const promptPath = path.join(promptDir as string, `append-${i}.md`);
+			fs.writeFileSync(promptPath, text, "utf-8");
+			args.push("--append-system-prompt", promptPath);
+		});
+	}
+
+	let structuredExt: Awaited<ReturnType<typeof createStructuredOutputExtension>> | null = null;
+	if (schema) {
+		structuredExt = await createStructuredOutputExtension(schema);
+		args.push("-e", structuredExt.path);
+	}
+	args.push(prompt);
+
+	const messages: Message[] = [];
+	const usage = emptyUsage();
+	const result: SubagentResult = {
+		ok: false,
+		outputText: "",
+		usage,
+		exitCode: 0,
+		aborted: false,
+	};
+	let stderr = "";
+	let timedOut = false;
+
+	try {
+		result.exitCode = await new Promise<number>((resolve) => {
+			const invocation = getPiInvocation(args);
+			const proc = spawn(invocation.command, invocation.args, {
+				cwd,
+				shell: false,
+				stdio: ["ignore", "pipe", "pipe"],
+				env: { ...process.env, [DEPTH_ENV]: String(depth + 1) },
+			});
+			let buffer = "";
+
+			const processLine = (line: string) => {
+				if (!line.trim()) return;
+				let event: { type?: string; message?: Message; toolName?: string };
+				try {
+					event = JSON.parse(line);
+				} catch {
+					return;
+				}
+
+				if (event.type === "message_end" && event.message) {
+					const msg = event.message;
+					messages.push(msg);
+					if (msg.role === "assistant") {
+						const assistant = msg as AssistantMessage;
+						usage.turns++;
+						if (assistant.usage) {
+							addUsage(usage, {
+								input: assistant.usage.input || 0,
+								output: assistant.usage.output || 0,
+								cacheRead: assistant.usage.cacheRead || 0,
+								cacheWrite: assistant.usage.cacheWrite || 0,
+								cost: assistant.usage.cost?.total || 0,
+								turns: 0,
+							});
+						}
+						if (!result.model && assistant.model) result.model = assistant.model;
+						if (assistant.stopReason) result.stopReason = assistant.stopReason;
+						if (assistant.errorMessage) result.errorMessage = assistant.errorMessage;
+
+						const text = lastAssistantText([msg]);
+						onEvent?.({ activity: text ? text.slice(0, 120) : "thinking...", usage: { ...usage } });
+					}
+				} else if (event.type === "tool_execution_start" && event.toolName) {
+					onEvent?.({ activity: `→ ${event.toolName}`, usage: { ...usage } });
+				}
+			};
+
+			proc.stdout.on("data", (data) => {
+				buffer += data.toString();
+				const lines = buffer.split("\n");
+				buffer = lines.pop() || "";
+				for (const line of lines) processLine(line);
+			});
+
+			proc.stderr.on("data", (data) => {
+				stderr += data.toString();
+			});
+
+			// NOTE: `proc.killed` only means a signal was *sent*, so it cannot
+			// detect a child that ignores SIGTERM; check exit status instead.
+			let killedByUs = false;
+			const terminate = () => {
+				killedByUs = true;
+				proc.kill("SIGTERM");
+				setTimeout(() => {
+					if (proc.exitCode === null && proc.signalCode === null) proc.kill("SIGKILL");
+				}, 5000).unref();
+			};
+
+			let timeoutTimer: NodeJS.Timeout | undefined;
+			if (timeoutMs && timeoutMs > 0) {
+				timeoutTimer = setTimeout(() => {
+					timedOut = true;
+					terminate();
+				}, timeoutMs);
+				timeoutTimer.unref();
+			}
+
+			proc.on("close", (code, killSignal) => {
+				if (timeoutTimer) clearTimeout(timeoutTimer);
+				if (buffer.trim()) processLine(buffer);
+				// A null code means signal death. Our own kills are reported via
+				// the aborted/timedOut flags; an external kill (OOM killer, manual
+				// SIGKILL) must read as failure, not as a success with partial output.
+				if (code === null && !killedByUs) {
+					result.errorMessage ??= `pi killed by signal ${killSignal ?? "unknown"}`;
+					resolve(1);
+					return;
+				}
+				resolve(code ?? 0);
+			});
+
+			proc.on("error", () => {
+				resolve(1);
+			});
+
+			if (signal) {
+				const killProc = () => {
+					result.aborted = true;
+					terminate();
+				};
+				if (signal.aborted) killProc();
+				else signal.addEventListener("abort", killProc, { once: true });
+			}
+		});
+	} finally {
+		structuredExt?.cleanup();
+		if (promptDir) fs.rmSync(promptDir, { recursive: true, force: true });
+	}
+
+	result.outputText = lastAssistantText(messages);
+
+	if (result.aborted) {
+		result.errorMessage ??= "Subagent aborted";
+		return result;
+	}
+
+	// A timed-out child dies via SIGTERM (close code null → 0), so without this
+	// check it could masquerade as a successful run with partial output.
+	if (timedOut) {
+		result.errorMessage = `Subagent timed out after ${timeoutMs}ms`;
+		return result;
+	}
+
+	const failed = result.exitCode !== 0 || result.stopReason === "error" || result.stopReason === "aborted";
+	if (failed) {
+		result.errorMessage ??= stderr.trim() || `pi exited with code ${result.exitCode}`;
+		return result;
+	}
+
+	if (schema) {
+		const structured = lastEmitResultArguments(messages);
+		if (structured === undefined) {
+			result.errorMessage = `Subagent finished without calling ${EMIT_RESULT_TOOL} (structured output missing)`;
+			return result;
+		}
+		result.structured = structured;
+	}
+
+	result.ok = true;
+	return result;
+}