feat: integrate vite runtime (#155)

packages/vite-node-miniflare/examples/remix/vite.config.ts

@@ -12,6 +12,7 @@ export default defineConfig({
         "react/jsx-dev-runtime",
         "react-dom",
         "react-dom/server.browser",
+        "@remix-run/server-runtime",
       ],
     },
   },

packages/vite-node-miniflare/src/client/vite-node.ts

@@ -4,16 +4,20 @@ import {
   proxyTinyRpc,
 } from "@hiogawa/tiny-rpc";
 import { tinyassert } from "@hiogawa/utils";
+import type { HMRPayload } from "vite";
 import type { ViteNodeRunnerOptions } from "vite-node";
 import { ViteNodeRunner } from "vite-node/client";
 import { installSourcemapsSupport } from "vite-node/source-map";
+import { ViteRuntime } from "vite/runtime";
 import type { ViteNodeRpc } from "..";
 import { __setDebug } from "./polyfills/debug";
 import { __setUnsafeEval } from "./polyfills/node-vm";
 
 export interface ViteNodeMiniflareClient {
   rpc: TinyRpcProxy<ViteNodeRpc>;
   runner: ViteNodeRunner;
+  runtime: ViteRuntime;
+  runtimeHMRHandler: (payload: HMRPayload) => void;
 }
 
 export function createViteNodeClient(options: {
@@ -29,6 +33,53 @@ export function createViteNodeClient(options: {
     adapter: httpClientAdapter({ url: options.serverRpcUrl }),
   });
 
+  let runtimeHMRHandler!: (payload: HMRPayload) => void;
+
+  const runtime = new ViteRuntime(
+    {
+      root: options.runnerOptions.root,
+      fetchModule(id, importer) {
+        return rpc.ssrFetchModule(id, importer);
+      },
+      sourcemapInterceptor: "prepareStackTrace",
+      hmr: {
+        connection: {
+          isReady() {
+            return true;
+          },
+          // TODO: only for custom event to server?
+          send(messages) {
+            console.log("[runtime.hmr.connection.send]", messages);
+          },
+          // TODO: for now, we fetch HMRPayload via separate rpc, so we just grab the callback and use it later.
+          onUpdate(callback) {
+            // this is called during ViteRuntime constructor
+            runtimeHMRHandler = callback;
+          },
+        },
+        logger: console,
+      },
+    },
+    {
+      async runViteModule(context, transformed, id) {
+        // do same as vite-node/client
+        // https://github.com/vitest-dev/vitest/blob/c6e04125fb4a0af2db8bd58ea193b965d50d415f/packages/vite-node/src/client.ts#L415
+        const codeDefinition = `'use strict';async (${Object.keys(context).join(
+          ","
+        )})=>{{`;
+        const code = `${codeDefinition}${transformed}\n}}`;
+        const fn = options.unsafeEval.eval(code, id);
+        await fn(...Object.values(context));
+        Object.freeze(context.__vite_ssr_exports__);
+      },
+
+      runExternalModule(filepath) {
+        console.error("[runExternalModule]", filepath);
+        throw new Error(`[runExternalModule] ${filepath}`);
+      },
+    }
+  );
+
   const runner = new ViteNodeRunner({
     ...options.runnerOptions,
     fetchModule(id) {
@@ -42,7 +93,8 @@ export function createViteNodeClient(options: {
   // Since Vitest's getSourceMap/extractSourceMap relies on `Buffer.from(mapString, 'base64').toString('utf-8')`,
   // we inject minimal Buffer polyfill temporary during this function.
   // https://github.com/vitest-dev/vitest/blob/8dabef860a3f51f5a4c4debc10faa1837fdcdd71/packages/vite-node/src/source-map.ts#L57-L62
-  installSourcemapsSupport({
+  // prettier-ignore
+  0 && installSourcemapsSupport({
     getSourceMap: (source) => {
       const teardown = setupBufferPolyfill();
       try {
@@ -53,7 +105,7 @@ export function createViteNodeClient(options: {
     },
   });
 
-  return { rpc, runner };
+  return { rpc, runner, runtime, runtimeHMRHandler };
 }
 
 function setupBufferPolyfill() {

packages/vite-node-miniflare/src/client/worker-entry.ts

@@ -8,6 +8,7 @@ interface Env {
   __VITE_NODE_SERVER_RPC_URL: string;
   __VITE_NODE_RUNNER_OPTIONS: any;
   __VITE_NODE_DEBUG: boolean;
+  __VITE_RUNTIME_HMR: boolean;
   __WORKER_ENTRY: string;
 }
 
@@ -24,6 +25,42 @@ export default {
         debug: env.__VITE_NODE_DEBUG,
       });
 
+      if (1) {
+        // fetch HMRPayload before execution
+        // TODO: listen HMRPayload event (birpc? websocket? SSE?)
+        const payloads = await client.rpc.getHMRPayloads();
+        for (const payload of payloads) {
+          if (env.__VITE_NODE_DEBUG) {
+            console.log("[HMRPayload]", payload);
+          }
+          // simple module tree invalidation when ssr hmr is disabled
+          if (!env.__VITE_RUNTIME_HMR && payload.type === "update") {
+            for (const update of payload.updates) {
+              // TODO: unwrapId?
+              const invalidated = client.runtime.moduleCache.invalidateDepTree([
+                update.path,
+              ]);
+              if (env.__VITE_NODE_DEBUG) {
+                console.log("[vite-node-miniflare] invalidateDepTree:", [
+                  ...invalidated,
+                ]);
+              }
+            }
+            continue;
+          }
+          await (client.runtimeHMRHandler(payload) as any as Promise<void>);
+        }
+
+        const workerEntry = await client.runtime.executeEntrypoint(
+          env.__WORKER_ENTRY
+        );
+        const workerEnv = {
+          ...env,
+          __VITE_NODE_MINIFLARE_CLIENT: client,
+        };
+        return await workerEntry.default.fetch(request, workerEnv, ctx);
+      }
+
       // invalidate modules similar to nuxt
       // https://github.com/nuxt/nuxt/blob/1de44a5a5ca5757d53a8b52c9809cbc027d2d246/packages/vite/src/runtime/vite-node.mjs#L21-L23
       const invalidatedModules = await client.rpc.getInvalidatedModules();

packages/vite-node-miniflare/src/server/plugin.ts

@@ -14,6 +14,8 @@ import { setupViteNodeServerRpc } from "./vite-node";
 export function vitePluginViteNodeMiniflare(pluginOptions: {
   entry: string;
   debug?: boolean;
+  // for now disable ssr hmr by default for react plugin
+  hmr?: boolean;
   // hooks to customize options
   miniflareOptions?: (options: MiniflareOptions) => void;
   viteNodeServerOptions?: (options: ViteNodeServerOptions) => void;
@@ -82,6 +84,7 @@ export function vitePluginViteNodeMiniflare(pluginOptions: {
             entry: pluginOptions.entry,
             rpcOrigin: ctx.url.origin,
             debug: pluginOptions.debug,
+            hmr: pluginOptions.hmr,
             viteNodeRunnerOptions,
           });
           pluginOptions.miniflareOptions?.(miniflareOptions);

packages/vite-node-miniflare/src/server/vite-node.ts

@@ -1,6 +1,12 @@
 import { exposeTinyRpc, httpServerAdapter } from "@hiogawa/tiny-rpc";
 import type { MiniflareOptions } from "miniflare";
-import { type ViteDevServer, normalizePath } from "vite";
+import {
+  type HMRPayload,
+  ServerHMRConnector,
+  type ViteDevServer,
+  fetchModule,
+  normalizePath,
+} from "vite";
 import type { ViteNodeRunnerOptions } from "vite-node";
 import type { ViteNodeServer } from "vite-node/server";
 import { WORKER_ENTRY_SCRIPT } from "../client/worker-entry-script";
@@ -10,9 +16,10 @@ import { WORKER_ENTRY_SCRIPT } from "../client/worker-entry-script";
 // prettier-ignore
 export type ViteNodeRpc =
   Pick<ViteNodeServer, "fetchModule" | "resolveId"> &
-  Pick<ViteDevServer, "transformIndexHtml"> &
+  Pick<ViteDevServer, "transformIndexHtml" | "ssrFetchModule"> &
   {
     getInvalidatedModules: () => string[];
+    getHMRPayloads: () => HMRPayload[];
   };
 
 export function setupViteNodeServerRpc(
@@ -25,16 +32,35 @@ export function setupViteNodeServerRpc(
   // https://github.com/nuxt/nuxt/blob/1de44a5a5ca5757d53a8b52c9809cbc027d2d246/packages/vite/src/vite-node.ts#L62
   const invalidatedModules = new Set<string>();
 
+  // for starter, collect HMRPayload with builtin ServerHMRConnector
+  // and let worker entry fetch them via rpc before rendering
+  const connector = new ServerHMRConnector(viteNodeServer.server);
+  let hmrPayloads: HMRPayload[] = [];
+  connector.onUpdate((payload) => {
+    hmrPayloads.push(payload);
+  });
+
   const rpcRoutes: ViteNodeRpc = {
     fetchModule: viteNodeServer.fetchModule.bind(viteNodeServer),
     resolveId: viteNodeServer.resolveId.bind(viteNodeServer),
     transformIndexHtml: viteNodeServer.server.transformIndexHtml,
+    ssrFetchModule: (id, importer) => {
+      // not using default `viteDevServer.ssrFetchModule` since its source map expects mysterious two empty lines,
+      // which doesn't exist in workerd's unsafe eval
+      // https://github.com/vitejs/vite/pull/12165#issuecomment-1910686678
+      return fetchModule(viteDevServer, id, importer);
+    },
     getInvalidatedModules: () => {
       // there must be at most one client to make use of this RPC
       const result = [...invalidatedModules];
       invalidatedModules.clear();
       return result;
     },
+    getHMRPayloads: () => {
+      const result = hmrPayloads;
+      hmrPayloads = [];
+      return result;
+    },
     // framework can utilize custom RPC to implement some features on main Vite process and expose them to Workerd
     // (e.g. Remix's DevServerHooks)
     ...options.customRpc,
@@ -64,6 +90,7 @@ export function setupViteNodeServerRpc(
     entry: string;
     rpcOrigin: string;
     debug?: boolean;
+    hmr?: boolean;
     viteNodeRunnerOptions: Partial<ViteNodeRunnerOptions>;
   }) {
     return {
@@ -85,6 +112,7 @@ export function setupViteNodeServerRpc(
         __VITE_NODE_SERVER_RPC_URL: options.rpcOrigin + rpcBase,
         __VITE_NODE_RUNNER_OPTIONS: options.viteNodeRunnerOptions as any,
         __VITE_NODE_DEBUG: options.debug ?? false,
+        __VITE_RUNTIME_HMR: options.hmr ?? false,
       },
     } satisfies MiniflareOptions;
   }