feat(core): Add content-based retries for JSON generation (#9264)

packages/core/src/utils/flashFallback.test.ts

@@ -86,7 +86,7 @@ describe('Retry Utility Fallback Integration', () => {
       maxAttempts: 2,
       initialDelayMs: 1,
       maxDelayMs: 10,
-      shouldRetry: (error: Error) => {
+      shouldRetryOnError: (error: Error) => {
         const status = (error as Error & { status?: number }).status;
         return status === 429;
       },
@@ -123,7 +123,7 @@ describe('Retry Utility Fallback Integration', () => {
         maxAttempts: 5,
         initialDelayMs: 10,
         maxDelayMs: 100,
-        shouldRetry: (error: Error) => {
+        shouldRetryOnError: (error: Error) => {
           const status = (error as Error & { status?: number }).status;
           return status === 429;
         },

packages/core/src/utils/retry.test.ts

@@ -137,10 +137,11 @@ describe('retryWithBackoff', () => {
     const mockFn = vi.fn(async () => {
       throw new NonRetryableError('Non-retryable error');
     });
-    const shouldRetry = (error: Error) => !(error instanceof NonRetryableError);
+    const shouldRetryOnError = (error: Error) =>
+      !(error instanceof NonRetryableError);
 
     const promise = retryWithBackoff(mockFn, {
-      shouldRetry,
+      shouldRetryOnError,
       initialDelayMs: 10,
     });
 

packages/core/src/utils/retry.ts

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import type { GenerateContentResponse } from '@google/genai';
 import { AuthType } from '../core/contentGenerator.js';
 import {
   isProQuotaExceededError,
@@ -18,7 +19,8 @@ export interface RetryOptions {
   maxAttempts: number;
   initialDelayMs: number;
   maxDelayMs: number;
-  shouldRetry: (error: Error) => boolean;
+  shouldRetryOnError: (error: Error) => boolean;
+  shouldRetryOnContent?: (content: GenerateContentResponse) => boolean;
   onPersistent429?: (
     authType?: string,
     error?: unknown,
@@ -30,7 +32,7 @@ const DEFAULT_RETRY_OPTIONS: RetryOptions = {
   maxAttempts: 5,
   initialDelayMs: 5000,
   maxDelayMs: 30000, // 30 seconds
-  shouldRetry: defaultShouldRetry,
+  shouldRetryOnError: defaultShouldRetry,
 };
 
 /**
@@ -88,7 +90,8 @@ export async function retryWithBackoff<T>(
     maxDelayMs,
     onPersistent429,
     authType,
-    shouldRetry,
+    shouldRetryOnError,
+    shouldRetryOnContent,
   } = {
     ...DEFAULT_RETRY_OPTIONS,
     ...cleanOptions,
@@ -101,7 +104,20 @@ export async function retryWithBackoff<T>(
   while (attempt < maxAttempts) {
     attempt++;
     try {
-      return await fn();
+      const result = await fn();
+
+      if (
+        shouldRetryOnContent &&
+        shouldRetryOnContent(result as GenerateContentResponse)
+      ) {
+        const jitter = currentDelay * 0.3 * (Math.random() * 2 - 1);
+        const delayWithJitter = Math.max(0, currentDelay + jitter);
+        await delay(delayWithJitter);
+        currentDelay = Math.min(maxDelayMs, currentDelay * 2);
+        continue;
+      }
+
+      return result;
     } catch (error) {
       const errorStatus = getErrorStatus(error);
 
@@ -191,7 +207,7 @@ export async function retryWithBackoff<T>(
       }
 
       // Check if we've exhausted retries or shouldn't retry
-      if (attempt >= maxAttempts || !shouldRetry(error as Error)) {
+      if (attempt >= maxAttempts || !shouldRetryOnError(error as Error)) {
         throw error;
       }