pi-ext

import { describe, expect, it } from "vitest";
import type { AllocatedAtom, SourceChunk } from "../core";
import {
	buildClassifierContext,
	buildSummarizerContext,
	CLASSIFIER_TOOL,
	MAX_PRIOR_CATALOG_ATOMS,
	MAX_PRIOR_CATALOG_CHARS,
} 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.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("bounds the prior catalog while retaining recent active atoms", () => {
		const prior = Array.from({ length: 200 }, (_, index) => ({
			...atom,
			key: `key-${index}`,
			text: `fact-${index}-${"x".repeat(600)}`,
			lastSeen: index,
		}));
		const context = buildClassifierContext([chunk], prior);
		const payload = JSON.parse(
			(context.messages[0].content as Array<{ text: string }>)[0].text,
		);
		expect(payload.priorCatalog.length).toBeLessThanOrEqual(
			MAX_PRIOR_CATALOG_ATOMS,
		);
		expect(JSON.stringify(payload.priorCatalog).length).toBeLessThanOrEqual(
			MAX_PRIOR_CATALOG_CHARS,
		);
		expect(payload.priorCatalog[0].key).toBe("key-199");
	});

	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");
	});
});