1import type { ExportTypes } from "../rolldown/export-types.ts";
2import { isExportTypes } from "../rolldown/export-types.ts";
3import { MODULE_REFERENCE_REGEX } from "../rolldown/plugins/index.ts";
4import assert from "node:assert";
5import { URL as NodeURL } from "node:url";
6import * as vite from "vite";
7import type { FetchFunctionOptions } from "vite/module-runner";
8import {
9 ENVIRONMENT_NAME_HEADER,
10 EXPORT_TYPES_EVENT,
11 INIT_PATH,
12 REQUEST_EXPORT_TYPES_EVENT,
13} from "./module-runner/constants.shared.ts";
14
15/** How long to wait for the Worker to report its export types before giving up. */
16const EXPORT_TYPES_TIMEOUT_MS = 10_000;
17
18export class DistilledDevEnvironment extends vite.DevEnvironment {
19 transport: HotChannel;
20
21 constructor(name: string, config: vite.ResolvedConfig) {
22 const transport = new HotChannel();
23 super(name, config, {
24 hot: true,
25 transport,
26 });
27 this.transport = transport;
28 }
29
30 async connect(address: string | globalThis.URL) {
31 const url = new NodeURL(address.toString());
32 url.protocol = "ws";
33 url.pathname = INIT_PATH;
34 // The Worker ambient declaration hides Bun's headers overload when both
35 // runtime type libraries are loaded by the consolidated package config.
36 const ws = new (WebSocket as unknown as {
37 new (
38 url: string,
39 options: { headers: Record<string, string> },
40 ): WebSocket;
41 })(url.toString(), {
42 headers: { [ENVIRONMENT_NAME_HEADER]: this.name },
43 });
44 await new Promise<void>((resolve, reject) => {
45 ws.addEventListener("open", () => {
46 resolve();
47 });
48 ws.addEventListener("error", (event) => {
49 // Depending on which global WebSocket type wins (bun-types vs
50 // @types/node's undici), the event may or may not carry `error`.
51 reject(
52 "error" in event
53 ? event.error
54 : new Error("WebSocket connection error"),
55 );
56 });
57 });
58 this.transport.ws = ws;
59 }
60 override async close(): Promise<void> {
61 await super.close();
62 // `transport.close()` is idempotent; make sure the module-runner socket is
63 // released even if Vite didn't close the hot channel itself.
64 this.transport.close();
65 }
66
67 /**
68 * Asks the Worker to evaluate the entry module and classify its exports.
69 *
70 * Resolves to `undefined` if the Worker could not evaluate the entry — the
71 * error is surfaced to the user by whichever request hits the entry next, and
72 * the dev server keeps whatever export types it already has.
73 */
74 async requestExportTypes(
75 timeoutMs: number = EXPORT_TYPES_TIMEOUT_MS,
76 ): Promise<ExportTypes | undefined> {
77 return await new Promise<ExportTypes | undefined>((resolve) => {
78 let timer: ReturnType<typeof setTimeout>;
79 const listener = (data: unknown) => {
80 clearTimeout(timer);
81 this.transport.off(EXPORT_TYPES_EVENT, listener);
82 resolve(isExportTypes(data) ? data : undefined);
83 };
84 timer = setTimeout(() => {
85 this.transport.off(EXPORT_TYPES_EVENT, listener);
86 resolve(undefined);
87 }, timeoutMs);
88 this.transport.on(EXPORT_TYPES_EVENT, listener);
89 this.transport.send({
90 type: "custom",
91 event: REQUEST_EXPORT_TYPES_EVENT,
92 });
93 });
94 }
95
96 override async fetchModule(
97 id: string,
98 importer?: string,
99 options?: FetchFunctionOptions,
100 ): Promise<vite.FetchResult> {
101 // Additional modules (CompiledWasm, Data, Text) are resolved to
102 // `__CLOUDFLARE_MODULE__...` ids and must be externalized so the module
103 // runner loads them via native `import()` → workerd's module fallback.
104 if (MODULE_REFERENCE_REGEX.test(id)) {
105 return {
106 externalize: id,
107 type: "module",
108 };
109 }
110 return super.fetchModule(id, importer, options);
111 }
112}
113
114class HotChannel implements vite.HotChannel {
115 #ws?: WebSocket;
116 queue?: Array<string>;
117 listeners = new Map<string, Set<vite.HotChannelListener>>();
118
119 /**
120 * Replaces the socket, which also happens when the Worker runtime is
121 * restarted mid-session.
122 *
123 * Messages are dispatched from the moment a socket is attached rather than
124 * from `listen()`: the dev server queries the Worker for its export types
125 * during `configureServer`, which is before Vite calls `listen()`.
126 */
127 set ws(ws: WebSocket) {
128 this.#ws?.removeEventListener("message", this.boundDispatch);
129 this.#ws = ws;
130 ws.addEventListener("message", this.boundDispatch);
131 if (this.queue) {
132 for (const message of this.queue) {
133 ws.send(message);
134 }
135 this.queue = undefined;
136 }
137 }
138
139 send(payload: vite.CustomPayload) {
140 const json = JSON.stringify(payload);
141 if (this.#ws) {
142 this.#ws.send(json);
143 } else {
144 this.queue ??= [];
145 this.queue.push(json);
146 }
147 }
148
149 on(event: string, listener: vite.HotChannelListener) {
150 const listeners = this.listeners.get(event) ?? new Set();
151 listeners.add(listener);
152 this.listeners.set(event, listeners);
153 }
154
155 off(event: string, listener: vite.HotChannelListener) {
156 this.listeners.get(event)?.delete(listener);
157 }
158
159 private boundDispatch = this.dispatch.bind(this);
160
161 listen() {
162 assert(this.#ws, "WebSocket is not connected");
163 // Already dispatching — see the `ws` setter.
164 }
165
166 close() {
167 // The channel can be closed before a runner ever connected (e.g. a dev
168 // server created and torn down during a framework's type generation), so
169 // tolerate a missing socket — and close it so the runtime connection does
170 // not leak across server restarts.
171 if (!this.#ws) return;
172 this.#ws.removeEventListener("message", this.boundDispatch);
173 this.#ws.close();
174 this.#ws = undefined;
175 }
176
177 private dispatch(event: MessageEvent) {
178 const payload = JSON.parse(event.data.toString()) as vite.CustomPayload;
179
180 const listeners = this.listeners.get(payload.event) ?? new Set();
181 for (const listener of listeners) {
182 listener(payload.data, this.client);
183 }
184 }
185
186 private client: vite.HotChannelClient = {
187 send: (payload) => {
188 assert(this.#ws, "WebSocket is not connected");
189
190 this.#ws.send(JSON.stringify(payload));
191 },
192 };
193}