pi-ext

import { describe, expect, it } from "vitest";
import {
	allocateAtoms,
	atomWeight,
	combineUsage,
	createSourceChunks,
	type DecayDetails,
	findPreviousDecayDetails,
	formatFileAppendix,
	isTransientFailure,
	type MemoryAtom,
	mergeFileLists,
	mergeMemory,
	missingChunks,
	parseDecayDetails,
	parseModelReference,
	type RecordChunk,
	reconcileRecords,
	renderFallbackSummary,
	type SourceChunk,
	validateRecordChunk,
} from "../core";

function user(text: string) {
	return { role: "user", content: [{ type: "text", text }], timestamp: 1 };
}

function record(
	chunk: number,
	of: number,
	key: string,
	overrides: Partial<RecordChunk["atoms"][number]> = {},
): RecordChunk {
	return {
		chunk,
		of,
		atoms: [
			{
				key,
				text: `fact ${key}`,
				kind: "context",
				status: "active",
				...overrides,
			},
		],
	};
}

function memory(overrides: Partial<MemoryAtom> = {}): MemoryAtom {
	return {
		key: "fact",
		text: "fact",
		kind: "context",
		firstSeen: 1,
		lastSeen: 1,
		seenCount: 1,
		status: "active",
		sourcePosition: 1,
		...overrides,
	};
}

function chunk(
	chunkNumber: number,
	of: number,
	source: SourceChunk["source"] = "history",
): SourceChunk {
	return {
		chunk: chunkNumber,
		of,
		source,
		firstPosition: chunkNumber,
		lastPosition: chunkNumber,
		text: "source",
		estimatedTokens: 1,
	};
}

const zeroUsage = {
	input: 0,
	output: 0,
	cacheRead: 0,
	cacheWrite: 0,
	totalTokens: 0,
	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};

describe("model references", () => {
	it("splits provider from model id", () => {
		expect(parseModelReference("provider/model/variant")).toEqual({
			provider: "provider",
			modelId: "model/variant",
		});
		expect(parseModelReference("broken")).toBeUndefined();
	});
});

describe("source chunking", () => {
	it("keeps legacy summary, history, and split prefix in distinct ordered chunks", () => {
		const chunks = createSourceChunks({
			legacySummary: "old checkpoint",
			messages: [user("history one"), user("history two")] as never,
			turnPrefixMessages: [user("split") as never],
			targetTokens: 100,
		});

		expect(chunks.map((item) => item.source)).toEqual([
			"legacy-summary",
			"history",
			"split-prefix",
		]);
		expect(chunks.map((item) => item.chunk)).toEqual([1, 2, 3]);
		expect(chunks.every((item) => item.of === 3)).toBe(true);
		expect(chunks[1].text).toContain("MESSAGE 1 SOURCE history");
		expect(chunks[1].text).toContain("MESSAGE 2 SOURCE history");
	});

	it("never splits one oversized message", () => {
		const chunks = createSourceChunks({
			messages: [user("x".repeat(1_000))] as never,
			turnPrefixMessages: [],
			targetTokens: 10,
		});
		expect(chunks).toHaveLength(1);
		expect(chunks[0].text).toContain("x".repeat(1_000));
	});

	it("keeps short exchanges separate and groups consecutive user inputs", () => {
		const chunks = createSourceChunks({
			messages: [
				user("first"),
				user("clarification"),
				{ role: "assistant", content: [{ type: "text", text: "done" }] },
				user("second"),
			] as never,
			turnPrefixMessages: [],
			targetTokens: 10_000,
		});
		expect(chunks).toHaveLength(2);
		expect(chunks[0].text).toContain("clarification");
		expect(chunks[0].text).toContain("done");
		expect(chunks[1].text).toContain("second");
		expect(chunks[1].text).not.toContain("first");
	});

	it("keeps multi-tool batches intact and separates steered requests", () => {
		const chunks = createSourceChunks({
			messages: [
				user("investigate"),
				{
					role: "assistant",
					content: [
						{ type: "text", text: "checking" },
						{ type: "toolCall", id: "one", name: "read", arguments: {} },
						{ type: "toolCall", id: "two", name: "read", arguments: {} },
					],
				},
				{
					role: "toolResult",
					toolCallId: "one",
					toolName: "read",
					content: [{ type: "text", text: "a".repeat(800) }],
				},
				{
					role: "toolResult",
					toolCallId: "two",
					toolName: "read",
					content: [{ type: "text", text: "b".repeat(800) }],
				},
				{
					role: "assistant",
					content: [{ type: "text", text: "second result is stale" }],
				},
				user("actually, stop"),
			] as never,
			turnPrefixMessages: [],
			targetTokens: 100,
		});
		expect(chunks.at(-2)?.text).toContain("b".repeat(800));
		expect(chunks.at(-2)?.text).toContain("a".repeat(800));
		expect(chunks.at(-2)?.text).toContain("second result is stale");
		expect(chunks.at(-1)?.text).toContain("actually, stop");
	});
});

describe("classifier records", () => {
	it("accepts more than 16 atoms without dropping facts", () => {
		const many = record(1, 1, "fact-0");
		many.atoms = Array.from({ length: 20 }, (_, index) => ({
			...many.atoms[0],
			key: `fact-${index}`,
		}));
		expect(validateRecordChunk(many)?.atoms).toHaveLength(20);
		expect(
			mergeMemory([], reconcileRecords([many], 1), [chunk(1, 1)]),
		).toHaveLength(20);
	});

	it("strictly validates schema and rejects unknown fields", () => {
		expect(validateRecordChunk(record(1, 1, "one"))).toBeDefined();
		expect(
			validateRecordChunk({ ...record(1, 1, "one"), extra: true }),
		).toBeUndefined();
		expect(
			validateRecordChunk({
				...record(1, 1, "one"),
				atoms: [{ ...record(1, 1, "one").atoms[0], extra: true }],
			}),
		).toBeUndefined();
		expect(
			validateRecordChunk({
				...record(1, 1, "one"),
				atoms: [{ ...record(1, 1, "one").atoms[0], sourcePosition: 99 }],
			}),
		).toBeUndefined();
	});

	it("sorts complete out-of-order records", () => {
		expect(
			reconcileRecords([record(2, 2, "two"), record(1, 2, "one")], 2),
		).toEqual([record(1, 2, "one"), record(2, 2, "two")]);
	});

	it("tolerates exact duplicates and reports missing chunks", () => {
		expect(
			reconcileRecords(
				[record(1, 2, "one"), record(1, 2, "one"), record(2, 2, "two")],
				2,
			),
		).toHaveLength(2);
		expect(missingChunks([record(2, 3, "two")], 3)).toEqual([1, 3]);
		expect(missingChunks([record(1, 1, "one")], 1)).toEqual([]);
	});

	it("rejects conflicts, missing chunks, wrong totals, and invalid indexes", () => {
		expect(() =>
			reconcileRecords([record(1, 2, "one"), record(1, 2, "again")], 2),
		).toThrow("reclassified chunk 1");
		expect(() => reconcileRecords([record(1, 2, "one")], 2)).toThrow(
			"omitted chunks: 2",
		);
		expect(() => reconcileRecords([record(1, 3, "one")], 2)).toThrow(
			"reported 3 chunks, expected 2",
		);
		expect(() => reconcileRecords([record(3, 2, "three")], 2)).toThrow(
			"invalid chunk 3",
		);
	});
});

describe("structured memory", () => {
	it("migrates legacy summary at epoch zero and counts semantic recurrence", () => {
		const records = [
			record(1, 3, "same", { text: "old wording" }),
			record(2, 3, "same", { text: "new wording" }),
			record(3, 3, "same", { text: "newest wording", kind: "decision" }),
		];
		const atoms = mergeMemory([], records, [
			chunk(1, 3, "legacy-summary"),
			chunk(2, 3),
			chunk(3, 3),
		]);
		expect(atoms).toEqual([
			expect.objectContaining({
				key: "same",
				text: "newest wording",
				kind: "decision",
				firstSeen: 0,
				lastSeen: 2,
				seenCount: 3,
				sourcePosition: 3,
			}),
		]);
	});

	it("does not count duplicate atoms inside one chunk as recurrence", () => {
		const duplicated = record(1, 1, "same");
		duplicated.atoms.push({ ...duplicated.atoms[0], text: "latest" });
		const atoms = mergeMemory([], [duplicated], [chunk(1, 1)]);
		expect(atoms[0]).toEqual(
			expect.objectContaining({ seenCount: 1, text: "latest" }),
		);
	});

	it("updates prior atoms and lets latest current status supersede", () => {
		const atoms = mergeMemory(
			[memory({ key: " Rule ", seenCount: 2, lastSeen: 4 })],
			[record(1, 1, "rule", { status: "superseded", text: "replaced" })],
			[chunk(1, 1)],
		);
		expect(atoms[0]).toEqual(
			expect.objectContaining({
				key: "rule",
				text: "replaced",
				seenCount: 3,
				lastSeen: 5,
				status: "superseded",
			}),
		);
	});

	it("weights recency and logarithmically bounded recurrence", () => {
		const newest = memory({ lastSeen: 4, seenCount: 1 });
		const oldRecurring = memory({ lastSeen: 1, seenCount: 16 });
		expect(atomWeight(newest, 4)).toBe(1);
		expect(atomWeight(oldRecurring, 4)).toBe(0.75);
	});

	it("records recency inside the current discarded history", () => {
		const atoms = mergeMemory(
			[],
			[record(1, 2, "old"), record(2, 2, "new")],
			[chunk(1, 2), chunk(2, 2)],
		);
		expect(atoms.map((atom) => [atom.key, atom.lastSeen])).toEqual([
			["old", 1],
			["new", 2],
		]);
		const allocated = new Map(
			allocateAtoms(atoms, 200).map((atom) => [atom.key, atom.allocation]),
		);
		expect(allocated.get("new")).toBeGreaterThan(allocated.get("old") ?? 0);
	});

	it("always allocates protected active atoms and favors newer ordinary atoms", () => {
		const allocated = allocateAtoms(
			[
				memory({ key: "law", kind: "constraint", lastSeen: 1 }),
				memory({ key: "old", lastSeen: 1 }),
				memory({ key: "new", lastSeen: 3 }),
				memory({ key: "gone", status: "superseded", lastSeen: 4 }),
			],
			200,
		);
		const byKey = new Map(allocated.map((atom) => [atom.key, atom]));
		expect(byKey.get("law")?.protected).toBe(true);
		expect(byKey.get("law")?.allocation).toBeGreaterThanOrEqual(24);
		expect(byKey.get("new")?.allocation).toBeGreaterThan(
			byKey.get("old")?.allocation ?? 0,
		);
		expect(byKey.has("gone")).toBe(false);
	});

	it("reads only valid details from the latest compaction boundary", () => {
		const details: DecayDetails = {
			decay: {
				version: 1,
				model: "p/m",
				elapsedMs: 10,
				chunkCount: 1,
				keepRecentTokens: 20_000,
				atoms: [memory()],
				readFiles: ["read"],
				modifiedFiles: ["write"],
			},
		};
		expect(parseDecayDetails(details)).toEqual(details);
		expect(
			findPreviousDecayDetails([
				{ type: "compaction", details, id: "1", parentId: null } as never,
			]),
		).toEqual(details.decay);
		expect(
			findPreviousDecayDetails([
				{ type: "compaction", details, id: "1", parentId: null } as never,
				{ type: "compaction", details: {}, id: "2", parentId: "1" } as never,
			]),
		).toBeUndefined();
	});
});

describe("deterministic appendices and usage", () => {
	it("merges files, moves later modifications out of read-only, and sorts", () => {
		const files = mergeFileLists(
			{ readFiles: ["z", "both"], modifiedFiles: ["old"] },
			{
				read: new Set(["a"]),
				written: new Set(["both"]),
				edited: new Set(["new"]),
			},
		);
		expect(files).toEqual({
			readFiles: ["a", "z"],
			modifiedFiles: ["both", "new", "old"],
		});
		expect(formatFileAppendix(files.readFiles, files.modifiedFiles)).toBe(
			"\n\n<read-files>\na\nz\n</read-files>\n\n<modified-files>\nboth\nnew\nold\n</modified-files>",
		);
	});

	it("combines every usage and cost field", () => {
		const combined = combineUsage(
			{ ...zeroUsage, input: 2, totalTokens: 2, reasoning: 1 },
			{
				...zeroUsage,
				output: 3,
				totalTokens: 3,
				cacheWrite1h: 2,
				cost: { ...zeroUsage.cost, output: 0.2, total: 0.2 },
			},
		);
		expect(combined).toEqual(
			expect.objectContaining({
				input: 2,
				output: 3,
				totalTokens: 5,
				reasoning: 1,
				cacheWrite1h: 2,
				cost: expect.objectContaining({ output: 0.2, total: 0.2 }),
			}),
		);
	});
});

describe("degraded summarizer", () => {
	function atom(overrides: Partial<MemoryAtom>): MemoryAtom {
		return {
			key: "k",
			text: "t",
			kind: "context",
			firstSeen: 1,
			lastSeen: 1,
			seenCount: 1,
			status: "active",
			sourcePosition: 0,
			...overrides,
		};
	}

	it("renders every heading and only classified facts", () => {
		const summary = renderFallbackSummary(
			allocateAtoms(
				[
					atom({ key: "goal", text: "Ship decay", kind: "goal" }),
					atom({ key: "blocked", text: "Waiting on auth", kind: "blocker" }),
					atom({ key: "path", text: "src/main.ts", kind: "exact" }),
					atom({ key: "dropped", text: "Old idea", status: "superseded" }),
				],
				4_000,
			),
		);
		expect(summary).toContain("Decay summarizer unavailable");
		for (const heading of [
			"## Goal",
			"## Constraints & Preferences",
			"## Progress",
			"### Done",
			"### In Progress",
			"### Blocked",
			"## Key Decisions",
			"## Next Steps",
			"## Critical Context",
		])
			expect(summary).toContain(heading);
		expect(summary).toContain("- Ship decay");
		expect(summary).toContain("- Waiting on auth");
		expect(summary).toContain("- src/main.ts");
		expect(summary).not.toContain("Old idea");
	});

	it("classifies transient failures without treating aborts as retryable", () => {
		for (const message of [
			"429 rate limit exceeded",
			"upstream returned 503",
			"socket hang up",
			"fetch failed",
			"Decay classifier stream ended without completion",
		])
			expect(isTransientFailure(new Error(message))).toBe(true);
		const abort = new Error("The operation was aborted");
		abort.name = "AbortError";
		expect(isTransientFailure(abort)).toBe(false);
		expect(
			isTransientFailure(new Error("Decay classifier called unknown tool: x")),
		).toBe(false);
		expect(
			isTransientFailure(new Error("Decay model authentication failed")),
		).toBe(false);
	});
});