feat(dev): expose APIs for client-server communication

docs/guide/api-hmr.md

@@ -123,3 +123,11 @@ The following HMR events are dispatched by Vite automatically:
 - `'vite:error'` when an error occurs (e.g. syntax error)
 
 Custom HMR events can also be sent from plugins. See [handleHotUpdate](./api-plugin#handlehotupdate) for more details.
+
+## `hot.send(event, payload)`
+
+Send custom events back to Vite's dev server.
+
+If called before connected, the data will be buffered and sent once the connection is established.
+
+See [Client-server Communication](/guide/api-plugin.html#client-server-communication) for more details.

docs/guide/api-plugin.md

@@ -480,7 +480,7 @@ export default defineConfig({
 
 Check out [Vite Rollup Plugins](https://vite-rollup-plugins.patak.dev) for a list of compatible official Rollup plugins with usage instructions.
 
-## Path normalization
+## Path Normalization
 
 Vite normalizes paths while resolving ids to use POSIX separators ( / ) while preserving the volume in Windows. On the other hand, Rollup keeps resolved paths untouched by default, so resolved ids have win32 separators ( \\ ) in Windows. However, Rollup plugins use a [`normalizePath` utility function](https://github.com/rollup/plugins/tree/master/packages/pluginutils#normalizepath) from `@rollup/pluginutils` internally, which converts separators to POSIX before performing comparisons. This means that when these plugins are used in Vite, the `include` and `exclude` config pattern and other similar paths against resolved ids comparisons work correctly.
 
@@ -492,3 +492,78 @@ import { normalizePath } from 'vite'
 normalizePath('foo\\bar') // 'foo/bar'
 normalizePath('foo/bar') // 'foo/bar'
 ```
+
+## Client-server Communication
+
+Since Vite 2.9, we provide some utilities for plugins to help handle the communication with clients.
+
+### Server to Client
+
+On the plugin side, we could use `server.ws.send` to boardcast events to all the clients:
+
+```js
+// vite.config.js
+export default defineConfig({
+  plugins: [
+    {
+      // ...
+      configureServer(server) {
+        server.ws.send({
+          type: 'custom',
+          event: 'greetings',
+          data: { msg: 'hello' }
+        })
+      }
+    }
+  ]
+})
+```
+
+And on the client side, we can use [`hot.on`](/guide/api-hmr.html#hot-on-event-cb) to listen to the events:
+
+```ts
+// client side
+if (import.meta.hot) {
+  import.meta.hot.on('greetings', (data) => {
+    console.log(data.msg) // hello
+  })
+}
+```
+
+### Client to Server
+
+To send events from the client to the server, we can use [`hot.send`](/guide/api-hmr.html#hot-send-event-payload):
+
+```ts
+// client side
+if (import.meta.hot) {
+  import.meta.hot.send({
+    event: 'from-client',
+    data: { msg: 'Hey!' }
+  })
+}
+```
+
+On the plugin side, we can use `server.ws.onMessage` listening to the events:
+
+```js
+// vite.config.js
+export default defineConfig({
+  plugins: [
+    {
+      // ...
+      configureServer(server) {
+        server.ws.onMessage('from-client', (data, client) => {
+          console.log('Message from client:', data.msg) // Hey!
+          // reply only to the client (if needed)
+          client.send({
+            type: 'custom',
+            event: 'ack',
+            data: { msg: 'Hi! I got your message!' }
+          })
+        })
+      }
+    }
+  ]
+})
+```

packages/playground/hmr/__tests__/hmr.spec.ts

@@ -123,11 +123,16 @@ if (!isBuild) {
     await untilUpdated(() => el.textContent(), 'edited')
   })
 
+  test('plugin client-server communication', async () => {
+    const el = await page.$('.custom-communication')
+    await untilUpdated(() => el.textContent(), '3')
+  })
+
   test('full-reload encodeURI path', async () => {
     await page.goto(
       viteTestUrl + '/unicode-path/中文-にほんご-한글-🌕🌖🌗/index.html'
     )
-    let el = await page.$('#app')
+    const el = await page.$('#app')
     expect(await el.textContent()).toBe('title')
     await editFile(
       'unicode-path/中文-にほんご-한글-🌕🌖🌗/index.html',

packages/playground/hmr/hmr.js

@@ -57,6 +57,12 @@ if (import.meta.hot) {
   import.meta.hot.on('foo', ({ msg }) => {
     text('.custom', msg)
   })
+
+  // send custom event to server to calculate 1 + 2
+  import.meta.hot.send({ event: 'remote-add', data: { a: 1, b: 2 } })
+  import.meta.hot.on('remote-add-result', ({ result }) => {
+    text('.custom-communication', result)
+  })
 }
 
 function text(el, text) {

packages/playground/hmr/index.html

@@ -5,5 +5,6 @@
 <div class="dep"></div>
 <div class="nested"></div>
 <div class="custom"></div>
+<div class="custom-communication"></div>
 <div class="css-prev"></div>
 <div class="css-post"></div>

packages/playground/hmr/vite.config.js

@@ -17,6 +17,15 @@ module.exports = {
             }
           })
         }
+      },
+      configureServer(server) {
+        server.ws.onMessage('remote-add', ({ a, b }, client) => {
+          client.send({
+            type: 'custom',
+            event: 'remote-add-result',
+            data: { result: a + b }
+          })
+        })
       }
     }
   ]

packages/vite/src/client/client.ts

@@ -6,7 +6,7 @@ import type {
   Update,
   UpdatePayload
 } from 'types/hmrPayload'
-import type { CustomEventName } from 'types/customEvent'
+import type { CustomEventName, CustomEventPayload } from 'types/customEvent'
 import { ErrorOverlay, overlayId } from './overlay'
 // eslint-disable-next-line node/no-missing-import
 import '@vite/env'
@@ -30,6 +30,7 @@ const socketHost = __HMR_PORT__
 
 const socket = new WebSocket(`${socketProtocol}://${socketHost}`, 'vite-hmr')
 const base = __BASE__ || '/'
+const messageBuffer: string[] = []
 
 function warnFailedFetch(err: Error, path: string | string[]) {
   if (!err.message.match('fetch')) {
@@ -59,9 +60,10 @@ async function handleMessage(payload: HMRPayload) {
   switch (payload.type) {
     case 'connected':
       console.log(`[vite] connected.`)
+      sendMessageBuffer()
       // proxy(nginx, docker) hmr ws maybe caused timeout,
       // so send ping package let ws keep alive.
-      setInterval(() => socket.send('ping'), __HMR_TIMEOUT__)
+      setInterval(() => socket.send('{"type":"ping"}'), __HMR_TIMEOUT__)
       break
     case 'update':
       notifyListeners('vite:beforeUpdate', payload)
@@ -361,6 +363,13 @@ async function fetchUpdate({ path, acceptedPath, timestamp }: Update) {
   }
 }
 
+function sendMessageBuffer() {
+  if (socket.readyState === 1) {
+    messageBuffer.forEach((msg) => socket.send(msg))
+    messageBuffer.length = 0
+  }
+}
+
 interface HotModule {
   id: string
   callbacks: HotCallback[]
@@ -478,6 +487,11 @@ export const createHotContext = (ownerPath: string) => {
       }
       addToMap(customListenersMap)
       addToMap(newListeners)
+    },
+
+    send: (payload: CustomEventPayload) => {
+      messageBuffer.push(JSON.stringify({ type: 'custom', ...payload }))
+      sendMessageBuffer()
     }
   }
 

packages/vite/src/node/index.ts

@@ -73,7 +73,11 @@ export type { TransformOptions as EsbuildTransformOptions } from 'esbuild'
 export type { ESBuildOptions, ESBuildTransformResult } from './plugins/esbuild'
 export type { Manifest, ManifestChunk } from './plugins/manifest'
 export type { ResolveOptions, InternalResolveOptions } from './plugins/resolve'
-export type { WebSocketServer } from './server/ws'
+export type {
+  WebSocketServer,
+  WebSocketClient,
+  WebSocketCustomListener
+} from './server/ws'
 export type { PluginContainer } from './server/pluginContainer'
 export type { ModuleGraph, ModuleNode, ResolvedUrl } from './server/moduleGraph'
 export type { SendOptions } from './server/send'

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

@@ -3,28 +3,66 @@ import type { Server } from 'http'
 import { STATUS_CODES } from 'http'
 import type { ServerOptions as HttpsServerOptions } from 'https'
 import { createServer as createHttpsServer } from 'https'
-import type { ServerOptions } from 'ws'
-import { WebSocketServer as WebSocket } from 'ws'
-import type { WebSocket as WebSocketTypes } from 'types/ws'
+import type { ServerOptions, WebSocket as WebSocketRaw } from 'ws'
+import { WebSocketServer as WebSocketServerRaw } from 'ws'
 import type { ErrorPayload, HMRPayload } from 'types/hmrPayload'
 import type { ResolvedConfig } from '..'
 import { isObject } from '../utils'
 import type { Socket } from 'net'
 export const HMR_HEADER = 'vite-hmr'
 
+export type WebSocketCustomListener<T> = (
+  data: T,
+  client: WebSocketClient
+) => void
+
 export interface WebSocketServer {
-  on: WebSocketTypes.Server['on']
-  off: WebSocketTypes.Server['off']
+  /**
+   * Get all connected clients.
+   */
+  clients: Set<WebSocketClient>
+  /**
+   * Boardcast events to all clients
+   */
   send(payload: HMRPayload): void
+  /**
+   * Disconnect all clients and terminate the server.
+   */
   close(): Promise<void>
+  /**
+   * Handle custom event emitted by `import.meta.hot.send`
+   */
+  onMessage<T>(event: string, listener: WebSocketCustomListener<T>): () => void
+  /**
+   * Listen to raw events from the WebSocket server.
+   * @advanced
+   */
+  on: WebSocketServerRaw['on']
+  /**
+   * Unregister listeners for raw WebSocket server events.
+   * @advanced
+   */
+  off: WebSocketServerRaw['off']
+}
+
+export interface WebSocketClient {
+  /**
+   * Send event to the client
+   */
+  send(payload: HMRPayload): void
+  /**
+   * The raw WebSocket instance
+   * @advanced
+   */
+  socket: WebSocketRaw
 }
 
 export function createWebSocketServer(
   server: Server | null,
   config: ResolvedConfig,
   httpsOptions?: HttpsServerOptions
 ): WebSocketServer {
-  let wss: WebSocket
+  let wss: WebSocketServerRaw
   let httpsServer: Server | undefined = undefined
 
   const hmr = isObject(config.server.hmr) && config.server.hmr
@@ -33,9 +71,11 @@ export function createWebSocketServer(
   // TODO: the main server port may not have been chosen yet as it may use the next available
   const portsAreCompatible = !hmrPort || hmrPort === config.server.port
   const wsServer = hmrServer || (portsAreCompatible && server)
+  const customListeners = new Map<string, Set<WebSocketCustomListener<any>>>()
+  const clientsMap = new WeakMap<WebSocketRaw, WebSocketClient>()
 
   if (wsServer) {
-    wss = new WebSocket({ noServer: true })
+    wss = new WebSocketServerRaw({ noServer: true })
     wsServer.on('upgrade', (req, socket, head) => {
       if (req.headers['sec-websocket-protocol'] === HMR_HEADER) {
         wss.handleUpgrade(req, socket as Socket, head, (ws) => {
@@ -76,10 +116,22 @@ export function createWebSocketServer(
     }
 
     // vite dev server in middleware mode
-    wss = new WebSocket(websocketServerOptions)
+    wss = new WebSocketServerRaw(websocketServerOptions)
   }
 
   wss.on('connection', (socket) => {
+    socket.on('message', (raw) => {
+      if (!customListeners.size) return
+      let parsed: any
+      try {
+        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 = getSocketClent(socket)
+      listeners.forEach((listener) => listener(parsed.data, client))
+    })
     socket.send(JSON.stringify({ type: 'connected' }))
     if (bufferedError) {
       socket.send(JSON.stringify(bufferedError))
@@ -96,15 +148,35 @@ export function createWebSocketServer(
     }
   })
 
+  // Provide a wrapper to the ws client so we can send messages in JSON format
+  // To be consistent with server.ws.send
+  function getSocketClent(socket: WebSocketRaw) {
+    if (!clientsMap.has(socket)) {
+      clientsMap.set(socket, {
+        send: (payload) => socket.send(JSON.stringify(payload)),
+        socket
+      })
+    }
+    return clientsMap.get(socket)!
+  }
+
   // On page reloads, if a file fails to compile and returns 500, the server
   // sends the error payload before the client connection is established.
   // If we have no open clients, buffer the error and send it to the next
   // connected client.
   let bufferedError: ErrorPayload | null = null
 
   return {
-    on: wss.on.bind(wss),
-    off: wss.off.bind(wss),
+    get on() {
+      return wss.on.bind(wss)
+    },
+    get off() {
+      return wss.off.bind(wss)
+    },
+    get clients() {
+      return new Set(Array.from(wss.clients).map(getSocketClent))
+    },
+
     send(payload: HMRPayload) {
       if (payload.type === 'error' && !wss.clients.size) {
         bufferedError = payload
@@ -143,6 +215,16 @@ export function createWebSocketServer(
           }
         })
       })
+    },
+
+    onMessage<T>(event: string, listener: WebSocketCustomListener<T>) {
+      if (!customListeners.has(event)) customListeners.set(event, new Set())
+      customListeners.get(event)!.add(listener)
+
+      const off = () => {
+        customListeners.get(event)?.delete(listener)
+      }
+      return off
     }
   }
 }

packages/vite/types/customEvent.d.ts

@@ -3,3 +3,9 @@ export type CustomEventName<T extends string> = (T extends `vite:${T}`
   ? never
   : T) &
   (`vite:${T}` extends T ? never : T)
+
+export interface CustomEventPayload {
+  type?: 'custom'
+  event: string
+  data?: any
+}