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

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());
+  const cleanupUnmount = () => instance.unmount();
+  registerCleanup(cleanupUnmount);
 
-  registerCleanup(setupTtyCheck());
+  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().