pi-ext

import {
	mkdirSync,
	mkdtempSync,
	readFileSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
	DefaultResourceLoader,
	SettingsManager,
} from "@earendil-works/pi-coding-agent";
import { afterEach, expect, it, vi } from "vitest";
import {
	formatExtensionOverview,
	measureExtensions,
	registrationSignals,
} from "./context-extensions.mjs";
import {
	inspectExtensionRegistrations,
	resourceOptions,
} from "./context-startup-session.mjs";

let directory: string | undefined;
afterEach(() => {
	vi.unstubAllEnvs();
	if (directory) rmSync(directory, { recursive: true, force: true });
	directory = undefined;
});

it("inventories actual maps, provider ownership, and future registrations without configs or tool bodies", () => {
	const extension = (path: string) => ({
		path,
		tools: new Map(),
		handlers: new Map(),
	});
	const one = extension("/extensions/one/index.ts");
	one.tools.set("example", { description: "PRIVATE BODY" });
	one.handlers.set("future_event", [() => {}]);
	const result = {
		extensions: [
			one,
			extension("/extensions/new.ts"),
			extension("<inline:measurement>"),
		],
		runtime: {
			pendingProviderRegistrations: [
				{ name: "custom", extensionPath: one.path, config: "PRIVATE CONFIG" },
			],
			pendingNativeProviderRegistrations: [
				{
					provider: { id: "native", other: "PRIVATE CONFIG" },
					extensionPath: one.path,
				},
			],
		},
	};
	const inventory = inspectExtensionRegistrations(result);
	expect(inventory).toEqual([
		{
			path: one.path,
			name: "one",
			tools: ["example"],
			hooks: ["future_event"],
			providers: ["custom", "native"],
		},
		{
			path: "/extensions/new.ts",
			name: "new",
			tools: [],
			hooks: [],
			providers: [],
		},
	]);
	expect(JSON.stringify(inventory)).not.toContain("PRIVATE");
	expect(
		inspectExtensionRegistrations({ ...result, runtime: {} }, inventory),
	).toEqual(inventory);
	expect(registrationSignals(inventory[0])).toEqual(["tools", "provider"]);
	expect(registrationSignals(inventory[1])).toEqual(["unknown"]);
	expect(
		registrationSignals({ tools: [], providers: [], hooks: ["future_event"] }),
	).toEqual(["other-hooks"]);
	expect(
		registrationSignals({
			tools: [],
			providers: [],
			hooks: [
				"before_agent_start",
				"context",
				"before_provider_request",
				"resources_discover",
				"input",
				"session_start",
			],
		}),
	).toEqual(["prompt", "context", "payload", "resources", "input", "startup"]);
});

it.each([true, false])(
	"omits modules before evaluation, retaining package skills and inline factories (keep file: %s)",
	async (keep) => {
		directory = mkdtempSync(join(tmpdir(), "pi-extension-profile-test-"));
		const cwd = join(directory, "project");
		const agentDir = join(directory, "agent");
		const packageDir = join(directory, "package");
		for (const path of [cwd, agentDir, packageDir]) mkdirSync(path);
		writeFileSync(
			join(packageDir, "package.json"),
			JSON.stringify({
				name: "fixture",
				pi: { extensions: ["omit.mjs"], skills: ["skill.md"] },
			}),
		);
		writeFileSync(
			join(packageDir, "omit.mjs"),
			"throw new Error('Excluded module executed'); export default () => {}; ",
		);
		writeFileSync(
			join(packageDir, "skill.md"),
			"---\nname: fixture\ndescription: Fixture skill\n---\nBody",
		);
		const kept = join(directory, "keep.mjs");
		writeFileSync(
			kept,
			"export default pi => pi.registerFlag('kept', { type: 'boolean', description: 'kept' });",
		);
		const settingsManager = SettingsManager.inMemory(
			{ packages: [packageDir] },
			{ projectTrusted: false },
		);
		const loader = new DefaultResourceLoader({
			cwd,
			agentDir,
			settingsManager,
			...resourceOptions(
				"headless+1msg",
				[
					{
						name: "builtin",
						factory(pi) {
							pi.registerFlag("inline", {
								type: "boolean",
								description: "inline",
							});
						},
					},
				],
				{ extensionPaths: keep ? [kept] : [] },
			),
		});
		await loader.reload();
		expect(loader.getExtensions().errors).toEqual([]);
		expect(
			loader.getExtensions().extensions.map((extension) => extension.path),
		).toEqual([...(keep ? [kept] : []), "<inline:builtin>"]);
		expect(loader.getExtensions().extensions.at(-1)?.flags.has("inline")).toBe(
			true,
		);
		expect(loader.getSkills().skills.map((skill) => skill.name)).toContain(
			"fixture",
		);
	},
);

function fixture() {
	directory = mkdtempSync(join(tmpdir(), "pi-extension-profile-test-"));
	const cwd = join(directory, "project with spaces");
	mkdirSync(cwd);
	const inventory = [
		{
			path: "/extensions/writer/index.ts",
			name: "writer",
			tools: ["write"],
			hooks: [],
			providers: [],
		},
		{
			path: "/extensions/observer/index.ts",
			name: "observer",
			tools: [],
			hooks: ["message_end"],
			providers: [],
		},
		{
			path: "/extensions/brand-new/index.ts",
			name: "brand-new",
			tools: [],
			hooks: [],
			providers: [],
		},
		{
			path: "/extensions/new-hook/index.ts",
			name: "new-hook",
			tools: [],
			hooks: ["future_event"],
			providers: [],
		},
	];
	const log = join(directory, "calls.jsonl");
	const worker = join(directory, "worker.mjs");
	writeFileSync(
		worker,
		`
		import fs from 'node:fs';
		import assert from 'node:assert/strict';
		const options = JSON.parse(fs.readFileSync(process.argv[2], 'utf8'));
		fs.appendFileSync(${JSON.stringify(log)}, JSON.stringify(options) + '\\n');
		assert.equal(process.cwd(), options.cwd);
		assert.equal(options.agentDir, options.sourceAgentDir);
		assert.equal(options.inspectExtensions, true);
		const inventory = ${JSON.stringify(inventory)};
		const mode = process.env.EXTENSION_PROFILE_TEST;
		const calls = fs.readFileSync(${JSON.stringify(log)}, 'utf8').trim().split('\\n').length;
		const omitted = inventory.find(entry => options.extensionPaths && !options.extensionPaths.includes(entry.path));
		const result = { kind: options.kind, cwd: options.cwd, version: 'test', provider: 'test', model: 'model', thinking: 'low', contextWindow: 100000,
			extensions: inventory.filter(entry => !options.extensionPaths || options.extensionPaths.includes(entry.path)),
			local: { system: 80, tools: omitted?.name === 'writer' ? 0 : 20, messages: 8, toolCount: 1 },
			...(options.kind === 'headless+0msg' ? {} : { inputTokens: omitted?.name === 'writer' ? 80 : omitted?.name === 'brand-new' ? 105 : 100 }),
		};
		if (mode === 'startup' && options.kind === 'headless+0msg') result.error = 'controlled startup failure';
		if (mode === 'inventory') delete result.extensions;
		if (mode === 'omitted' && omitted?.name === 'writer') result.error = 'controlled dependency failure';
		if (mode === 'mismatch' && omitted?.name === 'writer') result.model = 'changed';
		if (mode === 'baseline' && calls === 2) result.error = 'controlled baseline failure';
		if (mode === 'drift' && calls === inventory.length + 3) result.inputTokens++;
		if (mode === 'repeat' && calls === inventory.length + 3) result.error = 'controlled repeat failure';
		fs.writeFileSync(options.resultPath, JSON.stringify(result));
	`,
	);
	return {
		cwd,
		worker,
		inventory,
		calls: () =>
			readFileSync(log, "utf8")
				.trim()
				.split("\n")
				.map((line) => JSON.parse(line)),
	};
}

it("automatically measures every new extension, exactly one omission each, with pinned controls and cleanup", async () => {
	const { cwd, worker, inventory, calls } = fixture();
	const progress = vi.fn();
	const report = await measureExtensions(cwd, worker, progress);
	expect(report.stable).toBe(true);
	expect(report.inventory).toEqual(inventory);
	expect(report.runs.map((run) => run.extension.path).sort()).toEqual(
		inventory.map((entry) => entry.path).sort(),
	);
	const jobs = calls();
	expect(jobs).toHaveLength(inventory.length + 3);
	expect(jobs[0].extensionPaths).toBeUndefined();
	expect(jobs[1].extensionPaths).toEqual(inventory.map((entry) => entry.path));
	expect(jobs.at(-1).extensionPaths).toEqual(jobs[1].extensionPaths);
	for (const job of jobs.slice(1)) expect(job.selection.model).toBe("model");
	for (const [i, run] of report.runs.entries())
		expect(jobs[i + 2].extensionPaths).toEqual(
			inventory
				.filter((entry) => entry.path !== run.extension.path)
				.map((entry) => entry.path),
		);
	for (const job of jobs) expect(() => readFileSync(job.resultPath)).toThrow();
	expect(progress).toHaveBeenCalledTimes(inventory.length + 4);
	const text = formatExtensionOverview(report);
	expect(text).toMatch(/writer\s+\+20\s+\+0\s+\+20/);
	expect(text).toMatch(/brand-new\s+-5/);
	expect(text).toMatch(/observer\s+\+0/);
	expect(text).toContain("Deltas are not additive");
});

it.each(["startup", "inventory", "baseline"])(
	"blocks dependent requests after %s failure",
	async (mode) => {
		const { cwd, worker, calls } = fixture();
		vi.stubEnv("EXTENSION_PROFILE_TEST", mode);
		const report = await measureExtensions(cwd, worker);
		expect(report.error).toBeTruthy();
		expect(report.runs).toEqual([]);
		expect(calls()).toHaveLength(mode === "baseline" ? 2 : 1);
		expect(formatExtensionOverview(report)).toContain("failed");
	},
);

it.each(["omitted", "mismatch"])(
	"reports %s failure without zero savings or blocking independent omissions",
	async (mode) => {
		const { cwd, worker } = fixture();
		vi.stubEnv("EXTENSION_PROFILE_TEST", mode);
		const report = await measureExtensions(cwd, worker);
		expect(report.stable).toBe(true);
		expect(report.runs).toHaveLength(4);
		expect(formatExtensionOverview(report)).toMatch(/writer\s+FAIL/);
	},
);

it.each(["drift", "repeat"])(
	"withholds savings when final control has %s",
	async (mode) => {
		const { cwd, worker } = fixture();
		vi.stubEnv("EXTENSION_PROFILE_TEST", mode);
		const report = await measureExtensions(cwd, worker);
		expect(report.stable).toBe(false);
		expect(report.error).toContain("savings withheld");
		expect(formatExtensionOverview(report)).toMatch(/writer\s+n\/a/);
	},
);