Revert "fix: handle request retries and model fallback correctly" (#11164)

2026-07-16 21:10:40 -07:00 · 2025-10-14 16:30:59 -07:00
parent a75de5a4b4
commit b8e9a82cbf
13 changed files with 821 additions and 1013 deletions
@@ -7,14 +7,9 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { ApiError } from '@google/genai';
-import { AuthType } from '../core/contentGenerator.js';
 import type { HttpError } from './retry.js';
 import { retryWithBackoff } from './retry.js';
 import { setSimulate429 } from './testUtils.js';
-import {
-  TerminalQuotaError,
-  RetryableQuotaError,
-} from './googleQuotaErrors.js';

 // Helper to create a mock function that fails a certain number of times
 const createFailingFunction = (
@@ -104,26 +99,26 @@ describe('retryWithBackoff', () => {

    // Expect it to fail with the error from the 5th attempt.
    await Promise.all([
-      expect(promise).rejects.toThrow('Simulated error attempt 3'),
+      expect(promise).rejects.toThrow('Simulated error attempt 5'),
      vi.runAllTimersAsync(),
    ]);

-    expect(mockFn).toHaveBeenCalledTimes(3);
+    expect(mockFn).toHaveBeenCalledTimes(5);
  });

-  it('should default to 3 maxAttempts if options.maxAttempts is undefined', async () => {
-    // This function will fail more than 3 times to ensure all retries are used.
+  it('should default to 5 maxAttempts if options.maxAttempts is undefined', async () => {
+    // This function will fail more than 5 times to ensure all retries are used.
    const mockFn = createFailingFunction(10);

    const promise = retryWithBackoff(mockFn, { maxAttempts: undefined });

    // Expect it to fail with the error from the 5th attempt.
    await Promise.all([
-      expect(promise).rejects.toThrow('Simulated error attempt 3'),
+      expect(promise).rejects.toThrow('Simulated error attempt 5'),
      vi.runAllTimersAsync(),
    ]);

-    expect(mockFn).toHaveBeenCalledTimes(3);
+    expect(mockFn).toHaveBeenCalledTimes(5);
  });

  it('should not retry if shouldRetry returns false', async () => {
@@ -340,13 +335,15 @@ describe('retryWithBackoff', () => {
  });

  describe('Flash model fallback for OAuth users', () => {
-    it('should trigger fallback for OAuth personal users on TerminalQuotaError', async () => {
+    it('should trigger fallback for OAuth personal users after persistent 429 errors', async () => {
      const fallbackCallback = vi.fn().mockResolvedValue('gemini-2.5-flash');

      let fallbackOccurred = false;
      const mockFn = vi.fn().mockImplementation(async () => {
        if (!fallbackOccurred) {
-          throw new TerminalQuotaError('Daily limit reached', {} as any);
+          const error: HttpError = new Error('Rate limit exceeded');
+          error.status = 429;
+          throw error;
        }
        return 'success';
      });
@@ -354,9 +351,143 @@ describe('retryWithBackoff', () => {
      const promise = retryWithBackoff(mockFn, {
        maxAttempts: 3,
        initialDelayMs: 100,
-        onPersistent429: async (authType?: string, error?: unknown) => {
+        onPersistent429: async (authType?: string) => {
          fallbackOccurred = true;
-          return await fallbackCallback(authType, error);
+          return await fallbackCallback(authType);
+        },
+        authType: 'oauth-personal',
+      });
+
+      // Advance all timers to complete retries
+      await vi.runAllTimersAsync();
+
+      // Should succeed after fallback
+      await expect(promise).resolves.toBe('success');
+
+      // Verify callback was called with correct auth type
+      expect(fallbackCallback).toHaveBeenCalledWith('oauth-personal');
+
+      // Should retry again after fallback
+      expect(mockFn).toHaveBeenCalledTimes(3); // 2 initial attempts + 1 after fallback
+    });
+
+    it('should NOT trigger fallback for API key users', async () => {
+      const fallbackCallback = vi.fn();
+
+      const mockFn = vi.fn(async () => {
+        const error: HttpError = new Error('Rate limit exceeded');
+        error.status = 429;
+        throw error;
+      });
+
+      const promise = retryWithBackoff(mockFn, {
+        maxAttempts: 3,
+        initialDelayMs: 100,
+        onPersistent429: fallbackCallback,
+        authType: 'gemini-api-key',
+      });
+
+      // Handle the promise properly to avoid unhandled rejections
+      const resultPromise = promise.catch((error) => error);
+      await vi.runAllTimersAsync();
+      const result = await resultPromise;
+
+      // Should fail after all retries without fallback
+      expect(result).toBeInstanceOf(Error);
+      expect(result.message).toBe('Rate limit exceeded');
+
+      // Callback should not be called for API key users
+      expect(fallbackCallback).not.toHaveBeenCalled();
+    });
+
+    it('should reset attempt counter and continue after successful fallback', async () => {
+      let fallbackCalled = false;
+      const fallbackCallback = vi.fn().mockImplementation(async () => {
+        fallbackCalled = true;
+        return 'gemini-2.5-flash';
+      });
+
+      const mockFn = vi.fn().mockImplementation(async () => {
+        if (!fallbackCalled) {
+          const error: HttpError = new Error('Rate limit exceeded');
+          error.status = 429;
+          throw error;
+        }
+        return 'success';
+      });
+
+      const promise = retryWithBackoff(mockFn, {
+        maxAttempts: 3,
+        initialDelayMs: 100,
+        onPersistent429: fallbackCallback,
+        authType: 'oauth-personal',
+      });
+
+      await vi.runAllTimersAsync();
+
+      await expect(promise).resolves.toBe('success');
+      expect(fallbackCallback).toHaveBeenCalledOnce();
+    });
+
+    it('should continue with original error if fallback is rejected', async () => {
+      const fallbackCallback = vi.fn().mockResolvedValue(null); // User rejected fallback
+
+      const mockFn = vi.fn(async () => {
+        const error: HttpError = new Error('Rate limit exceeded');
+        error.status = 429;
+        throw error;
+      });
+
+      const promise = retryWithBackoff(mockFn, {
+        maxAttempts: 3,
+        initialDelayMs: 100,
+        onPersistent429: fallbackCallback,
+        authType: 'oauth-personal',
+      });
+
+      // Handle the promise properly to avoid unhandled rejections
+      const resultPromise = promise.catch((error) => error);
+      await vi.runAllTimersAsync();
+      const result = await resultPromise;
+
+      // Should fail with original error when fallback is rejected
+      expect(result).toBeInstanceOf(Error);
+      expect(result.message).toBe('Rate limit exceeded');
+      expect(fallbackCallback).toHaveBeenCalledWith(
+        'oauth-personal',
+        expect.any(Error),
+      );
+    });
+
+    it('should handle mixed error types (only count consecutive 429s)', async () => {
+      const fallbackCallback = vi.fn().mockResolvedValue('gemini-2.5-flash');
+      let attempts = 0;
+      let fallbackOccurred = false;
+
+      const mockFn = vi.fn().mockImplementation(async () => {
+        attempts++;
+        if (fallbackOccurred) {
+          return 'success';
+        }
+        if (attempts === 1) {
+          // First attempt: 500 error (resets consecutive count)
+          const error: HttpError = new Error('Server error');
+          error.status = 500;
+          throw error;
+        } else {
+          // Remaining attempts: 429 errors
+          const error: HttpError = new Error('Rate limit exceeded');
+          error.status = 429;
+          throw error;
+        }
+      });
+
+      const promise = retryWithBackoff(mockFn, {
+        maxAttempts: 5,
+        initialDelayMs: 100,
+        onPersistent429: async (authType?: string) => {
+          fallbackOccurred = true;
+          return await fallbackCallback(authType);
        },
        authType: 'oauth-personal',
      });
@@ -364,51 +495,9 @@ describe('retryWithBackoff', () => {
      await vi.runAllTimersAsync();

      await expect(promise).resolves.toBe('success');
-      expect(fallbackCallback).toHaveBeenCalledWith(
-        'oauth-personal',
-        expect.any(TerminalQuotaError),
-      );
-      expect(mockFn).toHaveBeenCalledTimes(2);
+
+      // Should trigger fallback after 2 consecutive 429s (attempts 2-3)
+      expect(fallbackCallback).toHaveBeenCalledWith('oauth-personal');
    });
-
-    it('should use retryDelayMs from RetryableQuotaError', async () => {
-      const setTimeoutSpy = vi.spyOn(global, 'setTimeout');
-      const mockFn = vi.fn().mockImplementation(async () => {
-        throw new RetryableQuotaError('Per-minute limit', {} as any, 12.345);
-      });
-
-      const promise = retryWithBackoff(mockFn, {
-        maxAttempts: 2,
-        initialDelayMs: 100,
-      });
-
-      // Attach the rejection expectation *before* running timers
-      // eslint-disable-next-line vitest/valid-expect
-      const assertionPromise = expect(promise).rejects.toThrow();
-      await vi.runAllTimersAsync();
-      await assertionPromise;
-
-      expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 12345);
-    });
-
-    it.each([[AuthType.USE_GEMINI], [AuthType.USE_VERTEX_AI], [undefined]])(
-      'should not trigger fallback for non-Google auth users (authType: %s) on TerminalQuotaError',
-      async (authType) => {
-        const fallbackCallback = vi.fn();
-        const mockFn = vi.fn().mockImplementation(async () => {
-          throw new TerminalQuotaError('Daily limit reached', {} as any);
-        });
-
-        const promise = retryWithBackoff(mockFn, {
-          maxAttempts: 3,
-          onPersistent429: fallbackCallback,
-          authType,
-        });
-
-        await expect(promise).rejects.toThrow('Daily limit reached');
-        expect(fallbackCallback).not.toHaveBeenCalled();
-        expect(mockFn).toHaveBeenCalledTimes(1);
-      },
-    );
  });
 });