diff --git a/packages/core/src/node/__tests__/plugins-server-lifecycle.test.ts b/packages/core/src/node/__tests__/plugins-server-lifecycle.test.ts new file mode 100644 index 000000000..49a61fb46 --- /dev/null +++ b/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>) + vi.mocked(createDevToolsMiddleware).mockResolvedValue({ close, middleware } as unknown as Awaited>) + + 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() + }) +}) diff --git a/packages/core/src/node/__tests__/ws.test.ts b/packages/core/src/node/__tests__/ws.test.ts new file mode 100644 index 000000000..0c93a07a1 --- /dev/null +++ b/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() + }) +}) diff --git a/packages/core/src/node/plugins/server.ts b/packages/core/src/node/plugins/server.ts index c287e9810..ba8d9150d 100644 --- a/packages/core/src/node/plugins/server.ts +++ b/packages/core/src/node/plugins/server.ts @@ -37,6 +37,7 @@ export function renderDockImportsMap(docks: Iterable): string export function DevToolsServer(): Plugin { let context: ViteDevToolsNodeContext + let close: (() => Promise) | 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) { diff --git a/packages/core/src/node/server.ts b/packages/core/src/node/server.ts index 71701d299..7adf995b1 100644 --- a/packages/core/src/node/server.ts +++ b/packages/core/src/node/server.ts @@ -12,6 +12,7 @@ import { createWsServer } from './ws' export interface DevToolsMiddleware { h3: H3 rpc: Awaited>['rpc'] + close: Awaited>['close'] middleware: NodeHandler getConnectionMeta: Awaited>['getConnectionMeta'] } @@ -19,7 +20,7 @@ export interface DevToolsMiddleware { export async function createDevToolsMiddleware(options: CreateWsServerOptions): Promise { 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, } diff --git a/packages/core/src/node/ws.ts b/packages/core/src/node/ws.ts index da5ffaf4d..6fd107776 100644 --- a/packages/core/src/node/ws.ts +++ b/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 | undefined + const close = () => closePromise ??= transport.close() + const getConnectionMeta = async (): Promise => { 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, } } diff --git a/test/__snapshots__/tsnapi/@vitejs/devtools/index.snapshot.d.ts b/test/__snapshots__/tsnapi/@vitejs/devtools/index.snapshot.d.ts index 568aac89b..263fabad3 100644 --- a/test/__snapshots__/tsnapi/@vitejs/devtools/index.snapshot.d.ts +++ b/test/__snapshots__/tsnapi/@vitejs/devtools/index.snapshot.d.ts @@ -5,6 +5,7 @@ export interface DevToolsMiddleware { h3: H3; rpc: Awaited>['rpc']; + close: Awaited>['close']; middleware: NodeHandler; getConnectionMeta: Awaited>['getConnectionMeta']; }