From d65eab01d2569474c7866c6e76ce8060a6ce1535 Mon Sep 17 00:00:00 2001 From: Shreya Keshive Date: Mon, 12 Jan 2026 15:39:08 -0500 Subject: [PATCH] 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 | 1 + packages/cli/src/ui/AppContainer.tsx | 31 ++++++- packages/cli/src/ui/auth/AuthDialog.test.tsx | 24 +++++ packages/cli/src/ui/auth/AuthDialog.tsx | 9 +- .../LoginWithGoogleRestartDialog.test.tsx | 87 +++++++++++++++++++ .../ui/auth/LoginWithGoogleRestartDialog.tsx | 45 ++++++++++ ...LoginWithGoogleRestartDialog.test.tsx.snap | 8 ++ .../cli/src/ui/components/DialogManager.tsx | 1 + .../cli/src/ui/contexts/UIActionsContext.tsx | 1 + packages/cli/src/ui/types.ts | 2 + 10 files changed, 207 insertions(+), 2 deletions(-) create mode 100644 packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx create mode 100644 packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx create mode 100644 packages/cli/src/ui/auth/__snapshots__/LoginWithGoogleRestartDialog.test.tsx.snap diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index e083918683..3f77acd7a7 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/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 = ( diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 19f4ed44f2..e98a6476a2 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/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 ( + { + setAuthContext({}); + setAuthState(AuthState.Updating); + }} + /> + ); + } + return ( diff --git a/packages/cli/src/ui/auth/AuthDialog.test.tsx b/packages/cli/src/ui/auth/AuthDialog.test.tsx index 16f0f9cbe8..66be01856d 100644 --- a/packages/cli/src/ui/auth/AuthDialog.test.tsx +++ b/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(); + 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(); + 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'; diff --git a/packages/cli/src/ui/auth/AuthDialog.tsx b/packages/cli/src/ui/auth/AuthDialog.tsx index b133acf52b..558927dcf2 100644 --- a/packages/cli/src/ui/auth/AuthDialog.tsx +++ b/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) => { diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx new file mode 100644 index 0000000000..5dd9d0c171 --- /dev/null +++ b/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( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); + + it('calls onDismiss when escape is pressed', () => { + render(); + 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(); + 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(); + }, + ); +}); diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.tsx new file mode 100644 index 0000000000..0418e3f3f3 --- /dev/null +++ b/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 ( + + + {message} Press 'r' to restart, or 'escape' to + choose a different auth method. + + + ); +}; diff --git a/packages/cli/src/ui/auth/__snapshots__/LoginWithGoogleRestartDialog.test.tsx.snap b/packages/cli/src/ui/auth/__snapshots__/LoginWithGoogleRestartDialog.test.tsx.snap new file mode 100644 index 0000000000..effd559184 --- /dev/null +++ b/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. │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 132b1a020e..6a2fc46568 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/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} /> ); diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index 120def8b1c..85839829f5 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/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(null); diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 096caf862a..4f9a970278 100644 --- a/packages/cli/src/ui/types.ts +++ b/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