feat(admin): prompt user to restart the CLI if they change auth to oauth mid-session or don't have auth type selected at start of session (#16426)

packages/cli/src/test-utils/render.tsx

@@ -171,6 +171,7 @@ const mockUIActions: UIActions = {
   handleApiKeyCancel: vi.fn(),
   setBannerVisible: vi.fn(),
   setEmbeddedShellFocused: vi.fn(),
+  setAuthContext: vi.fn(),
 };
 
 export const renderWithProviders = (

packages/cli/src/ui/AppContainer.tsx

@@ -126,6 +126,7 @@ import {
   WARNING_PROMPT_DURATION_MS,
   QUEUE_ERROR_DISPLAY_DURATION_MS,
 } from './constants.js';
+import { LoginWithGoogleRestartDialog } from './auth/LoginWithGoogleRestartDialog.js';
 
 function isToolExecuting(pendingHistoryItems: HistoryItemWithoutId[]) {
   return pendingHistoryItems.some((item) => {
@@ -468,6 +469,16 @@ export const AppContainer = (props: AppContainerProps) => {
     apiKeyDefaultValue,
     reloadApiKey,
   } = useAuthCommand(settings, config);
+  const [authContext, setAuthContext] = useState<{ requiresRestart?: boolean }>(
+    {},
+  );
+
+  useEffect(() => {
+    if (authState === AuthState.Authenticated && authContext.requiresRestart) {
+      setAuthState(AuthState.AwaitingGoogleLoginRestart);
+      setAuthContext({});
+    }
+  }, [authState, authContext, setAuthState]);
 
   const { proQuotaRequest, handleProQuotaChoice } = useQuotaAndFallback({
     config,
@@ -511,6 +522,11 @@ export const AppContainer = (props: AppContainerProps) => {
   const handleAuthSelect = useCallback(
     async (authType: AuthType | undefined, scope: LoadableSettingScope) => {
       if (authType) {
+        if (authType === AuthType.LOGIN_WITH_GOOGLE) {
+          setAuthContext({ requiresRestart: true });
+        } else {
+          setAuthContext({});
+        }
         await clearCachedCredentialFile();
         settings.setValue(scope, 'security.auth.selectedType', authType);
 
@@ -539,7 +555,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
       }
       setAuthState(AuthState.Authenticated);
     },
-    [settings, config, setAuthState, onAuthError],
+    [settings, config, setAuthState, onAuthError, setAuthContext],
   );
 
   const handleApiKeySubmit = useCallback(
@@ -1687,6 +1703,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
       handleApiKeyCancel,
       setBannerVisible,
       setEmbeddedShellFocused,
+      setAuthContext,
     }),
     [
       handleThemeSelect,
@@ -1722,9 +1739,21 @@ Logging in with Google... Restarting Gemini CLI to continue.
       handleApiKeyCancel,
       setBannerVisible,
       setEmbeddedShellFocused,
+      setAuthContext,
     ],
   );
 
+  if (authState === AuthState.AwaitingGoogleLoginRestart) {
+    return (
+      <LoginWithGoogleRestartDialog
+        onDismiss={() => {
+          setAuthContext({});
+          setAuthState(AuthState.Updating);
+        }}
+      />
+    );
+  }
+
   return (
     <UIStateContext.Provider value={uiState}>
       <UIActionsContext.Provider value={uiActions}>

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

@@ -72,6 +72,7 @@ describe('AuthDialog', () => {
     setAuthState: (state: AuthState) => void;
     authError: string | null;
     onAuthError: (error: string | null) => void;
+    setAuthContext: (context: { requiresRestart?: boolean }) => void;
   };
   const originalEnv = { ...process.env };
 
@@ -94,6 +95,7 @@ describe('AuthDialog', () => {
       setAuthState: vi.fn(),
       authError: null,
       onAuthError: vi.fn(),
+      setAuthContext: vi.fn(),
     };
   });
 
@@ -217,6 +219,28 @@ describe('AuthDialog', () => {
       expect(props.settings.setValue).not.toHaveBeenCalled();
     });
 
+    it('sets auth context with requiresRestart: true for LOGIN_WITH_GOOGLE', async () => {
+      mockedValidateAuthMethod.mockReturnValue(null);
+      renderWithProviders(<AuthDialog {...props} />);
+      const { onSelect: handleAuthSelect } =
+        mockedRadioButtonSelect.mock.calls[0][0];
+      await handleAuthSelect(AuthType.LOGIN_WITH_GOOGLE);
+
+      expect(props.setAuthContext).toHaveBeenCalledWith({
+        requiresRestart: true,
+      });
+    });
+
+    it('sets auth context with empty object for other auth types', async () => {
+      mockedValidateAuthMethod.mockReturnValue(null);
+      renderWithProviders(<AuthDialog {...props} />);
+      const { onSelect: handleAuthSelect } =
+        mockedRadioButtonSelect.mock.calls[0][0];
+      await handleAuthSelect(AuthType.USE_GEMINI);
+
+      expect(props.setAuthContext).toHaveBeenCalledWith({});
+    });
+
     it('skips API key dialog on initial setup if env var is present', async () => {
       mockedValidateAuthMethod.mockReturnValue(null);
       process.env['GEMINI_API_KEY'] = 'test-key-from-env';

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

@@ -31,6 +31,7 @@ interface AuthDialogProps {
   setAuthState: (state: AuthState) => void;
   authError: string | null;
   onAuthError: (error: string | null) => void;
+  setAuthContext: (context: { requiresRestart?: boolean }) => void;
 }
 
 export function AuthDialog({
@@ -39,6 +40,7 @@ export function AuthDialog({
   setAuthState,
   authError,
   onAuthError,
+  setAuthContext,
 }: AuthDialogProps): React.JSX.Element {
   const [exiting, setExiting] = useState(false);
   let items = [
@@ -116,6 +118,11 @@ export function AuthDialog({
         return;
       }
       if (authType) {
+        if (authType === AuthType.LOGIN_WITH_GOOGLE) {
+          setAuthContext({ requiresRestart: true });
+        } else {
+          setAuthContext({});
+        }
         await clearCachedCredentialFile();
 
         settings.setValue(scope, 'security.auth.selectedType', authType);
@@ -143,7 +150,7 @@ export function AuthDialog({
       }
       setAuthState(AuthState.Unauthenticated);
     },
-    [settings, config, setAuthState, exiting],
+    [settings, config, setAuthState, exiting, setAuthContext],
   );
 
   const handleAuthSelect = (authMethod: AuthType) => {

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

@@ -0,0 +1,87 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { render } from '../../test-utils/render.js';
+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';
+
+// Mocks
+vi.mock('../hooks/useKeypress.js', () => ({
+  useKeypress: vi.fn(),
+}));
+
+vi.mock('../../utils/cleanup.js', () => ({
+  runExitCleanup: vi.fn(),
+}));
+
+const mockedUseKeypress = useKeypress as Mock;
+const mockedRunExitCleanup = runExitCleanup as Mock;
+
+describe('LoginWithGoogleRestartDialog', () => {
+  const onDismiss = vi.fn();
+  const exitSpy = vi
+    .spyOn(process, 'exit')
+    .mockImplementation(() => undefined as never);
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    exitSpy.mockClear();
+    vi.useRealTimers();
+  });
+
+  it('renders correctly', () => {
+    const { lastFrame } = render(
+      <LoginWithGoogleRestartDialog onDismiss={onDismiss} />,
+    );
+    expect(lastFrame()).toMatchSnapshot();
+  });
+
+  it('calls onDismiss when escape is pressed', () => {
+    render(<LoginWithGoogleRestartDialog onDismiss={onDismiss} />);
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'escape',
+      sequence: '\u001b',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(onDismiss).toHaveBeenCalledTimes(1);
+  });
+
+  it.each(['r', 'R'])(
+    'calls runExitCleanup and process.exit when %s is pressed',
+    async (keyName) => {
+      vi.useFakeTimers();
+
+      render(<LoginWithGoogleRestartDialog onDismiss={onDismiss} />);
+      const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+      keypressHandler({
+        name: keyName,
+        sequence: keyName,
+        ctrl: false,
+        meta: false,
+        shift: false,
+        paste: false,
+      });
+
+      // Advance timers to trigger the setTimeout callback
+      await vi.runAllTimersAsync();
+
+      expect(mockedRunExitCleanup).toHaveBeenCalledTimes(1);
+      expect(exitSpy).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE);
+
+      vi.useRealTimers();
+    },
+  );
+});

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

@@ -0,0 +1,45 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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';
+
+interface LoginWithGoogleRestartDialogProps {
+  onDismiss: () => void;
+}
+
+export const LoginWithGoogleRestartDialog = ({
+  onDismiss,
+}: LoginWithGoogleRestartDialogProps) => {
+  useKeypress(
+    (key) => {
+      if (key.name === 'escape') {
+        onDismiss();
+      } else if (key.name === 'r' || key.name === 'R') {
+        setTimeout(async () => {
+          await runExitCleanup();
+          process.exit(RELAUNCH_EXIT_CODE);
+        }, 100);
+      }
+    },
+    { isActive: true },
+  );
+
+  const message =
+    'You have successfully logged in with Google. Gemini CLI needs to be restarted.';
+
+  return (
+    <Box borderStyle="round" borderColor={theme.status.warning} paddingX={1}>
+      <Text color={theme.status.warning}>
+        {message} Press &apos;r&apos; to restart, or &apos;escape&apos; to
+        choose a different auth method.
+      </Text>
+    </Box>
+  );
+};

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

@@ -0,0 +1,8 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`LoginWithGoogleRestartDialog > renders correctly 1`] = `
+"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
+│ You have successfully logged in with Google. Gemini CLI needs to be restarted. Press 'r' to      │
+│ restart, or 'escape' to choose a different auth method.                                          │
+╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
+`;

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

@@ -183,6 +183,7 @@ export const DialogManager = ({
           setAuthState={uiActions.setAuthState}
           authError={uiState.authError}
           onAuthError={uiActions.onAuthError}
+          setAuthContext={uiActions.setAuthContext}
         />
       </Box>
     );

packages/cli/src/ui/contexts/UIActionsContext.tsx

@@ -56,6 +56,7 @@ export interface UIActions {
   handleApiKeyCancel: () => void;
   setBannerVisible: (visible: boolean) => void;
   setEmbeddedShellFocused: (value: boolean) => void;
+  setAuthContext: (context: { requiresRestart?: boolean }) => void;
 }
 
 export const UIActionsContext = createContext<UIActions | null>(null);

packages/cli/src/ui/types.ts

@@ -30,6 +30,8 @@ export enum AuthState {
   AwaitingApiKeyInput = 'awaiting_api_key_input',
   // Successfully authenticated
   Authenticated = 'authenticated',
+  // Waiting for the user to restart after a Google login
+  AwaitingGoogleLoginRestart = 'awaiting_google_login_restart',
 }
 
 // Only defining the state enum needed by the UI