pi-ext

import {
	appendFileSync,
	mkdirSync,
	mkdtempSync,
	rmSync,
	statSync,
	utimesSync,
	writeFileSync,
} from "node:fs";
import { readdir, readFile, stat } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, type MockInstance, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import costExtension, { __test } from "../index.ts";

vi.mock("node:fs/promises", async (importOriginal) => {
	const actual = await importOriginal<typeof import("node:fs/promises")>();
	return {
		...actual,
		readFile: vi.fn(actual.readFile),
		readdir: vi.fn(actual.readdir),
		stat: vi.fn(actual.stat),
	};
});

const STATUS_ID = "cost";
const noModel = () => undefined;

/** Catalog lookup as the extension sees it: `ctx.modelRegistry.find`. */
async function registryFind(): Promise<{
	session: PiSession;
	find: ExtensionContext["modelRegistry"]["find"];
}> {
	let registry: ExtensionContext["modelRegistry"] | undefined;
	const probe = (pi: ExtensionAPI) => {
		pi.on("session_start", (_event, ctx) => {
			registry = ctx.modelRegistry;
		});
	};
	const s = await createPiSession({
		extensionFactories: [costExtension, probe],
	});
	await s.session.extensionRunner.emit({
		type: "session_start",
		reason: "startup",
	});
	if (!registry) throw new Error("probe never saw session_start");
	return { session: s, find: registry.find.bind(registry) };
}

function extensionPath(): string {
	return fileURLToPath(new URL("../index.ts", import.meta.url));
}

function costEntry(cost: number): string {
	return `${JSON.stringify({
		type: "message",
		message: { role: "assistant", usage: { cost: { total: cost } } },
	})}\n`;
}

function monthKey(): string {
	const now = new Date();
	return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`;
}

describe("cost contract", () => {
	let session: PiSession | undefined;
	let root: string | undefined;
	let statusSpy: MockInstance | undefined;

	afterEach(async () => {
		await session?.dispose();
		if (root) rmSync(root, { recursive: true, force: true });
		session = undefined;
		root = undefined;
		statusSpy = undefined;
		vi.restoreAllMocks();
		vi.clearAllMocks();
		vi.useRealTimers();
	});

	function sessionDirectory(): string {
		root = mkdtempSync(path.join(tmpdir(), "pi-cost-test-"));
		const dir = path.join(root, "project");
		mkdirSync(dir);
		return dir;
	}

	function sessionFile(dir: string, name: string, cost: number): string {
		const file = path.join(dir, `${monthKey()}-${name}.jsonl`);
		writeFileSync(file, costEntry(cost), "utf8");
		return file;
	}

	async function start(dir: string): Promise<PiSession> {
		session = await createPiSession({ extensionFactories: [costExtension] });
		const runner = session.session.extensionRunner;
		const ui = { ...runner.getUIContext() };
		statusSpy = vi.spyOn(ui, "setStatus");
		runner.setUIContext(ui, "print");
		vi.spyOn(session.session.sessionManager, "getSessionDir").mockReturnValue(
			dir,
		);
		await runner.emit({ type: "session_start", reason: "startup" });
		return session;
	}

	function statusCalls(): Array<[string, unknown?]> {
		return (statusSpy?.mock.calls ?? []) as Array<[string, unknown?]>;
	}

	function status(): string | undefined {
		return statusCalls()
			.filter((call) => call[0] === STATUS_ID)
			.at(-1)?.[1] as string | undefined;
	}

	it("loads in real Pi extension runtime", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath()] });
		expect(session.extensions).toHaveLength(1);
	});

	it("scans the monthly root and direct project directories only", async () => {
		const dir = sessionDirectory();
		sessionFile(dir, "session", 1.25);
		sessionFile(path.dirname(dir), "child", 2);
		const sibling = path.join(path.dirname(dir), "other-project");
		mkdirSync(sibling);
		sessionFile(sibling, "child", 3);
		const nested = path.join(dir, "nested");
		mkdirSync(nested);
		sessionFile(nested, "ignored", 100);
		writeFileSync(path.join(dir, "2000-01-old.jsonl"), costEntry(100));
		writeFileSync(path.join(dir, `${monthKey()}-ignored.txt`), costEntry(100));

		await expect(__test.calculateMonthlyCost(dir, noModel)).resolves.toBe(6.25);
		vi.mocked(readdir).mockClear();
		await expect(
			__test.calculateMonthlyCost(dir, noModel, () => false),
		).resolves.toBeUndefined();
		expect(readdir).not.toHaveBeenCalled();
	});

	it("reuses unchanged subtotals and rereads only new or changed files", async () => {
		const dir = sessionDirectory();
		const file = sessionFile(dir, "parent", 1);
		sessionFile(dir, "child", 2);
		const cache = new Map();
		const refresh = () =>
			__test.calculateMonthlyCost(dir, noModel, () => true, cache);
		await expect(refresh()).resolves.toBe(3);
		expect(readFile).toHaveBeenCalledTimes(2);

		vi.mocked(readFile).mockClear();
		await expect(refresh()).resolves.toBe(3);
		expect(readFile).not.toHaveBeenCalled();

		// Same-size rewrite: modification time must invalidate the subtotal.
		const before = statSync(file);
		writeFileSync(file, costEntry(4));
		utimesSync(file, before.atime, new Date(before.mtimeMs + 2000));
		await expect(refresh()).resolves.toBe(6);
		expect(readFile).toHaveBeenCalledExactlyOnceWith(file, "utf-8");

		// Changed size with unchanged mtime must also invalidate it.
		vi.mocked(readFile).mockClear();
		const modified = statSync(file);
		appendFileSync(file, costEntry(5));
		utimesSync(file, modified.atime, modified.mtime);
		sessionFile(dir, "new-child", 3);
		await expect(refresh()).resolves.toBe(14);
		expect(readFile).toHaveBeenCalledTimes(2);

		vi.mocked(readFile).mockClear();
		rmSync(file);
		await expect(refresh()).resolves.toBe(5);
		expect(cache.has(file)).toBe(false);
		expect(readFile).not.toHaveBeenCalled();
	});

	it("drops previous-month cache entries on rollover", async () => {
		vi.useFakeTimers({ toFake: ["Date"] });
		vi.setSystemTime(new Date(2026, 0, 31, 12));
		const dir = sessionDirectory();
		const old = sessionFile(dir, "old", 8);
		const cache = new Map();
		await expect(
			__test.calculateMonthlyCost(dir, noModel, () => true, cache),
		).resolves.toBe(8);
		vi.setSystemTime(new Date(2026, 1, 1, 12));
		const next = sessionFile(dir, "new", 2);
		await expect(
			__test.calculateMonthlyCost(dir, noModel, () => true, cache),
		).resolves.toBe(2);
		expect([...cache.keys()]).toEqual([next]);
		expect(cache.has(old)).toBe(false);
	});

	it("does not cache failed reads or stats, and retries next refresh", async () => {
		const dir = sessionDirectory();
		sessionFile(dir, "session", 2);
		const cache = new Map();
		const refresh = () =>
			__test.calculateMonthlyCost(dir, noModel, () => true, cache);
		vi.mocked(readFile).mockRejectedValueOnce(new Error("unreadable"));
		await expect(refresh()).resolves.toBe(0);
		expect(cache.size).toBe(0);
		await expect(refresh()).resolves.toBe(2);
		vi.mocked(stat).mockRejectedValueOnce(new Error("removed during scan"));
		await expect(refresh()).resolves.toBe(0);
		expect(cache.size).toBe(0);
		await expect(refresh()).resolves.toBe(2);
	});

	it("does not commit a scan invalidated during a file read", async () => {
		const dir = sessionDirectory();
		const file = sessionFile(dir, "session", 2);
		const cache = new Map();
		await __test.calculateMonthlyCost(dir, noModel, () => true, cache);
		const original = cache.get(file);
		appendFileSync(file, costEntry(3));
		let current = true;
		vi.mocked(readFile).mockImplementationOnce(async () => {
			current = false;
			return costEntry(5);
		});
		await expect(
			__test.calculateMonthlyCost(dir, noModel, () => current, cache),
		).resolves.toBeUndefined();
		expect(cache.get(file)).toBe(original);
	});

	it("refreshes persisted child costs after agent_end without rereading unchanged files", async () => {
		const dir = sessionDirectory();
		sessionFile(dir, "parent", 1);
		const s = await start(dir);
		await vi.waitFor(() => expect(status()).toContain("$1.00"));
		vi.mocked(readFile).mockClear();
		const child = sessionFile(dir, "child", 2);

		await s.session.extensionRunner.emit({
			type: "agent_end",
			messages: [],
		});
		await vi.waitFor(() => expect(status()).toContain("$3.00"));
		expect(readFile).toHaveBeenCalledExactlyOnceWith(child, "utf-8");
	});

	it("coalesces requests during a blocked scan without blocking agent_end", async () => {
		const dir = sessionDirectory();
		sessionFile(dir, "parent", 1);
		const s = await start(dir);
		let release!: () => void;
		const gate = new Promise<void>((resolve) => {
			release = resolve;
		});
		vi.mocked(readFile).mockClear();
		vi.mocked(readFile).mockImplementationOnce(async () => {
			await gate;
			return costEntry(1);
		});
		try {
			await vi.waitFor(() => expect(readFile).toHaveBeenCalledOnce());
			vi.mocked(readdir).mockClear();
			sessionFile(dir, "child", 2);
			for (let i = 0; i < 3; i++) {
				await s.session.extensionRunner.emit({
					type: "agent_end",
					messages: [],
				});
			}
			expect(readdir).not.toHaveBeenCalled();
		} finally {
			release();
		}
		await vi.waitFor(() => expect(status()).toContain("$3.00"));
		expect(readFile).toHaveBeenCalledTimes(2);
		// One follow-up scan: root plus its one project directory.
		expect(readdir).toHaveBeenCalledTimes(2);
	});

	it.each(["shutdown", "replacement"])(
		"discards an in-flight monthly result on session %s",
		async (event) => {
			const dir = sessionDirectory();
			sessionFile(dir, "parent", 9);
			const s = await start(dir);
			let release!: () => void;
			const gate = new Promise<void>((resolve) => {
				release = resolve;
			});
			vi.mocked(readFile).mockClear();
			vi.mocked(readFile).mockImplementationOnce(async () => {
				await gate;
				return costEntry(9);
			});
			let callsAfterShutdown = 0;
			try {
				await vi.waitFor(() => expect(readFile).toHaveBeenCalledOnce());
				await s.session.extensionRunner.emit({
					type: "session_shutdown",
					reason: "quit",
				});
				callsAfterShutdown = statusCalls().length;
				if (event === "replacement") {
					const nextDir = path.join(dir, "separate-root", "project");
					mkdirSync(nextDir, { recursive: true });
					sessionFile(nextDir, "replacement", 2);
					vi.spyOn(s.session.sessionManager, "getSessionDir").mockReturnValue(
						nextDir,
					);
					await s.session.extensionRunner.emit({
						type: "session_start",
						reason: "reload",
					});
				}
			} finally {
				release();
			}
			if (event === "replacement") {
				await vi.waitFor(() => expect(status()).toContain("$2.00"));
			} else {
				await new Promise((resolve) => setTimeout(resolve, 20));
				expect(statusCalls()).toHaveLength(callsAfterShutdown);
			}
			expect(
				statusCalls().some((call) => String(call[1]).includes("$9.00")),
			).toBe(false);
		},
	);

	it("prices zero-cost assistant messages from the model catalog", async () => {
		const probe = await registryFind();
		session = probe.session;
		const findModel = probe.find;
		const model = findModel("zai", "glm-4.7");
		expect(model?.cost.input).toBeGreaterThan(0);
		const usage = { input: 1_000_000, output: 0, cacheRead: 0, cacheWrite: 0 };
		const entries = [
			{
				type: "message",
				message: {
					role: "assistant",
					provider: "zai",
					model: "glm-4.7",
					usage,
				},
			},
			{
				type: "message",
				message: {
					role: "assistant",
					provider: "nobody",
					model: "unknown",
					usage,
				},
			},
			{
				type: "message",
				message: {
					role: "assistant",
					provider: "zai",
					model: "glm-4.7",
					usage: { ...usage, cost: { total: 0.5 } },
				},
			},
		];
		expect(__test.sumMessageCosts(entries, findModel)).toBeCloseTo(
			(model?.cost.input ?? Number.NaN) + 0.5,
		);
		expect(__test.sumMessageCosts(entries, noModel)).toBe(0.5);
	});

	it("charges cache writes through the catalog", async () => {
		const probe = await registryFind();
		session = probe.session;
		const findModel = probe.find;
		const model = findModel("anthropic", "claude-haiku-4-5");
		expect(model?.cost.cacheWrite).toBeGreaterThan(0);
		const entries = [
			{
				type: "message",
				message: {
					role: "assistant",
					provider: "anthropic",
					model: "claude-haiku-4-5",
					usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 1_000_000 },
				},
			},
		];
		expect(__test.sumMessageCosts(entries, findModel)).toBeCloseTo(
			model?.cost.cacheWrite ?? Number.NaN,
		);
	});

	it("sums session usage from native usage-bearing entries", () => {
		expect(
			__test.sessionUsage([
				{
					type: "message",
					message: {
						role: "assistant",
						usage: { input: 100, output: 20, cacheRead: 300 },
					},
				},
				{
					type: "message",
					message: { role: "toolResult", usage: { input: 10, output: 5 } },
				},
				{ type: "compaction", usage: { input: 50, cacheWrite: 40 } },
				{ type: "branch_summary", usage: { output: 7 } },
				{ type: "message", message: { role: "user", usage: { input: 999 } } },
			]),
		).toEqual({ input: 160, output: 32, cacheRead: 300, cacheWrite: 40 });
	});

	it("uses the latest assistant prompt for cache hit rate", () => {
		expect(
			__test.latestAssistantCacheHitRate([
				{
					type: "message",
					message: { role: "assistant", usage: { input: 90, cacheRead: 10 } },
				},
				{
					type: "message",
					message: {
						role: "assistant",
						usage: { input: 10, cacheRead: 80, cacheWrite: 10 },
					},
				},
			]),
		).toBe(80);
	});

	it("shows only input, output, and cache hit rate in the tui label", () => {
		expect(
			__test.tokenLabel(
				{
					input: 1100,
					output: 978,
					cacheRead: 108900,
					cacheWrite: 2100,
				},
				98,
			),
		).toBe("↑1.1k ↓978  98%");
	});

	it("scales large token counts into millions", () => {
		expect(
			__test.tokenLabel(
				{ input: 1_618_000, output: 448_000, cacheRead: 0, cacheWrite: 0 },
				100,
			),
		).toBe("↑1.6M ↓448k  100%");
	});

	it("spells out the non-tui label without glyphs", () => {
		expect(
			__test.plainLabel({
				usage: {
					input: 1_618_000,
					output: 448_000,
					cacheRead: 121_527_000,
					cacheWrite: 1_381_000,
				},
				cacheHitRate: 100,
				turnCost: 0.16,
				sessionCost: 70.7,
				monthlyCost: 3146.2,
			}),
		).toBe(
			"in 1.6M, out 448k, cache hit 100%, turn $0.16, session $70.7, month $3.1k",
		);
	});
});