pi-ext

/**
 * README demo runner: drives the real Pi TUI in a pty through a scene file,
 * writes an asciicast v2 with post-production holds, and renders a GIF with
 * agg. With `--assert-only` it runs the same scene without writing anything,
 * so a scene doubles as a terminal E2E assertion in //:check.
 *
 * Scene file: extensions/<name>/demo/scene.json
 *   {
 *     "title": "footer",
 *     "extensions": ["footer", "cost", "context-size"],   // dist packages to load; first = protagonist
 *     "cols": 100, "rows": 28,
 *     "turns": [{ "text": "…" }, { "text": "…", "tool": { "name": "read", "arguments": {…} } }],
 *     "steps": [
 *       { "expect": "regex", "hold": 1.5 },
 *       { "type": "/footer", "key": "enter", "expect": "regex", "hold": 2 },
 *       { "key": "ctrl-c" }
 *     ]
 *   }
 * Time between steps in the cast is `hold` seconds (default 6 after an
 * `expect`); real waiting is compressed to at most 50 ms per output chunk.
 */
import {
	existsSync,
	mkdirSync,
	readdirSync,
	readFileSync,
	rmSync,
	statSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join, relative } from "node:path";
import { fileURLToPath, pathToFileURL } from "node:url";
import { ROOT, run } from "../lib/common.ts";

export interface SceneStep {
	type?: string;
	/** Named key, or `alt-<char>` for an Alt chord. */
	key?: string;
	expect?: string;
	/** Seconds held on screen after this step in the playback timeline. */
	hold?: number;
	/** Milliseconds to wait for `expect` before failing; default 15000. */
	timeout?: number;
	/** Match `expect` only against output produced after this step's input,
	 * for example once an interactive program has handed the terminal back. */
	fresh?: boolean;
}

export interface Scene {
	title: string;
	extensions: string[];
	cols?: number;
	rows?: number;
	turns?: unknown[];
	steps: SceneStep[];
	/** Font list for agg; default JetBrains Mono with a Nerd Font fallback. */
	font?: string;
	/** Files seeded into the demo project directory before Pi starts. */
	files?: Record<string, string>;
	/** Commands (argv arrays) run in the project directory after `files`, for
	 * fixtures a file list cannot express, such as a Git repository. */
	setup?: string[][];
	/** Regexes (visible text) that are expected to appear although they look like errors. */
	allowErrors?: string[];
	/** Extra user settings for the sandbox agent dir, for example `ultra.subagentExtensions`. */
	settings?: Record<string, unknown>;
	/** Extra `pi` CLI arguments, for example an extension flag such as `--nushell`. */
	args?: string[];
}

/** Screen text that means the demo went wrong, whatever the scene expected. */
const ERROR_PATTERNS: readonly RegExp[] = [
	/No more faux responses queued/u,
	/\/bug sends a report to the developers/u,
	/\bError: /u,
	/Extension error/u,
	/Failed to load extension/u,
];

const KEYS: Record<NonNullable<SceneStep["key"]>, string> = {
	enter: "\r",
	escape: "\u001b",
	tab: "\t",
	"ctrl-c": "\u0003",
	"ctrl-d": "\u0004",
	"ctrl-j": "\n",
	"ctrl-o": "\u000f",
	up: "\u001b[A",
	down: "\u001b[B",
};

/** `alt-<char>` sends ESC + char, which is how terminals encode Alt chords. */
function keyBytes(key: string): string {
	const alt = /^alt-(.)$/u.exec(key);
	if (alt) return `\u001b${alt[1]}`;
	const bytes = KEYS[key as keyof typeof KEYS];
	if (!bytes) throw new Error(`unknown key "${key}"`);
	return bytes;
}

const DEFAULT_HOLD = 6;
const TYPE_DELAY_S = 0.045;

/** Strip ANSI so `expect` regexes match visible text. */
function visible(text: string): string {
	return text
		.replaceAll(/\u001b\[[0-9;?]*[A-Za-z]/g, "")
		.replaceAll(/\u001b\][^\u0007]*\u0007/g, "")
		.replaceAll(/\u001b[()][A-Z0-9]/g, "")
		.replaceAll(/\u001b[=>]/g, "");
}

function nuguTheme(): string {
	const palette = JSON.parse(
		readFileSync(join(ROOT, "assets", "nugu.json"), "utf8"),
	) as {
		dark: { content: Record<string, string>; ui: Record<string, string> };
	};
	const c = palette.dark.content;
	const u = palette.dark.ui;
	const hex = (v: string) => v.replace("#", "");
	// agg theme: bg, fg, then 16 ANSI colours. Mirrors nugu's ghostty mapping.
	return [
		hex(c["backdrop"] ?? "#121212"),
		hex(c["normal"] ?? "#ededed"),
		hex(c["backdrop"] ?? "#121212"),
		"dc5050",
		hex(c["focus"] ?? "#7ba1aa"),
		hex(c["important_local"] ?? "#b59ae4"),
		hex(c["focus"] ?? "#7ba1aa"),
		hex(c["accent"] ?? "#d151d1"),
		hex(c["important_global"] ?? "#dfd3f3"),
		hex(c["normal"] ?? "#ededed"),
		hex(c["unfocus"] ?? "#5d5d5d"),
		"dc5050",
		hex(u["focus"] ?? "#9aabae"),
		hex(u["important_local"] ?? "#c9b6eb"),
		hex(u["focus"] ?? "#9aabae"),
		hex(u["accent"] ?? "#d86dd8"),
		hex(u["important_global"] ?? "#f3effb"),
		hex(u["normal"] ?? "#dbdbdb"),
	].join(",");
}

export interface RunResult {
	castPath?: string;
	gifPath?: string;
	realSeconds: number;
	playbackSeconds: number;
}

export async function runScene(
	extensionName: string,
	options: { assertOnly: boolean },
): Promise<RunResult> {
	const extDir = join(ROOT, "extensions", extensionName);
	const scenePath = join(extDir, "demo", "scene.json");
	if (!existsSync(scenePath))
		throw new Error(`${extensionName}: no demo/scene.json`);
	const scene = JSON.parse(readFileSync(scenePath, "utf8")) as Scene;
	const cols = scene.cols ?? 100;
	const rows = scene.rows ?? 28;

	// Loaded packages come from dist so the demo shows what ships.
	const extensionArgs = [
		...scene.extensions.flatMap((name) => {
			const packageDir = join(ROOT, "dist", name);
			const packagePath = join(packageDir, "package.json");
			if (!existsSync(packagePath))
				throw new Error(
					`${extensionName}: dist/${name} missing; run mise run //:build-ext ${name}`,
				);
			const pkg = JSON.parse(readFileSync(packagePath, "utf8")) as {
				pi?: { extensions?: string[] };
			};
			const entryNames = pkg.pi?.extensions;
			if (!entryNames?.length)
				throw new Error(
					`${extensionName}: dist/${name} has no Pi extension entry`,
				);
			return entryNames.flatMap((entryName) => {
				const entry = join(packageDir, entryName);
				if (!existsSync(entry))
					throw new Error(
						`${extensionName}: dist/${name}/${entryName} missing`,
					);
				return ["-e", entry];
			});
		}),
	];

	// Isolated agent dir: no user settings, sessions, or auth leak into the demo.
	// A stable path so the footer shows `…/pi-demo/<name>` instead of a pid.
	const sandbox = join(tmpdir(), "pi-demo-sandbox", extensionName);
	const agentDir = join(sandbox, "agent");
	// The demo runs in a neutral project directory so the screen shows the
	// extension, not this repository; the folder is pre-trusted to skip Pi's
	// trust prompt.
	const project = join(tmpdir(), "pi-demo", extensionName);
	rmSync(sandbox, { recursive: true, force: true });
	rmSync(project, { recursive: true, force: true });
	mkdirSync(agentDir, { recursive: true });
	mkdirSync(project, { recursive: true });
	for (const [file, content] of Object.entries(scene.files ?? {})) {
		const target = join(project, file);
		mkdirSync(join(target, ".."), { recursive: true });
		writeFileSync(target, content);
	}
	for (const [command, ...args] of scene.setup ?? []) {
		if (!command) continue;
		run(command, args, { cwd: project, check: true, capture: true });
	}
	writeFileSync(
		join(agentDir, "settings.json"),
		`${JSON.stringify({ theme: "dark", ...scene.settings }, null, "\t")}\n`,
	);
	// The demo provider lives in the sandbox agent dir so Pi discovers it the
	// same way for the recorded session and for any child session an
	// extension such as ultra spawns with the same PI_CODING_AGENT_DIR.
	const providerDir = join(agentDir, "extensions", "demo-provider");
	mkdirSync(providerDir, { recursive: true });
	writeFileSync(
		join(providerDir, "index.ts"),
		`export { default } from ${JSON.stringify(fileURLToPath(new URL("./provider.ts", import.meta.url)))};\n`,
	);
	writeFileSync(
		join(agentDir, "trust.json"),
		`${JSON.stringify({ [project]: true }, null, "\t")}\n`,
	);

	const { spawn } = await import("@lydell/node-pty");
	const started = Date.now();
	const events: Array<[number, "o", string]> = [];
	let timeline = 0;
	let lastReal = started;
	let screen = "";
	/** Everything Pi printed, ANSI stripped, for the error scan. */
	let transcript = "";
	const reportPath = join(sandbox, "provider-report.json");
	const allowed = (scene.allowErrors ?? []).map((p) => new RegExp(p, "u"));

	// Launch the installed CLI directly; Windows cannot execute a pi.cmd shim in a PTY.
	const piPackage = new URL(
		"../package.json",
		import.meta.resolve("@earendil-works/pi-coding-agent"),
	);
	const piMetadata = JSON.parse(readFileSync(piPackage, "utf8")) as {
		bin: { pi: string };
	};
	const proc = spawn(
		process.execPath,
		[
			fileURLToPath(new URL(piMetadata.bin.pi, piPackage)),
			"--no-session",
			"--no-skills",
			"--no-prompt-templates",
			"--no-themes",
			"--provider",
			"demo",
			"--model",
			"demo-model",
			...extensionArgs,
			...(scene.args ?? []),
		],
		{
			name: "xterm-256color",
			cols,
			rows,
			cwd: project,
			env: {
				...process.env,
				TERM: "xterm-256color",
				COLORTERM: "truecolor",
				PI_CODING_AGENT_DIR: agentDir,
				PI_DEMO_SCENE: scenePath,
				PI_DEMO_REPORT: reportPath,
				PI_DEMO_TRACE: join(sandbox, "trace.txt"),
				// No version check, package update check, or telemetry in a recording.
				PI_OFFLINE: "1",
				// Deterministic external editor for ctx.ui.editor and Ctrl+G.
				VISUAL: "nvim",
				EDITOR: "nvim",
				HOME: join(sandbox, "home"),
				XDG_STATE_HOME: join(sandbox, "state"),
				XDG_CONFIG_HOME: join(sandbox, "config"),
				XDG_CACHE_HOME: join(sandbox, "cache"),
				XDG_DATA_HOME: join(sandbox, "data"),
			},
		},
	);
	proc.onData((data) => {
		const now = Date.now();
		timeline += Math.min((now - lastReal) / 1000, 0.05);
		lastReal = now;
		events.push([Number(timeline.toFixed(3)), "o", data]);
		screen = (screen + data).slice(-40_000);
		transcript += visible(data);
	});
	let exited = false;
	proc.onExit(() => {
		exited = true;
	});

	const waitFor = (pattern: string, ms: number) =>
		new Promise<void>((resolve, reject) => {
			const re = new RegExp(pattern, "u");
			const t0 = Date.now();
			const tick = () => {
				if (re.test(visible(screen))) return resolve();
				if (exited)
					return reject(new Error(`pi exited while waiting for /${pattern}/`));
				if (Date.now() - t0 > ms)
					return reject(
						new Error(
							`timeout waiting for /${pattern}/\n--- last screen ---\n${visible(screen).slice(-1500)}`,
						),
					);
				setTimeout(tick, 25);
			};
			tick();
		});
	const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

	try {
		// Pi renders its model before binding extension shortcuts; the loaded
		// resource list appears only after binding finishes.
		await waitFor("\\[Extensions\\]", 15_000);
		for (const step of scene.steps) {
			if (step.fresh) screen = "";
			if (step.type)
				for (const ch of step.type) {
					proc.write(ch);
					timeline += TYPE_DELAY_S;
					await sleep(12);
				}
			if (step.key) proc.write(keyBytes(step.key));
			if (step.expect) await waitFor(step.expect, step.timeout ?? 15_000);
			const hold = step.hold ?? (step.expect ? DEFAULT_HOLD : 0);
			timeline += hold;
			// Keep only recent output for matching; the cast keeps everything.
			screen = screen.slice(-12_000);
		}
		await sleep(200);
	} catch (error) {
		// Keep the whole visible transcript for diagnosis; the sandbox is temp.
		writeFileSync(join(sandbox, "transcript.txt"), transcript);
		throw new Error(
			`${error instanceof Error ? error.message : String(error)}\n--- full transcript: ${join(sandbox, "transcript.txt")}`,
		);
	} finally {
		if (!exited) proc.kill();
	}
	const realSeconds = (Date.now() - started) / 1000;

	// A scene passes only when nothing on screen looked like an error and the
	// model was called exactly as often as turns were scripted; otherwise an
	// unplanned turn silently eats a reply and the recording lies.
	for (const pattern of ERROR_PATTERNS) {
		const match = pattern.exec(transcript);
		if (!match) continue;
		if (allowed.some((a) => a.test(match[0]))) continue;
		const at = Math.max(0, match.index - 200);
		throw new Error(
			`screen showed an error: ${match[0]}\n--- context ---\n${transcript.slice(at, match.index + 300)}`,
		);
	}
	if (existsSync(reportPath)) {
		const report = JSON.parse(readFileSync(reportPath, "utf8")) as {
			queued: number;
			calls: number;
		};
		// Child sessions replay their own list; only the primary queue is
		// accounted here.
		if (report.calls !== report.queued)
			throw new Error(
				`scripted ${report.queued} model turn(s) but Pi made ${report.calls} model call(s); align scene.turns with the prompts the steps actually send`,
			);
	}
	if (options.assertOnly) return { realSeconds, playbackSeconds: timeline };

	const header = JSON.stringify({
		version: 2,
		width: cols,
		height: rows,
		timestamp: Math.floor(started / 1000),
		title: scene.title,
		env: { TERM: "xterm-256color" },
	});
	const cast = `${[header, ...events.map((e) => JSON.stringify(e))].join("\n")}\n`;
	const assets = join(extDir, "assets");
	mkdirSync(assets, { recursive: true });
	const castPath = join(assets, "demo.cast");
	const gifPath = join(assets, "demo.gif");
	writeFileSync(castPath, cast);
	run(
		"agg",
		[
			"--theme",
			nuguTheme(),
			"--font-family",
			scene.font ?? "JetBrains Mono,MesloLGM Nerd Font Mono",
			"--font-size",
			"16",
			"--idle-time-limit",
			// Preserve authored reading pauses; agg otherwise caps them at 5 s.
			"60",
			castPath,
			gifPath,
		],
		{ cwd: ROOT, check: true },
	);
	return { castPath, gifPath, realSeconds, playbackSeconds: timeline };
}

function main(argv: string[]): Promise<number> {
	const assertOnly = argv.includes("--assert-only");
	const names = argv.filter((a) => !a.startsWith("--"));
	const withScene = readdirScenes();
	// Release gates pass their whole selection; extensions without a scene are
	// not a failure there, only when someone asks to record them.
	const targets = names.length
		? assertOnly
			? names.filter((name) => withScene.includes(name))
			: names
		: withScene;
	return (async () => {
		let failed = 0;
		for (const name of targets) {
			try {
				const result = await runScene(name, { assertOnly });
				const where = result.gifPath
					? ` → ${relative(ROOT, result.gifPath)} (${(statSize(result.gifPath) / 1024).toFixed(0)} KB)`
					: "";
				console.log(
					`OK ${name}: ${result.realSeconds.toFixed(1)}s real, ${result.playbackSeconds.toFixed(1)}s playback${where}`,
				);
			} catch (error) {
				failed++;
				console.error(
					`FAIL ${name}: ${error instanceof Error ? error.message : String(error)}`,
				);
			}
		}
		return failed ? 1 : 0;
	})();
}

export function readdirScenes(): string[] {
	const root = join(ROOT, "extensions");
	return readdirSync(root)
		.filter((name) => existsSync(join(root, name, "demo", "scene.json")))
		.sort();
}

const statSize = (path: string): number => statSync(path).size;

if (import.meta.url === pathToFileURL(process.argv[1] ?? "").href)
	main(process.argv.slice(2)).then((code) => {
		process.exitCode = code;
	});