pi-ext

// Verbatim copy of the pre-optimization hot path (index.ts runBatch loop +
// execute payload build), used as the benchmark baseline and diff-checked
// against the optimized implementation. Do not "improve" this file.

import { mkdtemp, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { interpolate } from "../lib.ts";

const MAX_INLINE = 256 * 1024;

interface BidiResponse {
	id: number;
	method?: string;
	type?: string;
	result?: unknown;
	error?: string;
	message?: string;
	skipped?: boolean;
}

interface BatchCtx {
	pending: Map<number, (r: BidiResponse) => void>;
	ledger: {
		subscriptions: string[];
		intercepts: string[];
		preloadScripts: string[];
	};
	spoolDir: string;
}

function ledgerUpdate(
	b: BatchCtx,
	method: string,
	params: Record<string, unknown>,
	res: BidiResponse,
): void {
	const ok = res.type === "success";
	const result = (res.result ?? {}) as Record<string, unknown>;
	if (method === "session.subscribe" && ok) {
		for (const e of (params.events as string[]) ?? [])
			if (!b.ledger.subscriptions.includes(e)) b.ledger.subscriptions.push(e);
	} else if (method === "session.unsubscribe") {
		const gone = new Set((params.events as string[]) ?? []);
		b.ledger.subscriptions = b.ledger.subscriptions.filter((x) => !gone.has(x));
	} else if (
		method === "network.addIntercept" &&
		ok &&
		typeof result.intercept === "string"
	) {
		b.ledger.intercepts.push(result.intercept);
	} else if (method === "network.removeIntercept") {
		const ids = new Set(
			((params.intercepts as string[]) ?? []).concat(
				(params.intercept as string[]) ?? [],
			),
		);
		b.ledger.intercepts = b.ledger.intercepts.filter((x) => !ids.has(x));
	} else if (
		method === "script.addPreloadScript" &&
		ok &&
		typeof result.id === "string"
	) {
		b.ledger.preloadScripts.push(result.id);
	} else if (
		method === "script.removePreloadScript" &&
		typeof params.id === "string"
	) {
		b.ledger.preloadScripts = b.ledger.preloadScripts.filter(
			(x) => x !== params.id,
		);
	}
}

export async function runBatchBaseline(
	ws: { send: (s: string) => void },
	b: BatchCtx,
	frames: { id: number; method: string; params?: Record<string, unknown> }[],
	byId: Map<number, BidiResponse>,
	timeoutMs: number,
): Promise<BidiResponse[]> {
	const responses: BidiResponse[] = [];
	let aborted = false;
	for (const frame of frames) {
		if (aborted) {
			responses.push({ id: frame.id, method: frame.method, skipped: true });
			continue;
		}
		let wire: unknown;
		try {
			wire = {
				id: frame.id,
				method: frame.method,
				params: interpolate(frame.params ?? {}, byId),
			};
		} catch (e) {
			responses.push({
				id: frame.id,
				method: frame.method,
				type: "error",
				error: "pi interpolation",
				message: e instanceof Error ? e.message : String(e),
			});
			aborted = true;
			continue;
		}
		const res = await new Promise<BidiResponse>((resolve) => {
			const t = setTimeout(
				() =>
					resolve({
						id: frame.id,
						type: "timeout",
						message: `no response within ${timeoutMs}ms; side effects may have applied`,
					}),
				timeoutMs,
			);
			b.pending.set(frame.id, (r) => {
				clearTimeout(t);
				resolve(r);
			});
			ws.send(JSON.stringify(wire));
		});
		ledgerUpdate(
			b,
			frame.method,
			(frame.params ?? {}) as Record<string, unknown>,
			res,
		);
		byId.set(frame.id, res);
		const out: BidiResponse = { ...res, id: frame.id, method: frame.method };
		const size = JSON.stringify(out).length;
		if (size > MAX_INLINE && res.type === "success") {
			const file = join(
				b.spoolDir,
				`${frame.method.replace(/[^\w.-]/g, "_")}-${frame.id}.json`,
			);
			await writeFile(file, JSON.stringify(out.result ?? null));
			out.result = {
				spooled: file,
				size,
				note: "full result written to file; read selectively",
			} as unknown;
		}
		responses.push(out);
	}
	return responses;
}

export function payloadBaseline(
	session: unknown,
	responses: BidiResponse[],
	events: unknown[],
	wait: unknown,
	eventCap: number,
): string {
	const payload = {
		session,
		responses,
		events,
		...(events.length === eventCap
			? { note: `event buffer capped at ${eventCap}` }
			: {}),
		...(wait ? { wait } : {}),
	};
	return JSON.stringify(payload);
}

// ── fixture ───────────────────────────────────────────────────────────────

export interface Fixture {
	ws: { send: (s: string) => void };
	ctx: BatchCtx;
	frames: { id: number; method: string; params?: Record<string, unknown> }[];
	byId: Map<number, BidiResponse>;
}

// Batch shape a real agent call produces: many small frames, a few huge
// (screenshot/print-like) results, some params with {{N.result.*}} refs.
export async function makeFixture(
	_smallKb = 2,
	bigKb = 512,
	small = 22,
	big = 3,
): Promise<Fixture> {
	const pending = new Map<number, (r: BidiResponse) => void>();
	const frames: Fixture["frames"] = [];
	const resultsById = new Map<number, unknown>();
	let id = 1;
	frames.push({
		id: id,
		method: "session.subscribe",
		params: { events: ["network"] },
	});
	resultsById.set(id, {});
	id++;
	frames.push({ id: id, method: "browsingContext.create", params: {} });
	const ctxId = `ctx-${id}-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee`;
	resultsById.set(id, { context: ctxId, userContext: "default" });
	id++;
	frames.push({
		id: id,
		method: "browsingContext.navigate",
		params: {
			context: `{{2.result.context}}`,
			url: "https://example.com/",
			wait: "complete",
		},
	});
	resultsById.set(id, { context: ctxId, url: "https://example.com/" });
	id++;
	for (let i = 0; i < small; i++) {
		frames.push({
			id: id,
			method: "script.evaluate",
			params: {
				expression: `1+${i}`,
				awaitPromise: false,
				target: { context: "{{2.result.context}}" },
			},
		});
		resultsById.set(id, { type: "number", value: i });
		id++;
	}
	for (let i = 0; i < big; i++) {
		frames.push({
			id: id,
			method: "browsingContext.captureScreenshot",
			params: { context: "{{2.result.context}}" },
		});
		resultsById.set(id, {
			data: "iVBORw0KGgoAAAANSUhEUg".padEnd(bigKb * 1024, "A"),
		});
		id++;
	}
	const ws = {
		send: (s: string) => {
			const f = JSON.parse(s);
			queueMicrotask(() =>
				pending.get(f.id)?.({
					id: f.id,
					type: "success",
					result: resultsById.get(f.id),
				}),
			);
		},
	};
	const spoolDir = await mkdtemp(join(tmpdir(), "pi-firefox-bench-"));
	return {
		ws,
		ctx: {
			pending,
			ledger: { subscriptions: [], intercepts: [], preloadScripts: [] },
			spoolDir,
		},
		frames,
		byId: new Map(),
	};
}

export async function makeSpoolCtx(): Promise<BatchCtx> {
	const dir = await mkdtemp(join(tmpdir(), "pi-firefox-bench-"));
	return {
		pending: new Map(),
		ledger: { subscriptions: [], intercepts: [], preloadScripts: [] },
		spoolDir: dir,
	};
}