From ac4a79223abb108b01b4daab0945737987983097 Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Mon, 29 Sep 2025 12:27:15 -0700 Subject: [PATCH] feat(core): Add content-based retries for JSON generation (#9264) --- packages/core/src/core/baseLlmClient.test.ts | 109 ++++++++++++++---- packages/core/src/core/baseLlmClient.ts | 63 +++++----- packages/core/src/core/geminiChat.test.ts | 5 +- packages/core/src/core/geminiChat.ts | 2 +- packages/core/src/utils/flashFallback.test.ts | 4 +- packages/core/src/utils/retry.test.ts | 5 +- packages/core/src/utils/retry.ts | 26 ++++- 7 files changed, 145 insertions(+), 69 deletions(-) diff --git a/packages/core/src/core/baseLlmClient.test.ts b/packages/core/src/core/baseLlmClient.test.ts index c7671e86fd..467bcb86b8 100644 --- a/packages/core/src/core/baseLlmClient.test.ts +++ b/packages/core/src/core/baseLlmClient.test.ts @@ -36,7 +36,29 @@ vi.mock('../utils/errors.js', async (importOriginal) => { }); vi.mock('../utils/retry.js', () => ({ - retryWithBackoff: vi.fn(async (fn) => await fn()), + retryWithBackoff: vi.fn(async (fn, options) => { + // Default implementation - just call the function + const result = await fn(); + + // If shouldRetryOnContent is provided, test it but don't actually retry + // (unless we want to simulate retry exhaustion for testing) + if (options?.shouldRetryOnContent) { + const shouldRetry = options.shouldRetryOnContent(result); + if (shouldRetry) { + // Check if we need to simulate retry exhaustion (for error testing) + const responseText = result?.candidates?.[0]?.content?.parts?.[0]?.text; + if ( + !responseText || + responseText.trim() === '' || + responseText.includes('{"color": "blue"') + ) { + throw new Error('Retry attempts exhausted for invalid content'); + } + } + } + + return result; + }), })); const mockGenerateContent = vi.fn(); @@ -96,8 +118,14 @@ describe('BaseLlmClient', () => { expect(result).toEqual({ color: 'blue' }); - // Ensure the retry mechanism was engaged + // Ensure the retry mechanism was engaged with shouldRetryOnContent expect(retryWithBackoff).toHaveBeenCalledTimes(1); + expect(retryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + shouldRetryOnContent: expect.any(Function), + }), + ); // Validate the parameters passed to the underlying generator expect(mockGenerateContent).toHaveBeenCalledTimes(1); @@ -194,9 +222,12 @@ describe('BaseLlmClient', () => { await client.generateJson(options); expect(retryWithBackoff).toHaveBeenCalledTimes(1); - expect(retryWithBackoff).toHaveBeenCalledWith(expect.any(Function), { - maxAttempts: customMaxAttempts, - }); + expect(retryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + maxAttempts: customMaxAttempts, + }), + ); }); it('should call retryWithBackoff without maxAttempts when not provided', async () => { @@ -206,9 +237,44 @@ describe('BaseLlmClient', () => { // No maxAttempts in defaultOptions await client.generateJson(defaultOptions); - expect(retryWithBackoff).toHaveBeenCalledWith(expect.any(Function), { - maxAttempts: 5, - }); + expect(retryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + maxAttempts: 5, + }), + ); + }); + }); + + describe('generateJson - Content Validation and Retries', () => { + it('should validate content using shouldRetryOnContent function', async () => { + const mockResponse = createMockResponse('{"color": "blue"}'); + mockGenerateContent.mockResolvedValue(mockResponse); + + await client.generateJson(defaultOptions); + + // Verify that retryWithBackoff was called with shouldRetryOnContent + expect(retryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + shouldRetryOnContent: expect.any(Function), + }), + ); + + // Test the shouldRetryOnContent function behavior + const retryCall = vi.mocked(retryWithBackoff).mock.calls[0]; + const shouldRetryOnContent = retryCall[1]?.shouldRetryOnContent; + + // Valid JSON should not trigger retry + expect(shouldRetryOnContent!(mockResponse)).toBe(false); + + // Empty response should trigger retry + expect(shouldRetryOnContent!(createMockResponse(''))).toBe(true); + + // Invalid JSON should trigger retry + expect( + shouldRetryOnContent!(createMockResponse('{"color": "blue"')), + ).toBe(true); }); }); @@ -222,14 +288,14 @@ describe('BaseLlmClient', () => { const result = await client.generateJson(defaultOptions); expect(result).toEqual({ color: 'purple' }); - expect(logMalformedJsonResponse).toHaveBeenCalledTimes(1); expect(logMalformedJsonResponse).toHaveBeenCalledWith( mockConfig, expect.any(MalformedJsonResponseEvent), ); - // Validate the telemetry event content - const event = vi.mocked(logMalformedJsonResponse).mock - .calls[0][1] as MalformedJsonResponseEvent; + // Validate the telemetry event content - find the most recent call + const calls = vi.mocked(logMalformedJsonResponse).mock.calls; + const lastCall = calls[calls.length - 1]; + const event = lastCall[1] as MalformedJsonResponseEvent; expect(event.model).toBe('test-model'); }); @@ -247,38 +313,37 @@ describe('BaseLlmClient', () => { }); describe('generateJson - Error Handling', () => { - it('should throw and report error for empty response', async () => { + it('should throw and report error for empty response after retry exhaustion', async () => { mockGenerateContent.mockResolvedValue(createMockResponse('')); - // The final error message includes the prefix added by the client's outer catch block. await expect(client.generateJson(defaultOptions)).rejects.toThrow( - 'Failed to generate JSON content: API returned an empty response for generateJson.', + 'Failed to generate JSON content: Retry attempts exhausted for invalid content', ); // Verify error reporting details expect(reportError).toHaveBeenCalledTimes(1); expect(reportError).toHaveBeenCalledWith( expect.any(Error), - 'Error in generateJson: API returned an empty response.', + 'API returned invalid content (empty or unparsable JSON) after all retries.', defaultOptions.contents, - 'generateJson-empty-response', + 'generateJson-invalid-content', ); }); - it('should throw and report error for invalid JSON syntax', async () => { + it('should throw and report error for invalid JSON syntax after retry exhaustion', async () => { const invalidJson = '{"color": "blue"'; // missing closing brace mockGenerateContent.mockResolvedValue(createMockResponse(invalidJson)); await expect(client.generateJson(defaultOptions)).rejects.toThrow( - /^Failed to generate JSON content: Failed to parse API response as JSON:/, + 'Failed to generate JSON content: Retry attempts exhausted for invalid content', ); expect(reportError).toHaveBeenCalledTimes(1); expect(reportError).toHaveBeenCalledWith( expect.any(Error), - 'Failed to parse JSON response from generateJson.', - expect.objectContaining({ responseTextFailedToParse: invalidJson }), - 'generateJson-parse', + 'API returned invalid content (empty or unparsable JSON) after all retries.', + defaultOptions.contents, + 'generateJson-invalid-content', ); }); diff --git a/packages/core/src/core/baseLlmClient.ts b/packages/core/src/core/baseLlmClient.ts index 4dc86eddd5..e1b02164a8 100644 --- a/packages/core/src/core/baseLlmClient.ts +++ b/packages/core/src/core/baseLlmClient.ts @@ -9,6 +9,7 @@ import type { GenerateContentConfig, Part, EmbedContentParameters, + GenerateContentResponse, } from '@google/genai'; import type { Config } from '../config/config.js'; import type { ContentGenerator } from './contentGenerator.js'; @@ -107,54 +108,44 @@ export class BaseLlmClient { promptId, ); + const shouldRetryOnContent = (response: GenerateContentResponse) => { + const text = getResponseText(response)?.trim(); + if (!text) { + return true; // Retry on empty response + } + try { + JSON.parse(this.cleanJsonResponse(text, model)); + return false; + } catch (_e) { + return true; + } + }; + const result = await retryWithBackoff(apiCall, { + shouldRetryOnContent, maxAttempts: maxAttempts ?? DEFAULT_MAX_ATTEMPTS, }); - let text = getResponseText(result)?.trim(); - if (!text) { - const error = new Error( - 'API returned an empty response for generateJson.', - ); - await reportError( - error, - 'Error in generateJson: API returned an empty response.', - contents, - 'generateJson-empty-response', - ); - throw error; - } - - text = this.cleanJsonResponse(text, model); - - try { - return JSON.parse(text); - } catch (parseError) { - const error = new Error( - `Failed to parse API response as JSON: ${getErrorMessage(parseError)}`, - ); - await reportError( - parseError, - 'Failed to parse JSON response from generateJson.', - { - responseTextFailedToParse: text, - originalRequestContents: contents, - }, - 'generateJson-parse', - ); - throw error; - } + // If we are here, the content is valid (not empty and parsable). + return JSON.parse( + this.cleanJsonResponse(getResponseText(result)!.trim(), model), + ); } catch (error) { if (abortSignal.aborted) { throw error; } + // Check if the error is from exhausting retries, and report accordingly. if ( error instanceof Error && - (error.message === 'API returned an empty response for generateJson.' || - error.message.startsWith('Failed to parse API response as JSON:')) + error.message.includes('Retry attempts exhausted') ) { - // We perform this check so that we don't report these again. + await reportError( + error, + 'API returned invalid content (empty or unparsable JSON) after all retries.', + contents, + 'generateJson-invalid-content', + ); } else { await reportError( error, diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 067f0d160c..55d0af1400 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -911,7 +911,10 @@ describe('GeminiChat', () => { try { return await apiCall(); } catch (error) { - if (options?.shouldRetry && options.shouldRetry(error)) { + if ( + options?.shouldRetryOnError && + options.shouldRetryOnError(error) + ) { // Try again return await apiCall(); } diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 7a9d083dad..18353ced0e 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -376,7 +376,7 @@ export class GeminiChat { ) => await handleFallback(this.config, model, authType, error); const streamResponse = await retryWithBackoff(apiCall, { - shouldRetry: (error: unknown) => { + shouldRetryOnError: (error: unknown) => { if (error instanceof ApiError && error.message) { if (error.status === 400) return false; if (isSchemaDepthError(error.message)) return false; diff --git a/packages/core/src/utils/flashFallback.test.ts b/packages/core/src/utils/flashFallback.test.ts index 6d4330f1ad..8ef9665f42 100644 --- a/packages/core/src/utils/flashFallback.test.ts +++ b/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; }, diff --git a/packages/core/src/utils/retry.test.ts b/packages/core/src/utils/retry.test.ts index d310327827..6a011f9a7a 100644 --- a/packages/core/src/utils/retry.test.ts +++ b/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, }); diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts index 0e88fd74e3..54d91f7d2d 100644 --- a/packages/core/src/utils/retry.ts +++ b/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( maxDelayMs, onPersistent429, authType, - shouldRetry, + shouldRetryOnError, + shouldRetryOnContent, } = { ...DEFAULT_RETRY_OPTIONS, ...cleanOptions, @@ -101,7 +104,20 @@ export async function retryWithBackoff( 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( } // 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; }