pi-ext

import { describe, expect, it } from "vitest";
import type { AllocatedAtom, SourceChunk } from "../core";
import {
	buildClassifierContext,
	buildSummarizerContext,
	CLASSIFIER_TOOL,
	MAX_PRIOR_CATALOG_TOKENS,
} from "../prompts";

const chunk: SourceChunk = {
	chunk: 1,
	of: 1,
	source: "history",
	firstPosition: 0,
	lastPosition: 0,
	text: "untrusted source",
	estimatedTokens: 4,
};

const atom: AllocatedAtom = {
	key: "decision",
	text: "Use stdlib",
	kind: "decision",
	firstSeen: 1,
	lastSeen: 2,
	seenCount: 2,
	status: "active",
	sourcePosition: 3,
	allocation: 40,
	weight: 2,
	protected: true,
	recurring: true,
};

describe("Decay model contracts", () => {
	it("offers exactly one prefer-strict constrained classifier tool", () => {
		const context = buildClassifierContext([chunk], [atom], "focus here");
		expect(context.tools).toEqual([CLASSIFIER_TOOL]);
		expect(CLASSIFIER_TOOL.name).toBe("record_chunk");
		expect(CLASSIFIER_TOOL.constrainedSampling).toEqual({
			type: "json_schema",
			strict: "prefer",
		});
		expect(
			(CLASSIFIER_TOOL.parameters as { additionalProperties?: unknown })
				.additionalProperties,
		).toBe(false);
		expect(JSON.stringify(CLASSIFIER_TOOL.parameters)).not.toContain(
			"sourcePosition",
		);
		expect(context.systemPrompt).toContain("Never increase recurrence");
		expect(context.systemPrompt).toContain("reuse that catalog key exactly");
		const payload = JSON.parse(
			(context.messages[0]!.content as Array<{ text: string }>)[0]!.text,
		);
		expect(payload).toEqual({
			oneRunFocus: "focus here",
			priorCatalog: [
				{
					key: "decision",
					text: "Use stdlib",
					kind: "decision",
					status: "active",
				},
			],
			chunks: [expect.objectContaining({ chunk: 1, text: "untrusted source" })],
		});
	});

	it("uses an evidence-relevant token budget for the prior catalog", () => {
		const prior = Array.from({ length: 200 }, (_, index) => ({
			...atom,
			key: `key-${index}`,
			text: `fact-${index}-${"x".repeat(600)}`,
			lastSeen: index,
		}));
		const context = buildClassifierContext(
			[{ ...chunk, text: "key-0 needs identity continuity" }],
			prior,
		);
		const payload = JSON.parse(
			(context.messages[0]!.content as Array<{ text: string }>)[0]!.text,
		);
		expect(
			Math.ceil(JSON.stringify(payload.priorCatalog).length / 4),
		).toBeLessThanOrEqual(MAX_PRIOR_CATALOG_TOKENS);
		expect(payload.priorCatalog[0].key).toBe("key-0");
	});

	it("counts JSON brackets and delimiters for tiny catalog entries", () => {
		const context = buildClassifierContext(
			[chunk],
			Array.from({ length: 200 }, (_, index) => ({
				...atom,
				key: `k${index}`,
				text: "xx",
				kind: "context" as const,
			})),
		);
		const payload = JSON.parse(
			(context.messages[0]!.content as Array<{ text: string }>)[0]!.text,
		);
		const serialized = JSON.stringify(payload.priorCatalog);
		expect(payload.priorCatalog).toHaveLength(132);
		expect(Math.ceil(serialized.length / 4)).toBeLessThanOrEqual(
			MAX_PRIOR_CATALOG_TOKENS,
		);
	});

	it("includes an exact serialized catalog boundary but not its next entry", () => {
		const boundary = {
			key: "boundary",
			text: "",
			kind: "decision" as const,
			status: "active" as const,
		};
		const text = "x".repeat(8_190 - JSON.stringify(boundary).length);
		const context = buildClassifierContext(
			[chunk],
			[
				{ ...atom, ...boundary, text, lastSeen: 2 },
				{ ...atom, key: "next", text: "x", kind: "context", lastSeen: 1 },
			],
		);
		const payload = JSON.parse(
			(context.messages[0]!.content as Array<{ text: string }>)[0]!.text,
		);
		expect(payload.priorCatalog).toEqual([{ ...boundary, text }]);
		expect(Math.ceil(JSON.stringify(payload.priorCatalog).length / 4)).toBe(
			MAX_PRIOR_CATALOG_TOKENS,
		);
	});

	it("retains lexical identity candidates including tail evidence and superseded rules", () => {
		const context = buildClassifierContext(
			[
				{
					...chunk,
					text: `${"noise ".repeat(4_096)}release ci tag retry semantics`,
				},
			],
			[
				{
					...atom,
					key: "release-tag-failure-retry-semantics",
					text: "Retry release tag after failure",
					kind: "decision",
					status: "superseded",
					lastSeen: 1,
				},
				{
					...atom,
					key: "release-ci-tag-retry-semantics",
					text: "Do not retry CI release tag",
					kind: "decision",
					status: "active",
					lastSeen: 2,
				},
			],
		);
		const payload = JSON.parse(
			(context.messages[0]!.content as Array<{ text: string }>)[0]!.text,
		);
		expect(
			payload.priorCatalog.map((entry: { key: string }) => entry.key),
		).toEqual([
			"release-ci-tag-retry-semantics",
			"release-tag-failure-retry-semantics",
		]);
		expect(context.systemPrompt).toContain("including for paraphrases");
		expect(context.systemPrompt).toContain(
			"explicitly replaces or invalidates",
		);
	});

	it("keeps permanent policy in system text and one-run focus in untrusted payload", () => {
		const context = buildSummarizerContext([atom], "only this run");
		expect(context.systemPrompt).toContain("Local code has already decided");
		expect(context.systemPrompt).not.toContain("only this run");
		const text = (context.messages[0]!.content as Array<{ text: string }>)[0]!
			.text;
		expect(text).toContain('"oneRunFocus":"only this run"');
		expect(text).toContain('"allocation":40');
		expect(text).toContain("## Critical Context");
	});
});