@oh-my-pi/pi-coding-agent 3.20.1 → 3.24.0

package/src/core/tools/task/executor.ts

@@ -1,28 +1,27 @@
 /**
- * Subprocess execution for subagents.
+ * Worker execution for subagents.
  *
- * Spawns `omp` in JSON mode to execute tasks with isolated context.
- * Parses JSON events for progress tracking.
+ * Runs each subagent in a Bun Worker and forwards AgentEvents for progress tracking.
  */
 
-import { existsSync, unlinkSync, writeFileSync } from "node:fs";
-import { tmpdir } from "node:os";
-import * as path from "node:path";
+import { writeFileSync } from "node:fs";
+import type { AgentEvent } from "@oh-my-pi/pi-agent-core";
+import type { EventBus } from "../../event-bus";
 import { ensureArtifactsDir, getArtifactPaths } from "./artifacts";
 import { resolveModelPattern } from "./model-resolver";
-import { resolveOmpCommand } from "./omp-command";
 import { subprocessToolRegistry } from "./subprocess-tool-registry";
 import {
 	type AgentDefinition,
 	type AgentProgress,
 	MAX_OUTPUT_BYTES,
 	MAX_OUTPUT_LINES,
-	OMP_BLOCKED_AGENT_ENV,
-	OMP_SPAWNS_ENV,
 	type SingleResult,
+	TASK_SUBAGENT_EVENT_CHANNEL,
+	TASK_SUBAGENT_PROGRESS_CHANNEL,
 } from "./types";
+import type { SubagentWorkerRequest, SubagentWorkerResponse } from "./worker-protocol";
 
-/** Options for subprocess execution */
+/** Options for worker execution */
 export interface ExecutorOptions {
 	cwd: string;
 	agent: AgentDefinition;
@@ -36,6 +35,7 @@ export interface ExecutorOptions {
 	sessionFile?: string | null;
 	persistArtifacts?: boolean;
 	artifactsDir?: string;
+	eventBus?: EventBus;
 }
 
 /**
@@ -127,7 +127,7 @@ function getUsageTokens(usage: unknown): number {
 }
 
 /**
- * Run a single agent as a subprocess.
+ * Run a single agent in a worker.
  */
 export async function runSubprocess(options: ExecutorOptions): Promise<SingleResult> {
 	const { cwd, agent, task, index, context, modelOverride, signal, onProgress } = options;
@@ -168,33 +168,6 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 		};
 	}
 
-	// Write system prompt to temp file
-	const tempDir = tmpdir();
-	const promptFile = path.join(
-		tempDir,
-		`omp-agent-${agent.name}-${Date.now()}-${Math.random().toString(36).slice(2)}.md`,
-	);
-
-	try {
-		writeFileSync(promptFile, agent.systemPrompt, "utf-8");
-	} catch (err) {
-		return {
-			index,
-			agent: agent.name,
-			agentSource: agent.source,
-			task,
-			description: options.description,
-			exitCode: 1,
-			output: "",
-			stderr: `Failed to write prompt file: ${err}`,
-			truncated: false,
-			durationMs: Date.now() - startTime,
-			tokens: 0,
-			modelOverride,
-			error: `Failed to write prompt file: ${err}`,
-		};
-	}
-
 	// Build full task with context
 	const fullTask = context ? `${context}\n\n${task}` : task;
 
@@ -215,182 +188,197 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 		}
 	}
 
-	// Build args
-	const args: string[] = ["--mode", "json", "--non-interactive"];
-
-	// Add system prompt
-	args.push("--append-system-prompt", promptFile);
-
 	// Add tools if specified
+	let toolNames: string[] | undefined;
 	if (agent.tools && agent.tools.length > 0) {
-		let toolList = agent.tools;
+		toolNames = agent.tools;
 		// Auto-include task tool if spawns defined but task not in tools
-		if (agent.spawns !== undefined && !toolList.includes("task")) {
-			toolList = [...toolList, "task"];
+		if (agent.spawns !== undefined && !toolNames.includes("task")) {
+			toolNames = [...toolNames, "task"];
 		}
-		args.push("--tools", toolList.join(","));
 	}
 
 	// Resolve and add model
 	const resolvedModel = resolveModelPattern(modelOverride || agent.model);
-	if (resolvedModel) {
-		args.push("--model", resolvedModel);
-	}
-
-	// Add session options - use subtask-specific session file for real-time streaming
-	if (subtaskSessionFile) {
-		args.push("--session", subtaskSessionFile);
-	} else if (options.sessionFile) {
-		args.push("--session", options.sessionFile);
-	} else {
-		args.push("--no-session");
-	}
-
-	// Add task as prompt
-	args.push("--prompt", fullTask);
-
-	// Set up environment - block same-agent recursion unless explicitly recursive
-	const env = { ...process.env };
-	if (!agent.recursive) {
-		env[OMP_BLOCKED_AGENT_ENV] = agent.name;
-	}
-
-	// Propagate spawn restrictions to subprocess
-	if (agent.spawns === undefined) {
-		env[OMP_SPAWNS_ENV] = ""; // No spawns = deny all
-	} else if (agent.spawns === "*") {
-		env[OMP_SPAWNS_ENV] = "*";
-	} else {
-		env[OMP_SPAWNS_ENV] = agent.spawns.join(",");
-	}
+	const sessionFile = subtaskSessionFile ?? options.sessionFile ?? null;
+	const spawnsEnv = agent.spawns === undefined ? "" : agent.spawns === "*" ? "*" : agent.spawns.join(",");
 
-	// Spawn subprocess
-	const ompCommand = resolveOmpCommand();
-	const proc = Bun.spawn([ompCommand.cmd, ...ompCommand.args, ...args], {
-		cwd,
-		stdin: "ignore",
-		stdout: "pipe",
-		stderr: "pipe",
-		env,
-	});
+	const worker = new Worker(new URL("./worker.ts", import.meta.url), { type: "module" });
 
 	let output = "";
 	let stderr = "";
 	let finalOutput = "";
 	let resolved = false;
 	let pendingTermination = false; // Set when shouldTerminate fires, wait for message_end
-	const jsonlEvents: string[] = [];
+
+	// Accumulate usage incrementally from message_end events (no memory for streaming events)
+	const accumulatedUsage = {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	};
+	let hasUsage = false;
+
+	let abortSent = false;
+	const requestAbort = () => {
+		if (abortSent) return;
+		abortSent = true;
+		const abortMessage: SubagentWorkerRequest = { type: "abort" };
+		worker.postMessage(abortMessage);
+		setTimeout(() => {
+			if (!resolved) {
+				worker.terminate();
+			}
+		}, 2000);
+	};
 
 	// Handle abort signal
 	const onAbort = () => {
-		if (!resolved) {
-			proc.kill(15); // SIGTERM
-		}
+		if (!resolved) requestAbort();
 	};
 	if (signal) {
 		signal.addEventListener("abort", onAbort, { once: true });
 	}
 
-	// Parse JSON events from stdout
-	const reader = proc.stdout.getReader();
-	const decoder = new TextDecoder();
-	let buffer = "";
+	const emitProgress = () => {
+		progress.durationMs = Date.now() - startTime;
+		onProgress?.({ ...progress });
+		if (options.eventBus) {
+			options.eventBus.emit(TASK_SUBAGENT_PROGRESS_CHANNEL, {
+				index,
+				agent: agent.name,
+				agentSource: agent.source,
+				task,
+				progress: { ...progress },
+			});
+		}
+	};
 
-	const processLine = (line: string) => {
+	const getMessageContent = (message: unknown): unknown => {
+		if (message && typeof message === "object" && "content" in message) {
+			return (message as { content?: unknown }).content;
+		}
+		return undefined;
+	};
+
+	const getMessageUsage = (message: unknown): unknown => {
+		if (message && typeof message === "object" && "usage" in message) {
+			return (message as { usage?: unknown }).usage;
+		}
+		return undefined;
+	};
+
+	const processEvent = (event: AgentEvent) => {
 		if (resolved) return;
 
-		try {
-			const event = JSON.parse(line);
-			jsonlEvents.push(line);
-			const now = Date.now();
-
-			switch (event.type) {
-				case "tool_execution_start":
-					progress.toolCount++;
-					progress.currentTool = event.toolName;
-					progress.currentToolArgs = extractToolArgsPreview(event.toolArgs || event.args || {});
-					progress.currentToolStartMs = now;
-					break;
-
-				case "tool_execution_end": {
-					if (progress.currentTool) {
-						progress.recentTools.unshift({
-							tool: progress.currentTool,
-							args: progress.currentToolArgs || "",
-							endMs: now,
+		if (options.eventBus) {
+			options.eventBus.emit(TASK_SUBAGENT_EVENT_CHANNEL, {
+				index,
+				agent: agent.name,
+				agentSource: agent.source,
+				task,
+				event,
+			});
+		}
+
+		const now = Date.now();
+
+		switch (event.type) {
+			case "tool_execution_start":
+				progress.toolCount++;
+				progress.currentTool = event.toolName;
+				progress.currentToolArgs = extractToolArgsPreview(
+					(event as { toolArgs?: Record<string, unknown> }).toolArgs || event.args || {},
+				);
+				progress.currentToolStartMs = now;
+				break;
+
+			case "tool_execution_end": {
+				if (progress.currentTool) {
+					progress.recentTools.unshift({
+						tool: progress.currentTool,
+						args: progress.currentToolArgs || "",
+						endMs: now,
+					});
+					// Keep only last 5
+					if (progress.recentTools.length > 5) {
+						progress.recentTools.pop();
+					}
+				}
+				progress.currentTool = undefined;
+				progress.currentToolArgs = undefined;
+				progress.currentToolStartMs = undefined;
+
+				// Check for registered subagent tool handler
+				const handler = subprocessToolRegistry.getHandler(event.toolName);
+				const eventArgs = (event as { args?: Record<string, unknown> }).args ?? {};
+				if (handler) {
+					// Extract data using handler
+					if (handler.extractData) {
+						const data = handler.extractData({
+							toolName: event.toolName,
+							toolCallId: event.toolCallId,
+							args: eventArgs,
+							result: event.result,
+							isError: event.isError,
 						});
-						// Keep only last 5
-						if (progress.recentTools.length > 5) {
-							progress.recentTools.pop();
+						if (data !== undefined) {
+							progress.extractedToolData = progress.extractedToolData || {};
+							progress.extractedToolData[event.toolName] = progress.extractedToolData[event.toolName] || [];
+							progress.extractedToolData[event.toolName].push(data);
 						}
 					}
-					progress.currentTool = undefined;
-					progress.currentToolArgs = undefined;
-					progress.currentToolStartMs = undefined;
-
-					// Check for registered subprocess tool handler
-					const handler = subprocessToolRegistry.getHandler(event.toolName);
-					if (handler) {
-						// Extract data using handler
-						if (handler.extractData) {
-							const data = handler.extractData({
-								toolName: event.toolName,
-								toolCallId: event.toolCallId,
-								args: event.args,
-								result: event.result,
-								isError: event.isError,
-							});
-							if (data !== undefined) {
-								progress.extractedToolData = progress.extractedToolData || {};
-								progress.extractedToolData[event.toolName] = progress.extractedToolData[event.toolName] || [];
-								progress.extractedToolData[event.toolName].push(data);
-							}
-						}
 
-						// Check if handler wants to terminate subprocess
-						if (
-							handler.shouldTerminate?.({
-								toolName: event.toolName,
-								toolCallId: event.toolCallId,
-								args: event.args,
-								result: event.result,
-								isError: event.isError,
-							})
-						) {
-							// Don't kill immediately - wait for message_end to get token counts
-							pendingTermination = true;
-							// Safety timeout in case message_end never arrives
-							setTimeout(() => {
-								if (!resolved) {
-									resolved = true;
-									proc.kill(15); // SIGTERM
-								}
-							}, 2000);
-						}
+					// Check if handler wants to terminate worker
+					if (
+						handler.shouldTerminate?.({
+							toolName: event.toolName,
+							toolCallId: event.toolCallId,
+							args: eventArgs,
+							result: event.result,
+							isError: event.isError,
+						})
+					) {
+						// Don't terminate immediately - wait for message_end to get token counts
+						pendingTermination = true;
+						// Safety timeout in case message_end never arrives
+						setTimeout(() => {
+							if (!resolved) {
+								requestAbort();
+							}
+						}, 2000);
 					}
-					break;
 				}
+				break;
+			}
 
-				case "message_update": {
-					// Extract text for progress display only (replace, don't accumulate)
-					const updateContent = event.message?.content || event.content;
-					if (updateContent && Array.isArray(updateContent)) {
-						const allText: string[] = [];
-						for (const block of updateContent) {
-							if (block.type === "text" && block.text) {
-								const lines = block.text.split("\n").filter((l: string) => l.trim());
-								allText.push(...lines);
-							}
+			case "message_update": {
+				// Extract text for progress display only (replace, don't accumulate)
+				const updateContent =
+					getMessageContent(event.message) || (event as AgentEvent & { content?: unknown }).content;
+				if (updateContent && Array.isArray(updateContent)) {
+					const allText: string[] = [];
+					for (const block of updateContent) {
+						if (block.type === "text" && block.text) {
+							const lines = block.text.split("\n").filter((l: string) => l.trim());
+							allText.push(...lines);
 						}
-						// Show last 8 lines from current state (not accumulated)
-						progress.recentOutput = allText.slice(-8).reverse();
 					}
-					break;
+					// Show last 8 lines from current state (not accumulated)
+					progress.recentOutput = allText.slice(-8).reverse();
 				}
+				break;
+			}
 
-				case "message_end": {
-					// Extract final text content from completed message
-					const messageContent = event.message?.content || event.content;
+			case "message_end": {
+				// Extract text from assistant and toolResult messages (not user prompts)
+				const role = event.message?.role;
+				if (role === "assistant") {
+					const messageContent =
+						getMessageContent(event.message) || (event as AgentEvent & { content?: unknown }).content;
 					if (messageContent && Array.isArray(messageContent)) {
 						for (const block of messageContent) {
 							if (block.type === "text" && block.text) {
@@ -398,97 +386,120 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 							}
 						}
 					}
-					// Extract usage (prefer message.usage, fallback to event.usage)
-					const messageUsage = event.message?.usage || event.usage;
-					if (messageUsage) {
-						// Accumulate tokens across messages (not overwrite)
-						progress.tokens += getUsageTokens(messageUsage);
-					}
-					// If pending termination, now we have tokens - terminate
-					if (pendingTermination && !resolved) {
-						resolved = true;
-						proc.kill(15); // SIGTERM
+				}
+				// Extract and accumulate usage (prefer message.usage, fallback to event.usage)
+				const messageUsage = getMessageUsage(event.message) || (event as AgentEvent & { usage?: unknown }).usage;
+				if (messageUsage && typeof messageUsage === "object") {
+					// Only count assistant messages (not tool results, etc.)
+					if (
+						role === "assistant" &&
+						event.message?.stopReason !== "aborted" &&
+						event.message?.stopReason !== "error"
+					) {
+						const usageRecord = messageUsage as Record<string, number | undefined>;
+						const costRecord = (messageUsage as { cost?: Record<string, number | undefined> }).cost;
+						hasUsage = true;
+						accumulatedUsage.input += usageRecord.input ?? 0;
+						accumulatedUsage.output += usageRecord.output ?? 0;
+						accumulatedUsage.cacheRead += usageRecord.cacheRead ?? 0;
+						accumulatedUsage.cacheWrite += usageRecord.cacheWrite ?? 0;
+						accumulatedUsage.totalTokens += usageRecord.totalTokens ?? 0;
+						if (costRecord) {
+							accumulatedUsage.cost.input += costRecord.input ?? 0;
+							accumulatedUsage.cost.output += costRecord.output ?? 0;
+							accumulatedUsage.cost.cacheRead += costRecord.cacheRead ?? 0;
+							accumulatedUsage.cost.cacheWrite += costRecord.cacheWrite ?? 0;
+							accumulatedUsage.cost.total += costRecord.total ?? 0;
+						}
 					}
-					break;
+					// Accumulate tokens for progress display
+					progress.tokens += getUsageTokens(messageUsage);
+				}
+				// If pending termination, now we have tokens - terminate
+				if (pendingTermination && !resolved) {
+					requestAbort();
 				}
+				break;
+			}
 
-				case "agent_end":
-					// Extract final content from messages array
-					if (event.messages && Array.isArray(event.messages)) {
-						for (const msg of event.messages) {
-							if (msg.content && Array.isArray(msg.content)) {
-								for (const block of msg.content) {
-									if (block.type === "text" && block.text) {
-										finalOutput += block.text;
-									}
+			case "agent_end":
+				// Extract final content from assistant messages only (not user prompts)
+				if (event.messages && Array.isArray(event.messages)) {
+					for (const msg of event.messages) {
+						if ((msg as { role?: string })?.role !== "assistant") continue;
+						const messageContent = getMessageContent(msg);
+						if (messageContent && Array.isArray(messageContent)) {
+							for (const block of messageContent) {
+								if (block.type === "text" && block.text) {
+									finalOutput += block.text;
 								}
 							}
 						}
 					}
-					break;
-			}
-
-			progress.durationMs = now - startTime;
-			// Clone progress object before passing to callback to prevent mutation during render
-			onProgress?.({ ...progress });
-		} catch {
-			// Ignore non-JSON lines
+				}
+				break;
 		}
+
+		emitProgress();
 	};
 
-	// Read stdout asynchronously
-	const stdoutDone = (async () => {
-		try {
-			while (true) {
-				const { done, value } = await reader.read();
-				if (done) break;
-				buffer += decoder.decode(value, { stream: true });
-				const lines = buffer.split("\n");
-				buffer = lines.pop() || "";
-				for (const line of lines) {
-					processLine(line);
-				}
-			}
-			// Process remaining buffer
-			if (buffer.trim()) {
-				processLine(buffer);
-			}
-		} catch {
-			// Ignore read errors
-		}
-	})();
+	const startMessage: SubagentWorkerRequest = {
+		type: "start",
+		payload: {
+			cwd,
+			task: fullTask,
+			systemPrompt: agent.systemPrompt,
+			model: resolvedModel,
+			toolNames,
+			sessionFile,
+			spawnsEnv,
+		},
+	};
 
-	// Capture stderr - Bun.spawn returns ReadableStream, convert to text
-	const stderrDone = (async () => {
-		try {
-			const stderrReader = proc.stderr.getReader();
-			const stderrDecoder = new TextDecoder();
-			while (true) {
-				const { done, value } = await stderrReader.read();
-				if (done) break;
-				stderr += stderrDecoder.decode(value, { stream: true });
-			}
-		} catch {
-			// Ignore stderr read errors
-		}
-	})();
+	interface WorkerMessageEvent<T> {
+		data: T;
+	}
+	interface WorkerErrorEvent {
+		message: string;
+	}
 
-	// Wait for process and stream readers to finish
-	const exitCode = await proc.exited;
-	await Promise.all([stdoutDone, stderrDone]);
-	resolved = true;
+	const done = await new Promise<Extract<SubagentWorkerResponse, { type: "done" }>>((resolve) => {
+		const onMessage = (event: WorkerMessageEvent<SubagentWorkerResponse>) => {
+			const message = event.data;
+			if (!message || resolved) return;
+			if (message.type === "event") {
+				processEvent(message.event);
+				return;
+			}
+			if (message.type === "done") {
+				resolved = true;
+				resolve(message);
+			}
+		};
+		const onError = (event: WorkerErrorEvent) => {
+			if (resolved) return;
+			resolved = true;
+			resolve({
+				type: "done",
+				exitCode: 1,
+				durationMs: Date.now() - startTime,
+				error: event.message,
+			});
+		};
+		worker.addEventListener("message", onMessage);
+		worker.addEventListener("error", onError);
+		worker.postMessage(startMessage);
+	});
 
 	// Cleanup
 	if (signal) {
 		signal.removeEventListener("abort", onAbort);
 	}
+	worker.terminate();
 
-	try {
-		if (existsSync(promptFile)) {
-			unlinkSync(promptFile);
-		}
-	} catch {
-		// Ignore cleanup errors
+	const exitCode = done.exitCode;
+	if (done.error) {
+		stderr = done.error;
 	}
 
 	// Use final output if available, otherwise accumulated output
@@ -511,10 +522,9 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 	}
 
 	// Update final progress
-	const wasAborted = signal?.aborted ?? false;
+	const wasAborted = done.aborted || signal?.aborted || false;
 	progress.status = wasAborted ? "aborted" : exitCode === 0 ? "completed" : "failed";
-	progress.durationMs = Date.now() - startTime;
-	onProgress?.(progress);
+	emitProgress();
 
 	return {
 		index,
@@ -531,7 +541,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 		modelOverride,
 		error: exitCode !== 0 && stderr ? stderr : undefined,
 		aborted: wasAborted,
-		jsonlEvents,
+		usage: hasUsage ? accumulatedUsage : undefined,
 		artifactPaths,
 		extractedToolData: progress.extractedToolData,
 		outputMeta,