feat(core): add getCurrentRpcSession support on node side

Author: antfu

packages/core/src/node/host-functions.ts

@@ -1,12 +1,14 @@
-import type { DevToolsNodeContext, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, RpcFunctionsHost as RpcFunctionsHostType } from '@vitejs/devtools-kit'
+import type { DevToolsNodeContext, DevToolsNodeRpcSession, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, RpcFunctionsHost as RpcFunctionsHostType } from '@vitejs/devtools-kit'
 import type { BirpcGroup } from 'birpc'
+import type { AsyncLocalStorage } from 'node:async_hooks'
 import { RpcFunctionsCollectorBase } from 'birpc-x'
 
 export class RpcFunctionsHost extends RpcFunctionsCollectorBase<DevToolsRpcServerFunctions, DevToolsNodeContext> implements RpcFunctionsHostType {
   /**
    * @internal
    */
-  rpcGroup: BirpcGroup<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, false> = undefined!
+  _rpcGroup: BirpcGroup<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, false> = undefined!
+  _asyncStorage: AsyncLocalStorage<DevToolsNodeRpcSession> = undefined!
 
   constructor(context: DevToolsNodeContext) {
     super(context)
@@ -19,9 +21,15 @@ export class RpcFunctionsHost extends RpcFunctionsCollectorBase<DevToolsRpcServe
     name: T,
     ...args: Args
   ): Promise<(Awaited<ReturnType<DevToolsRpcClientFunctions[T]>> | undefined)[]> {
-    if (!this.rpcGroup)
-      throw new Error('RpcFunctionsHost.rpcGroup is not set, it likely to be an internal bug of Vite DevTools')
+    if (!this._rpcGroup)
+      throw new Error('RpcFunctionsHost] RpcGroup is not set, it likely to be an internal bug of Vite DevTools')
     // @ts-expect-error - BirpcGroup.broadcast.$callOptional is not typed correctly
-    return this.rpcGroup.broadcast.$callOptional<T>(name, ...args)
+    return this._rpcGroup.broadcast.$callOptional<T>(name, ...args)
+  }
+
+  getCurrentRpcSession(): DevToolsNodeRpcSession | undefined {
+    if (!this._asyncStorage)
+      throw new Error('RpcFunctionsHost] AsyncLocalStorage is not set, it likely to be an internal bug of Vite DevTools')
+    return this._asyncStorage.getStore()
   }
 }

packages/core/src/node/rpc/internal/docks-list.ts

@@ -18,10 +18,14 @@ export const docksList = defineRpcFunction({
     ]
 
     return {
-      handler: (): DevToolsDockEntry[] => [
-        ...Array.from(context.docks.values()),
-        ...builtinDocksEntries,
-      ],
+      handler(): DevToolsDockEntry[] {
+        // const session = context.rpc.getCurrentRpcSession()
+        // console.log('session', session?.meta.id)
+        return [
+          ...Array.from(context.docks.values()),
+          ...builtinDocksEntries,
+        ]
+      },
     }
   },
 })

packages/core/src/node/ws.ts

@@ -1,7 +1,8 @@
 /* eslint-disable no-console */
-import type { ConnectionMeta, DevToolsNodeContext, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from '@vitejs/devtools-kit'
+import type { ConnectionMeta, DevToolsNodeContext, DevToolsNodeRpcSession, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from '@vitejs/devtools-kit'
 import type { WebSocket } from 'ws'
 import type { RpcFunctionsHost } from './host-functions'
+import { AsyncLocalStorage } from 'node:async_hooks'
 import { createRpcServer } from '@vitejs/devtools-rpc'
 import { createWsRpcPreset } from '@vitejs/devtools-rpc/presets/ws/server'
 import c from 'ansis'
@@ -23,16 +24,18 @@ export async function createWsServer(options: CreateWsServerOptions) {
 
   const preset = createWsRpcPreset({
     port: port!,
-    onConnected: (ws) => {
+    onConnected: (ws, meta) => {
       wsClients.add(ws)
-      console.log(c.green`${MARK_CHECK} Websocket client connected`)
+      console.log(c.green`${MARK_CHECK} Websocket client [${meta.id}] connected`)
     },
-    onDisconnected: (ws) => {
+    onDisconnected: (ws, meta) => {
       wsClients.delete(ws)
-      console.log(c.red`${MARK_CHECK} Websocket client disconnected`)
+      console.log(c.red`${MARK_CHECK} Websocket client [${meta.id}] disconnected`)
     },
   })
 
+  const asyncStorage = new AsyncLocalStorage<DevToolsNodeRpcSession>()
+
   const rpcGroup = createRpcServer<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions>(
     rpcHost.functions,
     {
@@ -46,11 +49,26 @@ export async function createWsServer(options: CreateWsServerOptions) {
           console.error(c.red`⬢ RPC error on executing rpc`)
           console.error(error)
         },
+        resolver(name, fn) {
+          if (!fn)
+            return undefined
+          // eslint-disable-next-line ts/no-this-alias
+          const rpc = this
+          return async function (this: any, ...args) {
+            return await asyncStorage.run({
+              rpc,
+              meta: rpc.$meta,
+            }, async () => {
+              return (await fn).apply(this, args)
+            })
+          }
+        },
       },
     },
   )
 
-  rpcHost.rpcGroup = rpcGroup
+  rpcHost._rpcGroup = rpcGroup
+  rpcHost._asyncStorage = asyncStorage
 
   const getConnectionMeta = async (): Promise<ConnectionMeta> => {
     return {

packages/kit/src/types/rpc.ts

@@ -1,13 +1,35 @@
+import type { BirpcReturn } from 'birpc'
 import type { RpcFunctionsCollectorBase } from 'birpc-x'
+import type { WebSocket } from 'ws'
 import type { DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from './rpc-augments'
 import type { DevToolsNodeContext } from './vite-plugin'
 
+export interface DevToolsNodeRpcSessionMeta {
+  id: number
+  ws?: WebSocket
+}
+
+export interface DevToolsNodeRpcSession {
+  meta: DevToolsNodeRpcSessionMeta
+  rpc: BirpcReturn<DevToolsRpcClientFunctions, DevToolsRpcServerFunctions, false>
+}
+
 export type RpcFunctionsHost = RpcFunctionsCollectorBase<DevToolsRpcServerFunctions, DevToolsNodeContext> & {
+  /**
+   * Broadcast a message to all connected clients
+   */
   boardcast: <
     T extends keyof DevToolsRpcClientFunctions,
     Args extends Parameters<DevToolsRpcClientFunctions[T]>,
   >(
     name: T,
     ...args: Args
   ) => Promise<(Awaited<ReturnType<DevToolsRpcClientFunctions[T]>> | undefined)[]>
+
+  /**
+   * Get the current RPC client
+   *
+   * Available in RPC functions to get the current RPC client
+   */
+  getCurrentRpcSession: () => DevToolsNodeRpcSession | undefined
 }

packages/kit/src/types/vite-plugin.ts

@@ -19,14 +19,43 @@ export interface DevToolsPluginOptions {
 }
 
 export interface DevToolsNodeContext {
+  /**
+   * Current working directory of Vite DevTools
+   */
   readonly cwd: string
+  /**
+   * Current mode of Vite DevTools
+   * - 'dev' - when Vite DevTools is running in dev mode
+   * - 'build' - when Vite DevTools is running in build mode (no server)
+   */
   readonly mode: 'dev' | 'build'
+  /**
+   * Resolved Vite configuration
+   */
   readonly viteConfig: ResolvedConfig
+  /**
+   * Vite dev server instance (only available in dev mode)
+   */
   readonly viteServer?: ViteDevServer
+  /**
+   * RPC functions host, for registering server-side RPC functions and calling client-side RPC functions
+   */
   rpc: RpcFunctionsHost
+  /**
+   * Docks host, for registering dock entries
+   */
   docks: DevToolsDockHost
+  /**
+   * Views host, for registering static views
+   */
   views: DevToolsViewHost
+  /**
+   * Utils for the node context
+   */
   utils: DevToolsNodeUtils
+  /**
+   * Terminals host, for registering terminal sessions and streaming terminal output
+   */
   terminals: DevToolsTerminalHost
 }
 

packages/rpc/src/presets/ws/server.ts

@@ -1,3 +1,4 @@
+import type { DevToolsNodeRpcSessionMeta } from '@vitejs/devtools-kit'
 import type { BirpcGroup, BirpcOptions, ChannelOptions } from 'birpc'
 import type { WebSocket } from 'ws'
 import type { RpcServerPreset } from '..'
@@ -7,10 +8,12 @@ import { defineRpcServerPreset } from '..'
 
 export interface WebSocketRpcServerOptions {
   port: number
-  onConnected?: (ws: WebSocket) => void
-  onDisconnected?: (ws: WebSocket) => void
+  onConnected?: (ws: WebSocket, meta: DevToolsNodeRpcSessionMeta) => void
+  onDisconnected?: (ws: WebSocket, meta: DevToolsNodeRpcSessionMeta) => void
 }
 
+let id = 0
+
 function NOOP() {}
 
 export const createWsRpcPreset: RpcServerPreset<
@@ -43,6 +46,10 @@ export const createWsRpcPreset: RpcServerPreset<
     } = options ?? {}
 
     wss.on('connection', (ws) => {
+      const meta: DevToolsNodeRpcSessionMeta = {
+        id: id++,
+        ws,
+      }
       const channel: ChannelOptions = {
         post: (data) => {
           ws.send(data)
@@ -54,6 +61,7 @@ export const createWsRpcPreset: RpcServerPreset<
         },
         serialize,
         deserialize,
+        meta,
       }
 
       rpc.updateChannels((channels) => {
@@ -66,9 +74,9 @@ export const createWsRpcPreset: RpcServerPreset<
           if (index >= 0)
             channels.splice(index, 1)
         })
-        onDisconnected(ws)
+        onDisconnected(ws, meta)
       })
-      onConnected(ws)
+      onConnected(ws, meta)
     })
   }
 })