Forward TypeSafe image opt-in, multipart and zstd uploads for Jev

src/gateway/routes/jev.test.ts

@@ -2,25 +2,26 @@ import { describe, expect, test } from "bun:test";
 
 import { testDatabase } from "../../test/database";
 import { jevRoutes } from "./jev";
-import { createTestAccount, fakeFetch, onlyBillingRecord, post, testBilling } from "./test-harness";
+import { billingRecords, createTestAccount, fakeFetch, onlyBillingRecord, post, testBilling } from "./test-harness";
 
 const { sql } = await testDatabase();
 
 describe("jev route", () => {
   const { billing, settlements, settled } = testBilling(sql);
 
-  const call = async (label: string, upstream: () => Response, model?: string) => {
+  const call = async (label: string, upstream: () => Response, model?: string, headers: Record<string, string> = {}) => {
     const account = await createTestAccount(sql, `jev-${label}`);
-    const { fetch } = fakeFetch(upstream);
+    const { fetch, upstream: sent } = fakeFetch(upstream);
     const app = jevRoutes({ sql, billing, settlements, enforceIdv: false, typesafeApiKey: "ts-key", fetch });
     const response = await app.handle(
       post("/proxy/v1/jev/systemone", { input: "hi", ...(model ? { model } : {}) }, {
         authorization: `Bearer ${account.apiKey}`,
+        ...headers,
       }),
     );
     await response.text();
     await settled();
-    return { account, response, record: await onlyBillingRecord(sql, account.accountId) };
+    return { account, response, sent, record: await onlyBillingRecord(sql, account.accountId) };
   };
 
   test("a successful response records the versioned model it reports", async () => {
@@ -46,6 +47,116 @@ describe("jev route", () => {
     expect(record.event?.outcome).toBe("provider_error");
   });
 
+  test("forwards TypeSafe's danger zone opt-in headers and the gateway's own key", async () => {
+    const { sent } = await call("images", () => Response.json({ model: "jev-1.13.0", output: "hi" }), undefined, {
+      "TypeSafe-Danger-Zone-Images": "true",
+      "x-other": "dropped",
+    });
+    expect(sent[0]?.headers.get("typesafe-danger-zone-images")).toBe("true");
+    expect(sent[0]?.headers.get("x-other")).toBeNull();
+    expect(sent[0]?.headers.get("authorization")).toBe("Bearer ts-key");
+  });
+
+  test("forwards a multipart image upload with the model defaulted", async () => {
+    const account = await createTestAccount(sql, "jev-multipart");
+    let sent: Request | undefined;
+    const app = jevRoutes({
+      sql,
+      billing,
+      settlements,
+      enforceIdv: false,
+      typesafeApiKey: "ts-key",
+      fetch: async (input, init) => {
+        sent = new Request(String(input), init);
+        return Response.json({ model: "jev-1.13.0", output: "hi", usage: { input_tokens: 1_000_000, output_tokens: 0 } });
+      },
+    });
+    const image = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0xff, 0x00, 0xd8]);
+    const form = new FormData();
+    form.append("request", JSON.stringify({ state: { door: "d" }, assets: { d: { part: "door" } } }));
+    form.append("door", new File([image], "door.png", { type: "image/png" }));
+    const response = await app.handle(
+      new Request("http://gateway.test/proxy/v1/jev/v1/systemone", {
+        method: "POST",
+        headers: { authorization: `Bearer ${account.apiKey}`, "TypeSafe-Danger-Zone-Images": "true" },
+        body: form,
+      }),
+    );
+    await response.text();
+    await settled();
+
+    expect(response.status).toBe(200);
+    expect(sent?.headers.get("content-type")).toStartWith("multipart/form-data; boundary=");
+    expect(sent?.headers.get("typesafe-danger-zone-images")).toBe("true");
+    const upstream = await sent!.formData();
+    expect([...upstream.keys()]).toEqual(["request", "door"]);
+    expect(JSON.parse(String(upstream.get("request")))).toEqual({
+      model: "jev-latest",
+      state: { door: "d" },
+      assets: { d: { part: "door" } },
+    });
+    expect(new Uint8Array(await (upstream.get("door") as File).arrayBuffer())).toEqual(image);
+    expect((await onlyBillingRecord(sql, account.accountId)).actualCostUsd).toBe("0.042000000000");
+  });
+
+  test("decodes a zstd multipart upload as the TypeSafe SDK sends it", async () => {
+    const account = await createTestAccount(sql, "jev-zstd");
+    let sent: Request | undefined;
+    const app = jevRoutes({
+      sql,
+      billing,
+      settlements,
+      enforceIdv: false,
+      typesafeApiKey: "ts-key",
+      fetch: async (input, init) => {
+        sent = new Request(String(input), init);
+        return Response.json({ model: "jev-1.13.0", output: "hi", usage: { input_tokens: 1_000_000, output_tokens: 0 } });
+      },
+    });
+    const image = crypto.getRandomValues(new Uint8Array(4096));
+    const form = new FormData();
+    form.append("request", JSON.stringify({ state: { door: "d" }, assets: { d: { part: "door" } } }));
+    form.append("door", new File([image], "door.png", { type: "image/png" }));
+    const encoded = new Response(form);
+    const response = await app.handle(
+      new Request("http://gateway.test/proxy/v1/jev/v1/systemone", {
+        method: "POST",
+        headers: {
+          authorization: `Bearer ${account.apiKey}`,
+          "content-type": encoded.headers.get("content-type")!,
+          "content-encoding": "zstd",
+        },
+        body: new Uint8Array(Bun.zstdCompressSync(await encoded.bytes())),
+      }),
+    );
+    await response.text();
+    await settled();
+
+    expect(response.status).toBe(200);
+    expect(sent?.headers.get("content-encoding")).toBeNull();
+    const upstream = await sent!.formData();
+    expect(JSON.parse(String(upstream.get("request"))).model).toBe("jev-latest");
+    expect(new Uint8Array(await (upstream.get("door") as File).arrayBuffer())).toEqual(image);
+    expect((await onlyBillingRecord(sql, account.accountId)).actualCostUsd).toBe("0.042000000000");
+  });
+
+  test("rejects a zstd body that decodes past the body limit before dispatch", async () => {
+    const account = await createTestAccount(sql, "jev-zstd-bomb");
+    const { fetch, upstream } = fakeFetch(() => Response.json({ model: "jev-1.13.0", output: "hi" }));
+    const app = jevRoutes({ sql, billing, settlements, enforceIdv: false, typesafeApiKey: "ts-key", fetch, maxBodyBytes: 64 * 1024 });
+    const response = await app.handle(
+      new Request("http://gateway.test/proxy/v1/jev/systemone", {
+        method: "POST",
+        headers: { authorization: `Bearer ${account.apiKey}`, "content-type": "application/json", "content-encoding": "zstd" },
+        body: new Uint8Array(Bun.zstdCompressSync(JSON.stringify({ input: "a".repeat(1024 * 1024) }))),
+      }),
+    );
+
+    expect(response.status).toBe(413);
+    expect(upstream).toHaveLength(0);
+    expect(await billingRecords(sql, account.accountId)).toEqual([]);
+  });
+
   test("a success without usable usage is charged the hold", async () => {
     const { response, record } = await call("no-usage", () => Response.json({ model: "jev-1.13.0", output: "hi" }));
     expect(response.status).toBe(200);

src/gateway/routes/jev.ts

@@ -1,7 +1,10 @@
+import { promisify } from "node:util";
+import { zstdDecompress } from "node:zlib";
+
 import { Elysia } from "elysia";
 
 import { Usd } from "../../billing/money";
-import { executeJsonProvider } from "../../providers/json-provider";
+import { meterJsonResponse } from "../../providers/json-provider";
 import { forwardableHeaders } from "../../providers/response-headers";
 import {
   DEFAULT_JEV_MODEL,
@@ -33,6 +36,30 @@ export type JevRouteDependencies = MeteredRouteDependencies & {
    */
   inputPricePerMillionTokensUsd?: string;
   baseUrl?: string;
+  /** Largest request body after decompression; defaults to 20 MiB. */
+  maxBodyBytes?: number;
+};
+
+const unzstd = promisify(zstdDecompress);
+
+/**
+ * The TypeSafe SDK zstd-compresses bodies of 1 KiB or more, multipart
+ * included. Decoded up front so the JSON can be screened and image parts are
+ * forwarded intact; the cap stops a small body expanding without bound.
+ */
+const decodeBody = async (request: Request, maxBytes: number) => {
+  const encoding = request.headers.get("content-encoding")?.trim().toLowerCase();
+  if (!encoding || encoding === "identity") return request;
+  if (encoding !== "zstd") throw new HttpError(415, `Unsupported Content-Encoding ${encoding}`);
+  const body = await unzstd(await request.bytes(), { maxOutputLength: maxBytes }).catch((error: { code?: string }) => {
+    throw error.code === "ERR_BUFFER_TOO_LARGE"
+      ? new HttpError(413, `Request body exceeds the ${Math.floor(maxBytes / 1024 / 1024)} MiB limit`)
+      : new HttpError(400, "Request body is not valid zstd");
+  });
+  const headers = new Headers(request.headers);
+  headers.delete("content-encoding");
+  headers.delete("content-length");
+  return new Request(request.url, { method: request.method, headers, body: new Uint8Array(body), signal: request.signal });
 };
 
 /**
@@ -46,19 +73,34 @@ export const jevRoutes = (deps: JevRouteDependencies) => {
   const reservation = Usd.parse(deps.reservationUsd ?? "0.02");
   const inputPrice = Usd.parse(deps.inputPricePerMillionTokensUsd ?? "0.042");
   const baseUrl = (deps.baseUrl ?? "https://api.typesafe.ai").replace(/\/$/, "");
+  const maxBodyBytes = deps.maxBodyBytes ?? 20 * 1024 * 1024;
 
   const authorize = async (request: Request, rawBody: string) => {
     const principal = await authorizeProviderRequest(deps, rateLimiter, request, rawBody);
     return principal;
   };
 
-  const upstreamHeaders = () => ({
-    "content-type": "application/json",
-    authorization: `Bearer ${deps.typesafeApiKey}`,
-  });
+  const upstreamHeaders = (request: Request) => {
+    const headers: Record<string, string> = { authorization: `Bearer ${deps.typesafeApiKey}` };
+    for (const [name, value] of request.headers) {
+      if (name.startsWith("typesafe-danger-zone-")) headers[name] = value;
+    }
+    return headers;
+  };
 
-  const systemOne = async (request: Request) => {
-    const rawBody = await request.text();
+  const systemOne = async (received: Request) => {
+    const request = await decodeBody(received, maxBodyBytes);
+    const form = request.headers.get("content-type")?.startsWith("multipart/form-data")
+      ? await request.formData().catch(() => {
+          throw new HttpError(400, "Request body must be valid multipart/form-data");
+        })
+      : null;
+    const requestPart = form?.get("request");
+    const rawBody = !form
+      ? await request.text()
+      : typeof requestPart === "string"
+        ? requestPart
+        : ((await requestPart?.text()) ?? "");
     const principal = await authorize(request, rawBody);
 
     const body = parseJsonObject(rawBody);
@@ -68,6 +110,7 @@ export const jevRoutes = (deps: JevRouteDependencies) => {
     }
     body.model = model;
     const requestBody = JSON.stringify(body);
+    form?.set("request", requestBody);
 
     const input: ProviderRouteInput = {
       provider: TYPESAFE,
@@ -76,10 +119,15 @@ export const jevRoutes = (deps: JevRouteDependencies) => {
       estimatedCostUsd: reservation,
       uncertainChargeUsd: reservation,
       execute: async () => {
-        const metered = await executeJsonProvider({
-          fetch: deps.fetch,
-          url: `${baseUrl}/v1/systemone`,
-          init: { method: "POST", headers: upstreamHeaders(), body: requestBody },
+        const upstream = await (deps.fetch ?? fetch)(`${baseUrl}/v1/systemone`, {
+          method: "POST",
+          headers: form
+            ? upstreamHeaders(request)
+            : { ...upstreamHeaders(request), "content-type": "application/json" },
+          body: form ?? requestBody,
+        });
+        const metered = await meterJsonResponse(upstream, {
+          init: { body: requestBody },
           extractCost: (response) => jevCost(response, inputPrice),
           extractTokens: jevTokens,
         });
@@ -103,7 +151,7 @@ export const jevRoutes = (deps: JevRouteDependencies) => {
     await authorize(request, "");
     const upstream = await (deps.fetch ?? fetch)(`${baseUrl}/v1/models`, {
       method: "GET",
-      headers: upstreamHeaders(),
+      headers: upstreamHeaders(request),
       signal: request.signal,
     });
     const headers = forwardableHeaders(upstream.headers);

src/server.ts

@@ -136,6 +136,7 @@ export const createBackend = (env: Env): Backend => {
       ...metered,
       typesafeApiKey: env.typesafeApiKey,
       inputPricePerMillionTokensUsd: env.typesafeInputPricePerMillionUsd,
+      maxBodyBytes: env.maxRequestBodyBytes,
     }),
     moderationRoutes({
       ...metered,