pi-ext

import {
	mkdirSync,
	mkdtempSync,
	readdirSync,
	readFileSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import {
	InMemoryCredentialStore,
	type OAuthCredential,
} from "@earendil-works/pi-ai";
import {
	type AfterProviderResponseEvent,
	initTheme,
} from "@earendil-works/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { __test } from "../index.ts";

/** 4h50m ahead, matching the reset window the footer renders. */
const RESET_AT = Date.now() + 4 * 3_600_000 + 50 * 60_000 + 30_000;

describe("quota contract", () => {
	let session: PiSession | undefined;

	afterEach(async () => {
		vi.unstubAllGlobals();
		await session?.dispose();
		session = undefined;
	});

	it("loads in real Pi extension runtime", async () => {
		const extensionPath = fileURLToPath(
			new URL("../index.ts", import.meta.url),
		);
		session = await createPiSession({ extensionPaths: [extensionPath] });

		expect(session.extensions).toHaveLength(1);
		expect(session.extensions[0]!.path).toBe(extensionPath);
	}, 15000);

	it("writes its opt-in lifecycle trace inside the fixture sandbox", async () => {
		const root = mkdtempSync(path.join(tmpdir(), "pi-ext-quota-debug-"));
		try {
			const extensionPath = fileURLToPath(
				new URL("../index.ts", import.meta.url),
			);
			session = await createPiSession({
				cwd: root,
				extensionPaths: [extensionPath],
				env: { PI_QUOTA_DEBUG: "1" },
			});
			const runner = session.session.extensionRunner;
			await runner.emit({ type: "session_start", reason: "startup" });
			await runner.emit({ type: "session_shutdown", reason: "quit" });
			const stateHome = session.env["XDG_STATE_HOME"];
			if (!stateHome) throw new Error("missing sandbox state home");
			const directory = path.join(stateHome, "pi-ext", "debug", "quota");
			const [file] = readdirSync(directory);
			if (!file) throw new Error("missing debug file");
			const events = readFileSync(path.join(directory, file), "utf8")
				.trim()
				.split("\n")
				.map((line) => JSON.parse(line).event);
			expect(events).toEqual(
				expect.arrayContaining(["session.start", "session.shutdown"]),
			);
		} finally {
			await session?.dispose();
			session = undefined;
			rmSync(root, { recursive: true, force: true });
		}
	});

	function codexToken(accountId: string) {
		const claim = Buffer.from(
			JSON.stringify({
				"https://api.openai.com/auth": { chatgpt_account_id: accountId },
			}),
		).toString("base64url");
		return `header.${claim}.signature`;
	}

	function quotaCommandContext() {
		if (!session) throw new Error("missing session");
		initTheme("default", false);
		const runner = session.session.extensionRunner;
		const ui = {
			...runner.getUIContext(),
			setStatus: vi.fn(),
			setWidget: vi.fn(),
		};
		runner.setUIContext(ui, "tui");
		const ctx = Object.create(runner.createCommandContext());
		Object.defineProperty(ctx, "model", {
			value: { provider: "openai-codex" },
		});
		const command = runner.getCommand("quota");
		if (!command) throw new Error("missing /quota command");
		return { command, ctx };
	}

	function installCodexRefresh(
		refresh: (credential: OAuthCredential) => Promise<OAuthCredential>,
	) {
		if (!session) throw new Error("missing session");
		const runtime = session.session.modelRuntime;
		const provider = runtime.getProvider("openai-codex");
		if (!provider?.auth.oauth) throw new Error("missing native Codex OAuth");
		runtime.registerNativeProvider({
			...provider,
			auth: {
				...provider.auth,
				oauth: {
					...provider.auth.oauth,
					refresh: (credential, _signal) => refresh(credential),
					toAuth: async (credential) => ({ apiKey: credential.access }),
				},
			},
		});
	}

	it("uses active stored Codex OAuth token, never fixture auth files", async () => {
		const credentials = new InMemoryCredentialStore();
		const token = codexToken("acct-active");
		await credentials.modify("openai-codex", async () => ({
			type: "oauth",
			access: token,
			refresh: "refresh-active",
			expires: Date.now() + 3_600_000,
		}));
		const root = mkdtempSync(path.join(tmpdir(), "pi-ext-quota-auth-"));
		try {
			const extensionPath = fileURLToPath(
				new URL("../index.ts", import.meta.url),
			);
			session = await createPiSession({
				cwd: root,
				credentials,
				extensionPaths: [extensionPath],
			});
			writeFileSync(
				path.join(session.agentDir, "auth.json"),
				JSON.stringify({
					"openai-codex": {
						type: "oauth",
						access: codexToken("acct-default"),
						refresh: "refresh-default",
						expires: Date.now() + 3_600_000,
						accountId: "acct-default",
					},
				}),
			);
			const codexHome = path.join(session.env["HOME"]!, ".codex");
			mkdirSync(codexHome, { recursive: true });
			writeFileSync(
				path.join(codexHome, "auth.json"),
				JSON.stringify({
					tokens: {
						access_token: codexToken("acct-codex"),
						account_id: "acct-codex",
					},
				}),
			);
			const fetch = vi.fn(async (_input: string | URL, _init?: RequestInit) => {
				return Response.json({ rate_limit: {} });
			});
			vi.stubGlobal("fetch", fetch);

			const { command, ctx } = quotaCommandContext();
			await command.handler("", ctx);

			expect(fetch).toHaveBeenCalledTimes(1);
			const [url, init] = fetch.mock.calls[0]!;
			expect(String(url)).toBe("https://chatgpt.com/backend-api/wham/usage");
			const headers = new Headers(init?.headers);
			expect(headers.get("Authorization")).toBe(`Bearer ${token}`);
			expect(headers.get("ChatGPT-Account-Id")).toBe("acct-active");
			expect(session.errors).toEqual([]);
		} finally {
			await session?.dispose();
			session = undefined;
			rmSync(root, { recursive: true, force: true });
		}
	}, 15000);

	it.each([
		["malformed token", "not-a-jwt"],
		["missing Codex account claim", "header.e30.signature"],
		["no OAuth credential", undefined],
	] as const)(
		"does not request Codex quota with %s",
		async (_name, access) => {
			const credentials = new InMemoryCredentialStore();
			if (access) {
				await credentials.modify("openai-codex", async () => ({
					type: "oauth",
					access,
					refresh: "refresh",
					expires: Date.now() + 3_600_000,
				}));
			}
			const extensionPath = fileURLToPath(
				new URL("../index.ts", import.meta.url),
			);
			session = await createPiSession({
				credentials,
				extensionPaths: [extensionPath],
			});
			const fetch = vi.fn(async (input: string | URL) => {
				throw new Error(`unexpected network endpoint: ${input}`);
			});
			vi.stubGlobal("fetch", fetch);

			const { command, ctx } = quotaCommandContext();
			await command.handler("", ctx);

			expect(fetch).not.toHaveBeenCalled();
			expect(session.errors).toEqual([]);
		},
		15000,
	);

	it("refreshes expired Codex OAuth before polling with refreshed account", async () => {
		const credentials = new InMemoryCredentialStore();
		await credentials.modify("openai-codex", async () => ({
			type: "oauth",
			access: codexToken("acct-old"),
			refresh: "refresh-old",
			expires: Date.now() - 1,
			account_id: "acct-old-credential",
		}));
		const extensionPath = fileURLToPath(
			new URL("../index.ts", import.meta.url),
		);
		session = await createPiSession({
			credentials,
			extensionPaths: [extensionPath],
		});
		const refreshed = codexToken("acct-refreshed");
		const refresh = vi.fn(async () => ({
			type: "oauth" as const,
			access: refreshed,
			refresh: "refresh-new",
			expires: Date.now() + 3_600_000,
			account_id: "acct-old-credential",
		}));
		installCodexRefresh(refresh);
		const fetch = vi.fn(async (_input: string | URL, _init?: RequestInit) =>
			Response.json({ rate_limit: {} }),
		);
		vi.stubGlobal("fetch", fetch);

		const { command, ctx } = quotaCommandContext();
		await command.handler("", ctx);

		expect(refresh).toHaveBeenCalledTimes(1);
		expect(fetch).toHaveBeenCalledTimes(1);
		const [url, init] = fetch.mock.calls[0]!;
		expect(String(url)).toBe("https://chatgpt.com/backend-api/wham/usage");
		const headers = new Headers(init?.headers);
		expect(headers.get("Authorization")).toBe(`Bearer ${refreshed}`);
		expect(headers.get("ChatGPT-Account-Id")).toBe("acct-refreshed");
		expect(session.errors).toEqual([]);
	}, 15000);

	it("derives Codex account identity only from the supplied token", () => {
		const encode = (value: unknown) =>
			Buffer.from(JSON.stringify(value)).toString("base64url");
		const token = `header.${encode({ "https://api.openai.com/auth": { chatgpt_account_id: "acct-session" } })}.signature`;
		expect(__test.codexAccountId(token)).toBe("acct-session");
		expect(
			__test.codexAccountId(
				`header.${encode({ account_id: "acct-other" })}.signature`,
			),
		).toBeUndefined();
		expect(
			__test.codexAccountId(
				`header.${encode({ "https://api.openai.com/auth": { chatgpt_account_id: "\n" } })}.signature`,
			),
		).toBeUndefined();
		expect(__test.codexAccountId("malformed-token")).toBeUndefined();
	});

	it("parses stolen provider quota shapes", () => {
		expect(
			__test.parseCopilot({
				quota_snapshots: {
					premium_interactions: { percent_remaining: 75 },
					chat: { percent_remaining: 40 },
				},
			}),
		).toMatchObject({
			name: "GitHub Copilot",
			fiveH: { used: 25, total: 100, remaining: 75 },
			weekly: { used: 60, total: 100, remaining: 40 },
		});

		expect(
			__test.parseGemini({
				buckets: [
					{ modelId: "gemini-pro", remainingFraction: 0.25 },
					{ modelId: "gemini-flash", remainingFraction: 0.75 },
				],
			}),
		).toMatchObject({
			name: "Gemini",
			fiveH: { used: 75, total: 100, remaining: 25 },
			weekly: { used: 25, total: 100, remaining: 75 },
		});

		expect(
			__test.parseMinimax({
				base_resp: { status_code: 0, status_msg: "" },
				model_remains: [
					{
						model_name: "MiniMax-M2",
						current_interval_total_count: 100,
						current_interval_usage_count: 30,
						current_weekly_total_count: 1000,
						current_weekly_usage_count: 800,
					},
				],
			}),
		).toMatchObject({
			name: "MiniMax",
			fiveH: { used: 70, total: 100, remaining: 30 },
			weekly: { used: 200, total: 1000, remaining: 800 },
		});

		expect(
			__test.parseMinimax({
				base_resp: { status_code: 0, status_msg: "" },
				model_remains: [
					{
						model_name: "general",
						current_interval_total_count: 0,
						current_interval_usage_count: 0,
						current_interval_remaining_percent: 25,
						current_weekly_total_count: 0,
						current_weekly_usage_count: 0,
						current_weekly_remaining_percent: 80,
					},
				],
			}),
		).toMatchObject({
			name: "MiniMax",
			fiveH: { used: 75, total: 100, remaining: 25 },
			weekly: { used: 20, total: 100, remaining: 80 },
		});
	});

	it("reads quota from provider response headers", () => {
		const anthropic: AfterProviderResponseEvent = {
			type: "after_provider_response",
			status: 200,
			headers: {
				"content-type": "text/event-stream",
				"anthropic-ratelimit-unified-5h-utilization": "0.26",
				"anthropic-ratelimit-unified-5h-reset": "1764554400",
				"anthropic-ratelimit-unified-7d-utilization": "0.75",
				"anthropic-ratelimit-unified-7d-reset": "1764615600",
			},
		};
		expect(__test.usageFromResponse("anthropic", anthropic)).toEqual({
			name: "Anthropic",
			fiveH: { used: 26, total: 100, remaining: 74, resetAt: 1764554400000 },
			weekly: { used: 75, total: 100, remaining: 25, resetAt: 1764615600000 },
		});

		const codex: AfterProviderResponseEvent = {
			type: "after_provider_response",
			status: 200,
			headers: {
				"x-codex-primary-used-percent": "42",
				"x-codex-primary-reset-at": "1900000000",
				"x-codex-secondary-used-percent": "7",
			},
		};
		expect(__test.usageFromResponse("openai-codex", codex)).toEqual({
			name: "OpenAI Codex",
			fiveH: { used: 42, total: 100, remaining: 58, resetAt: 1900000000000 },
			weekly: { used: 7, total: 100, remaining: 93, resetAt: undefined },
		});

		// Wrong provider or missing quota headers leave the poll result alone.
		expect(__test.usageFromResponse("klaus", anthropic)).toBeUndefined();
		expect(__test.usageFromResponse("anthropic", codex)).toBeUndefined();
		expect(
			__test.usageFromResponse("anthropic", {
				headers: { "anthropic-ratelimit-unified-5h-utilization": "nope" },
			}),
		).toBeUndefined();
	});

	it("survives provider responses in the real runtime", async () => {
		const extensionPath = fileURLToPath(
			new URL("../index.ts", import.meta.url),
		);
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await session.session.extensionRunner.emit({
			type: "after_provider_response",
			status: 200,
			headers: { "anthropic-ratelimit-unified-5h-utilization": "0.5" },
		});
		expect(session.errors).toEqual([]);
	}, 15000);

	it("drops the provider name for the active provider only", () => {
		const anthropic = {
			name: "Anthropic",
			fiveH: { used: 4, total: 100, remaining: 96, resetAt: RESET_AT },
			weekly: { used: 0, total: 0, remaining: 0 },
		};

		expect(__test.segmentText(anthropic, true)).toEqual([
			"4% (4h50m)",
			"4%",
			"4%",
		]);
		expect(__test.segmentText(anthropic, false)).toEqual([
			"Anthropic 4% (4h50m)",
			"Anthropic 4%",
			"4%",
		]);
	});

	it("keeps the name in status fallbacks and spells out the non-tui label", () => {
		const providers = [
			{
				name: "Anthropic",
				fiveH: { used: 4, total: 100, remaining: 96, resetAt: RESET_AT },
				weekly: { used: 0, total: 0, remaining: 0 },
			},
		];

		expect(__test.quotaStatusLabel(providers, "Anthropic")).toBe(
			"Anthropic 4% (4h50m)",
		);
		expect(__test.quotaPlainLabel(providers, "Anthropic")).toBe(
			"Anthropic quota 4% used, resets in 4h50m",
		);
		expect(__test.quotaPlainLabel([], "Anthropic")).toBeUndefined();
	});
});