From eacbb3551ce12012a0ef9b9338b8b747b2e7161c Mon Sep 17 00:00:00 2001
From: Bryan Morgan <bryanmorgan@google.com>
Date: Wed, 25 Jun 2025 15:37:23 -0400
Subject: [PATCH] changed 429 failover from 3 consecutive to 2 for OAuth users
 (#1579)

---
 .../src/utils/flashFallback.integration.test.ts     | 13 ++++++-------
 packages/core/src/utils/retry.test.ts               |  2 +-
 packages/core/src/utils/retry.ts                    |  2 +-
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/packages/core/src/utils/flashFallback.integration.test.ts b/packages/core/src/utils/flashFallback.integration.test.ts
index 21c40296..cffe9aa5 100644
--- a/packages/core/src/utils/flashFallback.integration.test.ts
+++ b/packages/core/src/utils/flashFallback.integration.test.ts
@@ -50,16 +50,15 @@ describe('Flash Fallback Integration', () => {
     expect(result).toBe(true);
   });
 
-  it('should trigger fallback after 3 consecutive 429 errors for OAuth users', async () => {
+  it('should trigger fallback after 2 consecutive 429 errors for OAuth users', async () => {
     let fallbackCalled = false;
     let fallbackModel = '';
 
-    // Mock function that simulates exactly 3 429 errors, then succeeds after fallback
+    // Mock function that simulates exactly 2 429 errors, then succeeds after fallback
     const mockApiCall = vi
       .fn()
       .mockRejectedValueOnce(createSimulated429Error())
       .mockRejectedValueOnce(createSimulated429Error())
-      .mockRejectedValueOnce(createSimulated429Error())
       .mockResolvedValueOnce('success after fallback');
 
     // Mock fallback handler
@@ -69,9 +68,9 @@ describe('Flash Fallback Integration', () => {
       return fallbackModel;
     });
 
-    // Test with OAuth personal auth type, with maxAttempts = 3 to ensure fallback triggers
+    // Test with OAuth personal auth type, with maxAttempts = 2 to ensure fallback triggers
     const result = await retryWithBackoff(mockApiCall, {
-      maxAttempts: 3,
+      maxAttempts: 2,
       initialDelayMs: 1,
       maxDelayMs: 10,
       shouldRetry: (error: Error) => {
@@ -89,8 +88,8 @@ describe('Flash Fallback Integration', () => {
       AuthType.LOGIN_WITH_GOOGLE_PERSONAL,
     );
     expect(result).toBe('success after fallback');
-    // Should have: 3 failures, then fallback triggered, then 1 success after retry reset
-    expect(mockApiCall).toHaveBeenCalledTimes(4);
+    // 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 () => {
diff --git a/packages/core/src/utils/retry.test.ts b/packages/core/src/utils/retry.test.ts
index 39f62981..031c0991 100644
--- a/packages/core/src/utils/retry.test.ts
+++ b/packages/core/src/utils/retry.test.ts
@@ -426,7 +426,7 @@ describe('retryWithBackoff', () => {
 
       await expect(promise).resolves.toBe('success');
 
-      // Should trigger fallback after 4 consecutive 429s (attempts 2-5)
+      // Should trigger fallback after 2 consecutive 429s (attempts 2-3)
       expect(fallbackCallback).toHaveBeenCalledWith('oauth-personal');
     });
   });
diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts
index e0fc4ced..851db0e7 100644
--- a/packages/core/src/utils/retry.ts
+++ b/packages/core/src/utils/retry.ts
@@ -97,7 +97,7 @@ export async function retryWithBackoff<T>(
       if (attempt >= maxAttempts || !shouldRetry(error as Error)) {
         // If we have persistent 429s and a fallback callback for OAuth
         if (
-          consecutive429Count >= 3 &&
+          consecutive429Count >= 2 &&
           onPersistent429 &&
           (authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL ||
             authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE)