WRNexusJS/packages/ai/test/ai.test.ts

import { afterEach, expect, test } from "bun:test";
import {
  AIError,
  createAI,
  createAIClient,
  deterministicAIProvider,
  type AIProvider,
} from "../src/index.ts";

const originalFetch = globalThis.fetch;
afterEach(() => {
  globalThis.fetch = originalFetch;
});

test("requires an API key before making a request", async () => {
  const previous = process.env.ANTHROPIC_API_KEY;
  delete process.env.ANTHROPIC_API_KEY;
  await expect(createAI().generate("hello")).rejects.toBeInstanceOf(AIError);
  if (previous !== undefined) process.env.ANTHROPIC_API_KEY = previous;
});

test("builds a bounded messages request and joins text blocks", async () => {
  let request: RequestInit | undefined;
  globalThis.fetch = (async (_url, init) => {
    request = init;
    return Response.json({ content: [{ type: "text", text: "Hello" }, { type: "tool_use" }] });
  }) as typeof fetch;
  const text = await createAI({ apiKey: "secret", model: "test-model", maxTokens: 123 }).generate(
    "Hi",
  );
  expect(text).toBe("Hello");
  expect(JSON.parse(String(request?.body))).toMatchObject({
    model: "test-model",
    max_tokens: 123,
    messages: [{ role: "user", content: "Hi" }],
    stream: false,
  });
});

test("streams fragmented SSE and keeps a final event without a newline", async () => {
  const encoder = new TextEncoder();
  globalThis.fetch = (async () =>
    new Response(
      new ReadableStream({
        start(controller) {
          controller.enqueue(
            encoder.encode('data: {"type":"content_block_delta","delta":{"type":"text_'),
          );
          controller.enqueue(encoder.encode('delta","text":"A"}}\n'));
          controller.enqueue(
            encoder.encode(
              'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"B"}}',
            ),
          );
          controller.close();
        },
      }),
    )) as unknown as typeof fetch;
  const chunks: string[] = [];
  for await (const chunk of createAI({ apiKey: "secret" }).stream("Hi")) chunks.push(chunk);
  expect(chunks).toEqual(["A", "B"]);
});

test("generates and validates structured output with an offline provider", async () => {
  const client = createAIClient({
    providers: [deterministicAIProvider({ responses: ['```json\n{"ok":true}\n```'] })],
  });
  const result = await client.generateObject<{ ok: boolean }>("ignored", {
    validate: (value): value is { ok: boolean } =>
      typeof value === "object" && value !== null && "ok" in value,
  });
  expect(result.value).toEqual({ ok: true });
  expect(client.capabilities().deterministic?.structuredOutput).toBe(true);
});

test("retries transient failures and reports redacted metadata", async () => {
  let calls = 0;
  const events: unknown[] = [];
  const provider: AIProvider = {
    name: "retrying",
    async generate() {
      calls++;
      if (calls === 1) throw new AIError("temporary secret-token", 503, "overloaded");
      return { value: "done", provider: "retrying", usage: { totalTokens: 3, costUsd: 0.01 } };
    },
  };
  const result = await createAIClient({
    providers: [provider],
    retry: { attempts: 2, baseDelayMs: 0 },
    onAttempt: (event) => {
      events.push(event);
    },
  }).generate("private prompt");
  expect(result.value).toBe("done");
  expect(calls).toBe(2);
  expect(JSON.stringify(events)).not.toContain("private prompt");
  expect(JSON.stringify(events)).not.toContain("raw");
});

test("opens a provider circuit and falls back", async () => {
  let failedCalls = 0;
  const failing: AIProvider = {
    name: "failing",
    async generate() {
      failedCalls++;
      throw new AIError("down", 503);
    },
  };
  const client = createAIClient({
    providers: [failing, deterministicAIProvider({ responses: ["one", "two"] })],
    retry: { attempts: 1 },
    circuitBreaker: { failureThreshold: 1, resetAfterMs: 60_000 },
  });
  expect((await client.generate("first")).value).toBe("one");
  expect((await client.generate("second")).value).toBe("two");
  expect(failedCalls).toBe(1);
});

test("executes validated tool calls and forwards cancellation", async () => {
  const client = createAIClient({ providers: [deterministicAIProvider({ responses: ["ok"] })] });
  const result = {
    value: "",
    provider: "deterministic",
    toolCalls: [{ id: "1", name: "double", arguments: { value: 4 } }],
  };
  const output = await client.executeTools(result, [
    {
      name: "double",
      validate: (value): value is { value: number } =>
        typeof value === "object" && value !== null && "value" in value,
      execute: ({ value }) => value * 2,
    },
  ]);
  expect(output[0]?.value).toBe(8);
});