pi-ext

import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";

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

type Case = {
	user?: unknown;
	project?: unknown;
	untrusted?: boolean;
	env?: Record<string, string>;
};

const DEFAULTS_SUFFIX = "; using defaults 70% and 85%";

describe("context-size threshold resolution", () => {
	let t: PiSession | undefined;
	let root: string | undefined;
	afterEach(async () => {
		await t?.dispose();
		t = undefined;
		if (root) rmSync(root, { recursive: true, force: true });
		root = undefined;
	});

	/** Start a session and return the footer tone for each probed percent. */
	async function start(input: Case, probes: number[]) {
		root = mkdtempSync(path.join(tmpdir(), "pi-ext-context-size-settings-"));
		if ("project" in input) {
			mkdirSync(path.join(root, ".pi"), { recursive: true });
			writeFileSync(
				path.join(root, ".pi", "settings.json"),
				JSON.stringify({ "context-size": input.project }),
			);
		}

		const colors: unknown[] = [];
		// An in-process factory shares the runner event bus with the path-loaded
		// extension to record `footer:segment` events.
		const recorder = (pi: ExtensionAPI) => {
			pi.events.on("footer:segment", (segment: any) => {
				if (segment.id !== "context-size" || segment.text === undefined) return;
				colors.push(segment.color);
				segment.accept();
			});
		};
		t = await createPiSession({
			cwd: root,
			extensionPaths: [extensionPath],
			extensionFactories: [recorder],
			...(input.env ? { env: input.env } : {}),
			...("user" in input
				? { userSettings: { "context-size": input.user } }
				: {}),
		});
		if (input.untrusted) t.session.settingsManager.setProjectTrusted(false);

		// Install a recording UI so `hasUI` is true and `mode` is tui, capturing
		// `ctx.ui.notify` calls.
		const uiCalls: { args: unknown[] }[] = [];
		t.session.extensionRunner.setUIContext(
			{
				notify: (...args: unknown[]) => uiCalls.push({ args }),
				setStatus: () => {},
				theme: { fg: (_color: unknown, text: string) => text },
			} as never,
			"tui",
		);

		const runner = t.session.extensionRunner as any;
		let percent = 0;
		runner.getContextUsageFn = () => ({
			tokens: percent * 1_000,
			contextWindow: 100_000,
			percent,
		});
		await runner.emit({ type: "session_start", reason: "startup" });
		const warnings = uiCalls
			.filter((call) => call.args[1] === "warning")
			.map((call) => call.args[0]);

		const tones: unknown[] = [];
		for (const probe of probes) {
			percent = probe;
			colors.length = 0;
			await t.session.prompt("Probe the configured context thresholds.");
			tones.push(colors.at(-1));
		}
		return { tones, warnings };
	}

	const DEFAULT_TONES = ["muted", "warning", "warning", "error"];
	const DEFAULT_PROBES = [69, 70, 84, 85];

	it("keeps 70/85 when nothing is configured", async () => {
		const { tones, warnings } = await start({}, DEFAULT_PROBES);
		expect(tones).toEqual(DEFAULT_TONES);
		expect(warnings).toEqual([]);
	});

	it.each([
		[
			"user values",
			{ user: { warningPercent: 50, errorPercent: 60 } },
			[49, 50, 59, 60],
		],
		[
			"env over user settings",
			{
				user: { warningPercent: 50, errorPercent: 60 },
				env: {
					PI_CONTEXT_SIZE_WARNING_PERCENT: "40",
					PI_CONTEXT_SIZE_ERROR_PERCENT: "90.5",
				},
			},
			[39, 40, 90, 90.5],
		],
		[
			"trusted project merged over user",
			{
				user: { warningPercent: 50, errorPercent: 60 },
				project: { errorPercent: 95 },
			},
			[49, 50, 94, 95],
		],
		[
			"untrusted project ignored",
			{
				user: { warningPercent: 50, errorPercent: 60 },
				project: { warningPercent: 10, errorPercent: 20 },
				untrusted: true,
			},
			[49, 50, 59, 60],
		],
	] satisfies [string, Case, number[]][])(
		"applies %s",
		async (_name, input, probes) => {
			const { tones, warnings } = await start(input, probes);
			expect(tones).toEqual(DEFAULT_TONES);
			expect(warnings).toEqual([]);
		},
	);

	it.each([
		[
			"wrong type",
			{ user: { warningPercent: "50" } },
			"Invalid context-size.warningPercent from user settings",
		],
		[
			"out of range",
			{ user: { warningPercent: 50 }, project: { errorPercent: 101 } },
			"Invalid context-size.errorPercent from project settings: must be greater than 0 and at most 100",
		],
		[
			"zero",
			{ user: { warningPercent: 0, errorPercent: 60 } },
			"Invalid context-size.warningPercent from user settings: must be greater than 0 and at most 100",
		],
		[
			"non-numeric env",
			{
				user: { warningPercent: 50, errorPercent: 60 },
				env: { PI_CONTEXT_SIZE_ERROR_PERCENT: "high" },
			},
			"Invalid context-size.errorPercent from environment variable PI_CONTEXT_SIZE_ERROR_PERCENT",
		],
		[
			"warning not below error",
			{ user: { warningPercent: 60, errorPercent: 60 } },
			"context-size.warningPercent (60) must be below context-size.errorPercent (60)",
		],
	] satisfies [string, Case, string][])(
		"rejects %s, warns once and uses both defaults",
		async (_name, input, error) => {
			const { tones, warnings } = await start(input, DEFAULT_PROBES);
			expect(tones).toEqual(DEFAULT_TONES);
			expect(warnings).toEqual([`${error}${DEFAULTS_SUFFIX}`]);
		},
	);
});