aegida-console / lib / chat / gateway.ts
gateway.ts
Raw
import "server-only";

import type { ChatRequest } from "@/lib/chat/types";

const maximumEventBytes = 1_048_576;
const maximumErrorBytes = 8_192;
const maximumJsonDepth = 64;
const bearerPattern = /^Bearer [^\s,]+$/;
const richErrorKeys = ["message", "type", "param", "code"];
const richGateErrors: Record<
  string,
  Array<{ message: string; type: string; param: string | null }>
> = {
  quota_exceeded: [{
    message: "Quota exceeded",
    type: "rate_limit_error",
    param: null,
  }],
  model_not_allowed: [{
    message: "Model access denied",
    type: "permission_error",
    param: "model",
  }],
  model_not_found: [{
    message: "Model not found",
    type: "invalid_request_error",
    param: "model",
  }],
  provider_unavailable: [{
    message: "Inference provider is temporarily unavailable",
    type: "server_error",
    param: null,
  }],
  service_unavailable: [
    {
      message: "Inference service is temporarily unavailable",
      type: "server_error",
      param: null,
    },
    {
      message: "Model catalog is temporarily unavailable",
      type: "server_error",
      param: null,
    },
  ],
};

type GatewayStatus = 403 | 404 | 429 | 502 | 503;

export class GatewayError extends Error {
  constructor(
    readonly status: GatewayStatus,
    message: string,
  ) {
    super(message);
    this.name = "GatewayError";
  }
}

export async function streamGatewayResponse(
  request: ChatRequest,
  authorization: string,
  signal: AbortSignal,
): Promise<ReadableStream<Uint8Array>> {
  if (!bearerPattern.test(authorization)) {
    throw new GatewayError(503, "Aegida Gate не настроен");
  }

  const endpoint = `${gateOrigin()}/v1/chat/completions`;
  let response: Response;
  try {
    response = await fetch(endpoint, {
      method: "POST",
      headers: {
        Authorization: authorization,
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        model: request.model,
        stream: true,
        stream_options: { include_usage: true },
        messages: request.messages,
      }),
      cache: "no-store",
      redirect: "error",
      signal,
    });
  } catch (error) {
    if (signal.aborted || isAbortError(error)) throw abortError();
    throw unavailableError();
  }

  if (!response.ok) throw await mapGatewayError(response, signal);
  if (!response.body) throw unavailableError();

  return translateOpenAiSse(response.body, signal);
}

async function mapGatewayError(
  response: Response,
  signal: AbortSignal,
): Promise<GatewayError> {
  if (![403, 404, 429, 502, 503].includes(response.status)) {
    await response.body?.cancel().catch(() => undefined);
    return unavailableError();
  }

  let code: string | null = null;
  try {
    const body = await readBoundedErrorBody(response, signal);
    code = parseExactGateError(body);
    if (!code) return unavailableError();
  } catch (error) {
    if (signal.aborted || isAbortError(error)) throw abortError();
    return unavailableError();
  }

  if (response.status === 429 && code === "quota_exceeded") {
    return new GatewayError(429, "Лимит токенов исчерпан");
  }
  if (
    response.status === 403 &&
    (code === "subject_inactive" || code === "model_not_allowed")
  ) {
    return new GatewayError(403, "Нет доступа к модели");
  }
  if (
    response.status === 404 &&
    (code === "model_not_found" || code === "route_unavailable")
  ) {
    return new GatewayError(404, "Модель недоступна");
  }
  if (response.status === 502 && code === "provider_unavailable") {
    return unavailableError(502);
  }
  if (response.status === 503 && code === "service_unavailable") {
    return unavailableError(503);
  }
  return unavailableError();
}

function parseExactGateError(body: string): string | null {
  const value = parseDuplicateAwareJson(body);
  const error = isRecord(value) ? value.error : null;
  if (
    !isRecord(value) ||
    Object.keys(value).length !== 1 ||
    !isRecord(error) ||
    typeof error.code !== "string"
  ) {
    return null;
  }

  const errorKeys = Object.keys(error);
  if (errorKeys.length === 1 && errorKeys[0] === "code") {
    return error.code;
  }
  if (
    errorKeys.length !== richErrorKeys.length ||
    !errorKeys.every((key) => richErrorKeys.includes(key))
  ) {
    return null;
  }

  const expectedTuples = richGateErrors[error.code];
  if (
    !expectedTuples?.some((expected) =>
      error.message === expected.message &&
      error.type === expected.type &&
      error.param === expected.param)
  ) {
    return null;
  }
  return error.code;
}

async function readBoundedErrorBody(
  response: Response,
  signal: AbortSignal,
): Promise<string> {
  if (!response.body) throw unavailableError();
  const reader = response.body.getReader();
  const decoder = new TextDecoder("utf-8", { fatal: true });
  let body = "";
  let bytes = 0;
  let cancelPromise: Promise<void> | null = null;
  const cancelReader = (reason?: unknown) => {
    cancelPromise ??= reader.cancel(reason).catch(() => undefined);
    return cancelPromise;
  };
  const onAbort = () => {
    void cancelReader(abortError());
  };
  signal.addEventListener("abort", onAbort, { once: true });
  try {
    while (true) {
      if (signal.aborted) throw abortError();
      const { done, value } = await reader.read();
      if (signal.aborted) throw abortError();
      if (done) break;
      bytes += value.byteLength;
      if (bytes > maximumErrorBytes) {
        await cancelReader();
        throw unavailableError();
      }
      body += decoder.decode(value, { stream: true });
    }
    return body + decoder.decode();
  } catch (error) {
    await cancelReader(error);
    throw error;
  } finally {
    signal.removeEventListener("abort", onAbort);
    if (cancelPromise) await cancelPromise;
    reader.releaseLock();
  }
}

function parseDuplicateAwareJson(body: string): unknown {
  let index = 0;

  const fail = (): never => {
    throw new SyntaxError("Invalid JSON");
  };
  const skipWhitespace = () => {
    while (/[\t\n\r ]/.test(body[index] ?? "")) index += 1;
  };
  const parseString = (): string => {
    if (body[index] !== '"') fail();
    index += 1;
    let value = "";
    while (index < body.length) {
      const character = body[index];
      index += 1;
      if (character === '"') return value;
      if (character === "\\") {
        const escape = body[index];
        index += 1;
        const simpleEscapes: Record<string, string> = {
          '"': '"',
          "\\": "\\",
          "/": "/",
          b: "\b",
          f: "\f",
          n: "\n",
          r: "\r",
          t: "\t",
        };
        if (escape in simpleEscapes) {
          value += simpleEscapes[escape];
          continue;
        }
        if (escape !== "u") fail();
        const firstHex = body.slice(index, index + 4);
        if (!/^[0-9a-fA-F]{4}$/.test(firstHex)) fail();
        index += 4;
        const first = Number.parseInt(firstHex, 16);
        if (first >= 0xd800 && first <= 0xdbff) {
          if (body.slice(index, index + 2) !== "\\u") fail();
          index += 2;
          const secondHex = body.slice(index, index + 4);
          if (!/^[0-9a-fA-F]{4}$/.test(secondHex)) fail();
          index += 4;
          const second = Number.parseInt(secondHex, 16);
          if (second < 0xdc00 || second > 0xdfff) fail();
          value += String.fromCodePoint(
            0x10000 + ((first - 0xd800) << 10) + (second - 0xdc00),
          );
          continue;
        }
        if (first >= 0xdc00 && first <= 0xdfff) fail();
        value += String.fromCharCode(first);
        continue;
      }
      const code = character.charCodeAt(0);
      if (code <= 0x1f) fail();
      if (code >= 0xd800 && code <= 0xdbff) {
        const next = body.charCodeAt(index);
        if (next < 0xdc00 || next > 0xdfff) fail();
        value += character + body[index];
        index += 1;
        continue;
      }
      if (code >= 0xdc00 && code <= 0xdfff) fail();
      value += character;
    }
    return fail();
  };
  const parseValue = (depth: number): unknown => {
    if (depth > maximumJsonDepth) fail();
    skipWhitespace();
    const character = body[index];
    if (character === '"') return parseString();
    if (character === "{") {
      index += 1;
      skipWhitespace();
      const result: Record<string, unknown> = Object.create(null) as Record<string, unknown>;
      const keys = new Set<string>();
      if (body[index] === "}") {
        index += 1;
        return result;
      }
      while (true) {
        skipWhitespace();
        const key = parseString();
        if (keys.has(key)) fail();
        keys.add(key);
        skipWhitespace();
        if (body[index] !== ":") fail();
        index += 1;
        result[key] = parseValue(depth + 1);
        skipWhitespace();
        if (body[index] === "}") {
          index += 1;
          return result;
        }
        if (body[index] !== ",") fail();
        index += 1;
      }
    }
    if (character === "[") {
      index += 1;
      skipWhitespace();
      const result: unknown[] = [];
      if (body[index] === "]") {
        index += 1;
        return result;
      }
      while (true) {
        result.push(parseValue(depth + 1));
        skipWhitespace();
        if (body[index] === "]") {
          index += 1;
          return result;
        }
        if (body[index] !== ",") fail();
        index += 1;
      }
    }
    for (const [literal, value] of [
      ["true", true],
      ["false", false],
      ["null", null],
    ] as const) {
      if (body.startsWith(literal, index)) {
        index += literal.length;
        return value;
      }
    }
    const numberMatch = /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?/.exec(
      body.slice(index),
    );
    if (!numberMatch) return fail();
    const number = numberMatch[0];
    index += number.length;
    return Number(number);
  };

  const value = parseValue(0);
  skipWhitespace();
  if (index !== body.length) fail();
  return value;
}

function unavailableError(status: 502 | 503 = 502): GatewayError {
  return new GatewayError(status, "Сервис моделей недоступен");
}

function translateOpenAiSse(
  source: ReadableStream<Uint8Array>,
  signal: AbortSignal,
): ReadableStream<Uint8Array> {
  const reader = source.getReader();
  const decoder = new TextDecoder();
  const encoder = new TextEncoder();
  let cancelled = false;
  let removeAbortListener = () => {};

  return new ReadableStream<Uint8Array>({
    start(controller) {
      let lineBuffer = "";
      let eventBytes = 0;
      let dataLines: string[] = [];
      let doneSeen = false;

      const invalidStream = () =>
        new GatewayError(502, "Aegida Gate вернул некорректный поток");

      const dispatchEvent = () => {
        if (dataLines.length === 0) return;
        const payload = dataLines.join("\n");
        dataLines = [];
        if (payload === "[DONE]") {
          doneSeen = true;
          return;
        }
        if (doneSeen) throw invalidStream();

        let value: unknown;
        try {
          value = JSON.parse(payload) as unknown;
        } catch {
          throw invalidStream();
        }
        if (!isRecord(value)) throw invalidStream();
        if ("error" in value) throw invalidStream();
        const choices = value.choices;
        if (!Array.isArray(choices)) throw invalidStream();
        if (choices.length === 0) {
          if (!("usage" in value)) throw invalidStream();
          return;
        }

        let content = "";
        for (const choice of choices) {
          if (!isRecord(choice) || !isRecord(choice.delta)) throw invalidStream();
          const delta = choice.delta.content;
          if (delta === undefined || delta === null) continue;
          if (typeof delta !== "string") throw invalidStream();
          content += delta;
        }
        if (content) controller.enqueue(encoder.encode(content));
      };

      const processLine = (line: string, consumedDelimiter: string) => {
        eventBytes += encoder.encode(line + consumedDelimiter).byteLength;
        if (eventBytes > maximumEventBytes) throw invalidStream();
        if (line === "") {
          dispatchEvent();
          eventBytes = 0;
          return;
        }
        if (line.startsWith(":")) return;
        const colon = line.indexOf(":");
        const field = colon < 0 ? line : line.slice(0, colon);
        let value = colon < 0 ? "" : line.slice(colon + 1);
        if (value.startsWith(" ")) value = value.slice(1);
        if (field === "data") dataLines.push(value);
      };

      const processText = (text: string, final = false) => {
        lineBuffer += text;
        while (true) {
          const line = takeLine(lineBuffer, final);
          if (!line) break;
          lineBuffer = lineBuffer.slice(line.consumed.length);
          processLine(line.value, line.delimiter);
        }
        if (eventBytes + encoder.encode(lineBuffer).byteLength > maximumEventBytes) {
          throw invalidStream();
        }
        if (final) {
          if (lineBuffer) {
            processLine(lineBuffer, "");
            lineBuffer = "";
          }
          dispatchEvent();
        }
      };

      const onAbort = () => {
        void reader.cancel(abortError());
      };
      signal.addEventListener("abort", onAbort, { once: true });
      removeAbortListener = () => signal.removeEventListener("abort", onAbort);

      void (async () => {
        try {
          while (!doneSeen) {
            if (signal.aborted) throw abortError();
            const { done, value } = await reader.read();
            if (cancelled) return;
            if (signal.aborted) throw abortError();
            if (done) {
              processText(decoder.decode(), true);
              break;
            }
            processText(decoder.decode(value, { stream: true }));
          }
          if (!doneSeen) throw invalidStream();
          await reader.cancel();
          if (!cancelled) controller.close();
        } catch (error) {
          if (cancelled) return;
          await reader.cancel().catch(() => undefined);
          controller.error(
            signal.aborted || isAbortError(error)
              ? abortError()
              : error instanceof GatewayError
                ? error
                : new GatewayError(502, "Aegida Gate вернул некорректный поток"),
          );
        } finally {
          removeAbortListener();
        }
      })();
    },
    async cancel(reason) {
      cancelled = true;
      removeAbortListener();
      await reader.cancel(reason);
    },
  });
}

function takeLine(
  buffer: string,
  final: boolean,
): { value: string; delimiter: string; consumed: string } | null {
  for (let index = 0; index < buffer.length; index += 1) {
    const character = buffer[index];
    if (character === "\n") {
      return {
        value: buffer.slice(0, index),
        delimiter: "\n",
        consumed: buffer.slice(0, index + 1),
      };
    }
    if (character === "\r") {
      if (index === buffer.length - 1 && !final) return null;
      const delimiter = buffer[index + 1] === "\n" ? "\r\n" : "\r";
      return {
        value: buffer.slice(0, index),
        delimiter,
        consumed: buffer.slice(0, index + delimiter.length),
      };
    }
  }
  return null;
}

function gateOrigin(): string {
  const raw = process.env.AEGIDA_GATE_URL;
  if (!raw) throw new GatewayError(503, "Aegida Gate не настроен");
  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 GatewayError(503, "Aegida Gate не настроен");
  }
}

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);
}