pi-ext

import { AsyncLocalStorage } from "node:async_hooks";
import { execFile } from "node:child_process";
import { rmSync } from "node:fs";
import {
	appendFile,
	chmod,
	copyFile,
	cp,
	mkdir,
	mkdtemp,
	rm,
	writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import path, { delimiter } from "node:path";
import { promisify } from "node:util";
import {
	type Context,
	fauxAssistantMessage,
	fauxProvider,
	InMemoryCredentialStore,
} from "@earendil-works/pi-ai";
import type { ExtensionUIContext } from "@earendil-works/pi-coding-agent";
import { stripTerminalSequences } from "@earendil-works/pi-tui";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { cleanGitEnvironment } from "../git.js";
import type { Hunk, Plan } from "../types.js";

const repositorySeeds = new Map<string, Promise<string>>();

export async function copySeedRepository(
	name: string,
	initialize: (root: string) => Promise<void>,
): Promise<string> {
	let seed = repositorySeeds.get(name);
	if (!seed) {
		const creation = (async () => {
			const root = await mkdtemp(path.join(tmpdir(), `strata-seed-${name}-`));
			try {
				await initialize(root);
				process.once("exit", () =>
					rmSync(root, { recursive: true, force: true }),
				);
				return root;
			} catch (error) {
				await rm(root, { recursive: true, force: true });
				throw error;
			}
		})();
		seed = creation;
		repositorySeeds.set(name, creation);
		void creation.catch(() => {
			if (repositorySeeds.get(name) === creation) repositorySeeds.delete(name);
		});
	}

	const source = await seed;
	const root = await mkdtemp(path.join(tmpdir(), `strata-${name}-`));
	try {
		await rm(root, { recursive: true, force: true });
		await cp(source, root, { recursive: true });
		cleanupNative(() => rm(root, { recursive: true, force: true }));
		return root;
	} catch (error) {
		await rm(root, { recursive: true, force: true });
		throw error;
	}
}

const nativeScope = new AsyncLocalStorage<{
	signal: AbortSignal;
	pending: Set<Promise<unknown>>;
	cancelling: Promise<unknown>[];
	cleanup: AsyncDisposableStack;
}>();

// Used by Vitest aroundEach; no runner import, so Playwright can use this file too.
export async function nativeTest(
	run: () => Promise<void>,
	context: { signal: AbortSignal },
): Promise<void> {
	const controller = new AbortController();
	const scope = {
		signal: AbortSignal.any([context.signal, controller.signal]),
		pending: new Set<Promise<unknown>>(),
		cancelling: [] as Promise<unknown>[],
		cleanup: new AsyncDisposableStack(),
	};
	await nativeScope.run(scope, async () => {
		try {
			await run();
		} finally {
			controller.abort();
			while (scope.pending.size) await Promise.allSettled([...scope.pending]);
			await scope.cleanup.disposeAsync();
		}
	});
}

export function testSignal(signal?: AbortSignal): AbortSignal {
	const parent = nativeScope.getStore()?.signal;
	return parent && signal
		? AbortSignal.any([parent, signal])
		: (parent ?? signal ?? new AbortController().signal);
}

export { own as nativeOperation };

function own<T>(work: () => Promise<T>): Promise<T> {
	const scope = nativeScope.getStore();
	let pending: Promise<T>;
	try {
		pending = work();
	} catch (error) {
		pending = Promise.reject(error);
	}
	scope?.pending.add(pending);
	void pending.then(
		() => scope?.pending.delete(pending),
		() => scope?.pending.delete(pending),
	);
	return pending;
}

export function cleanupNative(cleanup: () => Promise<void>): void {
	nativeScope.getStore()?.cleanup.defer(cleanup);
}

// Register an in-flight cancellation so fake-executable cleanup drains it before
// restoring PATH, matching createFixture's abort-driven cancel handler.
export function pushCancelling(promise: Promise<unknown>): void {
	nativeScope.getStore()?.cancelling.push(promise);
}

const exec = promisify(execFile);
export function git(cwd: string, ...args: string[]): Promise<string> {
	return own(async () => {
		const signal = testSignal();
		signal.throwIfAborted();
		const command = exec("git", args, {
			cwd,
			signal,
			env: cleanGitEnvironment(),
		});
		const closed = new Promise<void>((resolve) =>
			command.child.once("close", () => resolve()),
		);
		try {
			return (await command).stdout.trim();
		} finally {
			await closed;
		}
	});
}

export function withFakeGh(
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	return own(() => fakeGh(program, check));
}

async function fakeGh(
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	testSignal().throwIfAborted();
	const root = await mkdtemp(path.join(tmpdir(), "strata-gh-"));
	try {
		const bin = path.join(root, "bin");
		await mkdir(bin);
		const executable = path.join(
			bin,
			process.platform === "win32" ? "gh.exe" : "gh",
		);
		await copyFile(process.execPath, executable);
		await chmod(executable, 0o755);
		const preload = path.join(root, "fake-gh.cjs");
		await writeFile(
			preload,
			`if (/^gh(?:\\.exe)?$/i.test(require("node:path").basename(process.execPath)) && process.env.STRATA_FAKE_GH_PRELOAD === __filename) {\nconst original = process.argv.slice(1); original[0] = require("node:path").basename(original[0]); process.argv = [process.execPath, "fake-gh", ...(process.execArgv.includes("-c") ? ["-c"] : []), ...original]; const args = process.argv.slice(2);\n${program}\nprocess.exit(process.exitCode ?? 0);\n}`,
		);
		const previousPath = process.env["PATH"];
		const previousOptions = process.env["NODE_OPTIONS"];
		const previousPreload = process.env["STRATA_FAKE_GH_PRELOAD"];
		process.env["PATH"] = `${bin}${delimiter}${previousPath ?? ""}`;
		process.env["STRATA_FAKE_GH_PRELOAD"] = preload;
		process.env["NODE_OPTIONS"] = [
			previousOptions,
			`--require=${JSON.stringify(preload)}`,
		]
			.filter(Boolean)
			.join(" ");
		try {
			await check();
		} finally {
			await Promise.allSettled(nativeScope.getStore()?.cancelling ?? []);
			if (previousPath === undefined) delete process.env["PATH"];
			else process.env["PATH"] = previousPath;
			if (previousOptions === undefined) delete process.env["NODE_OPTIONS"];
			else process.env["NODE_OPTIONS"] = previousOptions;
			if (previousPreload === undefined)
				delete process.env["STRATA_FAKE_GH_PRELOAD"];
			else process.env["STRATA_FAKE_GH_PRELOAD"] = previousPreload;
		}
	} finally {
		await rm(root, { recursive: true, force: true });
	}
}

export function withFakeExecutable(
	name: string,
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	return own(() => fakeExecutable(name, program, check));
}

async function fakeExecutable(
	name: string,
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	testSignal().throwIfAborted();
	const root = await mkdtemp(path.join(tmpdir(), "strata-executable-"));
	try {
		const bin = path.join(root, "bin");
		await mkdir(bin);
		const executable = path.join(
			bin,
			process.platform === "win32" ? `${name}.exe` : name,
		);
		await copyFile(process.execPath, executable);
		await chmod(executable, 0o755);
		const preload = path.join(root, "fake-executable.cjs");
		await writeFile(
			preload,
			`const executable = require("node:path").basename(process.execPath).toLowerCase();
if (executable === ${JSON.stringify(name.toLowerCase())} || executable === ${JSON.stringify(`${name.toLowerCase()}.exe`)}) {
const original = process.argv.slice(1); original[0] = require("node:path").basename(original[0]); process.argv = [process.execPath, ${JSON.stringify(name)}, ...(process.execArgv.includes("-c") ? ["-c"] : []), ...original]; const args = process.argv.slice(2);
${program}
process.exit(process.exitCode ?? 0);
}`,
		);
		const previousPath = process.env["PATH"];
		const previousOptions = process.env["NODE_OPTIONS"];
		process.env["PATH"] = `${bin}${delimiter}${previousPath ?? ""}`;
		process.env["NODE_OPTIONS"] = [
			previousOptions,
			`--require=${JSON.stringify(preload)}`,
		]
			.filter(Boolean)
			.join(" ");
		try {
			await check();
		} finally {
			await Promise.allSettled(nativeScope.getStore()?.cancelling ?? []);
			if (previousPath === undefined) delete process.env["PATH"];
			else process.env["PATH"] = previousPath;
			if (previousOptions === undefined) delete process.env["NODE_OPTIONS"];
			else process.env["NODE_OPTIONS"] = previousOptions;
		}
	} finally {
		await rm(root, { recursive: true, force: true });
	}
}

export function fixturePlan(hunks: Array<Pick<Hunk, "id" | "path">>): Plan {
	return {
		summary: "Bound retries before using them.",
		cohorts: [
			{
				title: "Payment retries",
				layers: hunks.map((hunk, index) => ({
					id: `layer-${index + 1}`,
					title: `Review ${hunk.path}`,
					summary: `Layer ${index + 1} explains the original captured change.`,
					hunks: [{ id: hunk.id, summary: `Explanation for ${hunk.path}.` }],
					flow: ["Contract", "Consumer", "Tests"],
				})),
			},
		],
	};
}

type Fixture = {
	t: {
		env: NodeJS.ProcessEnv;
		session: PiSession["session"];
		events: {
			ui: Array<{ method: string; args: unknown[] }>;
			messages: Array<
				| PiSession["session"]["messages"][number]
				| {
						role: string;
						customType?: unknown;
						content?: unknown;
						display?: unknown;
						details?: unknown;
				  }
			>;
		};
	};
	root: string;
	ctx: ReturnType<PiSession["session"]["extensionRunner"]["createContext"]>;
	runner: PiSession["session"]["extensionRunner"];
	command: { handler: (args: string, ctx: unknown) => Promise<void> };
	faux: ReturnType<typeof fauxProvider>;
	notifications: Array<{ text: string; severity: string }>;
	widgets: Array<{ key: string; content: unknown; options: unknown }>;
	prompts: string[];
	run: (args?: string) => Promise<void>;
	widgetLines: (width?: number) => string[];
	widgetText: (width?: number) => string;
	calls: number;
	url: string;
	close: () => Promise<void>;
};

export function openFixture(
	mode: "tui" | "rpc" | "json" = "tui",
): Promise<Fixture> {
	return own(() => createFixture(mode));
}

async function createFixture(mode: "tui" | "rpc" | "json") {
	const signal = testSignal();
	signal.throwIfAborted();
	const root = await copySeedRepository("pi", async (seed) => {
		await git(seed, "init", "--quiet");
		await appendFile(
			path.join(seed, ".git", "config"),
			"[user]\n\tname = Strata Test\n\temail = strata@example.invalid\n",
		);
		await writeFile(
			path.join(seed, "contract.ts"),
			"export const attempts = 1;\n",
		);
		await writeFile(
			path.join(seed, "consumer.ts"),
			"export const retry = false;\n",
		);
		await git(seed, "add", ".");
		await git(seed, "commit", "--quiet", "-m", "base");
	});
	let session: PiSession | undefined;
	let closing: Promise<void> | undefined;
	const close = () =>
		(closing ??= (async () => {
			try {
				await session?.dispose();
			} finally {
				await rm(root, { recursive: true, force: true });
			}
		})());
	cleanupNative(close);
	try {
		await writeFile(
			path.join(root, "contract.ts"),
			"export const attempts = 3;\n",
		);
		await writeFile(
			path.join(root, "consumer.ts"),
			"export const retry = true;\n",
		);
		session = await createPiSession({
			credentials: new InMemoryCredentialStore(),
			cwd: root,
			env: { STRATA_NO_OPEN: "1" },
			extensionPaths: [path.resolve(import.meta.dirname, "../index.ts")],
		});
		const active = session;
		// The owned fixture roots its agent and XDG sandbox under cwd; keep them out
		// of the repository worktree so clean-worktree checks stay green, since the
		// sandbox lives outside the repository entirely.
		await appendFile(
			path.join(root, ".git", "info", "exclude"),
			"agent/\n.home/\n.config/\n.cache/\n.test-state/\n.data/\n",
		);
		signal.throwIfAborted();
		const runner = active.session.extensionRunner;
		const emit = runner.emit.bind(runner);
		runner.emit = ((...args: Parameters<typeof emit>) =>
			own(() => emit(...args))) as typeof runner.emit;
		const faux = fauxProvider({
			provider: "strata-integration",
			models: [{ id: "review", maxTokens: 16_384 }],
		});
		let calls = 0;
		const prompts: string[] = [];
		const respond = (context: Context) => {
			calls++;
			const message = context.messages.findLast((item) => item.role === "user");
			const text =
				typeof message?.content === "string"
					? message.content
					: (message?.content ?? [])
							.filter((part) => part.type === "text")
							.map((part) => part.text)
							.join("");
			prompts.push(text);
			if (text.startsWith("Organize")) {
				const hunks = [
					...text.matchAll(
						/<file path="(contract\.ts|consumer\.ts)">\s*<hunk id="([^"]+)"/g,
					),
				].map(([, filePath, id]) => ({ id: id!, path: filePath! }));
				return fauxAssistantMessage(JSON.stringify(fixturePlan(hunks)));
			}
			return fauxAssistantMessage("Three attempts include the first call.");
		};
		faux.setResponses(Array.from({ length: 20 }, () => respond));
		const base = runner.createContext();
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("strata-integration", "review");
		if (!model) throw new Error("Fixture model missing");
		await active.session.setModel(model);
		const notifications: Array<{ text: string; severity: string }> = [];
		const widgets: Array<{ key: string; content: unknown; options: unknown }> =
			[];
		const uiCalls: Array<{ method: string; args: unknown[] }> = [];
		const statusMessages: Array<{
			role: string;
			customType?: unknown;
			content?: unknown;
			display?: unknown;
			details?: unknown;
		}> = [];
		// pi.sendMessage maps to AgentSession.sendCustomMessage. Intercept it so
		// print mode writes nothing to stdout while tests still observe the emitted
		// custom status messages.
		active.session.sendCustomMessage = (async (message) => {
			statusMessages.push({ role: "custom", ...message });
		}) as typeof active.session.sendCustomMessage;
		const capturingUi = {
			notify: (text: string, severity = "info") => {
				uiCalls.push({ method: "notify", args: [text, severity] });
				notifications.push({ text, severity });
			},
			setWidget: (key: string, content: unknown, options?: unknown) => {
				uiCalls.push({ method: "setWidget", args: [key, content, options] });
				widgets.push({ key, content, options });
			},
			setStatus: (key: string, text?: string) => {
				uiCalls.push({ method: "setStatus", args: [key, text] });
			},
			setFooter: () => {},
			confirm: async () => true,
		};
		if (mode === "json") runner.setUIContext(undefined, "json");
		else
			runner.setUIContext(capturingUi as unknown as ExtensionUIContext, mode);
		const ctx = Object.create(base);
		Object.defineProperty(ctx, "waitForIdle", { value: async () => {} });
		// `extensions` is private on the real runner; reach the registered command
		// structurally, exactly as the excluded contract test does at runtime.
		const command = (
			runner as unknown as {
				extensions: Array<{
					commands: Map<
						string,
						{ handler: (args: string, ctx: unknown) => Promise<void> }
					>;
				}>;
			}
		).extensions[0]!.commands.get("strata");
		if (!command) throw new Error("strata command missing");
		const scope = nativeScope.getStore();
		const cancel = () => {
			const stopped = own(() => command.handler("close", ctx));
			scope?.cancelling.push(stopped);
		};
		signal.throwIfAborted();
		signal.addEventListener("abort", cancel, { once: true });
		cleanupNative(async () => {
			signal.removeEventListener("abort", cancel);
		});
		const run = (args = "") =>
			own(async () => {
				signal.throwIfAborted();
				await command.handler(args, ctx);
				signal.throwIfAborted();
			});
		return {
			t: {
				env: process.env,
				session: active.session,
				events: {
					ui: uiCalls,
					get messages() {
						return [...active.session.messages, ...statusMessages];
					},
				},
			},
			root,
			ctx,
			runner,
			command,
			faux,
			notifications,
			widgets,
			prompts,
			run,
			widgetLines(width = 120): string[] {
				const content = widgets.findLast(
					(item) => item.key === "strata",
				)?.content;
				if (content === undefined) return [];
				if (Array.isArray(content)) return content as string[];
				if (typeof content !== "function") return [];
				const theme = {
					bold: (text: string) => text,
					fg: (_color: string, text: string) => text,
				};
				return content({}, theme).render(width);
			},
			widgetText(width = 120): string {
				return this.widgetLines(width).map(stripTerminalSequences).join("\n");
			},
			get calls() {
				return calls;
			},
			get url(): string {
				const line = this.widgetLines(160)[1] ?? "";
				const match = /\x1b]8;;([^\x1b]+)\x1b\\/.exec(line);
				if (match?.[1]) return match[1];
				const notice = notifications.findLast((item) =>
					item.text.startsWith("Strata review: "),
				);
				if (notice) return notice.text.slice("Strata review: ".length);
				const message = statusMessages.findLast(
					(item) =>
						item.role === "custom" &&
						item.customType === "strata-status" &&
						typeof item.content === "string" &&
						item.content.startsWith("Strata review: "),
				);
				if (message && typeof message.content === "string")
					return message.content.slice("Strata review: ".length);
				throw new Error(`No review URL: ${line}`);
			},
			close,
		};
	} catch (error) {
		await close();
		throw error;
	}
}