refactor(core): Unify retry logic and remove schema depth check (#10453)

2026-04-24 20:14:44 -07:00 · 2025-10-03 09:08:27 -07:00
parent 667ca6d241
commit 3b92f127a9
4 changed files with 92 additions and 102 deletions
@@ -6,6 +6,7 @@

 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { ApiError } from '@google/genai';
 import type { HttpError } from './retry.js';
 import { retryWithBackoff } from './retry.js';
 import { setSimulate429 } from './testUtils.js';
@@ -80,22 +81,13 @@ describe('retryWithBackoff', () => {
      initialDelayMs: 10,
    });

-    // 2. IMPORTANT: Attach the rejection expectation to the promise *immediately*.
-    //    This ensures a 'catch' handler is present before the promise can reject.
-    //    The result is a new promise that resolves when the assertion is met.
-    // eslint-disable-next-line vitest/valid-expect
-    const assertionPromise = expect(promise).rejects.toThrow(
-      'Simulated error attempt 3',
-    );
+    // 2. Run timers and await expectation in parallel.
+    await Promise.all([
+      expect(promise).rejects.toThrow('Simulated error attempt 3'),
+      vi.runAllTimersAsync(),
+    ]);

-    // 3. Now, advance the timers. This will trigger the retries and the
-    //    eventual rejection. The handler attached in step 2 will catch it.
-    await vi.runAllTimersAsync();
-
-    // 4. Await the assertion promise itself to ensure the test was successful.
-    await assertionPromise;
-
-    // 5. Finally, assert the number of calls.
+    // 3. Finally, assert the number of calls.
    expect(mockFn).toHaveBeenCalledTimes(3);
  });

@@ -106,12 +98,10 @@ describe('retryWithBackoff', () => {
    const promise = retryWithBackoff(mockFn);

    // Expect it to fail with the error from the 5th attempt.
-    // eslint-disable-next-line vitest/valid-expect
-    const assertionPromise = expect(promise).rejects.toThrow(
-      'Simulated error attempt 5',
-    );
-    await vi.runAllTimersAsync();
-    await assertionPromise;
+    await Promise.all([
+      expect(promise).rejects.toThrow('Simulated error attempt 5'),
+      vi.runAllTimersAsync(),
+    ]);

    expect(mockFn).toHaveBeenCalledTimes(5);
  });
@@ -123,12 +113,10 @@ describe('retryWithBackoff', () => {
    const promise = retryWithBackoff(mockFn, { maxAttempts: undefined });

    // Expect it to fail with the error from the 5th attempt.
-    // eslint-disable-next-line vitest/valid-expect
-    const assertionPromise = expect(promise).rejects.toThrow(
-      'Simulated error attempt 5',
-    );
-    await vi.runAllTimersAsync();
-    await assertionPromise;
+    await Promise.all([
+      expect(promise).rejects.toThrow('Simulated error attempt 5'),
+      vi.runAllTimersAsync(),
+    ]);

    expect(mockFn).toHaveBeenCalledTimes(5);
  });
@@ -161,7 +149,38 @@ describe('retryWithBackoff', () => {
    expect(mockFn).not.toHaveBeenCalled();
  });

-  it('should use default shouldRetry if not provided, retrying on 429', async () => {
+  it('should use default shouldRetry if not provided, retrying on ApiError 429', async () => {
+    const mockFn = vi.fn(async () => {
+      throw new ApiError({ message: 'Too Many Requests', status: 429 });
+    });
+
+    const promise = retryWithBackoff(mockFn, {
+      maxAttempts: 2,
+      initialDelayMs: 10,
+    });
+
+    await Promise.all([
+      expect(promise).rejects.toThrow('Too Many Requests'),
+      vi.runAllTimersAsync(),
+    ]);
+
+    expect(mockFn).toHaveBeenCalledTimes(2);
+  });
+
+  it('should use default shouldRetry if not provided, not retrying on ApiError 400', async () => {
+    const mockFn = vi.fn(async () => {
+      throw new ApiError({ message: 'Bad Request', status: 400 });
+    });
+
+    const promise = retryWithBackoff(mockFn, {
+      maxAttempts: 2,
+      initialDelayMs: 10,
+    });
+    await expect(promise).rejects.toThrow('Bad Request');
+    expect(mockFn).toHaveBeenCalledTimes(1);
+  });
+
+  it('should use default shouldRetry if not provided, retrying on generic error with status 429', async () => {
    const mockFn = vi.fn(async () => {
      const error = new Error('Too Many Requests') as any;
      error.status = 429;
@@ -173,20 +192,16 @@ describe('retryWithBackoff', () => {
      initialDelayMs: 10,
    });

-    // Attach the rejection expectation *before* running timers
-    const assertionPromise =
-      expect(promise).rejects.toThrow('Too Many Requests'); // eslint-disable-line vitest/valid-expect
-
-    // Run timers to trigger retries and eventual rejection
-    await vi.runAllTimersAsync();
-
-    // Await the assertion
-    await assertionPromise;
+    // Run timers and await expectation in parallel.
+    await Promise.all([
+      expect(promise).rejects.toThrow('Too Many Requests'),
+      vi.runAllTimersAsync(),
+    ]);

    expect(mockFn).toHaveBeenCalledTimes(2);
  });

-  it('should use default shouldRetry if not provided, not retrying on 400', async () => {
+  it('should use default shouldRetry if not provided, not retrying on generic error with status 400', async () => {
    const mockFn = vi.fn(async () => {
      const error = new Error('Bad Request') as any;
      error.status = 400;
@@ -242,11 +257,11 @@ describe('retryWithBackoff', () => {

    // We expect rejections as mockFn fails 5 times
    const promise1 = runRetry();
-    // Attach the rejection expectation *before* running timers
-    // eslint-disable-next-line vitest/valid-expect
-    const assertionPromise1 = expect(promise1).rejects.toThrow();
-    await vi.runAllTimersAsync(); // Advance for the delay in the first runRetry
-    await assertionPromise1;
+    // Run timers and await expectation in parallel.
+    await Promise.all([
+      expect(promise1).rejects.toThrow(),
+      vi.runAllTimersAsync(),
+    ]);

    const firstDelaySet = setTimeoutSpy.mock.calls.map(
      (call) => call[1] as number,
@@ -257,11 +272,11 @@ describe('retryWithBackoff', () => {
    mockFn = createFailingFunction(5); // Re-initialize with 5 failures

    const promise2 = runRetry();
-    // Attach the rejection expectation *before* running timers
-    // eslint-disable-next-line vitest/valid-expect
-    const assertionPromise2 = expect(promise2).rejects.toThrow();
-    await vi.runAllTimersAsync(); // Advance for the delay in the second runRetry
-    await assertionPromise2;
+    // Run timers and await expectation in parallel.
+    await Promise.all([
+      expect(promise2).rejects.toThrow(),
+      vi.runAllTimersAsync(),
+    ]);

    const secondDelaySet = setTimeoutSpy.mock.calls.map(
      (call) => call[1] as number,
@@ -5,6 +5,7 @@
 */

 import type { GenerateContentResponse } from '@google/genai';
+import { ApiError } from '@google/genai';
 import { AuthType } from '../core/contentGenerator.js';
 import {
  isProQuotaExceededError,
@@ -42,17 +43,19 @@ const DEFAULT_RETRY_OPTIONS: RetryOptions = {
 * @returns True if the error is a transient error, false otherwise.
 */
 function defaultShouldRetry(error: Error | unknown): boolean {
-  // Check for common transient error status codes either in message or a status property
-  if (error && typeof (error as { status?: number }).status === 'number') {
-    const status = (error as { status: number }).status;
-    if (status === 429 || (status >= 500 && status < 600)) {
-      return true;
-    }
+  // Priority check for ApiError
+  if (error instanceof ApiError) {
+    // Explicitly do not retry 400 (Bad Request)
+    if (error.status === 400) return false;
+    return error.status === 429 || (error.status >= 500 && error.status < 600);
  }
-  if (error instanceof Error && error.message) {
-    if (error.message.includes('429')) return true;
-    if (error.message.match(/5\d{2}/)) return true;
+
+  // Check for status using helper (handles other error shapes)
+  const status = getErrorStatus(error);
+  if (status !== undefined) {
+    return status === 429 || (status >= 500 && status < 600);
  }
+
  return false;
 }