feat(modelAvailabilityService): integrate model availability service into backend logic (#14470)

packages/core/src/fallback/handler.test.ts

@@ -17,6 +17,7 @@ import {
 import { handleFallback } from './handler.js';
 import type { Config } from '../config/config.js';
 import type { ModelAvailabilityService } from '../availability/modelAvailabilityService.js';
+import { createAvailabilityServiceMock } from '../availability/testUtils.js';
 import { AuthType } from '../core/contentGenerator.js';
 import {
   DEFAULT_GEMINI_FLASH_MODEL,
@@ -45,25 +46,20 @@ vi.mock('../utils/secure-browser-launcher.js', () => ({
   openBrowserSecurely: vi.fn(),
 }));
 
+// Mock debugLogger to prevent console pollution and allow spying
+vi.mock('../utils/debugLogger.js', () => ({
+  debugLogger: {
+    warn: vi.fn(),
+    error: vi.fn(),
+    log: vi.fn(),
+  },
+}));
+
 const MOCK_PRO_MODEL = DEFAULT_GEMINI_MODEL;
 const FALLBACK_MODEL = DEFAULT_GEMINI_FLASH_MODEL;
 const AUTH_OAUTH = AuthType.LOGIN_WITH_GOOGLE;
 const AUTH_API_KEY = AuthType.USE_GEMINI;
 
-function createAvailabilityMock(
-  result: ReturnType<ModelAvailabilityService['selectFirstAvailable']>,
-): ModelAvailabilityService {
-  return {
-    markTerminal: vi.fn(),
-    markHealthy: vi.fn(),
-    markRetryOncePerTurn: vi.fn(),
-    consumeStickyAttempt: vi.fn(),
-    snapshot: vi.fn(),
-    selectFirstAvailable: vi.fn().mockReturnValue(result),
-    resetTurn: vi.fn(),
-  } as unknown as ModelAvailabilityService;
-}
-
 const createMockConfig = (overrides: Partial<Config> = {}): Config =>
   ({
     isInFallbackMode: vi.fn(() => false),
@@ -75,8 +71,12 @@ const createMockConfig = (overrides: Partial<Config> = {}): Config =>
     setPreviewModelBypassMode: vi.fn(),
     fallbackHandler: undefined,
     getFallbackModelHandler: vi.fn(),
+    setActiveModel: vi.fn(),
     getModelAvailabilityService: vi.fn(() =>
-      createAvailabilityMock({ selectedModel: FALLBACK_MODEL, skipped: [] }),
+      createAvailabilityServiceMock({
+        selectedModel: FALLBACK_MODEL,
+        skipped: [],
+      }),
     ),
     getModel: vi.fn(() => MOCK_PRO_MODEL),
     getPreviewFeatures: vi.fn(() => false),
@@ -98,6 +98,12 @@ describe('handleFallback', () => {
     mockConfig = createMockConfig({
       fallbackModelHandler: mockHandler,
     });
+    // Explicitly set the property to ensure it's present for legacy checks
+    mockConfig.fallbackModelHandler = mockHandler;
+
+    // We mocked debugLogger, so we don't need to spy on console.error for handler failures
+    // But tests might check console.error usage in legacy code if any?
+    // The handler uses console.error in legacyHandleFallback.
     consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
     fallbackEventSpy = vi.spyOn(coreEvents, 'emitFallbackModeChanged');
   });
@@ -538,7 +544,7 @@ describe('handleFallback', () => {
 
     beforeEach(() => {
       vi.clearAllMocks();
-      availability = createAvailabilityMock({
+      availability = createAvailabilityServiceMock({
         selectedModel: DEFAULT_GEMINI_FLASH_MODEL,
         skipped: [],
       });
@@ -612,7 +618,17 @@ describe('handleFallback', () => {
 
         expect(result).toBe(true);
         expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled();
-        expect(policyConfig.setFallbackMode).toHaveBeenCalledWith(true);
+        expect(policyConfig.setActiveModel).toHaveBeenCalledWith(
+          DEFAULT_GEMINI_FLASH_MODEL,
+        );
+        // Silent actions should not trigger the legacy fallback mode (via activateFallbackMode),
+        // but setActiveModel might trigger it via legacy sync if it switches to Flash.
+        // However, the test requirement is "doesn't emit fallback mode".
+        // Since we are mocking setActiveModel, we can verify setFallbackMode isn't called *independently*.
+        // But setActiveModel is mocked, so it won't trigger side effects unless the implementation does.
+        // We verified setActiveModel is called.
+        // We verify setFallbackMode is NOT called (which would happen if activateFallbackMode was called).
+        expect(policyConfig.setFallbackMode).not.toHaveBeenCalled();
       } finally {
         chainSpy.mockRestore();
       }
@@ -707,5 +723,34 @@ describe('handleFallback', () => {
       expect(policyConfig.getModelAvailabilityService).toHaveBeenCalled();
       expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled();
     });
+
+    it('calls setActiveModel and logs telemetry when handler returns "retry_always"', async () => {
+      policyHandler.mockResolvedValue('retry_always');
+
+      const result = await handleFallback(
+        policyConfig,
+        MOCK_PRO_MODEL,
+        AUTH_OAUTH,
+      );
+
+      expect(result).toBe(true);
+      expect(policyConfig.setActiveModel).toHaveBeenCalledWith(FALLBACK_MODEL);
+      expect(policyConfig.setFallbackMode).not.toHaveBeenCalled();
+      // TODO: add logging expect statement
+    });
+
+    it('calls setActiveModel when handler returns "stop"', async () => {
+      policyHandler.mockResolvedValue('stop');
+
+      const result = await handleFallback(
+        policyConfig,
+        MOCK_PRO_MODEL,
+        AUTH_OAUTH,
+      );
+
+      expect(result).toBe(false);
+      expect(policyConfig.setActiveModel).toHaveBeenCalledWith(FALLBACK_MODEL);
+      // TODO: add logging expect statement
+    });
   });
 });

packages/core/src/fallback/handler.ts

@@ -12,17 +12,14 @@ import {
   PREVIEW_GEMINI_MODEL,
 } from '../config/models.js';
 import { logFlashFallback, FlashFallbackEvent } from '../telemetry/index.js';
-import { coreEvents } from '../utils/events.js';
 import { openBrowserSecurely } from '../utils/secure-browser-launcher.js';
 import { debugLogger } from '../utils/debugLogger.js';
 import { getErrorMessage } from '../utils/errors.js';
 import { ModelNotFoundError } from '../utils/httpErrors.js';
-import {
-  RetryableQuotaError,
-  TerminalQuotaError,
-} from '../utils/googleQuotaErrors.js';
+import { TerminalQuotaError } from '../utils/googleQuotaErrors.js';
+import { coreEvents } from '../utils/events.js';
 import type { FallbackIntent, FallbackRecommendation } from './types.js';
-import type { FailureKind } from '../availability/modelPolicy.js';
+import { classifyFailureKind } from '../availability/errorClassification.js';
 import {
   buildFallbackPolicyContext,
   resolvePolicyChain,
@@ -126,6 +123,9 @@ async function handlePolicyDrivenFallback(
     chain,
     failedModel,
   );
+
+  const failureKind = classifyFailureKind(error);
+
   if (!candidates.length) {
     return null;
   }
@@ -145,7 +145,7 @@ async function handlePolicyDrivenFallback(
     return null;
   }
 
-  const failureKind = classifyFailureKind(error);
+  // failureKind is already declared and calculated above
   const action = resolvePolicyAction(failureKind, selectedPolicy);
 
   if (action === 'silent') {
@@ -183,6 +183,7 @@ async function handlePolicyDrivenFallback(
       failedModel,
       fallbackModel,
       authType,
+      error, // Pass the error so processIntent can handle preview-specific logic
     );
   } catch (handlerError) {
     debugLogger.error('Fallback handler failed:', handlerError);
@@ -209,25 +210,42 @@ async function processIntent(
   authType?: string,
   error?: unknown,
 ): Promise<boolean> {
+  const isAvailabilityEnabled = config.isModelAvailabilityServiceEnabled();
+
   switch (intent) {
     case 'retry_always':
-      // If the error is non-retryable, e.g. TerminalQuota Error, trigger a regular fallback to flash.
-      // For all other errors, activate previewModel fallback.
-      if (
-        failedModel === PREVIEW_GEMINI_MODEL &&
-        !(error instanceof TerminalQuotaError)
-      ) {
-        activatePreviewModelFallbackMode(config);
+      if (isAvailabilityEnabled) {
+        // TODO(telemetry): Implement generic fallback event logging. Existing
+        // logFlashFallback is specific to a single Model.
+        config.setActiveModel(fallbackModel);
       } else {
-        activateFallbackMode(config, authType);
+        // If the error is non-retryable, e.g. TerminalQuota Error, trigger a regular fallback to flash.
+        // For all other errors, activate previewModel fallback.
+        if (
+          failedModel === PREVIEW_GEMINI_MODEL &&
+          !(error instanceof TerminalQuotaError)
+        ) {
+          activatePreviewModelFallbackMode(config);
+        } else {
+          activateFallbackMode(config, authType);
+        }
       }
       return true;
 
     case 'retry_once':
+      if (isAvailabilityEnabled) {
+        config.setActiveModel(fallbackModel);
+      }
       return true;
 
     case 'stop':
-      activateFallbackMode(config, authType);
+      if (isAvailabilityEnabled) {
+        // TODO(telemetry): Implement generic fallback event logging. Existing
+        // logFlashFallback is specific to a single Model.
+        config.setActiveModel(fallbackModel);
+      } else {
+        activateFallbackMode(config, authType);
+      }
       return false;
 
     case 'retry_later':
@@ -260,16 +278,3 @@ function activatePreviewModelFallbackMode(config: Config) {
     // We might want a specific event for Preview Model fallback, but for now we just set the mode.
   }
 }
-
-function classifyFailureKind(error?: unknown): FailureKind {
-  if (error instanceof TerminalQuotaError) {
-    return 'terminal';
-  }
-  if (error instanceof RetryableQuotaError) {
-    return 'transient';
-  }
-  if (error instanceof ModelNotFoundError) {
-    return 'not_found';
-  }
-  return 'unknown';
-}