1 files changed, 47 insertions, 3 deletions

diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts
index 1e7d5bcb..e0fc4ced 100644
--- a/packages/core/src/utils/retry.ts
+++ b/packages/core/src/utils/retry.ts
@@ -4,11 +4,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { AuthType } from '../core/contentGenerator.js';
+
 export interface RetryOptions {
   maxAttempts: number;
   initialDelayMs: number;
   maxDelayMs: number;
   shouldRetry: (error: Error) => boolean;
+  onPersistent429?: (authType?: string) => Promise<string | null>;
+  authType?: string;
 }
 
 const DEFAULT_RETRY_OPTIONS: RetryOptions = {
@@ -59,29 +63,69 @@ export async function retryWithBackoff<T>(
   fn: () => Promise<T>,
   options?: Partial<RetryOptions>,
 ): Promise<T> {
-  const { maxAttempts, initialDelayMs, maxDelayMs, shouldRetry } = {
+  const {
+    maxAttempts,
+    initialDelayMs,
+    maxDelayMs,
+    shouldRetry,
+    onPersistent429,
+    authType,
+  } = {
     ...DEFAULT_RETRY_OPTIONS,
     ...options,
   };
 
   let attempt = 0;
   let currentDelay = initialDelayMs;
+  let consecutive429Count = 0;
 
   while (attempt < maxAttempts) {
     attempt++;
     try {
       return await fn();
     } catch (error) {
+      const errorStatus = getErrorStatus(error);
+
+      // Track consecutive 429 errors
+      if (errorStatus === 429) {
+        consecutive429Count++;
+      } else {
+        consecutive429Count = 0;
+      }
+
+      // Check if we've exhausted retries or shouldn't retry
       if (attempt >= maxAttempts || !shouldRetry(error as Error)) {
+        // If we have persistent 429s and a fallback callback for OAuth
+        if (
+          consecutive429Count >= 3 &&
+          onPersistent429 &&
+          (authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL ||
+            authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE)
+        ) {
+          try {
+            const fallbackModel = await onPersistent429(authType);
+            if (fallbackModel) {
+              // Reset attempt counter and try with new model
+              attempt = 0;
+              consecutive429Count = 0;
+              currentDelay = initialDelayMs;
+              continue;
+            }
+          } catch (fallbackError) {
+            // If fallback fails, continue with original error
+            console.warn('Fallback to Flash model failed:', fallbackError);
+          }
+        }
         throw error;
       }
 
-      const { delayDurationMs, errorStatus } = getDelayDurationAndStatus(error);
+      const { delayDurationMs, errorStatus: delayErrorStatus } =
+        getDelayDurationAndStatus(error);
 
       if (delayDurationMs > 0) {
         // Respect Retry-After header if present and parsed
         console.warn(
-          `Attempt ${attempt} failed with status ${errorStatus ?? 'unknown'}. Retrying after explicit delay of ${delayDurationMs}ms...`,
+          `Attempt ${attempt} failed with status ${delayErrorStatus ?? 'unknown'}. Retrying after explicit delay of ${delayDurationMs}ms...`,
           error,
         );
         await delay(delayDurationMs);