fix(vite-plugin): use module fallback service to handle wasm imports

fixtures/tanstack-start/src/routeTree.gen.ts

@@ -11,6 +11,7 @@
 import { Route as rootRouteImport } from './routes/__root'
 import { Route as AboutRouteImport } from './routes/about'
 import { Route as IndexRouteImport } from './routes/index'
+import { Route as ApiWasmRouteImport } from './routes/api.wasm'
 import { Route as ApiDbRouteImport } from './routes/api.db'
 
 const AboutRoute = AboutRouteImport.update({
@@ -23,6 +24,11 @@ const IndexRoute = IndexRouteImport.update({
   path: '/',
   getParentRoute: () => rootRouteImport,
 } as any)
+const ApiWasmRoute = ApiWasmRouteImport.update({
+  id: '/api/wasm',
+  path: '/api/wasm',
+  getParentRoute: () => rootRouteImport,
+} as any)
 const ApiDbRoute = ApiDbRouteImport.update({
   id: '/api/db',
   path: '/api/db',
@@ -33,30 +39,34 @@ export interface FileRoutesByFullPath {
   '/': typeof IndexRoute
   '/about': typeof AboutRoute
   '/api/db': typeof ApiDbRoute
+  '/api/wasm': typeof ApiWasmRoute
 }
 export interface FileRoutesByTo {
   '/': typeof IndexRoute
   '/about': typeof AboutRoute
   '/api/db': typeof ApiDbRoute
+  '/api/wasm': typeof ApiWasmRoute
 }
 export interface FileRoutesById {
   __root__: typeof rootRouteImport
   '/': typeof IndexRoute
   '/about': typeof AboutRoute
   '/api/db': typeof ApiDbRoute
+  '/api/wasm': typeof ApiWasmRoute
 }
 export interface FileRouteTypes {
   fileRoutesByFullPath: FileRoutesByFullPath
-  fullPaths: '/' | '/about' | '/api/db'
+  fullPaths: '/' | '/about' | '/api/db' | '/api/wasm'
   fileRoutesByTo: FileRoutesByTo
-  to: '/' | '/about' | '/api/db'
-  id: '__root__' | '/' | '/about' | '/api/db'
+  to: '/' | '/about' | '/api/db' | '/api/wasm'
+  id: '__root__' | '/' | '/about' | '/api/db' | '/api/wasm'
   fileRoutesById: FileRoutesById
 }
 export interface RootRouteChildren {
   IndexRoute: typeof IndexRoute
   AboutRoute: typeof AboutRoute
   ApiDbRoute: typeof ApiDbRoute
+  ApiWasmRoute: typeof ApiWasmRoute
 }
 
 declare module '@tanstack/react-router' {
@@ -75,6 +85,13 @@ declare module '@tanstack/react-router' {
       preLoaderRoute: typeof IndexRouteImport
       parentRoute: typeof rootRouteImport
     }
+    '/api/wasm': {
+      id: '/api/wasm'
+      path: '/api/wasm'
+      fullPath: '/api/wasm'
+      preLoaderRoute: typeof ApiWasmRouteImport
+      parentRoute: typeof rootRouteImport
+    }
     '/api/db': {
       id: '/api/db'
       path: '/api/db'
@@ -89,6 +106,7 @@ const rootRouteChildren: RootRouteChildren = {
   IndexRoute: IndexRoute,
   AboutRoute: AboutRoute,
   ApiDbRoute: ApiDbRoute,
+  ApiWasmRoute: ApiWasmRoute,
 }
 export const routeTree = rootRouteImport
   ._addFileChildren(rootRouteChildren)

fixtures/tanstack-start/src/routes/api.wasm.ts

@@ -0,0 +1,21 @@
+import { createFileRoute } from "@tanstack/react-router";
+
+interface Instance {
+  exports: {
+    add(a: number, b: number): number;
+  };
+}
+
+export const Route = createFileRoute("/api/wasm")({
+  server: {
+    handlers: {
+      GET: async () => {
+        const mod = await import("../wasm-example.wasm").then(
+          (m) => WebAssembly.instantiate(m.default) as Promise<Instance>,
+        );
+        const result = mod.exports.add(1, 2);
+        return Response.json({ result });
+      },
+    },
+  },
+});

fixtures/tanstack-start/test/smoke.test.ts

@@ -29,5 +29,10 @@ for (const mode of Playwright.SERVER_METHODS) {
       const response = await server.fetchJson<[{ "?column?": number }]>("/api/db");
       expect(response).toMatchObject([{ "?column?": 1 }]);
     });
+
+    it("fetches WASM", async ({ server }) => {
+      const response = await server.fetchJson<{ result: number }>("/api/wasm");
+      expect(response).toMatchObject({ result: 3 });
+    });
   });
 }

packages/cloudflare-rolldown-plugin/src/plugins/additional-modules.ts

@@ -17,7 +17,7 @@ export const MODULE_RULES = [
 ] as const;
 
 const MODULE_REFERENCE_PATTERN = `__CLOUDFLARE_MODULE__(${MODULE_RULES.map((rule) => rule.type).join("|")})__(.*?)__CLOUDFLARE_MODULE__`;
-const MODULE_REFERENCE_REGEX = new RegExp(MODULE_REFERENCE_PATTERN);
+export const MODULE_REFERENCE_REGEX = new RegExp(MODULE_REFERENCE_PATTERN);
 const MODULE_REFERENCE_GLOBAL_REGEX = new RegExp(MODULE_REFERENCE_PATTERN, "g");
 
 export const additionalModulesPlugin = createPlugin("additional-modules", () => {

packages/cloudflare-runtime/src/Runtime.ts

@@ -149,6 +149,7 @@ export const RuntimeLive = Layer.effect(
                     localDisk: storage.name,
                   },
                   containerEngine,
+                  ...worker.unsafe,
                 },
               },
               ...config.services,

packages/cloudflare-runtime/src/RuntimeWorker.ts

@@ -2,6 +2,7 @@ import type { HyperdriveOrigin } from "./bindings/hyperdrive/HyperdriveOrigin.sh
 import type { QueueConsumer } from "./bindings/queue/QueueOptions.shared.ts";
 import type { ContainerImage } from "./Docker.ts";
 import type { BindingHook } from "./PluginContext.ts";
+import type * as WorkerdConfig from "./workerd/Config.ts";
 
 export interface RuntimeWorker<B extends BindingHooks = BindingHooks> {
   readonly name: string;
@@ -18,6 +19,7 @@ export interface RuntimeWorker<B extends BindingHooks = BindingHooks> {
    * (local or in other dev instances) deliver to it.
    */
   readonly queueConsumers?: ReadonlyArray<QueueConsumer>;
+  readonly unsafe?: Partial<WorkerdConfig.Worker>;
 }
 
 export type BindingHooks = ReadonlyArray<BindingHook<any>>;

packages/cloudflare-runtime/src/registry/Registry.ts

@@ -8,6 +8,7 @@ import * as MutableHashMap from "effect/MutableHashMap";
 import * as Option from "effect/Option";
 import * as Path from "effect/Path";
 import * as Schedule from "effect/Schedule";
+import * as Semaphore from "effect/Semaphore";
 import type * as Scope from "effect/Scope";
 import * as Stream from "effect/Stream";
 import * as SubscriptionRef from "effect/SubscriptionRef";
@@ -96,6 +97,15 @@ export const RegistryLive = Layer.effect(
     // immediately (the watcher only reports future changes).
     const ref = yield* SubscriptionRef.make(yield* readAll);
 
+    // Serialize disk snapshots and writes so a poll/watch refresh that started
+    // before `write` cannot overwrite the eager in-memory update with a stale
+    // empty snapshot.
+    const updateLock = yield* Semaphore.make(1);
+    const refresh = readAll.pipe(
+      Effect.flatMap((newValue) => SubscriptionRef.set(ref, newValue)),
+      updateLock.withPermits(1),
+    );
+
     // The `fileSystemSupportsWatcher` flag is set to false on Windows and true everywhere else.
     // The flag can be overridden using a ConfigProvider, e.g. for testing.
     // If the watcher is not supported, we fall back to polling every 100ms.
@@ -108,14 +118,10 @@ export const RegistryLive = Layer.effect(
             // subscription. `mapEffect` runs the reads sequentially, so a
             // watcher-triggered read cannot be overwritten by an older one.
             Stream.merge(Stream.succeed(undefined)),
-            Stream.mapEffect(() => readAll),
+            Stream.mapEffect(() => refresh),
           )
-        : Stream.fromEffect(readAll).pipe(Stream.repeat(Schedule.spaced(100)))
-    ).pipe(
-      Stream.tap((newValue) => SubscriptionRef.set(ref, newValue)),
-      Stream.runDrain,
-      Effect.forkScoped,
-    );
+        : Stream.fromEffect(refresh).pipe(Stream.repeat(Schedule.spaced(100)))
+    ).pipe(Stream.runDrain, Effect.forkScoped);
 
     return Registry.of({
       read: (subscribers) =>
@@ -133,6 +139,7 @@ export const RegistryLive = Layer.effect(
             // Immediately update the in-memory registry so it's available without waiting on IO.
             SubscriptionRef.update(ref, (map) => MutableHashMap.set(map, entry.scriptName, entry)),
           ),
+          updateLock.withPermits(1),
           Effect.tap(() => {
             // Remove the entry from the filesystem when the scope closes — but
             // only while the file still holds THIS write's content. A

packages/cloudflare-vite-plugin/src/dev-environment.ts

@@ -1,4 +1,4 @@
-import { MODULE_RULES } from "@distilled.cloud/cloudflare-rolldown-plugin/plugins";
+import { MODULE_REFERENCE_REGEX } from "@distilled.cloud/cloudflare-rolldown-plugin/plugins";
 import assert from "node:assert";
 import * as vite from "vite";
 import type { FetchFunctionOptions } from "vite/module-runner";
@@ -42,8 +42,10 @@ export class DistilledDevEnvironment extends vite.DevEnvironment {
     importer?: string,
     options?: FetchFunctionOptions,
   ): Promise<vite.FetchResult> {
-    // Additional modules (CompiledWasm, Data, Text)
-    if (MODULE_RULES.some((rule) => rule.pattern.test(id))) {
+    // Additional modules (CompiledWasm, Data, Text) are resolved to
+    // `__CLOUDFLARE_MODULE__...` ids and must be externalized so the module
+    // runner loads them via native `import()` → workerd's module fallback.
+    if (MODULE_REFERENCE_REGEX.test(id)) {
       return {
         externalize: id,
         type: "module",

packages/cloudflare-vite-plugin/src/dev-server.ts

@@ -1,3 +1,4 @@
+import { MODULE_REFERENCE_REGEX } from "@distilled.cloud/cloudflare-rolldown-plugin/plugins";
 import type { BindingHooks, Module } from "@distilled.cloud/cloudflare-runtime";
 import * as Runtime from "@distilled.cloud/cloudflare-runtime/Runtime";
 import * as RuntimeServices from "@distilled.cloud/cloudflare-runtime/RuntimeServices";
@@ -16,6 +17,8 @@ import * as FetchHttpClient from "effect/unstable/http/FetchHttpClient";
 import * as Headers from "effect/unstable/http/Headers";
 import * as HttpServerRequest from "effect/unstable/http/HttpServerRequest";
 import * as HttpServerResponse from "effect/unstable/http/HttpServerResponse";
+import * as NodeFs from "node:fs/promises";
+import * as NodeHttp from "node:http";
 import type * as vite from "vite";
 import * as ModuleRunnerWorker from "worker:./module-runner/module-runner.worker.ts";
 import * as WrapperWorker from "worker:./module-runner/wrapper.worker.ts";
@@ -68,9 +71,8 @@ export const createDefaultContext = async (): Promise<
       accountId: process.env.CLOUDFLARE_ACCOUNT_ID!,
     },
   }).pipe(
-    Layer.provide(
-      Layer.mergeAll(Credentials.fromEnv(), importPlatformServices, FetchHttpClient.layer),
-    ),
+    Layer.provideMerge(importPlatformServices),
+    Layer.provide(Layer.merge(Credentials.fromEnv(), FetchHttpClient.layer)),
     Layer.buildWithScope(scope),
     Effect.runPromise,
   );
@@ -80,12 +82,91 @@ const closeScope = async (scope: Scope.Scope) => {
   await Effect.runPromiseExit(Scope.closeUnsafe(scope, Exit.void) ?? Effect.void);
 };
 
+const makeModuleFallbackService = Effect.gen(function* () {
+  const server = NodeHttp.createServer(async (req, res) => {
+    try {
+      const request = await parseModuleFallbackRequest(req);
+      if (!request) {
+        res.writeHead(400, { "Content-Type": "text/plain" });
+        res.end("Invalid module fallback request");
+        return;
+      }
+
+      const candidate = request.rawSpecifier ?? request.specifier;
+      const match = MODULE_REFERENCE_REGEX.exec(candidate);
+      if (!match) {
+        res.writeHead(400, { "Content-Type": "text/plain" });
+        res.end(`No match for module: ${candidate}`);
+        return;
+      }
+
+      const [, moduleType, modulePath] = match;
+      if (!moduleType || !modulePath) {
+        res.writeHead(400, { "Content-Type": "text/plain" });
+        res.end(`Invalid module type or path: ${match[0]}`);
+        return;
+      }
+
+      let content: Buffer;
+      try {
+        content = await NodeFs.readFile(modulePath);
+      } catch {
+        res.writeHead(404, { "Content-Type": "text/plain" });
+        res.end(`Module not found: ${modulePath}`);
+        return;
+      }
+
+      switch (moduleType) {
+        case "CompiledWasm": {
+          res.writeHead(200, { "Content-Type": "application/json" });
+          res.end(JSON.stringify({ wasm: Array.from(content) }));
+          return;
+        }
+        case "Data": {
+          res.writeHead(200, { "Content-Type": "application/json" });
+          res.end(JSON.stringify({ data: Array.from(content) }));
+          return;
+        }
+        case "Text": {
+          res.writeHead(200, { "Content-Type": "application/json" });
+          res.end(JSON.stringify({ text: content.toString() }));
+          return;
+        }
+        default: {
+          res.writeHead(400, { "Content-Type": "text/plain" });
+          res.end(`Invalid module type: ${moduleType}`);
+        }
+      }
+    } catch (error) {
+      res.writeHead(500, { "Content-Type": "text/plain" });
+      res.end(error instanceof Error ? (error.stack ?? error.message) : String(error));
+    }
+  });
+
+  yield* Effect.callback<void>((resume) => {
+    server.listen(0, "127.0.0.1", () => resume(Effect.void));
+  });
+  yield* Effect.addFinalizer(() =>
+    Effect.callback<void>((resume) => {
+      server.close(() => resume(Effect.void));
+    }),
+  );
+
+  const address = server.address();
+  if (address === null || typeof address === "string") {
+    return yield* Effect.die(new Error("Module fallback server address unavailable"));
+  }
+  return `127.0.0.1:${address.port}`;
+});
+
 const serve = Effect.fn(function* <B extends BindingHooks = BindingHooks>(
   options: CloudflareVitePluginOptions<B>,
   entryEnvironment: EntryEnvironment,
   server: vite.ViteDevServer,
 ) {
   const runtime = yield* Runtime.Runtime;
+  const moduleFallback = yield* makeModuleFallbackService;
+
   const name = options.worker?.name ?? `vite-dev-${crypto.randomUUID()}`;
   return yield* runtime.start({
     name,
@@ -127,9 +208,64 @@ const serve = Effect.fn(function* <B extends BindingHooks = BindingHooks>(
     ],
     hyperdrives: options.worker?.hyperdrives,
     assets: options.worker?.assets,
+    unsafe: {
+      moduleFallback,
+      ...(options.worker?.unsafe ?? {}),
+    },
   });
 });
 
+type ModuleFallbackRequest =
+  | {
+      protocol: "v1";
+      type: "import" | "require";
+      specifier: string;
+      rawSpecifier?: string;
+      referrer?: string;
+    }
+  | {
+      protocol: "v2";
+      type: "import" | "require" | "internal";
+      specifier: string;
+      rawSpecifier?: string;
+      referrer?: string;
+      attributes?: Array<{ name: string; value: string }>;
+    };
+
+const parseModuleFallbackRequest = async (
+  req: NodeHttp.IncomingMessage,
+): Promise<ModuleFallbackRequest | undefined> => {
+  if (req.method === "GET" && req.headers["x-resolve-method"]) {
+    const url = new URL(req.url ?? "", "http://localhost");
+    const specifier = url.searchParams.get("specifier");
+    if (!specifier) {
+      return;
+    }
+    const resolveMethod = req.headers["x-resolve-method"];
+    return {
+      protocol: "v1",
+      type: resolveMethod === "require" ? "require" : "import",
+      specifier,
+      rawSpecifier: url.searchParams.get("rawSpecifier") ?? undefined,
+      referrer: url.searchParams.get("referrer") ?? undefined,
+    };
+  }
+
+  if (req.method === "POST") {
+    const chunks: Array<Buffer> = [];
+    for await (const chunk of req) {
+      chunks.push(typeof chunk === "string" ? Buffer.from(chunk) : chunk);
+    }
+    const json: unknown = JSON.parse(Buffer.concat(chunks).toString());
+    if (typeof json === "object" && json !== null && "specifier" in json) {
+      return {
+        protocol: "v2",
+        ...(json as Omit<Extract<ModuleFallbackRequest, { protocol: "v2" }>, "protocol">),
+      };
+    }
+  }
+};
+
 async function makeWorkerModules(options: CloudflareVitePluginOptions): Promise<Array<Module>> {
   const [moduleRunnerWorker, wrapperWorker] = await Promise.all([
     ModuleRunnerWorker.worker(),