pi-ext

import { mkdtemp, rm } from "node:fs/promises";
import { homedir, tmpdir } from "node:os";
import { join, resolve } from "node:path";
import { pathToFileURL } from "node:url";
import { comparable, runWorker } from "./context-startup.mjs";

// Signals prioritize runs, never certify that an extension is context-free.
export function registrationSignals(extension) {
	const signals = [];
	if (extension.tools.length) signals.push("tools");
	if (extension.providers.length) signals.push("provider");
	for (const [event, label] of [
		["before_agent_start", "prompt"],
		["context", "context"],
		["before_provider_request", "payload"],
		["resources_discover", "resources"],
		["input", "input"],
		["session_start", "startup"],
	])
		if (extension.hooks.includes(event)) signals.push(label);
	return signals.length
		? signals
		: [extension.hooks.length ? "other-hooks" : "unknown"];
}

export async function measureExtensions(
	cwd = process.cwd(),
	worker,
	progress = () => {},
) {
	cwd = resolve(cwd);
	const agentDir =
		process.env.PI_CODING_AGENT_DIR ?? join(homedir(), ".pi", "agent");
	const directory = await mkdtemp(join(tmpdir(), "pi-context-extensions-"));
	let index = 0;
	const run = (kind, selection, extensionPaths) =>
		runWorker(
			{
				kind,
				cwd,
				agentDir,
				sourceAgentDir: agentDir,
				selection,
				extensionPaths,
				inspectExtensions: true,
				resultPath: join(directory, `${index++}.result.json`),
			},
			worker,
		);
	try {
		progress("Discovering configured extensions and startup selection");
		const startup = await run("headless+0msg");
		if (startup.error) return { error: startup.error, inventory: [], runs: [] };
		const inventory = startup.extensions;
		if (
			!Array.isArray(inventory) ||
			inventory.some(
				(entry) =>
					typeof entry.path !== "string" ||
					!entry.path ||
					!Array.isArray(entry.tools) ||
					!Array.isArray(entry.hooks) ||
					!Array.isArray(entry.providers),
			) ||
			new Set(inventory.map((entry) => entry.path)).size !== inventory.length
		)
			return {
				error: "Missing or invalid runtime extension inventory",
				inventory: [],
				runs: [],
			};
		const paths = inventory.map((entry) => entry.path);
		progress(
			`${inventory.length} file extensions discovered; measuring all, including uncertain registrations`,
		);
		progress(`Full-set control (1/${inventory.length + 2})`);
		const baseline = await run("headless+1msg", startup, paths);
		if (baseline.error)
			return { error: baseline.error, inventory, baseline, runs: [] };
		const ordered = [...inventory].sort((a, b) => {
			const uncertain = (entry) =>
				["other-hooks", "unknown"].includes(registrationSignals(entry)[0]);
			return Number(uncertain(a)) - Number(uncertain(b));
		});
		const runs = [];
		for (const extension of ordered) {
			progress(
				`Without ${extension.name} (${runs.length + 2}/${inventory.length + 2}) [${registrationSignals(extension).join(", ")}]`,
			);
			const result = await run(
				"headless+1msg",
				startup,
				paths.filter((path) => path !== extension.path),
			);
			runs.push({ extension, result });
		}
		progress(
			`Repeated full-set control (${inventory.length + 2}/${inventory.length + 2})`,
		);
		const after = await run("headless+1msg", startup, paths);
		const stable =
			comparable(baseline, after) &&
			baseline.inputTokens === after.inputTokens &&
			JSON.stringify(baseline.local) === JSON.stringify(after.local);
		return {
			inventory,
			baseline,
			after,
			runs,
			stable,
			...(!stable
				? { error: "Full-set control failed or changed; savings withheld" }
				: {}),
		};
	} finally {
		await rm(directory, { recursive: true, force: true });
	}
}

export function formatExtensionOverview(report) {
	const { baseline, after, runs } = report;
	if (!baseline || baseline.error)
		return `Extension measurement failed: ${report.error}`;
	const n = (value) => value.toLocaleString("en-US");
	const signed = (value) => `${value >= 0 ? "+" : ""}${n(value)}`;
	const valid = (result) =>
		report.stable &&
		comparable(baseline, result) &&
		Number.isSafeInteger(result.inputTokens);
	const label = (extension) =>
		report.inventory.filter((entry) => entry.name === extension.name).length > 1
			? extension.path
			: extension.name;
	const width = Math.max(
		17,
		...report.inventory.map((entry) => label(entry).length),
	);
	const row = (cells) =>
		cells
			.map((cell, i) =>
				i ? String(cell).padStart(11) : String(cell).padEnd(width),
			)
			.join(" ");
	const lines = [
		`Pi ${baseline.version} · ${baseline.provider}/${baseline.model} · ${baseline.thinking}`,
		baseline.cwd,
		`Full set: ${n(baseline.inputTokens)} input tok; repeat: ${after?.error ? "FAIL" : n(after.inputTokens)}`,
		"",
		row([
			"extension removed",
			"saved tok",
			"~sys saved",
			"~tool saved",
			"~msg saved",
		]),
		"─".repeat(width + 48),
	];
	const ordered = [...runs].sort((a, b) => {
		if (valid(a.result) !== valid(b.result)) return valid(a.result) ? -1 : 1;
		return valid(a.result) ? a.result.inputTokens - b.result.inputTokens : 0;
	});
	for (const { extension, result } of ordered) {
		lines.push(
			row(
				valid(result)
					? [
							label(extension),
							signed(baseline.inputTokens - result.inputTokens),
							...["system", "tools", "messages"].map((key) =>
								signed(baseline.local[key] - result.local[key]),
							),
						]
					: [label(extension), result.error ? "FAIL" : "n/a", "", "", ""],
			),
		);
	}
	lines.push(
		"",
		"Positive = tokens saved by removal; negative = removal adds tokens.",
		"Input is provider-reported including cache; ~ = Pi estimates. Deltas are not additive.",
		"Same config/resources; only one extension module omitted before loading. Built-ins retained.",
		"All discovered file extensions measured, including unknown/future registrations. No name allowlist.",
	);
	if (report.error) lines.push(report.error);
	for (const { extension, result } of runs)
		if (result.error) lines.push(`${label(extension)}: ${result.error}`);
	return lines.join("\n");
}

if (
	process.argv[1] &&
	import.meta.url === pathToFileURL(resolve(process.argv[1])).href
) {
	try {
		const report = await measureExtensions(
			process.argv[2],
			undefined,
			(message) => console.error(message),
		);
		console.log(formatExtensionOverview(report));
		if (report.error || report.runs.some((run) => run.result.error))
			process.exitCode = 1;
	} catch {
		console.error(
			"Extension measurement failed: check installed Pi and measurement directory access.",
		);
		process.exitCode = 1;
	}
}