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<QuotaSnapshot> {
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<string> {
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<string, unknown>, 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<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}