feat(auth): improve API key authentication flow (#11760)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-15 15:50:35 -07:00 · 2025-10-29 18:58:08 -07:00
parent 6c8a48db13
commit 06035d5d43
25 changed files with 1216 additions and 76 deletions
--- a/packages/cli/src/ui/auth/ApiAuthDialog.test.tsx
+++ b/packages/cli/src/ui/auth/ApiAuthDialog.test.tsx
@@ -0,0 +1,125 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { render } from 'ink-testing-library';
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
+import { ApiAuthDialog } from './ApiAuthDialog.js';
+import { useKeypress } from '../hooks/useKeypress.js';
+import {
+  useTextBuffer,
+  type TextBuffer,
+} from '../components/shared/text-buffer.js';
+
+// Mocks
+vi.mock('../hooks/useKeypress.js', () => ({
+  useKeypress: vi.fn(),
+}));
+
+vi.mock('../components/shared/text-buffer.js', () => ({
+  useTextBuffer: vi.fn(),
+}));
+
+vi.mock('../contexts/UIStateContext.js', () => ({
+  useUIState: vi.fn(() => ({
+    mainAreaWidth: 80,
+  })),
+}));
+
+const mockedUseKeypress = useKeypress as Mock;
+const mockedUseTextBuffer = useTextBuffer as Mock;
+
+describe('ApiAuthDialog', () => {
+  const onSubmit = vi.fn();
+  const onCancel = vi.fn();
+  let mockBuffer: TextBuffer;
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+    mockBuffer = {
+      text: '',
+      lines: [''],
+      cursor: [0, 0],
+      visualCursor: [0, 0],
+      viewportVisualLines: [''],
+      handleInput: vi.fn(),
+      setText: vi.fn((newText) => {
+        mockBuffer.text = newText;
+        mockBuffer.viewportVisualLines = [newText];
+      }),
+    } as unknown as TextBuffer;
+    mockedUseTextBuffer.mockReturnValue(mockBuffer);
+  });
+
+  it('renders correctly', () => {
+    const { lastFrame } = render(
+      <ApiAuthDialog onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    expect(lastFrame()).toMatchSnapshot();
+  });
+
+  it('renders with a defaultValue', () => {
+    render(
+      <ApiAuthDialog
+        onSubmit={onSubmit}
+        onCancel={onCancel}
+        defaultValue="test-key"
+      />,
+    );
+    expect(mockedUseTextBuffer).toHaveBeenCalledWith(
+      expect.objectContaining({
+        initialText: 'test-key',
+        viewport: expect.objectContaining({
+          height: 4,
+        }),
+      }),
+    );
+  });
+
+  it('calls onSubmit when the text input is submitted', () => {
+    mockBuffer.text = 'submitted-key';
+    render(<ApiAuthDialog onSubmit={onSubmit} onCancel={onCancel} />);
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'return',
+      sequence: '\r',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(onSubmit).toHaveBeenCalledWith('submitted-key');
+  });
+
+  it('calls onCancel when the text input is cancelled', () => {
+    render(<ApiAuthDialog onSubmit={onSubmit} onCancel={onCancel} />);
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'escape',
+      sequence: '\u001b',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(onCancel).toHaveBeenCalled();
+  });
+
+  it('displays an error message', () => {
+    const { lastFrame } = render(
+      <ApiAuthDialog
+        onSubmit={onSubmit}
+        onCancel={onCancel}
+        error="Invalid API Key"
+      />,
+    );
+
+    expect(lastFrame()).toContain('Invalid API Key');
+  });
+});
--- a/packages/cli/src/ui/auth/ApiAuthDialog.tsx
+++ b/packages/cli/src/ui/auth/ApiAuthDialog.tsx
@@ -0,0 +1,97 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type React from 'react';
+import { Box, Text } from 'ink';
+import { theme } from '../semantic-colors.js';
+import { TextInput } from '../components/shared/TextInput.js';
+import { useTextBuffer } from '../components/shared/text-buffer.js';
+import { useUIState } from '../contexts/UIStateContext.js';
+
+interface ApiAuthDialogProps {
+  onSubmit: (apiKey: string) => void;
+  onCancel: () => void;
+  error?: string | null;
+  defaultValue?: string;
+}
+
+export function ApiAuthDialog({
+  onSubmit,
+  onCancel,
+  error,
+  defaultValue = '',
+}: ApiAuthDialogProps): React.JSX.Element {
+  const { mainAreaWidth } = useUIState();
+  const viewportWidth = mainAreaWidth - 8;
+
+  const buffer = useTextBuffer({
+    initialText: defaultValue || '',
+    initialCursorOffset: defaultValue?.length || 0,
+    viewport: {
+      width: viewportWidth,
+      height: 4,
+    },
+    isValidPath: () => false, // No path validation needed for API key
+    inputFilter: (text) =>
+      text.replace(/[^a-zA-Z0-9_-]/g, '').replace(/[\r\n]/g, ''),
+    singleLine: true,
+  });
+
+  const handleSubmit = (value: string) => {
+    onSubmit(value);
+  };
+
+  return (
+    <Box
+      borderStyle="round"
+      borderColor={theme.border.focused}
+      flexDirection="column"
+      padding={1}
+      width="100%"
+    >
+      <Text bold color={theme.text.primary}>
+        Enter Gemini API Key
+      </Text>
+      <Box marginTop={1} flexDirection="column">
+        <Text color={theme.text.primary}>
+          Please enter your Gemini API key. It will be securely stored in your
+          system keychain.
+        </Text>
+        <Text color={theme.text.secondary}>
+          You can get an API key from{' '}
+          <Text color={theme.text.link}>
+            https://aistudio.google.com/app/apikey
+          </Text>
+        </Text>
+      </Box>
+      <Box marginTop={1} flexDirection="row">
+        <Box
+          borderStyle="round"
+          borderColor={theme.border.default}
+          paddingX={1}
+          flexGrow={1}
+        >
+          <TextInput
+            buffer={buffer}
+            onSubmit={handleSubmit}
+            onCancel={onCancel}
+            placeholder="Paste your API key here"
+          />
+        </Box>
+      </Box>
+      {error && (
+        <Box marginTop={1}>
+          <Text color={theme.status.error}>{error}</Text>
+        </Box>
+      )}
+      <Box marginTop={1}>
+        <Text color={theme.text.secondary}>
+          (Press Enter to submit, Esc to cancel)
+        </Text>
+      </Box>
+    </Box>
+  );
+}
--- a/packages/cli/src/ui/auth/AuthDialog.test.tsx
+++ b/packages/cli/src/ui/auth/AuthDialog.test.tsx
@@ -191,7 +191,7 @@ describe('AuthDialog', () => {
        AuthType.USE_GEMINI,
      );
      expect(props.setAuthState).toHaveBeenCalledWith(
-        AuthState.Unauthenticated,
+        AuthState.AwaitingApiKeyInput,
      );
    });

--- a/packages/cli/src/ui/auth/AuthDialog.tsx
+++ b/packages/cli/src/ui/auth/AuthDialog.tsx
@@ -119,6 +119,10 @@ Logging in with Google... Please restart Gemini CLI to continue.
          process.exit(0);
        }
      }
+      if (authType === AuthType.USE_GEMINI) {
+        setAuthState(AuthState.AwaitingApiKeyInput);
+        return;
+      }
      setAuthState(AuthState.Unauthenticated);
    },
    [settings, config, setAuthState],
@@ -157,50 +161,52 @@ Logging in with Google... Please restart Gemini CLI to continue.
  return (
    <Box
      borderStyle="round"
-      borderColor={theme.border.default}
-      flexDirection="column"
+      borderColor={theme.border.focused}
+      flexDirection="row"
      padding={1}
      width="100%"
+      alignItems="flex-start"
    >
-      <Text bold color={theme.text.primary}>
-        Get started
-      </Text>
-      <Box marginTop={1}>
-        <Text color={theme.text.primary}>
-          How would you like to authenticate for this project?
+      <Text color={theme.text.accent}>? </Text>
+      <Box flexDirection="column" flexGrow={1}>
+        <Text bold color={theme.text.primary}>
+          Get started
        </Text>
-      </Box>
-      <Box marginTop={1}>
-        <RadioButtonSelect
-          items={items}
-          initialIndex={initialAuthIndex}
-          onSelect={handleAuthSelect}
-          onHighlight={() => {
-            onAuthError(null);
-          }}
-        />
-      </Box>
-      {authError && (
        <Box marginTop={1}>
-          <Text color={theme.status.error}>{authError}</Text>
+          <Text color={theme.text.primary}>
+            How would you like to authenticate for this project?
+          </Text>
+        </Box>
+        <Box marginTop={1}>
+          <RadioButtonSelect
+            items={items}
+            initialIndex={initialAuthIndex}
+            onSelect={handleAuthSelect}
+            onHighlight={() => {
+              onAuthError(null);
+            }}
+          />
+        </Box>
+        {authError && (
+          <Box marginTop={1}>
+            <Text color={theme.status.error}>{authError}</Text>
+          </Box>
+        )}
+        <Box marginTop={1}>
+          <Text color={theme.text.secondary}>(Use Enter to select)</Text>
+        </Box>
+        <Box marginTop={1}>
+          <Text color={theme.text.primary}>
+            Terms of Services and Privacy Notice for Gemini CLI
+          </Text>
+        </Box>
+        <Box marginTop={1}>
+          <Text color={theme.text.link}>
+            {
+              'https://github.com/google-gemini/gemini-cli/blob/main/docs/tos-privacy.md'
+            }
+          </Text>
        </Box>
-      )}
-      <Box marginTop={1}>
-        <Text color={theme.text.secondary}>
-          (Use Enter to select, Esc to close)
-        </Text>
-      </Box>
-      <Box marginTop={1}>
-        <Text color={theme.text.primary}>
-          Terms of Services and Privacy Notice for Gemini CLI
-        </Text>
-      </Box>
-      <Box marginTop={1}>
-        <Text color={theme.text.link}>
-          {
-            'https://github.com/google-gemini/gemini-cli/blob/main/docs/tos-privacy.md'
-          }
-        </Text>
      </Box>
    </Box>
  );
--- a/packages/cli/src/ui/auth/snapshots/ApiAuthDialog.test.tsx.snap
+++ b/packages/cli/src/ui/auth/snapshots/ApiAuthDialog.test.tsx.snap
@@ -0,0 +1,18 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`ApiAuthDialog > renders correctly 1`] = `
+"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
+│                                                                                                  │
+│ Enter Gemini API Key                                                                             │
+│                                                                                                  │
+│ Please enter your Gemini API key. It will be securely stored in your system keychain.            │
+│ You can get an API key from https://aistudio.google.com/app/apikey                               │
+│                                                                                                  │
+│ ╭──────────────────────────────────────────────────────────────────────────────────────────────╮ │
+│ │ Paste your API key here                                                                      │ │
+│ ╰──────────────────────────────────────────────────────────────────────────────────────────────╯ │
+│                                                                                                  │
+│ (Press Enter to submit, Esc to cancel)                                                           │
+│                                                                                                  │
+╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
+`;
--- a/packages/cli/src/ui/auth/useAuth.ts
+++ b/packages/cli/src/ui/auth/useAuth.ts
@@ -6,7 +6,12 @@

 import { useState, useEffect, useCallback } from 'react';
 import type { LoadedSettings } from '../../config/settings.js';
-import { AuthType, debugLogger, type Config } from '@google/gemini-cli-core';
+import {
+  AuthType,
+  type Config,
+  loadApiKey,
+  debugLogger,
+} from '@google/gemini-cli-core';
 import { getErrorMessage } from '@google/gemini-cli-core';
 import { AuthState } from '../types.js';
 import { validateAuthMethod } from '../../config/auth.js';
@@ -22,6 +27,10 @@ export function validateAuthMethodWithSettings(
  if (settings.merged.security?.auth?.useExternal) {
    return null;
  }
+  // If using Gemini API key, we don't validate it here as we might need to prompt for it.
+  if (authType === AuthType.USE_GEMINI) {
+    return null;
+  }
  return validateAuthMethod(authType);
 }

@@ -31,6 +40,9 @@ export const useAuthCommand = (settings: LoadedSettings, config: Config) => {
  );

  const [authError, setAuthError] = useState<string | null>(null);
+  const [apiKeyDefaultValue, setApiKeyDefaultValue] = useState<
+    string | undefined
+  >(undefined);

  const onAuthError = useCallback(
    (error: string | null) => {
@@ -42,6 +54,14 @@ export const useAuthCommand = (settings: LoadedSettings, config: Config) => {
    [setAuthError, setAuthState],
  );

+  const reloadApiKey = useCallback(async () => {
+    const storedKey = (await loadApiKey()) ?? '';
+    const envKey = process.env['GEMINI_API_KEY'] ?? '';
+    const key = storedKey || envKey;
+    setApiKeyDefaultValue(key);
+    return key; // Return the key for immediate use
+  }, []);
+
  useEffect(() => {
    (async () => {
      if (authState !== AuthState.Unauthenticated) {
@@ -59,6 +79,15 @@ export const useAuthCommand = (settings: LoadedSettings, config: Config) => {
        }
        return;
      }
+
+      if (authType === AuthType.USE_GEMINI) {
+        const key = await reloadApiKey(); // Use the unified function
+        if (!key) {
+          setAuthState(AuthState.AwaitingApiKeyInput);
+          return;
+        }
+      }
+
      const error = validateAuthMethodWithSettings(authType, settings);
      if (error) {
        onAuthError(error);
@@ -87,12 +116,22 @@ export const useAuthCommand = (settings: LoadedSettings, config: Config) => {
        onAuthError(`Failed to login. Message: ${getErrorMessage(e)}`);
      }
    })();
-  }, [settings, config, authState, setAuthState, setAuthError, onAuthError]);
+  }, [
+    settings,
+    config,
+    authState,
+    setAuthState,
+    setAuthError,
+    onAuthError,
+    reloadApiKey,
+  ]);

  return {
    authState,
    setAuthState,
    authError,
    onAuthError,
+    apiKeyDefaultValue,
+    reloadApiKey,
  };
 };