+69
-45
@@ -8,24 +8,30 @@ 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";
|
||||
const jsonValue = z.json();
|
||||
|
||||
function compactLimitIdentity(value: string) {
|
||||
return value.toLowerCase().replace(/[^a-z0-9]+/g, "");
|
||||
return value.toLowerCase().replaceAll(/[^a-z0-9]+/gu, "");
|
||||
}
|
||||
|
||||
function isSparkLimit(meteredFeature?: string, limitName?: string | null) {
|
||||
return (
|
||||
(meteredFeature !== undefined &&
|
||||
compactLimitIdentity(meteredFeature) === "codexbengalfox") ||
|
||||
(limitName != null && compactLimitIdentity(limitName) === "gpt53codexspark")
|
||||
(limitName !== null &&
|
||||
limitName !== undefined &&
|
||||
compactLimitIdentity(limitName) === "gpt53codexspark")
|
||||
);
|
||||
}
|
||||
|
||||
function hasReserveToken(value: string) {
|
||||
return /(^|[^a-z0-9])reserve($|[^a-z0-9])/iu.test(value);
|
||||
}
|
||||
|
||||
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))
|
||||
(limitName !== null && limitName !== undefined && hasReserveToken(limitName))
|
||||
);
|
||||
}
|
||||
|
||||
@@ -81,7 +87,10 @@ 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,
|
||||
latest:
|
||||
state.latest && windows !== undefined && windows.length > 0
|
||||
? { ...state.latest, windows }
|
||||
: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -89,9 +98,9 @@ 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(),
|
||||
limit_window_seconds: z.number().positive(),
|
||||
used_percent: z.number(),
|
||||
reset_at: z.number(),
|
||||
});
|
||||
const additionalRateLimitIdentity = z.object({
|
||||
limit_name: z.string().optional(),
|
||||
@@ -187,7 +196,7 @@ export function normalizeQuota(
|
||||
function accountID(credential: Credential.OAuth) {
|
||||
const metadataID = z
|
||||
.string()
|
||||
.safeParse(credential.metadata?.accountID ?? credential.metadata?.accountId);
|
||||
.safeParse(credential.metadata?.["accountID"] ?? credential.metadata?.["accountId"]);
|
||||
if (metadataID.success) return metadataID.data;
|
||||
const parts = credential.access.split(".");
|
||||
if (parts.length !== 3) return;
|
||||
@@ -200,7 +209,7 @@ function accountID(credential: Credential.OAuth) {
|
||||
.optional(),
|
||||
organizations: z.array(z.object({ id: z.string() })).optional(),
|
||||
})
|
||||
.safeParse(JSON.parse(Buffer.from(parts[1]!, "base64url").toString()));
|
||||
.safeParse(JSON.parse(Buffer.from(parts[1] ?? "", "base64url").toString()));
|
||||
if (!claims.success) return;
|
||||
return (
|
||||
claims.data.chatgpt_account_id ??
|
||||
@@ -250,23 +259,30 @@ export class QuotaTracker {
|
||||
return;
|
||||
}
|
||||
const id = accountID(credential);
|
||||
const account = id
|
||||
? await fingerprint(`account:${id}`)
|
||||
: (this.state?.account ?? (await fingerprint(`credential:${credential.refresh}`)));
|
||||
const account =
|
||||
id !== undefined && 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 (this.state?.account !== account) await this.load(account, generation);
|
||||
if (generation !== this.generation) return;
|
||||
if (!this.timer) this.timer = setInterval(() => void this.refresh(), this.pollMs);
|
||||
this.timer ??= setInterval(() => void this.refresh(), this.pollMs);
|
||||
return credential;
|
||||
}
|
||||
|
||||
private async load(account: string, generation: number) {
|
||||
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;
|
||||
this.state = stored?.account === account ? withoutExcludedState(stored) : { account };
|
||||
}
|
||||
|
||||
/** Reads state through a call so checks after `await` are not narrowed away. */
|
||||
private current(): StoredState | undefined {
|
||||
return this.state;
|
||||
}
|
||||
|
||||
private async clear() {
|
||||
this.state = undefined;
|
||||
if (this.timer) clearInterval(this.timer);
|
||||
@@ -297,32 +313,20 @@ export class QuotaTracker {
|
||||
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;
|
||||
}
|
||||
const response = await this.fetchUsage(credential);
|
||||
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 payload = jsonValue.parse(await response.json());
|
||||
if (generation !== this.generation || !this.current()) return;
|
||||
const parsed = quotaPayload.safeParse(payload);
|
||||
if (parsed.success && parsed.data.account_id) {
|
||||
if (
|
||||
parsed.success &&
|
||||
parsed.data.account_id !== undefined &&
|
||||
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 state = this.current();
|
||||
if (generation !== this.generation || !state) return;
|
||||
if (account !== state.account) this.state = { account };
|
||||
}
|
||||
const snapshot = normalizeQuota(payload);
|
||||
if (!snapshot) return;
|
||||
@@ -335,6 +339,26 @@ export class QuotaTracker {
|
||||
}
|
||||
}
|
||||
|
||||
private async fetchUsage(credential: Credential.OAuth) {
|
||||
const controller = new AbortController();
|
||||
this.controller = controller;
|
||||
const timeout = setTimeout(() => {
|
||||
controller.abort();
|
||||
}, TIMEOUT_MS);
|
||||
try {
|
||||
const headers = new Headers({
|
||||
Authorization: `Bearer ${credential.access}`,
|
||||
Accept: "application/json",
|
||||
});
|
||||
const id = accountID(credential);
|
||||
if (id !== undefined && id !== "") headers.set("ChatGPT-Account-Id", id);
|
||||
return await this.fetcher(USAGE_URL, { headers, signal: controller.signal });
|
||||
} finally {
|
||||
clearTimeout(timeout);
|
||||
if (this.controller === controller) this.controller = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private apply(snapshot: QuotaSnapshot) {
|
||||
if (!this.state) return;
|
||||
this.state.latest = snapshot;
|
||||
|
||||
Reference in New Issue
Block a user