aegida-console / lib / chat / quota.ts
quota.ts
Raw
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);
}