fix(core): close ws transport on vite server shutdown (#519)

Author: webfansplz

packages/core/src/node/__tests__/plugins-server-lifecycle.test.ts

@@ -0,0 +1,46 @@
+import type { ViteDevServer } from 'vite'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { createDevToolsContext } from '../context'
+import { DevToolsServer } from '../plugins/server'
+import { createDevToolsHub } from '../server'
+
+vi.mock('../context', () => ({
+  createDevToolsContext: vi.fn(),
+}))
+
+vi.mock('../server', () => ({
+  createDevToolsHub: vi.fn(),
+}))
+
+describe('devToolsServer lifecycle', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('closes the DevTools middleware with the Vite server', async () => {
+    const close = vi.fn(async () => {})
+    const middleware = vi.fn()
+    vi.mocked(createDevToolsContext).mockResolvedValue({} as Awaited<ReturnType<typeof createDevToolsContext>>)
+    vi.mocked(createDevToolsHub).mockResolvedValue({ close, middleware } as unknown as Awaited<ReturnType<typeof createDevToolsHub>>)
+
+    const viteDevServer = {
+      config: {
+        root: process.cwd(),
+        server: {},
+      },
+      middlewares: {
+        use: vi.fn(),
+      },
+    } as unknown as ViteDevServer
+    const plugin = DevToolsServer()
+
+    const { closeBundle, configureServer } = plugin
+    if (typeof configureServer !== 'function' || typeof closeBundle !== 'function')
+      throw new TypeError('Expected DevTools server lifecycle hooks')
+
+    await configureServer.call({} as never, viteDevServer)
+    await closeBundle.call({} as never)
+
+    expect(close).toHaveBeenCalledOnce()
+  })
+})

packages/core/src/node/plugins/server.ts

@@ -38,6 +38,7 @@ export function renderDockImportsMap(docks: Iterable<DevToolsDockEntry>): string
 
 export function DevToolsServer(): Plugin {
   let context: ViteDevToolsNodeContext
+  let close: (() => Promise<void>) | undefined
   return {
     name: 'vite:devtools:server',
     enforce: 'post',
@@ -49,7 +50,7 @@ export function DevToolsServer(): Plugin {
         ? '0.0.0.0'
         : viteDevServer.config.server.host || 'localhost'
 
-      const { middleware } = await createDevToolsHub({
+      const devtools = await createDevToolsHub({
         context,
         // Share Vite's HTTP server for a route-bound WS upgrade; fall back to a
         // side-car when Vite runs in middleware mode without its own server.
@@ -58,6 +59,7 @@ export function DevToolsServer(): Plugin {
         server: (viteDevServer.httpServer ?? undefined) as NodeHttpServer | undefined,
         host,
       })
+      close = devtools.close
       viteDevServer.middlewares.use((req, res, next) => {
         if (req.url === DEVTOOLS_MOUNT_PATH_NO_TRAILING_SLASH || req.url?.startsWith(`${DEVTOOLS_MOUNT_PATH_NO_TRAILING_SLASH}?`)) {
           res.statusCode = 302
@@ -70,7 +72,10 @@ export function DevToolsServer(): Plugin {
       })
       // The hub middleware answers the whole `/__devtools/` surface and
       // `next()`s outside its base, so mount it at the server root.
-      viteDevServer.middlewares.use(middleware)
+      viteDevServer.middlewares.use(devtools.middleware)
+    },
+    async closeBundle() {
+      await close?.()
     },
     resolveId(id) {
       if (id === DEVTOOLS_DOCK_IMPORTS_VIRTUAL_ID) {