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

5 files changed, 68 insertions, 53 deletions

diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx
index cd94311b..f66d8a5b 100644
--- a/packages/cli/src/ui/App.tsx
+++ b/packages/cli/src/ui/App.tsx
@@ -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(),
-      );
+        // 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);
+      }
 
-      // 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(
diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts
index bb074a71..e34880a6 100644
--- a/packages/core/src/config/config.test.ts
+++ b/packages/core/src/config/config.test.ts
@@ -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
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 268871ca..59f1e1ba 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -274,8 +274,8 @@ export class Config {
   }
 
   async refreshAuth(authMethod: AuthType) {
-    this.contentGeneratorConfig = await createContentGeneratorConfig(
-      this.model,
+    this.contentGeneratorConfig = createContentGeneratorConfig(
+      this,
       authMethod,
     );
 
diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts
index 92144aa4..c50678af 100644
--- a/packages/core/src/core/contentGenerator.test.ts
+++ b/packages/core/src/core/contentGenerator.test.ts
@@ -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();
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index d3434c23..109c0ffc 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -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;
   }