From 5fc8fea8d762d67d3bff14d00307138faff0bca5 Mon Sep 17 00:00:00 2001 From: Spencer Date: Fri, 10 Apr 2026 00:21:14 -0400 Subject: [PATCH] fix: resolve lifecycle memory leaks by cleaning up listeners and root closures (#25049) --- packages/cli/src/gemini.test.tsx | 3 +- packages/cli/src/gemini_cleanup.test.tsx | 2 + packages/cli/src/interactiveCli.tsx | 56 ++++++++++++++++--- packages/cli/src/ui/AppContainer.tsx | 18 +++++- packages/cli/src/utils/cleanup.ts | 14 +++++ .../services/shellExecutionService.test.ts | 7 +++ .../src/services/shellExecutionService.ts | 17 +++++- packages/core/src/tools/mcp-client.ts | 9 +-- 8 files changed, 108 insertions(+), 18 deletions(-) diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 7712d39bb2..5b31d153fe 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/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()), diff --git a/packages/cli/src/gemini_cleanup.test.tsx b/packages/cli/src/gemini_cleanup.test.tsx index 0fc43ba2bf..2df1ab4d82 100644 --- a/packages/cli/src/gemini_cleanup.test.tsx +++ b/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), }; diff --git a/packages/cli/src/interactiveCli.tsx b/packages/cli/src/interactiveCli.tsx index 0d73f95016..39411c19dd 100644 --- a/packages/cli/src/interactiveCli.tsx +++ b/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) { diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 0bba196e1d..eaf6fc3e75 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/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 diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts index 0b7c75941a..2f18bdee30 100644 --- a/packages/cli/src/utils/cleanup.ts +++ b/packages/cli/src/utils/cleanup.ts @@ -24,10 +24,24 @@ export function registerCleanup(fn: (() => void) | (() => Promise)) { cleanupFunctions.push(fn); } +export function removeCleanup(fn: (() => void) | (() => Promise)) { + 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(). diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts index a7b21ebefc..749150361f 100644 --- a/packages/core/src/services/shellExecutionService.test.ts +++ b/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 diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index 46b894426f..607d84ef8d 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/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(); + } } diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 7e1ba49b89..a7852050fc 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -812,7 +812,7 @@ type StatusChangeListener = ( serverName: string, status: MCPServerStatus, ) => void; -const statusChangeListeners: StatusChangeListener[] = []; +const statusChangeListeners: Set = new Set(); /** * Add a listener for MCP server status changes @@ -820,7 +820,7 @@ const statusChangeListeners: StatusChangeListener[] = []; 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); } /**