/** * @license * Copyright 2025 Google LLC * SPDX-License-Identifier: Apache-2.0 */ import { describe, it, expect, beforeEach, vi } from 'vitest'; import { Config } from '../config/config.js'; import { setSimulate429, disableSimulationAfterFallback, shouldSimulate429, createSimulated429Error, resetRequestCounter, } from './testUtils.js'; import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; import { retryWithBackoff } from './retry.js'; import { AuthType } from '../core/contentGenerator.js'; describe('Flash Fallback Integration', () => { let config: Config; beforeEach(() => { config = new Config({ sessionId: 'test-session', targetDir: '/test', debugMode: false, cwd: '/test', model: 'gemini-2.5-pro', }); // Reset simulation state for each test setSimulate429(false); resetRequestCounter(); }); it('should automatically accept fallback', async () => { // Set up a minimal flash fallback handler for testing const flashFallbackHandler = async (): Promise => true; config.setFlashFallbackHandler(flashFallbackHandler); // Call the handler directly to test const result = await config.flashFallbackHandler!( 'gemini-2.5-pro', DEFAULT_GEMINI_FLASH_MODEL, ); // Verify it automatically accepts expect(result).toBe(true); }); it('should trigger fallback after 2 consecutive 429 errors for OAuth users', async () => { let fallbackCalled = false; let fallbackModel = ''; // Mock function that simulates exactly 2 429 errors, then succeeds after fallback const mockApiCall = vi .fn() .mockRejectedValueOnce(createSimulated429Error()) .mockRejectedValueOnce(createSimulated429Error()) .mockResolvedValueOnce('success after fallback'); // Mock fallback handler const mockFallbackHandler = vi.fn(async (_authType?: string) => { fallbackCalled = true; fallbackModel = DEFAULT_GEMINI_FLASH_MODEL; return fallbackModel; }); // Test with OAuth personal auth type, with maxAttempts = 2 to ensure fallback triggers const result = await retryWithBackoff(mockApiCall, { maxAttempts: 2, initialDelayMs: 1, maxDelayMs: 10, shouldRetry: (error: Error) => { const status = (error as Error & { status?: number }).status; return status === 429; }, onPersistent429: mockFallbackHandler, authType: AuthType.LOGIN_WITH_GOOGLE, }); // Verify fallback was triggered expect(fallbackCalled).toBe(true); expect(fallbackModel).toBe(DEFAULT_GEMINI_FLASH_MODEL); expect(mockFallbackHandler).toHaveBeenCalledWith( AuthType.LOGIN_WITH_GOOGLE, expect.any(Error), ); expect(result).toBe('success after fallback'); // Should have: 2 failures, then fallback triggered, then 1 success after retry reset expect(mockApiCall).toHaveBeenCalledTimes(3); }); it('should not trigger fallback for API key users', async () => { let fallbackCalled = false; // Mock function that simulates 429 errors const mockApiCall = vi.fn().mockRejectedValue(createSimulated429Error()); // Mock fallback handler const mockFallbackHandler = vi.fn(async () => { fallbackCalled = true; return DEFAULT_GEMINI_FLASH_MODEL; }); // Test with API key auth type - should not trigger fallback try { await retryWithBackoff(mockApiCall, { maxAttempts: 5, initialDelayMs: 10, maxDelayMs: 100, shouldRetry: (error: Error) => { const status = (error as Error & { status?: number }).status; return status === 429; }, onPersistent429: mockFallbackHandler, authType: AuthType.USE_GEMINI, // API key auth type }); } catch (error) { // Expected to throw after max attempts expect((error as Error).message).toContain('Rate limit exceeded'); } // Verify fallback was NOT triggered for API key users expect(fallbackCalled).toBe(false); expect(mockFallbackHandler).not.toHaveBeenCalled(); }); it('should properly disable simulation state after fallback', () => { // Enable simulation setSimulate429(true); // Verify simulation is enabled expect(shouldSimulate429()).toBe(true); // Disable simulation after fallback disableSimulationAfterFallback(); // Verify simulation is now disabled expect(shouldSimulate429()).toBe(false); }); });