feat: add vite client connect events (#20978)

Author: sapphi-red

docs/guide/api-environment-plugins.md

@@ -227,6 +227,43 @@ export default defineConfig({
 
 The `applyToEnvironment` hook is called at config time, currently after `configResolved` due to projects in the ecosystem modifying the plugins in it. Environment plugins resolution may be moved before `configResolved` in the future.
 
+## Application-Plugin Communication
+
+`environment.hot` allows plugins to communicate with the code on the application side for a given environment. This is the equivalent of [the Client-server Communication feature](/guide/api-plugin#client-server-communication), but supports environments other than the client environment.
+
+:::warning Note
+
+Note that this feature is only available for environments that supports HMR.
+
+:::
+
+### Managing the Application Instances
+
+Be aware that there might be multiple application instances running in the same environment. For example, if you multiple tabs open in the browser, each tab is a separate application instance and have a separate connection to the server.
+
+When a new connection is established, a `vite:client:connect` event is emitted on the environment's `hot` instance. When the connection is closed, a `vite:client:disconnect` event is emitted.
+
+Each event handler receives the `NormalizedHotChannelClient` as the second argument. The client is an object with a `send` method that can be used to send messages to that specific application instance. The client reference is always the same for the same connection, so you can keep it to track the connection.
+
+### Example Usage
+
+The plugin side:
+
+```js
+configureServer(server) {
+  server.environments.ssr.hot.on('my:greetings', (data, client) => {
+    // do something with the data,
+    // and optionally send a response to that application instance
+    client.send('my:foo:reply', `Hello from server! You said: ${data}`)
+  })
+
+  // broadcast a message to all application instances
+  server.environments.ssr.hot.send('my:foo', 'Hello from server!')
+}
+```
+
+The application side is same with the Client-server Communication feature. You can use the `import.meta.hot` object to send messages to the plugin.
+
 ## Environment in Build Hooks
 
 In the same way as during dev, plugin hooks also receive the environment instance during build, replacing the `ssr` boolean.

docs/guide/api-environment-runtimes.md

@@ -316,27 +316,45 @@ import { createServer, RemoteEnvironmentTransport, DevEnvironment } from 'vite'
 function createWorkerEnvironment(name, config, context) {
   const worker = new Worker('./worker.js')
   const handlerToWorkerListener = new WeakMap()
+  const client = {
+    send(payload: HotPayload) {
+      worker.postMessage(payload)
+    },
+  }
 
   const workerHotChannel = {
     send: (data) => worker.postMessage(data),
     on: (event, handler) => {
-      if (event === 'connection') return
+      // client is already connected
+      if (event === 'vite:client:connect') return
+      if (event === 'vite:client:disconnect') {
+        const listener = () => {
+          handler(undefined, client)
+        }
+        handlerToWorkerListener.set(handler, listener)
+        worker.on('exit', listener)
+        return
+      }
 
       const listener = (value) => {
         if (value.type === 'custom' && value.event === event) {
-          const client = {
-            send(payload) {
-              worker.postMessage(payload)
-            },
-          }
           handler(value.data, client)
         }
       }
       handlerToWorkerListener.set(handler, listener)
       worker.on('message', listener)
     },
     off: (event, handler) => {
-      if (event === 'connection') return
+      if (event === 'vite:client:connect') return
+      if (event === 'vite:client:disconnect') {
+        const listener = handlerToWorkerListener.get(handler)
+        if (listener) {
+          worker.off('exit', listener)
+          handlerToWorkerListener.delete(handler)
+        }
+        return
+      }
+
       const listener = handlerToWorkerListener.get(handler)
       if (listener) {
         worker.off('message', listener)
@@ -363,6 +381,8 @@ await createServer({
 
 :::
 
+Make sure to implement the `vite:client:connect` / `vite:client:disconnect` events in the `on` / `off` methods when those methods exist. `vite:client:connect` event should be emitted when the connection is established, and `vite:client:disconnect` event should be emitted when the connection is closed. The `HotChannelClient` object passed to the event handler must have the same reference for the same connection.
+
 A different example using an HTTP request to communicate between the runner and the server:
 
 ```ts

packages/vite/src/node/server/hmr.ts

@@ -145,6 +145,9 @@ export interface NormalizedHotChannel<Api = any> {
       client: NormalizedHotChannelClient,
     ) => void,
   ): void
+  /**
+   * @deprecated use `vite:client:connect` event instead
+   */
   on(event: 'connection', listener: () => void): void
   /**
    * Unregister event listener
@@ -174,6 +177,10 @@ export const normalizeHotChannel = (
     (data: any, client: NormalizedHotChannelClient) => void | Promise<void>,
     (data: any, client: HotChannelClient) => void | Promise<void>
   >()
+  const normalizedClients = new WeakMap<
+    HotChannelClient,
+    NormalizedHotChannelClient
+  >()
 
   let invokeHandlers: InvokeMethods | undefined
   let listenerForInvokeHandler:
@@ -226,22 +233,24 @@ export const normalizeHotChannel = (
         data: any,
         client: HotChannelClient,
       ) => {
-        const normalizedClient: NormalizedHotChannelClient = {
-          send: (...args) => {
-            let payload: HotPayload
-            if (typeof args[0] === 'string') {
-              payload = {
-                type: 'custom',
-                event: args[0],
-                data: args[1],
+        if (!normalizedClients.has(client)) {
+          normalizedClients.set(client, {
+            send: (...args) => {
+              let payload: HotPayload
+              if (typeof args[0] === 'string') {
+                payload = {
+                  type: 'custom',
+                  event: args[0],
+                  data: args[1],
+                }
+              } else {
+                payload = args[0]
               }
-            } else {
-              payload = args[0]
-            }
-            client.send(payload)
-          },
+              client.send(payload)
+            },
+          })
         }
-        fn(data, normalizedClient)
+        fn(data, normalizedClients.get(client)!)
       }
       normalizedListenerMap.set(fn, listenerWithNormalizedClient)
 

packages/vite/src/node/server/ws.ts

@@ -280,6 +280,20 @@ export function createWebSocketServer(
     })
   }
 
+  const emitCustomEvent = <T extends string>(
+    event: T,
+    data: InferCustomEventPayload<T>,
+    socket: WebSocketRaw,
+  ) => {
+    const listeners = customListeners.get(event)
+    if (!listeners?.size) return
+
+    const client = getSocketClient(socket)
+    for (const listener of listeners) {
+      listener(data, client)
+    }
+  }
+
   wss.on('connection', (socket) => {
     socket.on('message', (raw) => {
       if (!customListeners.size) return
@@ -288,17 +302,20 @@ export function createWebSocketServer(
         parsed = JSON.parse(String(raw))
       } catch {}
       if (!parsed || parsed.type !== 'custom' || !parsed.event) return
-      const listeners = customListeners.get(parsed.event)
-      if (!listeners?.size) return
-      const client = getSocketClient(socket)
-      listeners.forEach((listener) => listener(parsed.data, client))
+      emitCustomEvent(parsed.event, parsed.data, socket)
     })
     socket.on('error', (err) => {
       config.logger.error(`${colors.red(`ws error:`)}\n${err.stack}`, {
         timestamp: true,
         error: err,
       })
     })
+    socket.on('close', () => {
+      emitCustomEvent('vite:client:disconnect', undefined, socket)
+    })
+
+    emitCustomEvent('vite:client:connect', undefined, socket)
+
     socket.send(JSON.stringify({ type: 'connected' }))
     if (bufferedError) {
       socket.send(JSON.stringify(bufferedError))

packages/vite/src/node/ssr/runtime/__tests__/server-worker-runner.spec.ts

@@ -4,35 +4,53 @@ import type { HotChannel, HotChannelListener, HotPayload } from 'vite'
 import { DevEnvironment } from '../../..'
 import { createServer } from '../../../server'
 
-const createWorkerTransport = (w: Worker): HotChannel => {
+const createWorkerTransport = (worker: Worker): HotChannel => {
   const handlerToWorkerListener = new WeakMap<
     HotChannelListener,
     (value: HotPayload) => void
   >()
+  const client = {
+    send(payload: HotPayload) {
+      worker.postMessage(payload)
+    },
+  }
 
   return {
-    send: (data) => w.postMessage(data),
+    send: (data) => worker.postMessage(data),
     on: (event: string, handler: HotChannelListener) => {
-      if (event === 'connection') return
+      // client is already connected
+      if (event === 'vite:client:connect') return
+      if (event === 'vite:client:disconnect') {
+        const listener = () => {
+          handler(undefined, client)
+        }
+        handlerToWorkerListener.set(handler, listener)
+        worker.on('exit', listener)
+        return
+      }
 
       const listener = (value: HotPayload) => {
         if (value.type === 'custom' && value.event === event) {
-          const client = {
-            send(payload: HotPayload) {
-              w.postMessage(payload)
-            },
-          }
           handler(value.data, client)
         }
       }
       handlerToWorkerListener.set(handler, listener)
-      w.on('message', listener)
+      worker.on('message', listener)
     },
     off: (event, handler: HotChannelListener) => {
-      if (event === 'connection') return
+      if (event === 'vite:client:connect') return
+      if (event === 'vite:client:disconnect') {
+        const listener = handlerToWorkerListener.get(handler)
+        if (listener) {
+          worker.off('exit', listener)
+          handlerToWorkerListener.delete(handler)
+        }
+        return
+      }
+
       const listener = handlerToWorkerListener.get(handler)
       if (listener) {
-        w.off('message', listener)
+        worker.off('message', listener)
         handlerToWorkerListener.delete(handler)
       }
     },

packages/vite/src/node/ssr/runtime/serverModuleRunner.ts

@@ -91,13 +91,23 @@ export const createServerModuleRunnerTransport = (options: {
   return {
     connect({ onMessage }) {
       options.channel.api!.outsideEmitter.on('send', onMessage)
+      options.channel.api!.innerEmitter.emit(
+        'vite:client:connect',
+        undefined,
+        hmrClient,
+      )
       onMessage({ type: 'connected' })
       handler = onMessage
     },
     disconnect() {
       if (handler) {
         options.channel.api!.outsideEmitter.off('send', handler)
       }
+      options.channel.api!.innerEmitter.emit(
+        'vite:client:disconnect',
+        undefined,
+        hmrClient,
+      )
     },
     send(payload) {
       if (payload.type !== 'custom') {

packages/vite/types/customEvent.d.ts

@@ -6,6 +6,7 @@ import type {
 } from './hmrPayload'
 
 export interface CustomEventMap {
+  // client events
   'vite:beforeUpdate': UpdatePayload
   'vite:afterUpdate': UpdatePayload
   'vite:beforePrune': PrunePayload
@@ -14,6 +15,10 @@ export interface CustomEventMap {
   'vite:invalidate': InvalidatePayload
   'vite:ws:connect': WebSocketConnectionPayload
   'vite:ws:disconnect': WebSocketConnectionPayload
+
+  // server events
+  'vite:client:connect': undefined
+  'vite:client:disconnect': undefined
 }
 
 export interface WebSocketConnectionPayload {

playground/hmr-ssr/__tests__/hmr-ssr.spec.ts

@@ -10,7 +10,12 @@ import {
   test,
   vi,
 } from 'vitest'
-import type { InlineConfig, RunnableDevEnvironment, ViteDevServer } from 'vite'
+import type {
+  InlineConfig,
+  Plugin,
+  RunnableDevEnvironment,
+  ViteDevServer,
+} from 'vite'
 import { createRunnableDevEnvironment, createServer } from 'vite'
 import type { ModuleRunner } from 'vite/module-runner'
 import {
@@ -47,12 +52,31 @@ const updated = (file: string, via?: string) => {
 
 if (!isBuild) {
   describe('hmr works correctly', () => {
+    const hotEventCounts = { connect: 0, disconnect: 0 }
+
     beforeAll(async () => {
-      await setupModuleRunner('/hmr.ts')
+      function hotEventsPlugin(): Plugin {
+        return {
+          name: 'hot-events',
+          configureServer(server) {
+            server.environments.ssr.hot.on(
+              'vite:client:connect',
+              () => hotEventCounts.connect++,
+            )
+            server.environments.ssr.hot.on(
+              'vite:client:disconnect',
+              () => hotEventCounts.disconnect++,
+            )
+          },
+        }
+      }
+
+      await setupModuleRunner('/hmr.ts', { plugins: [hotEventsPlugin()] })
     })
 
     test('should connect', async () => {
       expect(clientLogs).toContain('[vite] connected.')
+      expect(hotEventCounts).toStrictEqual({ connect: 1, disconnect: 0 })
     })
 
     test('self accept', async () => {

playground/hmr/__tests__/hmr.spec.ts

@@ -29,6 +29,25 @@ if (!isBuild) {
     browserLogs.length = 0
   })
 
+  const fetchHotEvents = async (): Promise<{
+    connectCount: number
+    disconnectCount: number
+  }> => {
+    const res = await fetch(viteTestUrl + '/hot-events-counts')
+    return res.json()
+  }
+  test('hot events', async () => {
+    expect(await fetchHotEvents()).toStrictEqual({
+      connectCount: 1,
+      disconnectCount: 0,
+    })
+    await untilBrowserLogAfter(() => page.reload(), [/connected/])
+    expect(await fetchHotEvents()).toStrictEqual({
+      connectCount: 2,
+      disconnectCount: 1,
+    })
+  })
+
   test('self accept', async () => {
     const el = await page.$('.app')
     await untilBrowserLogAfter(