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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 68 additions and 53 deletions

View File

@ -57,6 +57,7 @@ import {
EditorType,
FlashFallbackEvent,
logFlashFallback,
AuthType,
} from '@google/gemini-cli-core';
import { validateAuthMethod } from '../config/auth.js';
import { useLogger } from './hooks/useLogger.js';
@ -294,64 +295,70 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
): Promise<boolean> => {
let message: string;
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default)
const isPaidTier =
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
if (
config.getContentGeneratorConfig().authType ===
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
if (error && isProQuotaExceededError(error)) {
if (isPaidTier) {
message = `⚡ You have reached your daily ${currentModel} quota limit.
// Check if this is a Pro quota exceeded error
if (error && isProQuotaExceededError(error)) {
if (isPaidTier) {
message = `⚡ You have reached your daily ${currentModel} quota limit.
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`;
} else {
message = `⚡ You have reached your daily ${currentModel} quota limit.
} else {
message = `⚡ You have reached your daily ${currentModel} quota limit.
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
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`;
}
} else if (error && isGenericQuotaExceededError(error)) {
if (isPaidTier) {
message = `⚡ You have reached your daily quota limit.
}
} else if (error && isGenericQuotaExceededError(error)) {
if (isPaidTier) {
message = `⚡ You have reached your daily quota limit.
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`;
} else {
message = `⚡ You have reached your daily quota limit.
} else {
message = `⚡ You have reached your daily quota limit.
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
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`;
}
} else {
if (isPaidTier) {
// 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.
}
} else {
if (isPaidTier) {
// 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.
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`;
} else {
// 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.
} else {
// 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.
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
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`;
}
}
// 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
config.setModel(fallbackModel);
logFlashFallback(

View File

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

View File

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

View File

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

View File

@ -52,17 +52,17 @@ export type ContentGeneratorConfig = {
authType?: AuthType | undefined;
};
export async function createContentGeneratorConfig(
model: string | undefined,
export function createContentGeneratorConfig(
config: Config,
authType: AuthType | undefined,
): Promise<ContentGeneratorConfig> {
): ContentGeneratorConfig {
const geminiApiKey = process.env.GEMINI_API_KEY || undefined;
const googleApiKey = process.env.GOOGLE_API_KEY || undefined;
const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT || undefined;
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION || undefined;
// 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 = {
model: effectiveModel,
@ -80,10 +80,14 @@ export async function createContentGeneratorConfig(
if (authType === AuthType.USE_GEMINI && geminiApiKey) {
contentGeneratorConfig.apiKey = geminiApiKey;
contentGeneratorConfig.vertexai = false;
contentGeneratorConfig.model = await getEffectiveModel(
getEffectiveModel(
contentGeneratorConfig.apiKey,
contentGeneratorConfig.model,
);
).then((newModel) => {
if (newModel !== contentGeneratorConfig.model) {
config.flashFallbackHandler?.(contentGeneratorConfig.model, newModel);
}
});
return contentGeneratorConfig;
}