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",
] as const;
const WORKER = join(
	dirname(fileURLToPath(import.meta.url)),
	"context-startup-session.ts",
);

type EstimateTokens =
	typeof import("@earendil-works/pi-coding-agent").estimateTokens;
type TokenMessage = Parameters<EstimateTokens>[0];

interface ToolDefinition {
	name: string;
	description: string;
	parameters: unknown;
}

interface SnapshotInput {
	systemPrompt?: string;
	tools?: ToolDefinition[];
	messages?: TokenMessage[];
}

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

interface Measurement {
	kind?: string;
	error?: string;
	provider?: string | undefined;
	model?: string | undefined;
	thinking?: string;
	version?: string;
	cwd?: string;
	contextWindow?: number;
	inputTokens?: number;
	uncachedInputTokens?: number | undefined;
	cacheReadTokens?: number | undefined;
	cacheWriteTokens?: number | undefined;
	outputTokens?: number | undefined;
	local?: LocalSnapshot;
}

interface WorkerOptions {
	kind: string;
	cwd: string;
	sourceAgentDir?: string;
	agentDir?: string;
	selection?: Measurement;
	resultPath: string;
}

interface AssistantUsage {
	input?: number;
	cacheRead?: number;
	cacheWrite?: number;
	output?: number;
}

interface AssistantEventMessage {
	role?: string;
	stopReason?: string;
	provider?: string;
	model?: string;
	usage?: AssistantUsage;
}

interface PiEvent {
	type?: string;
	message?: AssistantEventMessage;
}

export function snapshot(
	{ systemPrompt = "", tools = [], messages = [] }: SnapshotInput,
	estimateTokens: EstimateTokens,
): LocalSnapshot {
	const textTokens = (content: string) =>
		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: string): Measurement {
	const messages: AssistantEventMessage[] = [];
	for (const line of stdout.split("\n")) {
		let event: PiEvent | undefined;
		try {
			event = JSON.parse(line) as PiEvent;
		} 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;
	if (message === undefined)
		throw new Error(`Expected one assistant response, got ${messages.length}`);
	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 as number) >= 0,
		) ||
		(input ?? 0) + (cacheRead ?? 0) + (cacheWrite ?? 0) === 0
	)
		throw new Error("Missing or invalid provider token usage");
	return {
		provider: message.provider,
		model: message.model,
		inputTokens: (input ?? 0) + (cacheRead ?? 0) + (cacheWrite ?? 0),
		uncachedInputTokens: input,
		cacheReadTokens: cacheRead,
		cacheWriteTokens: cacheWrite,
		outputTokens: output,
	};
}

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

export function formatOverview(results: Measurement[]): string {
	const get = (kind: string) => 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: number) => value.toLocaleString("en-US");
	const signed = (value: number) => `${value >= 0 ? "+" : ""}${n(value)}`;
	const widths = [labelWidth, 11, 10, 10, 10];
	const row = (cells: string[]) =>
		cells
			.map((cell, i) =>
				i
					? String(cell).padStart(widths[i] ?? labelWidth)
					: 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 as number)} 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 as LocalSnapshot).system),
							n((r.local as LocalSnapshot).tools),
							n((r.local as LocalSnapshot).messages),
						],
			);
		}),
		"",
	];
	const delta = (label: string, a?: Measurement, b?: Measurement) => {
		const ai = a?.inputTokens;
		const bi = b?.inputTokens;
		if (!comparable(a, b) || ai === undefined || bi === undefined)
			return `${label}: n/a (failed or mismatched runs)`;
		return `${label}: ${signed(ai - bi)} tok (${signed(Math.round(((ai - bi) / bi) * 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 h = (headless as Measurement).local as LocalSnapshot;
		const z = (zero as Measurement).local as LocalSnapshot;
		const changes = (["system", "tools", "messages"] as const).map(
			(key) => `${key} ${signed(h[key] - z[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: WorkerOptions,
	worker: string = WORKER,
): Promise<Measurement> {
	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"),
		) as Measurement;
		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: string = process.cwd(),
	worker: string = WORKER,
): Promise<Measurement[]> {
	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: Measurement[] = [];
	let selection: Measurement | undefined;
	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,
				resultPath,
				...(selection !== undefined ? { selection } : {}),
			};
			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;
	}
}