Consistently guard restarts against concurrent auto updates (#21016)

Author: scidomino

packages/cli/src/ui/AppContainer.tsx

@@ -129,7 +129,7 @@ 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 { RELAUNCH_EXIT_CODE } from '../utils/processUtils.js';
+import { relaunchApp } from '../utils/processUtils.js';
 import type { SessionInfo } from '../utils/sessionUtils.js';
 import { useMessageQueue } from './hooks/useMessageQueue.js';
 import { useMcpStatus } from './hooks/useMcpStatus.js';
@@ -781,13 +781,12 @@ export const AppContainer = (props: AppContainerProps) => {
           authType === AuthType.LOGIN_WITH_GOOGLE &&
           config.isBrowserLaunchSuppressed()
         ) {
-          await runExitCleanup();
           writeToStdout(`
 ----------------------------------------------------------------
 Logging in with Google... Restarting Gemini CLI to continue.
 ----------------------------------------------------------------
           `);
-          process.exit(RELAUNCH_EXIT_CODE);
+          await relaunchApp();
         }
       }
       setAuthState(AuthState.Authenticated);
@@ -2497,8 +2496,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
             });
           }
         }
-        await runExitCleanup();
-        process.exit(RELAUNCH_EXIT_CODE);
+        await relaunchApp();
       },
       handleNewAgentsSelect: async (choice: NewAgentsChoice) => {
         if (newAgents && choice === NewAgentsChoice.ACKNOWLEDGE) {

packages/cli/src/ui/auth/AuthDialog.tsx

@@ -21,9 +21,8 @@ import {
 } from '@google/gemini-cli-core';
 import { useKeypress } from '../hooks/useKeypress.js';
 import { AuthState } from '../types.js';
-import { runExitCleanup } from '../../utils/cleanup.js';
 import { validateAuthMethodWithSettings } from './useAuth.js';
-import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
+import { relaunchApp } from '../../utils/processUtils.js';
 
 interface AuthDialogProps {
   config: Config;
@@ -133,10 +132,7 @@ export function AuthDialog({
           config.isBrowserLaunchSuppressed()
         ) {
           setExiting(true);
-          setTimeout(async () => {
-            await runExitCleanup();
-            process.exit(RELAUNCH_EXIT_CODE);
-          }, 100);
+          setTimeout(relaunchApp, 100);
           return;
         }
 

packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx

@@ -9,7 +9,10 @@ import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
 import { LoginWithGoogleRestartDialog } from './LoginWithGoogleRestartDialog.js';
 import { useKeypress } from '../hooks/useKeypress.js';
 import { runExitCleanup } from '../../utils/cleanup.js';
-import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
+import {
+  RELAUNCH_EXIT_CODE,
+  _resetRelaunchStateForTesting,
+} from '../../utils/processUtils.js';
 import { type Config } from '@google/gemini-cli-core';
 
 // Mocks
@@ -38,6 +41,7 @@ describe('LoginWithGoogleRestartDialog', () => {
     vi.clearAllMocks();
     exitSpy.mockClear();
     vi.useRealTimers();
+    _resetRelaunchStateForTesting();
   });
 
   it('renders correctly', async () => {

packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx

@@ -8,8 +8,7 @@ import { type Config } from '@google/gemini-cli-core';
 import { Box, Text } from 'ink';
 import { theme } from '../semantic-colors.js';
 import { useKeypress } from '../hooks/useKeypress.js';
-import { runExitCleanup } from '../../utils/cleanup.js';
-import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
+import { relaunchApp } from '../../utils/processUtils.js';
 
 interface LoginWithGoogleRestartDialogProps {
   onDismiss: () => void;
@@ -36,8 +35,7 @@ export const LoginWithGoogleRestartDialog = ({
               });
             }
           }
-          await runExitCleanup();
-          process.exit(RELAUNCH_EXIT_CODE);
+          await relaunchApp();
         }, 100);
         return true;
       }

packages/cli/src/ui/components/DialogManager.tsx

@@ -21,8 +21,7 @@ import { ProQuotaDialog } from './ProQuotaDialog.js';
 import { ValidationDialog } from './ValidationDialog.js';
 import { OverageMenuDialog } from './OverageMenuDialog.js';
 import { EmptyWalletDialog } from './EmptyWalletDialog.js';
-import { runExitCleanup } from '../../utils/cleanup.js';
-import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
+import { relaunchApp } from '../../utils/processUtils.js';
 import { SessionBrowser } from './SessionBrowser.js';
 import { PermissionsModifyTrustDialog } from './PermissionsModifyTrustDialog.js';
 import { ModelDialog } from './ModelDialog.js';
@@ -231,10 +230,7 @@ export const DialogManager = ({
       <Box flexDirection="column">
         <SettingsDialog
           onSelect={() => uiActions.closeSettingsDialog()}
-          onRestartRequest={async () => {
-            await runExitCleanup();
-            process.exit(RELAUNCH_EXIT_CODE);
-          }}
+          onRestartRequest={relaunchApp}
           availableTerminalHeight={terminalHeight - staticExtraHeight}
         />
       </Box>

packages/cli/src/ui/components/FolderTrustDialog.test.tsx

@@ -246,7 +246,9 @@ describe('FolderTrustDialog', () => {
 
   it('should call relaunchApp when isRestarting is true', async () => {
     vi.useFakeTimers();
-    const relaunchApp = vi.spyOn(processUtils, 'relaunchApp');
+    const relaunchApp = vi
+      .spyOn(processUtils, 'relaunchApp')
+      .mockResolvedValue(undefined);
     const { waitUntilReady, unmount } = renderWithProviders(
       <FolderTrustDialog onSelect={vi.fn()} isRestarting={true} />,
     );
@@ -259,7 +261,9 @@ describe('FolderTrustDialog', () => {
 
   it('should not call relaunchApp if unmounted before timeout', async () => {
     vi.useFakeTimers();
-    const relaunchApp = vi.spyOn(processUtils, 'relaunchApp');
+    const relaunchApp = vi
+      .spyOn(processUtils, 'relaunchApp')
+      .mockResolvedValue(undefined);
     const { waitUntilReady, unmount } = renderWithProviders(
       <FolderTrustDialog onSelect={vi.fn()} isRestarting={true} />,
     );

packages/cli/src/ui/components/FolderTrustDialog.tsx

@@ -54,9 +54,7 @@ export const FolderTrustDialog: React.FC<FolderTrustDialogProps> = ({
   useEffect(() => {
     let timer: ReturnType<typeof setTimeout>;
     if (isRestarting) {
-      timer = setTimeout(async () => {
-        await relaunchApp();
-      }, 250);
+      timer = setTimeout(relaunchApp, 250);
     }
     return () => {
       if (timer) clearTimeout(timer);

packages/cli/src/ui/components/IdeTrustChangeDialog.test.tsx

@@ -62,7 +62,9 @@ describe('IdeTrustChangeDialog', () => {
   });
 
   it('calls relaunchApp when "r" is pressed', async () => {
-    const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp');
+    const relaunchAppSpy = vi
+      .spyOn(processUtils, 'relaunchApp')
+      .mockResolvedValue(undefined);
     const { stdin, waitUntilReady, unmount } = renderWithProviders(
       <IdeTrustChangeDialog reason="NONE" />,
     );
@@ -78,7 +80,9 @@ describe('IdeTrustChangeDialog', () => {
   });
 
   it('calls relaunchApp when "R" is pressed', async () => {
-    const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp');
+    const relaunchAppSpy = vi
+      .spyOn(processUtils, 'relaunchApp')
+      .mockResolvedValue(undefined);
     const { stdin, waitUntilReady, unmount } = renderWithProviders(
       <IdeTrustChangeDialog reason="CONNECTION_CHANGE" />,
     );
@@ -94,7 +98,9 @@ describe('IdeTrustChangeDialog', () => {
   });
 
   it('does not call relaunchApp when another key is pressed', async () => {
-    const relaunchAppSpy = vi.spyOn(processUtils, 'relaunchApp');
+    const relaunchAppSpy = vi
+      .spyOn(processUtils, 'relaunchApp')
+      .mockResolvedValue(undefined);
     const { stdin, waitUntilReady, unmount } = renderWithProviders(
       <IdeTrustChangeDialog reason="CONNECTION_CHANGE" />,
     );

packages/cli/src/utils/processUtils.test.ts

@@ -5,7 +5,11 @@
  */
 
 import { vi } from 'vitest';
-import { RELAUNCH_EXIT_CODE, relaunchApp } from './processUtils.js';
+import {
+  RELAUNCH_EXIT_CODE,
+  relaunchApp,
+  _resetRelaunchStateForTesting,
+} from './processUtils.js';
 import * as cleanup from './cleanup.js';
 import * as handleAutoUpdate from './handleAutoUpdate.js';
 
@@ -19,6 +23,10 @@ describe('processUtils', () => {
     .mockReturnValue(undefined as never);
   const runExitCleanup = vi.spyOn(cleanup, 'runExitCleanup');
 
+  beforeEach(() => {
+    _resetRelaunchStateForTesting();
+  });
+
   afterEach(() => vi.clearAllMocks());
 
   it('should wait for updates, run cleanup, and exit with the relaunch code', async () => {

packages/cli/src/utils/processUtils.ts

@@ -15,7 +15,16 @@ export const RELAUNCH_EXIT_CODE = 199;
 /**
  * Exits the process with a special code to signal that the parent process should relaunch it.
  */
+let isRelaunching = false;
+
+/** @internal only for testing */
+export function _resetRelaunchStateForTesting(): void {
+  isRelaunching = false;
+}
+
 export async function relaunchApp(): Promise<void> {
+  if (isRelaunching) return;
+  isRelaunching = true;
   await waitForUpdateCompletion();
   await runExitCleanup();
   process.exit(RELAUNCH_EXIT_CODE);