refactor: reuse Vite's websocket (#243)

package.json

@@ -40,6 +40,7 @@
     "simple-git-hooks": "^2.8.1",
     "typescript": "^5.0.4",
     "unocss": "^0.52.3",
+    "vite-hot-client": "^0.2.1",
     "vue-tsc": "^1.6.5"
   },
   "simple-git-hooks": {

packages/devtools/client/composables/npm.ts

@@ -13,19 +13,21 @@ export function usePackageUpdate(name: string, options?: NpmCommandOptions): Ret
 }
 
 export function useNuxtVersion() {
-  return useAsyncData('npm:check:nuxt', () => rpc.checkForUpdateFor('nuxt')).data
+  return useAsyncState('npm:check:nuxt', () => rpc.checkForUpdateFor('nuxt'))
 }
 
 export function satisfyNuxtVersion(range: string) {
   const nuxt = useNuxtVersion()
-  if (!nuxt?.value?.current)
-    return false
-  return semver.satisfies(nuxt.value.current, range)
+  return computed(() => {
+    if (!nuxt?.value?.current)
+      return false
+    return semver.satisfies(nuxt.value.current, range)
+  })
 }
 
 function getPackageUpdate(name: string, options?: NpmCommandOptions) {
   const nuxt = useNuxtApp()
-  const info = useAsyncData(`npm:check:${name}`, () => rpc.checkForUpdateFor(name)).data
+  const info = useAsyncState(`npm:check:${name}`, () => rpc.checkForUpdateFor(name))
 
   const state = ref<PackageUpdateState>('idle')
 

packages/devtools/client/composables/rpc.ts

@@ -1,14 +1,14 @@
 import { createBirpc } from 'birpc'
 import { parse, stringify } from 'flatted'
+import { createHotContext } from 'vite-hot-client'
 import type { ClientFunctions, ServerFunctions } from '../../src/types'
-
-const RECONNECT_INTERVAL = 2000
+import { WS_EVENT_NAME } from '../../src/constant'
 
 export const wsConnecting = ref(false)
 export const wsError = ref<any>()
 export const wsConnectingDebounced = useDebounce(wsConnecting, 2000)
 
-let connectPromise = connectWS()
+const connectPromise = connectVite()
 let onMessage: Function = () => {}
 
 export const clientFunctions = {
@@ -19,9 +19,11 @@ export const extendedRpcMap = new Map<string, any>()
 
 export const rpc = createBirpc<ServerFunctions>(clientFunctions, {
   post: async (d) => {
-    (await connectPromise).send(d)
+    (await connectPromise).send(WS_EVENT_NAME, d)
+  },
+  on: (fn) => {
+    onMessage = fn
   },
-  on: (fn) => { onMessage = fn },
   serialize: stringify,
   deserialize: parse,
   resolver(name, fn) {
@@ -38,33 +40,19 @@ export const rpc = createBirpc<ServerFunctions>(clientFunctions, {
   timeout: 120_000,
 })
 
-async function connectWS() {
-  const wsUrl = new URL('ws://host/__nuxt_devtools__/entry')
-  wsUrl.protocol = location.protocol === 'https:' ? 'wss:' : 'ws:'
-  wsUrl.host = location.host
+async function connectVite() {
+  const hot = await createHotContext()
 
-  const ws = new WebSocket(wsUrl.toString())
-  ws.addEventListener('message', e => onMessage(String(e.data)))
-  ws.addEventListener('error', (e) => {
-    console.error(e)
-    wsError.value = e
-  })
-  ws.addEventListener('close', () => {
-    // eslint-disable-next-line no-console
-    console.log('[nuxt-devtools] WebSocket closed, reconnecting...')
-    wsConnecting.value = true
-    setTimeout(async () => {
-      connectPromise = connectWS()
-    }, RECONNECT_INTERVAL)
+  if (!hot)
+    throw new Error('Unable to connect to devtools')
+
+  hot.on(WS_EVENT_NAME, (data) => {
+    onMessage(data)
   })
-  wsConnecting.value = true
-  if (ws.readyState !== WebSocket.OPEN)
-    await new Promise(resolve => ws.addEventListener('open', resolve))
 
-  // eslint-disable-next-line no-console
-  console.log('[nuxt-devtools] WebSocket connected.')
-  wsConnecting.value = false
-  wsError.value = null
+  // TODO:
+  // hot.on('vite:connect', (data) => {})
+  // hot.on('vite:disconnect', (data) => {})
 
-  return ws
+  return hot
 }

packages/devtools/client/composables/state-modules.ts

@@ -11,11 +11,11 @@ const ignoredModules = [
 ]
 
 export function useModulesList() {
-  return useAsyncData('modules-list', async () => {
+  return useAsyncState('getModulesList', async () => {
     const modules = await $fetch<ModuleStaticInfo[]>('https://cdn.jsdelivr.net/npm/@nuxt/modules@latest/modules.json')
     return modules
       .filter((m: ModuleStaticInfo) => !ignoredModules.includes(m.npm) && m.compatibility.nuxt.includes('^3'))
-  }).data
+  })
 }
 
 export function useInstalledModules() {

packages/devtools/client/pages/modules/analyze-build.vue

@@ -9,7 +9,7 @@ definePageMeta({
   layout: 'full',
   category: 'analyze',
   show() {
-    return () => satisfyNuxtVersion('^3.5.0')
+    return satisfyNuxtVersion('^3.5.0')
   },
 })
 

packages/devtools/client/pages/modules/storage.vue

@@ -13,10 +13,10 @@ const router = useRouter()
 const searchString = ref('')
 const newKey = ref('')
 const currentStorage = computed({
-  get(): string | undefined {
+  get(): string | number | undefined {
     return useRoute().query?.storage as string | undefined
   },
-  set(storage: string | undefined): void {
+  set(storage: string | number | undefined): void {
     router.replace({ query: { storage } })
   },
 })

packages/devtools/package.json

@@ -65,7 +65,6 @@
     "rc9": "^2.1.0",
     "semver": "^7.5.1",
     "sirv": "^2.0.3",
-    "tinyws": "^0.1.0",
     "unimport": "^3.0.7",
     "vite-plugin-inspect": "^0.7.28",
     "vite-plugin-vue-inspector": "^3.4.2",

packages/devtools/src/constant.ts

@@ -4,6 +4,7 @@ export const ROUTE_PATH = '/__nuxt_devtools__'
 export const ROUTE_ENTRY = `${ROUTE_PATH}/entry`
 export const ROUTE_CLIENT = `${ROUTE_PATH}/client`
 export const ROUTE_ANALYZE = `${ROUTE_PATH}/analyze`
+export const WS_EVENT_NAME = 'nuxt:devtools:rpc'
 
 export const defaultOptions: ModuleOptions = {
   enabled: undefined, // determine multiple conditions

packages/devtools/src/module-main.ts

@@ -1,8 +1,7 @@
 import { existsSync } from 'node:fs'
 import { join } from 'pathe'
 import type { Nuxt } from 'nuxt/schema'
-import { addPlugin, logger } from '@nuxt/kit'
-import { tinyws } from 'tinyws'
+import { addPlugin, addVitePlugin, logger } from '@nuxt/kit'
 import type { ViteDevServer } from 'vite'
 import { searchForWorkspaceRoot } from 'vite'
 import sirv from 'sirv'
@@ -11,7 +10,7 @@ import { version } from '../package.json'
 import type { ModuleOptions } from './types'
 import { setupRPC } from './server-rpc'
 import { clientDir, isGlobalInstall, packageDir, runtimeDir } from './dirs'
-import { ROUTE_ANALYZE, ROUTE_CLIENT, ROUTE_ENTRY } from './constant'
+import { ROUTE_ANALYZE, ROUTE_CLIENT } from './constant'
 
 export async function enableModule(options: ModuleOptions, nuxt: Nuxt) {
   // Disable in test mode
@@ -39,10 +38,12 @@ export async function enableModule(options: ModuleOptions, nuxt: Nuxt) {
   })
 
   const {
-    middleware: rpcMiddleware,
+    vitePlugin,
     ...ctx
   } = setupRPC(nuxt, options)
 
+  addVitePlugin(vitePlugin)
+
   const clientDirExists = existsSync(clientDir)
   const analyzeDir = join(nuxt.options.rootDir, '.nuxt/analyze')
 
@@ -63,8 +64,6 @@ export async function enableModule(options: ModuleOptions, nuxt: Nuxt) {
 
   // TODO: Use WS from nitro server when possible
   nuxt.hook('vite:serverCreated', (server: ViteDevServer) => {
-    server.middlewares.use(ROUTE_ENTRY, tinyws() as any)
-    server.middlewares.use(ROUTE_ENTRY, rpcMiddleware as any)
     server.middlewares.use(ROUTE_ANALYZE, sirv(analyzeDir, { single: false, dev: true }))
     // serve the front end in production
     if (clientDirExists)

packages/devtools/src/server-rpc/index.ts

@@ -1,12 +1,12 @@
-import type { TinyWSRequest } from 'tinyws'
-import type { NodeIncomingMessage, NodeServerResponse } from 'h3'
 import type { WebSocket } from 'ws'
 import { createBirpcGroup } from 'birpc'
 import type { ChannelOptions } from 'birpc'
 
 import { parse, stringify } from 'flatted'
 import type { Nuxt } from 'nuxt/schema'
+import type { Plugin } from 'vite'
 import type { ClientFunctions, ModuleOptions, NuxtDevtoolsServerContext, ServerFunctions } from '../types'
+import { WS_EVENT_NAME } from '../constant'
 import { setupStorageRPC } from './storage'
 import { setupAssetsRPC } from './assets'
 import { setupNpmRPC } from './npm'
@@ -86,36 +86,50 @@ export function setupRPC(nuxt: Nuxt, options: ModuleOptions) {
   } satisfies ServerFunctions)
 
   const wsClients = new Set<WebSocket>()
-  const middleware = async (req: NodeIncomingMessage & TinyWSRequest, _res: NodeServerResponse, next: Function) => {
-    // Handle WebSocket
-    if (req.ws) {
-      const ws = await req.ws()
-      wsClients.add(ws)
-      const channel: ChannelOptions = {
-        post: d => ws.send(d),
-        on: fn => ws.on('message', fn),
-        serialize: stringify,
-        deserialize: parse,
-      }
-      rpc.updateChannels((c) => {
-        c.push(channel)
-      })
-      ws.on('close', () => {
-        wsClients.delete(ws)
+
+  const vitePlugin: Plugin = {
+    name: 'nuxt:devtools:rpc',
+    configureServer(server) {
+      server.ws.on('connection', (ws) => {
+        wsClients.add(ws)
+        const channel: ChannelOptions = {
+          post: d => ws.send(JSON.stringify({
+            type: 'custom',
+            event: WS_EVENT_NAME,
+            data: d,
+          })),
+          on: (fn) => {
+            ws.on('message', (e) => {
+              try {
+                const data = JSON.parse(String(e)) || {}
+                if (data.type === 'custom' && data.event === WS_EVENT_NAME) {
+                  // console.log(data.data)
+                  fn(data.data)
+                }
+              }
+              catch {}
+            })
+          },
+          serialize: stringify,
+          deserialize: parse,
+        }
         rpc.updateChannels((c) => {
-          const index = c.indexOf(channel)
-          if (index >= 0)
-            c.splice(index, 1)
+          c.push(channel)
+        })
+        ws.on('close', () => {
+          wsClients.delete(ws)
+          rpc.updateChannels((c) => {
+            const index = c.indexOf(channel)
+            if (index >= 0)
+              c.splice(index, 1)
+          })
         })
       })
-    }
-    else {
-      next()
-    }
+    },
   }
 
   return {
-    middleware,
+    vitePlugin,
     ...ctx,
   }
 }