@nmzpy/pi-ember-stack 0.2.1 → 0.2.2

package/plugins/pi-custom-agents/subagent/extensions/runner.ts

@@ -1,33 +1,20 @@
 /**
  * SDK-based sub-agent runner for pi-subagent.
  *
- * Creates an in-process AgentSession via the pi SDK instead of spawning a
- * separate `pi` process. This eliminates cold-start overhead and allows
- * fine-grained control over token budget:
- *
- *   - The agent's system prompt is prepended with project context files
- *     (AGENTS.md from cwd and ancestors) so sub-agents share the same
- *     codebase conventions as the main agent.
- *   - No extensions, no skills, no prompt templates loaded.
- *   - Thinking disabled, compaction disabled, retry disabled.
- *   - In-memory session (no disk I/O).
- *   - Shared auth/model infrastructure (no re-connection).
- *
- * Estimated token savings vs process-spawn: ~4-11K tokens per invocation.
+ * Compared to process-spawn, this saves ~4-11K tokens per invocation by using
+ * the pi SDK directly with a minimal system prompt, no extensions, no skills,
+ * no prompt templates, no thinking, and no compaction. To avoid blocking the
+ * parent TUI thread while the subagent runs, each agent is executed in a
+ * Node worker_thread and messages its progress back.
  */
 
+import * as path from "node:path";
+import { Worker } from "node:worker_threads";
 import type { Message, Model } from "@earendil-works/pi-ai";
-import type { AgentMessage } from "@earendil-works/pi-agent-core";
 import {
-	AuthStorage,
-	createAgentSession,
-	createExtensionRuntime,
+	type AuthStorage,
 	getAgentDir,
-	loadProjectContextFiles,
-	ModelRegistry,
-	type ResourceLoader,
-	SessionManager,
-	SettingsManager,
+	type ModelRegistry,
 } from "@earendil-works/pi-coding-agent";
 
 // ---------------------------------------------------------------------------
@@ -41,7 +28,7 @@ const PARALLEL_TOOL_CALL_GUIDANCE = `
 When multiple independent tool calls are needed (e.g. reading several files,
 searching for different patterns), emit them all in a single response rather
 than one at a time. The runtime executes independent tool calls in parallel,
-so batching them saves round-trips and reduces latency.
+so batching saves round-trips and reduces latency.
 `;
 
 // ---------------------------------------------------------------------------
@@ -70,6 +57,123 @@ export interface SubAgentResult {
 	errorMessage?: string;
 }
 
+interface WorkerInput {
+	cwd: string;
+	agentDir: string;
+	authPath: string;
+	modelsPath: string | undefined;
+	systemPrompt: string;
+	fullSystemPrompt: string;
+	task: string;
+	tools: string[];
+	model: Model<any>;
+	agentName: string;
+	thinkingLevel: string;
+}
+
+// ---------------------------------------------------------------------------
+// Worker script — kept inline so the package stays self-contained. Worker
+// threads cannot load .ts via jiti, so the code is written as plain JS and
+// executed from a data: URL.
+// ---------------------------------------------------------------------------
+
+const WORKER_SCRIPT = `
+const { parentPort, workerData } = require("node:worker_threads");
+
+const { createAgentSession, getAgentDir, loadProjectContextFiles, AuthStorage, ModelRegistry, SessionManager, SettingsManager } = require("@earendil-works/pi-coding-agent");
+
+function post(type, payload) {
+	parentPort?.postMessage({ type, payload });
+}
+
+const result = {
+	agent: workerData.agentName || "subagent",
+	task: workerData.task,
+	exitCode: 0,
+	messages: [],
+	stderr: "",
+	usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+	model: workerData.model ? \`\${workerData.model.provider}/\${workerData.model.id}\` : undefined,
+};
+
+async function main() {
+	const authStorage = AuthStorage.create(workerData.authPath);
+	const modelRegistry = ModelRegistry.create(authStorage, workerData.modelsPath);
+	const settingsManager = SettingsManager.inMemory({ compaction: { enabled: false }, retry: { enabled: false } });
+
+	const resourceLoader = {
+		getExtensions: () => ({ extensions: [], errors: [], runtime: {} }),
+		getSkills: () => ({ skills: [], diagnostics: [] }),
+		getPrompts: () => ({ prompts: [], diagnostics: [] }),
+		getThemes: () => ({ themes: [], diagnostics: [] }),
+		getAgentsFiles: () => ({ agentsFiles: [] }),
+		getSystemPrompt: () => workerData.fullSystemPrompt,
+		getAppendSystemPrompt: () => [],
+		extendResources: () => {},
+		reload: async () => {},
+	};
+
+	const { session } = await createAgentSession({
+		cwd: workerData.cwd,
+		model: workerData.model,
+		thinkingLevel: workerData.thinkingLevel || "off",
+		authStorage,
+		modelRegistry,
+		resourceLoader,
+		tools: workerData.tools,
+		sessionManager: SessionManager.inMemory(workerData.cwd),
+		settingsManager,
+	});
+
+	const unsubscribe = session.subscribe((event) => {
+		if (event.type === "message_end") {
+			const msg = event.message;
+			if (msg.role === "assistant") {
+				result.usage.turns++;
+				if (msg.usage) {
+					result.usage.input += msg.usage.input || 0;
+					result.usage.output += msg.usage.output || 0;
+					result.usage.cacheRead += msg.usage.cacheRead || 0;
+					result.usage.cacheWrite += msg.usage.cacheWrite || 0;
+					result.usage.cost += msg.usage.cost?.total || 0;
+					result.usage.contextTokens = msg.usage.totalTokens || 0;
+				}
+				if (!result.model && msg.model) {
+					result.model = \`\${msg.provider || "?"}/\${msg.model}\`;
+				}
+				if (msg.stopReason) result.stopReason = msg.stopReason;
+				if (msg.errorMessage) result.errorMessage = msg.errorMessage;
+			}
+			result.messages.push(msg);
+			post("message", { ...result, messages: [...result.messages] });
+		}
+		if (event.type === "agent_end" && result.messages.length === 0 && event.messages) {
+			result.messages = event.messages;
+		}
+	});
+
+	try {
+		await session.prompt(workerData.task);
+		result.exitCode = result.stopReason === "aborted" ? 1 : 0;
+	} catch (error) {
+		result.exitCode = 1;
+		result.stopReason = result.stopReason || "error";
+		result.errorMessage = result.errorMessage || (error instanceof Error ? error.message : String(error));
+	} finally {
+		unsubscribe?.();
+		try { session.dispose(); } catch {}
+	}
+	post("done", result);
+}
+
+main().catch((err) => {
+	result.exitCode = 1;
+	result.stopReason = result.stopReason || "error";
+	result.errorMessage = result.errorMessage || (err instanceof Error ? err.message : String(err));
+	post("done", result);
+});
+`;
+
 // ---------------------------------------------------------------------------
 // Public API
 // ---------------------------------------------------------------------------
@@ -99,187 +203,147 @@ export async function runSubAgent(options: {
 		signal,
 		agentName = "subagent",
 		thinkingLevel = "off",
-		onUpdate,
 		onMessage,
 	} = options;
 
-	const result: SubAgentResult = {
-		agent: agentName,
-		task,
-		exitCode: 0,
-		messages: [],
-		stderr: "",
-		usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
-		model: `${model.provider}/${model.id}`,
-	};
+	if (signal?.aborted) {
+		return {
+			agent: agentName,
+			task,
+			exitCode: 1,
+			messages: [],
+			stderr: "",
+			usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+			stopReason: "aborted",
+			errorMessage: "Sub-agent aborted before start",
+		};
+	}
+
+	const agentDir = getAgentDir();
+	const authPath = getAuthStoragePath(authStorage);
+	const modelsPath = getModelsPath(modelRegistry);
 
-	// Load project context files (AGENTS.md from cwd and ancestors) so the
-	// sub-agent has the same codebase conventions as the main agent.
-	const contextFiles = loadProjectContextFiles({ cwd, agentDir: getAgentDir() });
+	const contextFiles = loadProjectContextFilesCompat({ cwd, agentDir });
 	const contextPrefix = contextFiles.length > 0
 		? contextFiles.map((f) => f.content).join("\n\n---\n\n") + "\n\n---\n\n"
 		: "";
 	const fullSystemPrompt = contextPrefix + systemPrompt + PARALLEL_TOOL_CALL_GUIDANCE;
 
-	// Build a minimal resource loader. The sub-agent sees the agent's system
-	// prompt plus project context files (AGENTS.md) — no extensions, no skills.
-	const resourceLoader: ResourceLoader = {
-		getExtensions: () => ({ extensions: [], errors: [], runtime: createExtensionRuntime() }),
-		getSkills: () => ({ skills: [], diagnostics: [] }),
-		getPrompts: () => ({ prompts: [], diagnostics: [] }),
-		getThemes: () => ({ themes: [], diagnostics: [] }),
-		getAgentsFiles: () => ({ agentsFiles: [] }),
-		getSystemPrompt: () => fullSystemPrompt,
-		getAppendSystemPrompt: () => [],
-		extendResources: () => {},
-		reload: async () => {},
+	const input: WorkerInput = {
+		cwd,
+		agentDir,
+		authPath,
+		modelsPath,
+		systemPrompt,
+		fullSystemPrompt,
+		task,
+		tools,
+		model,
+		agentName,
+		thinkingLevel,
 	};
 
-	const settingsManager = SettingsManager.inMemory({
-		compaction: { enabled: false },
-		retry: { enabled: false },
-	});
+	return new Promise<SubAgentResult>((resolve, reject) => {
+		let settled = false;
+		const worker = new Worker(WORKER_SCRIPT, { eval: true, workerData: input as any });
 
-	try {
-		if (signal?.aborted) {
-			result.exitCode = 1;
-			result.stopReason = "aborted";
-			result.errorMessage = "Sub-agent aborted before start";
-			return result;
+		function finish(result: SubAgentResult) {
+			if (settled) return;
+			settled = true;
+			worker.terminate().catch(() => {});
+			resolve(result);
 		}
 
-		const { session } = await createAgentSession({
-			cwd,
-			model,
-			thinkingLevel,
-			authStorage,
-			modelRegistry,
-			resourceLoader,
-			tools,
-			sessionManager: SessionManager.inMemory(cwd),
-			settingsManager,
-		});
-
-		let cleanupAbort: (() => void) | undefined;
-		let cleanupEventAbort: (() => void) | undefined;
-		try {
-			// Wire abort signal
-			if (signal) {
-				const onAbort = () => session.abort();
-				if (signal.aborted) {
-					// Already aborted — shortcut
-					result.exitCode = 1;
-					result.stopReason = "aborted";
-					result.errorMessage = "Sub-agent aborted before start";
-					onAbort();
-					return result;
-				}
-				signal.addEventListener("abort", onAbort, { once: true });
-				cleanupAbort = () => signal.removeEventListener("abort", onAbort);
+		worker.on("message", (message: { type: "message" | "done"; payload: SubAgentResult }) => {
+			if (message.type === "message") {
+				onMessage?.(message.payload);
+			} else if (message.type === "done") {
+				finish(message.payload);
 			}
+		});
 
-			// Collect all messages and usage stats from events
-			const eventPromise = new Promise<void>((resolve, reject) => {
-				let settled = false;
-				const finish = (fn: () => void) => {
-					if (settled) return;
-					settled = true;
-					fn();
-				};
-
-				const unsubscribe = session.subscribe((event) => {
-					try {
-						switch (event.type) {
-							case "message_end": {
-								const msg = event.message as AgentMessage;
-								if (msg.role === "assistant") {
-									result.usage.turns++;
-									if (msg.usage) {
-										result.usage.input += msg.usage.input || 0;
-										result.usage.output += msg.usage.output || 0;
-										result.usage.cacheRead += msg.usage.cacheRead || 0;
-										result.usage.cacheWrite += msg.usage.cacheWrite || 0;
-										result.usage.cost += msg.usage.cost?.total || 0;
-										result.usage.contextTokens = msg.usage.totalTokens || 0;
-									}
-									if (!result.model && msg.model) {
-										result.model = `${msg.provider || "?"}/${msg.model}`;
-									}
-									if (msg.stopReason) result.stopReason = msg.stopReason;
-									if (msg.errorMessage) result.errorMessage = msg.errorMessage;
-								}
-								// Collect all messages for extraction
-								result.messages.push(msg as unknown as Message);
-								if (onMessage) onMessage({ ...result, messages: [...result.messages] });
-								break;
-							}
-							case "agent_end": {
-								// agent_end carries all messages; use them if we haven't collected
-								if (result.messages.length === 0 && event.messages) {
-									result.messages = event.messages as unknown as Message[];
-								}
-								finish(() => {
-									unsubscribe();
-									resolve();
-								});
-								break;
-							}
-						}
-					} catch (err) {
-						finish(() => {
-							unsubscribe();
-							reject(err);
-						});
-					}
-				});
-
-				// Resolve on abort so the eventPromise doesn't hang
-				if (signal) {
-					const onAbortResolve = () => {
-						finish(() => {
-							result.exitCode = 1;
-							result.stopReason = "aborted";
-							if (!result.errorMessage) result.errorMessage = "Sub-agent aborted";
-							unsubscribe();
-							resolve();
-						});
-					};
-					signal.addEventListener("abort", onAbortResolve, { once: true });
-					cleanupEventAbort = () => signal.removeEventListener("abort", onAbortResolve);
-				}
+		worker.on("error", (err) => {
+			finish({
+				agent: agentName,
+				task,
+				exitCode: 1,
+				messages: [],
+				stderr: "",
+				usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+				stopReason: "error",
+				errorMessage: err instanceof Error ? err.message : String(err),
 			});
+		});
 
-			await Promise.race([
-				session.prompt(task),
-				eventPromise,
-			]);
+		worker.on("exit", (code) => {
+			if (settled) return;
+			finish({
+				agent: agentName,
+				task,
+				exitCode: 1,
+				messages: [],
+				stderr: "",
+				usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+				stopReason: "error",
+				errorMessage: `Sub-agent worker exited unexpectedly (code ${code})`,
+			});
+		});
 
-			if (result.stopReason !== "aborted") {
-				result.exitCode = 0;
-			}
-			return result;
-		} finally {
-			cleanupAbort?.();
-			cleanupEventAbort?.();
-			try {
-				session.dispose();
-			} catch {
-				// Best-effort cleanup
+		if (signal) {
+			const onAbort = () => {
+				if (settled) return;
+				worker.terminate().catch(() => {});
+				finish({
+					agent: agentName,
+					task,
+					exitCode: 1,
+					messages: [],
+					stderr: "",
+					usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+					stopReason: "aborted",
+					errorMessage: "Sub-agent aborted",
+				});
+			};
+			if (signal.aborted) {
+				onAbort();
+				return;
 			}
+			signal.addEventListener("abort", onAbort, { once: true });
 		}
-	} catch (err) {
-		const message = err instanceof Error ? err.message : String(err);
-		result.exitCode = 1;
-		result.errorMessage = message;
-		if (!result.stopReason) result.stopReason = "error";
-		return result;
-	}
+	});
 }
 
 // ---------------------------------------------------------------------------
 // Helpers
 // ---------------------------------------------------------------------------
 
+function getAuthStoragePath(authStorage: AuthStorage): string {
+	try {
+		const storage = authStorage as unknown as { authPath?: string; path?: string };
+		if (storage.authPath) return storage.authPath;
+		if (storage.path) return storage.path;
+	} catch {}
+	return path.join(getAgentDir(), "auth.json");
+}
+
+function getModelsPath(modelRegistry: ModelRegistry): string | undefined {
+	try {
+		const registry = modelRegistry as unknown as { modelsPath?: string; path?: string };
+		if (registry.modelsPath) return registry.modelsPath;
+		if (registry.path) return registry.path;
+	} catch {}
+	return undefined;
+}
+
+function loadProjectContextFilesCompat({ cwd, agentDir }: { cwd: string; agentDir: string }): { content: string }[] {
+	try {
+		const { loadProjectContextFiles } = require("@earendil-works/pi-coding-agent");
+		return loadProjectContextFiles({ cwd, agentDir });
+	} catch {
+		return [];
+	}
+}
+
 export function getFinalOutput(messages: Message[]): string {
 	// Prefer the last assistant message with non-empty text and NO tool calls (pure final answer).
 	for (let i = messages.length - 1; i >= 0; i--) {

package/plugins/pi-custom-agents/subagent/extensions/test/render.test.ts

@@ -0,0 +1,145 @@
+import { describe, expect, mock, test } from "bun:test";
+import { renderSubagentLayout, anySubagentRunning, SubagentCapLine } from "../render.ts";
+
+function makeTheme() {
+	const fg = mock((tag: string, text: string) => `[${tag}:${text}]`);
+	return {
+		fg,
+		bold: mock((s: string) => `*${s}*`),
+	};
+}
+
+function stripAnsi(text: string): string {
+	return text.replace(/\u001b\[[0-9;]*m/g, "");
+}
+
+function makeResult(agent: string, exitCode: number, failed = false) {
+	return {
+		agent,
+		task: "test",
+		exitCode,
+		messages: [] as any[],
+		stderr: "",
+		usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+		...(failed ? { stopReason: "error", errorMessage: "fail" } : {}),
+	} as any;
+}
+
+describe("SubagentCapLine", () => {
+	test("renders exactly the current width and disappears when hidden", () => {
+		let visible = true;
+		const cap = new SubagentCapLine(() => visible, (_color, text) => text);
+
+		expect(cap.render(37)).toEqual(["\u2500".repeat(37)]);
+		expect(cap.render(241)).toEqual(["\u2500".repeat(241)]);
+
+		visible = false;
+		expect(cap.render(80)).toEqual([]);
+	});
+});
+
+describe("renderSubagentLayout", () => {
+	test("running single mode uses the Thinking gradient without a bullet", () => {
+		const theme = makeTheme() as any;
+		const out = renderSubagentLayout({ agent: "coder", task: "do stuff" }, [], theme);
+		expect(stripAnsi(out)).toContain("coder");
+		expect(out).toContain("\u001b[38;2;");
+		expect(stripAnsi(out)).not.toContain("\u2022");
+		expect(out).not.toContain("\u2713");
+		expect(out).not.toContain("\u2717");
+		expect(out).not.toContain("subagent");
+		expect(out).not.toContain("[user]");
+		expect(out).not.toContain("\u23f3");
+	});
+
+	test("completed single mode shows checkmark", () => {
+		const theme = makeTheme() as any;
+		const out = renderSubagentLayout({ agent: "coder", task: "do stuff" }, [makeResult("coder", 0)], theme);
+		expect(stripAnsi(out)).toContain("coder");
+		expect(stripAnsi(out)).toContain("\u2713");
+		expect(stripAnsi(out)).not.toContain("\u2717");
+		expect(out).not.toContain("\u001b[38;2;");
+		expect(stripAnsi(out)).not.toContain("\u2022");
+	});
+
+	test("failed single mode shows X mark", () => {
+		const theme = makeTheme() as any;
+		const out = renderSubagentLayout({ agent: "coder", task: "do stuff" }, [makeResult("coder", 1, true)], theme);
+		expect(stripAnsi(out)).toContain("coder");
+		expect(stripAnsi(out)).toContain("\u2717");
+		expect(stripAnsi(out)).not.toContain("\u2713");
+		expect(out).not.toContain("\u001b[38;2;");
+		expect(stripAnsi(out)).not.toContain("\u2022");
+	});
+
+	test("failed single mode includes inline error text", () => {
+		const theme = makeTheme() as any;
+		const result = makeResult("coder", 1, true);
+		result.errorMessage = "timeout during read";
+		const out = renderSubagentLayout({ agent: "coder", task: "do stuff" }, [result], theme);
+		expect(stripAnsi(out)).toContain("timeout during read");
+	});
+
+	test("completed single mode does not include error text", () => {
+		const theme = makeTheme() as any;
+		const out = renderSubagentLayout({ agent: "coder", task: "do stuff" }, [makeResult("coder", 0)], theme);
+		expect(stripAnsi(out)).not.toContain("error");
+		expect(stripAnsi(out)).not.toContain("Error");
+	});
+
+	test("failed parallel mode shows inline error only on failed child", () => {
+		const theme = makeTheme() as any;
+		const args = { tasks: [{ agent: "coder", task: "a" }, { agent: "scout", task: "b" }] };
+		const failed = makeResult("coder", 1, true);
+		failed.errorMessage = "could not read file";
+		const out = renderSubagentLayout(args, [failed, makeResult("scout", 0)], theme);
+		const output = stripAnsi(out);
+		expect(output).toContain("coder");
+		expect(output).toContain("scout");
+		expect(output).toContain("could not read file");
+		const matches = output.split("could not read file").length - 1;
+		expect(matches).toBe(1);
+	});
+
+	test("parallel mode shows plain Subagents header with children", () => {
+		const theme = makeTheme() as any;
+		const args = { tasks: [{ agent: "coder", task: "a" }, { agent: "scout", task: "b" }] };
+		const out = renderSubagentLayout(args, [], theme);
+		const lines = out.split("\n");
+		const header = lines[0];
+		expect(stripAnsi(header)).toContain("Subagents");
+		expect(header).not.toContain("\u001b[38;2;");
+		expect(stripAnsi(out)).toContain("coder");
+		expect(stripAnsi(out)).toContain("scout");
+		expect(out).not.toContain("\u23f3");
+		expect(out).not.toContain("parallel");
+		expect(out).not.toContain("[user]");
+		expect(stripAnsi(out)).not.toContain("\u2022");
+	});
+
+	test("chain mode only shows started steps", () => {
+		const theme = makeTheme() as any;
+		const args = { chain: [{ agent: "scout", task: "a" }, { agent: "coder", task: "b" }] };
+		const out = renderSubagentLayout(args, [makeResult("scout", 0)], theme);
+		expect(out).toContain("scout");
+		expect(out).not.toContain("coder");
+	});
+
+	test("no hourglass glyphs anywhere", () => {
+		const theme = makeTheme() as any;
+		const args = { tasks: [{ agent: "coder", task: "a" }] };
+		const out = renderSubagentLayout(args, [makeResult("coder", -1)], theme);
+		expect(out).not.toContain("\u23f3");
+		expect(out).not.toContain("\u25d0");
+	});
+
+	test("anySubagentRunning true when exitCode is -1", () => {
+		const args = { tasks: [{ agent: "coder", task: "a" }] };
+		expect(anySubagentRunning(args, [makeResult("coder", -1)])).toBe(true);
+	});
+
+	test("anySubagentRunning false when all done", () => {
+		const args = { tasks: [{ agent: "coder", task: "a" }] };
+		expect(anySubagentRunning(args, [makeResult("coder", 0)])).toBe(false);
+	});
+});