1import { DurableObject } from "cloudflare:workers";
2import {
3 ModuleRunner,
4 ssrDynamicImportKey,
5 ssrModuleExportsKey,
6} from "vite/module-runner";
7import {
8 ENVIRONMENT_NAME_HEADER,
9 EXPORT_TYPES_EVENT,
10 INIT_PATH,
11 REQUEST_EXPORT_TYPES_EVENT,
12} from "./constants.shared.ts";
13import { stripInternalEnv, type Env } from "./env.worker.ts";
14
15declare global {
16 // This global variable is accessed by `@vitejs/plugin-rsc`
17 var __VITE_ENVIRONMENT_RUNNER_IMPORT__: (
18 environmentName: string,
19 id: string,
20 ) => Promise<unknown>;
21}
22
23/**
24 * Module imports have to run inside the module runner Durable Object's
25 * `IoContext`, but they are requested from other contexts (a Worker request,
26 * a dynamic `import()`). The requesting side registers the callback here,
27 * asks the object to run it over RPC by id, and reads the result by id once
28 * the RPC returns. Both sides share this module because they share one V8
29 * isolate.
30 *
31 * Every entry is removed once the caller has read its result, whether the
32 * callback succeeded or failed. A result is a module namespace, and a
33 * namespace keeps every module it (transitively) imported alive. Retaining
34 * the results would pin each previous module graph after an HMR update or a
35 * full runner reload until the isolate ran out of heap.
36 */
37export class CallbackRegistry {
38 private nextId = 0;
39 private readonly pending = new Map<number, () => Promise<unknown>>();
40 private readonly results = new Map<number, unknown>();
41
42 /**
43 * Registers `callback` and asks `execute` to run it under its id. Resolves
44 * with the callback's result once `execute` returns.
45 */
46 async run<T>(
47 execute: (id: number) => Promise<void>,
48 callback: () => Promise<T>,
49 ): Promise<T> {
50 const id = this.nextId++;
51 this.pending.set(id, callback);
52 try {
53 await execute(id);
54 return this.results.get(id) as T;
55 } finally {
56 this.pending.delete(id);
57 this.results.delete(id);
58 }
59 }
60
61 /** Runs the callback registered under `id` and stores its result. */
62 async execute(id: number): Promise<void> {
63 const callback = this.pending.get(id);
64 if (!callback) {
65 throw new Error(`No pending callback with id ${id}`);
66 }
67 this.results.set(id, await callback());
68 }
69
70 /** Entries still registered. Zero whenever no `run` is in flight. */
71 get size(): number {
72 return this.pending.size + this.results.size;
73 }
74}
75
76const callbacks = new CallbackRegistry();
77
78/** Runs `callback` inside the module runner Durable Object's `IoContext`. */
79const runInModuleRunner = <T>(
80 env: Env,
81 callback: () => Promise<T>,
82): Promise<T> =>
83 callbacks.run(
84 (id) =>
85 env.__DISTILLED_MODULE_RUNNER__.get("singleton").executeCallback(id),
86 callback,
87 );
88
89/**
90 * Retrieves a specific export from a Worker entry module using the module runner.
91 */
92export async function getWorkerEntryExport<T>(
93 env: Env,
94 exportName: string,
95): Promise<T> {
96 const module = await globalThis.__VITE_ENVIRONMENT_RUNNER_IMPORT__(
97 env.__DISTILLED_ENVIRONMENT__.environmentName,
98 env.__DISTILLED_ENVIRONMENT__.entryId,
99 );
100
101 const exportValue =
102 typeof module === "object" &&
103 module !== null &&
104 exportName in module &&
105 (module as Record<string, unknown>)[exportName];
106
107 if (!exportValue) {
108 throw new Error(
109 `"${env.__DISTILLED_ENVIRONMENT__.entryName}" does not define a "${exportName}" export.`,
110 );
111 }
112
113 return exportValue as T;
114}
115
116export class ModuleRunnerDO extends DurableObject<Env> {
117 private webSockets = new Map<string, WebSocket>();
118 private moduleRunners = new Map<string, ModuleRunner>();
119
120 async fetch(request: Request) {
121 const { pathname } = new URL(request.url);
122 if (pathname !== INIT_PATH) {
123 throw new Error(`Invalid path: ${pathname}`);
124 }
125 globalThis.__VITE_ENVIRONMENT_RUNNER_IMPORT__ = async (
126 environmentName: string,
127 id: string,
128 ) => {
129 const moduleRunner = this.moduleRunners.get(environmentName);
130 if (!moduleRunner) {
131 throw new NotInitializedError(environmentName);
132 }
133 return runInModuleRunner(this.env, () => moduleRunner.import(id));
134 };
135 const environmentName = request.headers.get(ENVIRONMENT_NAME_HEADER);
136 if (!environmentName) {
137 throw new Error(`Missing ${ENVIRONMENT_NAME_HEADER} header`);
138 }
139 const { 0: client, 1: server } = new WebSocketPair();
140 server.accept();
141
142 const moduleRunner = this.makeModuleRunner(server, environmentName);
143 this.moduleRunners.set(environmentName, moduleRunner);
144 // `send()` writes to this socket, which is how `import.meta.hot.send()` in
145 // user code reaches the dev server.
146 this.webSockets.set(environmentName, server);
147 server.addEventListener("message", ({ data }) => {
148 if (isRequestExportTypes(data)) {
149 void this.reportExportTypes(environmentName);
150 }
151 });
152 return new Response(null, { status: 101, webSocket: client });
153 }
154
155 send(environmentName: string, data: string): void {
156 const webSocket = this.webSockets.get(environmentName);
157 if (!webSocket) {
158 throw new NotInitializedError(environmentName);
159 }
160 webSocket.send(data);
161 }
162
163 /**
164 * Evaluates the Worker entry and reports how each of its exports should be
165 * wrapped. The dev server needs this before it can serve a request, so the
166 * import runs here rather than being driven from user code.
167 */
168 private async reportExportTypes(environmentName: string): Promise<void> {
169 const { entryId, exportTypesId } = this.env.__DISTILLED_ENVIRONMENT__;
170 let data: unknown;
171 try {
172 // Both imports run inside this object's IoContext, so they can go
173 // straight to the module runner instead of through `runInModuleRunner`.
174 const { getExportTypes } = (await this.import(
175 environmentName,
176 exportTypesId,
177 )) as {
178 getExportTypes: (module: unknown) => Record<string, string>;
179 };
180 data = getExportTypes(await this.import(environmentName, entryId));
181 } catch (error) {
182 // A reply still goes out so the dev server does not wait out its timeout.
183 // The entry failing to evaluate is reported to the user through the
184 // request that triggered it; the dev server keeps its current export
185 // types.
186 // oxlint-disable-next-line no-console
187 console.error(
188 "Failed to determine the Worker entry's export types:",
189 error,
190 );
191 data = null;
192 }
193 this.send(
194 environmentName,
195 JSON.stringify({ type: "custom", event: EXPORT_TYPES_EVENT, data }),
196 );
197 }
198
199 private async import(environmentName: string, id: string): Promise<unknown> {
200 const moduleRunner = this.moduleRunners.get(environmentName);
201 if (!moduleRunner) {
202 throw new NotInitializedError(environmentName);
203 }
204 return await moduleRunner.import(id);
205 }
206
207 async executeCallback(id: number): Promise<void> {
208 await callbacks.execute(id);
209 }
210
211 makeModuleRunner(webSocket: WebSocket, environmentName: string) {
212 const env = this.env;
213 return new ModuleRunner(
214 {
215 sourcemapInterceptor: "prepareStackTrace",
216 transport: {
217 connect({ onMessage }) {
218 webSocket.addEventListener("message", async ({ data }) => {
219 onMessage(JSON.parse(data.toString()));
220 });
221
222 onMessage({
223 type: "custom",
224 event: "vite:ws:connect",
225 data: { webSocket },
226 });
227 },
228 disconnect() {
229 webSocket.close();
230 },
231 async send(data) {
232 // We send messages via a binding to the Durable Object.
233 // This is because `import.meta.send` may be called within a Worker's request context.
234 // Directly using a WebSocket created in another context would be forbidden.
235 const stub = env.__DISTILLED_MODULE_RUNNER__.get("singleton");
236 stub.send(environmentName, JSON.stringify(data));
237 },
238 invoke: async (data) => {
239 const response = await env.__DISTILLED_INVOKE_MODULE__.fetch(
240 new Request("http://localhost", {
241 method: "POST",
242 headers: {
243 "content-type": "application/json",
244 [ENVIRONMENT_NAME_HEADER]: environmentName,
245 },
246 body: JSON.stringify(data),
247 }),
248 );
249 const result = await response.json<
250 { result: unknown } | { error: unknown }
251 >();
252
253 return result;
254 },
255 },
256 hmr: true,
257 },
258 {
259 runInlinedModule: async (context, transformed, module) => {
260 // Wrap dynamic imports to route deferred dynamic imports
261 // through the DO's IoContext.
262 const originalDynamicImport = context[ssrDynamicImportKey];
263 context[ssrDynamicImportKey] = (dep) => {
264 return runInModuleRunner(env, () => originalDynamicImport(dep));
265 };
266
267 // The trailing newline ensures a `//` comment on the last line of
268 // `transformed` (e.g. a sourceMappingURL comment preserved by
269 // vite-plus) cannot swallow the closing brace.
270 const code = `"use strict";async (${Object.keys(context).join(",")})=>{${transformed}\n}`;
271 try {
272 const fn = env.__DISTILLED_UNSAFE_EVAL__.eval(code, module.id);
273 await fn(...Object.values(context));
274 Object.seal(context[ssrModuleExportsKey]);
275 } catch (error) {
276 // oxlint-disable-next-line no-console
277 console.error(
278 `[vite-plugin] Failed to evaluate inlined module "${module.id}":`,
279 error,
280 );
281 throw error;
282 }
283 },
284 runExternalModule: async (filepath) => {
285 if (filepath === "cloudflare:workers") {
286 const { env, ...mod } = await import("cloudflare:workers");
287 return Object.seal({
288 ...mod,
289 env: stripInternalEnv(env as Env),
290 });
291 }
292 return await import(filepath);
293 },
294 },
295 );
296 }
297}
298
299function isRequestExportTypes(data: string | ArrayBuffer): boolean {
300 if (typeof data !== "string" || !data.includes(REQUEST_EXPORT_TYPES_EVENT)) {
301 return false;
302 }
303 try {
304 const payload: unknown = JSON.parse(data);
305 return (
306 typeof payload === "object" &&
307 payload !== null &&
308 (payload as { event?: unknown }).event === REQUEST_EXPORT_TYPES_EVENT
309 );
310 } catch {
311 return false;
312 }
313}
314
315class NotInitializedError extends Error {
316 constructor(environmentName: string) {
317 super(
318 `Module runner not initialized for environment: "${environmentName}". If this is a child environment, make sure to set \`childEnvironments: ["${environmentName}"]\` in the plugin config.`,
319 );
320 }
321}