bake: make import.meta.hot.on/off accept the documented vite: event names

src/runtime/bake/hmr-module.ts

@@ -228,17 +228,13 @@ export class HMRModule {
   }
 
   on(event: string, cb: HotEventHandler) {
-    // Vite compatibility, but favor using Bun's event names.
-    if (event.startsWith("vite:")) {
-      event = "bun:" + event.slice(4);
-    }
-
+    event = normalizeEventName(event);
     (eventHandlers[event] ??= []).push(cb);
     this.dispose(() => this.off(event, cb));
   }
 
   off(event: string, cb: HotEventHandler) {
-    const handlers = eventHandlers[event];
+    const handlers = eventHandlers[normalizeEventName(event)];
     if (!handlers) return;
     const index = handlers.indexOf(cb);
     if (index !== -1) {
@@ -817,6 +813,12 @@ function createAcceptArray(modules: string[], key: Id) {
   return arr;
 }
 
+/** `vite:*` is accepted as an alias of `bun:*` for Vite compatibility. Both
+ * `on` and `off` go through this, so a handler can be removed by either name. */
+function normalizeEventName(event: string): string {
+  return event.startsWith("vite:") ? "bun:" + event.slice("vite:".length) : event;
+}
+
 export function emitEvent(event: HMREvent, data: any) {
   const handlers = eventHandlers[event];
   if (!handlers) return;

test/bake/dev/hot.test.ts

@@ -433,9 +433,12 @@ devTest("import.meta.hot on/off events", {
     }),
     "index.ts": `
       console.log("Initial setup");
-      // Add event listener
+      import.meta.hot.on("bun:beforeUpdate", () => {
+        console.log("bun:beforeUpdate (initial)");
+      });
+      // "vite:*" names are aliases of the "bun:*" events
       import.meta.hot.on("vite:beforeUpdate", () => {
-        console.log("Before update event");
+        console.log("vite:beforeUpdate (initial)");
       });
       import.meta.hot.accept();
     `,
@@ -447,28 +450,69 @@ devTest("import.meta.hot on/off events", {
       "index.ts",
       `
         console.log("Updated setup");
-        // Events implementation is partial according to docs
         import.meta.hot.on("vite:beforeUpdate", () => {
-          console.log("Before update event 2");
+          console.log("vite:beforeUpdate (updated)");
+        });
+        import.meta.hot.on("vite:afterUpdate", () => {
+          console.log("vite:afterUpdate (updated)");
         });
         const handler = () => {
-          console.log("Another handler");
+          console.log("removed handler");
         };
         import.meta.hot.on("vite:beforeUpdate", handler);
-        // Remove the handler
         import.meta.hot.off("vite:beforeUpdate", handler);
         import.meta.hot.accept();
       `,
     );
-    await c.expectMessage("Updated setup");
+    // The initial module's listeners fire for the update that replaces it. The
+    // afterUpdate listener added by the new module fires once the update is done.
+    await c.expectMessage(
+      "bun:beforeUpdate (initial)",
+      "vite:beforeUpdate (initial)",
+      "Updated setup",
+      "vite:afterUpdate (updated)",
+    );
     await dev.write(
       "index.ts",
       `
         console.log("Third update");
         import.meta.hot.accept();
       `,
     );
-    await c.expectMessage("Third update");
+    // The initial module's listeners were removed when it was replaced, and the
+    // off() listener never fires. The replaced module's afterUpdate listener is
+    // removed before this update finishes.
+    await c.expectMessage("vite:beforeUpdate (updated)", "Third update");
+  },
+});
+devTest("import.meta.hot.off accepts bun: and vite: event names interchangeably", {
+  files: {
+    "index.html": emptyHtmlFile({
+      scripts: ["index.ts"],
+    }),
+    "index.ts": `
+      console.log("Initial setup");
+      const a = () => console.log("a");
+      import.meta.hot.on("bun:beforeUpdate", a);
+      import.meta.hot.off("vite:beforeUpdate", a);
+      const b = () => console.log("b");
+      import.meta.hot.on("vite:beforeUpdate", b);
+      import.meta.hot.off("bun:beforeUpdate", b);
+      import.meta.hot.on("vite:beforeUpdate", () => console.log("still listening"));
+      import.meta.hot.accept();
+    `,
+  },
+  async test(dev) {
+    await using c = await dev.client("/");
+    await c.expectMessage("Initial setup");
+    await dev.write(
+      "index.ts",
+      `
+        console.log("Updated setup");
+        import.meta.hot.accept();
+      `,
+    );
+    await c.expectMessage("still listening", "Updated setup");
   },
 });
 devTest("hmr forwards every merged inotify sub-path from a directory batch", {