fix(dts-plugin): handle EXIT before dev server init (#4429)

Co-authored-by: Zack Jackson <25274700+ScriptedAlchemy@users.noreply.github.com>
Co-authored-by: 晓 <zhouxiao.shaw@bytedance.com>

Author: 3 people

packages/dts-plugin/src/dev-worker/forkDevWorker.ts

@@ -23,6 +23,7 @@ import {
 
 import { DevWorkerOptions } from './DevWorker';
 import { getIpFromEntry } from './utils';
+import { handleDevWorkerMessage } from './handleWorkerMessage';
 
 interface Options extends DevWorkerOptions {
   name: string;
@@ -189,20 +190,7 @@ export async function forkDevWorker(
 }
 
 process.on('message', (message: rpc.RpcMessage) => {
-  fileLog(
-    `ChildProcess(${process.pid}), message: ${JSON.stringify(message)} `,
-    'forkDevWorker',
-    'info',
-  );
-  if (message.type === rpc.RpcGMCallTypes.EXIT) {
-    fileLog(
-      `ChildProcess(${process.pid}) SIGTERM, Federation DevServer will exit...`,
-      'forkDevWorker',
-      'error',
-    );
-    moduleServer?.exit();
-    process.exit(0);
-  }
+  handleDevWorkerMessage(message, { moduleServer, log: fileLog });
 });
 
 rpc.exposeRpc(forkDevWorker);

packages/dts-plugin/src/dev-worker/handleWorkerMessage.spec.ts

@@ -0,0 +1,63 @@
+import { describe, expect, it, vi, afterEach } from 'vitest';
+
+import { RpcGMCallTypes } from '../core/rpc/types';
+import { handleDevWorkerMessage } from './handleWorkerMessage';
+
+describe('handleDevWorkerMessage', () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('does not throw when EXIT arrives before module server is initialized', () => {
+    const processExit = vi.fn();
+
+    expect(() =>
+      handleDevWorkerMessage(
+        {
+          type: RpcGMCallTypes.EXIT,
+          id: 'exit-before-init',
+        },
+        { processExit, log: vi.fn() },
+      ),
+    ).not.toThrow();
+
+    expect(processExit).toHaveBeenCalledWith(0);
+  });
+
+  it('calls module server exit when present', () => {
+    const processExit = vi.fn();
+    const moduleServer = {
+      exit: vi.fn(),
+    };
+
+    handleDevWorkerMessage(
+      {
+        type: RpcGMCallTypes.EXIT,
+        id: 'exit-with-server',
+      },
+      { moduleServer, processExit, log: vi.fn() },
+    );
+
+    expect(moduleServer.exit).toHaveBeenCalledTimes(1);
+    expect(processExit).toHaveBeenCalledWith(0);
+  });
+
+  it('ignores non-exit messages', () => {
+    const processExit = vi.fn();
+    const moduleServer = {
+      exit: vi.fn(),
+    };
+
+    handleDevWorkerMessage(
+      {
+        type: RpcGMCallTypes.CALL,
+        id: 'message',
+        args: [],
+      },
+      { moduleServer, processExit, log: vi.fn() },
+    );
+
+    expect(moduleServer.exit).not.toHaveBeenCalled();
+    expect(processExit).not.toHaveBeenCalled();
+  });
+});

packages/dts-plugin/src/dev-worker/handleWorkerMessage.ts

@@ -0,0 +1,38 @@
+import { RpcGMCallTypes, type RpcMessage } from '../core/rpc/types';
+
+interface HandleDevWorkerMessageOptions {
+  moduleServer?: {
+    exit: () => void;
+  };
+  processExit?: (code?: number) => void;
+  pid?: number;
+  log?: (message: string, scope: string, level: string) => void;
+}
+
+export function handleDevWorkerMessage(
+  message: RpcMessage,
+  options: HandleDevWorkerMessageOptions = {},
+): void {
+  const {
+    moduleServer,
+    processExit = process.exit,
+    pid = process.pid,
+    log = () => undefined,
+  } = options;
+
+  log(
+    `ChildProcess(${pid}), message: ${JSON.stringify(message)} `,
+    'forkDevWorker',
+    'info',
+  );
+
+  if (message.type === RpcGMCallTypes.EXIT) {
+    log(
+      `ChildProcess(${pid}) SIGTERM, Federation DevServer will exit...`,
+      'forkDevWorker',
+      'error',
+    );
+    moduleServer?.exit();
+    processExit(0);
+  }
+}