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,
	type ExtensionToolContext,
	initTheme,
} from "@earendil-works/pi-coding-agent";
import {
	resetCapabilitiesCache,
	setCapabilities,
} from "@earendil-works/pi-tui";
import { afterEach, beforeAll, describe, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import type { MockupRequest, MockupResult } from "../mockup.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[] = [];
let piSession: PiSession | undefined;
beforeAll(() => initTheme("default", false));
afterEach(async () => {
	await piSession?.dispose();
	piSession = undefined;
	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");
}

type RenderCtx = Parameters<
	NonNullable<ReturnType<typeof createMockupRuntime>["tool"]["renderCall"]>
>[2];

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

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.exposure, "model-only");
		assert.strictEqual(tool.renderShell, undefined);
		assert.strictEqual(
			tool.promptSnippet,
			"Review HTML mockups 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("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" as const, 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;
		piSession = await createPiSession({});
		const createContext = (toolCallId: string): ExtensionToolContext => {
			if (!piSession) throw new Error("missing Pi session");
			const context = Object.create(
				piSession.session.extensionRunner.createToolContext(
					toolCallId,
					undefined,
				),
			) as ExtensionToolContext;
			Object.defineProperties(context, {
				mode: { value: "tui" },
				hasUI: { value: true },
				isIdle: { value: () => false },
				abort: { value: () => undefined },
				ui: {
					value: {
						notify: (message: string) => notifications.push(message),
						select: async () => {
							selections++;
							return "Open interactive browser";
						},
					},
				},
				sessionManager: {
					value: {
						getSessionId: () => "session",
						getSessionFile: () => undefined,
						getBranch: () => branch,
					},
				},
			});
			return context;
		};
		const ctx = createContext("one");
		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),
			createContext("two"),
		);
		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),
			createContext("three"),
		);
		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);
	});
});