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