pi-taskflow 0.0.1

package/extensions/runner.ts

@@ -0,0 +1,301 @@
+/**
+ * Subagent runner — spawns an isolated `pi --mode json -p` process for a single
+ * task and collects its structured output and usage. Adapted from the pi
+ * subagent extension's runSingleAgent.
+ */
+
+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 { Message } from "@earendil-works/pi-ai";
+import { withFileMutationQueue } from "@earendil-works/pi-coding-agent";
+import type { AgentConfig } from "./agents.ts";
+
+export interface UsageStats {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	cost: number;
+	contextTokens: number;
+	turns: number;
+}
+
+export function emptyUsage(): UsageStats {
+	return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 };
+}
+
+export interface RunResult {
+	agent: string;
+	task: string;
+	exitCode: number;
+	output: string;
+	stderr: string;
+	usage: UsageStats;
+	model?: string;
+	stopReason?: string;
+	errorMessage?: string;
+}
+
+export interface RunOptions {
+	model?: string;
+	thinking?: string;
+	tools?: string[];
+	cwd?: string;
+	signal?: AbortSignal;
+	onText?: (text: string) => void;
+}
+
+export function isFailed(r: RunResult): boolean {
+	return r.exitCode !== 0 || r.stopReason === "error" || r.stopReason === "aborted";
+}
+
+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 "";
+}
+
+async function writePromptToTempFile(agentName: string, prompt: string): Promise<{ dir: string; filePath: string }> {
+	const tmpDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-taskflow-"));
+	const safeName = agentName.replace(/[^\w.-]+/g, "_");
+	const filePath = path.join(tmpDir, `prompt-${safeName}.md`);
+	await withFileMutationQueue(filePath, async () => {
+		await fs.promises.writeFile(filePath, prompt, { encoding: "utf-8", mode: 0o600 });
+	});
+	return { dir: tmpDir, filePath };
+}
+
+function getPiInvocation(args: string[]): { command: string; args: string[] } {
+	// Explicit override (used by tests and unusual launch setups).
+	const override = process.env.PI_TASKFLOW_PI_BIN;
+	if (override) return { command: override, args };
+
+	const currentScript = process.argv[1];
+	const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
+	// Only re-exec the current script if it actually looks like the pi CLI entry.
+	const looksLikePi = currentScript ? /(?:^|[\\/])(?:cli|pi)\.(?:js|mjs|cjs)$/.test(currentScript) : false;
+	if (currentScript && !isBunVirtualScript && looksLikePi && 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 };
+}
+
+/**
+ * Run a single subagent task. Resolves the agent from `agents` by name and
+ * spawns an isolated pi process, returning structured output + usage.
+ */
+export async function runAgentTask(
+	defaultCwd: string,
+	agents: AgentConfig[],
+	agentName: string,
+	task: string,
+	opts: RunOptions,
+	globalThinking?: string,
+): Promise<RunResult> {
+	const agent = agents.find((a) => a.name === agentName);
+	if (!agent) {
+		const available = agents.map((a) => `"${a.name}"`).join(", ") || "none";
+		return {
+			agent: agentName,
+			task,
+			exitCode: 1,
+			output: "",
+			stderr: `Unknown agent: "${agentName}". Available: ${available}.`,
+			usage: emptyUsage(),
+			errorMessage: `Unknown agent: ${agentName}`,
+			stopReason: "error",
+		};
+	}
+
+	const model = opts.model ?? agent.model;
+	const thinking = opts.thinking ?? agent.thinking ?? globalThinking;
+	const tools = opts.tools ?? agent.tools;
+
+	const args: string[] = ["--mode", "json", "-p", "--no-session"];
+	if (model) args.push("--model", model);
+	if (thinking) args.push("--thinking", thinking);
+	if (tools && tools.length > 0) args.push("--tools", tools.join(","));
+
+	let tmpPromptDir: string | null = null;
+	let tmpPromptPath: string | null = null;
+
+	const messages: Message[] = [];
+	const result: RunResult = {
+		agent: agentName,
+		task,
+		exitCode: 0,
+		output: "",
+		stderr: "",
+		usage: emptyUsage(),
+		model,
+	};
+
+	try {
+		if (agent.systemPrompt.trim()) {
+			const tmp = await writePromptToTempFile(agent.name, agent.systemPrompt);
+			tmpPromptDir = tmp.dir;
+			tmpPromptPath = tmp.filePath;
+			args.push("--append-system-prompt", tmpPromptPath);
+		}
+		args.push(`Task: ${task}`);
+
+		let wasAborted = false;
+		const exitCode = await new Promise<number>((resolve) => {
+			const invocation = getPiInvocation(args);
+			const proc = spawn(invocation.command, invocation.args, {
+				cwd: opts.cwd ?? defaultCwd,
+				shell: false,
+				stdio: ["ignore", "pipe", "pipe"],
+			});
+			let buffer = "";
+
+			const processLine = (line: string) => {
+				if (!line.trim()) return;
+				let event: any;
+				try {
+					event = JSON.parse(line);
+				} catch {
+					return;
+				}
+				if (event.type === "message_end" && event.message) {
+					const msg = event.message as Message;
+					messages.push(msg);
+					if (msg.role === "assistant") {
+						result.usage.turns++;
+						const u = (msg as any).usage;
+						if (u) {
+							result.usage.input += u.input || 0;
+							result.usage.output += u.output || 0;
+							result.usage.cacheRead += u.cacheRead || 0;
+							result.usage.cacheWrite += u.cacheWrite || 0;
+							result.usage.cost += u.cost?.total || 0;
+							result.usage.contextTokens = u.totalTokens || 0;
+						}
+						if (!result.model && (msg as any).model) result.model = (msg as any).model;
+						if ((msg as any).stopReason) result.stopReason = (msg as any).stopReason;
+						if ((msg as any).errorMessage) result.errorMessage = (msg as any).errorMessage;
+						const text = getFinalOutput([msg]);
+						if (text && opts.onText) opts.onText(text);
+					}
+				}
+			};
+
+			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) => {
+				result.stderr += data.toString();
+			});
+			proc.on("close", (code) => {
+				if (buffer.trim()) processLine(buffer);
+				resolve(code ?? 0);
+			});
+			proc.on("error", () => resolve(1));
+
+			if (opts.signal) {
+				const kill = () => {
+					wasAborted = true;
+					proc.kill("SIGTERM");
+					setTimeout(() => {
+						if (!proc.killed) proc.kill("SIGKILL");
+					}, 5000);
+				};
+				if (opts.signal.aborted) kill();
+				else opts.signal.addEventListener("abort", kill, { once: true });
+			}
+		});
+
+		result.exitCode = exitCode;
+		result.output = getFinalOutput(messages);
+		if (wasAborted) {
+			result.stopReason = "aborted";
+			result.errorMessage = "Subagent was aborted";
+		}
+		if (isFailed(result) && !result.output) {
+			result.output = result.errorMessage || result.stderr || "(no output)";
+		}
+		return result;
+	} finally {
+		if (tmpPromptPath) {
+			try {
+				fs.unlinkSync(tmpPromptPath);
+			} catch {
+				/* ignore */
+			}
+		}
+		if (tmpPromptDir) {
+			try {
+				fs.rmdirSync(tmpPromptDir);
+			} catch {
+				/* ignore */
+			}
+		}
+	}
+}
+
+/** Run an array of items through `fn` with a bounded concurrency pool. */
+export async function mapWithConcurrencyLimit<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 = new Array(limit).fill(null).map(async () => {
+		while (true) {
+			const current = nextIndex++;
+			if (current >= items.length) return;
+			results[current] = await fn(items[current], current);
+		}
+	});
+	await Promise.all(workers);
+	return results;
+}
+
+export function formatTokens(count: number): string {
+	if (count < 1000) return count.toString();
+	if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
+	if (count < 1000000) return `${Math.round(count / 1000)}k`;
+	return `${(count / 1000000).toFixed(1)}M`;
+}
+
+export function formatUsage(usage: UsageStats, model?: string): string {
+	const parts: string[] = [];
+	if (usage.turns) parts.push(`${usage.turns} turn${usage.turns > 1 ? "s" : ""}`);
+	if (usage.input) parts.push(`↑${formatTokens(usage.input)}`);
+	if (usage.output) parts.push(`↓${formatTokens(usage.output)}`);
+	if (usage.cacheRead) parts.push(`R${formatTokens(usage.cacheRead)}`);
+	if (usage.cost) parts.push(`$${usage.cost.toFixed(4)}`);
+	if (model) parts.push(model);
+	return parts.join(" ");
+}
+
+export function aggregateUsage(usages: UsageStats[]): UsageStats {
+	const total = emptyUsage();
+	for (const u of usages) {
+		total.input += u.input;
+		total.output += u.output;
+		total.cacheRead += u.cacheRead;
+		total.cacheWrite += u.cacheWrite;
+		total.cost += u.cost;
+		total.turns += u.turns;
+	}
+	return total;
+}

package/extensions/runtime.ts

@@ -0,0 +1,298 @@
+/**
+ * Taskflow runtime — the orchestration engine.
+ *
+ * Resolves the phase DAG into topological layers and executes each phase by
+ * delegating to isolated subagents. Intermediate phase outputs live here (in
+ * RunState) and never enter the host conversation's context window — only the
+ * final phase output is returned to the caller.
+ *
+ * Supports resume: phases whose resolved input hash matches a cached completed
+ * result are skipped.
+ */
+
+import type { AgentConfig } from "./agents.ts";
+import { coerceArray, interpolate, type InterpolationContext, safeParse } from "./interpolate.ts";
+import { aggregateUsage, emptyUsage, isFailed, mapWithConcurrencyLimit, runAgentTask, type RunResult, type UsageStats } from "./runner.ts";
+import { dependenciesOf, finalPhase, type Phase, type Taskflow, topoLayers } from "./schema.ts";
+import { hashInput, type PhaseState, type RunState } from "./store.ts";
+
+export interface RuntimeDeps {
+	cwd: string;
+	agents: AgentConfig[];
+	globalThinking?: string;
+	signal?: AbortSignal;
+	/** Persist run state after each phase (for resume). */
+	persist?: (state: RunState) => void;
+	/** Live progress callback for TUI streaming. */
+	onProgress?: (state: RunState) => void;
+	/** Injectable task runner (defaults to spawning a real subagent). Enables testing. */
+	runTask?: typeof runAgentTask;
+}
+
+export interface RuntimeResult {
+	state: RunState;
+	finalOutput: string;
+	ok: boolean;
+	totalUsage: UsageStats;
+}
+
+function buildInterpolationContext(
+	state: RunState,
+	previousOutput: string | undefined,
+	locals?: Record<string, unknown>,
+): InterpolationContext {
+	const steps: Record<string, { output: string; json?: unknown }> = {};
+	for (const [id, ps] of Object.entries(state.phases)) {
+		if (ps.status === "done" && ps.output !== undefined) {
+			steps[id] = { output: ps.output, json: ps.json };
+		}
+	}
+	return { args: state.args, steps, previousOutput, locals };
+}
+
+function resultToPhaseState(id: string, r: RunResult, inputHash: string, parseJson: boolean): PhaseState {
+	const failed = isFailed(r);
+	return {
+		id,
+		status: failed ? "failed" : "done",
+		output: r.output,
+		json: parseJson && !failed ? safeParse(r.output) : undefined,
+		usage: r.usage,
+		model: r.model,
+		error: failed ? r.errorMessage || r.stderr || r.output : undefined,
+		inputHash,
+		endedAt: Date.now(),
+	};
+}
+
+/** Merge several sub-results into a single PhaseState (for map/parallel). */
+function mergePhaseState(
+	id: string,
+	results: RunResult[],
+	inputHash: string,
+	parseJson: boolean,
+): PhaseState {
+	const anyFailed = results.some(isFailed);
+	const usage = aggregateUsage(results.map((r) => r.usage));
+	// Combine outputs as a labelled list; also expose a JSON array of outputs.
+	const combinedText = results
+		.map((r, i) => `### [${i + 1}/${results.length}] ${r.agent}${isFailed(r) ? " (failed)" : ""}\n\n${r.output}`)
+		.join("\n\n---\n\n");
+	const jsonArray = parseJson ? results.map((r) => safeParse(r.output) ?? r.output) : undefined;
+	return {
+		id,
+		status: anyFailed ? "failed" : "done",
+		output: combinedText,
+		json: jsonArray,
+		usage,
+		error: anyFailed ? results.filter(isFailed).map((r) => `${r.agent}: ${r.errorMessage ?? r.stderr}`).join("; ") : undefined,
+		inputHash,
+		endedAt: Date.now(),
+	};
+}
+
+async function executePhase(
+	phase: Phase,
+	state: RunState,
+	deps: RuntimeDeps,
+	prior: PhaseState | undefined,
+): Promise<PhaseState> {
+	const type = phase.type ?? "agent";
+	const concurrency = phase.concurrency ?? state.def.concurrency ?? 8;
+	const previousOutput = lastCompletedOutput(state, phase);
+	const run = deps.runTask ?? runAgentTask;
+
+	const runOne = (agentName: string, task: string, _locals?: Record<string, unknown>) =>
+		run(
+			deps.cwd,
+			deps.agents,
+			agentName,
+			task,
+			{
+				model: phase.model,
+				thinking: phase.thinking,
+				tools: phase.tools,
+				cwd: phase.cwd,
+				signal: deps.signal,
+			},
+			deps.globalThinking,
+		);
+
+	const parseJson = phase.output === "json";
+
+	if (type === "agent" || type === "gate") {
+		const ctx = buildInterpolationContext(state, previousOutput);
+		const { text } = interpolate(phase.task ?? "", ctx);
+		const inputHash = hashInput(phase.id, phase.agent ?? "", text);
+		const cached = cachedPhase(prior, inputHash);
+		if (cached) return cached;
+
+		const r = await runOne(phase.agent ?? defaultAgent(deps), text);
+		return resultToPhaseState(phase.id, r, inputHash, parseJson);
+	}
+
+	if (type === "parallel") {
+		const ctx = buildInterpolationContext(state, previousOutput);
+		const branches = (phase.branches ?? []).map((b) => ({
+			agent: b.agent ?? phase.agent ?? defaultAgent(deps),
+			task: interpolate(b.task, ctx).text,
+		}));
+		const inputHash = hashInput(phase.id, JSON.stringify(branches));
+		const cached = cachedPhase(prior, inputHash);
+		if (cached) return cached;
+
+		const results = await mapWithConcurrencyLimit(branches, concurrency, (b) => runOne(b.agent, b.task));
+		return mergePhaseState(phase.id, results, inputHash, parseJson);
+	}
+
+	if (type === "map") {
+		const ctx = buildInterpolationContext(state, previousOutput);
+		const overResolved = interpolate(phase.over ?? "", ctx).text;
+		// `over` may itself be a placeholder that resolved to a JSON string.
+		const arr = coerceArray(safeParse(overResolved)) ?? coerceArray(directRef(phase.over ?? "", state));
+		if (!arr) {
+			return {
+				id: phase.id,
+				status: "failed",
+				error: `map phase '${phase.id}': 'over' (${phase.over}) did not resolve to an array`,
+				inputHash: hashInput(phase.id, "no-array"),
+				endedAt: Date.now(),
+				usage: emptyUsage(),
+			};
+		}
+		const loopVar = phase.as ?? "item";
+		const tasks = arr.map((item) => {
+			const localCtx = buildInterpolationContext(state, previousOutput, { [loopVar]: item });
+			return {
+				agent: phase.agent ?? defaultAgent(deps),
+				task: interpolate(phase.task ?? "", localCtx).text,
+			};
+		});
+		const inputHash = hashInput(phase.id, JSON.stringify(tasks));
+		const cached = cachedPhase(prior, inputHash);
+		if (cached) return cached;
+
+		const results = await mapWithConcurrencyLimit(tasks, concurrency, (t) => runOne(t.agent, t.task));
+		return mergePhaseState(phase.id, results, inputHash, parseJson);
+	}
+
+	if (type === "reduce") {
+		const ctx = buildInterpolationContext(state, previousOutput);
+		// Inputs for reduce come from `from` phases; interpolation already exposes them.
+		const { text } = interpolate(phase.task ?? "", ctx);
+		const inputHash = hashInput(phase.id, text);
+		const cached = cachedPhase(prior, inputHash);
+		if (cached) return cached;
+
+		const r = await runOne(phase.agent ?? defaultAgent(deps), text);
+		return resultToPhaseState(phase.id, r, inputHash, parseJson);
+	}
+
+	return {
+		id: phase.id,
+		status: "failed",
+		error: `Unknown phase type: ${type}`,
+		endedAt: Date.now(),
+		usage: emptyUsage(),
+	};
+}
+
+/** Resolve a `{steps.x.json}`-style ref directly to its parsed value (bypassing stringify). */
+function directRef(over: string, state: RunState): unknown {
+	const m = over.match(/^\{steps\.([a-zA-Z0-9_]+)\.(output|json)\}$/);
+	if (!m) return undefined;
+	const step = state.phases[m[1]];
+	if (!step || step.status !== "done") return undefined;
+	if (m[2] === "json") return step.json ?? safeParse(step.output ?? "");
+	return safeParse(step.output ?? "");
+}
+
+function lastCompletedOutput(state: RunState, phase: Phase): string | undefined {
+	const deps = dependenciesOf(phase);
+	for (let i = deps.length - 1; i >= 0; i--) {
+		const ps = state.phases[deps[i]];
+		if (ps?.status === "done") return ps.output;
+	}
+	return undefined;
+}
+
+function cachedPhase(prior: PhaseState | undefined, inputHash: string): PhaseState | null {
+	if (prior && prior.status === "done" && prior.inputHash === inputHash) {
+		return { ...prior, status: "done" };
+	}
+	return null;
+}
+
+function defaultAgent(deps: RuntimeDeps): string {
+	return deps.agents[0]?.name ?? "default";
+}
+
+/**
+ * Execute a full taskflow. Mutates and persists `state` as it progresses.
+ */
+export async function executeTaskflow(state: RunState, deps: RuntimeDeps): Promise<RuntimeResult> {
+	const def: Taskflow = state.def;
+	const layers = topoLayers(def.phases);
+
+	state.status = "running";
+	deps.persist?.(state);
+	deps.onProgress?.(state);
+
+	let aborted = false;
+
+	for (const layer of layers) {
+		if (deps.signal?.aborted) {
+			aborted = true;
+			break;
+		}
+		// Phases within a layer have no inter-dependencies → run concurrently.
+		const layerConcurrency = Math.max(1, def.concurrency ?? 8);
+		await mapWithConcurrencyLimit(layer, layerConcurrency, async (phase) => {
+			// Snapshot prior state BEFORE marking running, so resume cache checks work.
+			const prior = state.phases[phase.id];
+			// Skip if a dependency failed (unless this phase is optional).
+			const failedDep = dependenciesOf(phase).some((d) => state.phases[d]?.status === "failed");
+			if (failedDep) {
+				state.phases[phase.id] = {
+					id: phase.id,
+					status: "skipped",
+					error: "Upstream dependency failed",
+					endedAt: Date.now(),
+					usage: emptyUsage(),
+				};
+				deps.persist?.(state);
+				deps.onProgress?.(state);
+				return;
+			}
+
+			state.phases[phase.id] = {
+				...(state.phases[phase.id] ?? { id: phase.id }),
+				id: phase.id,
+				status: "running",
+				startedAt: Date.now(),
+			};
+			deps.onProgress?.(state);
+
+			const ps = await executePhase(phase, state, deps, prior);
+			state.phases[phase.id] = ps;
+			deps.persist?.(state);
+			deps.onProgress?.(state);
+		});
+	}
+
+	const fp = finalPhase(def.phases);
+	const finalState = state.phases[fp.id];
+	const anyFailed = Object.values(state.phases).some((p) => p.status === "failed");
+
+	state.status = aborted ? "paused" : anyFailed ? "failed" : "completed";
+	deps.persist?.(state);
+	deps.onProgress?.(state);
+
+	const totalUsage = aggregateUsage(Object.values(state.phases).map((p) => p.usage ?? emptyUsage()));
+	return {
+		state,
+		finalOutput: finalState?.output ?? "(no output)",
+		ok: state.status === "completed",
+		totalUsage,
+	};
+}