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; }; 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}`]); }, ); });