feat: neural edge surface on /v1/infer (WO15)

scripts/gen-models.mjs

@@ -8,6 +8,7 @@ const dir = process.env.INTERSCRIPT_ML_DIR ?? "../ml-models"
 const data = parse(readFileSync(`${dir}/models.yaml`, "utf8"))
 const fields = [
   "id",
+  "kind",
   "task",
   "scripts",
   "precision",

src/ml.ts

@@ -0,0 +1,132 @@
+/**
+ * WO15: the neural edge surface — kind-aware execution on the worker.
+ *
+ * Plane artifacts small enough for worker memory (budget below) run
+ * IN-WORKER: zip fetched from the Releases data channel (Cache API
+ * between calls), sha-verified, decoded by the interscript-ts plane
+ * runtime. Larger models answer 503 with a pointer to the
+ * client/runtime path — workerd caps isolate memory at 128 MB, and we
+ * say so rather than OOM.
+ */
+
+import { createPlaneModel } from "interscript/ml"
+import { getModel } from "./models.js"
+
+/** In-worker execution budget (bytes). 100 MB leaves headroom inside
+ * the 128 MB workerd isolate limit for the session + decode. */
+export const WORKER_MODEL_MAX_BYTES = 100 * 1024 * 1024
+
+export class MlError extends Error {
+  constructor(
+    readonly status: number,
+    readonly code: string,
+    message: string,
+  ) {
+    super(message)
+  }
+}
+
+export interface MlRequest {
+  readonly id: string
+  readonly text: string
+  readonly preserveDiacritics?: boolean
+}
+
+/** A resolved artifact: from the index, or supplied directly (tests,
+ * self-hosters with mirrors). */
+export interface PlaneArtifact {
+  readonly url: string
+  readonly sha256: string
+  readonly size?: number | undefined
+}
+
+interface ZipCache {
+  match(key: RequestInfo): Promise<Response | undefined>
+  put(key: RequestInfo, res: Response): Promise<void>
+}
+
+function cacheInterface(): ZipCache | undefined {
+  if (typeof caches === "undefined") return undefined
+  // workerd exposes caches.default; tests run without it (no cache)
+  return (caches as unknown as { default?: ZipCache }).default
+}
+
+async function fetchZip(url: string, sha256: string): Promise<Uint8Array> {
+  const cache = cacheInterface()
+  const key = `https://ml.interscript.internal/zip/${sha256}`
+  const hit = await cache?.match(key)
+  if (hit) return new Uint8Array(await hit.arrayBuffer())
+  const res = await fetch(url)
+  if (!res.ok) {
+    throw new MlError(502, "artifact_fetch_failed", `artifact fetch failed: ${res.status}`)
+  }
+  const bytes = new Uint8Array(await res.arrayBuffer())
+  const digest = await crypto.subtle.digest("SHA-256", bytes.buffer as ArrayBuffer)
+  const got = Array.from(new Uint8Array(digest))
+    .map((b) => b.toString(16).padStart(2, "0"))
+    .join("")
+  if (got !== sha256) {
+    throw new MlError(502, "artifact_sha_mismatch", `artifact sha256 mismatch: ${got}`)
+  }
+  await cache?.put(key, new Response(bytes))
+  return bytes
+}
+
+/** Errors that mean "not executable here, try the fallback path". */
+export function isEdgeIneligible(err: unknown): boolean {
+  return err instanceof MlError && err.status >= 500
+}
+
+export async function runModel(req: MlRequest): Promise<string> {
+  const entry = getModel(req.id)
+  if (!entry) {
+    throw new MlError(404, "unknown_model", `unknown model id '${req.id}'`)
+  }
+  const kind = (entry as { kind?: string }).kind ?? "imf-v1"
+  if (kind !== "plane") {
+    throw new MlError(
+      501, "kind_not_edge_executable",
+      `kind '${kind}' does not run in-worker yet; use the interscript runtime (py/ruby/ts)`,
+    )
+  }
+  if (!entry.url || !entry.sha256) {
+    throw new MlError(503, "artifact_unresolvable", "model has no release artifact")
+  }
+  const artifact: PlaneArtifact = {
+    url: entry.url, sha256: entry.sha256, size: entry.size ?? undefined,
+  }
+  const args: Parameters<typeof executePlane>[0] = {
+    artifact, text: req.text,
+  }
+  if (req.preserveDiacritics !== undefined) {
+    args.preserveDiacritics = req.preserveDiacritics
+  }
+  return executePlane(args)
+}
+
+export interface ExecutePlaneArgs {
+  readonly artifact: PlaneArtifact
+  readonly text: string
+  preserveDiacritics?: boolean
+}
+
+export async function executePlane(args: ExecutePlaneArgs): Promise<string> {
+  const { artifact, text } = args
+  if ((artifact.size ?? 0) > WORKER_MODEL_MAX_BYTES) {
+    throw new MlError(
+      503, "model_too_large_for_worker",
+      `artifact exceeds the in-worker budget (${WORKER_MODEL_MAX_BYTES} bytes); ` +
+        `workerd isolates cap at 128 MB. Run via the interscript runtime ` +
+        `(py/ruby/ts) or a self-hosted inference service.`,
+    )
+  }
+  const zip = await fetchZip(artifact.url, artifact.sha256)
+  const model = await createPlaneModel(zip)
+  try {
+    return await model.transform(text, {
+      preserveDiacritics: args.preserveDiacritics ?? false,
+    })
+  } finally {
+    await model.dispose()
+  }
+}

src/rest.ts

@@ -27,6 +27,7 @@ import { InputTooLongError, MapNotFoundError } from "./errors.js"
 import { INFER_TIMEOUT_MS, LIMITS } from "./limits.js"
 import { bundledSystemCodes } from "./engine.js"
 import { getModel, listModels, MODELS_INDEX_VERSION } from "./models.js"
+import { isEdgeIneligible, MlError, runModel } from "./ml.js"
 import { OPENAPI } from "./openapi.js"
 import { docsPage } from "./docs.js"
 
@@ -225,6 +226,27 @@ rest.post("/v1/infer", async (c) => {
   if (!getModel(model)) {
     return errorResponse(404, "model_not_found", `Couldn't locate ${model}`)
   }
+  // Edge-first (WO15): plane artifacts inside the worker memory budget
+  // execute IN-WORKER; anything else falls through to the proxy.
+  const preserve = (body as { preserve_diacritics?: boolean }).preserve_diacritics
+  try {
+    const req: { id: string; text: string; preserveDiacritics?: boolean } = {
+      id: model, text: input,
+    }
+    if (preserve !== undefined) req.preserveDiacritics = preserve
+    const output = await runModel(req)
+    return c.json({ model, input, output, executed: "in-worker" })
+  } catch (err) {
+    if (!isEdgeIneligible(err)) {
+      if (err instanceof MlError && err.status === 404) {
+        return errorResponse(404, "model_not_found", `Couldn't locate ${model}`)
+      }
+      if (err instanceof MlError && err.status === 502) {
+        return errorResponse(502, err.code, err.message)
+      }
+      throw err
+    }
+  }
   const { ML_ENDPOINT, ML_TOKEN } = (c.env ?? {}) as Record<string, string | undefined>
   if (!ML_ENDPOINT || !ML_TOKEN) {
     return errorResponse(

src/types/interscript-ml.d.ts

@@ -0,0 +1,17 @@
+/**
+ * Augmentation: the plane loader ships in the next interscript
+ * release (interscript-ts#99/#100). Until the dependency bumps, the
+ * worker declares the contract it consumes; vitest aliases
+ * "interscript/ml" to the local interscript-ts source so tests run
+ * the REAL implementation.
+ */
+declare module "interscript/ml" {
+  export interface PlaneTransformOptions {
+    preserveDiacritics?: boolean
+  }
+  export interface PlaneModelLike {
+    transform(input: string, opts?: PlaneTransformOptions): Promise<string>
+    dispose(): Promise<void>
+  }
+  export function createPlaneModel(zipBytes: Uint8Array): Promise<PlaneModelLike>
+}

test/ml.test.ts

@@ -0,0 +1,111 @@
+/**
+ * WO15: the neural edge surface — in-worker plane execution.
+ * Runs the REAL plane runtime (vitest alias -> interscript-ts source)
+ * against a REAL local HTTP server serving the tiny-plane fixture.
+ */
+
+import { createServer, type Server } from "node:http"
+import { readFileSync } from "node:fs"
+import { afterAll, beforeAll, describe, expect, it } from "vitest"
+import {
+  executePlane,
+  isEdgeIneligible,
+  MlError,
+  runModel,
+  WORKER_MODEL_MAX_BYTES,
+} from "../src/ml.js"
+import { createHash } from "node:crypto"
+
+const FIXTURE = new URL("./fixtures/tiny-plane.zip", import.meta.url)
+const bytes = new Uint8Array(readFileSync(FIXTURE))
+const sha = createHash("sha256").update(bytes).digest("hex")
+
+let server: Server
+let port = 0
+
+beforeAll(async () => {
+  server = createServer((req, res) => {
+    res.writeHead(200, { "content-type": "application/zip" })
+    res.end(Buffer.from(bytes))
+  })
+  await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve))
+  port = (server.address() as { port: number }).port
+})
+
+afterAll(async () => {
+  await new Promise<void>((resolve) => server.close(() => resolve()))
+})
+
+const artifact = () => ({ url: `http://127.0.0.1:${port}/tiny-plane.zip`, sha256: sha })
+
+describe("executePlane (real fetch -> verify -> session -> decode)", () => {
+  it("decodes the fixture through the full chain", async () => {
+    const out = await executePlane({ artifact: artifact(), text: "كتب" })
+    expect(out).toBe("كُتُبُ") // MASK(4) -> 1 -> 2: damma on k=2
+  })
+
+  it("preserve mode round-trips user diacritics", async () => {
+    const out = await executePlane({
+      artifact: artifact(), text: "كُتُبُ", preserveDiacritics: true,
+    })
+    expect(out).toBe("كُتُبُ")
+  })
+
+  it("rejects a sha mismatch loudly", async () => {
+    await expect(
+      executePlane({ artifact: { ...artifact(), sha256: "0".repeat(64) }, text: "كتب" }),
+    ).rejects.toThrow(/sha256 mismatch/)
+  })
+
+  it("refuses artifacts above the in-worker budget", async () => {
+    await expect(
+      executePlane({
+        artifact: { ...artifact(), size: WORKER_MODEL_MAX_BYTES + 1 },
+        text: "كتب",
+      }),
+    ).rejects.toThrow(/in-worker budget/)
+  })
+})
+
+describe("runModel dispatch", () => {
+  it("404s on unknown ids", async () => {
+    await expect(runModel({ id: "nope", text: "x" })).rejects.toMatchObject({
+      status: 404,
+    })
+  })
+
+  it("marks server-side ineligibility for the proxy fallback", () => {
+    const err = new MlError(503, "model_too_large_for_worker", "big")
+    expect(isEdgeIneligible(err)).toBe(true)
+    expect(isEdgeIneligible(new MlError(404, "unknown_model", "x"))).toBe(false)
+    expect(isEdgeIneligible(new Error("x"))).toBe(false)
+  })
+})
+
+describe("POST /v1/infer edge-first dispatch", () => {
+  it("404s unknown models (route contract)", async () => {
+    const { rest } = await import("../src/rest.js")
+    const res = await rest.request("/v1/infer", {
+      method: "POST",
+      body: JSON.stringify({ model: "does-not-exist", input: "x" }),
+    })
+    expect(res.status).toBe(404)
+    expect(((await res.json()) as { error: { code: string } }).error.code).toBe(
+      "model_not_found",
+    )
+  })
+
+  it("falls back to the proxy (503 unconfigured) for oversized plane artifacts", async () => {
+    const { rest } = await import("../src/rest.js")
+    // heb-diac-plane-2.0 is 416MB — over the in-worker budget, no
+    // ML_ENDPOINT configured in tests -> proxy path -> 503
+    const res = await rest.request("/v1/infer", {
+      method: "POST",
+      body: JSON.stringify({ model: "heb-diac-plane-2.0", input: "שלום" }),
+    })
+    expect(res.status).toBe(503)
+    expect(((await res.json()) as { error: { code: string } }).error.code).toBe(
+      "inference_unconfigured",
+    )
+  })
+})

vitest.config.ts

@@ -1,6 +1,17 @@
 import { defineConfig } from "vitest/config"
 
+import { fileURLToPath } from "node:url"
+
 export default defineConfig({
+  resolve: {
+    alias: {
+      // tests run the REAL plane runtime from the local interscript-ts
+      // checkout (the published package lags by one release)
+      "interscript/ml": fileURLToPath(
+        new URL("../interscript-ts/src/ml/index.ts", import.meta.url),
+      ),
+    },
+  },
   test: {
     exclude: ["**/node_modules/**", "**/dist/**", "vendor/**", "coverage/**"],
   },