pi-ext

import {
	chmod,
	mkdtemp,
	readdir,
	readFile,
	rm,
	stat,
	writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, describe, expect, it } from "vitest";
import {
	FileSessionStore,
	KlausCacheCorruptionError,
	loadCheckpoint,
	loadCheckpoints,
	MemorySessionStore,
	type OpaqueEntry,
	StagingSessionStore,
	type StoreKey,
	saveCheckpoint,
} from "../src/cache";

const roots: string[] = [];
const key = { projectKey: "project", sessionId: "session" };

async function temporaryRoot(label: string): Promise<string> {
	const root = await mkdtemp(join(tmpdir(), `klaus-${label}-`));
	roots.push(root);
	return root;
}

async function onlyFile(root: string, suffix: string): Promise<string> {
	const files = (await readdir(root)).filter((file) => file.endsWith(suffix));
	expect(files).toHaveLength(1);
	const file = files[0];
	if (!file) throw new Error(`Missing ${suffix} fixture file.`);
	return join(root, file);
}

const checkpoint = (id: string, overrides = {}) => ({
	sdkSessionId: `sdk-${id}`,
	fingerprint: `fingerprint-${id}`,
	position: `position-${id}`,
	piLeafId: `leaf-${id}`,
	messageCount: 2,
	protocol: 1 as const,
	...overrides,
});

afterEach(async () => {
	await Promise.all(
		roots.splice(0).map((root) => rm(root, { recursive: true, force: true })),
	);
});

describe("Klaus in-memory opaque session store", () => {
	it("round-trips entries and ignores repeated UUIDs across and within batches", async () => {
		const store = new MemorySessionStore();
		await store.append(key, [
			{ type: "user", uuid: "one", opaque: { value: 1 } },
			{ type: "assistant", uuid: "one", ignored: true },
			{ type: "assistant", uuid: "two" },
		]);
		await store.append(key, [
			{ type: "user", uuid: "one", opaque: { value: 2 } },
		]);

		expect(await store.load(key)).toEqual([
			{ type: "user", uuid: "one", opaque: { value: 1 } },
			{ type: "assistant", uuid: "two" },
		]);
	});

	it("appends every UUID-less opaque entry", async () => {
		const store = new MemorySessionStore();
		const marker = { type: "mode", value: "same" };
		await store.append(key, [marker, marker]);
		await store.append(key, [marker]);

		expect(await store.load(key)).toEqual([marker, marker, marker]);
	});

	it("isolates project, session, and subpath key dimensions", async () => {
		const store = new MemorySessionStore();
		const keys: StoreKey[] = [
			key,
			{ ...key, projectKey: "other" },
			{ ...key, sessionId: "other" },
			{ ...key, subpath: "agent-a" },
			{ ...key, subpath: "agent-b" },
		];
		await Promise.all(
			keys.map((entryKey, index) =>
				store.append(entryKey, [{ type: "entry", index }]),
			),
		);

		for (const [index, entryKey] of keys.entries()) {
			expect(await store.load(entryKey)).toEqual([{ type: "entry", index }]);
		}
		expect(
			await store.load({ projectKey: "missing", sessionId: "missing" }),
		).toBeNull();
	});

	it("deep-clones appended and loaded values", async () => {
		const store = new MemorySessionStore();
		const source = { type: "entry", uuid: "one", nested: { value: 1 } };
		await store.append(key, [source]);
		source.nested.value = 2;
		const loaded = await store.load(key);
		if (!loaded) throw new Error("Missing memory entry.");
		const nested = loaded[0]?.nested as { value: number };
		nested.value = 3;

		expect(await store.load(key)).toEqual([
			{ type: "entry", uuid: "one", nested: { value: 1 } },
		]);
	});

	it.each([null, [], "entry", {}, { type: 1 }] as const)(
		"rejects invalid opaque entry %j",
		async (value) => {
			const store = new MemorySessionStore();
			await expect(
				store.append(key, [value as unknown as OpaqueEntry]),
			).rejects.toThrow("invalid opaque session entry");
			expect(await store.load(key)).toBeNull();
		},
	);

	it("reports no subkeys because opaque hashes are intentionally not reversible", async () => {
		const store = new MemorySessionStore();
		await store.append({ ...key, subpath: "agent" }, [{ type: "entry" }]);
		expect(await store.listSubkeys(key)).toEqual([]);
	});
});

describe("Klaus staging session store", () => {
	it("keeps query writes staged until commit", async () => {
		const target = new MemorySessionStore();
		const staging = new StagingSessionStore(target);
		await staging.append(key, [{ type: "user", uuid: "one" }]);

		expect(await target.load(key)).toBeNull();
		expect(await staging.load(key)).toEqual([{ type: "user", uuid: "one" }]);
		await staging.commit();
		expect(await target.load(key)).toEqual([{ type: "user", uuid: "one" }]);
	});

	it("combines committed and matching staged batches without leaking other keys", async () => {
		const target = new MemorySessionStore();
		await target.append(key, [{ type: "base", uuid: "base" }]);
		const staging = new StagingSessionStore(target);
		await staging.append(key, [{ type: "staged", uuid: "one" }]);
		await staging.append({ ...key, subpath: "other" }, [
			{ type: "other", uuid: "two" },
		]);

		expect(await staging.load(key)).toEqual([
			{ type: "base", uuid: "base" },
			{ type: "staged", uuid: "one" },
		]);
		expect(await staging.load({ ...key, subpath: "other" })).toEqual([
			{ type: "other", uuid: "two" },
		]);
	});

	it("deep-clones staged keys and entries", async () => {
		const target = new MemorySessionStore();
		const staging = new StagingSessionStore(target);
		const mutableKey: StoreKey = { ...key, subpath: "before" };
		const entry = { type: "entry", nested: { value: 1 } };
		await staging.append(mutableKey, [entry]);
		mutableKey.subpath = "after";
		entry.nested.value = 2;

		expect(await staging.load({ ...key, subpath: "before" })).toEqual([
			{ type: "entry", nested: { value: 1 } },
		]);
	});

	it("commits each batch once even when commit is repeated", async () => {
		const target = new MemorySessionStore();
		const staging = new StagingSessionStore(target);
		await staging.append(key, [{ type: "marker" }]);
		await staging.commit();
		await staging.commit();

		expect(await target.load(key)).toEqual([{ type: "marker" }]);
	});
});

describe("Klaus disk opaque session store", () => {
	it("creates a pinned manifest and round-trips opaque NDJSON", async () => {
		const parent = await temporaryRoot("cache");
		const root = join(parent, "sessions");
		const store = new FileSessionStore(root);
		await store.append(key, [
			{ type: "assistant", uuid: "one", payload: "opaque" },
		]);

		expect(await store.load(key)).toEqual([
			{ type: "assistant", uuid: "one", payload: "opaque" },
		]);
		expect(
			JSON.parse(await readFile(join(root, "manifest.json"), "utf8")),
		).toEqual({ schema: 1, sdk: "0.3.281" });
		const transcript = await onlyFile(root, ".ndjson");
		expect(await readFile(transcript, "utf8")).toBe(
			'{"type":"assistant","uuid":"one","payload":"opaque"}\n',
		);
	});

	it.runIf(process.platform !== "win32")(
		"uses private directories, manifests, and transcript files",
		async () => {
			const parent = await temporaryRoot("permissions");
			const root = join(parent, "sessions");
			const store = new FileSessionStore(root);
			await store.append(key, [{ type: "assistant", uuid: "one" }]);

			expect((await stat(root)).mode & 0o777).toBe(0o700);
			expect((await stat(join(root, "manifest.json"))).mode & 0o777).toBe(
				0o600,
			);
			expect((await stat(await onlyFile(root, ".ndjson"))).mode & 0o777).toBe(
				0o600,
			);
		},
	);

	it("deduplicates UUIDs but preserves UUID-less records on disk", async () => {
		const root = join(await temporaryRoot("dedup"), "sessions");
		const store = new FileSessionStore(root);
		await store.append(key, [
			{ type: "entry", uuid: "one", value: 1 },
			{ type: "entry", uuid: "one", value: 2 },
			{ type: "marker", value: 1 },
		]);
		await store.append(key, [
			{ type: "entry", uuid: "one", value: 3 },
			{ type: "marker", value: 1 },
		]);

		expect(await store.load(key)).toEqual([
			{ type: "entry", uuid: "one", value: 1 },
			{ type: "marker", value: 1 },
			{ type: "marker", value: 1 },
		]);
	});

	it("serializes concurrent appends for one transcript without loss", async () => {
		const root = join(await temporaryRoot("concurrent-cache"), "sessions");
		const store = new FileSessionStore(root);
		await Promise.all(
			Array.from({ length: 40 }, (_, index) =>
				store.append(key, [{ type: "entry", uuid: `uuid-${index}`, index }]),
			),
		);
		const loaded = await store.load(key);

		expect(loaded).toHaveLength(40);
		expect(loaded?.map((entry) => entry.uuid).toSorted()).toEqual(
			Array.from({ length: 40 }, (_, index) => `uuid-${index}`).toSorted(),
		);
	});

	it("keeps concurrent transcripts isolated by key", async () => {
		const root = join(await temporaryRoot("parallel-keys"), "sessions");
		const store = new FileSessionStore(root);
		const left = { ...key, sessionId: "left" };
		const right = { ...key, sessionId: "right" };
		await Promise.all([
			store.append(left, [{ type: "entry", value: "left" }]),
			store.append(right, [{ type: "entry", value: "right" }]),
		]);

		expect(await store.load(left)).toEqual([{ type: "entry", value: "left" }]);
		expect(await store.load(right)).toEqual([
			{ type: "entry", value: "right" },
		]);
		expect(
			(await readdir(root)).filter((file) => file.endsWith(".ndjson")),
		).toHaveLength(2);
	});

	it("does not partially append a batch containing an invalid entry", async () => {
		const root = join(await temporaryRoot("invalid-batch"), "sessions");
		const store = new FileSessionStore(root);
		await expect(
			store.append(key, [
				{ type: "valid", uuid: "one" },
				null as unknown as OpaqueEntry,
			]),
		).rejects.toThrow("invalid opaque session entry");
		expect(await store.load(key)).toBeNull();
	});

	it("fails closed on malformed transcript JSON or entry shape", async () => {
		for (const text of ["not-json\n", '{"uuid":"missing-type"}\n']) {
			const root = join(await temporaryRoot("corrupt-transcript"), "sessions");
			const store = new FileSessionStore(root);
			await store.append(key, [{ type: "valid" }]);
			const transcript = await onlyFile(root, ".ndjson");
			await writeFile(transcript, text);
			await expect(store.load(key)).rejects.toBeInstanceOf(
				KlausCacheCorruptionError,
			);
			expect(await readFile(transcript, "utf8")).toBe(text);
		}
	});

	it.each([
		{ schema: 2, sdk: "0.3.281" },
		{ schema: 1, sdk: "0.3.238" },
		{},
		null,
	] as const)("rejects incompatible manifest %j", async (manifest) => {
		const root = join(await temporaryRoot("manifest"), "sessions");
		await import("node:fs/promises").then(({ mkdir }) => mkdir(root));
		await writeFile(join(root, "manifest.json"), JSON.stringify(manifest));
		const store = new FileSessionStore(root);

		await expect(store.prepare()).rejects.toThrow(
			"Klaus cache requires a tested migration",
		);
	});

	it("rejects a nonempty cache root without a manifest", async () => {
		const root = join(await temporaryRoot("missing-manifest"), "sessions");
		await import("node:fs/promises").then(async ({ mkdir }) => {
			await mkdir(root);
			await writeFile(join(root, "unknown"), "data");
		});

		await expect(new FileSessionStore(root).prepare()).rejects.toThrow(
			"Klaus cache requires a tested migration",
		);
	});

	it("makes prepare idempotent", async () => {
		const root = join(await temporaryRoot("prepare"), "sessions");
		const store = new FileSessionStore(root);
		await Promise.all([store.prepare(), store.prepare(), store.prepare()]);
		expect(await readdir(root)).toEqual(["manifest.json"]);
	});

	it("returns no subkeys because persisted names are opaque hashes", async () => {
		const root = join(await temporaryRoot("subkeys"), "sessions");
		const store = new FileSessionStore(root);
		await store.append({ ...key, subpath: "agent" }, [{ type: "entry" }]);
		expect(await store.listSubkeys(key)).toEqual([]);
	});
});

describe("Klaus checkpoints", () => {
	it("returns an empty lineage for a missing checkpoint file", async () => {
		const root = join(await temporaryRoot("missing-checkpoint"), "checkpoints");
		expect(await loadCheckpoints(root, "missing")).toEqual([]);
		expect(await loadCheckpoint(root, "missing")).toBeUndefined();
	});

	it("atomically round-trips private checkpoint arrays", async () => {
		const root = join(await temporaryRoot("checkpoint"), "checkpoints");
		await saveCheckpoint(root, "pi-session", checkpoint("one"));
		await saveCheckpoint(root, "pi-session", checkpoint("two"));

		expect(await loadCheckpoints(root, "pi-session")).toEqual([
			checkpoint("one"),
			checkpoint("two"),
		]);
		expect(await loadCheckpoint(root, "pi-session")).toEqual(checkpoint("two"));
		const files = await readdir(root);
		expect(files.some((file) => file.endsWith(".tmp"))).toBe(false);
	});

	it.runIf(process.platform !== "win32")(
		"writes checkpoint directories and files privately",
		async () => {
			const root = join(await temporaryRoot("checkpoint-mode"), "checkpoints");
			await saveCheckpoint(root, "pi-session", checkpoint("one"));
			expect((await stat(root)).mode & 0o777).toBe(0o700);
			expect((await stat(await onlyFile(root, ".json"))).mode & 0o777).toBe(
				0o600,
			);
		},
	);

	it("loads the legacy single-checkpoint representation", async () => {
		const root = join(await temporaryRoot("legacy-checkpoint"), "checkpoints");
		await saveCheckpoint(root, "pi-session", checkpoint("one"));
		const path = await onlyFile(root, ".json");
		await writeFile(path, JSON.stringify(checkpoint("legacy")));

		expect(await loadCheckpoints(root, "pi-session")).toEqual([
			checkpoint("legacy"),
		]);
	});

	it("replaces the same leaf and fingerprint instead of growing lineage", async () => {
		const root = join(await temporaryRoot("replace-checkpoint"), "checkpoints");
		const first = checkpoint("same");
		const replacement = { ...first, sdkSessionId: "sdk-new", messageCount: 9 };
		await saveCheckpoint(root, "pi-session", first);
		await saveCheckpoint(root, "pi-session", replacement);

		expect(await loadCheckpoints(root, "pi-session")).toEqual([replacement]);
	});

	it("keeps distinct fingerprints for one leaf and distinct leaves for one fingerprint", async () => {
		const root = join(await temporaryRoot("lineage-checkpoint"), "checkpoints");
		const first = checkpoint("first", { piLeafId: "shared-leaf" });
		const second = checkpoint("second", { piLeafId: "shared-leaf" });
		const third = checkpoint("third", {
			fingerprint: second.fingerprint,
			piLeafId: "other-leaf",
		});
		await saveCheckpoint(root, "pi-session", first);
		await saveCheckpoint(root, "pi-session", second);
		await saveCheckpoint(root, "pi-session", third);

		expect(await loadCheckpoints(root, "pi-session")).toEqual([
			first,
			second,
			third,
		]);
	});

	it("isolates checkpoint files by Pi session ID", async () => {
		const root = join(
			await temporaryRoot("checkpoint-isolation"),
			"checkpoints",
		);
		await Promise.all([
			saveCheckpoint(root, "session-a", checkpoint("a")),
			saveCheckpoint(root, "session-b", checkpoint("b")),
		]);

		expect(await loadCheckpoint(root, "session-a")).toEqual(checkpoint("a"));
		expect(await loadCheckpoint(root, "session-b")).toEqual(checkpoint("b"));
		expect(
			(await readdir(root)).filter((file) => file.endsWith(".json")),
		).toHaveLength(2);
	});

	it("serializes concurrent lineage writes without losing checkpoints", async () => {
		const root = join(
			await temporaryRoot("concurrent-checkpoints"),
			"checkpoints",
		);
		const expected = Array.from({ length: 30 }, (_, index) =>
			checkpoint(String(index)),
		);
		await Promise.all(
			expected.map((item) => saveCheckpoint(root, "pi-session", item)),
		);

		expect(await loadCheckpoints(root, "pi-session")).toEqual(expected);
	});

	it.each([
		null,
		{},
		{ sdkSessionId: 1, fingerprint: "f", messageCount: 1, protocol: 1 },
		{ sdkSessionId: "s", fingerprint: 1, messageCount: 1, protocol: 1 },
		{ sdkSessionId: "s", fingerprint: "f", messageCount: 1.5, protocol: 1 },
		{ sdkSessionId: "s", fingerprint: "f", messageCount: -1, protocol: 1 },
		{ sdkSessionId: "s", fingerprint: "f", messageCount: 1, protocol: 2 },
		{
			sdkSessionId: "s",
			fingerprint: "f",
			messageCount: 1,
			protocol: 1,
			position: 1,
		},
		{
			sdkSessionId: "s",
			fingerprint: "f",
			messageCount: 1,
			protocol: 1,
			piLeafId: 1,
		},
	] as const)("rejects invalid checkpoint %j", async (value) => {
		const root = join(await temporaryRoot("invalid-checkpoint"), "checkpoints");
		await saveCheckpoint(root, "pi-session", checkpoint("valid"));
		await writeFile(await onlyFile(root, ".json"), JSON.stringify(value));

		await expect(loadCheckpoints(root, "pi-session")).rejects.toThrow(
			"Klaus checkpoint is invalid",
		);
	});

	it("rejects malformed checkpoint JSON", async () => {
		const root = join(
			await temporaryRoot("malformed-checkpoint"),
			"checkpoints",
		);
		await saveCheckpoint(root, "pi-session", checkpoint("valid"));
		await writeFile(await onlyFile(root, ".json"), "{");
		await expect(loadCheckpoint(root, "pi-session")).rejects.toThrow();
	});

	it.runIf(process.platform !== "win32")(
		"replaces an insecure checkpoint file with mode 0600",
		async () => {
			const root = join(
				await temporaryRoot("checkpoint-replace-mode"),
				"checkpoints",
			);
			await saveCheckpoint(root, "pi-session", checkpoint("one"));
			const path = await onlyFile(root, ".json");
			await chmod(path, 0o644);
			await saveCheckpoint(root, "pi-session", checkpoint("two"));
			expect((await stat(path)).mode & 0o777).toBe(0o600);
		},
	);
});