From bd5eb8fcaef6e03ee6244ede331a892ce398952b Mon Sep 17 00:00:00 2001 From: Michael Date: Thu, 6 Feb 2025 08:11:01 +0100 Subject: [PATCH] feat: add retry decorator with rate limit handling (#1605) * fix: improve retry decorator with smart rate limit handling - Add handling of rate limit (429) errors - Implement retry timing based on response headers - Add exponential backoff when no headers present - Add a few unit tests Fixes #713 * Create modern-knives-tan.md * Improve readability in retry.ts --------- Co-authored-by: Michael Overhorst Co-authored-by: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com> --- .changeset/modern-knives-tan.md | 5 + src/api/providers/anthropic.ts | 2 + src/api/providers/deepseek.ts | 2 + src/api/providers/gemini.ts | 2 + src/api/providers/mistral.ts | 2 + src/api/providers/openai-native.ts | 2 + src/api/providers/openai.ts | 2 + src/api/providers/openrouter.ts | 2 + src/api/providers/vertex.ts | 2 + src/api/retry.test.ts | 216 +++++++++++++++++++++++++++++ src/api/retry.ts | 62 +++++++++ 11 files changed, 299 insertions(+) create mode 100644 .changeset/modern-knives-tan.md create mode 100644 src/api/retry.test.ts create mode 100644 src/api/retry.ts diff --git a/.changeset/modern-knives-tan.md b/.changeset/modern-knives-tan.md new file mode 100644 index 0000000000..d2aaa250b3 --- /dev/null +++ b/.changeset/modern-knives-tan.md @@ -0,0 +1,5 @@ +--- +"claude-dev": patch +--- + +Add automatic retry for rate limited requests diff --git a/src/api/providers/anthropic.ts b/src/api/providers/anthropic.ts index 8c3fd1b87d..3aff867b7f 100644 --- a/src/api/providers/anthropic.ts +++ b/src/api/providers/anthropic.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming" +import { withRetry } from "../retry" import { anthropicDefaultModelId, AnthropicModelId, anthropicModels, ApiHandlerOptions, ModelInfo } from "../../shared/api" import { ApiHandler } from "../index" import { ApiStream } from "../transform/stream" @@ -16,6 +17,7 @@ export class AnthropicHandler implements ApiHandler { }) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const model = this.getModel() let stream: AnthropicStream diff --git a/src/api/providers/deepseek.ts b/src/api/providers/deepseek.ts index 763e1ae68f..9049e646db 100644 --- a/src/api/providers/deepseek.ts +++ b/src/api/providers/deepseek.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" +import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, DeepSeekModelId, ModelInfo, deepSeekDefaultModelId, deepSeekModels } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" @@ -18,6 +19,7 @@ export class DeepSeekHandler implements ApiHandler { }) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const model = this.getModel() diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index 39c55548d1..b452af6dbb 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import { GoogleGenerativeAI } from "@google/generative-ai" +import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "../../shared/api" import { convertAnthropicMessageToGemini } from "../transform/gemini-format" @@ -17,6 +18,7 @@ export class GeminiHandler implements ApiHandler { this.client = new GoogleGenerativeAI(options.geminiApiKey) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const model = this.client.getGenerativeModel({ model: this.getModel().id, diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index 5be89b18f3..28c2331c60 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import { Mistral } from "@mistralai/mistralai" +import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, @@ -26,6 +27,7 @@ export class MistralHandler implements ApiHandler { }) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const stream = await this.client.chat.stream({ model: this.getModel().id, diff --git a/src/api/providers/openai-native.ts b/src/api/providers/openai-native.ts index 8a47ec4345..63f30bd61c 100644 --- a/src/api/providers/openai-native.ts +++ b/src/api/providers/openai-native.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" +import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, @@ -22,6 +23,7 @@ export class OpenAiNativeHandler implements ApiHandler { }) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { switch (this.getModel().id) { case "o1": diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index fd73abb567..c03b1d13ec 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI, { AzureOpenAI } from "openai" +import { withRetry } from "../retry" import { ApiHandlerOptions, azureOpenAiDefaultApiVersion, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api" import { ApiHandler } from "../index" import { convertToOpenAiMessages } from "../transform/openai-format" @@ -27,6 +28,7 @@ export class OpenAiHandler implements ApiHandler { } } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const modelId = this.options.openAiModelId ?? "" const isDeepseekReasoner = modelId.includes("deepseek-reasoner") diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index 6b8f40c5e7..45cbd2c7ac 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -2,6 +2,7 @@ import { Anthropic } from "@anthropic-ai/sdk" import axios from "axios" import delay from "delay" import OpenAI from "openai" +import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, ModelInfo, openRouterDefaultModelId, openRouterDefaultModelInfo } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" @@ -24,6 +25,7 @@ export class OpenRouterHandler implements ApiHandler { }) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const model = this.getModel() diff --git a/src/api/providers/vertex.ts b/src/api/providers/vertex.ts index c8f1efd873..286562ed45 100644 --- a/src/api/providers/vertex.ts +++ b/src/api/providers/vertex.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import { AnthropicVertex } from "@anthropic-ai/vertex-sdk" +import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "../../shared/api" import { ApiStream } from "../transform/stream" @@ -18,6 +19,7 @@ export class VertexHandler implements ApiHandler { }) } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const stream = await this.client.messages.create({ model: this.getModel().id, diff --git a/src/api/retry.test.ts b/src/api/retry.test.ts new file mode 100644 index 0000000000..43b8eaf3e9 --- /dev/null +++ b/src/api/retry.test.ts @@ -0,0 +1,216 @@ +import { describe, it } from "mocha" +import "should" +import { withRetry } from "./retry" + +describe("Retry Decorator", () => { + describe("withRetry", () => { + it("should not retry on success", async () => { + let callCount = 0 + class TestClass { + @withRetry() + async *successMethod() { + callCount++ + yield "success" + } + } + + const test = new TestClass() + const result = [] + for await (const value of test.successMethod()) { + result.push(value) + } + + callCount.should.equal(1) + result.should.deepEqual(["success"]) + }) + + it("should retry on rate limit (429) error", async () => { + let callCount = 0 + class TestClass { + @withRetry({ maxRetries: 2, baseDelay: 10, maxDelay: 100 }) + async *failMethod() { + callCount++ + if (callCount === 1) { + const error: any = new Error("Rate limit exceeded") + error.status = 429 + throw error + } + yield "success after retry" + } + } + + const test = new TestClass() + const result = [] + for await (const value of test.failMethod()) { + result.push(value) + } + + callCount.should.equal(2) + result.should.deepEqual(["success after retry"]) + }) + + it("should not retry on non-rate-limit errors", async () => { + let callCount = 0 + class TestClass { + @withRetry() + async *failMethod() { + callCount++ + throw new Error("Regular error") + } + } + + const test = new TestClass() + try { + for await (const _ of test.failMethod()) { + // Should not reach here + } + throw new Error("Should have thrown") + } catch (error: any) { + error.message.should.equal("Regular error") + callCount.should.equal(1) + } + }) + + it("should respect retry-after header with delta seconds", async () => { + let callCount = 0 + const startTime = Date.now() + class TestClass { + @withRetry({ maxRetries: 2, baseDelay: 1000 }) // Use large baseDelay to ensure header takes precedence + async *failMethod() { + callCount++ + if (callCount === 1) { + const error: any = new Error("Rate limit exceeded") + error.status = 429 + error.headers = { "retry-after": "0.01" } // 10ms delay + throw error + } + yield "success after retry" + } + } + + const test = new TestClass() + const result = [] + for await (const value of test.failMethod()) { + result.push(value) + } + + const duration = Date.now() - startTime + duration.should.be.approximately(10, 10) // Allow 10ms variance + callCount.should.equal(2) + result.should.deepEqual(["success after retry"]) + }) + + it("should respect retry-after header with Unix timestamp", async () => { + let callCount = 0 + const startTime = Date.now() + const retryTimestamp = Math.floor(Date.now() / 1000) + 0.01 // 10ms in the future + + class TestClass { + @withRetry({ maxRetries: 2, baseDelay: 1000 }) // Use large baseDelay to ensure header takes precedence + async *failMethod() { + callCount++ + if (callCount === 1) { + const error: any = new Error("Rate limit exceeded") + error.status = 429 + error.headers = { "retry-after": retryTimestamp.toString() } + throw error + } + yield "success after retry" + } + } + + const test = new TestClass() + const result = [] + for await (const value of test.failMethod()) { + result.push(value) + } + + const duration = Date.now() - startTime + duration.should.be.approximately(10, 10) // Allow 10ms variance + callCount.should.equal(2) + result.should.deepEqual(["success after retry"]) + }) + + it("should use exponential backoff when no retry-after header", async () => { + let callCount = 0 + const startTime = Date.now() + class TestClass { + @withRetry({ maxRetries: 2, baseDelay: 10, maxDelay: 100 }) + async *failMethod() { + callCount++ + if (callCount === 1) { + const error: any = new Error("Rate limit exceeded") + error.status = 429 + throw error + } + yield "success after retry" + } + } + + const test = new TestClass() + const result = [] + for await (const value of test.failMethod()) { + result.push(value) + } + + const duration = Date.now() - startTime + // First retry should be after baseDelay (10ms) + duration.should.be.approximately(10, 10) + callCount.should.equal(2) + result.should.deepEqual(["success after retry"]) + }) + + it("should respect maxDelay", async () => { + let callCount = 0 + const startTime = Date.now() + class TestClass { + @withRetry({ maxRetries: 3, baseDelay: 50, maxDelay: 10 }) + async *failMethod() { + callCount++ + if (callCount < 3) { + const error: any = new Error("Rate limit exceeded") + error.status = 429 + throw error + } + yield "success after retries" + } + } + + const test = new TestClass() + const result = [] + for await (const value of test.failMethod()) { + result.push(value) + } + + const duration = Date.now() - startTime + // Both retries should be capped at maxDelay (10ms each) + duration.should.be.approximately(20, 20) + callCount.should.equal(3) + result.should.deepEqual(["success after retries"]) + }) + + it("should throw after maxRetries attempts", async () => { + let callCount = 0 + class TestClass { + @withRetry({ maxRetries: 2, baseDelay: 10 }) + async *failMethod() { + callCount++ + const error: any = new Error("Rate limit exceeded") + error.status = 429 + throw error + } + } + + const test = new TestClass() + try { + for await (const _ of test.failMethod()) { + // Should not reach here + } + throw new Error("Should have thrown") + } catch (error: any) { + error.message.should.equal("Rate limit exceeded") + callCount.should.equal(2) // Initial attempt + 1 retry + } + }) + }) +}) diff --git a/src/api/retry.ts b/src/api/retry.ts new file mode 100644 index 0000000000..deeabfb365 --- /dev/null +++ b/src/api/retry.ts @@ -0,0 +1,62 @@ +interface RetryOptions { + maxRetries?: number + baseDelay?: number + maxDelay?: number +} + +const DEFAULT_OPTIONS: Required = { + maxRetries: 3, + baseDelay: 1_000, + maxDelay: 10_000, +} + +export function withRetry(options: RetryOptions = {}) { + const { maxRetries, baseDelay, maxDelay } = { ...DEFAULT_OPTIONS, ...options } + + return function (_target: any, _propertyKey: string, descriptor: PropertyDescriptor) { + const originalMethod = descriptor.value + + descriptor.value = async function* (...args: any[]) { + for (let attempt = 0; attempt < maxRetries; attempt++) { + try { + yield* originalMethod.apply(this, args) + return + } catch (error: any) { + const isRateLimit = error?.status === 429 + const isLastAttempt = attempt === maxRetries - 1 + + if (!isRateLimit || isLastAttempt) { + throw error + } + + // Get retry delay from header or calculate exponential backoff + // Check various rate limit headers + const retryAfter = + error.headers?.["retry-after"] || + error.headers?.["x-ratelimit-reset"] || + error.headers?.["ratelimit-reset"] + + let delay: number + if (retryAfter) { + // Handle both delta-seconds and Unix timestamp formats + const retryValue = parseInt(retryAfter, 10) + if (retryValue > Date.now() / 1000) { + // Unix timestamp + delay = retryValue * 1000 - Date.now() + } else { + // Delta seconds + delay = retryValue * 1000 + } + } else { + // Use exponential backoff if no header + delay = Math.min(maxDelay, baseDelay * Math.pow(2, attempt)) + } + + await new Promise((resolve) => setTimeout(resolve, delay)) + } + } + } + + return descriptor + } +}