2 files changed, 72 insertions, 4 deletions
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.
     }