feat: add quota tracking
This commit is contained in:
@@ -0,0 +1,397 @@
|
||||
import { expect, test } from "bun:test";
|
||||
import type { Schema } from "effect";
|
||||
import { normalizeQuota, QuotaTracker } from "../src/quota.js";
|
||||
import { formatQuotaWindow, loadQuotaStatus } from "../src/tui.js";
|
||||
|
||||
const plus = {
|
||||
plan_type: "plus",
|
||||
account_id: "account-one",
|
||||
rate_limit: {
|
||||
allowed: true,
|
||||
limit_reached: false,
|
||||
primary_window: {
|
||||
used_percent: 12.5,
|
||||
limit_window_seconds: 18_000,
|
||||
reset_after_seconds: 1_000,
|
||||
reset_at: 2_000_000_000,
|
||||
},
|
||||
secondary_window: {
|
||||
used_percent: 41,
|
||||
limit_window_seconds: 604_800,
|
||||
reset_after_seconds: 2_000,
|
||||
reset_at: 2_000_500_000,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
test("normalizes representative Plus and Pro payloads and hides absent windows", () => {
|
||||
expect(normalizeQuota(plus, 123)).toEqual({
|
||||
plan: "plus",
|
||||
fetchedAt: 123,
|
||||
windows: [
|
||||
{
|
||||
id: "codex:primary",
|
||||
label: null,
|
||||
seconds: 18_000,
|
||||
used: 12.5,
|
||||
resetsAt: 2_000_000_000_000,
|
||||
},
|
||||
{
|
||||
id: "codex:secondary",
|
||||
label: null,
|
||||
seconds: 604_800,
|
||||
used: 41,
|
||||
resetsAt: 2_000_500_000_000,
|
||||
},
|
||||
],
|
||||
});
|
||||
expect(normalizeQuota({ ...plus, plan_type: "pro" }, 123)?.windows).toHaveLength(2);
|
||||
expect(
|
||||
normalizeQuota({
|
||||
plan_type: "pro",
|
||||
rate_limit: { primary_window: plus.rate_limit.primary_window },
|
||||
})?.windows,
|
||||
).toEqual([
|
||||
{
|
||||
id: "codex:primary",
|
||||
label: null,
|
||||
seconds: 18_000,
|
||||
used: 12.5,
|
||||
resetsAt: 2_000_000_000_000,
|
||||
},
|
||||
]);
|
||||
expect(normalizeQuota({ plan_type: "plus", rate_limit: null })).toBeUndefined();
|
||||
});
|
||||
|
||||
test("normalizes additional named limits with stable IDs distinct from equal durations", () => {
|
||||
const normalized = normalizeQuota({
|
||||
...plus,
|
||||
additional_rate_limits: [
|
||||
{
|
||||
limit_name: "GPT-5.6 Luna",
|
||||
metered_feature: "codex_luna",
|
||||
rate_limit: { primary_window: plus.rate_limit.primary_window },
|
||||
},
|
||||
],
|
||||
});
|
||||
expect(normalized?.windows.map((window) => window.id)).toEqual([
|
||||
"codex:primary",
|
||||
"codex:secondary",
|
||||
"codex_luna:primary",
|
||||
]);
|
||||
expect(normalized?.windows[2]?.label).toBe("GPT-5.6 Luna");
|
||||
expect(
|
||||
normalizeQuota({
|
||||
plan_type: "pro",
|
||||
rate_limit: null,
|
||||
additional_rate_limits: [
|
||||
{
|
||||
limit_name: "GPT-5.6 Luna",
|
||||
metered_feature: "codex_luna",
|
||||
rate_limit: { primary_window: plus.rate_limit.primary_window },
|
||||
},
|
||||
],
|
||||
})?.windows.map((window) => window.id),
|
||||
).toEqual(["codex_luna:primary"]);
|
||||
});
|
||||
|
||||
function fixture(
|
||||
credential:
|
||||
| {
|
||||
type: "oauth";
|
||||
access: string;
|
||||
refresh: string;
|
||||
metadata?: { accountID?: string; accountId?: string };
|
||||
}
|
||||
| {
|
||||
type: "key";
|
||||
key: string;
|
||||
},
|
||||
payloads: unknown[],
|
||||
pollMs = 1_000_000,
|
||||
storageReads: Promise<Schema.Json | undefined>[] = [],
|
||||
) {
|
||||
const stored = new Map<string, Schema.Json>();
|
||||
const sessions = new Map([
|
||||
["root", {}],
|
||||
["child", { parentID: "root" }],
|
||||
]);
|
||||
let calls = 0;
|
||||
let storageCalls = 0;
|
||||
let changes = 0;
|
||||
let current: typeof credential | undefined = credential;
|
||||
const headers: Headers[] = [];
|
||||
const tracker = new QuotaTracker(
|
||||
// SAFETY: This focused fixture implements only the native calls exercised by
|
||||
// QuotaTracker; branded connection and session fields are irrelevant here.
|
||||
{
|
||||
integration: {
|
||||
connection: {
|
||||
async active() {
|
||||
// SAFETY: Only presence is observed; QuotaTracker passes this value back unchanged.
|
||||
return current ? ({ id: "connection" } as never) : undefined;
|
||||
},
|
||||
async resolve() {
|
||||
// SAFETY: The fixture covers the credential discriminant and fields QuotaTracker reads.
|
||||
return current as never;
|
||||
},
|
||||
},
|
||||
} as never,
|
||||
session: {
|
||||
async get({ sessionID }: { sessionID: string }) {
|
||||
// SAFETY: QuotaTracker reads only parentID from this focused fixture.
|
||||
return sessions.get(sessionID) as never;
|
||||
},
|
||||
} as never,
|
||||
storage: {
|
||||
async get(key: string) {
|
||||
storageCalls++;
|
||||
const deferred = storageReads.shift();
|
||||
if (deferred) return deferred;
|
||||
return stored.get(key);
|
||||
},
|
||||
async set(key: string, value: Schema.Json) {
|
||||
stored.set(key, value);
|
||||
},
|
||||
async remove(key: string) {
|
||||
stored.delete(key);
|
||||
},
|
||||
async scan() {
|
||||
return { entries: [] };
|
||||
},
|
||||
},
|
||||
},
|
||||
async () => {
|
||||
changes++;
|
||||
},
|
||||
async (_url, init) => {
|
||||
calls++;
|
||||
const requestHeaders = new Headers(init?.headers);
|
||||
headers.push(requestHeaders);
|
||||
expect(requestHeaders.get("authorization")).toStartWith("Bearer ");
|
||||
const payload = await (payloads.shift() ?? plus);
|
||||
if (payload instanceof Error) throw payload;
|
||||
return Response.json(payload);
|
||||
},
|
||||
pollMs,
|
||||
);
|
||||
return {
|
||||
tracker,
|
||||
calls: () => calls,
|
||||
storageCalls: () => storageCalls,
|
||||
changes: () => changes,
|
||||
headers,
|
||||
switch(credential: typeof current) {
|
||||
current = credential;
|
||||
},
|
||||
logout: () => {
|
||||
current = undefined;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
test("activates only for OAuth and establishes an awaited root baseline shared by children", async () => {
|
||||
const key = fixture({ type: "key", key: "secret" }, [plus]);
|
||||
await key.tracker.init();
|
||||
await key.tracker.start("root");
|
||||
expect(key.calls()).toBe(0);
|
||||
expect(await key.tracker.status("root")).toBeUndefined();
|
||||
await key.tracker.dispose();
|
||||
|
||||
const oauth = fixture(
|
||||
{
|
||||
type: "oauth",
|
||||
access: "opaque",
|
||||
refresh: "refresh",
|
||||
metadata: { accountID: "account-one" },
|
||||
},
|
||||
[
|
||||
plus,
|
||||
{
|
||||
...plus,
|
||||
rate_limit: {
|
||||
...plus.rate_limit,
|
||||
primary_window: { ...plus.rate_limit.primary_window, used_percent: 17.5 },
|
||||
},
|
||||
},
|
||||
],
|
||||
);
|
||||
await oauth.tracker.init();
|
||||
await oauth.tracker.start("root");
|
||||
expect(oauth.calls()).toBe(1);
|
||||
expect(oauth.headers[0]?.get("chatgpt-account-id")).toBe("account-one");
|
||||
await oauth.tracker.refresh();
|
||||
const child = await oauth.tracker.status("child");
|
||||
expect(child?.windows[0]?.sessionConsumed).toBe(5);
|
||||
expect(child?.windows[0]?.remaining).toBe(82.5);
|
||||
await oauth.tracker.dispose();
|
||||
});
|
||||
|
||||
test("polls, carries consumption across quota resets, tolerates errors, and clears on logout", async () => {
|
||||
const reset = {
|
||||
...plus,
|
||||
rate_limit: {
|
||||
...plus.rate_limit,
|
||||
primary_window: {
|
||||
...plus.rate_limit.primary_window,
|
||||
used_percent: 2,
|
||||
reset_at: 2_000_100_000,
|
||||
},
|
||||
},
|
||||
};
|
||||
const result = fixture(
|
||||
{
|
||||
type: "oauth",
|
||||
access: "opaque",
|
||||
refresh: "refresh",
|
||||
metadata: { accountID: "account-one" },
|
||||
},
|
||||
[plus, new Error("offline"), reset, reset, reset, reset],
|
||||
10,
|
||||
);
|
||||
await result.tracker.init();
|
||||
await result.tracker.start("root");
|
||||
await Bun.sleep(35);
|
||||
const status = await result.tracker.status("root");
|
||||
expect(result.calls()).toBeGreaterThanOrEqual(3);
|
||||
expect(status?.windows[0]?.sessionConsumed).toBe(2);
|
||||
result.logout();
|
||||
await result.tracker.authChanged(true);
|
||||
expect(await result.tracker.status("root")).toBeUndefined();
|
||||
const calls = result.calls();
|
||||
await result.tracker.dispose();
|
||||
await Bun.sleep(20);
|
||||
expect(result.calls()).toBe(calls);
|
||||
});
|
||||
|
||||
test("sidebar text places estimated percentage-point use beside latest remaining quota", () => {
|
||||
expect(
|
||||
formatQuotaWindow({
|
||||
seconds: 604_800,
|
||||
id: "codex:secondary",
|
||||
label: null,
|
||||
used: 25,
|
||||
remaining: 75,
|
||||
sessionConsumed: 2.25,
|
||||
resetsAt: 0,
|
||||
}),
|
||||
).toBe("7d: 75% left · ~2.3pp this session");
|
||||
});
|
||||
|
||||
test("inactive roots freeze, resume from a cached baseline on failure, and deactivate by family", async () => {
|
||||
const used = (percent: number) => ({
|
||||
...plus,
|
||||
rate_limit: {
|
||||
...plus.rate_limit,
|
||||
primary_window: { ...plus.rate_limit.primary_window, used_percent: percent },
|
||||
},
|
||||
});
|
||||
const result = fixture(
|
||||
{
|
||||
type: "oauth",
|
||||
access: "opaque",
|
||||
refresh: "refresh",
|
||||
metadata: { accountID: "account-one" },
|
||||
},
|
||||
[used(10), used(11), used(30), new Error("temporary"), used(30), used(33), used(35)],
|
||||
);
|
||||
await result.tracker.init();
|
||||
await result.tracker.start("root");
|
||||
await result.tracker.complete("root");
|
||||
expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(1);
|
||||
await result.tracker.refresh();
|
||||
expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(1);
|
||||
|
||||
await result.tracker.start("root");
|
||||
await result.tracker.start("child");
|
||||
await result.tracker.complete("root");
|
||||
await result.tracker.refresh();
|
||||
await result.tracker.complete("child");
|
||||
expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(6);
|
||||
await result.tracker.dispose();
|
||||
});
|
||||
|
||||
test("an auth switch detaches an in-flight old-account response and immediately fetches the new account", async () => {
|
||||
const pending = Promise.withResolvers<unknown>();
|
||||
const second = {
|
||||
...plus,
|
||||
account_id: "account-two",
|
||||
rate_limit: {
|
||||
...plus.rate_limit,
|
||||
primary_window: { ...plus.rate_limit.primary_window, used_percent: 70 },
|
||||
},
|
||||
};
|
||||
const result = fixture(
|
||||
{ type: "oauth", access: "old", refresh: "old", metadata: { accountID: "account-one" } },
|
||||
[
|
||||
pending.promise,
|
||||
second,
|
||||
{
|
||||
...second,
|
||||
rate_limit: {
|
||||
...second.rate_limit,
|
||||
primary_window: { ...second.rate_limit.primary_window, used_percent: 75 },
|
||||
},
|
||||
},
|
||||
],
|
||||
);
|
||||
await result.tracker.init();
|
||||
const old = result.tracker.start("root");
|
||||
while (result.calls() === 0) await Bun.sleep(1);
|
||||
result.switch({
|
||||
type: "oauth",
|
||||
access: "new",
|
||||
refresh: "new",
|
||||
metadata: { accountID: "account-two" },
|
||||
});
|
||||
await result.tracker.authChanged(true);
|
||||
expect(result.calls()).toBe(2);
|
||||
pending.resolve(plus);
|
||||
await old;
|
||||
expect((await result.tracker.status("root"))?.windows[0]?.used).toBe(70);
|
||||
expect(result.headers[1]?.get("chatgpt-account-id")).toBe("account-two");
|
||||
await result.tracker.refresh();
|
||||
expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(0);
|
||||
await result.tracker.dispose();
|
||||
});
|
||||
|
||||
test("an auth switch during an old account storage read cannot restore stale state or polling", async () => {
|
||||
const pending = Promise.withResolvers<Schema.Json | undefined>();
|
||||
const second = {
|
||||
...plus,
|
||||
account_id: "account-two",
|
||||
rate_limit: {
|
||||
...plus.rate_limit,
|
||||
primary_window: { ...plus.rate_limit.primary_window, used_percent: 64 },
|
||||
},
|
||||
};
|
||||
const result = fixture(
|
||||
{ type: "oauth", access: "old", refresh: "old", metadata: { accountID: "account-one" } },
|
||||
[second],
|
||||
1_000_000,
|
||||
[pending.promise],
|
||||
);
|
||||
const old = result.tracker.init();
|
||||
while (result.storageCalls() === 0) await Bun.sleep(1);
|
||||
result.switch({
|
||||
type: "oauth",
|
||||
access: "new",
|
||||
refresh: "new",
|
||||
metadata: { accountID: "account-two" },
|
||||
});
|
||||
await result.tracker.authChanged(true);
|
||||
expect(result.calls()).toBe(1);
|
||||
pending.resolve(undefined);
|
||||
await old;
|
||||
expect((await result.tracker.status("root"))?.windows[0]?.used).toBe(64);
|
||||
expect(result.headers[0]?.get("chatgpt-account-id")).toBe("account-two");
|
||||
await result.tracker.dispose();
|
||||
});
|
||||
|
||||
test("TUI quota loading converts RPC failures to an empty render state", async () => {
|
||||
expect(
|
||||
await loadQuotaStatus(async () => {
|
||||
throw new Error("disconnected");
|
||||
}),
|
||||
).toBeNull();
|
||||
});
|
||||
Reference in New Issue
Block a user