mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
fix: prevent task corruption during API retry cycles
This commit is contained in:
hannesrudolph
parent
0028c56711
commit
9edfe46a68
8 changed files with 1617 additions and 54 deletions
296
src/api/error-handling/UnifiedErrorHandler.test.ts
Normal file
296
src/api/error-handling/UnifiedErrorHandler.test.ts
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
import { describe, test, expect } from "vitest"
|
||||
import { UnifiedErrorHandler, ErrorContext } from "./UnifiedErrorHandler"
|
||||
|
||||
describe("UnifiedErrorHandler", () => {
|
||||
const createContext = (overrides: Partial<ErrorContext> = {}): ErrorContext => ({
|
||||
isStreaming: false,
|
||||
provider: "anthropic",
|
||||
modelId: "claude-3-sonnet",
|
||||
retryAttempt: 0,
|
||||
requestId: "test-request",
|
||||
...overrides,
|
||||
})
|
||||
|
||||
describe("error classification", () => {
|
||||
test("classifies HTTP 429 as THROTTLING", () => {
|
||||
const error = { status: 429, message: "Rate limit exceeded" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("THROTTLING")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies ThrottlingException as THROTTLING", () => {
|
||||
const error = { name: "ThrottlingException", message: "Request was throttled" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("THROTTLING")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies AccessDeniedException as ACCESS_DENIED", () => {
|
||||
const error = { name: "AccessDeniedException", message: "Access denied" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("ACCESS_DENIED")
|
||||
expect(result.shouldRetry).toBe(false)
|
||||
expect(result.shouldThrow).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies ResourceNotFoundException as NOT_FOUND", () => {
|
||||
const error = { name: "ResourceNotFoundException", message: "Resource not found" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("NOT_FOUND")
|
||||
expect(result.shouldRetry).toBe(false)
|
||||
expect(result.shouldThrow).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies ServiceUnavailableException as SERVICE_UNAVAILABLE", () => {
|
||||
const error = { name: "ServiceUnavailableException", message: "Service unavailable" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("SERVICE_UNAVAILABLE")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies ValidationException as INVALID_REQUEST", () => {
|
||||
const error = { name: "ValidationException", message: "Invalid request" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("INVALID_REQUEST")
|
||||
expect(result.shouldRetry).toBe(false)
|
||||
expect(result.shouldThrow).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies throttling patterns in message", () => {
|
||||
const error = new Error("too many requests, please wait")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("THROTTLING")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies rate limit patterns in message", () => {
|
||||
const error = new Error("rate limit exceeded, please wait")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("RATE_LIMITED")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies quota patterns in message", () => {
|
||||
const error = new Error("quota exceeded for this month")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("QUOTA_EXCEEDED")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies network errors", () => {
|
||||
const error = new Error("network connection failed")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("NETWORK_ERROR")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies timeout errors", () => {
|
||||
const error = new Error("request timed out")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("TIMEOUT")
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
|
||||
test("classifies generic errors", () => {
|
||||
const error = new Error("something went wrong")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("GENERIC")
|
||||
})
|
||||
|
||||
test("classifies unknown non-Error objects", () => {
|
||||
const error = "string error"
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.errorType).toBe("UNKNOWN")
|
||||
})
|
||||
})
|
||||
|
||||
describe("retry logic", () => {
|
||||
test("retries throttling errors up to max attempts", () => {
|
||||
const error = { status: 429, message: "Rate limit exceeded" }
|
||||
|
||||
// Should retry for first few attempts
|
||||
for (let attempt = 0; attempt < 5; attempt++) {
|
||||
const context = createContext({ retryAttempt: attempt })
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
}
|
||||
|
||||
// Should not retry after max attempts
|
||||
const contextMaxAttempts = createContext({ retryAttempt: 5 })
|
||||
const resultMaxAttempts = UnifiedErrorHandler.handle(error, contextMaxAttempts)
|
||||
expect(resultMaxAttempts.shouldRetry).toBe(false)
|
||||
})
|
||||
|
||||
test("does not retry non-retryable errors", () => {
|
||||
const error = { name: "AccessDeniedException", message: "Access denied" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.shouldRetry).toBe(false)
|
||||
})
|
||||
|
||||
test("retries service unavailable errors", () => {
|
||||
const error = new Error("service temporarily unavailable")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.shouldRetry).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("streaming context handling", () => {
|
||||
test("throws immediately for throttling in streaming context", () => {
|
||||
const error = { status: 429, message: "Rate limit exceeded" }
|
||||
const context = createContext({ isStreaming: true })
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.shouldThrow).toBe(true)
|
||||
expect(result.shouldRetry).toBe(true) // Still retryable, but should throw for proper handling
|
||||
})
|
||||
|
||||
test("provides stream chunks for non-throwing streaming errors", () => {
|
||||
const error = new Error("generic error")
|
||||
const context = createContext({ isStreaming: true })
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.streamChunks).toBeDefined()
|
||||
expect(result.streamChunks).toHaveLength(2)
|
||||
expect(result.streamChunks![0].type).toBe("text")
|
||||
expect(result.streamChunks![1].type).toBe("usage")
|
||||
})
|
||||
|
||||
test("does not provide stream chunks for non-streaming context", () => {
|
||||
const error = new Error("generic error")
|
||||
const context = createContext({ isStreaming: false })
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.streamChunks).toBeUndefined()
|
||||
})
|
||||
})
|
||||
|
||||
describe("retry delay calculation", () => {
|
||||
test("calculates exponential backoff", () => {
|
||||
const error = new Error("generic error message")
|
||||
|
||||
const context0 = createContext({ retryAttempt: 0 })
|
||||
const result0 = UnifiedErrorHandler.handle(error, context0)
|
||||
expect(result0.retryDelay).toBe(5) // base delay
|
||||
|
||||
const context1 = createContext({ retryAttempt: 1 })
|
||||
const result1 = UnifiedErrorHandler.handle(error, context1)
|
||||
expect(result1.retryDelay).toBe(10) // 5 * 2^1
|
||||
|
||||
const context2 = createContext({ retryAttempt: 2 })
|
||||
const result2 = UnifiedErrorHandler.handle(error, context2)
|
||||
expect(result2.retryDelay).toBe(20) // 5 * 2^2
|
||||
})
|
||||
|
||||
test("respects maximum delay", () => {
|
||||
const error = new Error("service unavailable")
|
||||
const context = createContext({ retryAttempt: 10 }) // Very high retry attempt
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.retryDelay).toBeLessThanOrEqual(600) // Max 10 minutes
|
||||
})
|
||||
|
||||
test("adjusts delay based on error type", () => {
|
||||
const baseRetryAttempt = 1
|
||||
|
||||
// Service unavailable gets longer delay
|
||||
const serviceError = { name: "ServiceUnavailableException", message: "Service unavailable" }
|
||||
const serviceContext = createContext({ retryAttempt: baseRetryAttempt })
|
||||
const serviceResult = UnifiedErrorHandler.handle(serviceError, serviceContext)
|
||||
|
||||
// Network error gets shorter delay
|
||||
const networkError = new Error("network connection failed")
|
||||
const networkContext = createContext({ retryAttempt: baseRetryAttempt })
|
||||
const networkResult = UnifiedErrorHandler.handle(networkError, networkContext)
|
||||
|
||||
expect(serviceResult.retryDelay).toBeGreaterThan(networkResult.retryDelay!)
|
||||
})
|
||||
|
||||
test("extracts provider-specific retry delay", () => {
|
||||
// Simulate Google Gemini retry info
|
||||
const error = {
|
||||
message: "Rate limit exceeded",
|
||||
errorDetails: [
|
||||
{
|
||||
"@type": "type.googleapis.com/google.rpc.RetryInfo",
|
||||
retryDelay: "30s",
|
||||
},
|
||||
],
|
||||
}
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.retryDelay).toBe(31) // 30s + 1s buffer
|
||||
})
|
||||
})
|
||||
|
||||
describe("error message formatting", () => {
|
||||
test("formats error message with context", () => {
|
||||
const error = new Error("Test error message")
|
||||
const context = createContext({
|
||||
provider: "anthropic",
|
||||
modelId: "claude-3-sonnet",
|
||||
retryAttempt: 2,
|
||||
})
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.formattedMessage).toContain("[anthropic:claude-3-sonnet]")
|
||||
expect(result.formattedMessage).toContain("Test error message")
|
||||
expect(result.formattedMessage).toContain("(Retry 2)")
|
||||
})
|
||||
|
||||
test("includes error type in formatted message", () => {
|
||||
const error = { status: 429, message: "Rate limit exceeded" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.formattedMessage).toContain("[THROTTLING]")
|
||||
})
|
||||
|
||||
test("handles non-Error objects", () => {
|
||||
const error = { someProperty: "not an Error object" }
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.formattedMessage).toContain("Unknown error")
|
||||
})
|
||||
|
||||
test("cleans up whitespace in error messages", () => {
|
||||
const error = new Error("Error with extra whitespace")
|
||||
const context = createContext()
|
||||
|
||||
const result = UnifiedErrorHandler.handle(error, context)
|
||||
expect(result.formattedMessage).toContain("Error with extra whitespace")
|
||||
})
|
||||
})
|
||||
})
|
||||
328
src/api/error-handling/UnifiedErrorHandler.ts
Normal file
328
src/api/error-handling/UnifiedErrorHandler.ts
Normal file
|
|
@ -0,0 +1,328 @@
|
|||
/**
|
||||
* UnifiedErrorHandler - Provides consistent error handling across streaming and non-streaming contexts
|
||||
*
|
||||
* This handler standardizes error classification, retry logic, and response formatting
|
||||
* to prevent inconsistent behavior during API retry cycles.
|
||||
*/
|
||||
|
||||
export interface ErrorContext {
|
||||
isStreaming: boolean
|
||||
provider: string
|
||||
modelId: string
|
||||
retryAttempt?: number
|
||||
requestId?: string
|
||||
}
|
||||
|
||||
export interface ErrorHandlerResponse {
|
||||
shouldRetry: boolean
|
||||
shouldThrow: boolean
|
||||
errorType: string
|
||||
formattedMessage: string
|
||||
retryDelay?: number
|
||||
streamChunks?: Array<{
|
||||
type: string
|
||||
text?: string
|
||||
inputTokens?: number
|
||||
outputTokens?: number
|
||||
}>
|
||||
}
|
||||
|
||||
export type ErrorType =
|
||||
| "THROTTLING"
|
||||
| "RATE_LIMITED"
|
||||
| "ACCESS_DENIED"
|
||||
| "NOT_FOUND"
|
||||
| "INVALID_REQUEST"
|
||||
| "SERVICE_UNAVAILABLE"
|
||||
| "TIMEOUT"
|
||||
| "NETWORK_ERROR"
|
||||
| "QUOTA_EXCEEDED"
|
||||
| "GENERIC"
|
||||
| "UNKNOWN"
|
||||
|
||||
export class UnifiedErrorHandler {
|
||||
/**
|
||||
* Main error handling entry point
|
||||
*/
|
||||
static handle(error: unknown, context: ErrorContext): ErrorHandlerResponse {
|
||||
const errorType = this.classifyError(error)
|
||||
const shouldRetry = this.shouldRetryError(errorType, context.retryAttempt)
|
||||
const shouldThrow = this.shouldThrowImmediately(errorType, context.isStreaming)
|
||||
const retryDelay = this.calculateRetryDelay(errorType, error, context.retryAttempt)
|
||||
|
||||
const formattedMessage = this.formatErrorMessage(error, errorType, context)
|
||||
|
||||
const response: ErrorHandlerResponse = {
|
||||
shouldRetry,
|
||||
shouldThrow,
|
||||
errorType,
|
||||
formattedMessage,
|
||||
retryDelay,
|
||||
}
|
||||
|
||||
// For streaming context, provide chunks when not throwing immediately
|
||||
if (context.isStreaming && !shouldThrow) {
|
||||
response.streamChunks = [
|
||||
{ type: "text", text: `Error: ${formattedMessage}` },
|
||||
{ type: "usage", inputTokens: 0, outputTokens: 0 },
|
||||
]
|
||||
}
|
||||
|
||||
return response
|
||||
}
|
||||
|
||||
/**
|
||||
* Classify error into standardized error types
|
||||
*/
|
||||
private static classifyError(error: unknown): ErrorType {
|
||||
// Handle null/undefined
|
||||
if (!error) return "UNKNOWN"
|
||||
|
||||
// Check for HTTP 429 (highest priority)
|
||||
if ((error as any).status === 429 || (error as any).$metadata?.httpStatusCode === 429) {
|
||||
return "THROTTLING"
|
||||
}
|
||||
|
||||
// Check for specific error names/types (AWS, etc.)
|
||||
const errorName = (error as any).name || ""
|
||||
const errorType = (error as any).__type || ""
|
||||
|
||||
if (errorName === "ThrottlingException" || errorType === "ThrottlingException") {
|
||||
return "THROTTLING"
|
||||
}
|
||||
|
||||
if (errorName === "ServiceUnavailableException" || errorType === "ServiceUnavailableException") {
|
||||
return "SERVICE_UNAVAILABLE"
|
||||
}
|
||||
|
||||
if (errorName === "AccessDeniedException" || errorType === "AccessDeniedException") {
|
||||
return "ACCESS_DENIED"
|
||||
}
|
||||
|
||||
if (errorName === "ResourceNotFoundException" || errorType === "ResourceNotFoundException") {
|
||||
return "NOT_FOUND"
|
||||
}
|
||||
|
||||
if (errorName === "ValidationException" || errorType === "ValidationException") {
|
||||
return "INVALID_REQUEST"
|
||||
}
|
||||
|
||||
// Pattern matching in error message (check both error.message and direct message property)
|
||||
const message = ((error as any).message || "").toLowerCase()
|
||||
|
||||
if (message) {
|
||||
// Throttling patterns (most specific first)
|
||||
if (this.matchesThrottlingPatterns(message)) {
|
||||
return "THROTTLING"
|
||||
}
|
||||
|
||||
// Rate limiting patterns
|
||||
if (this.matchesRateLimitPatterns(message)) {
|
||||
return "RATE_LIMITED"
|
||||
}
|
||||
|
||||
// Quota patterns
|
||||
if (this.matchesQuotaPatterns(message)) {
|
||||
return "QUOTA_EXCEEDED"
|
||||
}
|
||||
|
||||
// Service availability patterns
|
||||
if (this.matchesServiceUnavailablePatterns(message)) {
|
||||
return "SERVICE_UNAVAILABLE"
|
||||
}
|
||||
|
||||
// Access/permission patterns
|
||||
if (this.matchesAccessDeniedPatterns(message)) {
|
||||
return "ACCESS_DENIED"
|
||||
}
|
||||
|
||||
// Not found patterns
|
||||
if (this.matchesNotFoundPatterns(message)) {
|
||||
return "NOT_FOUND"
|
||||
}
|
||||
|
||||
// Network/timeout patterns
|
||||
if (this.matchesNetworkErrorPatterns(message)) {
|
||||
return "NETWORK_ERROR"
|
||||
}
|
||||
|
||||
if (this.matchesTimeoutPatterns(message)) {
|
||||
return "TIMEOUT"
|
||||
}
|
||||
}
|
||||
|
||||
// If it's an Error instance or has a message, classify as GENERIC
|
||||
// Otherwise classify as UNKNOWN
|
||||
if (error instanceof Error || (error as any).message) {
|
||||
return "GENERIC"
|
||||
}
|
||||
|
||||
return "UNKNOWN"
|
||||
}
|
||||
|
||||
/**
|
||||
* Determine if error should trigger a retry
|
||||
*/
|
||||
private static shouldRetryError(errorType: ErrorType, retryAttempt: number = 0): boolean {
|
||||
const MAX_RETRIES = 5
|
||||
|
||||
if (retryAttempt >= MAX_RETRIES) {
|
||||
return false
|
||||
}
|
||||
|
||||
const retryableTypes: ErrorType[] = [
|
||||
"THROTTLING",
|
||||
"RATE_LIMITED",
|
||||
"SERVICE_UNAVAILABLE",
|
||||
"TIMEOUT",
|
||||
"NETWORK_ERROR",
|
||||
"QUOTA_EXCEEDED",
|
||||
]
|
||||
|
||||
return retryableTypes.includes(errorType)
|
||||
}
|
||||
|
||||
/**
|
||||
* Determine if error should be thrown immediately (for proper retry handling)
|
||||
*/
|
||||
private static shouldThrowImmediately(errorType: ErrorType, isStreaming: boolean): boolean {
|
||||
// For throttling errors in streaming context, throw immediately for proper retry handling
|
||||
if ((errorType === "THROTTLING" || errorType === "RATE_LIMITED") && isStreaming) {
|
||||
return true
|
||||
}
|
||||
|
||||
// For other critical errors, throw immediately regardless of context
|
||||
const immediateThrowTypes: ErrorType[] = ["ACCESS_DENIED", "NOT_FOUND", "INVALID_REQUEST"]
|
||||
|
||||
return immediateThrowTypes.includes(errorType)
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculate appropriate retry delay based on error type
|
||||
*/
|
||||
private static calculateRetryDelay(errorType: ErrorType, error: unknown, retryAttempt: number = 0): number {
|
||||
// Default exponential backoff
|
||||
const baseDelay = 5 // seconds
|
||||
const maxDelay = 600 // 10 minutes
|
||||
|
||||
// Check for provider-specific retry information
|
||||
const providerDelay = this.extractProviderRetryDelay(error)
|
||||
if (providerDelay > 0) {
|
||||
return providerDelay
|
||||
}
|
||||
|
||||
// Calculate exponential backoff - for attempt 0, return base delay
|
||||
let exponentialDelay: number
|
||||
if (retryAttempt === 0) {
|
||||
exponentialDelay = baseDelay
|
||||
} else {
|
||||
exponentialDelay = Math.min(Math.ceil(baseDelay * Math.pow(2, retryAttempt)), maxDelay)
|
||||
}
|
||||
|
||||
// Adjust based on error type
|
||||
switch (errorType) {
|
||||
case "THROTTLING":
|
||||
case "RATE_LIMITED":
|
||||
return exponentialDelay
|
||||
case "SERVICE_UNAVAILABLE":
|
||||
return Math.min(exponentialDelay * 1.5, maxDelay) // Slightly longer for service issues
|
||||
case "QUOTA_EXCEEDED":
|
||||
return Math.min(exponentialDelay * 2, maxDelay) // Longer for quota issues
|
||||
case "NETWORK_ERROR":
|
||||
case "TIMEOUT":
|
||||
return Math.min(exponentialDelay * 0.5, maxDelay) // Shorter for network issues
|
||||
default:
|
||||
return exponentialDelay
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract provider-specific retry delay (e.g., Google Gemini retry info)
|
||||
*/
|
||||
private static extractProviderRetryDelay(error: unknown): number {
|
||||
if (!(error as any).errorDetails) return 0
|
||||
|
||||
// Google Gemini retry info
|
||||
const geminiRetryDetails = (error as any).errorDetails?.find(
|
||||
(detail: any) => detail["@type"] === "type.googleapis.com/google.rpc.RetryInfo",
|
||||
)
|
||||
|
||||
if (geminiRetryDetails?.retryDelay) {
|
||||
const match = geminiRetryDetails.retryDelay.match(/^(\d+)s$/)
|
||||
if (match) {
|
||||
return Number(match[1]) + 1 // Add 1 second buffer
|
||||
}
|
||||
}
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
/**
|
||||
* Format error message with context information
|
||||
*/
|
||||
private static formatErrorMessage(error: unknown, errorType: ErrorType, context: ErrorContext): string {
|
||||
let message = error instanceof Error ? error.message : "Unknown error"
|
||||
|
||||
// Clean up common noise in error messages
|
||||
message = message.replace(/\s+/g, " ").trim()
|
||||
|
||||
// Add context-specific information
|
||||
const contextInfo = `[${context.provider}:${context.modelId}]`
|
||||
|
||||
// Add retry information if applicable
|
||||
const retryInfo = context.retryAttempt ? ` (Retry ${context.retryAttempt})` : ""
|
||||
|
||||
// Add error type for debugging
|
||||
const typeInfo = `[${errorType}]`
|
||||
|
||||
return `${contextInfo} ${typeInfo} ${message}${retryInfo}`
|
||||
}
|
||||
|
||||
// Pattern matching helper methods
|
||||
private static matchesThrottlingPatterns(message: string): boolean {
|
||||
const patterns = [
|
||||
"throttl",
|
||||
"overloaded",
|
||||
"too many requests",
|
||||
"request limit",
|
||||
"concurrent requests",
|
||||
"bedrock is unable to process",
|
||||
]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesRateLimitPatterns(message: string): boolean {
|
||||
const patterns = ["rate", "limit", "please wait"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesQuotaPatterns(message: string): boolean {
|
||||
const patterns = ["quota exceeded", "quota", "billing", "credits"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesServiceUnavailablePatterns(message: string): boolean {
|
||||
const patterns = ["service unavailable", "busy", "temporarily unavailable", "server error"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesAccessDeniedPatterns(message: string): boolean {
|
||||
const patterns = ["access", "denied", "unauthorized", "forbidden", "permission"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesNotFoundPatterns(message: string): boolean {
|
||||
const patterns = ["not found", "does not exist", "invalid model"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesNetworkErrorPatterns(message: string): boolean {
|
||||
const patterns = ["network", "connection", "dns", "host", "socket"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
|
||||
private static matchesTimeoutPatterns(message: string): boolean {
|
||||
const patterns = ["timeout", "timed out", "deadline", "abort"]
|
||||
return patterns.some((pattern) => message.includes(pattern))
|
||||
}
|
||||
}
|
||||
|
|
@ -5,6 +5,7 @@ import type { ModelInfo } from "@roo-code/types"
|
|||
import type { ApiHandler, ApiHandlerCreateMessageMetadata } from "../index"
|
||||
import { ApiStream } from "../transform/stream"
|
||||
import { countTokens } from "../../utils/countTokens"
|
||||
import { UnifiedErrorHandler, ErrorContext, ErrorHandlerResponse } from "../error-handling/UnifiedErrorHandler"
|
||||
|
||||
/**
|
||||
* Base class for API providers that implements common functionality.
|
||||
|
|
@ -32,4 +33,34 @@ export abstract class BaseProvider implements ApiHandler {
|
|||
|
||||
return countTokens(content, { useWorker: true })
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle errors using the unified error handler
|
||||
*
|
||||
* @param error The error to handle
|
||||
* @param context Error context information
|
||||
* @returns Error handler response with retry/throw decisions
|
||||
*/
|
||||
protected handleError(error: unknown, context: ErrorContext): ErrorHandlerResponse {
|
||||
return UnifiedErrorHandler.handle(error, context)
|
||||
}
|
||||
|
||||
/**
|
||||
* Create error context for unified error handling
|
||||
*
|
||||
* @param isStreaming Whether the operation is streaming
|
||||
* @param retryAttempt Current retry attempt number
|
||||
* @param requestId Optional request identifier
|
||||
* @returns Error context object
|
||||
*/
|
||||
protected createErrorContext(isStreaming: boolean, retryAttempt?: number, requestId?: string): ErrorContext {
|
||||
const model = this.getModel()
|
||||
return {
|
||||
isStreaming,
|
||||
provider: model.id,
|
||||
modelId: model.id,
|
||||
retryAttempt,
|
||||
requestId,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
320
src/core/task/StreamStateManager.test.ts
Normal file
320
src/core/task/StreamStateManager.test.ts
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
import { describe, test, expect, beforeEach, vi } from "vitest"
|
||||
import { StreamStateManager } from "./StreamStateManager"
|
||||
import { ClineApiReqCancelReason } from "../../shared/ExtensionMessage"
|
||||
|
||||
// Mock Task class for testing
|
||||
class MockTask {
|
||||
public isStreaming = false
|
||||
public currentStreamingContentIndex = 0
|
||||
public assistantMessageContent: any[] = []
|
||||
public presentAssistantMessageLocked = false
|
||||
public presentAssistantMessageHasPendingUpdates = false
|
||||
public userMessageContent: any[] = []
|
||||
public userMessageContentReady = false
|
||||
public didRejectTool = false
|
||||
public didAlreadyUseTool = false
|
||||
public didCompleteReadingStream = false
|
||||
public didFinishAbortingStream = false
|
||||
public isWaitingForFirstChunk = false
|
||||
public abort = false
|
||||
public abandoned = false
|
||||
|
||||
public clineMessages: any[] = []
|
||||
public diffViewProvider = {
|
||||
isEditing: false,
|
||||
revertChanges: vi.fn().mockResolvedValue(undefined),
|
||||
reset: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
|
||||
// Mock private methods that StreamStateManager needs to access
|
||||
public saveClineMessages = vi.fn().mockResolvedValue(undefined)
|
||||
public addToApiConversationHistory = vi.fn().mockResolvedValue(undefined)
|
||||
}
|
||||
|
||||
describe("StreamStateManager", () => {
|
||||
let mockTask: MockTask
|
||||
let streamStateManager: StreamStateManager
|
||||
|
||||
beforeEach(() => {
|
||||
mockTask = new MockTask()
|
||||
streamStateManager = new StreamStateManager(mockTask as any)
|
||||
})
|
||||
|
||||
describe("initialization", () => {
|
||||
test("captures initial state correctly", () => {
|
||||
const snapshot = streamStateManager.getStreamStateSnapshot()
|
||||
|
||||
expect(snapshot.isStreaming).toBe(false)
|
||||
expect(snapshot.currentStreamingContentIndex).toBe(0)
|
||||
expect(snapshot.presentAssistantMessageLocked).toBe(false)
|
||||
expect(snapshot.presentAssistantMessageHasPendingUpdates).toBe(false)
|
||||
expect(snapshot.userMessageContentReady).toBe(false)
|
||||
expect(snapshot.didRejectTool).toBe(false)
|
||||
expect(snapshot.didAlreadyUseTool).toBe(false)
|
||||
expect(snapshot.didCompleteReadingStream).toBe(false)
|
||||
expect(snapshot.didFinishAbortingStream).toBe(false)
|
||||
expect(snapshot.isWaitingForFirstChunk).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("resetToInitialState", () => {
|
||||
test("resets all streaming state to initial values", async () => {
|
||||
// Modify state to non-initial values
|
||||
mockTask.isStreaming = true
|
||||
mockTask.currentStreamingContentIndex = 5
|
||||
mockTask.assistantMessageContent = [{ type: "text", content: "test" }]
|
||||
mockTask.presentAssistantMessageLocked = true
|
||||
mockTask.presentAssistantMessageHasPendingUpdates = true
|
||||
mockTask.userMessageContent = [{ type: "text", text: "user message" }]
|
||||
mockTask.userMessageContentReady = true
|
||||
mockTask.didRejectTool = true
|
||||
mockTask.didAlreadyUseTool = true
|
||||
mockTask.didCompleteReadingStream = true
|
||||
mockTask.didFinishAbortingStream = true
|
||||
mockTask.isWaitingForFirstChunk = true
|
||||
|
||||
// Reset state
|
||||
await streamStateManager.resetToInitialState()
|
||||
|
||||
// Verify all properties are reset
|
||||
expect(mockTask.isStreaming).toBe(false)
|
||||
expect(mockTask.currentStreamingContentIndex).toBe(0)
|
||||
expect(mockTask.assistantMessageContent).toHaveLength(0)
|
||||
expect(mockTask.presentAssistantMessageLocked).toBe(false)
|
||||
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false)
|
||||
expect(mockTask.userMessageContent).toHaveLength(0)
|
||||
expect(mockTask.userMessageContentReady).toBe(false)
|
||||
expect(mockTask.didRejectTool).toBe(false)
|
||||
expect(mockTask.didAlreadyUseTool).toBe(false)
|
||||
expect(mockTask.didCompleteReadingStream).toBe(false)
|
||||
expect(mockTask.didFinishAbortingStream).toBe(false)
|
||||
expect(mockTask.isWaitingForFirstChunk).toBe(false)
|
||||
})
|
||||
|
||||
test("reverts diff changes when editing", async () => {
|
||||
mockTask.diffViewProvider.isEditing = true
|
||||
|
||||
await streamStateManager.resetToInitialState()
|
||||
|
||||
expect(mockTask.diffViewProvider.revertChanges).toHaveBeenCalled()
|
||||
expect(mockTask.diffViewProvider.reset).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
test("continues reset even if diff operations fail", async () => {
|
||||
mockTask.diffViewProvider.isEditing = true
|
||||
mockTask.diffViewProvider.revertChanges.mockRejectedValue(new Error("Diff error"))
|
||||
mockTask.isStreaming = true
|
||||
|
||||
await streamStateManager.resetToInitialState()
|
||||
|
||||
// State should still be reset despite diff error
|
||||
expect(mockTask.isStreaming).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("abortStreamSafely", () => {
|
||||
test("performs comprehensive cleanup on abort", async () => {
|
||||
// Set up state that needs cleanup
|
||||
mockTask.isStreaming = true
|
||||
mockTask.diffViewProvider.isEditing = true
|
||||
mockTask.assistantMessageContent = [{ type: "text", content: "partial message" }]
|
||||
mockTask.clineMessages = [
|
||||
{ ts: Date.now(), type: "say", say: "api_req_started", text: "test", partial: true },
|
||||
]
|
||||
|
||||
await streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
// Verify cleanup was performed
|
||||
expect(mockTask.diffViewProvider.revertChanges).toHaveBeenCalled()
|
||||
expect(mockTask.didFinishAbortingStream).toBe(true)
|
||||
expect(mockTask.saveClineMessages).toHaveBeenCalled()
|
||||
expect(mockTask.addToApiConversationHistory).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
test("handles partial message cleanup", async () => {
|
||||
const partialMessage = {
|
||||
ts: Date.now(),
|
||||
type: "say",
|
||||
say: "api_req_started",
|
||||
text: "test",
|
||||
partial: true,
|
||||
}
|
||||
mockTask.clineMessages = [partialMessage]
|
||||
|
||||
await streamStateManager.abortStreamSafely("streaming_failed", "Connection error")
|
||||
|
||||
expect(partialMessage.partial).toBe(false)
|
||||
expect(mockTask.saveClineMessages).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
test("adds interruption message to conversation history", async () => {
|
||||
mockTask.assistantMessageContent = [{ type: "text", content: "This is a partial response" }]
|
||||
|
||||
await streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
expect(mockTask.addToApiConversationHistory).toHaveBeenCalledWith({
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "This is a partial response\n\n[Response interrupted by user]",
|
||||
},
|
||||
],
|
||||
})
|
||||
})
|
||||
|
||||
test("adds API error interruption message", async () => {
|
||||
mockTask.assistantMessageContent = [{ type: "text", content: "Partial response" }]
|
||||
|
||||
await streamStateManager.abortStreamSafely("streaming_failed", "API timeout")
|
||||
|
||||
expect(mockTask.addToApiConversationHistory).toHaveBeenCalledWith({
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "Partial response\n\n[Response interrupted by API Error]",
|
||||
},
|
||||
],
|
||||
})
|
||||
})
|
||||
|
||||
test("prevents concurrent abort operations", async () => {
|
||||
// Set up state that triggers diff cleanup
|
||||
mockTask.diffViewProvider.isEditing = true
|
||||
|
||||
// Start first abort
|
||||
const firstAbort = streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
// Start second abort immediately (should be ignored)
|
||||
const secondAbort = streamStateManager.abortStreamSafely("streaming_failed")
|
||||
|
||||
await Promise.all([firstAbort, secondAbort])
|
||||
|
||||
// Only one cleanup operation should have occurred
|
||||
expect(mockTask.diffViewProvider.revertChanges).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
test("ensures didFinishAbortingStream is always set", async () => {
|
||||
// Simulate error during cleanup
|
||||
mockTask.diffViewProvider.revertChanges.mockRejectedValue(new Error("Cleanup failed"))
|
||||
|
||||
await streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
expect(mockTask.didFinishAbortingStream).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("stream lifecycle management", () => {
|
||||
test("prepareForStreaming resets state and sets initial values", async () => {
|
||||
// Set dirty state
|
||||
mockTask.isStreaming = true
|
||||
mockTask.didCompleteReadingStream = true
|
||||
|
||||
await streamStateManager.prepareForStreaming()
|
||||
|
||||
expect(mockTask.isStreaming).toBe(false)
|
||||
expect(mockTask.isWaitingForFirstChunk).toBe(false)
|
||||
expect(mockTask.didCompleteReadingStream).toBe(false)
|
||||
expect(mockTask.didFinishAbortingStream).toBe(false)
|
||||
})
|
||||
|
||||
test("markStreamingStarted updates streaming state", () => {
|
||||
streamStateManager.markStreamingStarted()
|
||||
|
||||
expect(mockTask.isStreaming).toBe(true)
|
||||
expect(mockTask.isWaitingForFirstChunk).toBe(false)
|
||||
})
|
||||
|
||||
test("markStreamingCompleted updates completion state", () => {
|
||||
mockTask.isStreaming = true
|
||||
|
||||
streamStateManager.markStreamingCompleted()
|
||||
|
||||
expect(mockTask.isStreaming).toBe(false)
|
||||
expect(mockTask.didCompleteReadingStream).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("safety checks", () => {
|
||||
test("isStreamSafe returns true for safe conditions", () => {
|
||||
expect(streamStateManager.isStreamSafe()).toBe(true)
|
||||
})
|
||||
|
||||
test("isStreamSafe returns false when task is aborted", () => {
|
||||
mockTask.abort = true
|
||||
|
||||
expect(streamStateManager.isStreamSafe()).toBe(false)
|
||||
})
|
||||
|
||||
test("isStreamSafe returns false when task is abandoned", () => {
|
||||
mockTask.abandoned = true
|
||||
|
||||
expect(streamStateManager.isStreamSafe()).toBe(false)
|
||||
})
|
||||
|
||||
test("isStreamSafe returns false when aborting in progress", async () => {
|
||||
// Start an abort operation
|
||||
const abortPromise = streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
// Should not be safe during abort
|
||||
expect(streamStateManager.isStreamSafe()).toBe(false)
|
||||
|
||||
await abortPromise
|
||||
|
||||
// Should be safe again after abort completes
|
||||
expect(streamStateManager.isStreamSafe()).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("getStreamStateSnapshot", () => {
|
||||
test("returns current state snapshot", () => {
|
||||
mockTask.isStreaming = true
|
||||
mockTask.currentStreamingContentIndex = 3
|
||||
mockTask.userMessageContentReady = true
|
||||
|
||||
const snapshot = streamStateManager.getStreamStateSnapshot()
|
||||
|
||||
expect(snapshot.isStreaming).toBe(true)
|
||||
expect(snapshot.currentStreamingContentIndex).toBe(3)
|
||||
expect(snapshot.userMessageContentReady).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("forceCleanup", () => {
|
||||
test("performs emergency cleanup", () => {
|
||||
mockTask.isStreaming = true
|
||||
mockTask.assistantMessageContent = [{ type: "text", content: "test" }]
|
||||
mockTask.userMessageContent = [{ type: "text", text: "test" }]
|
||||
|
||||
streamStateManager.forceCleanup()
|
||||
|
||||
expect(mockTask.isStreaming).toBe(false)
|
||||
expect(mockTask.didFinishAbortingStream).toBe(true)
|
||||
expect(mockTask.assistantMessageContent).toHaveLength(0)
|
||||
expect(mockTask.userMessageContent).toHaveLength(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe("error handling", () => {
|
||||
test("continues operation when partial message cleanup fails", async () => {
|
||||
mockTask.clineMessages = [{ partial: true }]
|
||||
mockTask.saveClineMessages.mockRejectedValue(new Error("Save failed"))
|
||||
|
||||
// Should not throw
|
||||
await streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
expect(mockTask.didFinishAbortingStream).toBe(true)
|
||||
})
|
||||
|
||||
test("continues operation when history update fails", async () => {
|
||||
mockTask.assistantMessageContent = [{ type: "text", content: "test" }]
|
||||
mockTask.addToApiConversationHistory.mockRejectedValue(new Error("History failed"))
|
||||
|
||||
// Should not throw
|
||||
await streamStateManager.abortStreamSafely("user_cancelled")
|
||||
|
||||
expect(mockTask.didFinishAbortingStream).toBe(true)
|
||||
})
|
||||
})
|
||||
})
|
||||
254
src/core/task/StreamStateManager.ts
Normal file
254
src/core/task/StreamStateManager.ts
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
import { Anthropic } from "@anthropic-ai/sdk"
|
||||
import type { AssistantMessageContent } from "../assistant-message"
|
||||
import type { Task } from "./Task"
|
||||
import { ClineApiReqCancelReason } from "../../shared/ExtensionMessage"
|
||||
|
||||
/**
|
||||
* StreamState - Interface defining all streaming-related state properties
|
||||
*/
|
||||
export interface StreamState {
|
||||
isStreaming: boolean
|
||||
currentStreamingContentIndex: number
|
||||
assistantMessageContent: AssistantMessageContent[]
|
||||
presentAssistantMessageLocked: boolean
|
||||
presentAssistantMessageHasPendingUpdates: boolean
|
||||
userMessageContent: Anthropic.Messages.ContentBlockParam[]
|
||||
userMessageContentReady: boolean
|
||||
didRejectTool: boolean
|
||||
didAlreadyUseTool: boolean
|
||||
didCompleteReadingStream: boolean
|
||||
didFinishAbortingStream: boolean
|
||||
isWaitingForFirstChunk: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
* StreamStateManager - Manages comprehensive stream state during API calls
|
||||
*
|
||||
* This class provides atomic stream state management to prevent corruption during
|
||||
* retry cycles, ensuring proper cleanup and coordination between streaming operations.
|
||||
*/
|
||||
export class StreamStateManager {
|
||||
private task: Task
|
||||
private initialState: Partial<StreamState> = {}
|
||||
private isAborting: boolean = false
|
||||
|
||||
constructor(task: Task) {
|
||||
this.task = task
|
||||
this.captureInitialState()
|
||||
}
|
||||
|
||||
/**
|
||||
* Capture the initial clean state for reset operations
|
||||
*/
|
||||
private captureInitialState(): void {
|
||||
this.initialState = {
|
||||
isStreaming: false,
|
||||
currentStreamingContentIndex: 0,
|
||||
assistantMessageContent: [],
|
||||
presentAssistantMessageLocked: false,
|
||||
presentAssistantMessageHasPendingUpdates: false,
|
||||
userMessageContent: [],
|
||||
userMessageContentReady: false,
|
||||
didRejectTool: false,
|
||||
didAlreadyUseTool: false,
|
||||
didCompleteReadingStream: false,
|
||||
didFinishAbortingStream: false,
|
||||
isWaitingForFirstChunk: false,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Atomically reset all streaming state to initial clean state
|
||||
*/
|
||||
async resetToInitialState(): Promise<void> {
|
||||
// Ensure no concurrent abort operations
|
||||
if (this.isAborting) {
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
// Ensure any pending diff operations are reverted
|
||||
if (this.task.diffViewProvider.isEditing) {
|
||||
await this.task.diffViewProvider.revertChanges()
|
||||
}
|
||||
|
||||
// Reset all streaming state atomically
|
||||
Object.assign(this.task, this.initialState)
|
||||
|
||||
// Clear any partial content arrays
|
||||
this.task.assistantMessageContent.length = 0
|
||||
this.task.userMessageContent.length = 0
|
||||
|
||||
// Reset diff provider state
|
||||
await this.task.diffViewProvider.reset()
|
||||
} catch (error) {
|
||||
console.error("Error resetting stream state:", error)
|
||||
// Continue with state reset even if diff operations fail
|
||||
Object.assign(this.task, this.initialState)
|
||||
this.task.assistantMessageContent.length = 0
|
||||
this.task.userMessageContent.length = 0
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Safely abort a stream with comprehensive cleanup
|
||||
*/
|
||||
async abortStreamSafely(cancelReason: ClineApiReqCancelReason, streamingFailedMessage?: string): Promise<void> {
|
||||
// Prevent concurrent abort operations
|
||||
if (this.isAborting) {
|
||||
return
|
||||
}
|
||||
|
||||
this.isAborting = true
|
||||
|
||||
// Mark as aborting to prevent concurrent operations
|
||||
this.task.didFinishAbortingStream = false
|
||||
|
||||
try {
|
||||
// Revert any pending changes first
|
||||
if (this.task.diffViewProvider.isEditing) {
|
||||
await this.task.diffViewProvider.revertChanges()
|
||||
}
|
||||
|
||||
// Handle partial messages consistently
|
||||
await this.handlePartialMessageCleanup()
|
||||
|
||||
// Add interruption message to conversation history
|
||||
await this.addInterruptionToHistory(cancelReason, streamingFailedMessage)
|
||||
|
||||
// Reset stream state
|
||||
await this.resetToInitialState()
|
||||
} catch (error) {
|
||||
console.error("Error during stream abort:", error)
|
||||
// Ensure state is reset even if cleanup fails
|
||||
await this.resetToInitialState()
|
||||
} finally {
|
||||
// Always mark as finished aborting and reset abort flag
|
||||
this.task.didFinishAbortingStream = true
|
||||
this.isAborting = false
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle cleanup of partial messages in conversation history
|
||||
*/
|
||||
private async handlePartialMessageCleanup(): Promise<void> {
|
||||
try {
|
||||
const lastMessage = this.task.clineMessages.at(-1)
|
||||
if (lastMessage && lastMessage.partial) {
|
||||
lastMessage.partial = false
|
||||
// Use the public method or delegate to task
|
||||
await (this.task as any).saveClineMessages()
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error cleaning up partial messages:", error)
|
||||
// Don't throw - this is cleanup, continue with abort
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Add interruption message to API conversation history
|
||||
*/
|
||||
private async addInterruptionToHistory(
|
||||
cancelReason: ClineApiReqCancelReason,
|
||||
streamingFailedMessage?: string,
|
||||
): Promise<void> {
|
||||
try {
|
||||
// Reconstruct assistant message from current content
|
||||
let assistantMessage = ""
|
||||
for (const content of this.task.assistantMessageContent) {
|
||||
if (content.type === "text") {
|
||||
assistantMessage += content.content
|
||||
}
|
||||
}
|
||||
|
||||
if (assistantMessage) {
|
||||
const interruptionText = `\n\n[${
|
||||
cancelReason === "streaming_failed"
|
||||
? "Response interrupted by API Error"
|
||||
: "Response interrupted by user"
|
||||
}]`
|
||||
|
||||
await (this.task as any).addToApiConversationHistory({
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: assistantMessage + interruptionText,
|
||||
},
|
||||
],
|
||||
})
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error adding interruption to history:", error)
|
||||
// Don't throw - this is cleanup, continue with abort
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare for a new streaming operation
|
||||
*/
|
||||
async prepareForStreaming(): Promise<void> {
|
||||
// Ensure clean state before starting new stream
|
||||
await this.resetToInitialState()
|
||||
|
||||
// Set initial streaming state
|
||||
this.task.isStreaming = false // Will be set to true when stream starts
|
||||
this.task.isWaitingForFirstChunk = false
|
||||
this.task.didCompleteReadingStream = false
|
||||
this.task.didFinishAbortingStream = false
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark streaming as started
|
||||
*/
|
||||
markStreamingStarted(): void {
|
||||
this.task.isStreaming = true
|
||||
this.task.isWaitingForFirstChunk = false
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark streaming as completed
|
||||
*/
|
||||
markStreamingCompleted(): void {
|
||||
this.task.isStreaming = false
|
||||
this.task.didCompleteReadingStream = true
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if stream is in a safe state for operations
|
||||
*/
|
||||
isStreamSafe(): boolean {
|
||||
return !this.isAborting && !this.task.abort && !this.task.abandoned
|
||||
}
|
||||
|
||||
/**
|
||||
* Get current stream state snapshot for debugging
|
||||
*/
|
||||
getStreamStateSnapshot(): Partial<StreamState> {
|
||||
return {
|
||||
isStreaming: this.task.isStreaming,
|
||||
currentStreamingContentIndex: this.task.currentStreamingContentIndex,
|
||||
presentAssistantMessageLocked: this.task.presentAssistantMessageLocked,
|
||||
presentAssistantMessageHasPendingUpdates: this.task.presentAssistantMessageHasPendingUpdates,
|
||||
userMessageContentReady: this.task.userMessageContentReady,
|
||||
didRejectTool: this.task.didRejectTool,
|
||||
didAlreadyUseTool: this.task.didAlreadyUseTool,
|
||||
didCompleteReadingStream: this.task.didCompleteReadingStream,
|
||||
didFinishAbortingStream: this.task.didFinishAbortingStream,
|
||||
isWaitingForFirstChunk: this.task.isWaitingForFirstChunk,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Force cleanup - for emergency situations
|
||||
* @internal
|
||||
*/
|
||||
forceCleanup(): void {
|
||||
this.isAborting = false
|
||||
this.task.isStreaming = false
|
||||
this.task.didFinishAbortingStream = true
|
||||
this.task.assistantMessageContent.length = 0
|
||||
this.task.userMessageContent.length = 0
|
||||
}
|
||||
}
|
||||
|
|
@ -28,6 +28,7 @@ import { CloudService } from "@roo-code/cloud"
|
|||
// api
|
||||
import { ApiHandler, ApiHandlerCreateMessageMetadata, buildApiHandler } from "../../api"
|
||||
import { ApiStream } from "../../api/transform/stream"
|
||||
import { UnifiedErrorHandler, ErrorContext } from "../../api/error-handling/UnifiedErrorHandler"
|
||||
|
||||
// shared
|
||||
import { findLastIndex } from "../../shared/array"
|
||||
|
|
@ -88,6 +89,10 @@ import { getMessagesSinceLastSummary, summarizeConversation } from "../condense"
|
|||
import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
|
||||
import { restoreTodoListForTask } from "../tools/updateTodoListTool"
|
||||
|
||||
// State management
|
||||
import { TaskStateLock, GlobalRateLimitManager } from "./TaskStateLock"
|
||||
import { StreamStateManager } from "./StreamStateManager"
|
||||
|
||||
// Constants
|
||||
const MAX_EXPONENTIAL_BACKOFF_SECONDS = 600 // 10 minutes
|
||||
|
||||
|
|
@ -146,7 +151,6 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
// API
|
||||
readonly apiConfiguration: ProviderSettings
|
||||
api: ApiHandler
|
||||
private static lastGlobalApiRequestTime?: number
|
||||
private consecutiveAutoApprovedRequestsCount: number = 0
|
||||
|
||||
/**
|
||||
|
|
@ -154,7 +158,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
* @internal
|
||||
*/
|
||||
static resetGlobalApiRequestTime(): void {
|
||||
Task.lastGlobalApiRequestTime = undefined
|
||||
GlobalRateLimitManager.reset()
|
||||
}
|
||||
|
||||
toolRepetitionDetector: ToolRepetitionDetector
|
||||
|
|
@ -208,6 +212,9 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
didAlreadyUseTool = false
|
||||
didCompleteReadingStream = false
|
||||
|
||||
// Stream state management
|
||||
private streamStateManager: StreamStateManager
|
||||
|
||||
constructor({
|
||||
provider,
|
||||
apiConfiguration,
|
||||
|
|
@ -288,6 +295,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
}
|
||||
|
||||
this.toolRepetitionDetector = new ToolRepetitionDetector(this.consecutiveMistakeLimit)
|
||||
this.streamStateManager = new StreamStateManager(this)
|
||||
|
||||
onCreated?.(this)
|
||||
|
||||
|
|
@ -1326,17 +1334,8 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
this.didFinishAbortingStream = true
|
||||
}
|
||||
|
||||
// Reset streaming state.
|
||||
this.currentStreamingContentIndex = 0
|
||||
this.assistantMessageContent = []
|
||||
this.didCompleteReadingStream = false
|
||||
this.userMessageContent = []
|
||||
this.userMessageContentReady = false
|
||||
this.didRejectTool = false
|
||||
this.didAlreadyUseTool = false
|
||||
this.presentAssistantMessageLocked = false
|
||||
this.presentAssistantMessageHasPendingUpdates = false
|
||||
|
||||
// Reset streaming state using StreamStateManager
|
||||
await this.streamStateManager.resetToInitialState()
|
||||
await this.diffViewProvider.reset()
|
||||
|
||||
// Yields only if the first chunk is successful, otherwise will
|
||||
|
|
@ -1683,12 +1682,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
|
||||
// Use the shared timestamp so that subtasks respect the same rate-limit
|
||||
// window as their parent tasks.
|
||||
if (Task.lastGlobalApiRequestTime) {
|
||||
const now = Date.now()
|
||||
const timeSinceLastRequest = now - Task.lastGlobalApiRequestTime
|
||||
const rateLimit = apiConfiguration?.rateLimitSeconds || 0
|
||||
rateLimitDelay = Math.ceil(Math.max(0, rateLimit * 1000 - timeSinceLastRequest) / 1000)
|
||||
}
|
||||
rateLimitDelay = await GlobalRateLimitManager.calculateRateLimitDelay(apiConfiguration?.rateLimitSeconds || 0)
|
||||
|
||||
// Only show rate limiting message if we're not retrying. If retrying, we'll include the delay there.
|
||||
if (rateLimitDelay > 0 && retryAttempt === 0) {
|
||||
|
|
@ -1702,7 +1696,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
|
||||
// Update last request time before making the request so that subsequent
|
||||
// requests — even from new subtasks — will honour the provider's rate-limit.
|
||||
Task.lastGlobalApiRequestTime = Date.now()
|
||||
await GlobalRateLimitManager.updateLastRequestTime()
|
||||
|
||||
const systemPrompt = await this.getSystemPrompt()
|
||||
const { contextTokens } = this.getTokenUsage()
|
||||
|
|
@ -1795,36 +1789,24 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
this.isWaitingForFirstChunk = false
|
||||
} catch (error) {
|
||||
this.isWaitingForFirstChunk = false
|
||||
|
||||
// Use UnifiedErrorHandler for consistent error handling
|
||||
const errorContext: ErrorContext = {
|
||||
isStreaming: false, // First chunk failure, not streaming yet
|
||||
provider: this.api.getModel().id,
|
||||
modelId: this.api.getModel().id,
|
||||
retryAttempt,
|
||||
requestId: metadata.taskId,
|
||||
}
|
||||
|
||||
const errorResponse = UnifiedErrorHandler.handle(error, errorContext)
|
||||
|
||||
// note that this api_req_failed ask is unique in that we only present this option if the api hasn't streamed any content yet (ie it fails on the first chunk due), as it would allow them to hit a retry button. However if the api failed mid-stream, it could be in any arbitrary state where some tools may have executed, so that error is handled differently and requires cancelling the task entirely.
|
||||
if (autoApprovalEnabled && alwaysApproveResubmit) {
|
||||
let errorMsg
|
||||
|
||||
if (error.error?.metadata?.raw) {
|
||||
errorMsg = JSON.stringify(error.error.metadata.raw, null, 2)
|
||||
} else if (error.message) {
|
||||
errorMsg = error.message
|
||||
} else {
|
||||
errorMsg = "Unknown error"
|
||||
}
|
||||
|
||||
if (autoApprovalEnabled && alwaysApproveResubmit && errorResponse.shouldRetry) {
|
||||
const baseDelay = requestDelaySeconds || 5
|
||||
let exponentialDelay = Math.min(
|
||||
Math.ceil(baseDelay * Math.pow(2, retryAttempt)),
|
||||
MAX_EXPONENTIAL_BACKOFF_SECONDS,
|
||||
)
|
||||
|
||||
// If the error is a 429, and the error details contain a retry delay, use that delay instead of exponential backoff
|
||||
if (error.status === 429) {
|
||||
const geminiRetryDetails = error.errorDetails?.find(
|
||||
(detail: any) => detail["@type"] === "type.googleapis.com/google.rpc.RetryInfo",
|
||||
)
|
||||
if (geminiRetryDetails) {
|
||||
const match = geminiRetryDetails?.retryDelay?.match(/^(\d+)s$/)
|
||||
if (match) {
|
||||
exponentialDelay = Number(match[1]) + 1
|
||||
}
|
||||
}
|
||||
}
|
||||
let exponentialDelay =
|
||||
errorResponse.retryDelay ||
|
||||
Math.min(Math.ceil(baseDelay * Math.pow(2, retryAttempt)), MAX_EXPONENTIAL_BACKOFF_SECONDS)
|
||||
|
||||
// Wait for the greater of the exponential delay or the rate limit delay
|
||||
const finalDelay = Math.max(exponentialDelay, rateLimitDelay)
|
||||
|
|
@ -1833,7 +1815,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
for (let i = finalDelay; i > 0; i--) {
|
||||
await this.say(
|
||||
"api_req_retry_delayed",
|
||||
`${errorMsg}\n\nRetry attempt ${retryAttempt + 1}\nRetrying in ${i} seconds...`,
|
||||
`${errorResponse.formattedMessage}\n\nRetry attempt ${retryAttempt + 1}\nRetrying in ${i} seconds...`,
|
||||
undefined,
|
||||
true,
|
||||
)
|
||||
|
|
@ -1842,7 +1824,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
|
||||
await this.say(
|
||||
"api_req_retry_delayed",
|
||||
`${errorMsg}\n\nRetry attempt ${retryAttempt + 1}\nRetrying now...`,
|
||||
`${errorResponse.formattedMessage}\n\nRetry attempt ${retryAttempt + 1}\nRetrying now...`,
|
||||
undefined,
|
||||
false,
|
||||
)
|
||||
|
|
@ -1853,10 +1835,7 @@ export class Task extends EventEmitter<ClineEvents> {
|
|||
|
||||
return
|
||||
} else {
|
||||
const { response } = await this.ask(
|
||||
"api_req_failed",
|
||||
error.message ?? JSON.stringify(serializeError(error), null, 2),
|
||||
)
|
||||
const { response } = await this.ask("api_req_failed", errorResponse.formattedMessage)
|
||||
|
||||
if (response !== "yesButtonClicked") {
|
||||
// This will never happen since if noButtonClicked, we will
|
||||
|
|
|
|||
192
src/core/task/TaskStateLock.test.ts
Normal file
192
src/core/task/TaskStateLock.test.ts
Normal file
|
|
@ -0,0 +1,192 @@
|
|||
import { describe, test, expect, beforeEach, afterEach } from "vitest"
|
||||
import { TaskStateLock, GlobalRateLimitManager } from "./TaskStateLock"
|
||||
|
||||
describe("TaskStateLock", () => {
|
||||
afterEach(() => {
|
||||
TaskStateLock.clearAllLocks()
|
||||
})
|
||||
|
||||
test("acquire and release lock", async () => {
|
||||
const lockKey = "test-lock"
|
||||
|
||||
// Acquire lock
|
||||
const release = await TaskStateLock.acquire(lockKey)
|
||||
expect(TaskStateLock.isLocked(lockKey)).toBe(true)
|
||||
|
||||
// Release lock
|
||||
release()
|
||||
expect(TaskStateLock.isLocked(lockKey)).toBe(false)
|
||||
})
|
||||
|
||||
test("tryAcquire returns null when lock is held", async () => {
|
||||
const lockKey = "test-lock"
|
||||
|
||||
// Acquire lock
|
||||
const release = await TaskStateLock.acquire(lockKey)
|
||||
|
||||
// Try to acquire the same lock should fail
|
||||
const tryResult = TaskStateLock.tryAcquire(lockKey)
|
||||
expect(tryResult).toBeNull()
|
||||
|
||||
// Release and try again should succeed
|
||||
release()
|
||||
const tryResult2 = TaskStateLock.tryAcquire(lockKey)
|
||||
expect(tryResult2).not.toBeNull()
|
||||
|
||||
if (tryResult2) {
|
||||
tryResult2()
|
||||
}
|
||||
})
|
||||
|
||||
test("withLock executes function with exclusive access", async () => {
|
||||
const lockKey = "test-lock"
|
||||
let counter = 0
|
||||
|
||||
// Start two concurrent operations
|
||||
const promise1 = TaskStateLock.withLock(lockKey, async () => {
|
||||
const initialValue = counter
|
||||
await new Promise((resolve) => setTimeout(resolve, 10))
|
||||
counter = initialValue + 1
|
||||
return "result1"
|
||||
})
|
||||
|
||||
const promise2 = TaskStateLock.withLock(lockKey, async () => {
|
||||
const initialValue = counter
|
||||
await new Promise((resolve) => setTimeout(resolve, 10))
|
||||
counter = initialValue + 1
|
||||
return "result2"
|
||||
})
|
||||
|
||||
const [result1, result2] = await Promise.all([promise1, promise2])
|
||||
|
||||
// Both operations should complete but counter should be 2 (not corrupted)
|
||||
expect(counter).toBe(2)
|
||||
expect([result1, result2]).toEqual(["result1", "result2"])
|
||||
})
|
||||
|
||||
test("multiple different locks can be held simultaneously", async () => {
|
||||
const lock1 = "lock-1"
|
||||
const lock2 = "lock-2"
|
||||
|
||||
const release1 = await TaskStateLock.acquire(lock1)
|
||||
const release2 = await TaskStateLock.acquire(lock2)
|
||||
|
||||
expect(TaskStateLock.isLocked(lock1)).toBe(true)
|
||||
expect(TaskStateLock.isLocked(lock2)).toBe(true)
|
||||
|
||||
release1()
|
||||
release2()
|
||||
|
||||
expect(TaskStateLock.isLocked(lock1)).toBe(false)
|
||||
expect(TaskStateLock.isLocked(lock2)).toBe(false)
|
||||
})
|
||||
|
||||
test("clearAllLocks clears all active locks", async () => {
|
||||
const lock1 = "lock-1"
|
||||
const lock2 = "lock-2"
|
||||
|
||||
await TaskStateLock.acquire(lock1)
|
||||
await TaskStateLock.acquire(lock2)
|
||||
|
||||
expect(TaskStateLock.isLocked(lock1)).toBe(true)
|
||||
expect(TaskStateLock.isLocked(lock2)).toBe(true)
|
||||
|
||||
TaskStateLock.clearAllLocks()
|
||||
|
||||
expect(TaskStateLock.isLocked(lock1)).toBe(false)
|
||||
expect(TaskStateLock.isLocked(lock2)).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("GlobalRateLimitManager", () => {
|
||||
beforeEach(() => {
|
||||
GlobalRateLimitManager.reset()
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
GlobalRateLimitManager.reset()
|
||||
})
|
||||
|
||||
test("updateLastRequestTime sets current timestamp", async () => {
|
||||
const before = Date.now()
|
||||
const timestamp = await GlobalRateLimitManager.updateLastRequestTime()
|
||||
const after = Date.now()
|
||||
|
||||
expect(timestamp).toBeGreaterThanOrEqual(before)
|
||||
expect(timestamp).toBeLessThanOrEqual(after)
|
||||
|
||||
const retrieved = await GlobalRateLimitManager.getLastRequestTime()
|
||||
expect(retrieved).toBe(timestamp)
|
||||
})
|
||||
|
||||
test("getLastRequestTime returns undefined when not set", async () => {
|
||||
const result = await GlobalRateLimitManager.getLastRequestTime()
|
||||
expect(result).toBeUndefined()
|
||||
})
|
||||
|
||||
test("calculateRateLimitDelay returns 0 when no previous request", async () => {
|
||||
const delay = await GlobalRateLimitManager.calculateRateLimitDelay(5)
|
||||
expect(delay).toBe(0)
|
||||
})
|
||||
|
||||
test("calculateRateLimitDelay calculates correct delay", async () => {
|
||||
// Set a request time 2 seconds ago
|
||||
const now = Date.now()
|
||||
const twoSecondsAgo = now - 2000
|
||||
|
||||
// Manually set the timestamp by updating then overriding
|
||||
await GlobalRateLimitManager.updateLastRequestTime()
|
||||
// We need to access the private field for testing - using bracket notation
|
||||
;(GlobalRateLimitManager as any).lastApiRequestTime = twoSecondsAgo
|
||||
|
||||
// With 5 second rate limit, should need to wait ~3 more seconds
|
||||
const delay = await GlobalRateLimitManager.calculateRateLimitDelay(5)
|
||||
expect(delay).toBeGreaterThanOrEqual(2)
|
||||
expect(delay).toBeLessThanOrEqual(4) // Allow some timing variance
|
||||
})
|
||||
|
||||
test("calculateRateLimitDelay returns 0 when enough time has passed", async () => {
|
||||
// Set a request time 10 seconds ago
|
||||
const tenSecondsAgo = Date.now() - 10000
|
||||
|
||||
await GlobalRateLimitManager.updateLastRequestTime()
|
||||
;(GlobalRateLimitManager as any).lastApiRequestTime = tenSecondsAgo
|
||||
|
||||
// With 5 second rate limit, no delay needed
|
||||
const delay = await GlobalRateLimitManager.calculateRateLimitDelay(5)
|
||||
expect(delay).toBe(0)
|
||||
})
|
||||
|
||||
test("hasActiveRateLimit returns correct status", async () => {
|
||||
// Initially no rate limit
|
||||
expect(await GlobalRateLimitManager.hasActiveRateLimit()).toBe(false)
|
||||
|
||||
// After setting timestamp
|
||||
await GlobalRateLimitManager.updateLastRequestTime()
|
||||
expect(await GlobalRateLimitManager.hasActiveRateLimit()).toBe(true)
|
||||
|
||||
// After reset
|
||||
GlobalRateLimitManager.reset()
|
||||
expect(await GlobalRateLimitManager.hasActiveRateLimit()).toBe(false)
|
||||
})
|
||||
|
||||
test("concurrent operations maintain consistency", async () => {
|
||||
const promises = []
|
||||
|
||||
// Start multiple concurrent operations
|
||||
for (let i = 0; i < 10; i++) {
|
||||
promises.push(GlobalRateLimitManager.updateLastRequestTime())
|
||||
}
|
||||
|
||||
const timestamps = await Promise.all(promises)
|
||||
|
||||
// All timestamps should be valid and in ascending order
|
||||
for (let i = 1; i < timestamps.length; i++) {
|
||||
expect(timestamps[i]).toBeGreaterThanOrEqual(timestamps[i - 1])
|
||||
}
|
||||
|
||||
// The final timestamp should be the one stored
|
||||
const stored = await GlobalRateLimitManager.getLastRequestTime()
|
||||
expect(stored).toBe(Math.max(...timestamps))
|
||||
})
|
||||
})
|
||||
163
src/core/task/TaskStateLock.ts
Normal file
163
src/core/task/TaskStateLock.ts
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
/**
|
||||
* TaskStateLock - Provides atomic locking mechanisms for critical shared state
|
||||
*
|
||||
* This class prevents race conditions in shared state access during API retry cycles
|
||||
* by implementing a promise-based locking system that ensures sequential access to
|
||||
* critical resources like global rate limiting timestamps.
|
||||
*/
|
||||
|
||||
export class TaskStateLock {
|
||||
private static readonly locks = new Map<string, Promise<void>>()
|
||||
|
||||
/**
|
||||
* Acquire an exclusive lock for the given key
|
||||
* @param lockKey - Unique identifier for the resource being locked
|
||||
* @returns Promise that resolves to a release function
|
||||
*/
|
||||
static async acquire(lockKey: string): Promise<() => void> {
|
||||
// Wait for existing lock to be released
|
||||
while (TaskStateLock.locks.has(lockKey)) {
|
||||
await TaskStateLock.locks.get(lockKey)
|
||||
}
|
||||
|
||||
// Create new lock
|
||||
let releaseLock: () => void
|
||||
const lockPromise = new Promise<void>((resolve) => {
|
||||
releaseLock = resolve
|
||||
})
|
||||
|
||||
TaskStateLock.locks.set(lockKey, lockPromise)
|
||||
|
||||
return () => {
|
||||
TaskStateLock.locks.delete(lockKey)
|
||||
releaseLock!()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Try to acquire a lock without waiting
|
||||
* @param lockKey - Unique identifier for the resource being locked
|
||||
* @returns Release function if lock acquired, null if lock unavailable
|
||||
*/
|
||||
static tryAcquire(lockKey: string): (() => void) | null {
|
||||
if (TaskStateLock.locks.has(lockKey)) {
|
||||
return null // Lock not available
|
||||
}
|
||||
|
||||
let releaseLock: () => void
|
||||
const lockPromise = new Promise<void>((resolve) => {
|
||||
releaseLock = resolve
|
||||
})
|
||||
|
||||
TaskStateLock.locks.set(lockKey, lockPromise)
|
||||
|
||||
return () => {
|
||||
TaskStateLock.locks.delete(lockKey)
|
||||
releaseLock!()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a function with an exclusive lock
|
||||
* @param lockKey - Unique identifier for the resource being locked
|
||||
* @param fn - Function to execute while holding the lock
|
||||
* @returns Promise resolving to the function's return value
|
||||
*/
|
||||
static async withLock<T>(lockKey: string, fn: () => Promise<T> | T): Promise<T> {
|
||||
const release = await TaskStateLock.acquire(lockKey)
|
||||
try {
|
||||
return await fn()
|
||||
} finally {
|
||||
release()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a lock is currently active
|
||||
* @param lockKey - Unique identifier for the resource
|
||||
* @returns True if lock is active, false otherwise
|
||||
*/
|
||||
static isLocked(lockKey: string): boolean {
|
||||
return TaskStateLock.locks.has(lockKey)
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear all locks (for testing purposes)
|
||||
* @internal
|
||||
*/
|
||||
static clearAllLocks(): void {
|
||||
for (const [lockKey, lockPromise] of TaskStateLock.locks) {
|
||||
// Resolve all pending locks to prevent deadlocks
|
||||
lockPromise.then(() => {}).catch(() => {})
|
||||
}
|
||||
TaskStateLock.locks.clear()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* GlobalRateLimitManager - Manages atomic access to global rate limiting state
|
||||
*
|
||||
* Provides thread-safe operations for updating and reading the global API request
|
||||
* timestamp used across all tasks and subtasks for rate limiting.
|
||||
*/
|
||||
export class GlobalRateLimitManager {
|
||||
private static lastApiRequestTime?: number
|
||||
private static readonly LOCK_KEY = "global_rate_limit"
|
||||
|
||||
/**
|
||||
* Atomically update the last request time to the current timestamp
|
||||
* @returns The timestamp that was set
|
||||
*/
|
||||
static async updateLastRequestTime(): Promise<number> {
|
||||
return TaskStateLock.withLock(GlobalRateLimitManager.LOCK_KEY, () => {
|
||||
const now = Date.now()
|
||||
GlobalRateLimitManager.lastApiRequestTime = now
|
||||
return now
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Atomically read the last request time
|
||||
* @returns The last request timestamp, or undefined if never set
|
||||
*/
|
||||
static async getLastRequestTime(): Promise<number | undefined> {
|
||||
return TaskStateLock.withLock(GlobalRateLimitManager.LOCK_KEY, () => {
|
||||
return GlobalRateLimitManager.lastApiRequestTime
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Atomically calculate rate limit delay based on current time and rate limit
|
||||
* @param rateLimitSeconds - Rate limit in seconds
|
||||
* @returns Delay in seconds needed before next request
|
||||
*/
|
||||
static async calculateRateLimitDelay(rateLimitSeconds: number): Promise<number> {
|
||||
return TaskStateLock.withLock(GlobalRateLimitManager.LOCK_KEY, () => {
|
||||
if (!GlobalRateLimitManager.lastApiRequestTime) {
|
||||
return 0
|
||||
}
|
||||
|
||||
const now = Date.now()
|
||||
const timeSinceLastRequest = now - GlobalRateLimitManager.lastApiRequestTime
|
||||
return Math.ceil(Math.max(0, rateLimitSeconds * 1000 - timeSinceLastRequest) / 1000)
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset the global timestamp (for testing purposes)
|
||||
* @internal
|
||||
*/
|
||||
static reset(): void {
|
||||
GlobalRateLimitManager.lastApiRequestTime = undefined
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if rate limiting is active
|
||||
* @returns True if a previous request time exists
|
||||
*/
|
||||
static async hasActiveRateLimit(): Promise<boolean> {
|
||||
return TaskStateLock.withLock(GlobalRateLimitManager.LOCK_KEY, () => {
|
||||
return GlobalRateLimitManager.lastApiRequestTime !== undefined
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue