pi-ext

import { execFile } from "node:child_process";
import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { homedir, tmpdir } from "node:os";
import { dirname, join, resolve } from "node:path";
import { fileURLToPath, pathToFileURL } from "node:url";
import { promisify } from "node:util";

const exec = promisify(execFile);
export const PROMPT = "Reply only OK. Do not use tools.";
export const CASES = [
	"baseline+1msg",
	"headless+0msg",
	"headless+1msg",
	"headless+1msg+no_pi_config",
];
const WORKER = join(
	dirname(fileURLToPath(import.meta.url)),
	"context-startup-session.mjs",
);

export function snapshot(
	{ systemPrompt = "", tools = [], messages = [] },
	estimateTokens,
) {
	const textTokens = (content) =>
		estimateTokens({ role: "user", content, timestamp: 0 });
	const definitions = tools.map(({ name, description, parameters }) => ({
		name,
		description,
		parameters,
	}));
	return {
		system: textTokens(systemPrompt),
		tools: definitions.length ? textTokens(JSON.stringify(definitions)) : 0,
		messages: messages.reduce(
			(total, message) => total + estimateTokens(message),
			0,
		),
		toolCount: definitions.length,
	};
}

export function summarize(stdout) {
	const messages = [];
	for (const line of stdout.split("\n")) {
		let event;
		try {
			event = JSON.parse(line);
		} catch {
			continue;
		}
		if (event?.type === "message_end" && event.message?.role === "assistant")
			messages.push(event.message);
	}
	if (messages.length !== 1)
		throw new Error(`Expected one assistant response, got ${messages.length}`);
	const message = messages[0];
	if (message.stopReason !== "stop")
		throw new Error(
			"Pi did not complete a plain response; measurement rejected",
		);
	const { input, cacheRead, cacheWrite, output } = message.usage ?? {};
	if (
		![input, cacheRead, cacheWrite, output].every(
			(value) => Number.isSafeInteger(value) && value >= 0,
		) ||
		input + cacheRead + cacheWrite === 0
	)
		throw new Error("Missing or invalid provider token usage");
	return {
		provider: message.provider,
		model: message.model,
		inputTokens: input + cacheRead + cacheWrite,
		uncachedInputTokens: input,
		cacheReadTokens: cacheRead,
		cacheWriteTokens: cacheWrite,
		outputTokens: output,
	};
}

export function comparable(a, b) {
	return (
		a &&
		b &&
		!a.error &&
		!b.error &&
		["provider", "model", "thinking", "version", "cwd", "contextWindow"].every(
			(key) => a[key] !== undefined && a[key] === b[key],
		)
	);
}

export function formatOverview(results) {
	const get = (kind) => results.find((row) => row.kind === kind);
	const base = get(CASES[0]);
	const zero = get(CASES[1]);
	const headless = get(CASES[2]);
	const noPiConfig = get(CASES[3]);
	const labelWidth = Math.max(...CASES.map((kind) => kind.length));
	const n = (value) => value.toLocaleString("en-US");
	const signed = (value) => `${value >= 0 ? "+" : ""}${n(value)}`;
	const row = (cells) =>
		cells
			.map((cell, i) =>
				i
					? String(cell).padStart([labelWidth, 11, 10, 10, 10][i])
					: String(cell).padEnd(labelWidth),
			)
			.join(" ");
	const lines = [
		!zero || zero.error
			? "Pi context overview (incomplete)"
			: `Pi ${zero.version} · ${zero.provider}/${zero.model} · ${zero.thinking} · ${n(zero.contextWindow)} ctx`,
		zero?.cwd ?? "",
		"",
		row(["case", "input tok", "~sys tok", "~tools tok", "~msgs tok"]),
		"─".repeat(labelWidth + 45),
		...CASES.map((kind) => {
			const r = get(kind);
			return row(
				!r || r.error
					? [kind, "FAIL", "", "", ""]
					: [
							kind,
							r.inputTokens === undefined ? "n/a" : n(r.inputTokens),
							n(r.local.system),
							n(r.local.tools),
							n(r.local.messages),
						],
			);
		}),
		"",
	];
	const delta = (label, a, b) => {
		if (
			!comparable(a, b) ||
			a.inputTokens === undefined ||
			b.inputTokens === undefined
		)
			return `${label}: n/a (failed or mismatched runs)`;
		return `${label}: ${signed(a.inputTokens - b.inputTokens)} tok (${signed(Math.round(((a.inputTokens - b.inputTokens) / b.inputTokens) * 100))}%)`;
	};
	lines.push(delta("Config", headless, base));
	lines.push(delta("Pi text resources", headless, noPiConfig));
	lines.push(
		"no_pi_config keeps extension settings and project context; removes agent-dir instructions/skills/prompts.",
	);
	if (comparable(headless, zero)) {
		const changes = ["system", "tools", "messages"].map(
			(key) => `${key} ${signed(headless.local[key] - zero.local[key])}`,
		);
		lines.push(`First turn (~tok): ${changes.join(" · ")}`);
	}
	lines.push(
		"Tokens: input = provider incl. cache; ~ = Pi heuristic (chars/4).",
	);
	lines.push("Estimates need not sum to provider input.");
	lines.push("SDK startup · offline catalogs · ephemeral · read-only auth");
	for (const r of results) if (r.error) lines.push(`${r.kind}: ${r.error}`);
	return lines.join("\n");
}

export async function runWorker(options, worker = WORKER) {
	const configPath = `${options.resultPath}.options.json`;
	await writeFile(configPath, JSON.stringify(options), { mode: 0o600 });
	try {
		try {
			await exec(process.execPath, [worker, configPath], {
				cwd: options.cwd,
				timeout: 120_000,
				maxBuffer: 8 * 1024 * 1024,
			});
		} catch {
			/* Read the worker's sanitized failure, never echo subprocess logs. */
		}
		const result = JSON.parse(await readFile(options.resultPath, "utf8"));
		if (
			options.selection &&
			!result.error &&
			!comparable(result, options.selection)
		)
			return {
				kind: options.kind,
				error: "Model/thinking/version/cwd mismatch",
			};
		return result;
	} catch {
		return {
			kind: options.kind,
			error: "Process or measurement collection failed",
		};
	}
}

export async function measure(cwd = process.cwd(), worker = WORKER) {
	cwd = resolve(cwd);
	const sourceAgentDir =
		process.env.PI_CODING_AGENT_DIR ?? join(homedir(), ".pi", "agent");
	const directory = await mkdtemp(join(tmpdir(), "pi-context-overview-"));
	await mkdir(join(directory, "baseline"));
	const results = [];
	let selection;
	try {
		for (const kind of [CASES[1], CASES[0], CASES[2], CASES[3]]) {
			const resultPath = join(directory, `${kind}.result.json`);
			const options = {
				kind,
				cwd,
				sourceAgentDir,
				agentDir:
					kind === CASES[0] ? join(directory, "baseline") : sourceAgentDir,
				selection,
				resultPath,
			};
			const result = await runWorker(options, worker);
			results.push(result);
			if (kind === CASES[1]) {
				if (result.error) {
					for (const blocked of [CASES[0], CASES[2], CASES[3]])
						results.push({
							kind: blocked,
							error: "Blocked: startup model/config unresolved",
						});
					break;
				}
				selection = result;
			}
		}
		return results;
	} finally {
		await rm(directory, { recursive: true, force: true });
	}
}

if (
	process.argv[1] &&
	import.meta.url === pathToFileURL(resolve(process.argv[1])).href
) {
	try {
		const results = await measure(process.argv[2]);
		console.log(formatOverview(results));
		if (results.some((row) => row.error)) process.exitCode = 1;
	} catch {
		console.error(
			"Context overview failed: check installed Pi and measurement directory access.",
		);
		process.exitCode = 1;
	}
}