pi-ext

/**
 * Demo provider: a Pi extension loaded only while recording README demos.
 * Registers the `demo` provider with scripted assistant turns read from the
 * scene file named by `PI_DEMO_SCENE`, so recordings are deterministic and
 * cost nothing. The scene's `turns` array supplies one reply per user turn;
 * when it runs out, the last reply repeats.
 */
import { readFileSync, writeFileSync } from "node:fs";
import {
	createFauxCore,
	type FauxResponseStep,
	fauxAssistantMessage,
	fauxToolCall,
	type ToolCall,
} from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

export const DEMO_PROVIDER = "demo";
export const DEMO_MODEL = "demo-model";
/** A second, pricier model so model-switching demos have somewhere to go. */
export const DEMO_DEEP_MODEL = "demo-deep";

interface SceneTurn {
	text?: string;
	tool?: { name: string; arguments: ToolCall["arguments"] };
	/** Token usage reported for this turn so cost and context segments move. */
	usage?: { input?: number; output?: number; cacheRead?: number };
}

interface SceneFile {
	turns?: SceneTurn[];
	/** Replies for every child session an extension spawns in-process (for
	 * example ultra sub-agents); each child replays this list from the start,
	 * so concurrent children stay deterministic. */
	childTurns?: SceneTurn[];
}

interface Report {
	queued: number;
	calls: number;
	instances: number;
	childCalls: number;
}

/** `{slot.id}` and `{slot.path}` in a turn resolve to the most recent
 * `system_bundle_slot` result in the transcript, so a scripted write can
 * target an ID the-system generated at runtime. */
function substituteSlot(value: unknown, transcript: string): unknown {
	if (typeof value === "string") {
		if (
			!value.includes("{slot.") &&
			!value.includes("{spec.") &&
			!value.includes("{prompt:")
		)
			return value;
		// `{prompt:REGEX}`: first capture group of REGEX over the latest user
		// message, for example a fanout item path inside an ultra envelope.
		const withPrompt = value.replaceAll(
			/\{prompt:((?:[^{}\\]|\\.)+)\}/gu,
			(_m, pattern: string) =>
				new RegExp(pattern, "u").exec(
					transcript.split("\u0000")[1] ?? "",
				)?.[1] ?? "",
		);
		const id = /ID: (\S+)/u.exec(transcript)?.[1] ?? "";
		const path = /Path: (\S+)/u.exec(transcript)?.[1] ?? "";
		// `{spec.id}`: the most recent reserved spec, for plans that reference it.
		const spec =
			/Reserved spec bundle\.\s*Title: [^\n]*\s*ID: (\S+)/u.exec(
				transcript,
			)?.[1] ?? "";
		return withPrompt
			.replaceAll("{slot.id}", id)
			.replaceAll("{slot.path}", path)
			.replaceAll("{spec.id}", spec);
	}
	if (Array.isArray(value))
		return value.map((v) => substituteSlot(v, transcript));
	if (value && typeof value === "object")
		return Object.fromEntries(
			Object.entries(value).map(([k, v]) => [k, substituteSlot(v, transcript)]),
		);
	return value;
}

/** Transcript text of every bundle reservation so far, newest first, so
 * `{slot.*}` resolves to the latest and `{spec.id}` to the latest spec. */
function lastSlotText(context: {
	messages: Array<{ role: string; content: unknown }>;
}): string {
	const parts: string[] = [];
	let lastUser = "";
	for (let i = context.messages.length - 1; i >= 0; i--) {
		const message = context.messages[i];
		if (message?.role === "user" && !lastUser)
			lastUser = JSON.stringify(message.content);
		if (message?.role !== "toolResult") continue;
		const text = JSON.stringify(message.content).replaceAll("\\n", "\n");
		if (text.includes("Reserved ") && text.includes("ID: ")) parts.push(text);
	}
	// Slot texts, then a NUL, then the latest user message for `{prompt:…}`.
	return `${parts.join("\n")}\u0000${lastUser}`;
}

function readScene(): SceneFile {
	const path = process.env["PI_DEMO_SCENE"];
	if (!path) return { turns: [{ text: "Hello from the demo model." }] };
	return JSON.parse(readFileSync(path, "utf8")) as SceneFile;
}

export default function demoProvider(pi: ExtensionAPI): void {
	const scene = readScene();
	const turns = scene.turns ?? [];
	const core = createFauxCore({
		api: "openai",
		provider: DEMO_PROVIDER,
		models: [
			{ id: DEMO_MODEL, name: "Demo Model" },
			{ id: DEMO_DEEP_MODEL, name: "Demo Deep" },
		],
	});
	// The runner compares calls made with turns queued after the scene ends.
	// Pi instantiates this extension again for /new, ctx.newSession, and
	// every in-process child session, so instances coordinate through the
	// report file: the primary queue continues where the previous primary
	// stopped, and child sessions (when the scene declares childTurns) each
	// replay childTurns from the start.
	const reportPath = process.env["PI_DEMO_REPORT"];
	let previous: Report = { queued: 0, calls: 0, instances: 0, childCalls: 0 };
	if (reportPath) {
		try {
			previous = JSON.parse(readFileSync(reportPath, "utf8")) as Report;
		} catch {
			// first instance
		}
	}
	const isChild = Boolean(scene.childTurns) && previous.instances > 0;
	const queue = isChild
		? (scene.childTurns ?? [])
		: turns.slice(previous.calls);
	if (reportPath) {
		const write = () =>
			writeFileSync(
				reportPath,
				JSON.stringify({
					queued: turns.length,
					instances: previous.instances + 1,
					calls: isChild
						? previous.calls
						: previous.calls + core.state.callCount,
					childCalls: isChild
						? previous.childCalls + core.state.callCount
						: previous.childCalls,
				} satisfies Report),
			);
		write();
		pi.on("turn_end", write);
		pi.on("session_shutdown", write);
	}
	const build = (turn: SceneTurn, transcript: string) => {
		const text = substituteSlot(turn.text, transcript) as string | undefined;
		const message = turn.tool
			? fauxAssistantMessage([
					...(text ? [{ type: "text" as const, text }] : []),
					fauxToolCall(
						turn.tool.name,
						substituteSlot(
							turn.tool.arguments,
							transcript,
						) as ToolCall["arguments"],
					),
				])
			: fauxAssistantMessage(text ?? "");
		if (turn.usage) {
			const input = turn.usage.input ?? 0;
			const output = turn.usage.output ?? 0;
			const cacheRead = turn.usage.cacheRead ?? 0;
			message.usage = {
				...message.usage,
				input,
				output,
				cacheRead,
				cacheWrite: 0,
				totalTokens: input + output + cacheRead,
			};
		}
		return message;
	};
	// PI_DEMO_TRACE=<file>: append one line per model call with the system
	// prompt head and the latest user text, to align scene.turns with what Pi
	// actually asks (compaction and sub-agent calls are invisible on screen).
	const tracePath = process.env["PI_DEMO_TRACE"];
	const trace = (
		index: number,
		context: {
			systemPrompt?: string;
			messages: Array<{ role: string; content: unknown }>;
		},
	) => {
		if (!tracePath) return;
		const lastUser = [...context.messages]
			.reverse()
			.find((m) => m.role === "user");
		const head = (text: unknown) => JSON.stringify(text ?? "").slice(0, 160);
		writeFileSync(
			tracePath,
			`${isChild ? "child" : "main"} #${index + 1} system=${head(context.systemPrompt)} user=${head(lastUser?.content)}\n`,
			{ flag: "a" },
		);
	};
	core.setResponses(
		queue.map(
			(turn, index): FauxResponseStep =>
				(context) => {
					const typed = context as unknown as {
						systemPrompt?: string;
						messages: Array<{ role: string; content: unknown }>;
					};
					trace(index, typed);
					return build(turn, lastSlotText(typed));
				},
		),
	);
	pi.registerProvider(DEMO_PROVIDER, {
		apiKey: "demo",
		api: "openai",
		baseUrl: "https://demo.invalid",
		streamSimple: core.streamSimple,
		models: [
			{
				id: DEMO_MODEL,
				name: "Demo Model",
				reasoning: false,
				input: ["text"],
				cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
				contextWindow: 200_000,
				maxTokens: 8192,
			},
			{
				id: DEMO_DEEP_MODEL,
				name: "Demo Deep",
				reasoning: true,
				input: ["text"],
				cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
				contextWindow: 1_000_000,
				maxTokens: 32_000,
			},
		],
	});
}