WorkRoot-WRNexusJSComponent/test/studio-ai.test.ts

import { expect, test } from "bun:test";
import { CORE_DESIGN_SYSTEMS } from "@wrnexus/design-system";
import { designRevision, DEFAULT_DESIGN_BRIEF } from "@wrnexus/design-studio";
import { apimartProvider, studioAI } from "../app/lib/studio-ai.ts";
import { addStudioGuide } from "../app/lib/studio-guide.ts";

const proposals = ["2px", "4px", "8px"].map((value, index) => ({
  name: `Direction ${index + 1}`,
  rationale: "A different control shape.",
  patch: [{ op: "set", path: "/tokens/shape.control/value", value }],
}));
const completion = () =>
  Response.json({
    choices: [
      {
        finish_reason: "stop",
        message: { content: JSON.stringify({ alternatives: proposals }) },
      },
    ],
  });
const mockFetch = (
  fn: (input: string | URL | Request, init?: RequestInit) => Promise<Response>,
) => fn as typeof fetch;
const system = CORE_DESIGN_SYSTEMS.square;
const input = async (locks: string[] = []) => ({
  system,
  locks,
  brief: {
    ...DEFAULT_DESIGN_BRIEF,
    product: "A website",
    audience: "Customers",
  },
  kind: "directions",
  baseRevision: await designRevision(system, locks),
});
const request = async (body?: unknown, origin = "https://components.example") =>
  new Request("https://components.example/design-studio/api/generate", {
    method: "POST",
    headers: { "content-type": "application/json", origin },
    body: JSON.stringify(body ?? (await input())),
  });

test("APIMart uses the specified model, bearer authentication, non-streaming JSON and cancellation", async () => {
  const controller = new AbortController();
  const provider = apimartProvider(
    "test-secret",
    mockFetch(async (url, init) => {
      expect(String(url)).toBe("https://api.apimart.ai/v1/chat/completions");
      expect(new Headers(init?.headers).get("authorization")).toBe(
        "Bearer test-secret",
      );
      expect(init?.signal).toBe(controller.signal);
      const body = JSON.parse(init?.body as string);
      expect(body.model).toBe("gpt-6-luna");
      expect(body.stream).toBe(false);
      expect(body.max_tokens).toBe(4096);
      expect(body.messages).toEqual([
        { role: "system", content: "rules" },
        { role: "user", content: "brief" },
      ]);
      return completion();
    }),
  );
  const result = await provider.generate("brief", {
    system: "rules",
    signal: controller.signal,
  });
  expect(JSON.parse(result.value).alternatives).toHaveLength(3);
});

test("AI proposals use engine validation, preserve draft source, locks and revision checks", async () => {
  let calls = 0;
  const handler = studioAI(
    "test-secret",
    mockFetch(async () => {
      calls++;
      return completion();
    }),
  );
  const original = JSON.stringify(system);
  const response = await handler(await request());
  const batch = await response!.json();
  expect(batch.alternatives).toHaveLength(3);
  expect(batch.provenance.provider).toBe("APIMart");
  expect(batch.provenance.model).toBe("gpt-6-luna");
  expect(JSON.stringify(system)).toBe(original);
  const locked = await handler(
    await request(await input(["/tokens/shape.control"])),
  );
  expect((await locked!.json()).alternatives).toHaveLength(0);
  const before = calls;
  expect(
    (await handler(
      await request({ ...(await input()), baseRevision: "a".repeat(64) }),
    ))!.status,
  ).toBe(409);
  expect(calls).toBe(before);
});

test("metadata and errors never contain the key or raw upstream response", async () => {
  const handler = studioAI(
    "test-secret",
    mockFetch(
      async () =>
        new Response("test-secret upstream credentials", { status: 401 }),
    ),
  );
  const meta = await handler(
    new Request("https://components.example/design-studio/api/generation"),
  );
  const text = await meta!.text();
  expect(text).toContain("gpt-6-luna");
  expect(text).not.toContain("test-secret");
  const failure = await handler(await request());
  const body = await failure!.text();
  expect(body).not.toContain("test-secret");
  expect(body).not.toContain("upstream credentials");
  expect(body).toContain("WRN-STUDIO-PROVIDER");
});

test("same-origin checks, body bounds and local fallback remain enforced", async () => {
  let calls = 0;
  const handler = studioAI(
    "test-secret",
    mockFetch(async () => {
      calls++;
      return completion();
    }),
  );
  expect(
    (await handler(await request(await input(), "https://other.example")))!
      .status,
  ).toBe(403);
  expect(
    (await handler(await request({ huge: "x".repeat(2_000_001) })))!.status,
  ).toBe(413);
  expect(calls).toBe(0);
  expect(await studioAI()(await request())).toBeUndefined();
  expect(
    addStudioGuide(
      '<head><script type="module"></script></head><div id="diagnostics"></div>',
      false,
    ),
  ).toContain("APIMART_API_KEY");
  expect(addStudioGuide('<div id="diagnostics"></div>', true)).toContain(
    "APIMart AI is enabled",
  );
});

test("a public AI endpoint enforces a bounded generation rate", async () => {
  const handler = studioAI(
    "test-secret",
    mockFetch(async () => completion()),
  );
  for (let index = 0; index < 10; index++)
    expect((await handler(await request()))!.status).toBe(200);
  expect((await handler(await request()))!.status).toBe(429);
});