fix: resolve lifecycle memory leaks by cleaning up listeners and root closures (google-gemini#25049)

Author: spencer426

packages/cli/src/gemini.test.tsx

@@ -1428,12 +1428,13 @@ describe('startInteractiveUI', () => {
   vi.mock('./ui/utils/updateCheck.js', () => ({
     checkForUpdates: vi.fn(() => Promise.resolve(null)),
   }));
-
   vi.mock('./utils/cleanup.js', () => ({
     cleanupCheckpoints: vi.fn(() => Promise.resolve()),
     registerCleanup: vi.fn(),
+    removeCleanup: vi.fn(),
     runExitCleanup: vi.fn(),
     registerSyncCleanup: vi.fn(),
+    removeSyncCleanup: vi.fn(),
     registerTelemetryConfig: vi.fn(),
     setupSignalHandlers: vi.fn(),
     setupTtyCheck: vi.fn(() => vi.fn()),

packages/cli/src/gemini_cleanup.test.tsx

@@ -142,7 +142,9 @@ vi.mock('./utils/cleanup.js', async (importOriginal) => {
     ...actual,
     cleanupCheckpoints: vi.fn().mockResolvedValue(undefined),
     registerCleanup: vi.fn(),
+    removeCleanup: vi.fn(),
     registerSyncCleanup: vi.fn(),
+    removeSyncCleanup: vi.fn(),
     registerTelemetryConfig: vi.fn(),
     runExitCleanup: vi.fn().mockResolvedValue(undefined),
   };

packages/cli/src/interactiveCli.tsx

@@ -9,7 +9,11 @@ import { render } from 'ink';
 import { basename } from 'node:path';
 import { AppContainer } from './ui/AppContainer.js';
 import { ConsolePatcher } from './ui/utils/ConsolePatcher.js';
-import { registerCleanup, setupTtyCheck } from './utils/cleanup.js';
+import {
+  registerCleanup,
+  removeCleanup,
+  setupTtyCheck,
+} from './utils/cleanup.js';
 import {
   type StartupWarning,
   type Config,
@@ -89,7 +93,6 @@ export async function startInteractiveUI(
     debugMode: config.getDebugMode(),
   });
   consolePatcher.patch();
-  registerCleanup(consolePatcher.cleanup);
 
   const { stdout: inkStdout, stderr: inkStderr } = createWorkingStdio();
 
@@ -167,11 +170,11 @@ export async function startInteractiveUI(
     },
   );
 
+  let cleanupLineWrapping: (() => void) | undefined;
   if (useAlternateBuffer) {
     disableLineWrapping();
-    registerCleanup(() => {
-      enableLineWrapping();
-    });
+    cleanupLineWrapping = () => enableLineWrapping();
+    registerCleanup(cleanupLineWrapping);
   }
 
   checkForUpdates(settings)
@@ -185,9 +188,48 @@ export async function startInteractiveUI(
       }
     });
 
-  registerCleanup(() => instance.unmount());
-
-  registerCleanup(setupTtyCheck());
+  const cleanupUnmount = () => instance.unmount();
+  registerCleanup(cleanupUnmount);
+
+  const cleanupTtyCheck = setupTtyCheck();
+  registerCleanup(cleanupTtyCheck);
+
+  const cleanupConsolePatcher = () => consolePatcher.cleanup();
+  registerCleanup(cleanupConsolePatcher);
+
+  try {
+    await instance.waitUntilExit();
+  } finally {
+    try {
+      removeCleanup(cleanupConsolePatcher);
+      cleanupConsolePatcher();
+    } catch (e: unknown) {
+      debugLogger.error('Error cleaning up console patcher:', e);
+    }
+
+    try {
+      removeCleanup(cleanupUnmount);
+      instance.unmount();
+    } catch (e: unknown) {
+      debugLogger.error('Error unmounting Ink instance:', e);
+    }
+
+    try {
+      removeCleanup(cleanupTtyCheck);
+      cleanupTtyCheck();
+    } catch (e: unknown) {
+      debugLogger.error('Error in TTY cleanup:', e);
+    }
+
+    if (cleanupLineWrapping) {
+      try {
+        removeCleanup(cleanupLineWrapping);
+        cleanupLineWrapping();
+      } catch (e: unknown) {
+        debugLogger.error('Error restoring line wrapping:', e);
+      }
+    }
+  }
 }
 
 function setWindowTitle(title: string, settings: LoadedSettings) {

packages/cli/src/ui/AppContainer.tsx

@@ -136,7 +136,11 @@ import { type IdeIntegrationNudgeResult } from './IdeIntegrationNudge.js';
 import { appEvents, AppEvent, TransientMessageType } from '../utils/events.js';
 import { type UpdateObject } from './utils/updateCheck.js';
 import { setUpdateHandler } from '../utils/handleAutoUpdate.js';
-import { registerCleanup, runExitCleanup } from '../utils/cleanup.js';
+import {
+  registerCleanup,
+  removeCleanup,
+  runExitCleanup,
+} from '../utils/cleanup.js';
 import { relaunchApp } from '../utils/processUtils.js';
 import type { SessionInfo } from '../utils/sessionUtils.js';
 import { useMessageQueue } from './hooks/useMessageQueue.js';
@@ -519,7 +523,7 @@ export const AppContainer = (props: AppContainerProps) => {
         debugLogger.warn('Background summary generation failed:', e);
       });
     })();
-    registerCleanup(async () => {
+    const cleanupFn = async () => {
       // Turn off mouse scroll.
       disableMouseEvents();
 
@@ -535,7 +539,15 @@ export const AppContainer = (props: AppContainerProps) => {
 
       // Fire SessionEnd hook on cleanup (only if hooks are enabled)
       await config?.getHookSystem()?.fireSessionEndEvent(SessionEndReason.Exit);
-    });
+    };
+    registerCleanup(cleanupFn);
+
+    return () => {
+      removeCleanup(cleanupFn);
+      cleanupFn().catch((e: unknown) =>
+        debugLogger.error('Error during cleanup:', e),
+      );
+    };
     // Disable the dependencies check here. historyManager gets flagged
     // but we don't want to react to changes to it because each new history
     // item, including the ones from the start session hook will cause a

packages/cli/src/utils/cleanup.ts

@@ -24,10 +24,24 @@ export function registerCleanup(fn: (() => void) | (() => Promise<void>)) {
   cleanupFunctions.push(fn);
 }
 
+export function removeCleanup(fn: (() => void) | (() => Promise<void>)) {
+  const index = cleanupFunctions.indexOf(fn);
+  if (index !== -1) {
+    cleanupFunctions.splice(index, 1);
+  }
+}
+
 export function registerSyncCleanup(fn: () => void) {
   syncCleanupFunctions.push(fn);
 }
 
+export function removeSyncCleanup(fn: () => void) {
+  const index = syncCleanupFunctions.indexOf(fn);
+  if (index !== -1) {
+    syncCleanupFunctions.splice(index, 1);
+  }
+}
+
 /**
  * Resets the internal cleanup state for testing purposes.
  * This allows tests to run in isolation without vi.resetModules().

packages/core/src/services/shellExecutionService.test.ts

@@ -208,6 +208,7 @@ describe('ShellExecutionService', () => {
   beforeEach(() => {
     vi.clearAllMocks();
     ExecutionLifecycleService.resetForTest();
+    ShellExecutionService.resetForTest();
     mockSerializeTerminalToObject.mockReturnValue([]);
     mockIsBinary.mockReturnValue(false);
     mockPlatform.mockReturnValue('linux');
@@ -1247,6 +1248,8 @@ describe('ShellExecutionService child_process fallback', () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
+    ExecutionLifecycleService.resetForTest();
+    ShellExecutionService.resetForTest();
 
     mockIsBinary.mockReturnValue(false);
     mockPlatform.mockReturnValue('linux');
@@ -1662,6 +1665,8 @@ describe('ShellExecutionService execution method selection', () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
+    ExecutionLifecycleService.resetForTest();
+    ShellExecutionService.resetForTest();
     onOutputEventMock = vi.fn();
 
     // Mock for pty
@@ -1786,6 +1791,8 @@ describe('ShellExecutionService environment variables', () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
+    ExecutionLifecycleService.resetForTest();
+    ShellExecutionService.resetForTest();
     vi.resetModules(); // Reset modules to ensure process.env changes are fresh
 
     // Mock for pty

packages/core/src/services/shellExecutionService.ts

@@ -704,7 +704,10 @@ export class ShellExecutionService {
 
         const finalStrippedOutput = stripAnsi(combinedOutput).trim();
         const exitCode = code;
-        const exitSignal = signal ? os.constants.signals[signal] : null;
+        const exitSignal =
+          signal && os.constants.signals
+            ? (os.constants.signals[signal] ?? null)
+            : null;
 
         const resultPayload: ShellExecutionResult = {
           rawOutput: Buffer.from(''),
@@ -1503,4 +1506,16 @@ export class ShellExecutionService {
       signal: info.signal,
     }));
   }
+
+  /**
+   * Resets the internal state of the ShellExecutionService.
+   * This is intended for use in tests to ensure isolation.
+   */
+  static resetForTest(): void {
+    this.activePtys.clear();
+    this.activeChildProcesses.clear();
+    this.backgroundLogPids.clear();
+    this.backgroundLogStreams.clear();
+    this.backgroundProcessHistory.clear();
+  }
 }

packages/core/src/tools/mcp-client.ts

@@ -812,15 +812,15 @@ type StatusChangeListener = (
   serverName: string,
   status: MCPServerStatus,
 ) => void;
-const statusChangeListeners: StatusChangeListener[] = [];
+const statusChangeListeners: Set<StatusChangeListener> = new Set();
 
 /**
  * Add a listener for MCP server status changes
  */
 export function addMCPStatusChangeListener(
   listener: StatusChangeListener,
 ): void {
-  statusChangeListeners.push(listener);
+  statusChangeListeners.add(listener);
 }
 
 /**
@@ -829,10 +829,7 @@ export function addMCPStatusChangeListener(
 export function removeMCPStatusChangeListener(
   listener: StatusChangeListener,
 ): void {
-  const index = statusChangeListeners.indexOf(listener);
-  if (index !== -1) {
-    statusChangeListeners.splice(index, 1);
-  }
+  statusChangeListeners.delete(listener);
 }
 
 /**