abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-20 08:47:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Auth First Run (#1207)

Browse Source 
Co-authored-by: Tommaso Sciortino <sciortino@gmail.com>
Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
This commit is contained in:
matt korwel
2025-06-19 16:52:22 -07:00
committed by GitHub
parent c48fcaa8c3
commit 04518b52c0
37 changed files with 636 additions and 349 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
packages/cli/index.ts
View File

@@ -13,7 +13,7 @@ import { main } from './src/gemini.js';
main().catch((error) => {
console.error('An unexpected critical error occurred:');
if (error instanceof Error) {
console.error(error.message);
console.error(error.stack);
} else {
console.error(String(error));
}

44
packages/cli/src/config/auth.ts Normal file
View File

@@ -0,0 +1,44 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { AuthType } from '@gemini-cli/core';
import { loadEnvironment } from './config.js';
export const validateAuthMethod = (authMethod: string): string | null => {
loadEnvironment();
if (authMethod === AuthType.LOGIN_WITH_GOOGLE_PERSONAL) {
return null;
}
if (authMethod === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE) {
if (!process.env.GOOGLE_CLOUD_PROJECT) {
return 'GOOGLE_CLOUD_PROJECT environment variable not found. Add that to your .env and try again, no reload needed!';
}
return null;
}
if (authMethod === AuthType.USE_GEMINI) {
if (!process.env.GEMINI_API_KEY) {
return 'GEMINI_API_KEY environment variable not found. Add that to your .env and try again, no reload needed!';
}
return null;
}
if (authMethod === AuthType.USE_VERTEX_AI) {
if (!process.env.GOOGLE_API_KEY) {
return 'GOOGLE_API_KEY environment variable not found. Add that to your .env and try again, no reload needed!';
}
if (!process.env.GOOGLE_CLOUD_PROJECT) {
return 'GOOGLE_CLOUD_PROJECT environment variable not found. Add that to your .env and try again, no reload needed!';
}
if (!process.env.GOOGLE_CLOUD_LOCATION) {
return 'GOOGLE_CLOUD_LOCATION environment variable not found. Add that to your .env and try again, no reload needed!';
}
return null;
}
return 'Invalid auth method selected.';
};

42
packages/cli/src/config/config.test.ts
View File

@@ -247,48 +247,6 @@ describe('loadCliConfig telemetry', () => {
});
});
describe('API Key Handling', () => {
const originalEnv = { ...process.env };
const originalArgv = process.argv;
beforeEach(() => {
vi.resetAllMocks();
process.argv = ['node', 'script.js'];
});
afterEach(() => {
process.env = originalEnv;
process.argv = originalArgv;
});
it('should use GEMINI_API_KEY from env', async () => {
process.env.GEMINI_API_KEY = 'gemini-key';
delete process.env.GOOGLE_API_KEY;
const settings: Settings = {};
const result = await loadCliConfig(settings, [], 'test-session');
expect(result.getContentGeneratorConfig().apiKey).toBe('gemini-key');
});
it('should use GOOGLE_API_KEY and warn when both GOOGLE_API_KEY and GEMINI_API_KEY are set', async () => {
const consoleWarnSpy = vi
.spyOn(console, 'warn')
.mockImplementation(() => {});
process.env.GEMINI_API_KEY = 'gemini-key';
process.env.GOOGLE_API_KEY = 'google-key';
const settings: Settings = {};
const result = await loadCliConfig(settings, [], 'test-session');
expect(consoleWarnSpy).toHaveBeenCalledWith(
'[WARN]',
'Both GEMINI_API_KEY and GOOGLE_API_KEY are set. Using GOOGLE_API_KEY.',
);
expect(result.getContentGeneratorConfig().apiKey).toBe('google-key');
});
});
describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => {
beforeEach(() => {
vi.resetAllMocks();

60
packages/cli/src/config/config.ts
View File

@@ -13,7 +13,6 @@ import {
setGeminiMdFilename as setServerGeminiMdFilename,
getCurrentGeminiMdFilename,
ApprovalMode,
ContentGeneratorConfig,
GEMINI_CONFIG_DIR as GEMINI_DIR,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_EMBEDDING_MODEL,
@@ -21,7 +20,7 @@ import {
TelemetryTarget,
} from '@gemini-cli/core';
import { Settings } from './settings.js';
import { getEffectiveModel } from '../utils/modelCheck.js';
import { Extension } from './extension.js';
import { getCliVersion } from '../utils/version.js';
import * as dotenv from 'dotenv';
@@ -194,15 +193,12 @@ export async function loadCliConfig(
extensionContextFilePaths,
);
const contentGeneratorConfig = await createContentGeneratorConfig(argv);
const mcpServers = mergeMcpServers(settings, extensions);
const sandboxConfig = await loadSandboxConfig(settings, argv);
return new Config({
sessionId,
contentGeneratorConfig,
embeddingModel: DEFAULT_GEMINI_EMBEDDING_MODEL,
sandbox: sandboxConfig,
targetDir: process.cwd(),
@@ -242,6 +238,7 @@ export async function loadCliConfig(
cwd: process.cwd(),
fileDiscoveryService: fileService,
bugCommand: settings.bugCommand,
model: argv.model!,
});
}
@@ -262,59 +259,6 @@ function mergeMcpServers(settings: Settings, extensions: Extension[]) {
}
return mcpServers;
}
async function createContentGeneratorConfig(
argv: CliArgs,
): Promise<ContentGeneratorConfig> {
const geminiApiKey = process.env.GEMINI_API_KEY;
const googleApiKey = process.env.GOOGLE_API_KEY;
const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT;
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION;
const hasCodeAssist = process.env.GEMINI_CODE_ASSIST === 'true';
const hasGeminiApiKey = !!geminiApiKey;
const hasGoogleApiKey = !!googleApiKey;
const hasVertexProjectLocationConfig =
!!googleCloudProject && !!googleCloudLocation;
if (hasGeminiApiKey && hasGoogleApiKey) {
logger.warn(
'Both GEMINI_API_KEY and GOOGLE_API_KEY are set. Using GOOGLE_API_KEY.',
);
}
if (
!hasCodeAssist &&
!hasGeminiApiKey &&
!hasGoogleApiKey &&
!hasVertexProjectLocationConfig
) {
logger.error(
'No valid API authentication configuration found. Please set ONE of the following combinations in your environment variables or .env file:\n' +
'1. GEMINI_CODE_ASSIST=true (for Code Assist access).\n' +
'2. GEMINI_API_KEY (for Gemini API access).\n' +
'3. GOOGLE_API_KEY (for Gemini API or Vertex AI Express Mode access).\n' +
'4. GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION (for Vertex AI access).\n\n' +
'For Gemini API keys, visit: https://ai.google.dev/gemini-api/docs/api-key\n' +
'For Vertex AI authentication, visit: https://cloud.google.com/vertex-ai/docs/authentication\n' +
'The GOOGLE_GENAI_USE_VERTEXAI environment variable can also be set to true/false to influence service selection when ambiguity exists.',
);
process.exit(1);
}
const config: ContentGeneratorConfig = {
model: argv.model || DEFAULT_GEMINI_MODEL,
apiKey: googleApiKey || geminiApiKey || '',
vertexai: hasGeminiApiKey ? false : undefined,
codeAssist: hasCodeAssist,
};
if (config.apiKey) {
config.model = await getEffectiveModel(config.apiKey, config.model);
}
return config;
}
function findEnvFile(startDir: string): string | null {
let currentDir = path.resolve(startDir);
while (true) {

2
packages/cli/src/config/settings.ts
View File

@@ -12,6 +12,7 @@ import {
getErrorMessage,
BugCommandSettings,
TelemetrySettings,
AuthType,
} from '@gemini-cli/core';
import stripJsonComments from 'strip-json-comments';
import { DefaultLight } from '../ui/themes/default-light.js';
@@ -32,6 +33,7 @@ export interface AccessibilitySettings {
export interface Settings {
theme?: string;
selectedAuthType?: AuthType;
sandbox?: boolean | string;
coreTools?: string[];
excludeTools?: string[];

87
packages/cli/src/gemini.tsx
View File

@@ -25,7 +25,9 @@ import {
WriteFileTool,
sessionId,
logUserPrompt,
AuthType,
} from '@gemini-cli/core';
import { validateAuthMethod } from './config/auth.js';
export async function main() {
const workspaceRoot = process.cwd();
@@ -47,10 +49,6 @@ export async function main() {
const extensions = loadExtensions(workspaceRoot);
const config = await loadCliConfig(settings.merged, extensions, sessionId);
// When using Code Assist this triggers the Oauth login.
// Do this now, before sandboxing, so web redirect works.
await config.getGeminiClient().initialize();
// Initialize centralized FileDiscoveryService
config.getFileService();
if (config.getCheckpointEnabled()) {
@@ -73,6 +71,15 @@ export async function main() {
if (!process.env.SANDBOX) {
const sandboxConfig = config.getSandbox();
if (sandboxConfig) {
if (settings.merged.selectedAuthType) {
// Validate authentication here because the sandbox will interfere with the Oauth2 web redirect.
const err = validateAuthMethod(settings.merged.selectedAuthType);
if (err) {
console.error(err);
process.exit(1);
}
await config.refreshAuth(settings.merged.selectedAuthType);
}
await start_sandbox(sandboxConfig);
process.exit(0);
}
@@ -152,28 +159,58 @@ async function loadNonInteractiveConfig(
extensions: Extension[],
settings: LoadedSettings,
) {
if (config.getApprovalMode() === ApprovalMode.YOLO) {
// Since everything is being allowed we can use normal yolo behavior.
return config;
let finalConfig = config;
if (config.getApprovalMode() !== ApprovalMode.YOLO) {
// Everything is not allowed, ensure that only read-only tools are configured.
const existingExcludeTools = settings.merged.excludeTools || [];
const interactiveTools = [
ShellTool.Name,
EditTool.Name,
WriteFileTool.Name,
];
const newExcludeTools = [
...new Set([...existingExcludeTools, ...interactiveTools]),
];
const nonInteractiveSettings = {
...settings.merged,
excludeTools: newExcludeTools,
};
finalConfig = await loadCliConfig(
nonInteractiveSettings,
extensions,
config.getSessionId(),
);
}
// Everything is not allowed, ensure that only read-only tools are configured.
const existingExcludeTools = settings.merged.excludeTools || [];
const interactiveTools = [ShellTool.Name, EditTool.Name, WriteFileTool.Name];
const newExcludeTools = [
...new Set([...existingExcludeTools, ...interactiveTools]),
];
const nonInteractiveSettings = {
...settings.merged,
excludeTools: newExcludeTools,
};
const newConfig = await loadCliConfig(
nonInteractiveSettings,
extensions,
config.getSessionId(),
return await validateNonInterActiveAuth(
settings.merged.selectedAuthType,
finalConfig,
);
await newConfig.getGeminiClient().initialize();
return newConfig;
}
async function validateNonInterActiveAuth(
selectedAuthType: AuthType | undefined,
nonInteractiveConfig: Config,
) {
// making a special case for the cli. many headless environments might not have a settings.json set
// so if GEMINI_API_KEY is set, we'll use that. However since the oauth things are interactive anyway, we'll
// still expect that exists
if (!selectedAuthType && !process.env.GEMINI_API_KEY) {
console.error(
'Please set an Auth method in your .gemini/settings.json OR specify GEMINI_API_KEY env variable file before running',
);
process.exit(1);
}
selectedAuthType = selectedAuthType || AuthType.USE_GEMINI;
const err = validateAuthMethod(selectedAuthType);
if (err != null) {
console.error(err);
process.exit(1);
}
await nonInteractiveConfig.refreshAuth(selectedAuthType);
return nonInteractiveConfig;
}

19
packages/cli/src/ui/App.test.tsx
View File

@@ -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

42
packages/cli/src/ui/App.tsx
View File

@@ -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<string>('');
const [showHelp, setShowHelp] = useState<boolean>(false);
const [themeError, setThemeError] = useState<string | null>(null);
const [authError, setAuthError] = useState<string | null>(null);
const [editorError, setEditorError] = useState<string | null>(null);
const [footerHeight, setFooterHeight] = useState<number>(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}
/>
</Box>
) : isAuthDialogOpen ? (
<Box flexDirection="column">
{authError && (
<Box marginBottom={1}>
<Text color={Colors.AccentRed}>{authError}</Text>
</Box>
)}
<AuthDialog
onSelect={handleAuthSelect}
onHighlight={handleAuthHighlight}
settings={settings}
initialErrorMessage={authError}
/>
</Box>
) : isEditorDialogOpen ? (
<Box flexDirection="column">
{editorError && (

41
packages/cli/src/ui/components/AuthDialog.test.tsx Normal file
View File

@@ -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(
<AuthDialog
onSelect={() => {}}
onHighlight={() => {}}
settings={settings}
initialErrorMessage="GEMINI_API_KEY environment variable not found"
/>,
);
expect(lastFrame()).toContain(
'GEMINI_API_KEY environment variable not found',
);
});
});

94
packages/cli/src/ui/components/AuthDialog.tsx Normal file
View File

@@ -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<string | null>(
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 (
<Box
borderStyle="round"
borderColor={Colors.Gray}
flexDirection="column"
padding={1}
width="100%"
>
<Text bold>Select Auth Method</Text>
<RadioButtonSelect
items={authItems}
initialIndex={initialAuthIndex}
onSelect={handleAuthSelect}
onHighlight={onHighlight}
isFocused={true}
/>
{errorMessage && (
<Box marginTop={1}>
<Text color={Colors.AccentRed}>{errorMessage}</Text>
</Box>
)}
<Box marginTop={1}>
<Text color={Colors.Gray}>(Use Enter to select)</Text>
</Box>
</Box>
);
}

3
packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
View File

@@ -103,6 +103,7 @@ describe('useSlashCommandProcessor', () => {
let mockSetShowHelp: ReturnType<typeof vi.fn>;
let mockOnDebugMessage: ReturnType<typeof vi.fn>;
let mockOpenThemeDialog: ReturnType<typeof vi.fn>;
let mockOpenAuthDialog: ReturnType<typeof vi.fn>;
let mockOpenEditorDialog: ReturnType<typeof vi.fn>;
let mockPerformMemoryRefresh: ReturnType<typeof vi.fn>;
let mockSetQuittingMessages: ReturnType<typeof vi.fn>;
@@ -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,

9
packages/cli/src/ui/hooks/slashCommandProcessor.ts
View File

@@ -68,6 +68,7 @@ export const useSlashCommandProcessor = (
setShowHelp: React.Dispatch<React.SetStateAction<boolean>>,
onDebugMessage: (message: string) => void,
openThemeDialog: () => void,
openAuthDialog: () => void,
openEditorDialog: () => void,
performMemoryRefresh: () => Promise<void>,
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,

57
packages/cli/src/ui/hooks/useAuthCommand.ts Normal file
View File

@@ -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,
};
};

1
packages/cli/src/ui/hooks/useGeminiStream.test.tsx
View File

@@ -359,6 +359,7 @@ describe('useGeminiStream', () => {
props.handleSlashCommand,
props.shellModeActive,
() => 'vscode' as EditorType,
() => {},
),
{
initialProps: {

7
packages/cli/src/ui/hooks/useGeminiStream.ts
View File

@@ -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<string | null>(null);
const abortControllerRef = useRef<AbortController | null>(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,
],
);

154
packages/cli/src/utils/modelCheck.test.ts
View File

@@ -1,154 +0,0 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { getEffectiveModel } from './modelCheck.js';
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
} from '@gemini-cli/core';
// Mock global fetch
global.fetch = vi.fn();
// Mock AbortController
const mockAbort = vi.fn();
global.AbortController = vi.fn(() => ({
signal: { aborted: false }, // Start with not aborted
abort: mockAbort,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
})) as any;
describe('getEffectiveModel', () => {
const apiKey = 'test-api-key';
beforeEach(() => {
vi.useFakeTimers();
vi.clearAllMocks();
// Reset signal for each test if AbortController mock is more complex
global.AbortController = vi.fn(() => ({
signal: { aborted: false },
abort: mockAbort,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
})) as any;
});
afterEach(() => {
vi.restoreAllMocks();
vi.useRealTimers();
});
describe('when currentConfiguredModel is not DEFAULT_GEMINI_MODEL', () => {
it('should return the currentConfiguredModel without fetching', async () => {
const customModel = 'custom-model-name';
const result = await getEffectiveModel(apiKey, customModel);
expect(result).toEqual(customModel);
expect(fetch).not.toHaveBeenCalled();
});
});
describe('when currentConfiguredModel is DEFAULT_GEMINI_MODEL', () => {
it('should switch to DEFAULT_GEMINI_FLASH_MODEL if fetch returns 429', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({
ok: false,
status: 429,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_FLASH_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
expect(fetch).toHaveBeenCalledWith(
`https://generativelanguage.googleapis.com/v1beta/models/${DEFAULT_GEMINI_MODEL}:generateContent?key=${apiKey}`,
expect.any(Object),
);
});
it('should return DEFAULT_GEMINI_MODEL if fetch returns 200', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({
ok: true,
status: 200,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch returns a non-429 error status (e.g., 500)', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({
ok: false,
status: 500,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch throws a network error', async () => {
(fetch as vi.Mock).mockRejectedValueOnce(new Error('Network error'));
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch times out', async () => {
// Simulate AbortController's signal changing and fetch throwing AbortError
const abortControllerInstance = {
signal: { aborted: false }, // mutable signal
abort: vi.fn(() => {
abortControllerInstance.signal.aborted = true; // Use abortControllerInstance
}),
};
(global.AbortController as vi.Mock).mockImplementationOnce(
() => abortControllerInstance,
);
(fetch as vi.Mock).mockImplementationOnce(
async ({ signal }: { signal: AbortSignal }) => {
// Simulate the timeout advancing and abort being called
vi.advanceTimersByTime(2000);
if (signal.aborted) {
throw new DOMException('Aborted', 'AbortError');
}
// Should not reach here in a timeout scenario
return { ok: true, status: 200 };
},
);
const resultPromise = getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
// Ensure timers are advanced to trigger the timeout within getEffectiveModel
await vi.advanceTimersToNextTimerAsync(); // Or advanceTimersByTime(2000) if more precise control is needed
const result = await resultPromise;
expect(mockAbort).toHaveBeenCalledTimes(0); // setTimeout calls controller.abort(), not our direct mockAbort
expect(abortControllerInstance.abort).toHaveBeenCalledTimes(1);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should correctly pass API key and model in the fetch request', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({ ok: true, status: 200 });
const specificApiKey = 'specific-key-for-this-test';
await getEffectiveModel(specificApiKey, DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledWith(
`https://generativelanguage.googleapis.com/v1beta/models/${DEFAULT_GEMINI_MODEL}:generateContent?key=${specificApiKey}`,
expect.objectContaining({
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
contents: [{ parts: [{ text: 'test' }] }],
generationConfig: {
maxOutputTokens: 1,
temperature: 0,
topK: 1,
thinkingConfig: { thinkingBudget: 0, includeThoughts: false },
},
}),
}),
);
});
});
});

68
packages/cli/src/utils/modelCheck.ts
View File

@@ -1,68 +0,0 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
} from '@gemini-cli/core';
/**
* Checks if the default "pro" model is rate-limited and returns a fallback "flash"
* model if necessary. This function is designed to be silent.
* @param apiKey The API key to use for the check.
* @param currentConfiguredModel The model currently configured in settings.
* @returns An object indicating the model to use, whether a switch occurred,
* and the original model if a switch happened.
*/
export async function getEffectiveModel(
apiKey: string,
currentConfiguredModel: string,
): Promise<string> {
if (currentConfiguredModel !== DEFAULT_GEMINI_MODEL) {
// Only check if the user is trying to use the specific pro model we want to fallback from.
return currentConfiguredModel;
}
const modelToTest = DEFAULT_GEMINI_MODEL;
const fallbackModel = DEFAULT_GEMINI_FLASH_MODEL;
const endpoint = `https://generativelanguage.googleapis.com/v1beta/models/${modelToTest}:generateContent?key=${apiKey}`;
const body = JSON.stringify({
contents: [{ parts: [{ text: 'test' }] }],
generationConfig: {
maxOutputTokens: 1,
temperature: 0,
topK: 1,
thinkingConfig: { thinkingBudget: 0, includeThoughts: false },
},
});
const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), 2000); // 500ms timeout for the request
try {
const response = await fetch(endpoint, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body,
signal: controller.signal,
});
clearTimeout(timeoutId);
if (response.status === 429) {
console.log(
`[INFO] Your configured model (${modelToTest}) was temporarily unavailable. Switched to ${fallbackModel} for this session.`,
);
return fallbackModel;
}
// For any other case (success, other error codes), we stick to the original model.
return currentConfiguredModel;
} catch (_error) {
clearTimeout(timeoutId);
// On timeout or any other fetch error, stick to the original model.
return currentConfiguredModel;
}
}