pi-ext

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

const extensionPath = path.resolve(import.meta.dirname, "../index.ts");

interface NotifyCall {
	args: unknown[];
}

/**
 * Install a recording UI context so `hasUI` is true (so the extension sees an
 * interactive UI) and `ctx.ui.notify` calls are captured.
 */
function installRecordingUi(session: PiSession): NotifyCall[] {
	const calls: NotifyCall[] = [];
	session.session.extensionRunner.setUIContext(
		{ notify: (...args: unknown[]) => calls.push({ args }) } as never,
		"print",
	);
	return calls;
}

type Case = { user?: unknown; project?: unknown; untrusted?: boolean };

function compactionEvent(): SessionBeforeCompactEvent {
	return {
		type: "session_before_compact" as const,
		preparation: {
			firstKeptEntryId: "kept-entry",
			messagesToSummarize: [
				{
					role: "user",
					content: [{ type: "text", text: "old" }],
					timestamp: 1,
				},
			],
			turnPrefixMessages: [],
			isSplitTurn: false,
			tokensBefore: 1_000,
			fileOps: { read: new Set(), written: new Set(), edited: new Set() },
			settings: { enabled: true, reserveTokens: 16_384, keepRecentTokens: 0 },
		},
		branchEntries: [],
		reason: "manual" as const,
		willRetry: false,
		signal: new AbortController().signal,
	} as SessionBeforeCompactEvent;
}

describe("decay settings 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;
		vi.restoreAllMocks();
	});

	async function compact(input: Case) {
		root = mkdtempSync(path.join(tmpdir(), "pi-ext-decay-settings-"));
		if ("project" in input) {
			mkdirSync(path.join(root, ".pi"), { recursive: true });
			writeFileSync(
				path.join(root, ".pi", "settings.json"),
				JSON.stringify({ decay: input.project }),
			);
		}
		t = await createPiSession({
			cwd: root,
			extensionPaths: [extensionPath],
			...("user" in input ? { userSettings: { decay: input.user } } : {}),
		});
		if (input.untrusted) t.session.settingsManager.setProjectTrusted(false);
		const notifications = installRecordingUi(t);
		const faux = fauxProvider({
			provider: "openai",
			models: [
				"gpt-4o",
				"user",
				"project",
				"summarizer",
				"plain",
				"limited",
			].map((id) => ({
				id,
				reasoning: id !== "plain",
				maxTokens: 8_192,
			})),
		});
		const runner = t.session.extensionRunner;
		const registry = runner.createContext().modelRegistry;
		registry.registerProvider("openai", {
			api: faux.api,
			apiKey: "test-key",
			models: faux.models.map((model) => ({
				...model,
				thinkingLevelMap: {
					off: model.id === "limited" ? null : "off",
					minimal: model.id === "limited" ? null : "minimal",
					low: "low",
					medium: "medium",
					high: "high",
					xhigh: model.id === "limited" ? null : "xhigh",
					max: model.id === "limited" ? null : "max",
				},
			})),
			streamSimple: (model, context, options) =>
				faux.provider.streamSimple(model, context, options),
		});
		await registry.refresh({ allowNetwork: false });
		// The owned fixture's current model is its own faux model. Register and
		// select openai/gpt-4o so `ctx.model` matches.
		const current = registry.find("openai", "gpt-4o");
		if (!current) throw new Error("missing faux gpt-4o model");
		await t.session.setModel(current);
		const streamSimple = vi.spyOn(faux.provider, "streamSimple");
		faux.setResponses([
			fauxAssistantMessage(
				fauxToolCall("record_chunk", { chunk: 1, of: 1, atoms: [] }),
				{ stopReason: "toolUse" },
			),
			fauxAssistantMessage("## Goal\n\nDone."),
		]);
		const result = await runner.emit(compactionEvent());
		const warnings = notifications
			.filter((call) => call.args[1] === "warning")
			.map((call) => call.args[0]);
		return {
			compacted: result?.compaction !== undefined,
			models: streamSimple.mock.calls.map(([model]) => model.id),
			thinking: streamSimple.mock.calls.map(
				([, , options]) => options?.reasoning,
			),
			warnings,
		};
	}

	it.each([
		["omitted uses current", {}, ["gpt-4o", "gpt-4o"]],
		["user model", { user: { model: "openai/user" } }, ["user", "user"]],
		[
			"project model over user, user summarizer kept",
			{
				user: { model: "openai/user", summarizerModel: "openai/summarizer" },
				project: { model: "openai/project" },
			},
			["project", "summarizer"],
		],
		[
			"project null over user keeps explicit unset",
			{
				user: { model: "openai/user", summarizerModel: "openai/summarizer" },
				project: { model: null, summarizerModel: null },
			},
			["gpt-4o", "gpt-4o"],
		],
		[
			"untrusted project ignored",
			{
				user: { model: "openai/user" },
				project: { model: "openai/project" },
				untrusted: true,
			},
			["user", "user"],
		],
	] satisfies [string, Case, string[]][])(
		"%s",
		async (_name, input, expected) => {
			const { compacted, models, warnings } = await compact(input);
			expect(compacted).toBe(true);
			expect(models).toEqual(expected);
			expect(warnings).toEqual([]);
		},
	);

	it.each([
		["omitted keeps low", {}, ["low", "low"]],
		[
			"user thinking applies to both models",
			{
				user: {
					model: "openai/user",
					summarizerModel: "openai/summarizer",
					thinking: "high",
				},
			},
			["high", "high"],
		],
		[
			"trusted project thinking overrides user",
			{ user: { thinking: "high" }, project: { thinking: "medium" } },
			["medium", "medium"],
		],
		[
			"untrusted project thinking ignored",
			{
				user: { thinking: "high" },
				project: { thinking: "low" },
				untrusted: true,
			},
			["high", "high"],
		],
		[
			"unsupported maximum thinking clamps to high",
			{ user: { model: "openai/limited", thinking: "max" } },
			["high", "high"],
		],
		[
			"unsupported off thinking clamps to low",
			{ user: { model: "openai/limited", thinking: "off" } },
			["low", "low"],
		],
		[
			"non-thinking models receive no thinking option",
			{ user: { model: "openai/plain", thinking: "high" } },
			[undefined, undefined],
		],
	] satisfies [string, Case, (string | undefined)[]][])(
		"%s",
		async (_name, input, expected) => {
			const result = await compact(input);
			expect(result.compacted).toBe(true);
			expect(result.thinking).toEqual(expected);
			expect(result.warnings).toEqual([]);
		},
	);

	it.each(["off", "minimal", "low", "medium", "high", "xhigh", "max"])(
		"accepts Pi thinking level %s",
		async (thinking) => {
			const result = await compact({ user: { thinking } });
			expect(result.compacted).toBe(true);
			const expected = thinking === "off" ? undefined : thinking;
			expect(result.thinking).toEqual([expected, expected]);
			expect(result.warnings).toEqual([]);
		},
	);

	it.each(["invalid", "", "HIGH", null, 42, true, {}, []])(
		"rejects invalid thinking %j",
		async (thinking) => {
			const result = await compact({ user: { thinking } });
			expect(result.compacted).toBe(false);
			expect(result.models).toEqual([]);
			expect(result.warnings).toEqual([
				"Invalid decay.thinking from user settings: must be off, minimal, low, medium, high, xhigh, or max; using default compaction",
			]);
		},
	);

	it.each([
		[
			"invalid project thinking over valid user",
			{ user: { thinking: "high" }, project: { thinking: "invalid" } },
			"Invalid decay.thinking from project settings: must be off, minimal, low, medium, high, xhigh, or max; using default compaction",
		],
		[
			"numeric user model",
			{ user: { model: 42 } },
			"Decay model setting must be provider/model, null, or omitted; using default compaction",
		],
		[
			"blank project model over valid user",
			{ user: { model: "openai/user" }, project: { model: " " } },
			"Decay model setting must be provider/model, null, or omitted; using default compaction",
		],
		[
			"boolean project summarizer over valid user",
			{
				user: { summarizerModel: "openai/summarizer" },
				project: { summarizerModel: false },
			},
			"Decay summarizerModel setting must be provider/model, null, or omitted; using default compaction",
		],
	] satisfies [string, Case, string][])(
		"rejects %s without falling through",
		async (_name, input, message) => {
			const { compacted, models, warnings } = await compact(input);
			expect(compacted).toBe(false);
			expect(models).toEqual([]);
			expect(warnings).toEqual([message]);
		},
	);
});