pi-ext

import { readdirSync, readFileSync, writeFileSync } from "node:fs";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import {
	createCodemodeExtension,
	type ExtensionUIContext,
} from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { applyInitialization } from "../init.ts";

const ENTRY = fileURLToPath(new URL("../index.ts", import.meta.url));

const recordingUi = (
	select: (title: string, items: string[]) => string | undefined,
): ExtensionUIContext =>
	({
		select: async (title: string, items: string[]) => select(title, items),
		confirm: async () => false,
		input: async () => undefined,
		notify: () => {},
		onTerminalInput: () => () => {},
		setStatus: () => {},
		setWorkingMessage: () => {},
		setWorkingVisible: () => {},
		setWorkingIndicator: () => {},
		setHiddenThinkingLabel: () => {},
		setWidget: () => {},
		setFooter: () => {},
		setHeader: () => {},
		setTitle: () => {},
		custom: async () => undefined,
		pasteToEditor: () => {},
		setEditorText: () => {},
		getEditorText: () => "",
		editor: async () => undefined,
		addAutocompleteProvider: () => {},
		setEditorComponent: () => {},
		getEditorComponent: () => undefined,
		getAllThemes: () => [],
		getTheme: () => undefined,
		setTheme: () => ({ success: false, error: "UI not available" }),
		getToolsExpanded: () => false,
		setToolsExpanded: () => {},
	}) as unknown as ExtensionUIContext;

describe("the-system pi-test-harness", () => {
	let session: PiSession | undefined;

	afterEach(async () => {
		await session?.dispose();
		session = undefined;
	});

	it("loads in real Pi extension runtime", async () => {
		session = await createPiSession({ extensionPaths: [ENTRY] });

		expect(session.extensions).toHaveLength(1);
		expect(session.extensions[0]!.path).toBe(ENTRY);
		expect(session.session.getActiveToolNames()).not.toContain("mockup");
	});

	it("writes debug lifecycle records in its sandbox when enabled", async () => {
		session = await createPiSession({
			env: { PI_THE_SYSTEM_DEBUG: "1" },
			extensionPaths: [ENTRY],
		});
		const context = session.session.extensionRunner.createCommandContext();
		const [lifecycleExtension] = session.extensions;
		if (!lifecycleExtension) throw new Error("the-system extension not loaded");
		for (const handler of lifecycleExtension.handlers.get("session_start") ??
			[])
			await handler({ type: "session_start" }, context);
		for (const handler of lifecycleExtension.handlers.get("session_shutdown") ??
			[])
			await handler({ type: "session_shutdown" }, context);
		const directory = path.join(
			session.cwd,
			".test-state",
			"pi-ext",
			"debug",
			"the-system",
		);
		const records = readFileSync(
			path.join(directory, readdirSync(directory)[0]!),
			"utf8",
		)
			.trim()
			.split("\n")
			.map((line) => JSON.parse(line));
		expect(records.map((record) => record.event)).toEqual(
			expect.arrayContaining(["session.start", "session.shutdown"]),
		);
	});

	it("registers scoped tools deferred without declaring them to the model", async () => {
		session = await createPiSession({ extensionPaths: [ENTRY] });
		applyInitialization(session.cwd, "PX");
		const extension = session.extensions.find((item) => item.path === ENTRY);
		if (!extension) throw new Error("the-system extension not loaded");
		const sendUserMessage = vi
			.spyOn(session.session, "sendUserMessage")
			.mockResolvedValue();
		const systemCommand = extension.commands.get("system");
		if (!systemCommand) throw new Error("system command not registered");
		await systemCommand.handler(
			"spec let us discuss caching",
			session.session.extensionRunner.createCommandContext(),
		);
		expect(sendUserMessage).toHaveBeenCalledOnce();
		for (const name of [
			"system_bundle_slot",
			"system_board_propose",
			"system_board_update",
		]) {
			const tool = extension.tools.get(name)?.definition;
			expect(tool?.exposure).toBe("deferred");
			expect(tool?.annotations?.readOnlyHint).toBe(false);
			expect(session.session.getCallableToolNames()).toContain(name);
			expect(session.session.getActiveToolNames()).not.toContain(name);
			expect(session.session.systemPrompt).not.toContain(`- ${name}:`);
		}
	});

	it.each([true, false])(
		"preserves nested codemode mutation approval: %s",
		async (approve) => {
			session = await createPiSession({
				extensionPaths: [ENTRY],
				extensionFactories: [createCodemodeExtension()],
				userSettings: { defaultTools: ["+codemode"] },
				responses: [
					fauxAssistantMessage(
						fauxToolCall("codemode", {
							code: 'await tools.edit({ path: ".system/SYSTEM.md", edits: [{ oldText: "Approved.", newText: "Human-approved." }] });',
						}),
					),
					fauxAssistantMessage("approval requested"),
				],
			});
			applyInitialization(session.cwd, "PX");
			const extension = session.extensions.find((item) => item.path === ENTRY);
			if (!extension) throw new Error("the-system extension not loaded");
			const systemCommand = extension.commands.get("system");
			if (!systemCommand) throw new Error("system command not registered");
			vi.spyOn(session.session, "sendUserMessage").mockResolvedValue();
			await systemCommand.handler(
				"spec nested confirmation",
				session.session.extensionRunner.createCommandContext(),
			);
			writeFileSync(
				path.join(session.cwd, ".system", "SYSTEM.md"),
				"Approved.\n",
			);
			const nestedStarts: {
				toolName: string;
				parentToolCallId?: string | undefined;
			}[] = [];
			const unsubscribe = session.session.subscribe((event) => {
				if (event.type === "tool_execution_start" && event.toolName === "edit")
					nestedStarts.push({
						toolName: event.toolName,
						parentToolCallId: event.parentToolCallId,
					});
			});
			const selectCalls: { args: [string, string[]] }[] = [];
			session.session.extensionRunner.setUIContext(
				recordingUi((title, items) => {
					selectCalls.push({ args: [title, items] });
					return items[approve ? 0 : 1];
				}),
				"tui",
			);

			await session.session.prompt("run the scripted nested edit");
			unsubscribe();
			const [approval] = selectCalls;
			expect(approval?.args).toEqual([
				"The System ownership approval",
				[
					"Approve once: edit .system/SYSTEM.md",
					"Deny: keep .system/SYSTEM.md unchanged",
				],
			]);
			expect(nestedStarts).toEqual([
				{ toolName: "edit", parentToolCallId: expect.any(String) },
			]);
			expect(
				readFileSync(path.join(session.cwd, ".system", "SYSTEM.md"), "utf8"),
			).toBe(approve ? "Human-approved.\n" : "Approved.\n");
		},
	);
});