Retry on provider failures

This commit is contained in:
Matt Rubens 2024-12-14 22:53:19 -05:00
parent 130c5ff779
commit c9e3e60102
17 changed files with 1112 additions and 476 deletions

View file

@ -1,5 +1,9 @@
# Roo Cline Changelog
## [2.2.8]
- Exponential backoff and retry on errors from providers
## [2.2.7]
- More fixes to search/replace diffs

View file

@ -13,6 +13,7 @@ A fork of Cline, an autonomous coding agent, optimized for speed and flexibility
- Support for dragging and dropping images into chats
- Support for auto-approving MCP tools
- Support for enabling/disabling MCP servers
- Exponential backoff and retry on errors from providers
## Disclaimer

4
package-lock.json generated
View file

@ -1,12 +1,12 @@
{
"name": "roo-cline",
"version": "2.2.7",
"version": "2.2.8",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "roo-cline",
"version": "2.2.7",
"version": "2.2.8",
"dependencies": {
"@anthropic-ai/bedrock-sdk": "^0.10.2",
"@anthropic-ai/sdk": "^0.26.0",

View file

@ -3,7 +3,7 @@
"displayName": "Roo Cline",
"description": "A fork of Cline, an autonomous coding agent, with some added experimental configuration and automation features.",
"publisher": "RooVeterinaryInc",
"version": "2.2.7",
"version": "2.2.8",
"icon": "assets/icons/rocket.png",
"galleryBanner": {
"color": "#617A91",

View file

@ -8,6 +8,14 @@ import { Anthropic } from '@anthropic-ai/sdk'
jest.mock('openai')
jest.mock('axios')
jest.mock('delay', () => jest.fn(() => Promise.resolve()))
jest.mock('../../utils/retry', () => ({
withRetry: jest.fn().mockImplementation(async function*(fn) {
const generator = await fn();
for await (const chunk of generator) {
yield chunk;
}
})
}))
describe('OpenRouterHandler', () => {
const mockOptions: ApiHandlerOptions = {
@ -41,81 +49,300 @@ describe('OpenRouterHandler', () => {
})
})
test('getModel returns correct model info when options are provided', () => {
const handler = new OpenRouterHandler(mockOptions)
const result = handler.getModel()
expect(result).toEqual({
id: mockOptions.openRouterModelId,
info: mockOptions.openRouterModelInfo
describe('getModel', () => {
test('returns correct model info when options are provided', () => {
const handler = new OpenRouterHandler(mockOptions)
const result = handler.getModel()
expect(result).toEqual({
id: mockOptions.openRouterModelId,
info: mockOptions.openRouterModelInfo
})
})
test('returns default model info when options are not provided', () => {
const handlerWithoutModelOptions = new OpenRouterHandler({
openRouterApiKey: 'test-key'
})
const result = handlerWithoutModelOptions.getModel()
expect(result).toEqual({
id: 'anthropic/claude-3.5-sonnet:beta',
info: expect.any(Object)
})
})
})
test('createMessage generates correct stream chunks', async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
id: 'test-id',
choices: [{
delta: {
content: 'test response'
}
}]
describe('createMessage', () => {
test('applies correct formatting for Claude models', async () => {
const claudeOptions = {
...mockOptions,
openRouterModelId: 'anthropic/claude-3-haiku'
};
const handler = new OpenRouterHandler(claudeOptions);
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
id: 'test-id',
choices: [{
delta: {
content: 'test response'
}
}]
};
}
}
}
// Mock OpenAI chat.completions.create
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
// Mock axios.get for generation details
;(axios.get as jest.Mock).mockResolvedValue({
data: {
data: {
native_tokens_prompt: 10,
native_tokens_completion: 20,
total_cost: 0.001
}
}
})
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }]
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
};
for await (const chunk of generator) {
chunks.push(chunk)
}
// Verify stream chunks
expect(chunks).toHaveLength(2) // One text chunk and one usage chunk
expect(chunks[0]).toEqual({
type: 'text',
text: 'test response'
// Mock OpenAI chat.completions.create
const mockCreate = jest.fn().mockResolvedValue(mockStream);
(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any;
// Mock axios.get for generation details
(axios.get as jest.Mock).mockResolvedValue({
data: {
data: {
native_tokens_prompt: 10,
native_tokens_completion: 20,
total_cost: 0.001
}
}
});
const systemPrompt = 'test system prompt';
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user' as const, content: 'message 1' },
{ role: 'assistant' as const, content: 'response 1' },
{ role: 'user' as const, content: 'message 2' }
];
const generator = handler.createMessage(systemPrompt, messages);
for await (const _ of generator) { /* consume generator */ }
// Verify OpenAI client was called with correct Claude-specific parameters
const callArgs = mockCreate.mock.calls[0][0];
expect(callArgs.model).toBe('anthropic/claude-3-haiku');
expect(callArgs.max_tokens).toBeUndefined();
expect(callArgs.temperature).toBe(0);
expect(callArgs.stream).toBe(true);
// Verify system message has ephemeral cache control
const systemMessage = callArgs.messages[0];
expect(systemMessage.role).toBe('system');
expect(systemMessage.content[0]).toEqual({
type: 'text',
text: systemPrompt,
cache_control: { type: 'ephemeral' }
});
// Verify user messages have ephemeral cache control
const userMessages = callArgs.messages.filter((m: OpenAI.Chat.ChatCompletionMessageParam) => m.role === 'user');
const lastTwoUserMessages = userMessages.slice(-2);
lastTwoUserMessages.forEach((msg: OpenAI.Chat.ChatCompletionMessageParam) => {
expect(Array.isArray(msg.content)).toBe(true);
const content = msg.content as Array<{type: string; text: string; cache_control?: {type: string}}>;
const textParts = content.filter(part => part.type === 'text');
expect(textParts.length).toBeGreaterThan(0);
const lastTextPart = textParts[textParts.length - 1];
expect(lastTextPart.cache_control).toEqual({ type: 'ephemeral' });
});
})
expect(chunks[1]).toEqual({
type: 'usage',
inputTokens: 10,
outputTokens: 20,
totalCost: 0.001,
fullResponseText: 'test response'
test('generates correct stream chunks with default options', async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
id: 'test-id',
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
// Mock OpenAI chat.completions.create
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
// Mock axios.get for generation details
;(axios.get as jest.Mock).mockResolvedValue({
data: {
data: {
native_tokens_prompt: 10,
native_tokens_completion: 20,
total_cost: 0.001
}
}
})
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }]
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
for await (const chunk of generator) {
chunks.push(chunk)
}
// Verify stream chunks
expect(chunks).toHaveLength(3) // text chunk, usage chunk, and final newlines
expect(chunks[0]).toEqual({
type: 'text',
text: 'test response'
})
expect(chunks[1]).toEqual({
type: 'usage',
inputTokens: 10,
outputTokens: 20,
totalCost: 0.001,
fullResponseText: 'test response'
})
expect(chunks[2]).toEqual({
type: 'text',
text: '\n\n'
})
// Verify OpenAI client was called with correct parameters
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
model: mockOptions.openRouterModelId,
temperature: 0,
messages: expect.arrayContaining([
{ role: 'system', content: systemPrompt },
{ role: 'user', content: 'test message' }
]),
stream: true
}))
})
test('includes middle-out transform when enabled', async () => {
const optionsWithMiddleOut = {
...mockOptions,
openRouterUseMiddleOutTransform: true
}
const handler = new OpenRouterHandler(optionsWithMiddleOut)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
id: 'test-id',
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
// Mock OpenAI chat.completions.create
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
// Mock axios.get for generation details
;(axios.get as jest.Mock).mockResolvedValue({
data: {
data: {
native_tokens_prompt: 10,
native_tokens_completion: 20,
total_cost: 0.001
}
}
})
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }]
const generator = handler.createMessage(systemPrompt, messages)
for await (const _ of generator) { /* consume generator */ }
// Verify OpenAI client was called with middle-out transform
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
transforms: ['middle-out']
}))
})
test('handles generation details fetch failure gracefully', async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
id: 'test-id',
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
// Mock OpenAI chat.completions.create
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
// Mock axios.get to fail
;(axios.get as jest.Mock).mockRejectedValue(new Error('Network error'))
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }]
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
for await (const chunk of generator) {
chunks.push(chunk)
}
// Should still get text chunks even if usage info fails
expect(chunks).toHaveLength(2)
expect(chunks[0]).toEqual({
type: 'text',
text: 'test response'
})
expect(chunks[1]).toEqual({
type: 'text',
text: '\n\n'
})
})
test('handles API errors gracefully', async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
error: {
message: 'Test error',
code: 500
}
}
}
}
// Mock OpenAI chat.completions.create
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }]
const generator = handler.createMessage(systemPrompt, messages)
await expect(async () => {
for await (const _ of generator) { /* consume generator */ }
}).rejects.toThrow('OpenRouter API Error 500: Test error')
})
// Verify OpenAI client was called with correct parameters
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
model: mockOptions.openRouterModelId,
temperature: 0,
messages: expect.arrayContaining([
{ role: 'system', content: systemPrompt },
{ role: 'user', content: 'test message' }
]),
stream: true
}))
})
})

View file

@ -8,7 +8,9 @@ import {
ModelInfo,
} from "../../shared/api"
import { ApiHandler } from "../index"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import delay from "delay"
import { withRetry } from "../utils/retry"
export class AnthropicHandler implements ApiHandler {
private options: ApiHandlerOptions
@ -23,145 +25,164 @@ export class AnthropicHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
let stream: AnthropicStream<Anthropic.Beta.PromptCaching.Messages.RawPromptCachingBetaMessageStreamEvent>
const modelId = this.getModel().id
switch (modelId) {
// 'latest' alias does not support cache_control
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307": {
/*
The latest message will be the new user message, one before will be the assistant message from a previous request, and the user message before that will be a previously cached user message. So we need to mark the latest user message as ephemeral to cache it for the next request, and mark the second to last user message as ephemeral to let the server know the last message to retrieve from the cache for the current request..
*/
const userMsgIndices = messages.reduce(
(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
[] as number[],
)
const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
stream = await this.client.beta.promptCaching.messages.create(
{
const self = this;
const modelId = this.getModel().id;
const gen = withRetry(async () => {
let stream: AnthropicStream<Anthropic.Beta.PromptCaching.Messages.RawPromptCachingBetaMessageStreamEvent>
switch (modelId) {
// 'latest' alias does not support cache_control
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307": {
/*
The latest message will be the new user message, one before will be the assistant message from a previous request, and the user message before that will be a previously cached user message. So we need to mark the latest user message as ephemeral to cache it for the next request, and mark the second to last user message as ephemeral to let the server know the last message to retrieve from the cache for the current request..
*/
const userMsgIndices = messages.reduce(
(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
[] as number[],
)
const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
stream = await self.client.beta.promptCaching.messages.create(
{
model: modelId,
max_tokens: self.getModel().info.maxTokens || 8192,
temperature: 0,
system: [{ text: systemPrompt, type: "text", cache_control: { type: "ephemeral" } }], // setting cache breakpoint for system prompt so new tasks can reuse it
messages: messages.map((message, index) => {
if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) {
return {
...message,
content:
typeof message.content === "string"
? [
{
type: "text",
text: message.content,
cache_control: { type: "ephemeral" },
},
]
: message.content.map((content, contentIndex) =>
contentIndex === message.content.length - 1
? { ...content, cache_control: { type: "ephemeral" } }
: content,
),
}
}
return message
}),
// tools, // cache breakpoints go from tools > system > messages, and since tools dont change, we can just set the breakpoint at the end of system (this avoids having to set a breakpoint at the end of tools which by itself does not meet min requirements for haiku caching)
// tool_choice: { type: "auto" },
// tools: tools,
stream: true,
},
(() => {
// prompt caching: https://x.com/alexalbert__/status/1823751995901272068
// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
// https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393
switch (modelId) {
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307":
return {
headers: { "anthropic-beta": "prompt-caching-2024-07-31" },
}
default:
return undefined
}
})(),
)
break
}
default: {
stream = (await self.client.messages.create({
model: modelId,
max_tokens: this.getModel().info.maxTokens || 8192,
max_tokens: self.getModel().info.maxTokens || 8192,
temperature: 0,
system: [{ text: systemPrompt, type: "text", cache_control: { type: "ephemeral" } }], // setting cache breakpoint for system prompt so new tasks can reuse it
messages: messages.map((message, index) => {
if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) {
return {
...message,
content:
typeof message.content === "string"
? [
{
type: "text",
text: message.content,
cache_control: { type: "ephemeral" },
},
]
: message.content.map((content, contentIndex) =>
contentIndex === message.content.length - 1
? { ...content, cache_control: { type: "ephemeral" } }
: content,
),
}
}
return message
}),
// tools, // cache breakpoints go from tools > system > messages, and since tools dont change, we can just set the breakpoint at the end of system (this avoids having to set a breakpoint at the end of tools which by itself does not meet min requirements for haiku caching)
system: [{ text: systemPrompt, type: "text" }],
messages,
// tools,
// tool_choice: { type: "auto" },
// tools: tools,
stream: true,
},
(() => {
// prompt caching: https://x.com/alexalbert__/status/1823751995901272068
// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
// https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393
switch (modelId) {
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307":
return {
headers: { "anthropic-beta": "prompt-caching-2024-07-31" },
}
default:
return undefined
}
})(),
)
break
})) as any
break
}
}
default: {
stream = (await this.client.messages.create({
model: modelId,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: 0,
system: [{ text: systemPrompt, type: "text" }],
messages,
// tools,
// tool_choice: { type: "auto" },
stream: true,
})) as any
break
}
}
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
// tells us cache reads/writes/input/output
const usage = chunk.message.usage
yield {
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
cacheReadTokens: usage.cache_read_input_tokens || undefined,
}
break
case "message_delta":
// tells us stop_reason, stop_sequence, and output tokens along the way and at the end of the message
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
}
break
case "message_stop":
// no usage data, just an indicator that the message is done
break
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
// we may receive multiple text blocks, in which case just insert a line break between them
if (chunk.index > 0) {
yield {
type: "text",
text: "\n",
}
}
return (async function*() {
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
// tells us cache reads/writes/input/output
const usage = chunk.message.usage
yield {
type: "text",
text: chunk.content_block.text,
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
cacheReadTokens: usage.cache_read_input_tokens || undefined,
} as ApiStreamChunk;
break
case "message_delta":
// tells us stop_reason, stop_sequence, and output tokens along the way and at the end of the message
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
} as ApiStreamChunk;
break
case "message_stop":
// no usage data, just an indicator that the message is done
break
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
// we may receive multiple text blocks, in which case just insert a line break between them
if (chunk.index > 0) {
yield {
type: "text",
text: "\n",
} as ApiStreamChunk;
}
yield {
type: "text",
text: chunk.content_block.text,
} as ApiStreamChunk;
break
}
break
}
break
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
yield {
type: "text",
text: chunk.delta.text,
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
yield {
type: "text",
text: chunk.delta.text,
} as ApiStreamChunk;
break
}
break
case "content_block_stop":
break
}
break
case "content_block_stop":
break
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`Anthropic request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -2,7 +2,8 @@ import AnthropicBedrock from "@anthropic-ai/bedrock-sdk"
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandler } from "../"
import { ApiHandlerOptions, bedrockDefaultModelId, BedrockModelId, bedrockModels, ModelInfo } from "../../shared/api"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
// https://docs.anthropic.com/en/api/claude-on-amazon-bedrock
export class AwsBedrockHandler implements ApiHandler {
@ -25,6 +26,8 @@ export class AwsBedrockHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const self = this;
// cross region inference requires prefixing the model id with the region
let modelId: string
if (this.options.awsUseCrossRegionInference) {
@ -45,59 +48,77 @@ export class AwsBedrockHandler implements ApiHandler {
modelId = this.getModel().id
}
const stream = await this.client.messages.create({
model: modelId,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: 0,
system: systemPrompt,
messages,
stream: true,
})
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
const usage = chunk.message.usage
yield {
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
}
break
case "message_delta":
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
}
break
const gen = withRetry(async () => {
const stream = await self.client.messages.create({
model: modelId,
max_tokens: self.getModel().info.maxTokens || 8192,
temperature: 0,
system: systemPrompt,
messages,
stream: true,
})
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
if (chunk.index > 0) {
yield {
type: "text",
text: "\n",
}
}
return (async function*() {
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
// tells us cache reads/writes/input/output
const usage = chunk.message.usage
yield {
type: "text",
text: chunk.content_block.text,
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
} as ApiStreamChunk;
break
case "message_delta":
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
} as ApiStreamChunk;
break
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
if (chunk.index > 0) {
yield {
type: "text",
text: "\n",
} as ApiStreamChunk;
}
yield {
type: "text",
text: chunk.content_block.text,
} as ApiStreamChunk;
break
}
break
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
yield {
type: "text",
text: chunk.delta.text,
} as ApiStreamChunk;
break
}
break
}
break
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
yield {
type: "text",
text: chunk.delta.text,
}
break
}
break
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`Bedrock request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -3,7 +3,8 @@ import { GoogleGenerativeAI } from "@google/generative-ai"
import { ApiHandler } from "../"
import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "../../shared/api"
import { convertAnthropicMessageToGemini } from "../transform/gemini-format"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
export class GeminiHandler implements ApiHandler {
private options: ApiHandlerOptions
@ -18,30 +19,48 @@ export class GeminiHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const self = this;
const model = this.client.getGenerativeModel({
model: this.getModel().id,
systemInstruction: systemPrompt,
})
const result = await model.generateContentStream({
contents: messages.map(convertAnthropicMessageToGemini),
generationConfig: {
// maxOutputTokens: this.getModel().info.maxTokens,
temperature: 0,
},
})
for await (const chunk of result.stream) {
yield {
type: "text",
text: chunk.text(),
const gen = withRetry(async () => {
const result = await model.generateContentStream({
contents: messages.map(convertAnthropicMessageToGemini),
generationConfig: {
// maxOutputTokens: this.getModel().info.maxTokens,
temperature: 0,
},
})
return (async function*() {
for await (const chunk of result.stream) {
yield {
type: "text",
text: chunk.text(),
} as ApiStreamChunk;
}
const response = await result.response
yield {
type: "usage",
inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
} as ApiStreamChunk;
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`Gemini request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
}
});
const response = await result.response
yield {
type: "usage",
inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -3,7 +3,8 @@ import OpenAI from "openai"
import { ApiHandler } from "../"
import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
export class LmStudioHandler implements ApiHandler {
private options: ApiHandlerOptions
@ -18,32 +19,48 @@ export class LmStudioHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const self = this;
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...convertToOpenAiMessages(messages),
]
try {
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
const gen = withRetry(async () => {
const stream = await self.client.chat.completions.create({
model: self.getModel().id,
messages: openAiMessages,
temperature: 0,
stream: true,
})
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
}).catch(error => {
// LM Studio doesn't return an error code/body for now
throw new Error(
"Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Cline's prompts.",
)
});
return (async function*() {
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
} as ApiStreamChunk;
}
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`LM Studio request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
} catch (error) {
// LM Studio doesn't return an error code/body for now
throw new Error(
"Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Cline's prompts.",
)
});
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -3,7 +3,8 @@ import OpenAI from "openai"
import { ApiHandler } from "../"
import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
export class OllamaHandler implements ApiHandler {
private options: ApiHandlerOptions
@ -18,25 +19,49 @@ export class OllamaHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const self = this;
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...convertToOpenAiMessages(messages),
]
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
messages: openAiMessages,
temperature: 0,
stream: true,
})
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
const gen = withRetry(async () => {
const stream = await self.client.chat.completions.create({
model: self.getModel().id,
messages: openAiMessages,
temperature: 0,
stream: true,
}).catch(error => {
// Check if it's a connection error, which likely means Ollama isn't running
if (error instanceof Error && error.message.includes('ECONNREFUSED')) {
throw new Error('Could not connect to Ollama. Please make sure Ollama is running on your system.');
}
throw error;
});
return (async function*() {
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
} as ApiStreamChunk;
}
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`Ollama request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -9,7 +9,8 @@ import {
openAiNativeModels,
} from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
export class OpenAiNativeHandler implements ApiHandler {
private options: ApiHandlerOptions
@ -23,52 +24,87 @@ export class OpenAiNativeHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const self = this;
switch (this.getModel().id) {
case "o1-preview":
case "o1-mini": {
// o1 doesnt support streaming, non-1 temp, or system prompt
const response = await this.client.chat.completions.create({
model: this.getModel().id,
messages: [{ role: "user", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
})
yield {
type: "text",
text: response.choices[0]?.message.content || "",
}
yield {
type: "usage",
inputTokens: response.usage?.prompt_tokens || 0,
outputTokens: response.usage?.completion_tokens || 0,
}
break
}
default: {
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
// max_completion_tokens: this.getModel().info.maxTokens,
temperature: 0,
messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
stream: true,
stream_options: { include_usage: true },
})
const gen = withRetry(async () => {
// o1 doesnt support streaming, non-1 temp, or system prompt
const response = await self.client.chat.completions.create({
model: self.getModel().id,
messages: [{ role: "user", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
});
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
return (async function*() {
yield {
type: "text",
text: delta.content,
}
}
// contains a null value except for the last chunk which contains the token usage statistics for the entire request
if (chunk.usage) {
text: response.choices[0]?.message.content || "",
} as ApiStreamChunk;
yield {
type: "usage",
inputTokens: chunk.usage.prompt_tokens || 0,
outputTokens: chunk.usage.completion_tokens || 0,
}
inputTokens: response.usage?.prompt_tokens || 0,
outputTokens: response.usage?.completion_tokens || 0,
} as ApiStreamChunk;
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`OpenAI Native request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
break;
}
default: {
const gen = withRetry(async () => {
const stream = await self.client.chat.completions.create({
model: self.getModel().id,
// max_completion_tokens: this.getModel().info.maxTokens,
temperature: 0,
messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
stream: true,
stream_options: { include_usage: true },
});
return (async function*() {
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
} as ApiStreamChunk;
}
// contains a null value except for the last chunk which contains the token usage statistics for the entire request
if (chunk.usage) {
yield {
type: "usage",
inputTokens: chunk.usage.prompt_tokens || 0,
outputTokens: chunk.usage.completion_tokens || 0,
} as ApiStreamChunk;
}
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`OpenAI Native request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
}
}

View file

@ -8,7 +8,8 @@ import {
} from "../../shared/api"
import { ApiHandler } from "../index"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
export class OpenAiHandler implements ApiHandler {
private options: ApiHandlerOptions
@ -49,22 +50,41 @@ export class OpenAiHandler implements ApiHandler {
requestOptions.stream_options = { include_usage: true }
}
const stream = await this.client.chat.completions.create(requestOptions)
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
}
}
if (chunk.usage) {
yield {
type: "usage",
inputTokens: chunk.usage.prompt_tokens || 0,
outputTokens: chunk.usage.completion_tokens || 0,
const self = this;
const gen = withRetry(async () => {
const stream = await self.client.chat.completions.create(requestOptions);
return (async function*() {
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
} as ApiStreamChunk;
}
if (chunk.usage) {
yield {
type: "usage",
inputTokens: chunk.usage.prompt_tokens || 0,
outputTokens: chunk.usage.completion_tokens || 0,
} as ApiStreamChunk;
}
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`OpenAI request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -6,6 +6,7 @@ import { ApiHandlerOptions, ModelInfo, openRouterDefaultModelId, openRouterDefau
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStreamChunk, ApiStreamUsageChunk } from "../transform/stream"
import delay from "delay"
import { withRetry } from "../utils/retry"
// Add custom interface for OpenRouter params
interface OpenRouterChatCompletionParams extends OpenAI.Chat.ChatCompletionCreateParamsStreaming {
@ -18,6 +19,7 @@ interface OpenRouterApiStreamUsageChunk extends ApiStreamUsageChunk {
}
export class OpenRouterHandler implements ApiHandler {
private static requestCount = 0
private options: ApiHandlerOptions
private client: OpenAI
@ -27,21 +29,28 @@ export class OpenRouterHandler implements ApiHandler {
baseURL: "https://openrouter.ai/api/v1",
apiKey: this.options.openRouterApiKey,
defaultHeaders: {
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline", // Optional, for including your app on openrouter.ai rankings.
"X-Title": "Roo-Cline", // Optional. Shows in rankings on openrouter.ai.
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline",
"X-Title": "Roo-Cline",
},
})
}
getModel(): { id: string; info: ModelInfo } {
const modelId = this.options.openRouterModelId
const modelInfo = this.options.openRouterModelInfo
if (modelId && modelInfo) {
return { id: modelId, info: modelInfo }
}
return { id: openRouterDefaultModelId, info: openRouterDefaultModelInfo }
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): AsyncGenerator<ApiStreamChunk> {
// Convert Anthropic messages to OpenAI format
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...convertToOpenAiMessages(messages),
]
// prompt caching: https://openrouter.ai/docs/prompt-caching
// this is specifically for claude models (some models may 'support prompt caching' automatically without this)
switch (this.getModel().id) {
case "anthropic/claude-3.5-sonnet":
case "anthropic/claude-3.5-sonnet:beta":
@ -66,17 +75,13 @@ export class OpenRouterHandler implements ApiHandler {
},
],
}
// Add cache_control to the last two user messages
// (note: this works because we only ever add one user message at a time, but if we added multiple we'd need to mark the user message before the last assistant message)
const lastTwoUserMessages = openAiMessages.filter((msg) => msg.role === "user").slice(-2)
lastTwoUserMessages.forEach((msg) => {
if (typeof msg.content === "string") {
msg.content = [{ type: "text", text: msg.content }]
}
if (Array.isArray(msg.content)) {
// NOTE: this is fine since env details will always be added at the end. but if it weren't there, and the user added a image_url type message, it would pop a text part before it and then move it after to the end.
let lastTextPart = msg.content.filter((part) => part.type === "text").pop()
if (!lastTextPart) {
lastTextPart = { type: "text", text: "..." }
msg.content.push(lastTextPart)
@ -86,12 +91,8 @@ export class OpenRouterHandler implements ApiHandler {
}
})
break
default:
break
}
// Not sure how openrouter defaults max tokens when no value is provided, but the anthropic api requires this value and since they offer both 4096 and 8192 variants, we should ensure 8192.
// (models usually default to max tokens allowed)
let maxTokens: number | undefined
switch (this.getModel().id) {
case "anthropic/claude-3.5-sonnet":
@ -105,83 +106,85 @@ export class OpenRouterHandler implements ApiHandler {
maxTokens = 8_192
break
}
// https://openrouter.ai/docs/transforms
let fullResponseText = "";
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
max_tokens: maxTokens,
temperature: 0,
messages: openAiMessages,
stream: true,
// This way, the transforms field will only be included in the parameters when openRouterUseMiddleOutTransform is true.
...(this.options.openRouterUseMiddleOutTransform && { transforms: ["middle-out"] })
} as OpenRouterChatCompletionParams);
const self = this;
let genId: string | undefined
const gen = withRetry(async () => {
const stream = await self.client.chat.completions.create({
model: self.getModel().id,
max_tokens: maxTokens,
temperature: 0,
messages: openAiMessages,
stream: true,
...(self.options.openRouterUseMiddleOutTransform && { transforms: ["middle-out"] })
} as OpenRouterChatCompletionParams);
for await (const chunk of stream as unknown as AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>) {
// openrouter returns an error object instead of the openai sdk throwing an error
if ("error" in chunk) {
const error = chunk.error as { message?: string; code?: number }
console.error(`OpenRouter API Error: ${error?.code} - ${error?.message}`)
throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`)
}
let genId: string | undefined;
if (!genId && chunk.id) {
genId = chunk.id
}
return (async function*() {
for await (const chunk of stream) {
if ("error" in chunk) {
const error = chunk.error as { message?: string; code?: number }
console.error(`OpenRouter API Error: ${error?.code} - ${error?.message}`)
throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`)
}
const delta = chunk.choices[0]?.delta
if (delta?.content) {
fullResponseText += delta.content;
if (!genId && chunk.id) {
genId = chunk.id
}
const delta = chunk.choices[0]?.delta
if (delta?.content) {
fullResponseText += delta.content;
yield {
type: "text",
text: delta.content,
} as ApiStreamChunk;
}
}
await delay(500)
try {
const response = await axios.get(`https://openrouter.ai/api/v1/generation?id=${genId}`, {
headers: {
Authorization: `Bearer ${self.options.openRouterApiKey}`,
},
timeout: 5_000,
})
const generation = response.data?.data
console.log("OpenRouter generation details:", response.data)
yield {
type: "usage",
inputTokens: generation?.native_tokens_prompt || 0,
outputTokens: generation?.native_tokens_completion || 0,
totalCost: generation?.total_cost || 0,
fullResponseText
} as ApiStreamChunk;
} catch (error) {
console.error("Error fetching OpenRouter generation details:", error)
}
// Add newlines before starting content
yield {
type: "text",
text: delta.content,
text: "\n\n"
} as ApiStreamChunk;
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`OpenRouter request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
// if (chunk.usage) {
// yield {
// type: "usage",
// inputTokens: chunk.usage.prompt_tokens || 0,
// outputTokens: chunk.usage.completion_tokens || 0,
// }
// }
});
for await (const chunk of gen) {
yield chunk;
}
await delay(500) // FIXME: necessary delay to ensure generation endpoint is ready
try {
const response = await axios.get(`https://openrouter.ai/api/v1/generation?id=${genId}`, {
headers: {
Authorization: `Bearer ${this.options.openRouterApiKey}`,
},
timeout: 5_000, // this request hangs sometimes
})
const generation = response.data?.data
console.log("OpenRouter generation details:", response.data)
yield {
type: "usage",
// cacheWriteTokens: 0,
// cacheReadTokens: 0,
// openrouter generation endpoint fails often
inputTokens: generation?.native_tokens_prompt || 0,
outputTokens: generation?.native_tokens_completion || 0,
totalCost: generation?.total_cost || 0,
fullResponseText
} as OpenRouterApiStreamUsageChunk;
} catch (error) {
// ignore if fails
console.error("Error fetching OpenRouter generation details:", error)
}
}
getModel(): { id: string; info: ModelInfo } {
const modelId = this.options.openRouterModelId
const modelInfo = this.options.openRouterModelInfo
if (modelId && modelInfo) {
return { id: modelId, info: modelInfo }
}
return { id: openRouterDefaultModelId, info: openRouterDefaultModelInfo }
}
}

View file

@ -2,7 +2,8 @@ import { Anthropic } from "@anthropic-ai/sdk"
import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
import { ApiHandler } from "../"
import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "../../shared/api"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamChunk } from "../transform/stream"
import { withRetry } from "../utils/retry"
// https://docs.anthropic.com/en/api/claude-on-vertex-ai
export class VertexHandler implements ApiHandler {
@ -19,59 +20,78 @@ export class VertexHandler implements ApiHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const stream = await this.client.messages.create({
model: this.getModel().id,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: 0,
system: systemPrompt,
messages,
stream: true,
})
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
const usage = chunk.message.usage
yield {
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
}
break
case "message_delta":
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
}
break
const self = this;
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
if (chunk.index > 0) {
yield {
type: "text",
text: "\n",
}
}
const gen = withRetry(async () => {
const stream = await self.client.messages.create({
model: self.getModel().id,
max_tokens: self.getModel().info.maxTokens || 8192,
temperature: 0,
system: systemPrompt,
messages,
stream: true,
})
return (async function*() {
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
const usage = chunk.message.usage
yield {
type: "text",
text: chunk.content_block.text,
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
} as ApiStreamChunk;
break
case "message_delta":
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
} as ApiStreamChunk;
break
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
if (chunk.index > 0) {
yield {
type: "text",
text: "\n",
} as ApiStreamChunk;
}
yield {
type: "text",
text: chunk.content_block.text,
} as ApiStreamChunk;
break
}
break
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
yield {
type: "text",
text: chunk.delta.text,
} as ApiStreamChunk;
break
}
break
}
break
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
yield {
type: "text",
text: chunk.delta.text,
}
break
}
break
}
})();
}, {
maxRetries: 5,
initialDelayMs: 2000,
onRetry: (error, attempt, delayMs) => {
console.log(`Vertex request failed (attempt ${attempt})`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
}
});
for await (const chunk of gen) {
yield chunk;
}
}

View file

@ -1,5 +1,5 @@
export type ApiStream = AsyncGenerator<ApiStreamChunk>
export type ApiStreamChunk = ApiStreamTextChunk | ApiStreamUsageChunk
export type ApiStreamChunk = ApiStreamTextChunk | ApiStreamUsageChunk | ApiStreamStatusChunk
export interface ApiStreamTextChunk {
type: "text"
@ -14,3 +14,8 @@ export interface ApiStreamUsageChunk {
cacheReadTokens?: number
totalCost?: number // openrouter
}
export interface ApiStreamStatusChunk {
type: "status"
text: string
}

View file

@ -0,0 +1,166 @@
import { withRetry } from '../retry';
import { ApiStreamChunk, ApiStreamTextChunk } from '../../transform/stream';
import delay from 'delay';
jest.mock('delay');
describe('withRetry', () => {
beforeEach(() => {
jest.clearAllMocks();
(delay as jest.Mock).mockResolvedValue(undefined);
});
it('should complete successfully with no retries', async () => {
const mockChunks: ApiStreamTextChunk[] = [
{ type: 'text', text: 'chunk1' },
{ type: 'text', text: 'chunk2' }
];
const operation = async () => {
const generator = async function* () {
for (const chunk of mockChunks) {
yield chunk;
}
};
return generator();
};
const result: ApiStreamChunk[] = [];
for await (const chunk of withRetry(operation)) {
result.push(chunk);
}
expect(result).toEqual(mockChunks);
expect(delay).not.toHaveBeenCalled();
});
it('should retry on failure and eventually succeed', async () => {
let attempts = 0;
const mockChunks: ApiStreamTextChunk[] = [
{ type: 'text', text: 'success' }
];
const operation = async () => {
const generator = async function* () {
attempts++;
if (attempts < 3) {
throw new Error('Temporary failure');
}
yield mockChunks[0];
};
return generator();
};
const onRetry = jest.fn();
const result: ApiStreamChunk[] = [];
for await (const chunk of withRetry(operation, {
maxRetries: 5,
initialDelayMs: 1000,
onRetry
})) {
result.push(chunk);
}
expect(attempts).toBe(3);
expect(onRetry).toHaveBeenCalledTimes(2);
expect(delay).toHaveBeenCalledTimes(2);
expect(delay).toHaveBeenNthCalledWith(1, 1000); // First retry
expect(delay).toHaveBeenNthCalledWith(2, 2000); // Second retry with exponential backoff
// Filter out retry status messages
const actualResults = result.filter(
(chunk): chunk is ApiStreamTextChunk =>
chunk.type === 'text' && chunk.text === 'success'
);
expect(actualResults).toEqual(mockChunks);
});
it('should throw after max retries exceeded', async () => {
const operation = async () => {
const generator = async function* () {
throw new Error('Persistent failure');
};
return generator();
};
const onRetry = jest.fn();
const generator = withRetry(operation, {
maxRetries: 3,
initialDelayMs: 1000,
onRetry
});
const result: ApiStreamChunk[] = [];
await expect(async () => {
for await (const chunk of generator) {
result.push(chunk);
}
}).rejects.toThrow('Persistent failure');
expect(onRetry).toHaveBeenCalledTimes(3);
expect(delay).toHaveBeenCalledTimes(3);
expect(delay).toHaveBeenNthCalledWith(1, 1000);
expect(delay).toHaveBeenNthCalledWith(2, 2000);
expect(delay).toHaveBeenNthCalledWith(3, 4000);
// Should have yielded retry status messages
expect(result.length).toBe(3);
result.forEach(chunk => {
expect(chunk.type).toBe('text');
if (chunk.type === 'text') {
expect(chunk.text).toMatch(/Request failed\. Retrying in \d+ seconds\.\.\./);
}
});
});
it('should use default options when not provided', async () => {
let attempts = 0;
const operation = async () => {
const generator = async function* () {
attempts++;
if (attempts === 1) {
throw new Error('First attempt failure');
}
yield { type: 'text' as const, text: 'success' };
};
return generator();
};
const result: ApiStreamChunk[] = [];
for await (const chunk of withRetry(operation)) {
result.push(chunk);
}
expect(attempts).toBe(2);
expect(delay).toHaveBeenCalledTimes(1);
expect(delay).toHaveBeenCalledWith(2000); // Default initialDelayMs
});
it('should handle custom onRetry callback', async () => {
const customOnRetry = jest.fn();
let attempts = 0;
const operation = async () => {
const generator = async function* () {
attempts++;
if (attempts === 1) {
throw new Error('Test error');
}
yield { type: 'text' as const, text: 'success' };
};
return generator();
};
for await (const chunk of withRetry(operation, { onRetry: customOnRetry })) {
// Consume chunks
}
expect(customOnRetry).toHaveBeenCalledTimes(1);
expect(customOnRetry).toHaveBeenCalledWith(
expect.any(Error),
1,
2000
);
});
});

51
src/api/utils/retry.ts Normal file
View file

@ -0,0 +1,51 @@
import delay from "delay"
import { ApiStreamChunk } from "../transform/stream"
export interface RetryOptions {
maxRetries?: number;
initialDelayMs?: number;
onRetry?: (error: unknown, attempt: number, delayMs: number) => void;
}
export async function* withRetry(
operation: () => Promise<AsyncGenerator<ApiStreamChunk>>,
options: RetryOptions = {}
): AsyncGenerator<ApiStreamChunk> {
const {
maxRetries = 5,
initialDelayMs = 2000,
onRetry = (error, attempt, delayMs) => {
console.log(`Operation failed, attempt ${attempt}/${maxRetries}`);
console.log(`Error:`, error);
console.log(`Retrying in ${delayMs}ms...`);
},
} = options;
let attempt = 0;
while (true) {
try {
const stream = await operation();
for await (const chunk of stream) {
yield chunk;
}
return;
} catch (error) {
attempt++;
if (attempt > maxRetries) {
console.log(`Max retries (${maxRetries}) exceeded, giving up`);
throw error;
}
const delayMs = initialDelayMs * Math.pow(2, attempt - 1);
onRetry(error, attempt, delayMs);
yield {
type: "text",
text: `Request failed. Retrying in ${delayMs/1000} seconds... (attempt ${attempt}/${maxRetries})\n\n`
};
await delay(delayMs);
}
}
}