Skip to content

Commit cba2720

Browse files
SandyTao520SandyTao520
authored
Run model availability check in the background to speed up startup (google-gemini#4256)
1 parent d622e59 commit cba2720

File tree

5 files changed

+69
-54
lines changed

5 files changed

+69
-54
lines changed

‎packages/cli/src/ui/App.tsx‎

Lines changed: 44 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ import {
5757
EditorType,
5858
FlashFallbackEvent,
5959
logFlashFallback,
60+
AuthType,
6061
} from '@google/gemini-cli-core';
6162
import { validateAuthMethod } from '../config/auth.js';
6263
import { useLogger } from './hooks/useLogger.js';
@@ -294,64 +295,70 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
294295
): Promise<boolean> => {
295296
let message: string;
296297

297-
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default)
298-
const isPaidTier =
299-
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
300-
301-
// Check if this is a Pro quota exceeded error
302-
if (error && isProQuotaExceededError(error)) {
303-
if (isPaidTier) {
304-
message = `⚡ You have reached your daily ${currentModel} quota limit.
298+
if (
299+
config.getContentGeneratorConfig().authType ===
300+
AuthType.LOGIN_WITH_GOOGLE
301+
) {
302+
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default)
303+
const isPaidTier =
304+
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
305+
306+
// Check if this is a Pro quota exceeded error
307+
if (error && isProQuotaExceededError(error)) {
308+
if (isPaidTier) {
309+
message = `⚡ You have reached your daily ${currentModel} quota limit.
305310
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
306311
⚡ 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`;
307-
} else {
308-
message = `⚡ You have reached your daily ${currentModel} quota limit.
312+
} else {
313+
message = `⚡ You have reached your daily ${currentModel} quota limit.
309314
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
310315
⚡ 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
311316
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
312317
⚡ You can switch authentication methods by typing /auth`;
313-
}
314-
} else if (error && isGenericQuotaExceededError(error)) {
315-
if (isPaidTier) {
316-
message = `⚡ You have reached your daily quota limit.
318+
}
319+
} else if (error && isGenericQuotaExceededError(error)) {
320+
if (isPaidTier) {
321+
message = `⚡ You have reached your daily quota limit.
317322
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
318323
⚡ 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`;
319-
} else {
320-
message = `⚡ You have reached your daily quota limit.
324+
} else {
325+
message = `⚡ You have reached your daily quota limit.
321326
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
322327
⚡ 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
323328
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
324329
⚡ You can switch authentication methods by typing /auth`;
325-
}
326-
} else {
327-
if (isPaidTier) {
328-
// Default fallback message for other cases (like consecutive 429s)
329-
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
330+
}
331+
} else {
332+
if (isPaidTier) {
333+
// Default fallback message for other cases (like consecutive 429s)
334+
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
330335
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
331336
⚡ 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`;
332-
} else {
333-
// Default fallback message for other cases (like consecutive 429s)
334-
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
337+
} else {
338+
// Default fallback message for other cases (like consecutive 429s)
339+
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
335340
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
336341
⚡ 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
337342
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
338343
⚡ You can switch authentication methods by typing /auth`;
344+
}
339345
}
340-
}
341346

342-
// Add message to UI history
343-
addItem(
344-
{
345-
type: MessageType.INFO,
346-
text: message,
347-
},
348-
Date.now(),
349-
);
347+
// Add message to UI history
348+
addItem(
349+
{
350+
type: MessageType.INFO,
351+
text: message,
352+
},
353+
Date.now(),
354+
);
355+
356+
// Set the flag to prevent tool continuation
357+
setModelSwitchedFromQuotaError(true);
358+
// Set global quota error flag to prevent Flash model calls
359+
config.setQuotaErrorOccurred(true);
360+
}
350361

351-
// Set the flag to prevent tool continuation
352-
setModelSwitchedFromQuotaError(true);
353-
// Set global quota error flag to prevent Flash model calls
354-
config.setQuotaErrorOccurred(true);
355362
// Switch model for future use but return false to stop current retry
356363
config.setModel(fallbackModel);
357364
logFlashFallback(

‎packages/core/src/config/config.test.ts‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -151,14 +151,12 @@ describe('Server Config (config.ts)', () => {
151151
apiKey: 'test-key',
152152
};
153153

154-
(createContentGeneratorConfig as Mock).mockResolvedValue(
155-
mockContentConfig,
156-
);
154+
(createContentGeneratorConfig as Mock).mockReturnValue(mockContentConfig);
157155

158156
await config.refreshAuth(authType);
159157

160158
expect(createContentGeneratorConfig).toHaveBeenCalledWith(
161-
MODEL, // Should be called with the original model 'gemini-pro'
159+
config,
162160
authType,
163161
);
164162
// Verify that contentGeneratorConfig is updated with the new model

‎packages/core/src/config/config.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -274,8 +274,8 @@ export class Config {
274274
}
275275

276276
async refreshAuth(authMethod: AuthType) {
277-
this.contentGeneratorConfig = await createContentGeneratorConfig(
278-
this.model,
277+
this.contentGeneratorConfig = createContentGeneratorConfig(
278+
this,
279279
authMethod,
280280
);
281281

‎packages/core/src/core/contentGenerator.test.ts‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,12 +64,18 @@ describe('createContentGenerator', () => {
6464

6565
describe('createContentGeneratorConfig', () => {
6666
const originalEnv = process.env;
67+
const mockConfig = {
68+
getModel: vi.fn().mockReturnValue('gemini-pro'),
69+
setModel: vi.fn(),
70+
flashFallbackHandler: vi.fn(),
71+
} as unknown as Config;
6772

6873
beforeEach(() => {
6974
// Reset modules to re-evaluate imports and environment variables
7075
vi.resetModules();
7176
// Restore process.env before each test
7277
process.env = { ...originalEnv };
78+
vi.clearAllMocks();
7379
});
7480

7581
afterAll(() => {
@@ -80,7 +86,7 @@ describe('createContentGeneratorConfig', () => {
8086
it('should configure for Gemini using GEMINI_API_KEY when set', async () => {
8187
process.env.GEMINI_API_KEY = 'env-gemini-key';
8288
const config = await createContentGeneratorConfig(
83-
undefined,
89+
mockConfig,
8490
AuthType.USE_GEMINI,
8591
);
8692
expect(config.apiKey).toBe('env-gemini-key');
@@ -90,7 +96,7 @@ describe('createContentGeneratorConfig', () => {
9096
it('should not configure for Gemini if GEMINI_API_KEY is empty', async () => {
9197
process.env.GEMINI_API_KEY = '';
9298
const config = await createContentGeneratorConfig(
93-
undefined,
99+
mockConfig,
94100
AuthType.USE_GEMINI,
95101
);
96102
expect(config.apiKey).toBeUndefined();
@@ -100,7 +106,7 @@ describe('createContentGeneratorConfig', () => {
100106
it('should configure for Vertex AI using GOOGLE_API_KEY when set', async () => {
101107
process.env.GOOGLE_API_KEY = 'env-google-key';
102108
const config = await createContentGeneratorConfig(
103-
undefined,
109+
mockConfig,
104110
AuthType.USE_VERTEX_AI,
105111
);
106112
expect(config.apiKey).toBe('env-google-key');
@@ -111,7 +117,7 @@ describe('createContentGeneratorConfig', () => {
111117
process.env.GOOGLE_CLOUD_PROJECT = 'env-gcp-project';
112118
process.env.GOOGLE_CLOUD_LOCATION = 'env-gcp-location';
113119
const config = await createContentGeneratorConfig(
114-
undefined,
120+
mockConfig,
115121
AuthType.USE_VERTEX_AI,
116122
);
117123
expect(config.vertexai).toBe(true);
@@ -123,7 +129,7 @@ describe('createContentGeneratorConfig', () => {
123129
process.env.GOOGLE_CLOUD_PROJECT = '';
124130
process.env.GOOGLE_CLOUD_LOCATION = '';
125131
const config = await createContentGeneratorConfig(
126-
undefined,
132+
mockConfig,
127133
AuthType.USE_VERTEX_AI,
128134
);
129135
expect(config.apiKey).toBeUndefined();

‎packages/core/src/core/contentGenerator.ts‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -52,17 +52,17 @@ export type ContentGeneratorConfig = {
5252
authType?: AuthType | undefined;
5353
};
5454

55-
export async function createContentGeneratorConfig(
56-
model: string | undefined,
55+
export function createContentGeneratorConfig(
56+
config: Config,
5757
authType: AuthType | undefined,
58-
): Promise<ContentGeneratorConfig> {
58+
): ContentGeneratorConfig {
5959
const geminiApiKey = process.env.GEMINI_API_KEY || undefined;
6060
const googleApiKey = process.env.GOOGLE_API_KEY || undefined;
6161
const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT || undefined;
6262
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION || undefined;
6363

6464
// Use runtime model from config if available, otherwise fallback to parameter or default
65-
const effectiveModel = model || DEFAULT_GEMINI_MODEL;
65+
const effectiveModel = config.getModel() || DEFAULT_GEMINI_MODEL;
6666

6767
const contentGeneratorConfig: ContentGeneratorConfig = {
6868
model: effectiveModel,
@@ -80,10 +80,14 @@ export async function createContentGeneratorConfig(
8080
if (authType === AuthType.USE_GEMINI && geminiApiKey) {
8181
contentGeneratorConfig.apiKey = geminiApiKey;
8282
contentGeneratorConfig.vertexai = false;
83-
contentGeneratorConfig.model = await getEffectiveModel(
83+
getEffectiveModel(
8484
contentGeneratorConfig.apiKey,
8585
contentGeneratorConfig.model,
86-
);
86+
).then((newModel) => {
87+
if (newModel !== contentGeneratorConfig.model) {
88+
config.flashFallbackHandler?.(contentGeneratorConfig.model, newModel);
89+
}
90+
});
8791

8892
return contentGeneratorConfig;
8993
}

0 commit comments

Comments
 (0)