pi-auto-reinject/test/kept-window.test.ts

import type { SessionEntry } from "@earendil-works/pi-coding-agent";
import { describe, expect, it } from "vitest";
import {
	filterSkillsNeedingReinject,
	filterSkillsNeedingReinjectByKept,
	getKeptEntries,
	skillsPresentInKeptWindow,
} from "../src/kept";
import type { TrackedSkill } from "../src/state";

const ts = "2024-12-03T14:00:00.000Z";

function userMessage(id: string, content: string): SessionEntry {
	return {
		type: "message",
		id,
		parentId: null,
		timestamp: ts,
		message: { role: "user", content },
	} as SessionEntry;
}

function assistantMessage(id: string, content: string): SessionEntry {
	return {
		type: "message",
		id,
		parentId: null,
		timestamp: ts,
		message: { role: "assistant", content: [{ type: "text", text: content }] },
	} as SessionEntry;
}

function compactionEntry(id: string, firstKeptEntryId: string): SessionEntry {
	return {
		type: "compaction",
		id,
		parentId: null,
		timestamp: ts,
		summary: "summary",
		firstKeptEntryId,
		tokensBefore: 1000,
	} as SessionEntry;
}

function reinjectCustomMessage(id: string, content: string): SessionEntry {
	return {
		type: "custom_message",
		id,
		parentId: null,
		timestamp: ts,
		customType: "skill-reinject:inject",
		content,
		display: true,
	} as SessionEntry;
}

const skillBlock = (name: string) =>
	`<skill name="${name}" location="/skills/${name}/SKILL.md">\nbody\n</skill>`;

const trackedSkills: TrackedSkill[] = [
	{
		name: "alpha",
		filePath: "/skills/alpha/SKILL.md",
		baseDir: "/skills/alpha",
		firstSeenAt: 1,
		lastSeenAt: 1,
		sources: ["slash"],
	},
	{
		name: "beta",
		filePath: "/skills/beta/SKILL.md",
		baseDir: "/skills/beta",
		firstSeenAt: 2,
		lastSeenAt: 2,
		sources: ["skill-block"],
	},
];

describe("getKeptEntries", () => {
	const branch: SessionEntry[] = [
		userMessage("e1", "old"),
		userMessage("e2", "kept start"),
		assistantMessage("e3", "reply"),
		compactionEntry("e4", "e2"),
		userMessage("e5", "after compact"),
	];

	it("slices from firstKeptEntryId through tail", () => {
		expect(getKeptEntries(branch, "e2").map((entry) => entry.id)).toEqual(["e2", "e3", "e4", "e5"]);
	});

	it("returns empty when firstKeptEntryId is missing", () => {
		expect(getKeptEntries(branch, "missing")).toEqual([]);
		expect(getKeptEntries([], "e1")).toEqual([]);
	});
});

describe("skillsPresentInKeptWindow", () => {
	it("detects skill blocks in kept user messages", () => {
		const kept = [userMessage("k1", skillBlock("alpha")), assistantMessage("k2", "ignored")];
		expect(skillsPresentInKeptWindow(kept, ["alpha", "beta"])).toEqual(new Set(["alpha"]));
	});

	it("returns empty when no tracked skills or no blocks", () => {
		expect(skillsPresentInKeptWindow([userMessage("k1", "plain")], ["alpha"])).toEqual(new Set());
		expect(skillsPresentInKeptWindow([userMessage("k1", skillBlock("alpha"))], [])).toEqual(new Set());
		expect(skillsPresentInKeptWindow([], ["alpha", "beta"])).toEqual(new Set());
	});

	it("detects skill blocks in skill-reinject:inject custom messages", () => {
		const kept = [reinjectCustomMessage("inj1", skillBlock("beta"))];
		expect(skillsPresentInKeptWindow(kept, ["alpha", "beta"])).toEqual(new Set(["beta"]));
	});

	it("ignores reinject custom outside kept slice when slicing branch", () => {
		const branch: SessionEntry[] = [
			reinjectCustomMessage("inj0", skillBlock("alpha")),
			userMessage("e2", "kept start"),
			compactionEntry("e4", "e2"),
		];
		const kept = getKeptEntries(branch, "e2");
		expect(skillsPresentInKeptWindow(kept, ["alpha", "beta"])).toEqual(new Set());
		expect(skillsPresentInKeptWindow(kept, ["alpha"])).toEqual(new Set());
	});

	it("counts reinject custom inside kept slice", () => {
		const branch: SessionEntry[] = [
			userMessage("e2", "kept start"),
			reinjectCustomMessage("inj1", skillBlock("alpha")),
			compactionEntry("e4", "e2"),
		];
		const kept = getKeptEntries(branch, "e2");
		expect(skillsPresentInKeptWindow(kept, ["alpha", "beta"])).toEqual(new Set(["alpha"]));
	});
});

describe("filterSkillsNeedingReinjectByKept", () => {
	it("returns tracked skills absent from kept window regardless of registration", () => {
		const keptPresent = new Set(["alpha"]);
		expect(filterSkillsNeedingReinjectByKept(trackedSkills, keptPresent)).toEqual(["beta"]);
	});

	it("preserves tracked order including unregistered names", () => {
		const keptPresent = new Set<string>();
		expect(filterSkillsNeedingReinjectByKept(trackedSkills, keptPresent)).toEqual(["alpha", "beta"]);
	});

	it("returns empty when all tracked skills are in kept window", () => {
		const keptPresent = new Set(["alpha", "beta"]);
		expect(filterSkillsNeedingReinjectByKept(trackedSkills, keptPresent)).toEqual([]);
	});
});

describe("filterSkillsNeedingReinject", () => {
	it("returns registered tracked skills absent from kept window", () => {
		const keptPresent = new Set(["alpha"]);
		const registered = new Set(["alpha", "beta"]);
		expect(filterSkillsNeedingReinject(trackedSkills, keptPresent, registered)).toEqual(["beta"]);
	});

	it("excludes unregistered skills and preserves tracked order", () => {
		const keptPresent = new Set<string>();
		const registered = new Set(["alpha", "beta"]);
		expect(filterSkillsNeedingReinject(trackedSkills, keptPresent, registered)).toEqual(["alpha", "beta"]);
		expect(filterSkillsNeedingReinject(trackedSkills, keptPresent, new Set(["beta"]))).toEqual(["beta"]);
	});

	it("returns empty when all tracked skills are in kept window", () => {
		const keptPresent = new Set(["alpha", "beta"]);
		const registered = new Set(["alpha", "beta"]);
		expect(filterSkillsNeedingReinject(trackedSkills, keptPresent, registered)).toEqual([]);
	});
});