pi-ext

import assert from "node:assert/strict";
import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { initTheme } from "@earendil-works/pi-coding-agent";
import { chromium } from "@playwright/test";
import { afterEach, test } from "vitest";
import chromeExtension from "../index.ts";

const fixtureRoots: string[] = [];
afterEach(async () => {
	await Promise.all(
		fixtureRoots
			.splice(0)
			.map((root) => rm(root, { recursive: true, force: true })),
	);
});

test("Pi registration, /chrome toggle, raw command, screenshot, agent close", {
	timeout: 60000,
}, async () => {
	initTheme("default", false);
	const fixtureRoot = await mkdtemp(join(tmpdir(), "pi-chrome-extension-e2e-"));
	fixtureRoots.push(fixtureRoot);
	const artifactDir = join(fixtureRoot, "captures");
	await mkdir(join(fixtureRoot, ".pi"));
	await writeFile(
		join(fixtureRoot, ".pi", "settings.json"),
		JSON.stringify({
			"chrome-cdp": { artifactDir, executablePath: chromium.executablePath() },
		}),
	);
	const handlers: Record<string, Function> = {};
	let tool:
		| {
				execute: Function;
				renderCall: Function;
				renderResult: Function;
				description: string;
				promptSnippet: string;
				promptGuidelines: string[];
				parameters: { properties: Record<string, unknown> };
				constrainedSampling: unknown;
		  }
		| undefined;
	let names: string[] = [];
	let chromeFlag = false;
	const notices: string[] = [];
	const pi = {
		getActiveTools: () => names,
		setActiveTools: (next: string[]) => {
			names = next;
		},
		registerCommand: (name: string, command: { handler: Function }) => {
			handlers[name] = command.handler;
		},
		registerTool: (value: NonNullable<typeof tool>) => {
			tool = value;
		},
		registerFlag: (name: string, options: { type: string }) => {
			assert.equal(name, "chrome");
			assert.equal(options.type, "boolean");
		},
		getFlag: (name: string) => {
			assert.equal(name, "chrome");
			return chromeFlag;
		},
		on: (name: string, fn: Function) => {
			handlers[name] = fn;
		},
	};
	chromeExtension(pi as never);
	const ctx = {
		cwd: fixtureRoot,
		isProjectTrusted: () => true,
		mode: "test",
		ui: { notify: (message: string) => notices.push(message) },
	};
	await handlers.session_start(undefined, ctx);
	assert.equal(names.includes("chrome_cdp"), false);
	chromeFlag = true;
	await handlers.session_start(undefined, ctx);
	assert.equal(names.includes("chrome_cdp"), true);
	chromeFlag = false;
	await handlers.session_start(undefined, ctx);
	assert.equal(names.includes("chrome_cdp"), false);
	await handlers.chrome("", ctx);
	assert.equal(names.includes("chrome_cdp"), true);
	assert.equal(handlers.resources_discover, undefined);
	assert.ok(tool);
	assert.match(tool.description, /status\.protocol/);
	assert.match(tool.promptSnippet, /raw CDP/);
	assert.ok(tool.parameters.properties.operations);
	assert.deepEqual(tool.constrainedSampling, {
		type: "json_schema",
		strict: "prefer",
	});
	assert.match(tool.description, /chromedevtools\.github\.io/);
	assert.ok(
		tool.promptGuidelines.some((guideline) => guideline.includes("sessionId")),
	);
	const theme = {
		fg: (_color: string, value: string) => value,
		bold: (value: string) => value,
	};
	const call = tool.renderCall(
		{
			operations: [
				{
					kind: "cdp",
					method: "Browser.getVersion",
					params: { secret: "private" },
				},
				{ kind: "status" },
			],
		},
		theme,
		{ expanded: true },
	);
	assert.match(call.render(100).join(" "), /Browser.getVersion/);
	assert.doesNotMatch(call.render(100).join(" "), /private/);
	const partial = tool.renderResult(
		{
			content: [],
			details: {
				progress: { completed: 1, total: 2, active: "wait", elapsedMs: 1200 },
			},
		},
		{ expanded: false, isPartial: true },
		theme,
		{ isError: false },
	);
	assert.match(partial.render(100).join(" "), /1\/2 ops; wait \(1s\)/);
	let file: string | undefined;
	try {
		const updates: {
			content: { text: string }[];
			details: { progress: { completed: number; elapsedMs: number } };
		}[] = [];
		const res = await tool.execute(
			"test",
			{
				operations: [
					{ kind: "cdp", method: "Browser.getVersion", label: "version" },
					{ kind: "screenshot", sessionId: "{{missing.result.sessionId}}" },
				],
			},
			new AbortController().signal,
			(update: (typeof updates)[number]) => updates.push(update),
			ctx,
		);
		assert.ok(updates.length >= 2);
		assert.equal(updates[0].details.progress.completed, 0);
		assert.ok(
			updates.some((update) => update.details.progress.completed === 1),
		);
		assert.ok(
			updates.every((update) => !JSON.stringify(update).includes("Chrome/")),
		);
		const output = JSON.parse(res.content[0].text);
		assert.match(JSON.stringify(output.output[0]), /Chrome/);
		assert.ok(output.failure);
		const capture = await tool.execute(
			"test",
			{
				operations: [
					{ kind: "screenshot", sessionId: output.status.initialSession },
				],
			},
			new AbortController().signal,
			undefined,
			ctx,
		);
		assert.equal(capture.content[1].type, "image");
		file = JSON.parse(capture.content[0].text).output[0].file;
		assert.ok(file?.includes("chrome-cdp-"));
		assert.ok(file?.startsWith(artifactDir));
		const compact = tool.renderResult(
			capture,
			{ expanded: false, isPartial: false },
			theme,
			{ isError: false },
		);
		const expanded = tool.renderResult(
			capture,
			{ expanded: true, isPartial: false },
			theme,
			{ isError: false },
		);
		assert.match(compact.render(100).join(" "), /✓ 1\/1 operation · 1 file/);
		assert.doesNotMatch(compact.render(100).join(" "), /chrome-cdp-/);
		assert.match(expanded.render(100).join(" "), /chrome-cdp-/);
		assert.match(expanded.render(100).join(" "), /Result preview/);
		const waited: typeof updates = [];
		const waitResult = await tool.execute(
			"test",
			{
				operations: [
					{ kind: "cdp", method: "Browser.getVersion", label: "waitProbe" },
					{
						kind: "wait",
						after: "waitProbe",
						event: "Page.neverOccurs",
						sessionId: output.status.initialSession,
						timeoutMs: 1150,
					},
				],
			},
			new AbortController().signal,
			(update: (typeof updates)[number]) => waited.push(update),
			ctx,
		);
		assert.match(
			JSON.parse(waitResult.content[0].text).failure,
			/wait timed out/,
		);
		assert.ok(
			waited.some((update) => update.details.progress.elapsedMs >= 900),
		);
		const closed = await tool.execute(
			"test",
			{ operations: [{ kind: "close" }] },
			new AbortController().signal,
			undefined,
			ctx,
		);
		assert.equal(JSON.parse(closed.content[0].text).output[0].closed, true);
		await tool.execute(
			"test",
			{ operations: [{ kind: "cdp", method: "Browser.getVersion" }] },
			new AbortController().signal,
			undefined,
			ctx,
		);
		await handlers.chrome("close", ctx);
		assert.equal(notices.at(-1), "Chrome stopped; tool off");
		const alreadyClosed = await tool.execute(
			"test",
			{ operations: [{ kind: "close" }] },
			new AbortController().signal,
			undefined,
			ctx,
		);
		assert.equal(JSON.parse(alreadyClosed.content[0].text).closed, true);
	} finally {
		await handlers.session_shutdown();
		if (file) await rm(file, { force: true });
	}
	await handlers.chrome("close", ctx);
	assert.equal(names.includes("chrome_cdp"), false);
});