From cc7670abaffeda1338cf3acfef2bc41a38c223a0 Mon Sep 17 00:00:00 2001
From: Anthony Fu <anthonyfu117@hotmail.com>
Date: Wed, 27 Nov 2024 12:52:27 +0800
Subject: [PATCH] fix(server): skip hot channel client normalization for
 wsServer  (#18782)

Co-authored-by: sapphi-red <49056869+sapphi-red@users.noreply.github.com>
---
 packages/vite/src/node/server/hmr.ts |  5 +++--
 packages/vite/src/node/server/ws.ts  | 23 ++++++++++++++++++-----
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/packages/vite/src/node/server/hmr.ts b/packages/vite/src/node/server/hmr.ts
index 265775e5e40c0b..fa15d0895803bf 100644
--- a/packages/vite/src/node/server/hmr.ts
+++ b/packages/vite/src/node/server/hmr.ts
@@ -173,6 +173,7 @@ export interface NormalizedHotChannel<Api = any> {
 export const normalizeHotChannel = (
   channel: HotChannel,
   enableHmr: boolean,
+  normalizeClient = true,
 ): NormalizedHotChannel => {
   const normalizedListenerMap = new WeakMap<
     (data: any, client: NormalizedHotChannelClient) => void | Promise<void>,
@@ -225,7 +226,7 @@ export const normalizeHotChannel = (
       event: string,
       fn: (data: any, client: NormalizedHotChannelClient) => void,
     ) => {
-      if (event === 'connection') {
+      if (event === 'connection' || !normalizeClient) {
         channel.on?.(event, fn as () => void)
         return
       }
@@ -260,7 +261,7 @@ export const normalizeHotChannel = (
       listenersForEvents.get(event)!.add(listenerWithNormalizedClient)
     },
     off: (event: string, fn: () => void) => {
-      if (event === 'connection') {
+      if (event === 'connection' || !normalizeClient) {
         channel.off?.(event, fn as () => void)
         return
       }
diff --git a/packages/vite/src/node/server/ws.ts b/packages/vite/src/node/server/ws.ts
index 70e21f3e9e6ce5..8b8081f43bfbb7 100644
--- a/packages/vite/src/node/server/ws.ts
+++ b/packages/vite/src/node/server/ws.ts
@@ -9,11 +9,12 @@ import colors from 'picocolors'
 import type { WebSocket as WebSocketRaw } from 'ws'
 import { WebSocketServer as WebSocketServerRaw_ } from 'ws'
 import type { WebSocket as WebSocketTypes } from 'dep-types/ws'
-import type { ErrorPayload } from 'types/hmrPayload'
+import type { ErrorPayload, HotPayload } from 'types/hmrPayload'
 import type { InferCustomEventPayload } from 'types/customEvent'
-import type { HotChannelClient, ResolvedConfig } from '..'
+import type { ResolvedConfig } from '..'
 import { isObject } from '../utils'
-import { type NormalizedHotChannel, normalizeHotChannel } from './hmr'
+import type { NormalizedHotChannel, NormalizedHotChannelClient } from './hmr'
+import { normalizeHotChannel } from './hmr'
 import type { HttpServer } from '.'
 
 /* In Bun, the `ws` module is overridden to hook into the native code. Using the bundled `js` version
@@ -66,7 +67,7 @@ export interface WebSocketServer extends NormalizedHotChannel {
   clients: Set<WebSocketClient>
 }
 
-export interface WebSocketClient extends HotChannelClient {
+export interface WebSocketClient extends NormalizedHotChannelClient {
   /**
    * The raw WebSocket instance
    * @advanced
@@ -255,7 +256,17 @@ export function createWebSocketServer(
   function getSocketClient(socket: WebSocketRaw) {
     if (!clientsMap.has(socket)) {
       clientsMap.set(socket, {
-        send: (payload) => {
+        send: (...args: any[]) => {
+          let payload: HotPayload
+          if (typeof args[0] === 'string') {
+            payload = {
+              type: 'custom',
+              event: args[0],
+              data: args[1],
+            }
+          } else {
+            payload = args[0]
+          }
           socket.send(JSON.stringify(payload))
         },
         socket,
@@ -329,6 +340,8 @@ export function createWebSocketServer(
       },
     },
     config.server.hmr !== false,
+    // Don't normalize client as we already handles the send, and to keep `.socket`
+    false,
   )
   return {
     ...normalizedHotChannel,