pi-ext

import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import contextSizeExtension from "../index";

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: TestSession | undefined;
	let root: string | undefined;
	afterEach(async () => {
		await t?.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		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-"));
		const write = (dir: string, value: unknown) => {
			mkdirSync(dir, { recursive: true });
			writeFileSync(
				path.join(dir, "settings.json"),
				JSON.stringify({ "context-size": value }),
			);
		};
		if ("user" in input)
			write(path.join(root, ".test-home/.pi/agent"), input.user);
		if ("project" in input) write(path.join(root, ".pi"), input.project);

		const colors: unknown[] = [];
		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 createTestSession({
			cwd: root,
			env: input.env,
			extensionFactories: [contextSizeExtension, recorder],
		});
		if (input.untrusted) t.session.settingsManager.setProjectTrusted(false);

		const runner = t.session.extensionRunner as any;
		runner.mode = "tui";
		let percent = 0;
		runner.getContextUsageFn = () => ({
			tokens: percent * 1_000,
			contextWindow: 100_000,
			percent,
		});
		const notifyBefore = t.events.uiCallsFor("notify").length;
		await runner.emit({ type: "session_start", reason: "startup" });
		const warnings = t.events
			.uiCallsFor("notify")
			.slice(notifyBefore)
			.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 runner.emit({ type: "turn_end", turnIndex: 0, message: {} });
			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}`]);
		},
	);
});