pi-ext

import { EventEmitter } from "node:events";
import { existsSync, mkdtempSync, readdirSync } from "node:fs";
import { chmod, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import { initTheme, keyHint } from "@earendil-works/pi-coding-agent";
import { afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import { launch, teardown } from "../index.ts";

const extensionPath = path.resolve(import.meta.dirname, "../index.ts");

describe("firefox-bidi pi-test-harness", () => {
	let t: TestSession | undefined;

	beforeAll(() => initTheme("default", false));

	afterEach(() => {
		t?.dispose();
		t = undefined;
	});

	it("loads the tool inert, the command with completions, and the bundled skill", async () => {
		t = await createTestSession({ extensions: [extensionPath] });
		const runner = t.session.extensionRunner;
		expect(runner.extensions).toHaveLength(1);
		expect(runner.extensions[0]?.path).toBe(extensionPath);

		const tool = runner.extensions[0]?.tools.get("firefox_bidi")?.definition;
		expect(tool).toBeDefined();
		expect(tool?.description).toContain("WebDriver BiDi");
		expect(tool?.promptGuidelines).toHaveLength(2);

		// tool ships inert: not advertised until /firefox
		expect(t.session.getActiveToolNames()).not.toContain("firefox_bidi");
		expect(t.session.systemPrompt).not.toContain(
			"Drive Firefox with raw WebDriver BiDi frames",
		);

		const command = runner.getCommand("firefox");
		expect(command).toBeDefined();
		expect(command?.getArgumentCompletions?.("")).toEqual([
			{ value: "close", label: "close" },
		]);
		expect(command?.getArgumentCompletions?.("clo")).toEqual([
			{ value: "close", label: "close" },
		]);
		expect(command?.getArgumentCompletions?.("x")).toBeNull();

		const skill = t.session.resourceLoader
			.getSkills()
			.skills.find((candidate) => candidate.name === "firefox-bidi");
		expect(skill).toBeDefined();
	});

	it("renders compact status, partial progress, and expanded frame metadata", async () => {
		t = await createTestSession({ extensions: [extensionPath] });
		const tool =
			t.session.extensionRunner.extensions[0]?.tools.get(
				"firefox_bidi",
			)?.definition;
		expect(tool?.executionMode).toBe("sequential");
		expect(tool?.promptSnippet).toContain("Firefox");
		const theme = {
			fg: (_color: string, text: string) => text,
			bold: (text: string) => text,
		};
		const args = {
			frames: [
				{ method: "session.new", params: { marker: "private-value" } },
				{ method: "script.evaluate" },
			],
		};
		const call = (expanded: boolean) =>
			tool
				?.renderCall?.(args as never, theme as never, { expanded } as never)
				.render(100)
				.join("\n");
		expect(call(false)).toContain(" Firefox · session.new → script.evaluate");
		expect(call(false)).toContain(keyHint("app.tools.expand", "to expand"));
		expect(call(false)).not.toContain("private-value");
		expect(call(true)).toContain("2. script.evaluate");
		expect(call(true)).not.toContain("to expand");

		const details = {
			sessionId: "session-1",
			responses: 2,
			events: 3,
			frames: [
				{ id: 1, method: "session.new", status: "success" },
				{ id: 2, method: "script.evaluate", status: "error" },
			],
			wait: { timedOut: true, method: "browsingContext.load" },
		};
		const result = {
			content: [{ type: "text" as const, text: "raw data" }],
			details,
		};
		const render = (expanded: boolean, isPartial: boolean) =>
			tool
				?.renderResult?.(
					result,
					{ expanded, isPartial } as never,
					theme as never,
					{} as never,
				)
				.render(100)
				.join("\n");
		expect(render(false, false)).toContain("✗ 2 frames · 1 failed/skipped");
		expect(render(false, false)).not.toContain("3 events");
		expect(render(false, false)).not.toContain("script.evaluate");
		expect(render(false, false)).toContain("1 failed/skipped");
		expect(render(true, false)).toContain("2 script.evaluate: error");
		expect(render(true, false)).toContain("3 events");
		expect(render(true, false)).toContain("raw data");
		expect(render(true, false)).toContain("session: session-1");
		expect(render(false, true)).toContain("running");
		expect(render(false, false)).toContain("wait timed out");

		const data = {
			...result,
			content: [
				{
					type: "text" as const,
					text: JSON.stringify({
						session: { id: "session-1" },
						responses: [{ id: 2, result: { value: "page title" } }],
						events: [{ method: "noise" }],
					}),
				},
			],
		};
		const preview = (expanded: boolean) =>
			tool
				?.renderResult?.(
					data,
					{ expanded, isPartial: false } as never,
					theme as never,
					{} as never,
				)
				.render(100)
				.join("\n");
		expect(preview(false)).not.toContain("page title");
		expect(preview(true)).toContain("page title");
		expect(preview(true)).not.toContain("noise");
	});

	it("activates via /firefox and deactivates via /firefox close without launching", async () => {
		t = await createTestSession({
			extensions: [extensionPath],
			env: { FIREFOX_BIN: "/nonexistent/firefox-bin" },
		});
		const runner = t.session.extensionRunner;
		const command = runner.getCommand("firefox");
		const context = {
			...runner.createCommandContext(),
			mode: "print" as const,
			hasUI: false,
		};

		expect(t.session.getActiveToolNames()).not.toContain("firefox_bidi");
		await command?.handler("", context);
		expect(t.session.getActiveToolNames()).toContain("firefox_bidi");
		await command?.handler("close", context);
		expect(t.session.getActiveToolNames()).not.toContain("firefox_bidi");
	});

	it("activates on session_start with the PI_FIREFOX=1 dev hook", async () => {
		t = await createTestSession({
			extensions: [extensionPath],
			env: { PI_FIREFOX: "1", FIREFOX_BIN: "/nonexistent/firefox-bin" },
		});
		expect(t.session.getActiveToolNames()).toContain("firefox_bidi");
	});

	it("fails fast with a clear error when the firefox binary is missing", async () => {
		t = await createTestSession({
			extensions: [extensionPath],
			env: { FIREFOX_BIN: "/nonexistent/firefox-bin" },
		});
		const runner = t.session.extensionRunner;
		const command = runner.getCommand("firefox");
		const context = {
			...runner.createCommandContext(),
			mode: "print" as const,
			hasUI: false,
		};
		await command?.handler("", context);

		const tool = runner.extensions[0]?.tools.get("firefox_bidi")?.definition;
		await expect(
			tool?.execute(
				"call-1",
				{ frames: [{ method: "session.status" }] },
				undefined,
				undefined,
				context,
			),
		).rejects.toThrow(/firefox not found/);
	});

	it("retains profile and spool when fake child survives termination", async () => {
		const profileDir = mkdtempSync(path.join(tmpdir(), "pi-firefox-refuse-"));
		const spoolDir = mkdtempSync(path.join(tmpdir(), "pi-firefox-refuse-"));
		const child = Object.assign(new EventEmitter(), {
			exitCode: null,
			signalCode: null,
			kill: vi.fn(() => false),
		});
		try {
			await expect(
				teardown({
					child,
					profileDir,
					spoolDir,
					ws: null,
					waiters: [],
					spoolBytes: 0,
				} as never),
			).rejects.toThrow(/profile retained/);
			expect(child.kill).toHaveBeenCalledTimes(2);
			expect(existsSync(profileDir)).toBe(true);
			expect(existsSync(spoolDir)).toBe(true);
		} finally {
			await rm(profileDir, { recursive: true, force: true });
			await rm(spoolDir, { recursive: true, force: true });
		}
	}, 10_000);

	it("retains ownership after a failed launch with unkillable child", async () => {
		const child = Object.assign(new EventEmitter(), {
			exitCode: null,
			signalCode: null,
			kill: vi.fn(() => false),
			stdout: new EventEmitter(),
			stderr: new EventEmitter(),
		});
		const originalTimeout = globalThis.setTimeout;
		const timer = vi
			.spyOn(globalThis, "setTimeout")
			.mockImplementation(((
				fn: (...args: unknown[]) => void,
				ms?: number,
				...args: unknown[]
			) =>
				originalTimeout(
					fn,
					ms === 30_000 ? 5 : ms,
					...args,
				)) as typeof setTimeout);
		let survivor: Parameters<typeof teardown>[0] = null;
		try {
			await expect(
				launch(
					{},
					(b) => {
						survivor = b;
					},
					(() => child) as never,
					async () => false,
				),
			).rejects.toThrow(/still alive/);
			if (!survivor) throw new Error("lost survivor ownership");
			expect(existsSync(survivor.profileDir)).toBe(true);
			child.exitCode = 0;
			await teardown(survivor);
			expect(existsSync(survivor.profileDir)).toBe(false);
		} finally {
			timer.mockRestore();
			if (survivor) {
				await rm(survivor.profileDir, { recursive: true, force: true });
				await rm(survivor.spoolDir, { recursive: true, force: true });
			}
		}
	});

	it.runIf(process.platform !== "win32")(
		"rolls back profile and spool on process spawn error",
		async () => {
			const dir = mkdtempSync(path.join(tmpdir(), "pi-firefox-fake-"));
			const before = new Set(
				readdirSync(tmpdir()).filter(
					(name) =>
						name.startsWith("pi-firefox-") &&
						!name.startsWith("pi-firefox-fake-"),
				),
			);
			t = await createTestSession({
				extensions: [extensionPath],
				env: { FIREFOX_BIN: dir },
			});
			const runner = t.session.extensionRunner;
			const tool = runner.extensions[0]?.tools.get("firefox_bidi")?.definition;
			await expect(
				tool?.execute(
					"spawn-error",
					{ frames: [] },
					undefined,
					undefined,
					runner.createCommandContext(),
				),
			).rejects.toThrow(/EACCES/);
			const after = readdirSync(tmpdir()).filter(
				(name) =>
					name.startsWith("pi-firefox-") &&
					!name.startsWith("pi-firefox-fake-"),
			);
			expect(after.filter((name) => !before.has(name))).toEqual([]);
		},
	);

	it.runIf(process.platform !== "win32")(
		"rolls back listener timeout and terminates fake browser before profile deletion",
		async () => {
			const dir = mkdtempSync(path.join(tmpdir(), "pi-firefox-fake-"));
			const bin = path.join(dir, "firefox");
			await writeFile(bin, "#!/bin/sh\nexec sleep 1000\n");
			await chmod(bin, 0o755);
			const before = new Set(
				readdirSync(tmpdir()).filter(
					(name) =>
						name.startsWith("pi-firefox-") &&
						!name.startsWith("pi-firefox-fake-"),
				),
			);
			const originalTimeout = globalThis.setTimeout;
			const timer = vi
				.spyOn(globalThis, "setTimeout")
				.mockImplementation(((
					fn: (...args: unknown[]) => void,
					ms?: number,
					...args: unknown[]
				) =>
					originalTimeout(
						fn,
						ms === 30_000 ? 20 : ms,
						...args,
					)) as typeof setTimeout);
			try {
				t = await createTestSession({
					extensions: [extensionPath],
					env: { FIREFOX_BIN: bin },
				});
				const runner = t.session.extensionRunner;
				const tool =
					runner.extensions[0]?.tools.get("firefox_bidi")?.definition;
				await expect(
					tool?.execute(
						"listener-timeout",
						{ frames: [] },
						undefined,
						undefined,
						runner.createCommandContext(),
					),
				).rejects.toThrow(/no BiDi listener/);
				const after = readdirSync(tmpdir()).filter(
					(name) =>
						name.startsWith("pi-firefox-") &&
						!name.startsWith("pi-firefox-fake-"),
				);
				expect(after.filter((name) => !before.has(name))).toEqual([]);
			} finally {
				timer.mockRestore();
			}
		},
	);

	it.runIf(process.platform !== "win32")(
		"serializes tool and close, rejects incompatible launch options, and waits causally",
		async () => {
			const dir = mkdtempSync(path.join(tmpdir(), "pi-firefox-fake-"));
			const bin = path.join(dir, "firefox");
			await writeFile(
				bin,
				"#!/bin/sh\necho 'WebDriver BiDi listening on ws://127.0.0.1:9922' >&2\nexec sleep 1000\n",
			);
			await chmod(bin, 0o755);
			const original = globalThis.WebSocket;
			class Socket {
				static OPEN = 1;
				readyState = 1;
				onopen?: () => void;
				onclose?: () => void;
				onmessage?: (message: { data: string }) => void;
				onerror?: () => void;
				static sockets: Socket[] = [];
				constructor(_url: string) {
					Socket.sockets.push(this);
					queueMicrotask(() => this.onopen?.());
				}
				send(text: string) {
					const frame = JSON.parse(text);
					queueMicrotask(() =>
						this.onmessage?.({
							data: JSON.stringify({
								id: frame.id,
								type: "success",
								result:
									frame.method === "script.evaluate"
										? { value: "x".repeat(300 * 1024) }
										: {},
							}),
						}),
					);
				}
				close() {
					this.readyState = 3;
					this.onclose?.();
				}
				event(method: string, context: string) {
					this.onmessage?.({
						data: JSON.stringify({
							type: "event",
							method,
							params: { context },
						}),
					});
				}
			}
			(globalThis as { WebSocket: unknown }).WebSocket = Socket;
			let retainedFile: string | undefined;
			try {
				t = await createTestSession({
					extensions: [extensionPath],
					env: { FIREFOX_BIN: bin },
				});
				const runner = t.session.extensionRunner;
				const tool =
					runner.extensions[0]?.tools.get("firefox_bidi")?.definition;
				if (!tool) throw new Error("missing firefox_bidi tool");
				const ctx = runner.createCommandContext();
				const execute = (args: Record<string, unknown>) =>
					tool.execute("test", args as never, undefined, undefined, ctx);
				await execute({ frames: [{ id: 40, method: "session.status" }] });
				const other = await createTestSession({
					extensions: [extensionPath],
					env: { FIREFOX_BIN: bin },
				});
				try {
					const second = other.session.extensionRunner;
					const otherTool =
						second.extensions[0]?.tools.get("firefox_bidi")?.definition;
					await otherTool?.execute(
						"separate",
						{ frames: [{ id: 40, method: "session.status" }] },
						undefined,
						undefined,
						second.createCommandContext(),
					);
				} finally {
					await other.session.extensionRunner
						.getCommand("firefox")
						?.handler(
							"close",
							other.session.extensionRunner.createCommandContext(),
						);
					other.dispose();
				}
				const firstSocket = Socket.sockets[0];
				firstSocket?.event("browsingContext.load", "old");
				await expect(
					execute({ frames: [{ id: 40, method: "session.status" }] }),
				).rejects.toThrow(/reused frame id/);
				await expect(
					execute({ frames: [], launch: { access: true } }),
				).rejects.toThrow(/different launch options/);
				const pending = execute({
					frames: [],
					wait: "browsingContext.load",
					waitContext: "new",
					waitTimeoutMs: 100,
				});
				await new Promise((resolve) => setTimeout(resolve, 5));
				firstSocket?.event("browsingContext.load", "wrong");
				firstSocket?.event("browsingContext.load", "new");
				const result = await pending;
				expect(JSON.parse(result.content[0].text).wait.params.context).toBe(
					"new",
				);
				const many = execute({
					frames: [],
					wait: ["browsingContext.load", "browsingContext.load"],
					waitTimeoutMs: 100,
				});
				await new Promise((resolve) => setTimeout(resolve, 5));
				firstSocket?.event("browsingContext.load", "first");
				firstSocket?.event("browsingContext.load", "second");
				expect(
					JSON.parse((await many).content[0].text).wait.map(
						(event: { params: { context: string } }) => event.params.context,
					),
				).toEqual(["first", "second"]);
				const oversized = execute({
					frames: [],
					wait: "log.entryAdded",
					waitTimeoutMs: 100,
				});
				await new Promise((resolve) => setTimeout(resolve, 5));
				firstSocket?.onmessage?.({
					data: JSON.stringify({
						type: "event",
						method: "log.entryAdded",
						params: {
							context: "secret".repeat(100_000),
							data: "🚀".repeat(10_000),
						},
					}),
				});
				const trimmed = await oversized;
				expect(JSON.parse(trimmed.content[0].text).wait.params.truncated).toBe(
					true,
				);
				expect(JSON.stringify(trimmed.details).length).toBeLessThan(10_000);
				if (!firstSocket) throw new Error("missing first socket");
				firstSocket.readyState = 3;
				await execute({ frames: [] });
				const current = Socket.sockets.at(-1);
				if (!current) throw new Error("missing replacement socket");
				const replacementWait = execute({
					frames: [],
					wait: "browsingContext.load",
					waitTimeoutMs: 100,
				});
				await new Promise((resolve) => setTimeout(resolve, 5));
				firstSocket?.event("browsingContext.load", "stale");
				firstSocket?.close();
				current.event("browsingContext.load", "current");
				expect(
					JSON.parse((await replacementWait).content[0].text).wait.params
						.context,
				).toBe("current");
				const big = await execute({ frames: [{ method: "script.evaluate" }] });
				const file = JSON.parse(big.content[0].text).responses[0].result
					.spooled as string;
				retainedFile = file;
				expect(existsSync(file)).toBe(true);
				const command = runner.getCommand("firefox");
				await command?.handler("close", ctx);
				expect(existsSync(file)).toBe(true);
				await runner.emit({ type: "session_start", reason: "new" });
				expect(existsSync(file)).toBe(false);
				const nextBig = await execute({
					frames: [{ method: "script.evaluate" }],
					launch: { access: true },
				});
				retainedFile = JSON.parse(nextBig.content[0].text).responses[0].result
					.spooled as string;
				await command?.handler("close", ctx);
				expect(existsSync(retainedFile)).toBe(true);
			} finally {
				await t?.session.extensionRunner
					.getCommand("firefox")
					?.handler("close", t.session.extensionRunner.createCommandContext());
				(globalThis as { WebSocket: unknown }).WebSocket = original;
			}
			if (!retainedFile) throw new Error("missing retained spool file");
			await t?.session.extensionRunner.emit({
				type: "session_shutdown",
				reason: "quit",
			});
			expect(existsSync(retainedFile)).toBe(false);
			t?.dispose();
			t = undefined;
		},
	);

	it.runIf(process.platform !== "win32")(
		"surfaces immediate browser exit through a controlled fake firefox on PATH",
		async () => {
			const fakeDir = mkdtempSync(path.join(tmpdir(), "pi-firefox-fake-"));
			const fakeBin = path.join(fakeDir, "fake-firefox");
			await writeFile(fakeBin, "#!/bin/sh\nexit 7\n");
			await chmod(fakeBin, 0o755);

			t = await createTestSession({
				extensions: [extensionPath],
				env: { FIREFOX_BIN: fakeBin },
			});
			const runner = t.session.extensionRunner;
			const command = runner.getCommand("firefox");
			const context = {
				...runner.createCommandContext(),
				mode: "print" as const,
				hasUI: false,
			};
			await command?.handler("", context);

			const tool = runner.extensions[0]?.tools.get("firefox_bidi")?.definition;
			await expect(
				tool?.execute(
					"call-2",
					{ frames: [{ method: "session.status" }] },
					undefined,
					undefined,
					context,
				),
			).rejects.toThrow(/firefox exited immediately \(code 7\)/);
		},
	);
});