Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
6.0 kB · 156 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157import { mkdtempSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { registerUsageFeed, unregisterUsageFeed } from "@earendil-works/pi-ai";import { afterEach, describe, expect, it, vi } from "vitest";import { AuthStorage } from "../src/core/auth-storage.js";import { initTheme } from "../src/modes/interactive/theme/theme.js";import { fetchUsageReports, renderUsageReportText, runUsageCommand, sessionPathFor, type UsageReport,} from "../src/modes/interactive/usage-report.js";
initTheme("dark");
const oauth = (access: string, email: string, accountId: string) => ({ type: "oauth" as const, access, refresh: `refresh-${accountId}`, expires: Date.now() + 60_000, email, accountId,});
describe("usage reports", () => { afterEach(() => { vi.restoreAllMocks(); unregisterUsageFeed("session-provider"); });
it("fetches and keeps quota data separate for every Codex account", async () => { const auth = AuthStorage.inMemory({ "openai-codex": [ oauth("access-one", "one@example.com", "account-one"), oauth("access-two", "two@example.com", "account-two"), ], }); const fetchImpl = vi.fn(async (_input: string | URL, init?: RequestInit) => { const account = (init?.headers as Record<string, string>)["ChatGPT-Account-Id"]; return new Response( JSON.stringify({ plan_type: "plus", rate_limit: { primary_window: { used_percent: account === "account-one" ? 20 : 80, limit_window_seconds: 18_000 }, }, }), { status: 200, headers: { "content-type": "application/json" } }, ); });
const result = await fetchUsageReports(auth, { fetchImpl, nowMs: 1_700_000_000_000 }); expect(result.reports).toHaveLength(2); expect(result.reports.map((report) => report.email)).toEqual(["one@example.com", "two@example.com"]); expect(result.reports.map((report) => report.limits[0]?.usedPercent)).toEqual([20, 80]); expect(result.reports.map((report) => report.limits[0]?.status)).toEqual(["ok", "ok"]); expect(result.skipped).toEqual([]); });
it("reports a provider with no feed, and why, instead of dropping it", async () => { const auth = AuthStorage.inMemory({ "no-feed-provider": [{ type: "api_key", key: "key" }] }); const result = await fetchUsageReports(auth, { fetchImpl: vi.fn() }); expect(result.reports).toEqual([]); expect(result.skipped).toEqual([{ provider: "no-feed-provider", reason: "no usage feed" }]); const text = renderUsageReportText(result, undefined, 100, { includeSession: false }); expect(text).toContain("No quota feed"); expect(text).toContain("no usage feed"); });
it("hands a captured console session to a feed that asks for one", async () => { const dir = mkdtempSync(join(tmpdir(), "usage-session-")); vi.spyOn(await import("../src/config.js"), "getAgentDir").mockReturnValue(dir); writeFileSync(join(dir, "session-provider-session.json"), JSON.stringify({ token: "captured" }));
registerUsageFeed({ provider: "session-provider", providerName: "Session Provider", credentialKinds: ["session"], advance: () => ({ report: { limits: [{ label: "Primary", usedPercent: 25, status: "ok" }] } }), }); const auth = AuthStorage.inMemory({ "session-provider": [{ type: "api_key", key: "unused" }] }); const result = await fetchUsageReports(auth, { fetchImpl: vi.fn() }); expect(result.reports).toHaveLength(1); expect(result.reports[0]?.limits[0]?.usedPercent).toBe(25); expect(sessionPathFor("session-provider")).toBe(join(dir, "session-provider-session.json")); });
it("asks its host to run the provider's own CLI, and treats a refusal as data", async () => { const auth = AuthStorage.inMemory({}); // The alibaba feed prefers the Bailian CLI; with no `bl` on this machine the refusal has to come // back as a response the feed can route around, never as a thrown error. const result = await fetchUsageReports(auth, { fetchImpl: vi.fn() }); expect(result.reports).toEqual([]);
const refused = await runUsageCommand({ kind: "command", command: "rm", args: ["-rf", "/"] }); expect(refused.status).toBe(126); expect(refused.body).toContain("will not run");
const missing = await runUsageCommand({ kind: "command", command: "bl", args: ["usage", "token-plan"] }); expect(missing.status).not.toBe(0); expect(typeof missing.body).toBe("string"); });
it("renders limits, balance, notes and unknown percentages", () => { const reports: UsageReport[] = [ { provider: "deepseek", providerName: "DeepSeek", fetchedAt: Date.now(), balance: "$11.83", notes: ["Extra usage is off"], limits: [{ label: "Primary", status: "unknown" }], }, { provider: "anthropic", providerName: "Anthropic", fetchedAt: Date.now(), email: "a@example.com", planType: "team_standard", limits: [{ label: "Weekly", usedPercent: 94, windowLabel: "7d window", status: "warning" }], }, ]; const text = renderUsageReportText({ reports, skipped: [] }, undefined, 100, { includeSession: false }); expect(text).toContain("DeepSeek"); // The renderer colours the label, so assert on the parts either side of the escape. expect(text).toContain("Balance"); expect(text).toContain("$11.83"); expect(text).toContain("Extra usage is off"); expect(text).toContain("no percentage reported"); expect(text).toContain("Weekly (7d window)"); expect(text).toContain("94.0% used"); expect(text).toContain("team_standard"); });
it("shows a fetch failure as data on that provider's row", () => { const text = renderUsageReportText( { reports: [ { provider: "openai-codex", providerName: "ChatGPT Codex", fetchedAt: Date.now(), limits: [], error: "OAuth token unavailable; run /login", }, ], skipped: [], }, undefined, 100, { includeSession: false }, ); expect(text).toContain("Unable to fetch quota: OAuth token unavailable; run /login"); });});