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

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 { createDevToolsMiddleware } from '../server'
+
+vi.mock('../context', () => ({
+  createDevToolsContext: vi.fn(),
+}))
+
+vi.mock('../server', () => ({
+  createDevToolsMiddleware: 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(createDevToolsMiddleware).mockResolvedValue({ close, middleware } as unknown as Awaited<ReturnType<typeof createDevToolsMiddleware>>)
+
+    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/__tests__/ws.test.ts

@@ -0,0 +1,41 @@
+import type { ResolvedConfig } from 'vite'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { createDevToolsContext } from '../context'
+import { createWsServer } from '../ws'
+
+const mocks = vi.hoisted(() => ({
+  close: vi.fn(async () => {}),
+}))
+
+vi.mock('devframe/rpc/transports/ws-server', () => ({
+  attachWsRpcTransport: vi.fn(() => ({
+    close: mocks.close,
+  })),
+}))
+
+describe('createWsServer', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('closes its transport once', async () => {
+    const context = await createDevToolsContext({
+      root: process.cwd(),
+      command: 'build',
+      plugins: [],
+      server: {},
+    } as unknown as ResolvedConfig)
+    const server = await createWsServer({
+      cwd: process.cwd(),
+      websocket: {
+        host: 'localhost',
+        port: 7812,
+      },
+      context,
+    })
+
+    await Promise.all([server.close(), server.close()])
+
+    expect(mocks.close).toHaveBeenCalledOnce()
+  })
+})

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

@@ -37,6 +37,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',
@@ -48,13 +49,14 @@ export function DevToolsServer(): Plugin {
         ? '0.0.0.0'
         : viteDevServer.config.server.host || 'localhost'
 
-      const { middleware } = await createDevToolsMiddleware({
+      const devtools = await createDevToolsMiddleware({
         cwd: viteDevServer.config.root,
         websocket: {
           host,
         },
         context,
       })
+      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
@@ -65,7 +67,10 @@ export function DevToolsServer(): Plugin {
 
         next()
       })
-      viteDevServer.middlewares.use(DEVTOOLS_MOUNT_PATH, middleware)
+      viteDevServer.middlewares.use(DEVTOOLS_MOUNT_PATH, devtools.middleware)
+    },
+    async closeBundle() {
+      await close?.()
     },
     resolveId(id) {
       if (id === DEVTOOLS_DOCK_IMPORTS_VIRTUAL_ID) {

packages/core/src/node/server.ts

@@ -12,14 +12,15 @@ import { createWsServer } from './ws'
 export interface DevToolsMiddleware {
   h3: H3
   rpc: Awaited<ReturnType<typeof createWsServer>>['rpc']
+  close: Awaited<ReturnType<typeof createWsServer>>['close']
   middleware: NodeHandler
   getConnectionMeta: Awaited<ReturnType<typeof createWsServer>>['getConnectionMeta']
 }
 
 export async function createDevToolsMiddleware(options: CreateWsServerOptions): Promise<DevToolsMiddleware> {
   const h3 = new H3()
 
-  const { rpc, getConnectionMeta } = await createWsServer(options)
+  const { rpc, close, getConnectionMeta } = await createWsServer(options)
 
   // Hand the host the live connection-meta getter so each mounted devframe's
   // `mountConnectionMeta` middleware can serve it at the devframe's own base
@@ -58,6 +59,7 @@ export async function createDevToolsMiddleware(options: CreateWsServerOptions):
   return {
     h3,
     rpc,
+    close,
     middleware: toNodeHandler(h3),
     getConnectionMeta,
   }

packages/core/src/node/ws.ts

@@ -154,7 +154,7 @@ export async function createWsServer(options: CreateWsServerOptions) {
   // narrow cast rather than waiting on that.
   const allowedOrigins = (context.viteConfig.devtools?.config as DevToolsConfig | undefined)?.allowedOrigins
 
-  attachWsRpcTransport(rpcGroup, {
+  const transport = attachWsRpcTransport(rpcGroup, {
     ...binding,
     allowedOrigins,
     definitions: rpcHost.definitions,
@@ -199,6 +199,9 @@ export async function createWsServer(options: CreateWsServerOptions) {
   rpcHost._rpcGroup = rpcGroup
   rpcHost._asyncStorage = asyncStorage
 
+  let closePromise: Promise<void> | undefined
+  const close = () => closePromise ??= transport.close()
+
   const getConnectionMeta = async (): Promise<ConnectionMeta> => {
     const jsonSerializableMethods: string[] = []
     for (const def of rpcHost.definitions.values()) {
@@ -218,6 +221,7 @@ export async function createWsServer(options: CreateWsServerOptions) {
     port,
     rpc: rpcGroup,
     rpcHost,
+    close,
     getConnectionMeta,
   }
 }

test/__snapshots__/tsnapi/@vitejs/devtools/index.snapshot.d.ts

@@ -5,6 +5,7 @@
 export interface DevToolsMiddleware {
   h3: H3;
   rpc: Awaited<ReturnType<typeof createWsServer>>['rpc'];
+  close: Awaited<ReturnType<typeof createWsServer>>['close'];
   middleware: NodeHandler;
   getConnectionMeta: Awaited<ReturnType<typeof createWsServer>>['getConnectionMeta'];
 }