Refactor Authentication Components and Hooks (#7750)

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

@@ -0,0 +1,258 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { renderWithProviders } from '../../test-utils/render.js';
+import {
+  describe,
+  it,
+  expect,
+  vi,
+  beforeEach,
+  afterEach,
+  type Mock,
+} from 'vitest';
+import { AuthDialog } from './AuthDialog.js';
+import { AuthType, type Config } from '@google/gemini-cli-core';
+import type { LoadedSettings } from '../../config/settings.js';
+import { SettingScope } from '../../config/settings.js';
+import { AuthState } from '../types.js';
+import { RadioButtonSelect } from '../components/shared/RadioButtonSelect.js';
+import { useKeypress } from '../hooks/useKeypress.js';
+import { validateAuthMethodWithSettings } from './useAuth.js';
+import { runExitCleanup } from '../../utils/cleanup.js';
+import { clearCachedCredentialFile } from '@google/gemini-cli-core';
+import { Text } from 'ink';
+
+// Mocks
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
+  return {
+    ...actual,
+    clearCachedCredentialFile: vi.fn(),
+  };
+});
+
+vi.mock('../../utils/cleanup.js', () => ({
+  runExitCleanup: vi.fn(),
+}));
+
+vi.mock('./useAuth.js', () => ({
+  validateAuthMethodWithSettings: vi.fn(),
+}));
+
+vi.mock('../hooks/useKeypress.js', () => ({
+  useKeypress: vi.fn(),
+}));
+
+vi.mock('../components/shared/RadioButtonSelect.js', () => ({
+  RadioButtonSelect: vi.fn(({ items, initialIndex }) => (
+    <>
+      {items.map((item: { value: string; label: string }, index: number) => (
+        <Text key={item.value}>
+          {index === initialIndex ? '(selected)' : '(not selected)'}{' '}
+          {item.label}
+        </Text>
+      ))}
+    </>
+  )),
+}));
+
+const mockedUseKeypress = useKeypress as Mock;
+const mockedRadioButtonSelect = RadioButtonSelect as Mock;
+const mockedValidateAuthMethod = validateAuthMethodWithSettings as Mock;
+const mockedRunExitCleanup = runExitCleanup as Mock;
+const mockedClearCachedCredentialFile = clearCachedCredentialFile as Mock;
+
+describe('AuthDialog', () => {
+  let props: {
+    config: Config;
+    settings: LoadedSettings;
+    setAuthState: (state: AuthState) => void;
+    authError: string | null;
+    onAuthError: (error: string) => void;
+  };
+  const originalEnv = { ...process.env };
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+    process.env = {};
+
+    props = {
+      config: {
+        isBrowserLaunchSuppressed: vi.fn().mockReturnValue(false),
+      } as unknown as Config,
+      settings: {
+        merged: {
+          security: {
+            auth: {},
+          },
+        },
+        setValue: vi.fn(),
+      } as unknown as LoadedSettings,
+      setAuthState: vi.fn(),
+      authError: null,
+      onAuthError: vi.fn(),
+    };
+  });
+
+  afterEach(() => {
+    process.env = originalEnv;
+  });
+
+  it('shows Cloud Shell option when in Cloud Shell environment', () => {
+    process.env['CLOUD_SHELL'] = 'true';
+    renderWithProviders(<AuthDialog {...props} />);
+    const items = mockedRadioButtonSelect.mock.calls[0][0].items;
+    expect(items).toContainEqual({
+      label: 'Use Cloud Shell user credentials',
+      value: AuthType.CLOUD_SHELL,
+    });
+  });
+
+  it('filters auth types when enforcedType is set', () => {
+    props.settings.merged.security!.auth!.enforcedType = AuthType.USE_GEMINI;
+    renderWithProviders(<AuthDialog {...props} />);
+    const items = mockedRadioButtonSelect.mock.calls[0][0].items;
+    expect(items).toHaveLength(1);
+    expect(items[0].value).toBe(AuthType.USE_GEMINI);
+  });
+
+  it('sets initial index to 0 when enforcedType is set', () => {
+    props.settings.merged.security!.auth!.enforcedType = AuthType.USE_GEMINI;
+    renderWithProviders(<AuthDialog {...props} />);
+    const { initialIndex } = mockedRadioButtonSelect.mock.calls[0][0];
+    expect(initialIndex).toBe(0);
+  });
+
+  it('selects initial auth type from settings', () => {
+    props.settings.merged.security!.auth!.selectedType = AuthType.USE_VERTEX_AI;
+    renderWithProviders(<AuthDialog {...props} />);
+    const { items, initialIndex } = mockedRadioButtonSelect.mock.calls[0][0];
+    expect(items[initialIndex].value).toBe(AuthType.USE_VERTEX_AI);
+  });
+
+  it('selects initial auth type from GEMINI_DEFAULT_AUTH_TYPE env var', () => {
+    process.env['GEMINI_DEFAULT_AUTH_TYPE'] = AuthType.USE_GEMINI;
+    renderWithProviders(<AuthDialog {...props} />);
+    const { items, initialIndex } = mockedRadioButtonSelect.mock.calls[0][0];
+    expect(items[initialIndex].value).toBe(AuthType.USE_GEMINI);
+  });
+
+  it('selects initial auth type from GEMINI_API_KEY env var', () => {
+    process.env['GEMINI_API_KEY'] = 'test-key';
+    renderWithProviders(<AuthDialog {...props} />);
+    const { items, initialIndex } = mockedRadioButtonSelect.mock.calls[0][0];
+    expect(items[initialIndex].value).toBe(AuthType.USE_GEMINI);
+  });
+
+  it('defaults to Login with Google', () => {
+    renderWithProviders(<AuthDialog {...props} />);
+    const { items, initialIndex } = mockedRadioButtonSelect.mock.calls[0][0];
+    expect(items[initialIndex].value).toBe(AuthType.LOGIN_WITH_GOOGLE);
+  });
+
+  describe('handleAuthSelect', () => {
+    it('calls onAuthError if validation fails', () => {
+      mockedValidateAuthMethod.mockReturnValue('Invalid method');
+      renderWithProviders(<AuthDialog {...props} />);
+      const { onSelect: handleAuthSelect } =
+        mockedRadioButtonSelect.mock.calls[0][0];
+      handleAuthSelect(AuthType.USE_GEMINI);
+
+      expect(mockedValidateAuthMethod).toHaveBeenCalledWith(
+        AuthType.USE_GEMINI,
+        props.settings,
+      );
+      expect(props.onAuthError).toHaveBeenCalledWith('Invalid method');
+      expect(props.settings.setValue).not.toHaveBeenCalled();
+    });
+
+    it('calls onSelect if validation passes', async () => {
+      mockedValidateAuthMethod.mockReturnValue(null);
+      renderWithProviders(<AuthDialog {...props} />);
+      const { onSelect: handleAuthSelect } =
+        mockedRadioButtonSelect.mock.calls[0][0];
+      await handleAuthSelect(AuthType.USE_GEMINI);
+
+      expect(mockedValidateAuthMethod).toHaveBeenCalledWith(
+        AuthType.USE_GEMINI,
+        props.settings,
+      );
+      expect(props.onAuthError).not.toHaveBeenCalled();
+      expect(mockedClearCachedCredentialFile).toHaveBeenCalled();
+      expect(props.settings.setValue).toHaveBeenCalledWith(
+        SettingScope.User,
+        'security.auth.selectedType',
+        AuthType.USE_GEMINI,
+      );
+      expect(props.setAuthState).toHaveBeenCalledWith(
+        AuthState.Unauthenticated,
+      );
+    });
+
+    it('exits process for Login with Google when browser is suppressed', async () => {
+      const exitSpy = vi
+        .spyOn(process, 'exit')
+        .mockImplementation(() => undefined as never);
+      const logSpy = vi.spyOn(console, 'log').mockImplementation(() => {});
+      vi.mocked(props.config.isBrowserLaunchSuppressed).mockReturnValue(true);
+      mockedValidateAuthMethod.mockReturnValue(null);
+
+      renderWithProviders(<AuthDialog {...props} />);
+      const { onSelect: handleAuthSelect } =
+        mockedRadioButtonSelect.mock.calls[0][0];
+      await handleAuthSelect(AuthType.LOGIN_WITH_GOOGLE);
+
+      expect(mockedRunExitCleanup).toHaveBeenCalled();
+      expect(logSpy).toHaveBeenCalledWith(
+        expect.stringContaining('Please restart Gemini CLI'),
+      );
+      expect(exitSpy).toHaveBeenCalledWith(0);
+
+      exitSpy.mockRestore();
+      logSpy.mockRestore();
+    });
+  });
+
+  it('displays authError when provided', () => {
+    props.authError = 'Something went wrong';
+    const { lastFrame } = renderWithProviders(<AuthDialog {...props} />);
+    expect(lastFrame()).toContain('Something went wrong');
+  });
+
+  describe('useKeypress', () => {
+    it('does nothing on escape if authError is present', () => {
+      props.authError = 'Some error';
+      renderWithProviders(<AuthDialog {...props} />);
+      const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+      keypressHandler({ name: 'escape' });
+      expect(props.onAuthError).not.toHaveBeenCalled();
+      expect(props.setAuthState).not.toHaveBeenCalled();
+    });
+
+    it('calls onAuthError on escape if no auth method is set', () => {
+      props.settings.merged.security!.auth!.selectedType = undefined;
+      renderWithProviders(<AuthDialog {...props} />);
+      const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+      keypressHandler({ name: 'escape' });
+      expect(props.onAuthError).toHaveBeenCalledWith(
+        'You must select an auth method to proceed. Press Ctrl+C twice to exit.',
+      );
+    });
+
+    it('calls onSelect(undefined) on escape if auth method is set', () => {
+      props.settings.merged.security!.auth!.selectedType = AuthType.USE_GEMINI;
+      renderWithProviders(<AuthDialog {...props} />);
+      const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+      keypressHandler({ name: 'escape' });
+      expect(props.setAuthState).toHaveBeenCalledWith(
+        AuthState.Unauthenticated,
+      );
+      expect(props.settings.setValue).not.toHaveBeenCalled();
+    });
+  });
+});

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

@@ -0,0 +1,188 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type React from 'react';
+import { useCallback } from 'react';
+import { Box, Text } from 'ink';
+import { Colors } from '../colors.js';
+import { RadioButtonSelect } from '../components/shared/RadioButtonSelect.js';
+import type { LoadedSettings } from '../../config/settings.js';
+import { SettingScope } from '../../config/settings.js';
+import {
+  AuthType,
+  clearCachedCredentialFile,
+  type Config,
+} 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';
+
+interface AuthDialogProps {
+  config: Config;
+  settings: LoadedSettings;
+  setAuthState: (state: AuthState) => void;
+  authError: string | null;
+  onAuthError: (error: string) => void;
+}
+
+export function AuthDialog({
+  config,
+  settings,
+  setAuthState,
+  authError,
+  onAuthError,
+}: AuthDialogProps): React.JSX.Element {
+  let items = [
+    {
+      label: 'Login with Google',
+      value: AuthType.LOGIN_WITH_GOOGLE,
+    },
+    ...(process.env['CLOUD_SHELL'] === 'true'
+      ? [
+          {
+            label: 'Use Cloud Shell user credentials',
+            value: AuthType.CLOUD_SHELL,
+          },
+        ]
+      : []),
+    {
+      label: 'Use Gemini API Key',
+      value: AuthType.USE_GEMINI,
+    },
+    { label: 'Vertex AI', value: AuthType.USE_VERTEX_AI },
+  ];
+
+  if (settings.merged.security?.auth?.enforcedType) {
+    items = items.filter(
+      (item) => item.value === settings.merged.security?.auth?.enforcedType,
+    );
+  }
+
+  let defaultAuthType = null;
+  const defaultAuthTypeEnv = process.env['GEMINI_DEFAULT_AUTH_TYPE'];
+  if (
+    defaultAuthTypeEnv &&
+    Object.values(AuthType).includes(defaultAuthTypeEnv as AuthType)
+  ) {
+    defaultAuthType = defaultAuthTypeEnv as AuthType;
+  }
+
+  let initialAuthIndex = items.findIndex((item) => {
+    if (settings.merged.security?.auth?.selectedType) {
+      return item.value === settings.merged.security.auth.selectedType;
+    }
+
+    if (defaultAuthType) {
+      return item.value === defaultAuthType;
+    }
+
+    if (process.env['GEMINI_API_KEY']) {
+      return item.value === AuthType.USE_GEMINI;
+    }
+
+    return item.value === AuthType.LOGIN_WITH_GOOGLE;
+  });
+  if (settings.merged.security?.auth?.enforcedType) {
+    initialAuthIndex = 0;
+  }
+
+  const onSelect = useCallback(
+    async (authType: AuthType | undefined, scope: SettingScope) => {
+      if (authType) {
+        await clearCachedCredentialFile();
+
+        settings.setValue(scope, 'security.auth.selectedType', authType);
+        if (
+          authType === AuthType.LOGIN_WITH_GOOGLE &&
+          config.isBrowserLaunchSuppressed()
+        ) {
+          runExitCleanup();
+          console.log(
+            `
+----------------------------------------------------------------
+Logging in with Google... Please restart Gemini CLI to continue.
+----------------------------------------------------------------
+            `,
+          );
+          process.exit(0);
+        }
+      }
+      setAuthState(AuthState.Unauthenticated);
+    },
+    [settings, config, setAuthState],
+  );
+
+  const handleAuthSelect = (authMethod: AuthType) => {
+    const error = validateAuthMethodWithSettings(authMethod, settings);
+    if (error) {
+      onAuthError(error);
+    } else {
+      onSelect(authMethod, SettingScope.User);
+    }
+  };
+
+  useKeypress(
+    (key) => {
+      if (key.name === 'escape') {
+        // Prevent exit if there is an error message.
+        // This means they user is not authenticated yet.
+        if (authError) {
+          return;
+        }
+        if (settings.merged.security?.auth?.selectedType === undefined) {
+          // Prevent exiting if no auth method is set
+          onAuthError(
+            'You must select an auth method to proceed. Press Ctrl+C twice to exit.',
+          );
+          return;
+        }
+        onSelect(undefined, SettingScope.User);
+      }
+    },
+    { isActive: true },
+  );
+
+  return (
+    <Box
+      borderStyle="round"
+      borderColor={Colors.Gray}
+      flexDirection="column"
+      padding={1}
+      width="100%"
+    >
+      <Text bold>Get started</Text>
+      <Box marginTop={1}>
+        <Text>How would you like to authenticate for this project?</Text>
+      </Box>
+      <Box marginTop={1}>
+        <RadioButtonSelect
+          items={items}
+          initialIndex={initialAuthIndex}
+          onSelect={handleAuthSelect}
+        />
+      </Box>
+      {authError && (
+        <Box marginTop={1}>
+          <Text color={Colors.AccentRed}>{authError}</Text>
+        </Box>
+      )}
+      <Box marginTop={1}>
+        <Text color={Colors.Gray}>(Use Enter to select)</Text>
+      </Box>
+      <Box marginTop={1}>
+        <Text>Terms of Services and Privacy Notice for Gemini CLI</Text>
+      </Box>
+      <Box marginTop={1}>
+        <Text color={Colors.AccentBlue}>
+          {
+            'https://github.com/google-gemini/gemini-cli/blob/main/docs/tos-privacy.md'
+          }
+        </Text>
+      </Box>
+    </Box>
+  );
+}

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

@@ -0,0 +1,63 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type React from 'react';
+import { useState, useEffect } from 'react';
+import { Box, Text } from 'ink';
+import Spinner from 'ink-spinner';
+import { Colors } from '../colors.js';
+import { useKeypress } from '../hooks/useKeypress.js';
+
+interface AuthInProgressProps {
+  onTimeout: () => void;
+}
+
+export function AuthInProgress({
+  onTimeout,
+}: AuthInProgressProps): React.JSX.Element {
+  const [timedOut, setTimedOut] = useState(false);
+
+  useKeypress(
+    (key) => {
+      if (key.name === 'escape' || (key.ctrl && key.name === 'c')) {
+        onTimeout();
+      }
+    },
+    { isActive: true },
+  );
+
+  useEffect(() => {
+    const timer = setTimeout(() => {
+      setTimedOut(true);
+      onTimeout();
+    }, 180000);
+
+    return () => clearTimeout(timer);
+  }, [onTimeout]);
+
+  return (
+    <Box
+      borderStyle="round"
+      borderColor={Colors.Gray}
+      flexDirection="column"
+      padding={1}
+      width="100%"
+    >
+      {timedOut ? (
+        <Text color={Colors.AccentRed}>
+          Authentication timed out. Please try again.
+        </Text>
+      ) : (
+        <Box>
+          <Text>
+            <Spinner type="dots" /> Waiting for auth... (Press ESC or CTRL+C to
+            cancel)
+          </Text>
+        </Box>
+      )}
+    </Box>
+  );
+}

packages/cli/src/ui/auth/useAuth.ts

@@ -0,0 +1,96 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useState, useEffect, useCallback } from 'react';
+import type { LoadedSettings } from '../../config/settings.js';
+import { AuthType, type Config } from '@google/gemini-cli-core';
+import { getErrorMessage } from '@google/gemini-cli-core';
+import { AuthState } from '../types.js';
+import { validateAuthMethod } from '../../config/auth.js';
+
+export function validateAuthMethodWithSettings(
+  authType: AuthType,
+  settings: LoadedSettings,
+): string | null {
+  const enforcedType = settings.merged.security?.auth?.enforcedType;
+  if (enforcedType && enforcedType !== authType) {
+    return `Authentication is enforced to be ${enforcedType}, but you are currently using ${authType}.`;
+  }
+  if (settings.merged.security?.auth?.useExternal) {
+    return null;
+  }
+  return validateAuthMethod(authType);
+}
+
+export const useAuthCommand = (settings: LoadedSettings, config: Config) => {
+  const [authState, setAuthState] = useState<AuthState>(
+    AuthState.Unauthenticated,
+  );
+
+  const [authError, setAuthError] = useState<string | null>(null);
+
+  const onAuthError = useCallback(
+    (error: string) => {
+      setAuthError(error);
+      setAuthState(AuthState.Updating);
+    },
+    [setAuthError, setAuthState],
+  );
+
+  useEffect(() => {
+    (async () => {
+      if (authState !== AuthState.Unauthenticated) {
+        return;
+      }
+
+      const authType = settings.merged.security?.auth?.selectedType;
+      if (!authType) {
+        if (process.env['GEMINI_API_KEY']) {
+          onAuthError(
+            'Existing API key detected (GEMINI_API_KEY). Select "Gemini API Key" option to use it.',
+          );
+        } else {
+          onAuthError('No authentication method selected.');
+        }
+        return;
+      }
+      const error = validateAuthMethodWithSettings(authType, settings);
+      if (error) {
+        onAuthError(error);
+        return;
+      }
+
+      const defaultAuthType = process.env['GEMINI_DEFAULT_AUTH_TYPE'];
+      if (
+        defaultAuthType &&
+        !Object.values(AuthType).includes(defaultAuthType as AuthType)
+      ) {
+        onAuthError(
+          `Invalid value for GEMINI_DEFAULT_AUTH_TYPE: "${defaultAuthType}". ` +
+            `Valid values are: ${Object.values(AuthType).join(', ')}.`,
+        );
+        return;
+      }
+
+      try {
+        await config.refreshAuth(authType);
+
+        console.log(`Authenticated via "${authType}".`);
+        setAuthError(null);
+        setAuthState(AuthState.Authenticated);
+      } catch (e) {
+        onAuthError(`Failed to login. Message: ${getErrorMessage(e)}`);
+      }
+    })();
+  }, [settings, config, authState, setAuthState, setAuthError, onAuthError]);
+
+  return {
+    authState,
+    setAuthState,
+    authError,
+    onAuthError,
+  };
+};