feat: add quota tracking
This commit is contained in:
+429
@@ -0,0 +1,429 @@
|
||||
import type { Plugin } from "@opencode/plugin";
|
||||
import type { Credential } from "@opencode/schema/credential";
|
||||
import type { Schema } from "effect";
|
||||
import { z } from "zod";
|
||||
|
||||
const INTEGRATION_ID = "openai";
|
||||
const USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
|
||||
const POLL_MS = 60_000;
|
||||
const TIMEOUT_MS = 5_000;
|
||||
const STORAGE_KEY = "chatgpt-quota";
|
||||
|
||||
export interface QuotaWindow {
|
||||
id: string;
|
||||
label: string | null;
|
||||
seconds: number;
|
||||
used: number;
|
||||
resetsAt: number;
|
||||
}
|
||||
|
||||
export interface QuotaSnapshot {
|
||||
plan: string;
|
||||
fetchedAt: number;
|
||||
windows: QuotaWindow[];
|
||||
}
|
||||
|
||||
export interface QuotaStatus {
|
||||
plan: string;
|
||||
fetchedAt: number;
|
||||
windows: Array<QuotaWindow & { remaining: number; sessionConsumed: number }>;
|
||||
}
|
||||
|
||||
const storedState = z.object({
|
||||
account: z.string(),
|
||||
latest: z
|
||||
.object({
|
||||
plan: z.string(),
|
||||
fetchedAt: z.number(),
|
||||
windows: z.array(
|
||||
z.object({
|
||||
id: z.string(),
|
||||
label: z.string().nullable(),
|
||||
seconds: z.number(),
|
||||
used: z.number(),
|
||||
resetsAt: z.number(),
|
||||
}),
|
||||
),
|
||||
})
|
||||
.optional(),
|
||||
sessions: z.record(
|
||||
z.string(),
|
||||
z.record(
|
||||
z.string(),
|
||||
z.object({ resetsAt: z.number(), lastUsed: z.number(), consumed: z.number() }),
|
||||
),
|
||||
),
|
||||
});
|
||||
type StoredState = z.infer<typeof storedState>;
|
||||
|
||||
type QuotaContext = Pick<Plugin.Context, "integration" | "session" | "storage">;
|
||||
type QuotaFetch = (input: string | URL | Request, init?: RequestInit) => Promise<Response>;
|
||||
|
||||
const quotaWindow = z.object({
|
||||
limit_window_seconds: z.number().finite().positive(),
|
||||
used_percent: z.number().finite(),
|
||||
reset_at: z.number().finite(),
|
||||
});
|
||||
const quotaPayload = z.object({
|
||||
plan_type: z.string(),
|
||||
account_id: z.string().optional(),
|
||||
rate_limit: z
|
||||
.object({
|
||||
allowed: z.boolean().optional(),
|
||||
primary_window: quotaWindow.nullish(),
|
||||
secondary_window: quotaWindow.nullish(),
|
||||
})
|
||||
.nullish(),
|
||||
additional_rate_limits: z
|
||||
.array(
|
||||
z.object({
|
||||
limit_name: z.string(),
|
||||
metered_feature: z.string(),
|
||||
rate_limit: z
|
||||
.object({
|
||||
allowed: z.boolean().optional(),
|
||||
primary_window: quotaWindow.nullish(),
|
||||
secondary_window: quotaWindow.nullish(),
|
||||
})
|
||||
.nullish(),
|
||||
}),
|
||||
)
|
||||
.nullish(),
|
||||
});
|
||||
|
||||
function normalizeWindow(
|
||||
id: string,
|
||||
item: z.infer<typeof quotaWindow>,
|
||||
label?: string,
|
||||
): QuotaWindow {
|
||||
return {
|
||||
id,
|
||||
label: label ?? null,
|
||||
seconds: item.limit_window_seconds,
|
||||
used: Math.max(0, Math.min(100, item.used_percent)),
|
||||
resetsAt: item.reset_at < 10_000_000_000 ? item.reset_at * 1000 : item.reset_at,
|
||||
};
|
||||
}
|
||||
|
||||
/** Normalize valid windows from the documented Codex usage payload. */
|
||||
export function normalizeQuota(
|
||||
input: Schema.Json,
|
||||
fetchedAt = Date.now(),
|
||||
): QuotaSnapshot | undefined {
|
||||
const parsed = quotaPayload.safeParse(input);
|
||||
if (!parsed.success) return;
|
||||
const plan = parsed.data.plan_type.toLowerCase();
|
||||
const windows: QuotaWindow[] = [];
|
||||
if (parsed.data.rate_limit)
|
||||
for (const [position, window] of [
|
||||
["primary", parsed.data.rate_limit.primary_window],
|
||||
["secondary", parsed.data.rate_limit.secondary_window],
|
||||
] as const)
|
||||
if (window) windows.push(normalizeWindow(`codex:${position}`, window));
|
||||
for (const additional of parsed.data.additional_rate_limits ?? []) {
|
||||
if (!additional.rate_limit) continue;
|
||||
for (const [position, window] of [
|
||||
["primary", additional.rate_limit.primary_window],
|
||||
["secondary", additional.rate_limit.secondary_window],
|
||||
] as const)
|
||||
if (window)
|
||||
windows.push(
|
||||
normalizeWindow(
|
||||
`${additional.metered_feature}:${position}`,
|
||||
window,
|
||||
additional.limit_name,
|
||||
),
|
||||
);
|
||||
}
|
||||
if (windows.length === 0) return;
|
||||
return { plan, fetchedAt, windows };
|
||||
}
|
||||
|
||||
function accountID(credential: Credential.OAuth) {
|
||||
const metadataID = z
|
||||
.string()
|
||||
.safeParse(credential.metadata?.accountID ?? credential.metadata?.accountId);
|
||||
if (metadataID.success) return metadataID.data;
|
||||
const parts = credential.access.split(".");
|
||||
if (parts.length !== 3) return;
|
||||
try {
|
||||
const claims = z
|
||||
.object({
|
||||
chatgpt_account_id: z.string().optional(),
|
||||
"https://api.openai.com/auth": z
|
||||
.object({ chatgpt_account_id: z.string().optional() })
|
||||
.optional(),
|
||||
organizations: z.array(z.object({ id: z.string() })).optional(),
|
||||
})
|
||||
.safeParse(JSON.parse(Buffer.from(parts[1]!, "base64url").toString()));
|
||||
if (!claims.success) return;
|
||||
return (
|
||||
claims.data.chatgpt_account_id ??
|
||||
claims.data["https://api.openai.com/auth"]?.chatgpt_account_id ??
|
||||
claims.data.organizations?.[0]?.id
|
||||
);
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
async function fingerprint(value: string) {
|
||||
const bytes = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(value));
|
||||
return Buffer.from(bytes).toString("hex");
|
||||
}
|
||||
|
||||
export class QuotaTracker {
|
||||
private state: StoredState | undefined;
|
||||
private timer: ReturnType<typeof setInterval> | undefined;
|
||||
private request: Promise<QuotaSnapshot | undefined> | undefined;
|
||||
private readonly requests = new Set<Promise<QuotaSnapshot | undefined>>();
|
||||
private readonly active = new Map<string, Set<string>>();
|
||||
private controller: AbortController | undefined;
|
||||
private generation = 0;
|
||||
private disposed = false;
|
||||
|
||||
constructor(
|
||||
private readonly context: QuotaContext,
|
||||
private readonly changed: () => Promise<void>,
|
||||
private readonly fetcher: QuotaFetch = fetch,
|
||||
private readonly pollMs = POLL_MS,
|
||||
) {}
|
||||
|
||||
async init() {
|
||||
try {
|
||||
await this.resolveCredential();
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
private async resolveCredential(generation = this.generation) {
|
||||
const connection = await this.context.integration.connection.active(INTEGRATION_ID);
|
||||
if (generation !== this.generation) return;
|
||||
const credential = connection
|
||||
? await this.context.integration.connection.resolve(connection)
|
||||
: undefined;
|
||||
if (generation !== this.generation) return;
|
||||
if (credential?.type !== "oauth") {
|
||||
if (this.state) await this.clear();
|
||||
return;
|
||||
}
|
||||
const id = accountID(credential);
|
||||
const account = id
|
||||
? await fingerprint(`account:${id}`)
|
||||
: (this.state?.account ?? (await fingerprint(`credential:${credential.refresh}`)));
|
||||
if (generation !== this.generation) return;
|
||||
if (this.state?.account !== account) {
|
||||
const parsed = storedState.safeParse(await this.context.storage.get(STORAGE_KEY));
|
||||
if (generation !== this.generation) return;
|
||||
const stored = parsed.success ? parsed.data : undefined;
|
||||
if (generation !== this.generation) return;
|
||||
this.state = stored?.account === account ? stored : { account, sessions: {} };
|
||||
}
|
||||
if (generation !== this.generation) return;
|
||||
if (!this.timer) this.timer = setInterval(() => void this.refresh(), this.pollMs);
|
||||
return credential;
|
||||
}
|
||||
|
||||
private async clear() {
|
||||
this.state = undefined;
|
||||
if (this.timer) clearInterval(this.timer);
|
||||
this.timer = undefined;
|
||||
await this.context.storage.remove(STORAGE_KEY);
|
||||
await this.changed();
|
||||
}
|
||||
|
||||
async refresh(): Promise<QuotaSnapshot | undefined> {
|
||||
const generation = this.generation;
|
||||
if (this.disposed) return;
|
||||
if (this.request) return this.request;
|
||||
const request = this.query();
|
||||
this.request = request;
|
||||
this.requests.add(request);
|
||||
try {
|
||||
const result = await request;
|
||||
if (generation !== this.generation) return;
|
||||
return result ?? this.state?.latest;
|
||||
} finally {
|
||||
this.requests.delete(request);
|
||||
if (this.request === request) this.request = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private async query() {
|
||||
const generation = this.generation;
|
||||
try {
|
||||
const credential = await this.resolveCredential(generation);
|
||||
if (!credential || !this.state || generation !== this.generation) return;
|
||||
const controller = new AbortController();
|
||||
this.controller = controller;
|
||||
const timeout = setTimeout(() => controller.abort(), TIMEOUT_MS);
|
||||
let response: Response;
|
||||
try {
|
||||
const headers = new Headers({
|
||||
Authorization: `Bearer ${credential.access}`,
|
||||
Accept: "application/json",
|
||||
});
|
||||
const id = accountID(credential);
|
||||
if (id) headers.set("ChatGPT-Account-Id", id);
|
||||
response = await this.fetcher(USAGE_URL, { headers, signal: controller.signal });
|
||||
} finally {
|
||||
clearTimeout(timeout);
|
||||
if (this.controller === controller) this.controller = undefined;
|
||||
}
|
||||
if (!response.ok || generation !== this.generation) return;
|
||||
// SAFETY: JSON.parse semantics from Response.json produce only JSON values;
|
||||
// quotaPayload validates all fields used below.
|
||||
const payload = (await response.json()) as Schema.Json;
|
||||
if (generation !== this.generation || !this.state) return;
|
||||
const parsed = quotaPayload.safeParse(payload);
|
||||
if (parsed.success && parsed.data.account_id) {
|
||||
const account = await fingerprint(`account:${parsed.data.account_id}`);
|
||||
if (generation !== this.generation || !this.state) return;
|
||||
if (account !== this.state.account) this.state = { account, sessions: {} };
|
||||
}
|
||||
const snapshot = normalizeQuota(payload);
|
||||
if (!snapshot) return;
|
||||
this.apply(snapshot);
|
||||
await this.persist();
|
||||
await this.changed();
|
||||
return snapshot;
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
private apply(snapshot: QuotaSnapshot) {
|
||||
if (!this.state) return;
|
||||
for (const root of this.active.keys()) {
|
||||
const tracked = (this.state.sessions[root] ??= {});
|
||||
for (const window of snapshot.windows) {
|
||||
const previous = tracked[window.id];
|
||||
if (!previous) {
|
||||
tracked[window.id] = {
|
||||
resetsAt: window.resetsAt,
|
||||
lastUsed: window.used,
|
||||
consumed: 0,
|
||||
};
|
||||
continue;
|
||||
}
|
||||
if (previous.resetsAt !== window.resetsAt) {
|
||||
previous.resetsAt = window.resetsAt;
|
||||
previous.lastUsed = window.used;
|
||||
previous.consumed += window.used;
|
||||
} else if (window.used >= previous.lastUsed) {
|
||||
previous.consumed += window.used - previous.lastUsed;
|
||||
previous.lastUsed = window.used;
|
||||
}
|
||||
}
|
||||
}
|
||||
this.state.latest = snapshot;
|
||||
}
|
||||
|
||||
private async root(sessionID: string) {
|
||||
let current = sessionID;
|
||||
const seen = new Set<string>();
|
||||
while (!seen.has(current)) {
|
||||
seen.add(current);
|
||||
const session = await this.context.session.get({ sessionID: current });
|
||||
if (!session.parentID) return current;
|
||||
current = session.parentID;
|
||||
}
|
||||
return sessionID;
|
||||
}
|
||||
|
||||
async start(sessionID: string) {
|
||||
const generation = this.generation;
|
||||
try {
|
||||
const root = await this.root(sessionID);
|
||||
if (generation !== this.generation) return;
|
||||
const family = this.active.get(root);
|
||||
if (family?.has(sessionID)) return;
|
||||
const snapshot = await this.refresh();
|
||||
if (!snapshot || !this.state || generation !== this.generation) return;
|
||||
if (!family) {
|
||||
const previous = this.state.sessions[root] ?? {};
|
||||
this.state.sessions[root] = Object.fromEntries(
|
||||
snapshot.windows.map((window) => [
|
||||
window.id,
|
||||
{
|
||||
resetsAt: window.resetsAt,
|
||||
lastUsed: window.used,
|
||||
consumed: previous[window.id]?.consumed ?? 0,
|
||||
},
|
||||
]),
|
||||
);
|
||||
}
|
||||
await this.persist();
|
||||
if (generation !== this.generation) return;
|
||||
await this.changed();
|
||||
if (generation !== this.generation) return;
|
||||
const active = family ?? new Set<string>();
|
||||
active.add(sessionID);
|
||||
this.active.set(root, active);
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
async complete(sessionID: string) {
|
||||
try {
|
||||
const root = await this.root(sessionID);
|
||||
const family = this.active.get(root);
|
||||
if (!family?.has(sessionID)) return;
|
||||
await this.refresh();
|
||||
family.delete(sessionID);
|
||||
if (family.size === 0) this.active.delete(root);
|
||||
await this.persist();
|
||||
await this.changed();
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
async status(sessionID: string): Promise<QuotaStatus | undefined> {
|
||||
if (!this.state?.latest) return;
|
||||
const tracked = this.state.sessions[await this.root(sessionID)];
|
||||
return {
|
||||
plan: this.state.latest.plan,
|
||||
fetchedAt: this.state.latest.fetchedAt,
|
||||
windows: this.state.latest.windows.map((window) => ({
|
||||
...window,
|
||||
remaining: Math.max(0, 100 - window.used),
|
||||
sessionConsumed: tracked?.[window.id]?.consumed ?? 0,
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
private async persist() {
|
||||
if (!this.state) return;
|
||||
const value: Schema.Json = this.state.latest
|
||||
? {
|
||||
account: this.state.account,
|
||||
sessions: this.state.sessions,
|
||||
latest: this.state.latest,
|
||||
}
|
||||
: { account: this.state.account, sessions: this.state.sessions };
|
||||
await this.context.storage.set(STORAGE_KEY, value);
|
||||
}
|
||||
|
||||
async authChanged(switched = false) {
|
||||
if (switched) {
|
||||
this.generation++;
|
||||
this.controller?.abort();
|
||||
this.controller = undefined;
|
||||
this.request = undefined;
|
||||
if (this.state) await this.clear();
|
||||
}
|
||||
await this.refresh();
|
||||
}
|
||||
|
||||
async dispose() {
|
||||
this.disposed = true;
|
||||
this.generation++;
|
||||
this.controller?.abort();
|
||||
if (this.timer) clearInterval(this.timer);
|
||||
this.timer = undefined;
|
||||
await Promise.all(this.requests);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user