[ide-mode] Close all open diffs when the CLI gets closed (#5792)

5 files changed, 117 insertions, 5 deletions

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<void>)> = [];
 
-export function registerCleanup(fn: () => void) {
+export function registerCleanup(fn: (() => void) | (() => Promise<void>)) {
   cleanupFunctions.push(fn);
 }
 
-export function runExitCleanup() {
+export async function runExitCleanup() {
   for (const fn of cleanupFunctions) {
     try {
-      fn();
+      await fn();
     } catch (_) {
       // Ignore errors during cleanup.
     }