feat: implement dynamic retry timing for Google Gemini API

- Enhanced error handling in GeminiHandler to preserve original error structure
- Added logic to distinguish between rate limits and quota exhaustion
- Preserved retry delay information from Gemini API responses
- Added new localized error messages for better user feedback
- Added comprehensive tests for the new retry logic

Fixes #6680
This commit is contained in:
Roo Code 2025-08-04 19:19:50 +00:00
parent 7ca4901024
commit ea054e81ae
3 changed files with 139 additions and 3 deletions

View file

@ -102,6 +102,67 @@ describe("GeminiHandler", () => {
}
}).rejects.toThrow()
})
it("should handle rate limit errors with retry information", async () => {
const mockError: any = new Error("Rate limit exceeded")
mockError.status = 429
mockError.errorDetails = [
{
"@type": "type.googleapis.com/google.rpc.RetryInfo",
retryDelay: "40s",
},
]
;(handler["client"].models.generateContentStream as any).mockRejectedValue(mockError)
const stream = handler.createMessage(systemPrompt, mockMessages)
await expect(async () => {
for await (const _chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow(t("common:errors.gemini.rate_limit"))
})
it("should handle quota exhaustion errors", async () => {
const mockError: any = new Error(
"You exceeded your current quota, please check your plan and billing details",
)
mockError.status = 429
mockError.message = "You exceeded your current quota, please check your plan and billing details"
;(handler["client"].models.generateContentStream as any).mockRejectedValue(mockError)
const stream = handler.createMessage(systemPrompt, mockMessages)
await expect(async () => {
for await (const _chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow(t("common:errors.gemini.quota_exhausted"))
})
it("should preserve error details for retry logic", async () => {
const mockError: any = new Error("Rate limit exceeded")
mockError.status = 429
mockError.errorDetails = [
{
"@type": "type.googleapis.com/google.rpc.RetryInfo",
retryDelay: "60s",
},
]
;(handler["client"].models.generateContentStream as any).mockRejectedValue(mockError)
const stream = handler.createMessage(systemPrompt, mockMessages)
try {
for await (const _chunk of stream) {
// Should throw before yielding any chunks
}
} catch (error: any) {
expect(error.status).toBe(429)
expect(error.errorDetails).toBeDefined()
expect(error.errorDetails[0].retryDelay).toBe("60s")
}
})
})
describe("completePrompt", () => {
@ -134,6 +195,33 @@ describe("GeminiHandler", () => {
)
})
it("should handle rate limit errors in completePrompt", async () => {
const mockError: any = new Error("Rate limit exceeded")
mockError.status = 429
mockError.errorDetails = [
{
"@type": "type.googleapis.com/google.rpc.RetryInfo",
retryDelay: "30s",
},
]
;(handler["client"].models.generateContent as any).mockRejectedValue(mockError)
await expect(handler.completePrompt("Test prompt")).rejects.toThrow(t("common:errors.gemini.rate_limit"))
})
it("should handle quota exhaustion errors in completePrompt", async () => {
const mockError: any = new Error(
"You exceeded your current quota, please check your plan and billing details",
)
mockError.status = 429
mockError.message = "You exceeded your current quota, please check your plan and billing details"
;(handler["client"].models.generateContent as any).mockRejectedValue(mockError)
await expect(handler.completePrompt("Test prompt")).rejects.toThrow(
t("common:errors.gemini.quota_exhausted"),
)
})
it("should handle empty response", async () => {
// Mock the response with empty text
;(handler["client"].models.generateContent as any).mockResolvedValue({

View file

@ -153,7 +153,30 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
totalCost: this.calculateCost({ info, inputTokens, outputTokens, cacheReadTokens }),
}
}
} catch (error) {
} catch (error: any) {
// Preserve the original error structure for retry logic
if (error.status === 429) {
// Check if this is a rate limit or quota exhaustion
const errorMessage = error.message || ""
const isQuotaExhausted = errorMessage.includes("quota") || errorMessage.includes("billing")
// Create an enhanced error that preserves the original structure
const enhancedError = new Error(
isQuotaExhausted
? t("common:errors.gemini.quota_exhausted", { error: errorMessage })
: t("common:errors.gemini.rate_limit", { error: errorMessage }),
)
// Preserve the original error properties for retry logic
Object.assign(enhancedError, {
status: error.status,
errorDetails: error.errorDetails,
message: error.message,
})
throw enhancedError
}
if (error instanceof Error) {
throw new Error(t("common:errors.gemini.generate_stream", { error: error.message }))
}
@ -235,7 +258,30 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
}
return text
} catch (error) {
} catch (error: any) {
// Preserve the original error structure for retry logic
if (error.status === 429) {
// Check if this is a rate limit or quota exhaustion
const errorMessage = error.message || ""
const isQuotaExhausted = errorMessage.includes("quota") || errorMessage.includes("billing")
// Create an enhanced error that preserves the original structure
const enhancedError = new Error(
isQuotaExhausted
? t("common:errors.gemini.quota_exhausted", { error: errorMessage })
: t("common:errors.gemini.rate_limit", { error: errorMessage }),
)
// Preserve the original error properties for retry logic
Object.assign(enhancedError, {
status: error.status,
errorDetails: error.errorDetails,
message: error.message,
})
throw enhancedError
}
if (error instanceof Error) {
throw new Error(t("common:errors.gemini.generate_complete_prompt", { error: error.message }))
}

View file

@ -89,7 +89,9 @@
"gemini": {
"generate_stream": "Gemini generate context stream error: {{error}}",
"generate_complete_prompt": "Gemini completion error: {{error}}",
"sources": "Sources:"
"sources": "Sources:",
"rate_limit": "Gemini API rate limit exceeded. The system will automatically retry with the appropriate delay.",
"quota_exhausted": "Gemini API quota exhausted: {{error}}. Please check your plan and billing details at https://ai.google.dev/gemini-api/docs/rate-limits"
},
"cerebras": {
"authenticationFailed": "Cerebras API authentication failed. Please check your API key is valid and not expired.",