Run model availability check in the background to speed up startup (#4256)

This commit is contained in:
Sandy Tao
2025-07-15 21:13:30 -07:00
committed by GitHub
parent d622e596a1
commit cba272082d
5 changed files with 68 additions and 53 deletions

View File

@@ -57,6 +57,7 @@ import {
EditorType, EditorType,
FlashFallbackEvent, FlashFallbackEvent,
logFlashFallback, logFlashFallback,
AuthType,
} from '@google/gemini-cli-core'; } from '@google/gemini-cli-core';
import { validateAuthMethod } from '../config/auth.js'; import { validateAuthMethod } from '../config/auth.js';
import { useLogger } from './hooks/useLogger.js'; import { useLogger } from './hooks/useLogger.js';
@@ -294,64 +295,70 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
): Promise<boolean> => { ): Promise<boolean> => {
let message: string; let message: string;
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default) if (
const isPaidTier = config.getContentGeneratorConfig().authType ===
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD; AuthType.LOGIN_WITH_GOOGLE
) {
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default)
const isPaidTier =
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
// Check if this is a Pro quota exceeded error // Check if this is a Pro quota exceeded error
if (error && isProQuotaExceededError(error)) { if (error && isProQuotaExceededError(error)) {
if (isPaidTier) { if (isPaidTier) {
message = `⚡ You have reached your daily ${currentModel} quota limit. message = `⚡ You have reached your daily ${currentModel} quota limit.
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session. ⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`; ⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
} else { } else {
message = `⚡ You have reached your daily ${currentModel} quota limit. message = `⚡ You have reached your daily ${currentModel} quota limit.
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session. ⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist ⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key ⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
⚡ You can switch authentication methods by typing /auth`; ⚡ You can switch authentication methods by typing /auth`;
} }
} else if (error && isGenericQuotaExceededError(error)) { } else if (error && isGenericQuotaExceededError(error)) {
if (isPaidTier) { if (isPaidTier) {
message = `⚡ You have reached your daily quota limit. message = `⚡ You have reached your daily quota limit.
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session. ⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`; ⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
} else { } else {
message = `⚡ You have reached your daily quota limit. message = `⚡ You have reached your daily quota limit.
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session. ⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist ⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key ⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
⚡ You can switch authentication methods by typing /auth`; ⚡ You can switch authentication methods by typing /auth`;
} }
} else { } else {
if (isPaidTier) { if (isPaidTier) {
// Default fallback message for other cases (like consecutive 429s) // Default fallback message for other cases (like consecutive 429s)
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session. message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit ⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`; ⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
} else { } else {
// Default fallback message for other cases (like consecutive 429s) // Default fallback message for other cases (like consecutive 429s)
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session. message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit ⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist ⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key ⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
⚡ You can switch authentication methods by typing /auth`; ⚡ You can switch authentication methods by typing /auth`;
}
} }
// Add message to UI history
addItem(
{
type: MessageType.INFO,
text: message,
},
Date.now(),
);
// Set the flag to prevent tool continuation
setModelSwitchedFromQuotaError(true);
// Set global quota error flag to prevent Flash model calls
config.setQuotaErrorOccurred(true);
} }
// Add message to UI history
addItem(
{
type: MessageType.INFO,
text: message,
},
Date.now(),
);
// Set the flag to prevent tool continuation
setModelSwitchedFromQuotaError(true);
// Set global quota error flag to prevent Flash model calls
config.setQuotaErrorOccurred(true);
// Switch model for future use but return false to stop current retry // Switch model for future use but return false to stop current retry
config.setModel(fallbackModel); config.setModel(fallbackModel);
logFlashFallback( logFlashFallback(

View File

@@ -151,14 +151,12 @@ describe('Server Config (config.ts)', () => {
apiKey: 'test-key', apiKey: 'test-key',
}; };
(createContentGeneratorConfig as Mock).mockResolvedValue( (createContentGeneratorConfig as Mock).mockReturnValue(mockContentConfig);
mockContentConfig,
);
await config.refreshAuth(authType); await config.refreshAuth(authType);
expect(createContentGeneratorConfig).toHaveBeenCalledWith( expect(createContentGeneratorConfig).toHaveBeenCalledWith(
MODEL, // Should be called with the original model 'gemini-pro' config,
authType, authType,
); );
// Verify that contentGeneratorConfig is updated with the new model // Verify that contentGeneratorConfig is updated with the new model

View File

@@ -274,8 +274,8 @@ export class Config {
} }
async refreshAuth(authMethod: AuthType) { async refreshAuth(authMethod: AuthType) {
this.contentGeneratorConfig = await createContentGeneratorConfig( this.contentGeneratorConfig = createContentGeneratorConfig(
this.model, this,
authMethod, authMethod,
); );

View File

@@ -64,12 +64,18 @@ describe('createContentGenerator', () => {
describe('createContentGeneratorConfig', () => { describe('createContentGeneratorConfig', () => {
const originalEnv = process.env; const originalEnv = process.env;
const mockConfig = {
getModel: vi.fn().mockReturnValue('gemini-pro'),
setModel: vi.fn(),
flashFallbackHandler: vi.fn(),
} as unknown as Config;
beforeEach(() => { beforeEach(() => {
// Reset modules to re-evaluate imports and environment variables // Reset modules to re-evaluate imports and environment variables
vi.resetModules(); vi.resetModules();
// Restore process.env before each test // Restore process.env before each test
process.env = { ...originalEnv }; process.env = { ...originalEnv };
vi.clearAllMocks();
}); });
afterAll(() => { afterAll(() => {
@@ -80,7 +86,7 @@ describe('createContentGeneratorConfig', () => {
it('should configure for Gemini using GEMINI_API_KEY when set', async () => { it('should configure for Gemini using GEMINI_API_KEY when set', async () => {
process.env.GEMINI_API_KEY = 'env-gemini-key'; process.env.GEMINI_API_KEY = 'env-gemini-key';
const config = await createContentGeneratorConfig( const config = await createContentGeneratorConfig(
undefined, mockConfig,
AuthType.USE_GEMINI, AuthType.USE_GEMINI,
); );
expect(config.apiKey).toBe('env-gemini-key'); expect(config.apiKey).toBe('env-gemini-key');
@@ -90,7 +96,7 @@ describe('createContentGeneratorConfig', () => {
it('should not configure for Gemini if GEMINI_API_KEY is empty', async () => { it('should not configure for Gemini if GEMINI_API_KEY is empty', async () => {
process.env.GEMINI_API_KEY = ''; process.env.GEMINI_API_KEY = '';
const config = await createContentGeneratorConfig( const config = await createContentGeneratorConfig(
undefined, mockConfig,
AuthType.USE_GEMINI, AuthType.USE_GEMINI,
); );
expect(config.apiKey).toBeUndefined(); expect(config.apiKey).toBeUndefined();
@@ -100,7 +106,7 @@ describe('createContentGeneratorConfig', () => {
it('should configure for Vertex AI using GOOGLE_API_KEY when set', async () => { it('should configure for Vertex AI using GOOGLE_API_KEY when set', async () => {
process.env.GOOGLE_API_KEY = 'env-google-key'; process.env.GOOGLE_API_KEY = 'env-google-key';
const config = await createContentGeneratorConfig( const config = await createContentGeneratorConfig(
undefined, mockConfig,
AuthType.USE_VERTEX_AI, AuthType.USE_VERTEX_AI,
); );
expect(config.apiKey).toBe('env-google-key'); expect(config.apiKey).toBe('env-google-key');
@@ -111,7 +117,7 @@ describe('createContentGeneratorConfig', () => {
process.env.GOOGLE_CLOUD_PROJECT = 'env-gcp-project'; process.env.GOOGLE_CLOUD_PROJECT = 'env-gcp-project';
process.env.GOOGLE_CLOUD_LOCATION = 'env-gcp-location'; process.env.GOOGLE_CLOUD_LOCATION = 'env-gcp-location';
const config = await createContentGeneratorConfig( const config = await createContentGeneratorConfig(
undefined, mockConfig,
AuthType.USE_VERTEX_AI, AuthType.USE_VERTEX_AI,
); );
expect(config.vertexai).toBe(true); expect(config.vertexai).toBe(true);
@@ -123,7 +129,7 @@ describe('createContentGeneratorConfig', () => {
process.env.GOOGLE_CLOUD_PROJECT = ''; process.env.GOOGLE_CLOUD_PROJECT = '';
process.env.GOOGLE_CLOUD_LOCATION = ''; process.env.GOOGLE_CLOUD_LOCATION = '';
const config = await createContentGeneratorConfig( const config = await createContentGeneratorConfig(
undefined, mockConfig,
AuthType.USE_VERTEX_AI, AuthType.USE_VERTEX_AI,
); );
expect(config.apiKey).toBeUndefined(); expect(config.apiKey).toBeUndefined();

View File

@@ -52,17 +52,17 @@ export type ContentGeneratorConfig = {
authType?: AuthType | undefined; authType?: AuthType | undefined;
}; };
export async function createContentGeneratorConfig( export function createContentGeneratorConfig(
model: string | undefined, config: Config,
authType: AuthType | undefined, authType: AuthType | undefined,
): Promise<ContentGeneratorConfig> { ): ContentGeneratorConfig {
const geminiApiKey = process.env.GEMINI_API_KEY || undefined; const geminiApiKey = process.env.GEMINI_API_KEY || undefined;
const googleApiKey = process.env.GOOGLE_API_KEY || undefined; const googleApiKey = process.env.GOOGLE_API_KEY || undefined;
const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT || undefined; const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT || undefined;
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION || undefined; const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION || undefined;
// Use runtime model from config if available, otherwise fallback to parameter or default // Use runtime model from config if available, otherwise fallback to parameter or default
const effectiveModel = model || DEFAULT_GEMINI_MODEL; const effectiveModel = config.getModel() || DEFAULT_GEMINI_MODEL;
const contentGeneratorConfig: ContentGeneratorConfig = { const contentGeneratorConfig: ContentGeneratorConfig = {
model: effectiveModel, model: effectiveModel,
@@ -80,10 +80,14 @@ export async function createContentGeneratorConfig(
if (authType === AuthType.USE_GEMINI && geminiApiKey) { if (authType === AuthType.USE_GEMINI && geminiApiKey) {
contentGeneratorConfig.apiKey = geminiApiKey; contentGeneratorConfig.apiKey = geminiApiKey;
contentGeneratorConfig.vertexai = false; contentGeneratorConfig.vertexai = false;
contentGeneratorConfig.model = await getEffectiveModel( getEffectiveModel(
contentGeneratorConfig.apiKey, contentGeneratorConfig.apiKey,
contentGeneratorConfig.model, contentGeneratorConfig.model,
); ).then((newModel) => {
if (newModel !== contentGeneratorConfig.model) {
config.flashFallbackHandler?.(contentGeneratorConfig.model, newModel);
}
});
return contentGeneratorConfig; return contentGeneratorConfig;
} }