Update HMR for next-api (#5814)

### Description

This makes several changes to our HMR protocol to make play nicely with
Next's already existing websocket and protocol:
1. Use a `turbopack-` prefix in our message's `type`
2. Allow any compatible `sendMessage` and `addMessageListener`
implementation (so we can reuse Next's websocket handlers)
3. Create a new "dev client" entrypoint so that we can precompile with
NCC
4. Receive a pre-parsed `data` from the websocket's `message.data`
string.



### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->


Closes WEB-1452

crates/turbopack-cli/js/src/entry/client.ts

@@ -1,9 +1,14 @@
 import { connect } from "@vercel/turbopack-ecmascript-runtime/dev/client/hmr-client";
-import { connectHMR } from "@vercel/turbopack-ecmascript-runtime/dev/client/websocket";
+import {
+  connectHMR,
+  addMessageListener,
+  sendMessage,
+} from "@vercel/turbopack-ecmascript-runtime/dev/client/websocket";
 
 export function initializeHMR(options: { assetPrefix: string }) {
   connect({
-    assetPrefix: options.assetPrefix,
+    addMessageListener,
+    sendMessage,
   });
   connectHMR({
     assetPrefix: options.assetPrefix,

crates/turbopack-ecmascript-hmr-protocol/src/lib.rs

@@ -27,12 +27,14 @@ impl Display for ResourceIdentifier {
 }
 
 #[derive(Deserialize)]
-#[serde(tag = "type", rename_all = "camelCase")]
+#[serde(tag = "type")]
 pub enum ClientMessage {
+    #[serde(rename = "turbopack-subscribe")]
     Subscribe {
         #[serde(flatten)]
         resource: ResourceIdentifier,
     },
+    #[serde(rename = "turbopack-unsubscribe")]
     Unsubscribe {
         #[serde(flatten)]
         resource: ResourceIdentifier,

crates/turbopack-ecmascript-runtime/js/package.json

@@ -14,6 +14,7 @@
     "check:dev-runtime-none": "tsc -p src/dev/runtime/none"
   },
   "exports": {
+    ".": "./src/main.js",
     "./*": "./src/*.ts"
   },
   "dependencies": {

crates/turbopack-ecmascript-runtime/js/src/dev/client/hmr-client.ts

@@ -3,21 +3,32 @@
 /// <reference path="../runtime/base/protocol.d.ts" />
 /// <reference path="../runtime/base/extensions.d.ts" />
 
-import { addEventListener, sendMessage } from "./websocket";
+import {
+  addMessageListener as turboSocketAddMessageListener,
+  sendMessage as turboSocketSendMessage,
+} from "./websocket";
+type SendMessage = typeof import("./websocket").sendMessage;
 
 export type ClientOptions = {
-  assetPrefix: string;
+  addMessageListener: typeof import("./websocket").addMessageListener;
+  sendMessage: SendMessage;
 };
 
-export function connect({ assetPrefix }: ClientOptions) {
-  addEventListener((event) => {
-    switch (event.type) {
-      case "connected":
-        handleSocketConnected();
+export function connect({
+  // TODO(WEB-1465) Remove this backwards compat fallback once
+  // vercel/next.js#54586 is merged.
+  addMessageListener = turboSocketAddMessageListener,
+  // TODO(WEB-1465) Remove this backwards compat fallback once
+  // vercel/next.js#54586 is merged.
+  sendMessage = turboSocketSendMessage,
+}: ClientOptions) {
+  addMessageListener((msg) => {
+    switch (msg.type) {
+      case "turbopack-connected":
+        handleSocketConnected(sendMessage);
         break;
-      case "message":
-        const msg: ServerMessage = JSON.parse(event.message.data);
-        handleSocketMessage(msg);
+      default:
+        handleSocketMessage(msg.data as ServerMessage);
         break;
     }
   });
@@ -28,13 +39,13 @@ export function connect({ assetPrefix }: ClientOptions) {
   }
   globalThis.TURBOPACK_CHUNK_UPDATE_LISTENERS = {
     push: ([chunkPath, callback]: [ChunkPath, UpdateCallback]) => {
-      subscribeToChunkUpdate(chunkPath, callback);
+      subscribeToChunkUpdate(chunkPath, sendMessage, callback);
     },
   };
 
   if (Array.isArray(queued)) {
     for (const [chunkPath, callback] of queued) {
-      subscribeToChunkUpdate(chunkPath, callback);
+      subscribeToChunkUpdate(chunkPath, sendMessage, callback);
     }
   }
 }
@@ -46,7 +57,7 @@ type UpdateCallbackSet = {
 
 const updateCallbackSets: Map<ResourceKey, UpdateCallbackSet> = new Map();
 
-function sendJSON(message: ClientMessage) {
+function sendJSON(sendMessage: SendMessage, message: ClientMessage) {
   sendMessage(JSON.stringify(message));
 }
 
@@ -59,23 +70,26 @@ function resourceKey(resource: ResourceIdentifier): ResourceKey {
   });
 }
 
-function subscribeToUpdates(resource: ResourceIdentifier): () => void {
-  sendJSON({
-    type: "subscribe",
+function subscribeToUpdates(
+  sendMessage: SendMessage,
+  resource: ResourceIdentifier
+): () => void {
+  sendJSON(sendMessage, {
+    type: "turbopack-subscribe",
     ...resource,
   });
 
   return () => {
-    sendJSON({
-      type: "unsubscribe",
+    sendJSON(sendMessage, {
+      type: "turbopack-unsubscribe",
       ...resource,
     });
   };
 }
 
-function handleSocketConnected() {
+function handleSocketConnected(sendMessage: SendMessage) {
   for (const key of updateCallbackSets.keys()) {
-    subscribeToUpdates(JSON.parse(key));
+    subscribeToUpdates(sendMessage, JSON.parse(key));
   }
 }
 
@@ -515,29 +529,39 @@ function handleSocketMessage(msg: ServerMessage) {
   }
 }
 
-export function subscribeToChunkUpdate(
+function subscribeToChunkUpdate(
   chunkPath: ChunkPath,
+  sendMessage: SendMessage,
   callback: UpdateCallback
 ): () => void {
   return subscribeToUpdate(
     {
       path: chunkPath,
     },
+    sendMessage,
     callback
   );
 }
 
 export function subscribeToUpdate(
   resource: ResourceIdentifier,
+  sendMessage: SendMessage,
   callback: UpdateCallback
 ) {
+  // TODO(WEB-1465) Remove this backwards compat fallback once
+  // vercel/next.js#54586 is merged.
+  if (callback === undefined) {
+    callback = sendMessage;
+    sendMessage = turboSocketSendMessage;
+  }
+
   const key = resourceKey(resource);
   let callbackSet: UpdateCallbackSet;
   const existingCallbackSet = updateCallbackSets.get(key);
   if (!existingCallbackSet) {
     callbackSet = {
       callbacks: new Set([callback]),
-      unsubscribe: subscribeToUpdates(resource),
+      unsubscribe: subscribeToUpdates(sendMessage, resource),
     };
     updateCallbackSets.set(key, callbackSet);
   } else {

crates/turbopack-ecmascript-runtime/js/src/dev/client/index.ts

@@ -0,0 +1,2 @@
+export * from "./hmr-client";
+export * from "./websocket";

crates/turbopack-ecmascript-runtime/js/src/dev/client/websocket.ts

@@ -1,7 +1,7 @@
 // Adapted from https://github.com/vercel/next.js/blob/canary/packages/next/client/dev/error-overlay/websocket.ts
 
 let source: WebSocket;
-const eventCallbacks: ((event: WebsocketEvent) => void)[] = [];
+const eventCallbacks: ((msg: WebSocketMessage) => void)[] = [];
 
 // TODO: add timeout again
 // let lastActivity = Date.now()
@@ -17,16 +17,16 @@ function getSocketProtocol(assetPrefix: string): string {
   return protocol === "http:" ? "ws" : "wss";
 }
 
-type WebsocketEvent =
+export type WebSocketMessage =
   | {
-      type: "connected";
+      type: "turbopack-connected";
     }
   | {
-      type: "message";
-      message: MessageEvent;
+      type: "turbopack-message";
+      data: Record<string, any>;
     };
 
-export function addEventListener(cb: (event: WebsocketEvent) => void) {
+export function addMessageListener(cb: (msg: WebSocketMessage) => void) {
   eventCallbacks.push(cb);
 }
 
@@ -51,10 +51,9 @@ export function connectHMR(options: HMROptions) {
     console.log("[HMR] connecting...");
 
     function handleOnline() {
+      const connected = { type: "turbopack-connected" as const };
       eventCallbacks.forEach((cb) => {
-        cb({
-          type: "connected",
-        });
+        cb(connected);
       });
 
       if (options.log) console.log("[HMR] connected");
@@ -64,11 +63,12 @@ export function connectHMR(options: HMROptions) {
     function handleMessage(event: MessageEvent) {
       // lastActivity = Date.now()
 
+      const message = {
+        type: "turbopack-message" as const,
+        data: JSON.parse(event.data),
+      };
       eventCallbacks.forEach((cb) => {
-        cb({
-          type: "message",
-          message: event,
-        });
+        cb(message);
       });
     }
 

crates/turbopack-ecmascript-runtime/js/src/dev/runtime/base/protocol.d.ts

@@ -91,11 +91,11 @@ type ResourceIdentifier = {
 };
 
 type ClientMessageSubscribe = {
-  type: "subscribe";
+  type: "turbopack-subscribe";
 } & ResourceIdentifier;
 
 type ClientMessageUnsubscribe = {
-  type: "unsubscribe";
+  type: "turbopack-unsubscribe";
 } & ResourceIdentifier;
 
 type ClientMessage = ClientMessageSubscribe | ClientMessageUnsubscribe;

crates/turbopack-ecmascript-runtime/js/src/main.js

@@ -0,0 +1 @@
+// required for NCC to work