MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-14 08:01:02 -07:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Sandy Tao
2025-09-29 12:27:15 -07:00
committed by GitHub
parent 042288e72c
commit ac4a79223a
7 changed files with 145 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
packages/core/src/core/baseLlmClient.test.ts
View File

@@ -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',
);
});

63
packages/core/src/core/baseLlmClient.ts
View File

@@ -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,

5
packages/core/src/core/geminiChat.test.ts
View File

@@ -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();
}

2
packages/core/src/core/geminiChat.ts
View File

@@ -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;

4
packages/core/src/utils/flashFallback.test.ts
View File

@@ -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;
},

5
packages/core/src/utils/retry.test.ts
View File

@@ -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,
});

26
packages/core/src/utils/retry.ts
View File

@@ -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;
}