pi-ext

import assert from "node:assert";
import {
	access,
	mkdtemp,
	readdir,
	readFile,
	rm,
	stat,
	utimes,
	writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import {
	type ExtensionContext,
	initTheme,
} from "@earendil-works/pi-coding-agent";
import {
	resetCapabilitiesCache,
	setCapabilities,
} from "@earendil-works/pi-tui";
import { afterEach, beforeAll, describe, it, vi } from "vitest";
import type { MockupRequest, MockupResult } from "../mockup.ts";
import {
	type MockupPreview,
	type MockupPreviewManager,
	PreviewCancelledError,
} from "../mockup-preview.ts";
import {
	createMockupRuntime,
	formatContextCapsule,
	type MockupToolDetails,
	mockupCacheRoot,
	persistMockupIteration,
	pruneMockupCache,
} from "../mockup-tool.ts";

const request: MockupRequest = {
	title: "Layout",
	question: "Which layout?",
	options: [
		{ id: "dense", label: "Dense", html: "<section>dense</section>" },
		{ id: "calm", label: "Calm", html: "<section>calm</section>" },
	],
};
const result: MockupResult = {
	action: "iterate",
	selected: ["dense"],
	notes: "keep density",
	annotations: [{ optionId: "dense", x: 0.25, y: 0.5, note: "tighten" }],
	review: {
		submittedAt: "2026-01-02T03:04:05.000Z",
		userAgent: "test-browser/1",
		viewport: { width: 1280, height: 720, devicePixelRatio: 1 },
	},
};

const roots: string[] = [];
beforeAll(() => initTheme("default", false));
afterEach(async () => {
	await Promise.all(
		roots.splice(0).map((root) => rm(root, { recursive: true })),
	);
});

async function temporaryRoot(): Promise<string> {
	const root = await mkdtemp(path.join(tmpdir(), "pi-mockup-test-"));
	roots.push(root);
	return root;
}

async function exists(file: string): Promise<boolean> {
	try {
		await access(file);
		return true;
	} catch {
		return false;
	}
}

async function size(root: string): Promise<number> {
	let bytes = 0;
	for (const entry of await readdir(root, { withFileTypes: true })) {
		const child = path.join(root, entry.name);
		if (entry.isDirectory()) bytes += await size(child);
		else if (entry.isFile()) bytes += (await stat(child)).size;
	}
	return bytes;
}

function plainTheme() {
	return {
		fg: (_style: string, text: string) => text,
		bold: (text: string) => text,
	} as never;
}

function rendered(component: { render(width: number): string[] }): string {
	return component.render(120).join("\n");
}

function renderContext(
	expanded: boolean,
	lastComponent?: unknown,
	state: Partial<{
		isPartial: boolean;
		argsComplete: boolean;
		executionStarted: boolean;
	}> = {},
) {
	return {
		expanded,
		lastComponent,
		isError: false,
		isPartial: false,
		argsComplete: true,
		executionStarted: true,
		...state,
	} as never;
}

describe("mockup tool contract", () => {
	it("keeps concise prompt metadata while guidance lives in a skill", async () => {
		const tool = createMockupRuntime({ openBrowser: () => false }).tool;
		const guidancePath = path.resolve(
			import.meta.dirname,
			"../guidance/mockup-html/SKILL.md",
		);

		assert.strictEqual(tool.label, "mockup");
		assert.strictEqual(
			tool.promptSnippet,
			"Preview HTML mockups in Pi or review them interactively in a browser",
		);
		assert.notStrictEqual(tool.promptSnippet, tool.description);
		assert.deepStrictEqual(tool.promptGuidelines ?? [], []);
		assert.match(await readFile(guidancePath, "utf8"), /name: mockup-html/);
		assert.strictEqual(tool.parameters.properties.options.minItems, 1);
		assert.strictEqual(tool.parameters.properties.options.maxItems, 4);
		assert.strictEqual(
			tool.parameters.properties.options.items.properties.html.description,
			"HTML fragment; scripts are allowed.",
		);
	});

	it("renders a compact call and reserves option details for expansion", () => {
		const tool = createMockupRuntime({ openBrowser: () => false }).tool;
		assert.ok(tool.renderCall);
		const collapsedComponent = tool.renderCall(
			request,
			plainTheme(),
			renderContext(false),
		);
		const collapsed = rendered(collapsedComponent);
		const expandedComponent = tool.renderCall(
			request,
			plainTheme(),
			renderContext(true, collapsedComponent),
		);
		const expanded = rendered(expandedComponent);

		assert.strictEqual(expandedComponent, collapsedComponent);
		assert.match(collapsed, /^mockup Layout · 2 options/);
		assert.match(collapsed, /to expand/);
		assert.doesNotMatch(collapsed, /Which layout\?/);
		assert.match(expanded, /Which layout\?/);
		assert.match(expanded, /dense Dense/);
		assert.match(expanded, /calm Calm/);
		assert.doesNotMatch(expanded, /<section>/);

		const generating = rendered(
			tool.renderCall(
				request,
				plainTheme(),
				renderContext(false, undefined, {
					isPartial: true,
					argsComplete: false,
					executionStarted: false,
				}),
			),
		);
		assert.match(generating, /generating HTML · \d+ chars/);

		const queued = rendered(
			tool.renderCall(
				request,
				plainTheme(),
				renderContext(false, undefined, {
					isPartial: true,
					argsComplete: true,
					executionStarted: false,
				}),
			),
		);
		assert.match(queued, /waiting to run/);

		const restored = rendered(
			tool.renderCall(
				request,
				plainTheme(),
				renderContext(false, undefined, {
					isPartial: false,
					argsComplete: false,
					executionStarted: false,
				}),
			),
		);
		assert.doesNotMatch(restored, /generating|waiting to run/);
	});

	it("previews inside Pi without putting frame bytes in tool results", async () => {
		setCapabilities({ images: "kitty", trueColor: true, hyperlinks: true });
		try {
			let selected = 0;
			let startedWith: MockupRequest | undefined;
			const preview: MockupPreview = {
				labels: request.options.map((option) => option.label),
				done: Promise.resolve(),
				current: () => [
					"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=",
					"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=",
				],
				cropped: () => [false, true],
				expired: () => false,
				version: () => 0,
				subscribe: () => () => undefined,
				close: () => undefined,
			};
			const manager: MockupPreviewManager = {
				start: async (_id, value) => {
					startedWith = value;
					return preview;
				},
				get: () => preview,
				close: async () => undefined,
			};
			const ctx = {
				mode: "tui",
				ui: {
					select: async () => {
						selected++;
						return "Show previews in Pi";
					},
					notify: () => undefined,
				},
				sessionManager: {
					getSessionId: () => "preview-session",
					getBranch: () => [],
				},
			} as unknown as ExtensionContext;
			const runtime = createMockupRuntime({ previewManager: manager });
			const updates: Array<{ details?: unknown }> = [];
			const output = await runtime.tool.execute(
				"preview-call",
				request,
				undefined,
				(update) => updates.push(update),
				ctx,
			);

			assert.strictEqual(selected, 1);
			assert.deepStrictEqual(startedWith, { ...request, multi: false });
			assert.deepStrictEqual(
				updates.map((update) => update.details),
				[
					{ kind: "preview", status: "starting" },
					{ kind: "preview", status: "running" },
				],
			);
			assert.deepStrictEqual(output.details, {
				kind: "preview",
				status: "complete",
			});
			assert.doesNotMatch(JSON.stringify(output), /data:image|iVBOR/);
			assert.strictEqual(output.terminate, true);

			assert.ok(runtime.tool.renderResult);
			const hidden = runtime.tool.renderResult(
				output,
				{ expanded: false, isPartial: false },
				plainTheme(),
				{
					...renderContext(false),
					toolCallId: "preview-call",
					showImages: false,
					invalidate: () => undefined,
				} as never,
			);
			assert.match(rendered(hidden), /Inline image previews are disabled/);
			const visible = runtime.tool.renderResult(
				output,
				{ expanded: false, isPartial: false },
				plainTheme(),
				{
					...renderContext(false, hidden),
					toolCallId: "preview-call",
					showImages: true,
					invalidate: () => undefined,
				} as never,
			);
			assert.strictEqual(visible, hidden);
			assert.match(rendered(visible), /\x1b_G/);
			assert.match(rendered(visible), /Preview cropped · use browser review/);
			const cancelled = runtime.tool.renderResult(
				{
					content: [{ type: "text", text: "cancelled" }],
					details: { kind: "cancelled" },
				},
				{ expanded: false, isPartial: false },
				plainTheme(),
				{
					...renderContext(false, visible),
					toolCallId: "preview-call",
				} as never,
			);
			assert.match(rendered(cancelled), /Mockup cancelled/);
		} finally {
			resetCapabilitiesCache();
		}
	});

	it("does not offer browser fallback when preview startup is cancelled", async () => {
		setCapabilities({ images: "kitty", trueColor: true, hyperlinks: true });
		try {
			let selections = 0;
			let notifications = 0;
			const manager: MockupPreviewManager = {
				start: async () => {
					throw new PreviewCancelledError();
				},
				get: () => undefined,
				close: async () => undefined,
			};
			const ctx = {
				mode: "tui",
				ui: {
					select: async () => {
						selections++;
						return "Show previews in Pi";
					},
					notify: () => notifications++,
				},
				sessionManager: {
					getSessionId: () => "cancelled-preview-session",
					getBranch: () => [],
				},
			} as unknown as ExtensionContext;
			const output = await createMockupRuntime({
				previewManager: manager,
			}).tool.execute("cancelled-preview", request, undefined, undefined, ctx);

			assert.strictEqual(selections, 1);
			assert.strictEqual(notifications, 0);
			assert.deepStrictEqual(output.details, { kind: "cancelled" });
		} finally {
			resetCapabilitiesCache();
		}
	});

	it("keeps feedback compact until expanded and renders errors", () => {
		const tool = createMockupRuntime({ openBrowser: () => false }).tool;
		assert.ok(tool.renderResult);
		const toolResult = {
			content: [{ type: "text" as const, text: "structured result" }],
			details: {
				...result,
				sessionId: "session",
				chainId: "chain",
				iteration: 1,
				url: "http://127.0.0.1:1234/?k=secret",
				cacheManifest: "/cache/manifest.json",
				history: [],
				contextCapsule: "capsule",
			},
		};
		const collapsedComponent = tool.renderResult(
			toolResult,
			{ expanded: false, isPartial: false },
			plainTheme(),
			renderContext(false),
		);
		const collapsed = rendered(collapsedComponent);
		const expandedComponent = tool.renderResult(
			toolResult,
			{ expanded: true, isPartial: false },
			plainTheme(),
			renderContext(true, collapsedComponent),
		);
		const expanded = rendered(expandedComponent);

		assert.strictEqual(expandedComponent, collapsedComponent);
		assert.match(collapsed, /^↻ iterate · selected dense · 1 annotation/);
		assert.match(collapsed, /to expand/);
		assert.doesNotMatch(collapsed, /keep density|tighten|cache/);
		assert.match(expanded, /notes: keep density/);
		assert.match(expanded, /dense @ 25%,50% tighten/);
		assert.match(expanded, /http:\/\/127\.0\.0\.1:1234\/\?k=secret/);
		assert.match(expanded, /cache: \/cache\/manifest\.json/);

		const waiting = rendered(
			tool.renderResult(
				{
					content: [{ type: "text", text: "Waiting for browser feedback" }],
					details: {
						kind: "waiting",
						status: "waiting",
						url: "http://127.0.0.1:4321/?k=secret",
						iteration: 2,
						browserFailed: true,
					},
				},
				{ expanded: false, isPartial: true },
				plainTheme(),
				renderContext(false),
			),
		);
		assert.match(waiting, /Browser did not open/);
		assert.match(waiting, /http:\/\/127\.0\.0\.1:4321\/\?k=secret/);

		const accepted = rendered(
			tool.renderResult(
				{
					...toolResult,
					details: {
						...toolResult.details,
						action: "accept",
						selected: [],
						annotations: [],
					},
				},
				{ expanded: false, isPartial: false },
				plainTheme(),
				renderContext(false),
			),
		);
		assert.match(accepted, /^✓ accepted · selected none · 0 annotations/);

		const errorResult = {
			content: [{ type: "text", text: "browser failed\nEND-BROWSER-DETAILS" }],
			details: {},
		};
		const error = rendered(
			tool.renderResult(
				errorResult,
				{ expanded: false, isPartial: false },
				plainTheme(),
				{ ...renderContext(false), isError: true } as never,
			),
		);
		const expandedError = rendered(
			tool.renderResult(
				errorResult,
				{ expanded: true, isPartial: false },
				plainTheme(),
				{ ...renderContext(true), isError: true } as never,
			),
		);
		assert.match(error, /^browser failed/);
		assert.match(error, /to expand/);
		assert.doesNotMatch(error, /END-BROWSER-DETAILS/);
		assert.match(expandedError, /END-BROWSER-DETAILS/);
	});
});

describe("mockup runtime lifecycle", () => {
	it("pins active browser iterations and prompts again after abandonment", async () => {
		setCapabilities({ images: "kitty", trueColor: true, hyperlinks: true });
		const root = await temporaryRoot();
		const branch: unknown[] = [];
		const opened: string[] = [];
		const notifications: string[] = [];
		let selections = 0;
		const ctx = {
			mode: "tui",
			hasUI: true,
			isIdle: () => false,
			abort: () => undefined,
			ui: {
				notify: (message: string) => notifications.push(message),
				select: async () => {
					selections++;
					return "Open interactive browser";
				},
			},
			sessionManager: {
				getSessionId: () => "session",
				getSessionFile: () => undefined,
				getBranch: () => branch,
			},
		} as unknown as ExtensionContext;
		const runtime = createMockupRuntime({
			cacheRoot: root,
			openBrowser: () => true,
			browserOpener: async (url) => {
				opened.push(url);
				return true;
			},
		});

		const updates: Array<{ details?: unknown }> = [];
		const firstExecution = runtime.tool.execute(
			"one",
			request,
			undefined,
			(update) => updates.push(update),
			ctx,
		);
		await vi.waitFor(() => assert.ok(updates.at(-1)?.details));
		const firstProgress = updates.at(-1)?.details as
			| { url?: string }
			| undefined;
		assert.ok(firstProgress?.url);
		await submit(firstProgress.url, "iterate");
		const first = await firstExecution;
		const firstDetails = first.details as MockupToolDetails;
		assert.strictEqual(firstDetails.url, firstProgress.url);
		assert.deepStrictEqual(
			updates.map(
				(update) => (update.details as { status?: string } | undefined)?.status,
			),
			["opening", "waiting"],
		);
		const measuredProgress = updates.at(-1)?.details as
			| { timings?: { serverMs?: number; openerMs?: number } }
			| undefined;
		assert.strictEqual(typeof measuredProgress?.timings?.serverMs, "number");
		assert.strictEqual(typeof measuredProgress?.timings?.openerMs, "number");
		assert.match(
			(first.content[0] as { text: string }).text,
			/Apply the feedback and call mockup again/,
		);
		branch.push({
			type: "message",
			message: {
				role: "toolResult",
				toolName: "mockup",
				details: firstDetails,
			},
		});

		updates.length = 0;
		const secondExecution = runtime.tool.execute(
			"two",
			request,
			undefined,
			(update) => updates.push(update),
			ctx,
		);
		await vi.waitFor(() => assert.ok(updates.at(-1)?.details));
		const secondProgress = updates.at(-1)?.details as
			| { url?: string }
			| undefined;
		assert.ok(secondProgress?.url);
		assert.strictEqual(secondProgress.url, firstProgress.url);
		await submit(secondProgress.url, "iterate");
		const second = await secondExecution;
		const secondDetails = second.details as MockupToolDetails;
		branch.push({
			type: "message",
			message: {
				role: "toolResult",
				toolName: "mockup",
				details: secondDetails,
			},
		});
		assert.deepStrictEqual(opened, [firstProgress.url, secondProgress.url]);

		runtime.agentSettled(ctx);
		updates.length = 0;
		const thirdExecution = runtime.tool.execute(
			"three",
			request,
			undefined,
			(update) => updates.push(update),
			ctx,
		);
		await vi.waitFor(() => assert.ok(updates.at(-1)?.details));
		const thirdProgress = updates.at(-1)?.details as
			| { url?: string }
			| undefined;
		assert.ok(thirdProgress?.url);
		assert.notStrictEqual(thirdProgress.url, secondProgress.url);
		await submit(thirdProgress.url, "accept");
		const third = await thirdExecution;
		const thirdDetails = third.details as MockupToolDetails;
		assert.strictEqual(thirdDetails.chainId, firstDetails.chainId);
		assert.strictEqual(thirdDetails.iteration, 3);
		assert.deepStrictEqual(opened, [
			firstProgress.url,
			secondProgress.url,
			thirdProgress.url,
		]);
		assert.strictEqual(selections, 2);
		assert.deepStrictEqual(notifications, []);
		resetCapabilitiesCache();
	});
});

async function submit(
	url: string,
	action: MockupResult["action"],
): Promise<void> {
	const response = await fetch(new URL("/api/submit", url), {
		method: "POST",
		headers: {
			Authorization: `Bearer ${new URL(url).searchParams.get("k")}`,
			"Content-Type": "application/json",
		},
		body: JSON.stringify({
			action,
			selected: [],
			notes: "",
			annotations: [],
			client: {
				userAgent: "mockup-tool-test/1",
				viewport: { width: 1280, height: 720, devicePixelRatio: 1 },
			},
		}),
	});
	assert.strictEqual(response.status, 200);
}

describe("mockup context cache", () => {
	it("uses the native OS cache location", () => {
		assert.strictEqual(
			mockupCacheRoot("linux", { XDG_CACHE_HOME: "/cache" }, "/home/o"),
			path.join("/cache", "pi", "mockups"),
		);
		assert.strictEqual(
			mockupCacheRoot(
				"win32",
				{ LOCALAPPDATA: "C:\\Users\\o\\AppData\\Local" },
				"C:\\Users\\o",
			),
			path.join("C:\\Users\\o\\AppData\\Local", "pi", "cache", "mockups"),
		);
		assert.strictEqual(
			mockupCacheRoot("darwin", {}, "/Users/o"),
			path.join("/Users/o", "Library", "Caches", "pi", "mockups"),
		);
	});

	it("persists exact fragments, request, feedback, metadata, and lineage", async () => {
		const root = await temporaryRoot();
		const sessionFile = path.join(root, "session.jsonl");
		await writeFile(sessionFile, "");
		const first = await persistMockupIteration(root, {
			sessionId: "session/one",
			sessionFile,
			chainId: "chain-one",
			iteration: 1,
			toolCallId: "tool/one",
			request,
			result,
		});
		const second = await persistMockupIteration(root, {
			sessionId: "session/one",
			sessionFile,
			chainId: "chain-one",
			iteration: 2,
			toolCallId: "tool-two",
			request,
			result: { ...result, action: "accept", selected: [] },
			previousManifest: first,
		});

		assert.strictEqual(
			await readFile(path.join(path.dirname(first), "dense.html"), "utf8"),
			request.options[0].html,
		);
		assert.deepStrictEqual(
			JSON.parse(
				await readFile(path.join(path.dirname(first), "request.json"), "utf8"),
			),
			request,
		);
		assert.deepStrictEqual(
			JSON.parse(
				await readFile(path.join(path.dirname(first), "feedback.json"), "utf8"),
			),
			result,
		);
		const manifest = JSON.parse(await readFile(second, "utf8"));
		assert.strictEqual(manifest.previousManifest, first);
		assert.strictEqual(manifest.options[0].html, "dense.html");
	});

	it("formats current feedback first and preserves earlier iteration context", () => {
		const capsule = formatContextCapsule("chain", [
			{
				iteration: 1,
				question: request.question,
				options: request.options.map(({ id, label }) => ({ id, label })),
				action: "iterate",
				selected: ["dense"],
				notes: "keep density",
				annotations: result.annotations,
				cacheManifest: "/cache/one/manifest.json",
			},
			{
				iteration: 2,
				question: "Now?",
				options: request.options.map(({ id, label }) => ({ id, label })),
				action: "accept",
				selected: [],
				notes: "none selected intentionally",
				annotations: [],
				cacheManifest: "/cache/two/manifest.json",
			},
		]);
		assert.ok(
			capsule.indexOf("Current: iteration 2, accept") <
				capsule.indexOf("Earlier iterations"),
		);
		assert.match(capsule, /Selected: none/);
		assert.match(capsule, /never infer rejection from non-selection/);
		assert.match(
			capsule,
			/1 iterate; selected dense \(Dense\); notes: keep density/,
		);
		assert.match(capsule, /annotation dense: tighten/);
	});

	it("injects cached context only into its owning session", async () => {
		const root = await temporaryRoot();
		const capsule = "owned mockup context";
		const details = {
			...result,
			sessionId: "owner",
			chainId: "chain",
			iteration: 1,
			url: "http://127.0.0.1:1234/?k=secret",
			history: [],
			contextCapsule: capsule,
		};
		const context = (sessionId: string) =>
			({
				isIdle: () => true,
				sessionManager: {
					getSessionId: () => sessionId,
					getBranch: () => [
						{
							type: "message",
							message: {
								role: "toolResult",
								toolName: "mockup",
								details,
							},
						},
					],
				},
			}) as unknown as ExtensionContext;
		const runtime = createMockupRuntime({
			cacheRoot: root,
			openBrowser: () => false,
		});
		assert.strictEqual(runtime.contextCapsule(context("owner")), capsule);
		assert.strictEqual(
			runtime.contextCapsule(context("new-session")),
			undefined,
		);
	});

	it("removes orphaned and expired entries while protecting the current session", async () => {
		const root = await temporaryRoot();
		const sessionFile = path.join(root, "live.jsonl");
		await writeFile(sessionFile, "");
		const live = await persistMockupIteration(root, {
			sessionId: "live",
			sessionFile,
			chainId: "chain",
			iteration: 1,
			toolCallId: "one",
			request,
			result,
		});
		const orphan = await persistMockupIteration(root, {
			sessionId: "orphan",
			sessionFile: path.join(root, "missing.jsonl"),
			chainId: "chain",
			iteration: 1,
			toolCallId: "one",
			request,
			result,
		});
		const old = new Date(Date.now() - 10_000);
		await utimes(path.dirname(live), old, old);

		await pruneMockupCache(root, "live", Date.now(), {
			ttlMs: 1,
			maxBytes: Number.MAX_SAFE_INTEGER,
		});
		assert.strictEqual(await exists(live), true);
		assert.strictEqual(await exists(orphan), false);
		await pruneMockupCache(root, undefined, Date.now(), {
			ttlMs: 1,
			maxBytes: Number.MAX_SAFE_INTEGER,
		});
		assert.strictEqual(await exists(live), false);
	});

	it("evicts least-recently-used iterations above the size cap", async () => {
		const root = await temporaryRoot();
		const sessionFile = path.join(root, "live.jsonl");
		await writeFile(sessionFile, "");
		const oldest = await persistMockupIteration(root, {
			sessionId: "live",
			sessionFile,
			chainId: "chain",
			iteration: 1,
			toolCallId: "old",
			request,
			result,
		});
		const newest = await persistMockupIteration(root, {
			sessionId: "live",
			sessionFile,
			chainId: "chain",
			iteration: 2,
			toolCallId: "new",
			request,
			result,
		});
		await utimes(path.dirname(oldest), new Date(1000), new Date(1000));
		await utimes(path.dirname(newest), new Date(2000), new Date(2000));
		const newestBytes = await size(path.dirname(newest));

		await pruneMockupCache(root, undefined, Date.now(), {
			ttlMs: Number.MAX_SAFE_INTEGER,
			maxBytes: newestBytes,
		});
		assert.strictEqual(await exists(oldest), false);
		assert.strictEqual(await exists(newest), true);
	});
});