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 { afterEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
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),
	};
});

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 pi-test-harness", () => {
	let t: TestSession | undefined;
	let root: string | undefined;

	afterEach(async () => {
		await t?.session.extensionRunner.emit({
			type: "session_shutdown",
			reason: "quit",
		});
		t?.dispose();
		if (root) rmSync(root, { recursive: true, force: true });
		t = undefined;
		root = 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) {
		t = await createTestSession({ extensionFactories: [costExtension] });
		vi.spyOn(t.session.sessionManager, "getSessionDir").mockReturnValue(dir);
		await t.session.extensionRunner.emit({
			type: "session_start",
			reason: "startup",
		});
		return t;
	}

	function status(): string | undefined {
		return t?.events.uiCallsFor("setStatus").at(-1)?.args[1] as
			| string
			| undefined;
	}

	it("loads in real Pi extension runtime", async () => {
		t = await createTestSession({ extensionFactories: [costExtension] });
		expect(t.session.extensionRunner.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)).resolves.toBe(6.25);
		vi.mocked(readdir).mockClear();
		await expect(
			__test.calculateMonthlyCost(dir, () => 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, () => 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, () => true, cache),
		).resolves.toBe(8);
		vi.setSystemTime(new Date(2026, 1, 1, 12));
		const next = sessionFile(dir, "new", 2);
		await expect(
			__test.calculateMonthlyCost(dir, () => 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, () => 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, () => 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, () => 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 session = await start(dir);
		await vi.waitFor(() => expect(status()).toContain("$1.00"));
		vi.mocked(readFile).mockClear();
		const child = sessionFile(dir, "child", 2);

		await session.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 session = 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 session.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 session = 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 session.session.extensionRunner.emit({
					type: "session_shutdown",
					reason: "quit",
				});
				callsAfterShutdown = session.events.uiCallsFor("setStatus").length;
				if (event === "replacement") {
					const nextDir = path.join(dir, "separate-root", "project");
					mkdirSync(nextDir, { recursive: true });
					sessionFile(nextDir, "replacement", 2);
					vi.spyOn(
						session.session.sessionManager,
						"getSessionDir",
					).mockReturnValue(nextDir);
					await session.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(session.events.uiCallsFor("setStatus")).toHaveLength(
					callsAfterShutdown,
				);
			}
			expect(
				session.events
					.uiCallsFor("setStatus")
					.some((call) => String(call.args[1]).includes("$9.00")),
			).toBe(false);
		},
	);

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