fix(tui): guarantee exit after drained teardown (#111015)

src/tui/tui.test.ts

@@ -7,8 +7,10 @@ import { MALFORMED_STREAMING_FRAGMENT_ERROR_MESSAGE } from "../shared/assistant-
 import { withEnv } from "../test-utils/env.js";
 import { getSlashCommands, parseCommand } from "./commands.js";
 import {
+  beginTuiShutdown,
   createBackspaceDeduper,
   createDeferredTuiFinish,
+  createTuiSignalHandlers,
   drainAndStopTuiSafely,
   installTuiTerminalLossExitHandler,
   isIgnorableTuiStopError,
@@ -426,8 +428,22 @@ describe("resolveTuiCtrlCAction", () => {
 });
 
 describe("TUI shutdown safety", () => {
+  const beginTestShutdown = (overrides: Partial<Parameters<typeof beginTuiShutdown>[0]> = {}) =>
+    beginTuiShutdown({
+      stopClient: vi.fn(),
+      stopTui: vi.fn(),
+      stopStatusTimeout: vi.fn(),
+      requestFinish: vi.fn(),
+      forceExit: vi.fn(),
+      hardExitMs: 2000,
+      keepHardExitArmed: true,
+      onError: vi.fn(),
+      ...overrides,
+    });
+
   afterEach(() => {
     vi.useRealTimers();
+    vi.restoreAllMocks();
   });
 
   it("drains terminal input before stopping the TUI", async () => {
@@ -547,6 +563,98 @@ describe("TUI shutdown safety", () => {
     expect(finish).toHaveBeenCalledTimes(1);
   });
 
+  it("forces process exit when gateway teardown never settles", async () => {
+    vi.useFakeTimers();
+    const exit = vi.spyOn(process, "exit").mockImplementation((() => undefined) as never);
+    const requestFinish = vi.fn();
+    const timer = beginTestShutdown({
+      stopClient: () => new Promise<void>(() => {}),
+      requestFinish,
+      forceExit: () => process.exit(130),
+    });
+
+    expect((timer as NodeJS.Timeout).hasRef()).toBe(false);
+    await vi.advanceTimersByTimeAsync(1999);
+    expect(exit).not.toHaveBeenCalled();
+    await vi.advanceTimersByTimeAsync(1);
+    expect(exit).toHaveBeenCalledWith(130);
+    expect(requestFinish).not.toHaveBeenCalled();
+  });
+
+  it("forces process exit after SIGTERM when gateway teardown never settles", async () => {
+    vi.useFakeTimers();
+    const exit = vi.spyOn(process, "exit").mockImplementation((() => undefined) as never);
+    const requestExit = vi.fn(() => {
+      beginTestShutdown({
+        stopClient: () => new Promise<void>(() => {}),
+        forceExit: () => process.exit(130),
+      });
+    });
+    const { sigtermHandler } = createTuiSignalHandlers({
+      handleCtrlC: vi.fn(),
+      requestExit,
+    });
+
+    sigtermHandler();
+    expect(requestExit).toHaveBeenCalledOnce();
+    await vi.advanceTimersByTimeAsync(2000);
+    expect(exit).toHaveBeenCalledWith(130);
+  });
+
+  it("keeps the force-exit deadline armed after already-drained teardown settles", async () => {
+    vi.useFakeTimers();
+    const forceExit = vi.fn();
+    const requestFinish = vi.fn();
+    beginTestShutdown({
+      requestFinish,
+      forceExit,
+    });
+
+    await vi.advanceTimersByTimeAsync(0);
+    expect(requestFinish).toHaveBeenCalledOnce();
+    expect(forceExit).not.toHaveBeenCalled();
+    await vi.advanceTimersByTimeAsync(2000);
+    expect(forceExit).toHaveBeenCalledOnce();
+  });
+
+  it("completes healthy shutdown promptly without waiting for the force-exit deadline", async () => {
+    vi.useFakeTimers();
+    const calls: string[] = [];
+    const forceExit = vi.fn();
+    beginTestShutdown({
+      stopClient: async () => {
+        calls.push("client");
+      },
+      stopTui: async () => {
+        calls.push("tui");
+      },
+      stopStatusTimeout: () => {
+        calls.push("status");
+      },
+      requestFinish: () => {
+        calls.push("finish");
+      },
+      forceExit,
+    });
+
+    await vi.advanceTimersByTimeAsync(0);
+    expect(calls).toEqual(["client", "tui", "status", "finish"]);
+    expect(forceExit).not.toHaveBeenCalled();
+  });
+
+  it("cancels the hard-exit deadline for embedded TUI callers after clean shutdown", async () => {
+    vi.useFakeTimers();
+    const forceExit = vi.fn();
+    beginTestShutdown({
+      forceExit,
+      keepHardExitArmed: false,
+      onError: vi.fn(),
+    });
+
+    await vi.advanceTimersByTimeAsync(2000);
+    expect(forceExit).not.toHaveBeenCalled();
+  });
+
   it("does not keep a clean standalone TUI alive for the watchdog deadline", () => {
     const unref = vi.fn();
     const setTimeoutFn = vi.fn((_callback: () => void, delayMs: number) => {

src/tui/tui.ts

@@ -418,6 +418,60 @@ type TuiProcessExitTimer = {
 
 type TuiProcessExitTimeout = (callback: () => void, delayMs: number) => TuiProcessExitTimer;
 
+type TuiShutdownTask = () => void | Promise<void>;
+
+export function beginTuiShutdown(params: {
+  stopClient: TuiShutdownTask;
+  stopTui: TuiShutdownTask;
+  stopStatusTimeout: () => void;
+  requestFinish: () => void;
+  forceExit: () => void;
+  hardExitMs: number;
+  keepHardExitArmed?: boolean;
+  onError: (error: unknown) => void;
+  clearTimeoutFn?: (timer: TuiProcessExitTimer) => void;
+  setTimeoutFn?: TuiProcessExitTimeout;
+}): TuiProcessExitTimer {
+  const setTimeoutFn =
+    params.setTimeoutFn ??
+    ((callback, timeoutMs) => setTimeout(callback, timeoutMs) as unknown as TuiProcessExitTimer);
+  const hardExitTimer = setTimeoutFn(params.forceExit, params.hardExitMs);
+  hardExitTimer.unref?.();
+  void Promise.resolve()
+    .then(params.stopClient)
+    .then(params.stopTui)
+    .finally(() => {
+      if (params.keepHardExitArmed !== true) {
+        const clearTimeoutFn =
+          params.clearTimeoutFn ??
+          ((timer) => clearTimeout(timer as unknown as ReturnType<typeof setTimeout>));
+        clearTimeoutFn(hardExitTimer);
+      }
+      params.stopStatusTimeout();
+    })
+    .catch(params.onError)
+    .finally(params.requestFinish);
+
+  // For the standalone command, settled teardown is not proof that runTui
+  // returned. Its unref keeps clean exits fast while preserving the deadline.
+  return hardExitTimer;
+}
+
+export function createTuiSignalHandlers(params: {
+  handleCtrlC: () => void;
+  requestExit: () => void;
+}): {
+  sigintHandler: () => void;
+  sigtermHandler: () => void;
+  sighupHandler: () => void;
+} {
+  return {
+    sigintHandler: params.handleCtrlC,
+    sigtermHandler: params.requestExit,
+    sighupHandler: params.requestExit,
+  };
+}
+
 export async function drainAndStopTuiSafely(tui: DrainableTui): Promise<void> {
   if (typeof tui.terminal?.drainInput === "function") {
     try {
@@ -1375,19 +1429,15 @@ export async function runTui(opts: RunTuiOptions): Promise<TuiResult> {
     };
     pluginApprovals?.dispose();
     taskSuggestions?.dispose();
-    const hardExitTimer = setTimeout(
+    beginTuiShutdown({
+      stopClient: () => client.stop(),
+      stopTui: () => drainAndStopTuiSafely(tui),
+      stopStatusTimeout,
+      requestFinish: deferredFinish.requestFinish,
       forceExit,
-      resolveTuiShutdownHardExitMs({ localMode: isLocalMode }),
-    );
-    hardExitTimer.unref?.();
-    void Promise.resolve()
-      .then(() => client.stop())
-      .then(() => drainAndStopTuiSafely(tui))
-      .finally(() => {
-        clearTimeout(hardExitTimer);
-        stopStatusTimeout();
-      })
-      .catch((err: unknown) => {
+      hardExitMs: resolveTuiShutdownHardExitMs({ localMode: isLocalMode }),
+      keepHardExitArmed: opts.forceProcessExitOnReturn === true,
+      onError: (err) => {
         if (!isTuiTerminalLossError(err)) {
           try {
             process.stderr.write(`openclaw tui shutdown failed: ${String(err)}\n`);
@@ -1395,10 +1445,8 @@ export async function runTui(opts: RunTuiOptions): Promise<TuiResult> {
             // Best effort only; exit must still complete.
           }
         }
-      })
-      .finally(() => {
-        deferredFinish.requestFinish();
-      });
+      },
+    });
   };
   const exitAwareClient = client as TuiBackend & {
     setRequestExitHandler?: (handler: () => void) => void;
@@ -1690,15 +1738,10 @@ export async function runTui(opts: RunTuiOptions): Promise<TuiResult> {
   updateHeader();
   setConnectionStatus(isLocalMode ? "starting local runtime" : "connecting");
   updateFooter();
-  const sigintHandler = () => {
-    handleCtrlC();
-  };
-  const sigtermHandler = () => {
-    requestExit();
-  };
-  const sighupHandler = () => {
-    requestExit();
-  };
+  const { sigintHandler, sigtermHandler, sighupHandler } = createTuiSignalHandlers({
+    handleCtrlC,
+    requestExit,
+  });
   process.on("SIGINT", sigintHandler);
   process.on("SIGTERM", sigtermHandler);
   process.on("SIGHUP", sighupHandler);