import "server-only"; import { isModelId, type ModelId } from "@/lib/chat/types"; const maximumQuotaBytes = 1_048_576; const bearerPattern = /^Bearer [^\s,]+$/; const quotaKeys = [ "object", "model", "available_tokens", "unlimited", "exhausted", "resets_at", "as_of", ]; export type QuotaSnapshot = { object: "aegida.quota"; model: ModelId; availableTokens: number | null; unlimited: boolean; exhausted: boolean; resetsAt: string | null; asOf: string; }; export type QuotaErrorReason = "unauthorized" | "unavailable"; export class QuotaError extends Error { constructor(readonly reason: QuotaErrorReason = "unavailable") { super("Aegida Gate quota is unavailable"); this.name = "QuotaError"; } } export async function getGateQuota( model: ModelId, authorization: string, signal: AbortSignal, ): Promise { if (!isModelId(model) || !bearerPattern.test(authorization)) { throw new QuotaError(); } let response: Response; try { response = await fetch( `${gateOrigin()}/v1/aegida/quota?model=${encodeURIComponent(model)}`, { method: "GET", headers: { Accept: "application/json", Authorization: authorization }, cache: "no-store", redirect: "error", signal, }, ); } catch (error) { if (signal.aborted || isAbortError(error)) throw abortError(); throw new QuotaError(); } if (response.status === 401) throw new QuotaError("unauthorized"); if (!response.ok) throw new QuotaError(); try { return parseGateQuota( JSON.parse(await readBoundedText(response, maximumQuotaBytes)) as unknown, model, ); } catch (error) { if (error instanceof QuotaError) throw error; if (signal.aborted || isAbortError(error)) throw abortError(); throw new QuotaError(); } } export function parseGateQuota(value: unknown, model: ModelId): QuotaSnapshot { if ( !isRecord(value) || !hasOnlyKeys(value, quotaKeys) || value.object !== "aegida.quota" || value.model !== model || typeof value.unlimited !== "boolean" || typeof value.exhausted !== "boolean" || (value.resets_at !== null && !isRfc3339(value.resets_at)) || !isRfc3339(value.as_of) ) { throw new QuotaError(); } const tokens = value.available_tokens; if (value.unlimited) { if (tokens !== null || value.exhausted || value.resets_at !== null) { throw new QuotaError(); } } else if ( !Number.isSafeInteger(tokens) || (tokens as number) < 0 || value.exhausted !== (tokens === 0) ) { throw new QuotaError(); } return { object: "aegida.quota", model, availableTokens: tokens as number | null, unlimited: value.unlimited, exhausted: value.exhausted, resetsAt: value.resets_at, asOf: value.as_of, }; } function isRfc3339(value: unknown): value is string { if (typeof value !== "string") return false; const match = /^(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})(?:\.\d{1,9})?(?:Z|([+-])(\d{2}):(\d{2}))$/.exec(value); if (!match) return false; const [, year, month, day, hour, minute, second, , offsetHour, offsetMinute] = match; const yearValue = Number(year); const monthValue = Number(month); const dayValue = Number(day); return ( monthValue >= 1 && monthValue <= 12 && dayValue >= 1 && dayValue <= new Date(Date.UTC(yearValue, monthValue, 0)).getUTCDate() && Number(hour) <= 23 && Number(minute) <= 59 && Number(second) <= 59 && (offsetHour === undefined || (Number(offsetHour) <= 23 && Number(offsetMinute) <= 59)) ); } function gateOrigin(): string { const raw = process.env.AEGIDA_GATE_URL; if (!raw) throw new QuotaError(); try { const url = new URL(raw); if ( (url.protocol !== "http:" && url.protocol !== "https:") || !url.hostname || url.username || url.password || (url.pathname !== "" && url.pathname !== "/") || url.search || url.hash ) { throw new Error("invalid origin"); } return url.origin; } catch { throw new QuotaError(); } } async function readBoundedText(response: Response, maximum: number): Promise { if (!response.body) throw new QuotaError(); const reader = response.body.getReader(); const decoder = new TextDecoder(); let text = ""; let bytes = 0; while (true) { const { done, value } = await reader.read(); if (done) break; bytes += value.byteLength; if (bytes > maximum) { await reader.cancel(); throw new QuotaError(); } text += decoder.decode(value, { stream: true }); } return text + decoder.decode(); } function hasOnlyKeys(value: Record, allowed: string[]): boolean { const keys = Object.keys(value); return keys.length === allowed.length && keys.every((key) => allowed.includes(key)); } function abortError(): DOMException { return new DOMException("The operation was aborted", "AbortError"); } function isAbortError(value: unknown): boolean { return value instanceof DOMException && value.name === "AbortError"; } function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); }