Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/pfui/__tests__/harness.test.ts

Raw
import { execFileSync } from "node:child_process";
import {
	cpSync,
	mkdirSync,
	mkdtempSync,
	readdirSync,
	readFileSync,
	rmSync,
	symlinkSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import {
	type BuildSystemPromptOptions,
	CONFIG_DIR_NAME,
	DefaultPackageManager,
	type ExtensionAPI,
	getAgentDir,
	getDocsPath,
	getExamplesPath,
	getPackageDir,
	getReadmePath,
	SettingsManager,
	VERSION,
} from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createTestSession,
	type TestSession,
	type TestSessionOptions,
} from "../../../test/harness";
import pfui from "../index";

const extensionRoot = path.resolve(import.meta.dirname, "..");
const repositoryRoot = path.resolve(extensionRoot, "../..");
const temporaryPaths: string[] = [];
let t: TestSession | undefined;
let shutdown = false;

function tempDirectory(prefix: string): string {
	const directory = mkdtempSync(path.join(tmpdir(), prefix));
	temporaryPaths.push(directory);
	return directory;
}

function copyPfuiEntrypoint(directory: string): void {
	cpSync(
		path.join(extensionRoot, "index.ts"),
		path.join(directory, "index.ts"),
	);
	mkdirSync(path.join(directory, "src"), { recursive: true });
	cpSync(
		path.join(extensionRoot, "src", "pi-ext-debug.ts"),
		path.join(directory, "src", "pi-ext-debug.ts"),
	);
	cpSync(
		path.join(extensionRoot, "src", "debug.ts"),
		path.join(directory, "src", "debug.ts"),
	);
	symlinkSync(
		path.join(repositoryRoot, "node_modules"),
		path.join(directory, "node_modules"),
		"junction",
	);
}

function copiedPfui(prefix: string): string {
	const directory = tempDirectory(prefix);
	copyPfuiEntrypoint(directory);
	return directory;
}

async function createPfuiSession(
	options: TestSessionOptions = {},
): Promise<TestSession> {
	const cwd = options.cwd ?? tempDirectory("pfui-harness-");
	mkdirSync(path.join(cwd, ".test-home", CONFIG_DIR_NAME, "agent"), {
		recursive: true,
	});
	return createTestSession({
		...options,
		cwd,
		extensionFactories: options.extensionFactories ?? [pfui],
	});
}

async function stopSession(): Promise<void> {
	if (!t || shutdown) return;
	await t.session.extensionRunner.emit({
		type: "session_shutdown",
		reason: "quit",
	});
	shutdown = true;
}

function baseOptions(): BuildSystemPromptOptions {
	return {
		cwd: "/work/project",
		selectedTools: ["read"],
		toolSnippets: { read: "Read files" },
		promptGuidelines: [],
	};
}

async function emittedOptions(
	options: BuildSystemPromptOptions,
): Promise<BuildSystemPromptOptions> {
	if (!t) throw new Error("Test session is not initialized");
	const result = await t.session.extensionRunner.emitBeforeAgentStart(
		"user prompt",
		undefined,
		options,
	);
	return result.systemPromptOptions;
}

afterEach(async () => {
	await stopSession();
	t?.dispose();
	t = undefined;
	shutdown = false;
	delete (globalThis as Record<symbol, unknown>)[
		Symbol.for("@bugabinga/pi-ext-pfui/state")
	];
	for (const directory of temporaryPaths.splice(0)) {
		rmSync(directory, { recursive: true, force: true });
	}
});

describe.sequential("pfui pi-test-harness", () => {
	it("writes enabled lifecycle debug records in an isolated sandbox", async () => {
		const root = mkdtempSync(path.join(tmpdir(), "pi-ext-pfui-debug-"));
		const original = {
			debug: process.env.PI_PFUI_DEBUG,
			state: process.env.XDG_STATE_HOME,
		};
		try {
			process.env.PI_PFUI_DEBUG = "1";
			process.env.XDG_STATE_HOME = path.join(root, "state");
			vi.resetModules();
			const { default: extension } = await import("../index.ts");
			t = await createTestSession({
				cwd: root,
				env: {
					PI_PFUI_DEBUG: "1",
					XDG_STATE_HOME: path.join(root, "state"),
				},
				extensionFactories: [extension],
			});
			await t.session.extensionRunner.emit({
				type: "session_start",
				reason: "startup",
			});
			await emittedOptions({ ...baseOptions(), customPrompt: "CUSTOM" });
			await t.session.extensionRunner.emit({
				type: "session_shutdown",
				reason: "quit",
			});
			const events = readdirSync(
				path.join(root, "state", "pi-ext", "debug", "pfui"),
			).flatMap((file) =>
				readFileSync(
					path.join(root, "state", "pi-ext", "debug", "pfui", file),
					"utf8",
				)
					.trim()
					.split("\n")
					.map((line) => JSON.parse(line)),
			);
			expect(events).toEqual(
				expect.arrayContaining([
					expect.objectContaining({
						extension: "pfui",
						event: "session.start",
					}),
					expect.objectContaining({
						extension: "pfui",
						event: "session.shutdown",
					}),
					expect.objectContaining({
						event: "prompt.apply",
						outcome: "custom_prompt",
					}),
				]),
			);
		} finally {
			t?.dispose();
			t = undefined;
			shutdown = false;
			if (original.debug === undefined) delete process.env.PI_PFUI_DEBUG;
			else process.env.PI_PFUI_DEBUG = original.debug;
			if (original.state === undefined) delete process.env.XDG_STATE_HOME;
			else process.env.XDG_STATE_HOME = original.state;
			rmSync(root, { recursive: true, force: true });
		}
	});

	it("loads through Pi once, first in the generated package list, while the workspace root exposes nothing", async () => {
		t = await createPfuiSession();

		expect(t.session.extensionRunner.extensions).toHaveLength(1);
		const manifest = JSON.parse(
			readFileSync(path.join(repositoryRoot, "package.json"), "utf8"),
		) as { pi?: unknown };
		expect(manifest.pi).toEqual({});

		const packageManager = new DefaultPackageManager({
			cwd: t.cwd,
			agentDir: path.join(t.cwd, ".test-home", CONFIG_DIR_NAME, "agent"),
			settingsManager: SettingsManager.inMemory(),
		});
		const resolved = await packageManager.resolveExtensionSources(
			[repositoryRoot],
			{ temporary: true },
		);
		expect(resolved.extensions).toEqual([]);

		// Pi loads one generated package per extension; pfui must come first so it
		// owns the system prompt before other extensions contribute sections.
		const entries = JSON.parse(
			execFileSync(
				process.execPath,
				["--no-warnings", "scripts/build/settings-entries.ts"],
				{ cwd: repositoryRoot, encoding: "utf8" },
			),
		) as string[];
		const suffix = (entry: string) => entry.split(path.sep).slice(-2).join("/");
		expect(suffix(entries[0] ?? "")).toBe("dist/pfui");
		expect(suffix(entries[1] ?? "")).toBe("dist/footer");
		expect(
			entries.filter((entry) => suffix(entry) === "dist/pfui"),
		).toHaveLength(1);
		expect(new Set(entries).size).toBe(entries.length);
	});

	it("composes the fallback with dynamic tools and deduplicated guidelines", async () => {
		t = await createPfuiSession();
		const options = baseOptions();
		options.selectedTools = ["write", "read", "write"];
		options.toolSnippets = {
			read: "Read current files",
			write: "Write complete files",
			inactive: "Must stay hidden",
		};
		options.promptGuidelines = [
			"  Keep paths explicit.  ",
			"Keep paths explicit.",
			"",
			"Verify changes.",
		];

		const result = await emittedOptions(options);
		const identity = readFileSync(
			path.join(extensionRoot, "SYSTEM.md"),
			"utf8",
		).trim();

		expect(result.customPrompt).toBe(identity);
		expect(result.sections?.tools).toContain(
			"- write: Write complete files\n- read: Read current files",
		);
		expect(result.sections?.tools).not.toContain("inactive");
		expect(result.sections?.general_guidelines).toBe(
			"Keep paths explicit.\nVerify changes.",
		);
		expect(result.sections?.tool_guidelines).toBe("");
	});

	it("renders guidelines contributed by later extensions", async () => {
		let rendered = "";
		const contributeGuidelines = (pi: ExtensionAPI) => {
			pi.on("before_agent_start", (event) => {
				event.systemPromptOptions.promptGuidelines.push(
					"Late general guidance",
				);
				event.systemPromptOptions.toolGuidelines.read = ["Late read guidance"];
			});
		};
		const observePrompt = (pi: ExtensionAPI) => {
			pi.on("before_agent_start", (event) => {
				rendered = event.systemPrompt;
			});
		};
		t = await createPfuiSession({
			extensionFactories: [pfui, contributeGuidelines, observePrompt],
		});

		await emittedOptions(baseOptions());

		expect(rendered).toContain(`<general_guidelines>
Late general guidance
</general_guidelines>`);
		expect(rendered).toContain(`<tool_guidelines>
<tool name="read">
Late read guidance
</tool>
</tool_guidelines>`);
	});

	it("passes a configured custom prompt through unchanged", async () => {
		t = await createPfuiSession();
		const options = { ...baseOptions(), customPrompt: "CUSTOM" };

		const result = await emittedOptions(options);
		expect(result.customPrompt).toBe("CUSTOM");
		expect(result.sections?.tools).toBeUndefined();
		expect(result.sections?.general_guidelines).toBeUndefined();
		expect(result.sections?.tool_guidelines).toBeUndefined();
	});

	it("leaves the fallback available for later extension composition", async () => {
		const appendLater = (pi: ExtensionAPI) => {
			pi.on("before_agent_start", (event) => ({
				systemPrompt: `${event.systemPrompt}\n\nLATER EXTENSION`,
			}));
		};
		t = await createPfuiSession({
			extensionFactories: [pfui, appendLater],
		});

		const result = await emittedOptions(baseOptions());
		expect(result.forceSystemPrompt).toContain(
			"You are the spirit in the machine.\nBound inside Pi, give the user's will form.",
		);
		expect(result.forceSystemPrompt).toContain("<tools>");
		expect(result.forceSystemPrompt).not.toContain("<rules>");
		expect(result.forceSystemPrompt).toMatch(/<\/tools>\n\nLATER EXTENSION$/);
	});

	it("preserves Pi prompt section formatting and normalizes cwd", async () => {
		let rendered = "";
		const observePrompt = (pi: ExtensionAPI) => {
			pi.on("before_agent_start", (event) => {
				rendered = event.systemPrompt;
			});
		};
		t = await createPfuiSession({
			extensionFactories: [pfui, observePrompt],
		});
		const skill = t.session.resourceLoader
			.getSkills()
			.skills.find(
				(candidate: { name: string }) => candidate.name === "pi-knowledge",
			);
		expect(skill).toBeDefined();
		const options: BuildSystemPromptOptions = {
			...baseOptions(),
			cwd: "C:\\work\\project",
			appendSystemPrompt: "APPENDED",
			contextFiles: [{ path: "/repo/AGENTS.md", content: "PROJECT RULE" }],
			skills: [skill],
		};

		const result = await emittedOptions(options);
		const identity = readFileSync(
			path.join(extensionRoot, "SYSTEM.md"),
			"utf8",
		).trim();
		expect(result.customPrompt).toBe(identity);
		expect(result.appendSystemPrompt).toBe("APPENDED");
		expect(result.contextFiles).toEqual([
			{ path: "/repo/AGENTS.md", content: "PROJECT RULE" },
		]);
		expect(result.skills).toEqual([skill]);
		expect(result.sections?.tools).toContain("- read: Read files");
		expect(result.sections?.general_guidelines).toBe("");
		expect(result.sections?.tool_guidelines).toBe("");
		expect(rendered).not.toContain("<general_guidelines>");
		expect(rendered).not.toContain("<tool_guidelines>");
		expect(result.sections?.cwd).toBeUndefined();
		expect(result.sections?.docs).toBeUndefined();
	});

	it("omits skills when neither read nor bash is active", async () => {
		t = await createPfuiSession();
		const skill = t.session.resourceLoader
			.getSkills()
			.skills.find(
				(candidate: { name: string }) => candidate.name === "pi-knowledge",
			);
		const result = await emittedOptions({
			...baseOptions(),
			selectedTools: ["powershell"],
			toolSnippets: { powershell: "Run PowerShell" },
			skills: [skill],
		});

		expect(result.skills).toEqual([skill]);
		expect(result.sections?.tools).not.toContain("pi-knowledge");
		expect(result.sections?.general_guidelines).toBe(
			"Use PowerShell for file operations like listing, searching, and finding files",
		);
		expect(result.sections?.tool_guidelines).toBe("");
	});

	it("separates general and tool guidelines without escaping their text", async () => {
		t = await createPfuiSession();
		const result = await emittedOptions({
			...baseOptions(),
			selectedTools: ["edit", "angel"],
			toolSnippets: { edit: "Edit files", angel: "Consult Angel" },
			toolGuidelines: {
				edit: ["Edit carefully", "Shared <rule> & text", "Edit carefully"],
				angel: ["Shared <rule> & text", "Ask a focused question"],
			},
			promptGuidelines: ["Global guidance", "Global guidance"],
		});

		expect(result.sections?.general_guidelines).toBe("Global guidance");
		expect(result.sections?.tool_guidelines).toBe(`<tool name="edit">
Edit carefully
Shared <rule> & text
</tool>

<tool name="angel">
Shared <rule> & text
Ask a focused question
</tool>`);
	});

	it("discovers the bundled pi-knowledge skill with its contract", async () => {
		t = await createPfuiSession();
		const skill = t.session.resourceLoader
			.getSkills()
			.skills.find(
				(candidate: { name: string }) => candidate.name === "pi-knowledge",
			);

		expect(skill?.filePath).toBe(
			path.join(extensionRoot, "skills", "pi-knowledge", "SKILL.md"),
		);
		const source = readFileSync(skill.filePath, "utf8");
		expect(source).toContain("Requires Pi and either bash or PowerShell.");
		expect(source).toContain(
			"questions about Pi and for implementing or changing Pi",
		);
		expect(source).toContain("$PI_CONTEXT_DOCS_PATH");
		expect(source).toContain("$env:PI_CONTEXT_DOCS_PATH");
		expect(source).toContain("Never dump the complete environment");
		expect(source).toContain("If neither Bash nor PowerShell is available");
		expect(source).toContain("Read every selected Markdown file completely");
		expect(source).toContain("If a mapped file moved or is absent");
	});

	it("keeps a valid bundled skill when Pi reports only warnings", async () => {
		const warned = tempDirectory("pfui-warned-");
		temporaryPaths.push(warned);
		copyPfuiEntrypoint(warned);
		cpSync(
			path.join(extensionRoot, "SYSTEM.md"),
			path.join(warned, "SYSTEM.md"),
		);
		mkdirSync(path.join(warned, "skills", "pi-knowledge"), {
			recursive: true,
		});
		writeFileSync(
			path.join(warned, "skills", "pi-knowledge", "SKILL.md"),
			`---\nname: pi-knowledge\ndescription: "${"x".repeat(1025)}"\n---\n\n# Pi knowledge\n`,
		);
		t = await createTestSession({
			cwd: tempDirectory("pfui-warned-cwd-"),
			extensions: [path.join(warned, "index.ts")],
		});

		expect(
			t.session.resourceLoader
				.getSkills()
				.skills.some(
					(candidate: { name: string }) => candidate.name === "pi-knowledge",
				),
		).toBe(true);
	});

	it("preserves Pi 0.86 tool guidelines and custom sections", async () => {
		t = await createPfuiSession();
		const options: BuildSystemPromptOptions = {
			...baseOptions(),
			toolGuidelines: {
				read: ["Read before editing."],
			},
			sections: {
				custom_context: "Keep the change focused.",
			},
		};

		const result = await emittedOptions(options);

		expect(t.events.uiCallsFor("notify")).toHaveLength(0);
		expect(result.sections?.general_guidelines).toBe("");
		expect(result.sections?.tool_guidelines).toBe(`<tool name="read">
Read before editing.
</tool>`);
		expect(result.sections?.custom_context).toBe("Keep the change focused.");
	});

	it("exports public Pi context and cleans only unchanged owned values", async () => {
		let packageDirAtShutdown: string | undefined;
		let docsPathAtShutdown: string | undefined;
		let versionAtShutdown: string | undefined;
		const observeShutdown = (pi: ExtensionAPI) => {
			pi.on("session_shutdown", () => {
				packageDirAtShutdown = process.env.PI_CONTEXT_PACKAGE_DIR;
				docsPathAtShutdown = process.env.PI_CONTEXT_DOCS_PATH;
				versionAtShutdown = process.env.PI_CONTEXT_VERSION;
			});
		};
		t = await createPfuiSession({
			extensionFactories: [pfui, observeShutdown],
			env: { PI_CONTEXT_VERSION: VERSION },
		});
		const packageDir = getPackageDir();
		const expected = {
			PI_CONTEXT_PACKAGE_DIR: packageDir,
			PI_CONTEXT_AGENT_DIR: getAgentDir(),
			PI_CONTEXT_README_PATH: getReadmePath(),
			PI_CONTEXT_DOCS_PATH: getDocsPath(),
			PI_CONTEXT_EXAMPLES_PATH: getExamplesPath(),
			PI_CONTEXT_CHANGELOG_PATH: path.join(packageDir, "CHANGELOG.md"),
			PI_CONTEXT_PACKAGE_JSON_PATH: path.join(packageDir, "package.json"),
			PI_CONTEXT_VERSION: VERSION,
			PI_CONTEXT_CONFIG_DIR_NAME: CONFIG_DIR_NAME,
		};
		for (const [name, value] of Object.entries(expected)) {
			expect(process.env[name]).toBe(value);
		}

		process.env.PI_CONTEXT_DOCS_PATH = "kept-after-shutdown";
		await stopSession();

		expect(packageDirAtShutdown).toBeUndefined();
		expect(docsPathAtShutdown).toBe("kept-after-shutdown");
		expect(versionAtShutdown).toBe(VERSION);
	});

	it("preserves a conflicting context variable and warns", async () => {
		t = await createPfuiSession({
			env: { PI_CONTEXT_VERSION: "external-version" },
		});

		expect(process.env.PI_CONTEXT_VERSION).toBe("external-version");
		const notifications = t.events.uiCallsFor("notify");
		expect(notifications).toHaveLength(1);
		expect(String(notifications[0]?.args[0])).toContain(
			"PI_CONTEXT_VERSION already has a different value",
		);

		await stopSession();
		expect(process.env.PI_CONTEXT_VERSION).toBe("external-version");
	});

	it("omits a missing resolved resource path and warns", async () => {
		t = await createTestSession({
			cwd: tempDirectory("pfui-missing-path-cwd-"),
			extensionFactories: [pfui],
		});

		expect(process.env.PI_CONTEXT_AGENT_DIR).toBeUndefined();
		const notifications = t.events.uiCallsFor("notify");
		expect(notifications).toHaveLength(1);
		expect(String(notifications[0]?.args[0])).toContain(
			"PI_CONTEXT_AGENT_DIR was omitted",
		);
	});

	it("preserves Pi's prompt when the bundled fallback is missing", async () => {
		const broken = copiedPfui(".pfui-missing-fallback-");
		cpSync(path.join(extensionRoot, "skills"), path.join(broken, "skills"), {
			recursive: true,
		});
		t = await createTestSession({
			cwd: tempDirectory("pfui-missing-fallback-cwd-"),
			extensions: [path.join(broken, "index.ts")],
		});

		const result = await emittedOptions(baseOptions());
		expect(result.customPrompt).toBeUndefined();
		const warnings = t.events
			.uiCallsFor("notify")
			.map((notification) => String(notification.args[0]))
			.join("\n");
		expect(warnings).toContain("Bundled SYSTEM.md is unavailable");
	});

	it.each([
		["missing", undefined, "is missing"],
		["invalid", "---\nname: pi-knowledge\n---\ninvalid\n", "is invalid"],
	])("omits a %s bundled skill and warns", async (_case, source, message) => {
		const broken = copiedPfui(`.pfui-${_case}-skill-`);
		cpSync(
			path.join(extensionRoot, "SYSTEM.md"),
			path.join(broken, "SYSTEM.md"),
		);
		if (source) {
			mkdirSync(path.join(broken, "skills", "pi-knowledge"), {
				recursive: true,
			});
			writeFileSync(
				path.join(broken, "skills", "pi-knowledge", "SKILL.md"),
				source,
			);
		}
		t = await createTestSession({
			cwd: tempDirectory(`pfui-${_case}-skill-cwd-`),
			extensions: [path.join(broken, "index.ts")],
		});

		expect(
			t.session.resourceLoader
				.getSkills()
				.skills.some(
					(candidate: { name: string }) => candidate.name === "pi-knowledge",
				),
		).toBe(false);
		const warnings = t.events
			.uiCallsFor("notify")
			.map((notification) => String(notification.args[0]))
			.join("\n");
		expect(warnings).toContain(`Bundled pi-knowledge skill ${message}`);
	});
});