279 lines
9.2 KiB
TypeScript
279 lines
9.2 KiB
TypeScript
import { expect, test } from "bun:test";
|
|
import type { Schema } from "effect";
|
|
import { normalizeQuota, QuotaTracker } from "../src/quota/quota.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 Plus and Pro payloads, hiding absent windows and excluded Spark and reserve limits", () => {
|
|
{
|
|
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();
|
|
}
|
|
{
|
|
const normalized = normalizeQuota({
|
|
...plus,
|
|
additional_rate_limits: [
|
|
{
|
|
limit_name: "gpt 5_3 codex spark",
|
|
metered_feature: "CODEX-BENGALFOX",
|
|
rate_limit: {
|
|
primary_window: { used_percent: "invalid" },
|
|
},
|
|
},
|
|
{
|
|
limit_name: "Luna Reserve",
|
|
metered_feature: "codex_luna",
|
|
rate_limit: { primary_window: { used_percent: "invalid" } },
|
|
},
|
|
{
|
|
limit_name: "Luna overflow",
|
|
metered_feature: "codex_luna_reserve",
|
|
rate_limit: { primary_window: { used_percent: "invalid" } },
|
|
},
|
|
{
|
|
limit_name: "GPT-5.6 Luna",
|
|
metered_feature: "codex_luna",
|
|
rate_limit: { primary_window: plus.rate_limit.primary_window },
|
|
},
|
|
{
|
|
limit_name: "Preserved capacity",
|
|
metered_feature: "codex_preserved",
|
|
rate_limit: { primary_window: plus.rate_limit.primary_window },
|
|
},
|
|
],
|
|
});
|
|
|
|
expect(normalized?.windows.map((window) => window.id)).toEqual([
|
|
"codex:primary",
|
|
"codex:secondary",
|
|
"codex_luna:primary",
|
|
"codex_preserved:primary",
|
|
]);
|
|
expect(normalized?.windows.slice(2).map((window) => window.label)).toEqual([
|
|
"GPT-5.6 Luna",
|
|
"Preserved capacity",
|
|
]);
|
|
}
|
|
});
|
|
|
|
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>();
|
|
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,
|
|
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,
|
|
stored,
|
|
switch(next: typeof current) {
|
|
current = next;
|
|
},
|
|
logout: () => {
|
|
current = undefined;
|
|
},
|
|
};
|
|
}
|
|
|
|
test("activates only for OAuth and refreshes quota at initialization and request start", async () => {
|
|
const key = fixture({ type: "key", key: "secret" }, [plus]);
|
|
await key.tracker.init();
|
|
await key.tracker.start("root");
|
|
expect(key.calls()).toBe(0);
|
|
expect(key.tracker.status()).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();
|
|
expect(oauth.calls()).toBe(1);
|
|
expect(oauth.tracker.status()?.windows[0]?.remaining).toBe(87.5);
|
|
await oauth.tracker.start("root");
|
|
expect(oauth.calls()).toBe(2);
|
|
expect(oauth.headers[0]?.get("chatgpt-account-id")).toBe("account-one");
|
|
const status = oauth.tracker.status();
|
|
expect(status?.windows[0]?.remaining).toBe(82.5);
|
|
expect(status?.windows[0]).not.toHaveProperty("sessionConsumed");
|
|
expect(oauth.stored.get("chatgpt-quota")).not.toHaveProperty("sessions");
|
|
await oauth.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 },
|
|
},
|
|
},
|
|
],
|
|
);
|
|
const old = result.tracker.init();
|
|
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(result.tracker.status()?.windows[0]?.used).toBe(70);
|
|
expect(result.headers[1]?.get("chatgpt-account-id")).toBe("account-two");
|
|
await result.tracker.refresh();
|
|
expect(result.tracker.status()?.windows[0]).not.toHaveProperty("sessionConsumed");
|
|
await result.tracker.dispose();
|
|
});
|