pi-ext

import {
	fauxProvider,
	InMemoryCredentialStore,
	InMemoryModelsStore,
} from "@earendil-works/pi-ai";
import {
	createAgentSession,
	DefaultResourceLoader,
	estimateTokens,
	ModelRuntime,
	type SessionBeforeCompactEvent,
	SessionManager,
	SettingsManager,
	sessionEntryToContextMessages,
} from "@earendil-works/pi-coding-agent";
import { describe, expect, it } from "vitest";
import {
	type DecayDetails,
	estimateKeptTokens,
	parseDecayDetails,
} from "../core";

const settings = {
	enabled: true,
	reserveTokens: 16_384,
	keepRecentTokens: 20_000,
};

function user(tokens: number, marker: string) {
	return {
		role: "user" as const,
		content: [
			{
				type: "text" as const,
				text: marker + "x".repeat(tokens * 4 - marker.length),
			},
		],
		timestamp: 1,
	};
}

function assistant(tokens: number, marker: string) {
	return {
		role: "assistant" as const,
		content: [
			{
				type: "text" as const,
				text: marker + "x".repeat(tokens * 4 - marker.length),
			},
		],
		api: "test",
		provider: "test",
		model: "test",
		usage: {
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason: "stop" as const,
		timestamp: 1,
	};
}

async function prepare(manager: SessionManager) {
	const settingsManager = SettingsManager.inMemory({ compaction: settings });
	const faux = fauxProvider({
		models: [{ id: "retention", contextWindow: 128_000, maxTokens: 16_384 }],
	});
	const modelRuntime = await ModelRuntime.create({
		credentials: new InMemoryCredentialStore(),
		modelsPath: null,
		modelsStore: new InMemoryModelsStore(),
		refreshOnCreate: false,
	});
	modelRuntime.registerNativeProvider(faux.provider);

	let captured: SessionBeforeCompactEvent | undefined;
	const resourceLoader = new DefaultResourceLoader({
		cwd: process.cwd(),
		agentDir: process.cwd(),
		settingsManager,
		noExtensions: true,
		noSkills: true,
		noPromptTemplates: true,
		noThemes: true,
		noContextFiles: true,
		extensionFactories: [
			(pi) => {
				pi.on("session_before_compact", (event) => {
					captured = event;
					return { cancel: true };
				});
			},
		],
	});
	await resourceLoader.reload();
	const { session } = await createAgentSession({
		model: faux.getModel("retention"),
		modelRuntime,
		noTools: "all",
		resourceLoader,
		sessionManager: manager,
		settingsManager,
	});
	try {
		await expect(session.compact()).rejects.toThrow("Compaction cancelled");
	} finally {
		session.dispose();
	}

	expect(faux.state.callCount).toBe(0);
	if (!captured) throw new Error("session_before_compact was not captured");
	return captured;
}

function projectedSuffix(event: SessionBeforeCompactEvent) {
	const boundary = event.branchEntries.findIndex(
		(entry) => entry.id === event.preparation.firstKeptEntryId,
	);
	return event.branchEntries
		.slice(boundary)
		.flatMap((entry) =>
			entry.type === "compaction" ? [] : sessionEntryToContextMessages(entry),
		);
}

function settleAndCompare(
	manager: SessionManager,
	event: SessionBeforeCompactEvent,
	summary = "new summary",
) {
	const suffix = projectedSuffix(event);
	const estimated = estimateKeptTokens(event);
	expect(estimated).toBe(
		suffix.reduce((total, message) => total + estimateTokens(message), 0),
	);
	manager.appendCompaction(
		summary,
		event.preparation.firstKeptEntryId,
		event.preparation.tokensBefore,
	);
	const rebuilt = manager.buildSessionContext().messages;
	expect(rebuilt[0]).toMatchObject({ role: "compactionSummary", summary });
	expect(rebuilt.slice(1)).toEqual(suffix);
	return { estimated, suffix, rebuilt };
}

describe("settled Decay retention accounting", () => {
	it("keeps normal-boundary message objects verbatim and may exceed the 20k target", async () => {
		const manager = SessionManager.inMemory();
		manager.appendMessage(user(12_000, "old"));
		const firstKept = user(12_000, "kept-a");
		const newest = user(12_000, "kept-b");
		manager.appendMessage(firstKept);
		manager.appendMessage(newest);
		const event = await prepare(manager);
		expect(event.preparation.isSplitTurn).toBe(false);

		const result = settleAndCompare(manager, event);
		expect(result.estimated).toBe(24_000);
		expect(result.estimated).toBeGreaterThan(settings.keepRecentTokens);
		expect(result.suffix).toHaveLength(2);
		expect(result.suffix[0]).toBe(firstKept);
		expect(result.suffix[1]).toBe(newest);
	});

	it("keeps split-turn suffix message objects verbatim", async () => {
		const manager = SessionManager.inMemory();
		manager.appendMessage(user(100, "request"));
		const firstKept = assistant(12_000, "kept-a");
		const newest = assistant(12_000, "kept-b");
		manager.appendMessage(firstKept);
		manager.appendMessage(newest);
		const event = await prepare(manager);
		expect(event.preparation.isSplitTurn).toBe(true);

		const result = settleAndCompare(manager, event);
		expect(result.estimated).toBe(24_000);
		expect(result.suffix[0]).toBe(firstKept);
		expect(result.suffix[1]).toBe(newest);
	});

	it("can stay below 20k at a tool-result cut and excludes metadata", async () => {
		const manager = SessionManager.inMemory();
		manager.appendMessage(user(1_000, "old"));
		manager.appendMessage({
			...assistant(1, "a"),
			content: [
				{ type: "toolCall", id: "call-1", name: "read", arguments: {} },
			],
		} as never);
		manager.appendMessage({
			role: "toolResult",
			toolCallId: "call-1",
			toolName: "read",
			content: [{ type: "text", text: "x".repeat(21_000 * 4) }],
			isError: false,
			timestamp: 1,
		});
		const metadataId = manager.appendCustomEntry("retention-test", {
			notContext: "x".repeat(40_000),
		});
		const newest = user(100, "new");
		manager.appendMessage(newest);
		const event = await prepare(manager);
		expect(event.preparation.firstKeptEntryId).toBe(metadataId);

		const result = settleAndCompare(manager, event);
		expect(result.estimated).toBe(100);
		expect(result.estimated).toBeLessThan(settings.keepRecentTokens);
		expect(result.suffix).toEqual([newest]);
	});

	it("drops an older compaction summary when the later boundary precedes it", async () => {
		const manager = SessionManager.inMemory();
		for (let index = 1; index <= 4; index++)
			manager.appendMessage(user(8_000, `turn-${index}`));
		const first = (await prepare(manager)).preparation;
		manager.appendCompaction(
			"s".repeat(4_000),
			first.firstKeptEntryId,
			first.tokensBefore,
		);
		manager.appendMessage(user(4_000, "later"));
		const secondEvent = await prepare(manager);
		const boundaryIndex = secondEvent.branchEntries.findIndex(
			(entry) => entry.id === secondEvent.preparation.firstKeptEntryId,
		);
		const olderCompactionIndex = secondEvent.branchEntries.findIndex(
			(entry) => entry.type === "compaction",
		);
		expect(boundaryIndex).toBeLessThan(olderCompactionIndex);

		const result = settleAndCompare(manager, secondEvent);
		expect(
			result.suffix.some((message) => message.role === "compactionSummary"),
		).toBe(false);
	});

	it("cuts before a repeated-compaction checkpoint and drops the older summary", async () => {
		const manager = SessionManager.inMemory();
		for (let index = 1; index <= 4; index++)
			manager.appendMessage(user(8_000, `turn-${index}`));
		const first = (await prepare(manager)).preparation;
		manager.appendCompaction(
			"older checkpoint",
			first.firstKeptEntryId,
			first.tokensBefore,
		);
		manager.appendMessage(user(12_000, "later-a"));
		manager.appendMessage(user(12_000, "later-b"));
		const secondEvent = await prepare(manager);
		const boundaryIndex = secondEvent.branchEntries.findIndex(
			(entry) => entry.id === secondEvent.preparation.firstKeptEntryId,
		);
		const olderCompactionIndex = secondEvent.branchEntries.findIndex(
			(entry) => entry.type === "compaction",
		);
		expect(boundaryIndex).toBeLessThan(olderCompactionIndex);

		const result = settleAndCompare(manager, secondEvent);
		expect(result.estimated).toBe(24_000);
		expect(
			result.suffix.some((message) => message.role === "compactionSummary"),
		).toBe(false);
	});

	it("returns undefined when the prepared boundary is absent", async () => {
		const manager = SessionManager.inMemory();
		manager.appendMessage(user(12_000, "old"));
		manager.appendMessage(user(12_000, "kept-a"));
		manager.appendMessage(user(12_000, "kept-b"));
		const preparedEvent = await prepare(manager);
		const event = {
			...preparedEvent,
			preparation: {
				...preparedEvent.preparation,
				firstKeptEntryId: "missing",
			},
		};
		expect(estimateKeptTokens(event)).toBeUndefined();
	});

	it("parses an optional valid estimate and rejects invalid present values", () => {
		const legacy: DecayDetails = {
			decay: {
				version: 1,
				model: "p/m",
				elapsedMs: 1,
				chunkCount: 1,
				keepRecentTokens: 20_000,
				atoms: [],
				readFiles: [],
				modifiedFiles: [],
			},
		};
		expect(parseDecayDetails(legacy)).toEqual(legacy);
		const withEstimate: DecayDetails = {
			decay: { ...legacy.decay, estimatedKeptTokens: 24_000 },
		};
		expect(parseDecayDetails(withEstimate)).toEqual(withEstimate);
		for (const invalid of [-1, 1.5, Number.NaN, "20000", null]) {
			expect(
				parseDecayDetails({
					decay: { ...legacy.decay, estimatedKeptTokens: invalid },
				}),
			).toBeUndefined();
		}
	});
});

// These are pure accounting tests. They exercise Pi's real preparation and session
// rebuilding APIs, but intentionally do not invoke Decay's extension hook or any model.