From 3af4913ef3f00de71744de551a568aa713a3beec Mon Sep 17 00:00:00 2001 From: christine betts Date: Fri, 8 Aug 2025 15:38:30 +0000 Subject: [PATCH] [ide-mode] Close all open diffs when the CLI gets closed (#5792) --- packages/cli/src/ui/App.tsx | 3 + .../ui/hooks/slashCommandProcessor.test.ts | 39 +++++++++++ .../cli/src/ui/hooks/slashCommandProcessor.ts | 4 +- packages/cli/src/utils/cleanup.test.ts | 68 +++++++++++++++++++ packages/cli/src/utils/cleanup.ts | 8 +-- packages/core/src/config/config.ts | 2 +- packages/core/src/ide/ide-client.ts | 9 ++- 7 files changed, 126 insertions(+), 7 deletions(-) create mode 100644 packages/cli/src/utils/cleanup.test.ts diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index a25b7a56..58a40b93 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -122,6 +122,9 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { const [idePromptAnswered, setIdePromptAnswered] = useState(false); const currentIDE = config.getIdeClient().getCurrentIde(); + useEffect(() => { + registerCleanup(() => config.getIdeClient().disconnect()); + }, [config]); const shouldShowIdePrompt = config.getIdeModeFeature() && currentIDE && diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index a37af262..37407689 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -60,6 +60,14 @@ vi.mock('../contexts/SessionContext.js', () => ({ useSessionStats: vi.fn(() => ({ stats: {} })), })); +const { mockRunExitCleanup } = vi.hoisted(() => ({ + mockRunExitCleanup: vi.fn(), +})); + +vi.mock('../../utils/cleanup.js', () => ({ + runExitCleanup: mockRunExitCleanup, +})); + import { act, renderHook, waitFor } from '@testing-library/react'; import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest'; import { useSlashCommandProcessor } from './slashCommandProcessor.js'; @@ -405,6 +413,37 @@ describe('useSlashCommandProcessor', () => { vi.useRealTimers(); } }); + + it('should call runExitCleanup when handling a "quit" action', async () => { + const quitAction = vi + .fn() + .mockResolvedValue({ type: 'quit', messages: [] }); + const command = createTestCommand({ + name: 'exit', + action: quitAction, + }); + const result = setupProcessorHook([command]); + + await waitFor(() => + expect(result.current.slashCommands).toHaveLength(1), + ); + + vi.useFakeTimers(); + + try { + await act(async () => { + await result.current.handleSlashCommand('/exit'); + }); + + await act(async () => { + await vi.advanceTimersByTimeAsync(200); + }); + + expect(mockRunExitCleanup).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); }); it('should handle "submit_prompt" action returned from a file-based command', async () => { diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index cfe4b385..9f4bbf90 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -18,6 +18,7 @@ import { ToolConfirmationOutcome, } from '@google/gemini-cli-core'; import { useSessionStats } from '../contexts/SessionContext.js'; +import { runExitCleanup } from '../../utils/cleanup.js'; import { Message, MessageType, @@ -370,7 +371,8 @@ export const useSlashCommandProcessor = ( } case 'quit': setQuittingMessages(result.messages); - setTimeout(() => { + setTimeout(async () => { + await runExitCleanup(); process.exit(0); }, 100); return { type: 'handled' }; diff --git a/packages/cli/src/utils/cleanup.test.ts b/packages/cli/src/utils/cleanup.test.ts new file mode 100644 index 00000000..0b254bac --- /dev/null +++ b/packages/cli/src/utils/cleanup.test.ts @@ -0,0 +1,68 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi } from 'vitest'; +import { registerCleanup, runExitCleanup } from './cleanup'; + +describe('cleanup', () => { + const originalCleanupFunctions = global['cleanupFunctions']; + + beforeEach(() => { + // Isolate cleanup functions for each test + global['cleanupFunctions'] = []; + }); + + afterAll(() => { + // Restore original cleanup functions + global['cleanupFunctions'] = originalCleanupFunctions; + }); + + it('should run a registered synchronous function', async () => { + const cleanupFn = vi.fn(); + registerCleanup(cleanupFn); + + await runExitCleanup(); + + expect(cleanupFn).toHaveBeenCalledTimes(1); + }); + + it('should run a registered asynchronous function', async () => { + const cleanupFn = vi.fn().mockResolvedValue(undefined); + registerCleanup(cleanupFn); + + await runExitCleanup(); + + expect(cleanupFn).toHaveBeenCalledTimes(1); + }); + + it('should run multiple registered functions', async () => { + const syncFn = vi.fn(); + const asyncFn = vi.fn().mockResolvedValue(undefined); + + registerCleanup(syncFn); + registerCleanup(asyncFn); + + await runExitCleanup(); + + expect(syncFn).toHaveBeenCalledTimes(1); + expect(asyncFn).toHaveBeenCalledTimes(1); + }); + + it('should continue running cleanup functions even if one throws an error', async () => { + const errorFn = vi.fn(() => { + throw new Error('Test Error'); + }); + const successFn = vi.fn(); + + registerCleanup(errorFn); + registerCleanup(successFn); + + await runExitCleanup(); + + expect(errorFn).toHaveBeenCalledTimes(1); + expect(successFn).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/cli/src/utils/cleanup.ts b/packages/cli/src/utils/cleanup.ts index 628b881c..1200b6da 100644 --- a/packages/cli/src/utils/cleanup.ts +++ b/packages/cli/src/utils/cleanup.ts @@ -8,16 +8,16 @@ import { promises as fs } from 'fs'; import { join } from 'path'; import { getProjectTempDir } from '@google/gemini-cli-core'; -const cleanupFunctions: Array<() => void> = []; +const cleanupFunctions: Array<(() => void) | (() => Promise)> = []; -export function registerCleanup(fn: () => void) { +export function registerCleanup(fn: (() => void) | (() => Promise)) { cleanupFunctions.push(fn); } -export function runExitCleanup() { +export async function runExitCleanup() { for (const fn of cleanupFunctions) { try { - fn(); + await fn(); } catch (_) { // Ignore errors during cleanup. } diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 473ab5a6..a06c4505 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -684,7 +684,7 @@ export class Config { await this.ideClient.connect(); logIdeConnection(this, new IdeConnectionEvent(IdeConnectionType.SESSION)); } else { - this.ideClient.disconnect(); + await this.ideClient.disconnect(); } } diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index 5ffcc2e3..6d8be7fa 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -175,7 +175,14 @@ export class IdeClient { } } - disconnect() { + async disconnect() { + if (this.state.status === IDEConnectionStatus.Disconnected) { + return; + } + for (const filePath of this.diffResponses.keys()) { + await this.closeDiff(filePath); + } + this.diffResponses.clear(); this.setState( IDEConnectionStatus.Disconnected, 'IDE integration disabled. To enable it again, run /ide enable.',