curriculum-project-hub/hub/test/unit/provider-model-catalog.test.ts

import { describe, expect, it, vi, beforeEach } from "vitest";
import { ProviderModelCatalog, type ProviderModelEntry } from "../../src/connections/providerModelCatalog.js";
import { LocalSecretEnvelope } from "../../src/security/secretEnvelope.js";
const TEST_KEY = Buffer.alloc(32, "k");
const secrets = new LocalSecretEnvelope({
  activeKeyId: "test-active",
  keys: new Map([["test-active", TEST_KEY]]),
});

function makeEncryptedPayload(baseUrl: string, authToken: string, anthropicApiKey = "") {
  const payload = { schemaVersion: 1, baseUrl, authToken, anthropicApiKey };
  return secrets.encryptJson({ purpose: "provider-connection", organizationId: "org-1", connectionId: "conn-1", secretVersionId: "sv-1" }, payload);
}

function mockPrismaWithConnection(overrides: Partial<{
  status: string;
  retiredAt: Date | null;
  providerId: string;
}> = {}) {
  const envelope = makeEncryptedPayload("https://openrouter.ai/api", "sk-test-key", "sk-ant-test");
  return {
    organizationProviderConnection: {
      findFirst: vi.fn().mockResolvedValue({
        id: "conn-1",
        organizationId: "org-1",
        providerId: overrides.providerId ?? "openrouter",
        status: overrides.status ?? "ACTIVE",
        activeSecretVersion: {
          id: "sv-1",
          connectionId: "conn-1",
          envelopeVersion: 1,
          keyId: "test-active",
          envelope,
          retiredAt: overrides.retiredAt ?? null,
        },
      }),
    },
  };
}

function mockFetchResponse(models: Array<{ id: string; name: string; supported_parameters: string[] }>) {
  return vi.fn().mockResolvedValue({
    ok: true,
    status: 200,
    json: () => Promise.resolve({ data: models }),
  });
}

describe("ProviderModelCatalog", () => {
  beforeEach(() => {
    vi.useFakeTimers();
  });

  it("fetches tool-capable models from the provider API", async () => {
    const prisma = mockPrismaWithConnection();
    const fetchImpl = mockFetchResponse([
      { id: "anthropic/claude-sonnet-5", name: "Claude Sonnet 5", supported_parameters: ["tools", "temperature"] },
      { id: "openai/gpt-4o", name: "GPT-4o", supported_parameters: ["tools", "temperature"] },
      { id: "meta/llama-3", name: "Llama 3", supported_parameters: ["temperature"] },
    ]);
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    const models = await catalog.listModels("org-1");

    expect(models).toHaveLength(3);
    expect(models.map((m: ProviderModelEntry) => m.id)).toEqual([
      "anthropic/claude-sonnet-5",
      "openai/gpt-4o",
      "meta/llama-3",
    ]);
    expect(models[0]).toMatchObject({ label: "Claude Sonnet 5", toolCapable: true });
    expect(models[2]).toMatchObject({ label: "Llama 3", toolCapable: false });
    expect(fetchImpl).toHaveBeenCalledTimes(1);
    const calledUrl = String(fetchImpl.mock.calls[0][0]);
    expect(calledUrl).toContain("supported_parameters=tools");
  });

  it("returns cached results on subsequent calls within TTL", async () => {
    const prisma = mockPrismaWithConnection();
    const fetchImpl = mockFetchResponse([
      { id: "anthropic/claude-sonnet-5", name: "Claude Sonnet 5", supported_parameters: ["tools"] },
    ]);
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    await catalog.listModels("org-1");
    await catalog.listModels("org-1");

    expect(fetchImpl).toHaveBeenCalledTimes(1);
  });

  it("re-fetches after TTL expires", async () => {
    const prisma = mockPrismaWithConnection();
    const fetchImpl = mockFetchResponse([
      { id: "anthropic/claude-sonnet-5", name: "Claude Sonnet 5", supported_parameters: ["tools"] },
    ]);
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    await catalog.listModels("org-1");
    vi.advanceTimersByTime(61_000);
    await catalog.listModels("org-1");

    expect(fetchImpl).toHaveBeenCalledTimes(2);
  });

  it("returns empty list when no ACTIVE provider connection exists", async () => {
    const prisma = {
      organizationProviderConnection: {
        findFirst: vi.fn().mockResolvedValue(null),
      },
    };
    const fetchImpl = mockFetchResponse([]);
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    const models = await catalog.listModels("org-1");

    expect(models).toEqual([]);
    expect(fetchImpl).not.toHaveBeenCalled();
  });

  it("invalidate forces a re-fetch on next call", async () => {
    const prisma = mockPrismaWithConnection();
    const fetchImpl = mockFetchResponse([
      { id: "anthropic/claude-sonnet-5", name: "Claude Sonnet 5", supported_parameters: ["tools"] },
    ]);
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    await catalog.listModels("org-1");
    catalog.invalidate("org-1");
    await catalog.listModels("org-1");

    expect(fetchImpl).toHaveBeenCalledTimes(2);
  });

  it("sends auth token and anthropic API key in headers", async () => {
    const prisma = mockPrismaWithConnection();
    const fetchImpl = mockFetchResponse([
      { id: "anthropic/claude-sonnet-5", name: "Claude Sonnet 5", supported_parameters: ["tools"] },
    ]);
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    await catalog.listModels("org-1");

    const init = fetchImpl.mock.calls[0][1] as RequestInit;
    const headers = init.headers as Record<string, string>;
    expect(headers.authorization).toBe("Bearer sk-test-key");
    expect(headers["x-api-key"]).toBe("sk-ant-test");
  });

  it("throws on non-200 response", async () => {
    const prisma = mockPrismaWithConnection();
    const fetchImpl = vi.fn().mockResolvedValue({ ok: false, status: 401 });
    const catalog = new ProviderModelCatalog(prisma as never, secrets, 60_000, fetchImpl as never);

    await expect(catalog.listModels("org-1")).rejects.toThrow("provider models request failed: status 401");
  });
});