coreci-chat/packages/byom/tests/router.test.ts

/**
 * router.test.ts — BYOM routing shim (REQ-008, REQ-009, Edge 1, G-001).
 *
 * Cases:
 *   - routeInference success → POSTs to the configured endpoint with the bearer key, returns the parsed body
 *   - no validated endpoint → ByomUnconfiguredError
 *   - endpoint unreachable (fetch throws) → ByomUnreachableError
 *   - endpoint returns 401/403 → ByomUnreachableError
 *   - endpoint returns 5xx → ByomUnreachableError
 *
 * Uses PGlite + migration + LocalEncryptedProvider for the endpoint registry,
 * and mocks `fetch` for the outbound inference call.
 */

import { describe, it, expect, beforeAll, beforeEach, afterEach, vi } from "vitest";
import { readFile } from "node:fs/promises";
import { join } from "node:path";
import { createDb } from "@coreci/db";
import { setDbClient } from "@coreci/db";
import { LocalEncryptedProvider, type SecretProvider } from "@coreci/secrets";
import { saveEndpoint, BYOM_SECRET_NAME } from "../src/repository.js";
import { routeInference, ByomUnconfiguredError, ByomUnreachableError } from "../src/router.js";
import type { ChatCompletionRequest, ChatCompletionResponse } from "../src/types.js";

const T1 = "00000000-0000-0000-0000-000000000001";
const MASTER_KEY = "test-master-key-for-byom-router-tests-32+";

async function runMigration(db: { exec: (t: string) => Promise<void> }): Promise<void> {
  const sql = await readFile(
    join(import.meta.dirname, "..", "..", "db", "migrations", "0001_init.sql"),
    "utf8",
  );
  await db.exec(sql);
}

function jsonResponse(body: unknown, status = 200): Response {
  return new Response(JSON.stringify(body), {
    status,
    headers: { "Content-Type": "application/json" },
  });
}

/** A fetch mock that returns a fresh OK response on every call (Response bodies are single-use). */
function okFetchMock(body: unknown = OK_RESPONSE): ReturnType<typeof vi.fn> {
  return vi.fn().mockImplementation(() => Promise.resolve(jsonResponse(body)));
}

const OK_RESPONSE: ChatCompletionResponse = {
  id: "chatcmpl-routed",
  choices: [{ index: 0, message: { role: "assistant", content: "hello" }, finish_reason: "stop" }],
  usage: { prompt_tokens: 5, completion_tokens: 1, total_tokens: 6 },
};

const PAYLOAD: ChatCompletionRequest = {
  model: "gpt-test",
  messages: [{ role: "user", content: "hi" }],
};

describe("BYOM routing shim (REQ-008, REQ-009, G-001)", () => {
  let db: Awaited<ReturnType<typeof createDb>>;
  let secrets: SecretProvider;

  beforeAll(async () => {
    db = await createDb({ mode: "pglite" });
    setDbClient(db);
    await runMigration(db);
  });

  beforeEach(async () => {
    const secretsDir = join(import.meta.dirname, ".secrets-test", `router-${Math.random().toString(36).slice(2)}`);
    secrets = new LocalEncryptedProvider({ baseDir: secretsDir, masterKey: MASTER_KEY });
    await db.query(`INSERT INTO tenants (id, name) VALUES ($1, 'T1') ON CONFLICT DO NOTHING`, [T1]);

    // Clean leftover byom_endpoints + audit_log rows for T1 (reset the chain).
    await db.query("ALTER TABLE byom_endpoints DISABLE ROW LEVEL SECURITY");
    await db.query(`DELETE FROM byom_endpoints WHERE tenant_id = $1`, [T1]);
    await db.query("ALTER TABLE byom_endpoints ENABLE ROW LEVEL SECURITY");
    await db.query("ALTER TABLE audit_log DISABLE ROW LEVEL SECURITY");
    await db.query(`DELETE FROM audit_log WHERE tenant_id = $1`, [T1]);
    await db.query("ALTER TABLE audit_log ENABLE ROW LEVEL SECURITY");
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it("routes inference to the configured endpoint and returns the parsed response", async () => {
    // Save a validated endpoint (mock fetch for the validate-on-save call).
    vi.stubGlobal("fetch", okFetchMock());
    await saveEndpoint(db, secrets, T1, "https://byom.example.com", "sk-route");

    // Now route a real inference call. fetch is still mocked to the OK body.
    const out = await routeInference(db, secrets, T1, PAYLOAD);
    expect(out.id).toBe("chatcmpl-routed");
    expect(out.choices[0]?.message.content).toBe("hello");

    // Verify the outbound call used the bearer key + the /v1/chat/completions URL.
    const calls = (fetch as unknown as ReturnType<typeof vi.fn>).mock.calls;
    // find the call to /v1/chat/completions with our payload (the validate call also posts there)
    const routed = calls.find(([u, init]) => {
      const body = (init as RequestInit).body as string;
      return u === "https://byom.example.com/v1/chat/completions" && body.includes('"gpt-test"');
    }) as [string, RequestInit] | undefined;
    expect(routed).toBeDefined();
    expect((routed![1].headers as Record<string, string>).Authorization).toBe("Bearer sk-route");
  });

  it("throws ByomUnconfiguredError when no validated endpoint exists (REQ-009)", async () => {
    vi.stubGlobal("fetch", vi.fn());
    await expect(routeInference(db, secrets, T1, PAYLOAD)).rejects.toBeInstanceOf(ByomUnconfiguredError);
    // No inference attempted (fetch not called).
    expect(fetch).not.toHaveBeenCalled();
  });

  it("throws ByomUnreachableError when the endpoint is unreachable (Edge 1)", async () => {
    vi.stubGlobal("fetch", okFetchMock());
    await saveEndpoint(db, secrets, T1, "https://byom.example.com", "sk-unreachable");

    // Now make the routed call throw (simulate network failure).
    vi.stubGlobal("fetch", vi.fn().mockRejectedValue(new Error("ECONNREFUSED")));
    await expect(routeInference(db, secrets, T1, PAYLOAD)).rejects.toBeInstanceOf(ByomUnreachableError);
  });

  it("throws ByomUnreachableError on 401/403 (treated as unreachable for the operator)", async () => {
    vi.stubGlobal("fetch", okFetchMock());
    await saveEndpoint(db, secrets, T1, "https://byom.example.com", "sk-revoked");

    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response("unauthorized", { status: 401 })));
    await expect(routeInference(db, secrets, T1, PAYLOAD)).rejects.toBeInstanceOf(ByomUnreachableError);
  });

  it("throws ByomUnreachableError on 5xx", async () => {
    vi.stubGlobal("fetch", okFetchMock());
    await saveEndpoint(db, secrets, T1, "https://byom.example.com", "sk-500");

    vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response("bad gateway", { status: 502 })));
    await expect(routeInference(db, secrets, T1, PAYLOAD)).rejects.toBeInstanceOf(ByomUnreachableError);
  });
});