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 { afterEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";

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

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

function compactionEvent() {
	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,
			previousSummary: undefined,
			fileOps: { read: new Set(), written: new Set(), edited: new Set() },
			settings: { enabled: true, reserveTokens: 16_384, keepRecentTokens: 0 },
		},
		branchEntries: [],
		customInstructions: undefined,
		reason: "manual" as const,
		willRetry: false,
		signal: new AbortController().signal,
	};
}

describe("decay model settings resolution", () => {
	let t: TestSession | undefined;
	let root: string | undefined;
	afterEach(() => {
		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-"));
		const write = (dir: string, value: unknown) => {
			mkdirSync(dir, { recursive: true });
			writeFileSync(
				path.join(dir, "settings.json"),
				JSON.stringify({ decay: 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);
		t = await createTestSession({ cwd: root, extensions: [extensionPath] });
		if (input.untrusted) t.session.settingsManager.setProjectTrusted(false);
		const faux = fauxProvider({
			provider: "openai",
			models: ["gpt-4o", "user", "project", "summarizer"].map((id) => ({
				id,
				maxTokens: 8_192,
			})),
		});
		const runner = t.session.extensionRunner;
		runner.createContext().modelRegistry.registerProvider(faux.provider);
		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 = t.events
			.uiCallsFor("notify")
			.filter((call) => call.args[1] === "warning")
			.map((call) => call.args[0]);
		return {
			compacted: result?.compaction !== undefined,
			models: streamSimple.mock.calls.map(([model]) => model.id),
			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([
		[
			"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]);
		},
	);
});