aegida-console / lib / chat / gateway.test.ts
gateway.test.ts
Raw
// @vitest-environment node

import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

vi.mock("server-only", () => ({}));

import { GatewayError, streamGatewayResponse } from "@/lib/chat/gateway";
import type { ChatRequest } from "@/lib/chat/types";

const request: ChatRequest = {
  model: "auto",
  messages: [{ role: "user", content: "hello" }],
};

beforeEach(() => {
  vi.stubEnv("AEGIDA_GATE_URL", "https://gate.example.test");
  vi.stubEnv("AI_GATEWAY_API_KEY", "shared-key-must-not-be-used");
});

afterEach(() => {
  vi.unstubAllEnvs();
  vi.unstubAllGlobals();
});

describe("streamGatewayResponse", () => {
  it("relays the exact user token and translates arbitrarily split OpenAI SSE to plain UTF-8", async () => {
    const sse = [
      ": keepalive\r\n\r\n",
      "data: {\"choices\":[{\"delta\":\r\n",
      "data: {\"content\":\"При\"}}]}\r\n\r\n",
      "data: {\"choices\":[{\"delta\":{\"content\":\"вет\"}}]}\n\n",
      "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n",
      "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
      "data: [DONE]\n\n",
    ].join("");
    const fetchMock = vi.fn().mockResolvedValue(
      new Response(splitEveryFewBytes(sse), {
        headers: { "Content-Type": "text/event-stream" },
      }),
    );
    vi.stubGlobal("fetch", fetchMock);
    const signal = new AbortController().signal;

    const stream = await streamGatewayResponse(request, "Bearer exact-user-jwt", signal);

    expect(await new Response(stream).text()).toBe("Привет");
    expect(fetchMock).toHaveBeenCalledWith(
      "https://gate.example.test/v1/chat/completions",
      {
        method: "POST",
        headers: {
          Authorization: "Bearer exact-user-jwt",
          "Content-Type": "application/json",
        },
        body: JSON.stringify({
          model: "auto",
          stream: true,
          stream_options: { include_usage: true },
          messages: [{ role: "user", content: "hello" }],
        }),
        cache: "no-store",
        redirect: "error",
        signal,
      },
    );
    expect(JSON.stringify(fetchMock.mock.calls[0])).not.toContain("shared-key-must-not-be-used");
  });

  it.each([
    ["malformed JSON", "data: {private-malformed-sentinel}\n\ndata: [DONE]\n\n"],
    [
      "an OpenAI error envelope",
      "data: {\"error\":{\"message\":\"private-error-sentinel\"}}\n\n",
    ],
    [
      "an invalid delta",
      "data: {\"choices\":[{\"delta\":{\"content\":42}}]}\n\ndata: [DONE]\n\n",
    ],
    [
      "EOF before DONE",
      "data: {\"choices\":[{\"delta\":{\"content\":\"partial-private-prompt\"}}]}\n\n",
    ],
  ])("fails safely for %s", async (_case, body) => {
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body)));

    const stream = await streamGatewayResponse(
      request,
      "Bearer exact-user-jwt",
      new AbortController().signal,
    );
    let error: unknown;
    try {
      await new Response(stream).text();
    } catch (caught) {
      error = caught;
    }
    expect(error).toBeInstanceOf(GatewayError);
    expect(String(error)).not.toContain("private-");
    expect(String(error)).not.toContain("hello");
  });

  it("rejects an SSE event larger than one MiB", async () => {
    const oversized = `data: ${"x".repeat(1_048_577)}\n\n`;
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(oversized)));

    const stream = await streamGatewayResponse(
      request,
      "Bearer exact-user-jwt",
      new AbortController().signal,
    );

    await expect(new Response(stream).text()).rejects.toMatchObject({
      status: 502,
      message: "Aegida Gate вернул некорректный поток",
    });
  });

  it.each([
    [429, "quota_exceeded", 429, "Лимит токенов исчерпан"],
    [403, "subject_inactive", 403, "Нет доступа к модели"],
    [403, "model_not_allowed", 403, "Нет доступа к модели"],
    [404, "model_not_found", 404, "Модель недоступна"],
    [404, "route_unavailable", 404, "Модель недоступна"],
  ])(
    "maps Gate %i %s to a safe user-facing error",
    async (gateStatus, code, status, message) => {
      vi.stubGlobal(
        "fetch",
        vi.fn().mockResolvedValue(
          Response.json({ error: { code } }, { status: gateStatus }),
        ),
      );

      await expect(
        streamGatewayResponse(
          request,
          "Bearer exact-user-jwt",
          new AbortController().signal,
        ),
      ).rejects.toMatchObject({ status, message });
    },
  );

  it.each([
    [
      429,
      {
        message: "Quota exceeded",
        type: "rate_limit_error",
        param: null,
        code: "quota_exceeded",
      },
      429,
      "Лимит токенов исчерпан",
    ],
    [
      403,
      {
        message: "Model access denied",
        type: "permission_error",
        param: "model",
        code: "model_not_allowed",
      },
      403,
      "Нет доступа к модели",
    ],
    [
      404,
      {
        message: "Model not found",
        type: "invalid_request_error",
        param: "model",
        code: "model_not_found",
      },
      404,
      "Модель недоступна",
    ],
    [
      502,
      {
        message: "Inference provider is temporarily unavailable",
        type: "server_error",
        param: null,
        code: "provider_unavailable",
      },
      502,
      "Сервис моделей недоступен",
    ],
    [
      503,
      {
        message: "Inference service is temporarily unavailable",
        type: "server_error",
        param: null,
        code: "service_unavailable",
      },
      503,
      "Сервис моделей недоступен",
    ],
    [
      503,
      {
        message: "Model catalog is temporarily unavailable",
        type: "server_error",
        param: null,
        code: "service_unavailable",
      },
      503,
      "Сервис моделей недоступен",
    ],
  ])(
    "maps the exact Gate OpenAI %i envelope without surfacing its message",
    async (gateStatus, errorBody, status, message) => {
      vi.stubGlobal(
        "fetch",
        vi.fn().mockResolvedValue(
          Response.json({ error: errorBody }, { status: gateStatus }),
        ),
      );

      let error: unknown;
      try {
        await streamGatewayResponse(
          request,
          "Bearer exact-user-jwt",
          new AbortController().signal,
        );
      } catch (caught) {
        error = caught;
      }

      expect(error).toMatchObject({ status, message });
      expect(String(error)).not.toContain(errorBody.message);
    },
  );

  it.each([
    ["malformed JSON", "{private-malformed-sentinel"],
    ["an unknown code", '{"error":{"code":"private_unknown_code"}}'],
    ["a code/status mismatch", '{"error":{"code":"quota_exceeded"}}', 403],
    [
      "a duplicate code field",
      '{"error":{"code":"private_unknown_code","code":"quota_exceeded"}}',
    ],
    [
      "nested escaped-equivalent duplicate code fields",
      '{"error":{"co\\u0064e":"private_unknown_code","code":"quota_exceeded"}}',
    ],
    [
      "escaped-equivalent duplicate outer fields",
      '{"err\\u006fr":{"code":"private_unknown_code"},"error":{"code":"quota_exceeded"}}',
    ],
    [
      "an extra secret-bearing error field",
      '{"error":{"code":"quota_exceeded","credential":"private-provider-secret"}}',
    ],
    [
      "an extra secret-bearing outer field",
      '{"error":{"code":"quota_exceeded"},"route_id":"private-route-secret"}',
    ],
    [
      "a rich envelope with a secret-bearing message mismatch",
      '{"error":{"message":"private-provider-secret","type":"server_error","param":null,"code":"service_unavailable"}}',
      503,
    ],
    [
      "a rich envelope with an extra request identifier",
      '{"error":{"message":"Inference service is temporarily unavailable","type":"server_error","param":null,"code":"service_unavailable","request_id":"private-request-secret"}}',
      503,
    ],
    [
      "a rich envelope with a mismatched parameter",
      '{"error":{"message":"Model not found","type":"invalid_request_error","param":null,"code":"model_not_found"}}',
      404,
    ],
    [
      "an invalid escaped surrogate",
      '{"error":{"message":"Quota exceeded","type":"rate_limit_error","param":null,"co\\ud800de":"quota_exceeded"}}',
    ],
    [
      "an extra JSON value after an otherwise safe envelope",
      '{"error":{"code":"quota_exceeded"}} true',
    ],
  ])("normalizes %s without leaking the Gate body", async (_case, body, status = 429) => {
    vi.stubGlobal(
      "fetch",
      vi.fn().mockResolvedValue(new Response(body, { status })),
    );

    let error: unknown;
    try {
      await streamGatewayResponse(
        request,
        "Bearer exact-user-jwt",
        new AbortController().signal,
      );
    } catch (caught) {
      error = caught;
    }

    expect(error).toBeInstanceOf(GatewayError);
    expect(error).toMatchObject({ status: 502, message: "Сервис моделей недоступен" });
    expect(String(error)).not.toContain("private-");
    expect(String(error)).not.toContain("hello");
  });

  it("cancels and normalizes a Gate error body larger than 8 KiB", async () => {
    const cancel = vi.fn();
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(new Uint8Array(8_193));
      },
      cancel,
    });
    vi.stubGlobal(
      "fetch",
      vi.fn().mockResolvedValue(new Response(body, { status: 429 })),
    );

    await expect(
      streamGatewayResponse(
        request,
        "Bearer exact-user-jwt",
        new AbortController().signal,
      ),
    ).rejects.toMatchObject({
      status: 502,
      message: "Сервис моделей недоступен",
    });
    expect(cancel).toHaveBeenCalledOnce();
    expect(body.locked).toBe(false);
  });

  it("cancels and releases a Gate error reader exactly once on fatal UTF-8", async () => {
    const cancel = vi.fn();
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(Uint8Array.of(0xff));
      },
      cancel,
    });
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body, { status: 429 })));

    await expect(
      streamGatewayResponse(
        request,
        "Bearer exact-user-jwt",
        new AbortController().signal,
      ),
    ).rejects.toMatchObject({ status: 502, message: "Сервис моделей недоступен" });

    expect(cancel).toHaveBeenCalledOnce();
    expect(body.locked).toBe(false);
  });

  it("cancels and releases a Gate error reader exactly once on abort", async () => {
    const cancel = vi.fn();
    const body = new ReadableStream<Uint8Array>({ cancel });
    const response = new Response(body, { status: 429 });
    const controller = new AbortController();
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(response));

    const pending = streamGatewayResponse(
      request,
      "Bearer exact-user-jwt",
      controller.signal,
    );
    const rejection = expect(pending).rejects.toMatchObject({ name: "AbortError" });
    await vi.waitFor(() => expect(body.locked).toBe(true));
    controller.abort();
    await rejection;

    expect(cancel).toHaveBeenCalledOnce();
    expect(body.locked).toBe(false);
  });

  it("releases a malformed Gate error body after successful EOF", async () => {
    const cancel = vi.fn();
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(new TextEncoder().encode("{private-malformed-sentinel"));
        controller.close();
      },
      cancel,
    });
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body, { status: 429 })));

    await expect(
      streamGatewayResponse(
        request,
        "Bearer exact-user-jwt",
        new AbortController().signal,
      ),
    ).rejects.toMatchObject({ status: 502, message: "Сервис моделей недоступен" });

    expect(cancel).not.toHaveBeenCalled();
    expect(body.locked).toBe(false);
  });

  it("releases an exact Gate error body after successful EOF without cancelling it", async () => {
    const cancel = vi.fn();
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(
          new TextEncoder().encode('{"error":{"code":"quota_exceeded"}}'),
        );
        controller.close();
      },
      cancel,
    });
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body, { status: 429 })));

    await expect(
      streamGatewayResponse(
        request,
        "Bearer exact-user-jwt",
        new AbortController().signal,
      ),
    ).rejects.toMatchObject({ status: 429, message: "Лимит токенов исчерпан" });

    expect(cancel).not.toHaveBeenCalled();
    expect(body.locked).toBe(false);
  });

  it.each([
    [502, "provider_unavailable"],
    [503, "service_unavailable"],
  ])("preserves an exact minimal safe %i infrastructure status", async (status, code) => {
    vi.stubGlobal(
      "fetch",
      vi.fn().mockResolvedValue(
        Response.json({ error: { code } }, { status }),
      ),
    );

    await expect(
      streamGatewayResponse(
        request,
        "Bearer exact-user-jwt",
        new AbortController().signal,
      ),
    ).rejects.toMatchObject({
      status,
      message: "Сервис моделей недоступен",
    });
  });

  it("cancels the upstream reader when the browser stops reading", async () => {
    const cancel = vi.fn();
    const upstream = new ReadableStream<Uint8Array>({
      start() {},
      cancel,
    });
    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(upstream)));
    const stream = await streamGatewayResponse(
      request,
      "Bearer exact-user-jwt",
      new AbortController().signal,
    );

    await stream.cancel("browser disconnected");

    expect(cancel).toHaveBeenCalledWith("browser disconnected");
  });

  it("requires a canonical user Bearer token and a configured Gate origin", async () => {
    const fetchMock = vi.fn();
    vi.stubGlobal("fetch", fetchMock);

    await expect(
      streamGatewayResponse(request, "shared-key", new AbortController().signal),
    ).rejects.toMatchObject({ status: 503 });
    vi.stubEnv("AEGIDA_GATE_URL", "https://gate.example.test/private/path");
    await expect(
      streamGatewayResponse(request, "Bearer user", new AbortController().signal),
    ).rejects.toMatchObject({ status: 503 });
    expect(fetchMock).not.toHaveBeenCalled();
  });
});

function splitEveryFewBytes(text: string): ReadableStream<Uint8Array> {
  const bytes = new TextEncoder().encode(text);
  const sizes = [1, 2, 7, 3, 11, 5];
  return new ReadableStream({
    start(controller) {
      let offset = 0;
      let index = 0;
      while (offset < bytes.byteLength) {
        const end = Math.min(offset + sizes[index % sizes.length], bytes.byteLength);
        controller.enqueue(bytes.slice(offset, end));
        offset = end;
        index += 1;
      }
      controller.close();
    },
  });
}