pi-plans 0.1.2 → 0.2.0

package/src/subagent.ts

@@ -9,6 +9,20 @@ import * as fs from "node:fs";
 import * as os from "node:os";
 import * as path from "node:path";
 
+export type SubagentProgressEvent =
+	| { type: "process"; phase: "started" | "exited"; code?: number }
+	| { type: "turn"; phase: "start" | "end" }
+	| { type: "message"; phase: "start" | "update" | "end"; role: string; text: string }
+	| {
+			type: "tool";
+			phase: "start" | "update" | "end";
+			toolCallId: string;
+			toolName: string;
+			detail: string;
+			isError?: boolean;
+	  }
+	| { type: "stderr"; text: string };
+
 export interface SubagentOptions {
 	systemPrompt: string;
 	task: string;
@@ -19,6 +33,8 @@ export interface SubagentOptions {
 	tools?: string[];
 	signal?: AbortSignal;
 	timeoutMs?: number;
+	/** Optional normalized progress sink. Exceptions from the sink are ignored. */
+	onProgress?: (event: SubagentProgressEvent) => void;
 }
 
 export interface SubagentResult {
@@ -28,6 +44,8 @@ export interface SubagentResult {
 	errorMessage?: string;
 	stderr: string;
 	turns: number;
+	cancelled?: boolean;
+	timedOut?: boolean;
 }
 
 /** Strip YAML frontmatter from an agent definition file. */
@@ -54,43 +72,132 @@ export function getPiInvocation(args: string[]): { command: string; args: string
 
 interface MessageLike {
 	role: string;
-	content: Array<{ type: string; text?: string }>;
+	content?: Array<{ type: string; text?: string; thinking?: string }> | string;
 	model?: string;
 }
 
+interface RawSubagentEvent {
+	type?: unknown;
+	message?: MessageLike;
+	toolCallId?: unknown;
+	toolName?: unknown;
+	args?: unknown;
+	partialResult?: unknown;
+	result?: unknown;
+	isError?: unknown;
+	assistantMessageEvent?: unknown;
+}
+
 function finalOutput(messages: MessageLike[]): string {
 	for (let i = messages.length - 1; i >= 0; i--) {
 		const message = messages[i];
 		if (message.role === "assistant") {
-			const text = message.content
-				.filter((part) => part.type === "text")
-				.map((part) => part.text ?? "")
-				.join("\n")
-				.trim();
+			const text = messageText(message, "text").trim();
 			if (text) return text;
 		}
 	}
 	return "";
 }
 
+function messageText(message: MessageLike | undefined, kind: "text" | "thinking" = "text"): string {
+	if (!message?.content) return "";
+	if (typeof message.content === "string") return kind === "text" ? message.content : "";
+	return message.content
+		.filter((part) => part.type === kind)
+		.map((part) => (kind === "thinking" ? part.thinking ?? part.text ?? "" : part.text ?? ""))
+		.join("\n");
+}
+
+function messageEventText(value: unknown): string {
+	if (!value || typeof value !== "object") return "";
+	const event = value as { delta?: unknown; content?: unknown };
+	return typeof event.delta === "string" ? event.delta : typeof event.content === "string" ? event.content : "";
+}
+
+function preview(value: unknown, maxLength = 240): string {
+	if (typeof value === "string") return value.slice(0, maxLength);
+	if (value === undefined) return "";
+	try {
+		const text = JSON.stringify(value);
+		return text ? text.slice(0, maxLength) : "";
+	} catch {
+		return String(value).slice(0, maxLength);
+	}
+}
+
+/**
+ * Normalize the JSONL events emitted by `pi --mode json` into the small
+ * protocol consumed by the refinement overlay. Unknown events are ignored so
+ * adding progress support cannot make result parsing version-fragile.
+ */
+export function normalizeSubagentEvent(value: unknown): SubagentProgressEvent | undefined {
+	if (!value || typeof value !== "object") return undefined;
+	const event = value as RawSubagentEvent;
+	if (typeof event.type !== "string") return undefined;
+
+	if (event.type === "turn_start") return { type: "turn", phase: "start" };
+	if (event.type === "turn_end") return { type: "turn", phase: "end" };
+
+	if (event.type === "message_start" || event.type === "message_update" || event.type === "message_end") {
+		if (!event.message && !event.assistantMessageEvent) return undefined;
+		const role = event.message?.role ?? "assistant";
+		return {
+			type: "message",
+			phase: event.type.slice("message_".length) as "start" | "update" | "end",
+			role,
+			text: messageText(event.message) || messageEventText(event.assistantMessageEvent),
+		};
+	}
+
+	const toolEvent =
+		event.type === "tool_execution_start"
+			? "start"
+			: event.type === "tool_execution_update"
+				? "update"
+				: event.type === "tool_execution_end" || event.type === "tool_result_end"
+					? "end"
+					: null;
+	if (toolEvent) {
+		return {
+			type: "tool",
+			phase: toolEvent,
+			toolCallId: typeof event.toolCallId === "string" ? event.toolCallId : "unknown",
+			toolName: typeof event.toolName === "string" ? event.toolName : "tool",
+			detail: preview(event.args ?? event.partialResult ?? event.result),
+			...(typeof event.isError === "boolean" ? { isError: event.isError } : {}),
+		};
+	}
+
+	return undefined;
+}
+
+function emitProgress(options: SubagentOptions, event: SubagentProgressEvent): void {
+	try {
+		options.onProgress?.(event);
+	} catch {
+		// A display sink must not be able to fail the child runner.
+	}
+}
+
 const DEFAULT_TIMEOUT_MS = 15 * 60 * 1000;
 
 export async function runPiSubagent(options: SubagentOptions): Promise<SubagentResult> {
 	const tools = options.tools ?? ["read", "grep", "find", "ls"];
-	const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-plans-subagent-"));
-	const promptFile = path.join(tmpDir, "system-prompt.md");
-	fs.writeFileSync(promptFile, options.systemPrompt, { encoding: "utf8", mode: 0o600 });
-
-	const args: string[] = ["--mode", "json", "-p", "--no-session", "--tools", tools.join(",")];
-	if (options.model) args.push("--model", options.model);
-	args.push("--append-system-prompt", promptFile);
-	args.push(`Task: ${options.task}`);
-
+	let tmpDir = "";
 	const messages: MessageLike[] = [];
 	let stderr = "";
-	let wasAborted = false;
+	let termination: "abort" | "timeout" | null = null;
 
 	try {
+		tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-plans-subagent-"));
+		const promptFile = path.join(tmpDir, "system-prompt.md");
+		fs.writeFileSync(promptFile, options.systemPrompt, { encoding: "utf8", mode: 0o600 });
+
+		const args: string[] = ["--mode", "json", "-p", "--no-session", "--tools", tools.join(",")];
+		if (options.model) args.push("--model", options.model);
+		args.push("--append-system-prompt", promptFile);
+		args.push(`Task: ${options.task}`);
+
 		const exitCode = await new Promise<number>((resolve) => {
 			const invocation = getPiInvocation(args);
 			const proc = spawn(invocation.command, invocation.args, {
@@ -99,20 +206,61 @@ export async function runPiSubagent(options: SubagentOptions): Promise<SubagentR
 				stdio: ["ignore", "pipe", "pipe"],
 			});
 			let buffer = "";
+			let closed = false;
+			let settled = false;
+			let timeoutTimer: ReturnType<typeof setTimeout> | undefined;
+			let forceKillTimer: ReturnType<typeof setTimeout> | undefined;
+
+			const onAbort = () => {
+				if (termination === null) termination = "abort";
+				killProc();
+			};
+
+			const cleanup = () => {
+				if (timeoutTimer) clearTimeout(timeoutTimer);
+				if (forceKillTimer) clearTimeout(forceKillTimer);
+				options.signal?.removeEventListener("abort", onAbort);
+			};
+
+			const finish = (code: number) => {
+				if (settled) return;
+				settled = true;
+				closed = true;
+				cleanup();
+				emitProgress(options, { type: "process", phase: "exited", code });
+				resolve(code);
+			};
 
 			const processLine = (line: string) => {
 				if (!line.trim()) return;
-				let event: { type?: string; message?: MessageLike };
+				let event: RawSubagentEvent;
 				try {
-					event = JSON.parse(line);
+					event = JSON.parse(line) as RawSubagentEvent;
 				} catch {
 					return;
 				}
+				const progress = normalizeSubagentEvent(event);
+				if (progress) emitProgress(options, progress);
 				if ((event.type === "message_end" || event.type === "tool_result_end") && event.message) {
 					messages.push(event.message);
 				}
 			};
 
+			const killProc = () => {
+				if (settled || termination === null) return;
+				if (!proc.killed) proc.kill("SIGTERM");
+				if (!forceKillTimer) {
+					forceKillTimer = setTimeout(() => {
+						try {
+							if (!closed) proc.kill("SIGKILL");
+						} catch {
+							// The process already exited.
+						}
+					}, 5000);
+				}
+			};
+
+			emitProgress(options, { type: "process", phase: "started" });
 			proc.stdout.on("data", (data) => {
 				buffer += data.toString();
 				const lines = buffer.split("\n");
@@ -120,78 +268,57 @@ export async function runPiSubagent(options: SubagentOptions): Promise<SubagentR
 				for (const line of lines) processLine(line);
 			});
 			proc.stderr.on("data", (data) => {
-				stderr += data.toString();
+				const text = data.toString();
+				stderr += text;
+				emitProgress(options, { type: "stderr", text });
 			});
 			proc.on("close", (code) => {
 				if (buffer.trim()) processLine(buffer);
-				resolve(code ?? 0);
+				finish(code ?? 0);
+			});
+			proc.on("error", (error) => {
+				stderr += error instanceof Error ? error.message : String(error);
+				finish(1);
 			});
-			proc.on("error", () => resolve(1));
-
-			const killProc = () => {
-				wasAborted = true;
-				proc.kill("SIGTERM");
-				setTimeout(() => {
-					try {
-						if (!proc.killed) proc.kill("SIGKILL");
-					} catch {
-						/* already gone */
-					}
-				}, 5000);
-			};
 
-			const timer = setTimeout(killProc, options.timeoutMs ?? DEFAULT_TIMEOUT_MS);
-			const onAbort = () => killProc();
+			timeoutTimer = setTimeout(() => {
+				if (termination === null) termination = "timeout";
+				killProc();
+			}, options.timeoutMs ?? DEFAULT_TIMEOUT_MS);
 			if (options.signal) {
-				if (options.signal.aborted) killProc();
+				if (options.signal.aborted) onAbort();
 				else options.signal.addEventListener("abort", onAbort, { once: true });
 			}
-			proc.on("close", () => {
-				clearTimeout(timer);
-				options.signal?.removeEventListener("abort", onAbort);
-			});
 		});
 
-		if (wasAborted) {
-			return {
-				ok: false,
-				output: finalOutput(messages),
-				stderr,
-				turns: messages.filter((m) => m.role === "assistant").length,
-				errorMessage: "Subagent was aborted",
-			};
+		const turns = messages.filter((message) => message.role === "assistant").length;
+		const output = finalOutput(messages);
+		if (termination === "abort") {
+			return { ok: false, output, stderr, turns, cancelled: true, errorMessage: "Subagent was aborted" };
+		}
+		if (termination === "timeout") {
+			return { ok: false, output, stderr, turns, timedOut: true, errorMessage: "Subagent timed out" };
 		}
 		if (exitCode !== 0) {
-			return {
-				ok: false,
-				output: finalOutput(messages),
-				stderr,
-				turns: messages.filter((m) => m.role === "assistant").length,
-				errorMessage: `pi exited with code ${exitCode}`,
-			};
+			return { ok: false, output, stderr, turns, errorMessage: `pi exited with code ${exitCode}` };
 		}
-		const output = finalOutput(messages);
 		if (!output) {
-			return {
-				ok: false,
-				output: "",
-				stderr,
-				turns: messages.filter((m) => m.role === "assistant").length,
-				errorMessage: "subagent produced no final output",
-			};
+			return { ok: false, output: "", stderr, turns, errorMessage: "subagent produced no final output" };
 		}
 		return {
 			ok: true,
 			output,
-			model: [...messages].reverse().find((m) => m.role === "assistant" && m.model)?.model,
+			model: [...messages].reverse().find((message) => message.role === "assistant" && message.model)?.model,
 			stderr,
-			turns: messages.filter((m) => m.role === "assistant").length,
+			turns,
 		};
 	} finally {
-		try {
-			fs.rmSync(tmpDir, { recursive: true, force: true });
-		} catch {
-			/* best effort */
+		if (tmpDir) {
+			try {
+				fs.rmSync(tmpDir, { recursive: true, force: true });
+			} catch {
+				/* best effort */
+			}
 		}
 	}
 }

package/tests/autocomplete.test.ts

@@ -0,0 +1,142 @@
+/** Tests for run-scoped Auto-complete and planning continuation. */
+
+import * as assert from "node:assert/strict";
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { after, before, describe, it } from "node:test";
+import {
+	autoCompleteStatus,
+	disableAutoComplete,
+	enableAutoComplete,
+	isAutoCompleteEnabled,
+	markPlanWritten,
+	recordAskChoice,
+	registerAutoCompleteTurnHandlers,
+	restoreAutoCompleteFromSession,
+	setAutoCompleteApi,
+} from "../src/autocomplete.ts";
+import { initState, setRunStatus, startRun } from "../src/state.ts";
+
+let root: string;
+
+function makeRun(name: string): { workdir: string; runId: string } {
+	const workdir = path.join(root, name);
+	fs.mkdirSync(workdir, { recursive: true });
+	initState(workdir);
+	const { run } = startRun(workdir, { topic: name, skill: "plan-small", requestText: "test" });
+	return { workdir, runId: run.run_id };
+}
+
+function makeContext(workdir: string, sessionManager: any = {}): any {
+	return {
+		cwd: workdir,
+		hasUI: true,
+		mode: "tui",
+		sessionManager,
+		ui: {
+			select: async (_question: string, options: string[]) => options.at(-1),
+			input: async () => "typed",
+		},
+	};
+}
+
+before(() => {
+	root = fs.mkdtempSync(path.join(os.tmpdir(), "pi-plans-autocomplete-"));
+});
+
+after(() => {
+	setAutoCompleteApi(null);
+	fs.rmSync(root, { recursive: true, force: true });
+});
+
+describe("Auto-complete state", () => {
+	it("enables only for the active planning run and restores only while planning", () => {
+		const { workdir, runId } = makeRun("restore");
+		const entries: any[] = [];
+		setAutoCompleteApi({ appendEntry: (customType: string, data: unknown) => entries.push({ type: "custom", customType, data }) } as any);
+		const ctx = makeContext(workdir);
+
+		assert.equal(enableAutoComplete(ctx), true);
+		assert.equal(autoCompleteStatus(ctx), "enabled");
+		assert.equal(entries.at(-1)?.data.runId, runId);
+
+		const restored = makeContext(workdir);
+		restoreAutoCompleteFromSession(restored, entries);
+		assert.equal(isAutoCompleteEnabled(restored), true);
+
+		setRunStatus(workdir, runId, "done");
+		const finished = makeContext(workdir);
+		restoreAutoCompleteFromSession(finished, entries);
+		assert.equal(autoCompleteStatus(finished), "disabled");
+
+		const stoppedEntries = [...entries, { type: "custom", customType: "pi-plans-autocomplete", data: { runId, enabled: false } }];
+		const stopped = makeContext(workdir);
+		restoreAutoCompleteFromSession(stopped, stoppedEntries);
+		assert.equal(autoCompleteStatus(stopped), "disabled");
+	});
+
+	it("clears explicitly and marks plan writes as a continuation boundary", () => {
+		const { workdir } = makeRun("clear");
+		const entries: any[] = [];
+		setAutoCompleteApi({ appendEntry: (customType: string, data: unknown) => entries.push({ customType, data }) } as any);
+		const ctx = makeContext(workdir);
+		enableAutoComplete(ctx);
+		markPlanWritten(ctx);
+		recordAskChoice(ctx, true);
+		assert.equal(disableAutoComplete(ctx, "test"), true);
+		assert.equal(autoCompleteStatus(ctx), "disabled");
+		assert.equal(entries.at(-1)?.data.enabled, false);
+	});
+});
+
+describe("Auto-complete continuation", () => {
+	it("queues one follow-up after an early stop and none after a natural next question", async () => {
+		const { workdir } = makeRun("continuation");
+		const session = {};
+		const sent: any[] = [];
+		const handlers = new Map<string, Function[]>();
+		const pi: any = {
+			on: (name: string, handler: Function) => handlers.set(name, [...(handlers.get(name) ?? []), handler]),
+			appendEntry: () => {},
+			sendUserMessage: async (content: string, options: unknown) => sent.push({ content, options }),
+		};
+		setAutoCompleteApi(pi);
+		registerAutoCompleteTurnHandlers(pi);
+		const ctx = makeContext(workdir, session);
+		enableAutoComplete(ctx);
+		await handlers.get("turn_start")?.[0]?.({}, ctx);
+		recordAskChoice(ctx, true);
+		await handlers.get("turn_end")?.[0]?.({ message: { role: "assistant" } }, ctx);
+		await handlers.get("turn_end")?.[0]?.({ message: { role: "assistant" } }, ctx);
+		assert.equal(sent.length, 1);
+		assert.equal((sent[0]!.options as any).deliverAs, "followUp");
+
+		await handlers.get("turn_start")?.[0]?.({}, ctx);
+		recordAskChoice(ctx, true);
+		markPlanWritten(ctx);
+		await handlers.get("turn_end")?.[0]?.({ message: { role: "assistant" } }, ctx);
+		assert.equal(sent.length, 1);
+
+		await handlers.get("turn_start")?.[0]?.({}, ctx);
+		recordAskChoice(ctx, true);
+		recordAskChoice(ctx, false);
+		await handlers.get("turn_end")?.[0]?.({ message: { role: "assistant" } }, ctx);
+		assert.equal(sent.length, 1);
+	});
+});
+
+describe("ask_choice Auto-complete wiring", () => {
+	it("routes eligible questions through the run-scoped mode", () => {
+		const source = fs.readFileSync(path.join(process.cwd(), "tools", "ask-choice.ts"), "utf8");
+		assert.match(source, /isAutoCompleteEnabled/);
+		assert.match(source, /recordAskChoice\(ctx, true\)/);
+		assert.match(source, /autoComplete && selected\.startsWith\("Auto-complete"\)/);
+		const indexSource = fs.readFileSync(path.join(process.cwd(), "index.ts"), "utf8");
+		assert.match(indexSource, /plans-autocomplete-stop/);
+		assert.match(indexSource, /autoCompleteStatus\(ctx\)/);
+		assert.match(indexSource, /restoreAutoCompleteFromSession/);
+		const execSource = fs.readFileSync(path.join(process.cwd(), "src", "exec.ts"), "utf8");
+		assert.match(execSource, /customType === AUTOCOMPLETE_ENTRY/);
+	});
+});

package/tests/compaction.test.ts

@@ -0,0 +1,74 @@
+import * as assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import {
+	extractReadRecords,
+	formatReadRecord,
+	legalFirstKeptEntryIndex,
+	mergeCompactionDetails,
+	planIAwareCompaction,
+	currentIExceedsTrigger,
+} from "../src/compaction.ts";
+
+function textEntry(id: string, text: string, tokens = 100) {
+	return { id, type: "message", tokens, message: { role: "assistant", content: [{ type: "text", text }] } };
+}
+
+describe("I-aware compaction policy", () => {
+	it("keeps a legal current-I suffix and never starts at a tool result", () => {
+		const entries = [
+			textEntry("i1", "[I-001:current] completed"),
+			{ id: "call", type: "message", tokens: 100, message: { role: "assistant", content: [{ type: "toolCall", id: "read-1", name: "read", arguments: { path: "/repo/a.ts", offset: 4, limit: 2 } }] } },
+			{ id: "result", type: "message", tokens: 100, message: { role: "toolResult", toolCallId: "read-1", content: [{ type: "text", text: "private source output that must be bounded" }] } },
+			textEntry("i2", "[I-002:current] active"),
+			{ id: "u", type: "message", tokens: 100, message: { role: "user", content: [{ type: "text", text: "latest question" }] } },
+			textEntry("a", "latest answer"),
+		];
+		const plan = planIAwareCompaction({ entries, currentI: "I-002", knownIIds: ["I-001", "I-002"], contextWindow: 1000, tokensBefore: 600 });
+		assert.equal(plan.currentI, "I-002");
+		assert.equal(plan.slices.filter((slice) => slice.id !== null).length, 2);
+		assert.ok(plan.firstKeptEntryId);
+		assert.notEqual(plan.firstKeptEntryId, "result");
+		assert.ok(plan.summaryEntries.every((entry) => !plan.keptEntries.includes(entry)));
+		assert.equal(legalFirstKeptEntryIndex(entries, 2), 1);
+	});
+
+	it("starts a new current-I slice after a prior compaction snapshot", () => {
+		const entries = [
+			{ id: "old-summary", type: "compaction", details: { currentI: "I-002" } },
+			textEntry("u", "new question", 100),
+			textEntry("a", "new answer", 100),
+		];
+		const plan = planIAwareCompaction({ entries, currentI: "I-002", knownIIds: ["I-001", "I-002"], contextWindow: 1000, tokensBefore: 200 });
+		assert.equal(plan.currentStartIndex, 1);
+		assert.equal(plan.slices.at(-1)?.id, "I-002");
+		assert.equal(plan.slices.at(-1)?.current, true);
+		assert.equal(plan.slices.at(-1)?.entries[0]?.id, "u");
+	});	it("extracts bounded paired Read records and merges by path/range", () => {
+		const raw = "x".repeat(500);
+		const entries = [
+			{ id: "call", type: "message", message: { role: "assistant", content: [{ type: "toolCall", id: "read-1", name: "read", arguments: { path: "/repo/a.ts", offset: 7, limit: 3 } }] } },
+			{ id: "result", type: "message", message: { role: "toolResult", toolCallId: "read-1", content: [{ type: "text", text: raw }] } },
+		];
+		const records = extractReadRecords(entries);
+		assert.equal(records.length, 1);
+		assert.equal(records[0]?.range, "7-9");
+		assert.equal(records[0]?.formatted, formatReadRecord(records[0]!));
+		assert.match(records[0]?.formatted ?? "", /^Read: \/repo\/a\.ts line 7-9 Extracted information summary: /);
+		assert.ok((records[0]?.formatted.length ?? 0) < raw.length);
+		const merged = mergeCompactionDetails(
+			{ readRecords: records },
+			{ readRecords: [{ ...records[0]!, summary: "new extraction", formatted: "" }] },
+		);
+		assert.equal(merged.readRecords?.length, 1);
+		assert.equal(merged.readRecords?.[0]?.summary, "new extraction");
+	});
+
+	it("uses strict trigger and records a hard floor when the retained suffix cannot fit", () => {
+		assert.equal(currentIExceedsTrigger(200, 1000), false);
+		assert.equal(currentIExceedsTrigger(201, 1000), true);
+		const entries = [textEntry("i1", "[I-001:current] " + "work ".repeat(40), 900), textEntry("u", "latest", 200), textEntry("a", "answer", 200)];
+		const plan = planIAwareCompaction({ entries, currentI: "I-001", contextWindow: 1000, tokensBefore: 1300 });
+		assert.equal(plan.metrics.targetMet, false);
+		assert.ok(plan.metrics.hardFloorReason);
+	});
+});