Add Workers AI model catalogue and multi-shape inference adapter

workers/api/scripts/smoke_api.mjs

@@ -316,6 +316,47 @@ let modelSlug = null;
     JSON.stringify(r.data).slice(0, 200),
   );
 }
+// Workers AI model ids are an allowlist, not free text. Before the
+// catalogue existed any string was forwarded to env.AI.run(), so a typo
+// published fine and failed only after the renderer had been debited.
+let waModelSlug = null;
+{
+  const r = await api("POST", "/v1/models", {
+    title: `Smoke WA Model ${RUN}`,
+    kind: "image",
+    provider: "workers_ai",
+    visibility: "private",
+  });
+  waModelSlug = r.data?.model?.slug ?? null;
+  record(
+    "POST /v1/models (workers_ai image)",
+    r.status === 201 && !!waModelSlug,
+    JSON.stringify(r.data).slice(0, 160),
+  );
+}
+{
+  const r = await api("POST", `/v1/models/${waModelSlug}/versions`, {
+    // One character short of the real id.
+    provider_model_id: "@cf/black-forest-labs/flux-1-schnel",
+    price_tokens: 10,
+  });
+  record(
+    "workers_ai version rejects an unlisted model id",
+    r.status === 400 && r.data?.error === "model_not_in_catalogue",
+    JSON.stringify(r.data).slice(0, 160),
+  );
+}
+{
+  const r = await api("POST", `/v1/models/${waModelSlug}/versions`, {
+    provider_model_id: "@cf/black-forest-labs/flux-1-schnell",
+    price_tokens: 10,
+  });
+  record(
+    "workers_ai version accepts a catalogued model id",
+    r.status === 201 && r.data?.version?.version === 1,
+    JSON.stringify(r.data).slice(0, 160),
+  );
+}
 {
   const r = await api("GET", `/v1/models/${modelSlug}`);
   record(

workers/api/src/ai/inference.ts

@@ -1,6 +1,7 @@
 import Anthropic from "@anthropic-ai/sdk";
 import type { Env } from "../types";
 import type { ModelKind, ModelProvider } from "../db/render";
+import { getCatalogueEntry, type CatalogueEntry } from "./workersAiCatalogue";
 
 /**
  * Model-execution layer for the render-token service.
@@ -28,11 +29,23 @@ export interface InferenceInput {
    * provider.
    */
   weightsRef?: string | null;
+  /**
+   * Source image for the img2img / inpaint / vision lanes. Raw bytes —
+   * the adapter converts to whatever encoding the target model wants.
+   */
+  sourceImage?: Uint8Array | null;
+  /** Inpainting mask. Required by the inpaint family, ignored elsewhere. */
+  maskImage?: Uint8Array | null;
 }
 
 export interface InferenceOutput {
-  outputKind: "code" | "image";
-  /** Present for `code` jobs — the generated source, inline. */
+  /**
+   * `text` covers the companion lane (vision reverse-prompting, prompt
+   * expansion). Deliberately NOT added to `ModelKind` — companion
+   * models are an internal assist, not something a creator publishes.
+   */
+  outputKind: "code" | "image" | "text";
+  /** Present for `code` and `text` jobs — the generated text, inline. */
   text: string | null;
   /** Present for `image` jobs — raw bytes for the caller to persist to R2. */
   bytes: Uint8Array | null;
@@ -72,6 +85,21 @@ export function isRenderMock(env: Env): boolean {
  * so scripts/smoke_api.mjs and local `wrangler dev` runs are hermetic.
  */
 async function runMock(input: InferenceInput): Promise<InferenceOutput> {
+  // Companion (vision / instruct) models return text, and are reached by
+  // provider model id rather than by ModelKind — so the mock has to look
+  // them up the same way the real path does, or mock runs would claim an
+  // image where production returns prose.
+  const mockEntry =
+    input.provider === "workers_ai"
+      ? getCatalogueEntry(input.providerModelId)
+      : undefined;
+  if (mockEntry?.outputKind === "text") {
+    const text =
+      `mock ${mockEntry.family} output — model=${mockEntry.id} ` +
+      `seed=${input.seed}: ${input.prompt.slice(0, 120)}`;
+    const outputHash = await sha256Hex(new TextEncoder().encode(text));
+    return { outputKind: "text", text, bytes: null, contentType: null, outputHash };
+  }
   if (input.kind === "code") {
     const text =
       `// mock render — provider=${input.provider} model=${input.providerModelId}\n` +
@@ -159,11 +187,180 @@ async function runAnthropic(
 }
 
 /**
- * Cloudflare Workers AI — image generation (reference art / textures).
- * Output feeds the studio's working material only; it is never eligible
- * for the frozen-artifact bundle (enforced in the freeze pipeline, not
- * here — this function just runs the model).
+ * Cloudflare Workers AI.
+ *
+ * Dispatches on the catalogue entry rather than assuming one shape:
+ * Workers AI models disagree on both request and response format, and
+ * the disagreement does not follow vendor or task (see the header of
+ * ./workersAiCatalogue.ts). An unlisted model is refused here rather
+ * than being forwarded to `env.AI.run()` on trust.
+ *
+ * Image output feeds the studio's working material only; it is never
+ * eligible for the frozen-artifact bundle (enforced in the freeze
+ * pipeline, not here — this function just runs the model).
+ */
+
+/** Decode a base64 string into bytes. */
+function base64ToBytes(b64: string): Uint8Array {
+  const binary = atob(b64);
+  return Uint8Array.from(binary, (ch) => ch.charCodeAt(0));
+}
+
+/**
+ * Caller params, narrowed to the keys this model accepts. Unknown keys
+ * are dropped: some models reject them and others ignore them silently,
+ * and a caller should not be able to discover which by experiment.
+ */
+export function filterParams(
+  entry: CatalogueEntry,
+  params: Record<string, unknown> | null,
+): Record<string, unknown> {
+  const out: Record<string, unknown> = { ...entry.defaults };
+  if (!params) return out;
+  for (const key of entry.paramAllowlist) {
+    if (params[key] !== undefined) out[key] = params[key];
+  }
+  return out;
+}
+
+/** Numeric seed where the seed is numeric; undefined otherwise. */
+function numericSeed(seed: string): number | undefined {
+  return /^\d+$/.test(seed) ? parseInt(seed, 10) : undefined;
+}
+
+/**
+ * Build the JSON body for a model. Families differ in more than
+ * parameters: the SD lanes want pixel arrays, vision wants an image
+ * alongside a prompt, and instruct wants a chat transcript.
+ */
+export function buildJsonInput(
+  entry: CatalogueEntry,
+  input: InferenceInput,
+): Record<string, unknown> {
+  const params = filterParams(entry, input.params);
+  const seed = numericSeed(input.seed);
+
+  if (entry.family === "instruct") {
+    const messages: { role: string; content: string }[] = [];
+    if (input.systemPrompt) {
+      messages.push({ role: "system", content: input.systemPrompt });
+    }
+    messages.push({ role: "user", content: input.prompt });
+    return { ...params, messages };
+  }
+
+  if (entry.family === "vision") {
+    if (!input.sourceImage) {
+      throw new InferenceError(
+        `${entry.label} requires a source image`,
+        "missing_source_image",
+      );
+    }
+    return { ...params, prompt: input.prompt, image: Array.from(input.sourceImage) };
+  }
+
+  const body: Record<string, unknown> = { ...params, prompt: input.prompt };
+  if (seed !== undefined) body.seed = seed;
+
+  if (entry.family === "img2img" || entry.family === "inpaint") {
+    if (!input.sourceImage) {
+      throw new InferenceError(
+        `${entry.label} requires a source image`,
+        "missing_source_image",
+      );
+    }
+    // Documented as an array of 8-bit unsigned integers, not raw bytes.
+    body.image = Array.from(input.sourceImage);
+  }
+  if (entry.family === "inpaint") {
+    if (!input.maskImage) {
+      throw new InferenceError(
+        `${entry.label} requires a mask`,
+        "missing_mask",
+      );
+    }
+    body.mask = Array.from(input.maskImage);
+  }
+  return body;
+}
+
+/**
+ * FLUX.2 models take multipart form-data even for a prompt-only call.
+ * FormData exposes neither its serialized body nor its boundary, so it
+ * is round-tripped through a Response to get both.
+ */
+function buildMultipartInput(
+  entry: CatalogueEntry,
+  input: InferenceInput,
+): { multipart: { body: unknown; contentType: string } } {
+  const params = filterParams(entry, input.params);
+  const seed = numericSeed(input.seed);
+  const form = new FormData();
+  form.append("prompt", input.prompt);
+  for (const [key, value] of Object.entries(params)) {
+    if (value !== undefined && value !== null) form.append(key, String(value));
+  }
+  if (seed !== undefined) form.append("seed", String(seed));
+  if (input.sourceImage) {
+    form.append(
+      "input_image_0",
+      new Blob([input.sourceImage], { type: "image/png" }),
+    );
+  }
+  const serialized = new Response(form);
+  const contentType = serialized.headers.get("content-type");
+  if (!contentType) {
+    throw new InferenceError("could not serialize multipart body", "provider_error");
+  }
+  return { multipart: { body: serialized.body, contentType } };
+}
+
+/**
+ * Decode whatever the model returned into bytes or text, per the
+ * catalogue's declared response shape.
  */
+export async function decodeOutput(
+  entry: CatalogueEntry,
+  result: unknown,
+): Promise<{ bytes: Uint8Array | null; text: string | null }> {
+  if (entry.responseShape === "binary") {
+    if (result instanceof ReadableStream) {
+      const buf = await new Response(result as ReadableStream).arrayBuffer();
+      return { bytes: new Uint8Array(buf), text: null };
+    }
+    if (result instanceof Uint8Array) return { bytes: result, text: null };
+    if (result instanceof ArrayBuffer) {
+      return { bytes: new Uint8Array(result), text: null };
+    }
+    throw new InferenceError(
+      `${entry.id} did not return binary output`,
+      "provider_error",
+    );
+  }
+
+  if (entry.responseShape === "json_base64_image") {
+    const image = (result as { image?: unknown } | null)?.image;
+    if (typeof image !== "string" || !image) {
+      throw new InferenceError(
+        `${entry.id} did not return a base64 image`,
+        "provider_error",
+      );
+    }
+    return { bytes: base64ToBytes(image), text: null };
+  }
+
+  // json_text
+  const field = entry.textField ?? "response";
+  const value = (result as Record<string, unknown> | null)?.[field];
+  if (typeof value !== "string" || !value.trim()) {
+    throw new InferenceError(
+      `${entry.id} did not return text in "${field}"`,
+      "provider_error",
+    );
+  }
+  return { bytes: null, text: value.trim() };
+}
+
 async function runWorkersAi(
   env: Env,
   input: InferenceInput,
@@ -174,37 +371,49 @@ async function runWorkersAi(
       "provider_unconfigured",
     );
   }
+  const entry = getCatalogueEntry(input.providerModelId);
+  if (!entry) {
+    throw new InferenceError(
+      `${input.providerModelId} is not an allowlisted Workers AI model`,
+      "model_not_in_catalogue",
+    );
+  }
+
+  const body =
+    entry.requestShape === "multipart"
+      ? buildMultipartInput(entry, input)
+      : buildJsonInput(entry, input);
+
   let result: unknown;
   try {
-    result = await env.AI.run(input.providerModelId, {
-      prompt: input.prompt,
-      seed: /^\d+$/.test(input.seed) ? parseInt(input.seed, 10) : undefined,
-    });
+    // The binding's per-model overloads can't be satisfied by a model id
+    // only known at runtime; the catalogue is what guarantees the body
+    // matches the model.
+    result = await (env.AI as unknown as {
+      run(model: string, body: unknown): Promise<unknown>;
+    }).run(entry.id, body);
   } catch (e) {
     throw new InferenceError(
       e instanceof Error ? e.message : "workers ai request failed",
       "provider_error",
     );
   }
-  // Workers AI image models return a ReadableStream (binary) response.
-  let bytes: Uint8Array;
-  if (result instanceof ReadableStream) {
-    const buf = await new Response(result).arrayBuffer();
-    bytes = new Uint8Array(buf);
-  } else if (result instanceof Uint8Array) {
-    bytes = result;
-  } else {
-    throw new InferenceError(
-      "unexpected Workers AI response shape",
-      "provider_error",
-    );
+
+  const { bytes, text } = await decodeOutput(entry, result);
+
+  if (entry.outputKind === "text") {
+    const outputHash = await sha256Hex(new TextEncoder().encode(text ?? ""));
+    return { outputKind: "text", text, bytes: null, contentType: null, outputHash };
+  }
+  if (!bytes) {
+    throw new InferenceError(`${entry.id} returned no image bytes`, "provider_error");
   }
   const outputHash = await sha256Hex(bytes);
   return {
     outputKind: "image",
     text: null,
     bytes,
-    contentType: "image/png",
+    contentType: entry.contentType,
     outputHash,
   };
 }