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.ts";

type WorkerConfig = Parameters<typeof runWorker>[0];

interface ExtensionEntry {
	name: string;
	path: string;
	tools: unknown[];
	hooks: string[];
	providers: unknown[];
}

interface LocalSnapshot {
	system: number;
	tools: number;
	messages: number;
	toolCount?: number;
}

interface MeasurementLike {
	kind?: string;
	error?: string;
	provider?: string | undefined;
	model?: string | undefined;
	thinking?: string;
	version?: string;
	cwd?: string;
	inputTokens?: number;
	local?: LocalSnapshot;
	extensions?: ExtensionEntry[];
}

interface Run {
	extension: ExtensionEntry;
	result: MeasurementLike;
}

interface Report {
	error?: string;
	inventory: ExtensionEntry[];
	baseline?: MeasurementLike;
	after?: MeasurementLike;
	runs: Run[];
	stable?: boolean | undefined;
}

// Signals prioritize runs, never certify that an extension is context-free.
export function registrationSignals(extension: ExtensionEntry): string[] {
	const signals: string[] = [];
	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"],
	] as [string, string][])
		if (extension.hooks.includes(event)) signals.push(label);
	return signals.length
		? signals
		: [extension.hooks.length ? "other-hooks" : "unknown"];
}

export async function measureExtensions(
	cwd: string = process.cwd(),
	worker?: string,
	progress: (message: string) => void = () => {},
): Promise<Report> {
	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: string,
		selection?: MeasurementLike,
		extensionPaths?: string[],
	): Promise<MeasurementLike> => {
		const options: WorkerConfig & {
			extensionPaths?: string[];
			inspectExtensions?: boolean;
		} = {
			kind,
			cwd,
			agentDir,
			sourceAgentDir: agentDir,
			inspectExtensions: true,
			resultPath: join(directory, `${index++}.result.json`),
			...(selection !== undefined ? { selection } : {}),
			...(extensionPaths !== undefined ? { extensionPaths } : {}),
		};
		return runWorker(options, 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: ExtensionEntry) => {
				const [first] = registrationSignals(entry);
				return (
					first !== undefined && ["other-hooks", "unknown"].includes(first)
				);
			};
			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: Report): string {
	const { baseline, after, runs } = report;
	if (!baseline || baseline.error)
		return `Extension measurement failed: ${report.error}`;
	const n = (value: number) => value.toLocaleString("en-US");
	const signed = (value: number) => `${value >= 0 ? "+" : ""}${n(value)}`;
	const valid = (result: MeasurementLike) =>
		report.stable &&
		comparable(baseline, result) &&
		Number.isSafeInteger(result.inputTokens);
	const label = (extension: ExtensionEntry) =>
		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: string[]) =>
		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 ?? Number.NaN)} input tok; repeat: ${after?.error ? "FAIL" : n(after?.inputTokens ?? Number.NaN)}`,
		"",
		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 ?? 0) - (b.result.inputTokens ?? 0)
			: 0;
	});
	for (const { extension, result } of ordered) {
		lines.push(
			row(
				valid(result)
					? [
							label(extension),
							signed((baseline.inputTokens ?? 0) - (result.inputTokens ?? 0)),
							...(["system", "tools", "messages"] as const).map((key) =>
								signed(
									(baseline.local?.[key] ?? 0) - (result.local?.[key] ?? 0),
								),
							),
						]
					: [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;
	}
}