chore: reorg files
This commit is contained in:
@@ -0,0 +1,399 @@
|
||||
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";
|
||||
|
||||
function compactLimitIdentity(value: string) {
|
||||
return value.toLowerCase().replace(/[^a-z0-9]+/g, "");
|
||||
}
|
||||
|
||||
function isSparkLimit(meteredFeature?: string, limitName?: string | null) {
|
||||
return (
|
||||
(meteredFeature !== undefined &&
|
||||
compactLimitIdentity(meteredFeature) === "codexbengalfox") ||
|
||||
(limitName != null && compactLimitIdentity(limitName) === "gpt53codexspark")
|
||||
);
|
||||
}
|
||||
|
||||
function isReserveLimit(meteredFeature?: string, limitName?: string | null) {
|
||||
const hasReserveToken = (value: string) => /(^|[^a-z0-9])reserve($|[^a-z0-9])/i.test(value);
|
||||
return (
|
||||
(meteredFeature !== undefined && hasReserveToken(meteredFeature)) ||
|
||||
(limitName != null && hasReserveToken(limitName))
|
||||
);
|
||||
}
|
||||
|
||||
function isExcludedLimit(meteredFeature?: string, limitName?: string | null) {
|
||||
return isSparkLimit(meteredFeature, limitName) || isReserveLimit(meteredFeature, limitName);
|
||||
}
|
||||
|
||||
function isExcludedWindow(window: Pick<QuotaWindow, "id" | "label">) {
|
||||
return isExcludedLimit(window.id.split(":", 1)[0], window.label);
|
||||
}
|
||||
|
||||
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 }>;
|
||||
}
|
||||
|
||||
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(),
|
||||
});
|
||||
type StoredState = z.infer<typeof storedState>;
|
||||
|
||||
function withoutExcludedState(state: StoredState): StoredState {
|
||||
const windows = state.latest?.windows.filter((window) => !isExcludedWindow(window));
|
||||
return {
|
||||
...state,
|
||||
latest: state.latest && windows?.length ? { ...state.latest, windows } : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
type QuotaContext = Pick<Plugin.Context, "integration" | "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 additionalRateLimitIdentity = z.object({
|
||||
limit_name: z.string().optional(),
|
||||
metered_feature: z.string().optional(),
|
||||
});
|
||||
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.preprocess(
|
||||
(value) =>
|
||||
Array.isArray(value)
|
||||
? value.filter((item) => {
|
||||
const identity = additionalRateLimitIdentity.safeParse(item);
|
||||
return (
|
||||
!identity.success ||
|
||||
!isExcludedLimit(identity.data.metered_feature, identity.data.limit_name)
|
||||
);
|
||||
})
|
||||
: value,
|
||||
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 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() {
|
||||
await this.refresh();
|
||||
}
|
||||
|
||||
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 value = await this.context.storage.get(STORAGE_KEY);
|
||||
const parsed = storedState.safeParse(value);
|
||||
if (generation !== this.generation) return;
|
||||
const stored = parsed.success ? parsed.data : undefined;
|
||||
if (generation !== this.generation) return;
|
||||
this.state = stored?.account === account ? withoutExcludedState(stored) : { account };
|
||||
}
|
||||
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 };
|
||||
}
|
||||
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;
|
||||
this.state.latest = snapshot;
|
||||
}
|
||||
|
||||
async start(sessionID: string) {
|
||||
if (this.active.has(sessionID)) return;
|
||||
const generation = this.generation;
|
||||
const snapshot = await this.refresh();
|
||||
if (!snapshot || generation !== this.generation) return;
|
||||
this.active.add(sessionID);
|
||||
}
|
||||
|
||||
async complete(sessionID: string) {
|
||||
if (!this.active.delete(sessionID)) return;
|
||||
await this.refresh();
|
||||
}
|
||||
|
||||
status(): QuotaStatus | undefined {
|
||||
if (!this.state?.latest) return;
|
||||
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),
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
private async persist() {
|
||||
if (!this.state) return;
|
||||
const value: Schema.Json = this.state.latest
|
||||
? {
|
||||
account: this.state.account,
|
||||
latest: this.state.latest,
|
||||
}
|
||||
: { account: this.state.account };
|
||||
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;
|
||||
this.active.clear();
|
||||
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