From 04518b52c0ddcd5ae1192763c55e472add218b3c Mon Sep 17 00:00:00 2001 From: matt korwel Date: Thu, 19 Jun 2025 16:52:22 -0700 Subject: Auth First Run (#1207) Co-authored-by: Tommaso Sciortino Co-authored-by: N. Taylor Mullen --- packages/cli/src/ui/App.test.tsx | 19 +++-- packages/cli/src/ui/App.tsx | 42 ++++++++++ packages/cli/src/ui/components/AuthDialog.test.tsx | 41 ++++++++++ packages/cli/src/ui/components/AuthDialog.tsx | 94 ++++++++++++++++++++++ .../cli/src/ui/hooks/slashCommandProcessor.test.ts | 3 + packages/cli/src/ui/hooks/slashCommandProcessor.ts | 9 +++ packages/cli/src/ui/hooks/useAuthCommand.ts | 57 +++++++++++++ packages/cli/src/ui/hooks/useGeminiStream.test.tsx | 1 + packages/cli/src/ui/hooks/useGeminiStream.ts | 7 +- 9 files changed, 266 insertions(+), 7 deletions(-) create mode 100644 packages/cli/src/ui/components/AuthDialog.test.tsx create mode 100644 packages/cli/src/ui/components/AuthDialog.tsx create mode 100644 packages/cli/src/ui/hooks/useAuthCommand.ts (limited to 'packages/cli/src/ui') diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 0ebaa34d..dca24b5c 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -145,6 +145,15 @@ vi.mock('./hooks/useGeminiStream', () => ({ })), })); +vi.mock('./hooks/useAuthCommand', () => ({ + useAuthCommand: vi.fn(() => ({ + isAuthDialogOpen: false, + openAuthDialog: vi.fn(), + handleAuthSelect: vi.fn(), + handleAuthHighlight: vi.fn(), + })), +})); + vi.mock('./hooks/useLogger', () => ({ useLogger: vi.fn(() => ({ getPreviousUserMessages: vi.fn().mockResolvedValue([]), @@ -176,7 +185,9 @@ describe('App UI', () => { }; const workspaceSettingsFile: SettingsFile = { path: '/workspace/.gemini/settings.json', - settings, + settings: { + ...settings, + }, }; return new LoadedSettings(userSettingsFile, workspaceSettingsFile, []); }; @@ -184,10 +195,6 @@ describe('App UI', () => { beforeEach(() => { const ServerConfigMocked = vi.mocked(ServerConfig, true); mockConfig = new ServerConfigMocked({ - contentGeneratorConfig: { - apiKey: 'test-key', - model: 'test-model', - }, embeddingModel: 'test-embedding-model', sandbox: undefined, targetDir: '/test/dir', @@ -197,7 +204,7 @@ describe('App UI', () => { showMemoryUsage: false, sessionId: 'test-session-id', cwd: '/tmp', - // Provide other required fields for ConfigParameters if necessary + model: 'model', }) as unknown as MockServerConfig; // Ensure the getShowMemoryUsage mock function is specifically set up if not covered by constructor mock diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index a9c5f0e7..c481ebd3 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -20,6 +20,7 @@ import { useTerminalSize } from './hooks/useTerminalSize.js'; import { useGeminiStream } from './hooks/useGeminiStream.js'; import { useLoadingIndicator } from './hooks/useLoadingIndicator.js'; import { useThemeCommand } from './hooks/useThemeCommand.js'; +import { useAuthCommand } from './hooks/useAuthCommand.js'; import { useEditorSettings } from './hooks/useEditorSettings.js'; import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js'; import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js'; @@ -31,6 +32,7 @@ import { ShellModeIndicator } from './components/ShellModeIndicator.js'; import { InputPrompt } from './components/InputPrompt.js'; import { Footer } from './components/Footer.js'; import { ThemeDialog } from './components/ThemeDialog.js'; +import { AuthDialog } from './components/AuthDialog.js'; import { EditorSettingsDialog } from './components/EditorSettingsDialog.js'; import { Colors } from './colors.js'; import { Help } from './components/Help.js'; @@ -51,6 +53,7 @@ import { isEditorAvailable, EditorType, } from '@gemini-cli/core'; +import { validateAuthMethod } from '../config/auth.js'; import { useLogger } from './hooks/useLogger.js'; import { StreamingContext } from './contexts/StreamingContext.js'; import { @@ -101,6 +104,7 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => { const [debugMessage, setDebugMessage] = useState(''); const [showHelp, setShowHelp] = useState(false); const [themeError, setThemeError] = useState(null); + const [authError, setAuthError] = useState(null); const [editorError, setEditorError] = useState(null); const [footerHeight, setFooterHeight] = useState(0); const [corgiMode, setCorgiMode] = useState(false); @@ -129,6 +133,23 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => { handleThemeHighlight, } = useThemeCommand(settings, setThemeError, addItem); + const { + isAuthDialogOpen, + openAuthDialog, + handleAuthSelect, + handleAuthHighlight, + } = useAuthCommand(settings, setAuthError, config); + + useEffect(() => { + if (settings.merged.selectedAuthType) { + const error = validateAuthMethod(settings.merged.selectedAuthType); + if (error) { + setAuthError(error); + openAuthDialog(); + } + } + }, [settings.merged.selectedAuthType, openAuthDialog, setAuthError]); + const { isEditorDialogOpen, openEditorDialog, @@ -197,6 +218,7 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => { setShowHelp, setDebugMessage, openThemeDialog, + openAuthDialog, openEditorDialog, performMemoryRefresh, toggleCorgiMode, @@ -306,6 +328,11 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => { return editorType as EditorType; }, [settings, openEditorDialog]); + const onAuthError = useCallback(() => { + setAuthError('reauth required'); + openAuthDialog(); + }, [openAuthDialog, setAuthError]); + const { streamingState, submitQuery, @@ -322,6 +349,7 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => { handleSlashCommand, shellModeActive, getPreferredEditor, + onAuthError, ); pendingHistoryItems.push(...pendingGeminiHistoryItems); const { elapsedTime, currentLoadingPhrase } = @@ -557,6 +585,20 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => { terminalWidth={mainAreaWidth} /> + ) : isAuthDialogOpen ? ( + + {authError && ( + + {authError} + + )} + + ) : isEditorDialogOpen ? ( {editorError && ( diff --git a/packages/cli/src/ui/components/AuthDialog.test.tsx b/packages/cli/src/ui/components/AuthDialog.test.tsx new file mode 100644 index 00000000..a5f46d93 --- /dev/null +++ b/packages/cli/src/ui/components/AuthDialog.test.tsx @@ -0,0 +1,41 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { render } from 'ink-testing-library'; +import { AuthDialog } from './AuthDialog.js'; +import { LoadedSettings } from '../../config/settings.js'; +import { AuthType } from '@gemini-cli/core'; + +describe('AuthDialog', () => { + it('should show an error if the initial auth type is invalid', () => { + const settings: LoadedSettings = new LoadedSettings( + { + settings: { + selectedAuthType: AuthType.USE_GEMINI, + }, + path: '', + }, + { + settings: {}, + path: '', + }, + [], + ); + + const { lastFrame } = render( + {}} + onHighlight={() => {}} + settings={settings} + initialErrorMessage="GEMINI_API_KEY environment variable not found" + />, + ); + + expect(lastFrame()).toContain( + 'GEMINI_API_KEY environment variable not found', + ); + }); +}); diff --git a/packages/cli/src/ui/components/AuthDialog.tsx b/packages/cli/src/ui/components/AuthDialog.tsx new file mode 100644 index 00000000..b16529cf --- /dev/null +++ b/packages/cli/src/ui/components/AuthDialog.tsx @@ -0,0 +1,94 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import React, { useState } from 'react'; +import { Box, Text, useInput } from 'ink'; +import { Colors } from '../colors.js'; +import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; +import { LoadedSettings, SettingScope } from '../../config/settings.js'; +import { AuthType } from '@gemini-cli/core'; +import { validateAuthMethod } from '../../config/auth.js'; + +interface AuthDialogProps { + onSelect: (authMethod: string | undefined, scope: SettingScope) => void; + onHighlight: (authMethod: string | undefined) => void; + settings: LoadedSettings; + initialErrorMessage?: string | null; +} + +export function AuthDialog({ + onSelect, + onHighlight, + settings, + initialErrorMessage, +}: AuthDialogProps): React.JSX.Element { + const [errorMessage, setErrorMessage] = useState( + initialErrorMessage || null, + ); + const authItems = [ + { + label: 'Login with Google Personal Account', + value: AuthType.LOGIN_WITH_GOOGLE_PERSONAL, + }, + { label: 'Gemini API Key', value: AuthType.USE_GEMINI }, + { + label: 'Login with GCP Project and Google Work Account', + value: AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE, + }, + { label: 'Vertex AI', value: AuthType.USE_VERTEX_AI }, + ]; + + let initialAuthIndex = authItems.findIndex( + (item) => item.value === settings.merged.selectedAuthType, + ); + + if (initialAuthIndex === -1) { + initialAuthIndex = 0; + } + + const handleAuthSelect = (authMethod: string) => { + const error = validateAuthMethod(authMethod); + if (error) { + setErrorMessage(error); + } else { + setErrorMessage(null); + onSelect(authMethod, SettingScope.User); + } + }; + + useInput((_input, key) => { + if (key.escape) { + onSelect(undefined, SettingScope.User); + } + }); + + return ( + + Select Auth Method + + {errorMessage && ( + + {errorMessage} + + )} + + (Use Enter to select) + + + ); +} diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 7c750af1..04931c7f 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -103,6 +103,7 @@ describe('useSlashCommandProcessor', () => { let mockSetShowHelp: ReturnType; let mockOnDebugMessage: ReturnType; let mockOpenThemeDialog: ReturnType; + let mockOpenAuthDialog: ReturnType; let mockOpenEditorDialog: ReturnType; let mockPerformMemoryRefresh: ReturnType; let mockSetQuittingMessages: ReturnType; @@ -120,6 +121,7 @@ describe('useSlashCommandProcessor', () => { mockSetShowHelp = vi.fn(); mockOnDebugMessage = vi.fn(); mockOpenThemeDialog = vi.fn(); + mockOpenAuthDialog = vi.fn(); mockOpenEditorDialog = vi.fn(); mockPerformMemoryRefresh = vi.fn().mockResolvedValue(undefined); mockSetQuittingMessages = vi.fn(); @@ -171,6 +173,7 @@ describe('useSlashCommandProcessor', () => { mockSetShowHelp, mockOnDebugMessage, mockOpenThemeDialog, + mockOpenAuthDialog, mockOpenEditorDialog, mockPerformMemoryRefresh, mockCorgiMode, diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 0e622f23..ee7b55cb 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -68,6 +68,7 @@ export const useSlashCommandProcessor = ( setShowHelp: React.Dispatch>, onDebugMessage: (message: string) => void, openThemeDialog: () => void, + openAuthDialog: () => void, openEditorDialog: () => void, performMemoryRefresh: () => Promise, toggleCorgiMode: () => void, @@ -197,6 +198,13 @@ export const useSlashCommandProcessor = ( openThemeDialog(); }, }, + { + name: 'auth', + description: 'change the auth method', + action: (_mainCommand, _subCommand, _args) => { + openAuthDialog(); + }, + }, { name: 'editor', description: 'set external editor preference', @@ -907,6 +915,7 @@ Add any other context about the problem here. setShowHelp, refreshStatic, openThemeDialog, + openAuthDialog, openEditorDialog, clearItems, performMemoryRefresh, diff --git a/packages/cli/src/ui/hooks/useAuthCommand.ts b/packages/cli/src/ui/hooks/useAuthCommand.ts new file mode 100644 index 00000000..a9b1cb1e --- /dev/null +++ b/packages/cli/src/ui/hooks/useAuthCommand.ts @@ -0,0 +1,57 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useState, useCallback, useEffect } from 'react'; +import { LoadedSettings, SettingScope } from '../../config/settings.js'; +import { AuthType, Config, clearCachedCredentialFile } from '@gemini-cli/core'; + +async function performAuthFlow(authMethod: AuthType, config: Config) { + await config.refreshAuth(authMethod); + console.log(`Authenticated via "${authMethod}".`); +} + +export const useAuthCommand = ( + settings: LoadedSettings, + setAuthError: (error: string | null) => void, + config: Config, +) => { + const [isAuthDialogOpen, setIsAuthDialogOpen] = useState( + settings.merged.selectedAuthType === undefined, + ); + + useEffect(() => { + if (!isAuthDialogOpen) { + performAuthFlow(settings.merged.selectedAuthType as AuthType, config); + } + }, [isAuthDialogOpen, settings, config]); + + const openAuthDialog = useCallback(() => { + setIsAuthDialogOpen(true); + }, []); + + const handleAuthSelect = useCallback( + async (authMethod: string | undefined, scope: SettingScope) => { + if (authMethod) { + await clearCachedCredentialFile(); + settings.setValue(scope, 'selectedAuthType', authMethod); + } + setIsAuthDialogOpen(false); + setAuthError(null); + }, + [settings, setAuthError], + ); + + const handleAuthHighlight = useCallback((_authMethod: string | undefined) => { + // For now, we don't do anything on highlight. + }, []); + + return { + isAuthDialogOpen, + openAuthDialog, + handleAuthSelect, + handleAuthHighlight, + }; +}; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 96dd6aef..36f420e4 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -359,6 +359,7 @@ describe('useGeminiStream', () => { props.handleSlashCommand, props.shellModeActive, () => 'vscode' as EditorType, + () => {}, ), { initialProps: { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 6d92af0d..4049c884 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -22,6 +22,7 @@ import { GitService, EditorType, ThoughtSummary, + isAuthError, } from '@gemini-cli/core'; import { type Part, type PartListUnion } from '@google/genai'; import { @@ -87,6 +88,7 @@ export const useGeminiStream = ( >, shellModeActive: boolean, getPreferredEditor: () => EditorType | undefined, + onAuthError: () => void, ) => { const [initError, setInitError] = useState(null); const abortControllerRef = useRef(null); @@ -496,7 +498,9 @@ export const useGeminiStream = ( setPendingHistoryItem(null); } } catch (error: unknown) { - if (!isNodeError(error) || error.name !== 'AbortError') { + if (isAuthError(error)) { + onAuthError(); + } else if (!isNodeError(error) || error.name !== 'AbortError') { addItem( { type: MessageType.ERROR, @@ -522,6 +526,7 @@ export const useGeminiStream = ( setInitError, geminiClient, startNewTurn, + onAuthError, ], ); -- cgit v1.2.3