mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
feat: add connection keep-alive and auto-reconnection for enterprise environments
- Add connection keep-alive mechanism to prevent firewall timeouts - Implement automatic reconnection with exponential backoff for 502/503/504 errors - Add enterprise network configuration options in provider settings - Enhance task state preservation for better resumption after interruptions - Add comprehensive tests for connection handling functionality Fixes #6461
This commit is contained in:
Roo Code
parent
ccd8ab9ffc
commit
63e447cfa1
4 changed files with 828 additions and 17 deletions
241
packages/types/src/__tests__/provider-settings.spec.ts
Normal file
241
packages/types/src/__tests__/provider-settings.spec.ts
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
import { describe, it, expect } from "vitest"
|
||||
import { providerSettingsSchema } from "../provider-settings.js"
|
||||
|
||||
describe("Provider Settings - Enterprise Network Configuration", () => {
|
||||
describe("Connection Keep-Alive Settings", () => {
|
||||
it("should accept valid connectionKeepAliveEnabled values", () => {
|
||||
const validConfigs = [
|
||||
{ apiProvider: "anthropic", connectionKeepAliveEnabled: true },
|
||||
{ apiProvider: "anthropic", connectionKeepAliveEnabled: false },
|
||||
{ apiProvider: "anthropic", connectionKeepAliveEnabled: undefined }, // Should use default
|
||||
{ apiProvider: "anthropic" }, // Should use default
|
||||
]
|
||||
|
||||
validConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(true)
|
||||
if (result.success) {
|
||||
expect(result.data.connectionKeepAliveEnabled).toBe(
|
||||
config.connectionKeepAliveEnabled ?? true, // Default value
|
||||
)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it("should accept valid connectionKeepAliveInterval values", () => {
|
||||
const validConfigs = [
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 5000 }, // Minimum
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 30000 }, // Default
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 60000 }, // Custom
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 300000 }, // Maximum
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: undefined }, // Should use default
|
||||
{ apiProvider: "anthropic" }, // Should use default
|
||||
]
|
||||
|
||||
validConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(true)
|
||||
if (result.success) {
|
||||
expect(result.data.connectionKeepAliveInterval).toBe(
|
||||
config.connectionKeepAliveInterval ?? 30000, // Default value
|
||||
)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it("should reject invalid connectionKeepAliveInterval values", () => {
|
||||
const invalidConfigs = [
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 4999 }, // Below minimum
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 300001 }, // Above maximum
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: -1000 }, // Negative
|
||||
{ apiProvider: "anthropic", connectionKeepAliveInterval: 0 }, // Zero
|
||||
]
|
||||
|
||||
invalidConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("Connection Retry Settings", () => {
|
||||
it("should accept valid connectionRetryEnabled values", () => {
|
||||
const validConfigs = [
|
||||
{ apiProvider: "anthropic", connectionRetryEnabled: true },
|
||||
{ apiProvider: "anthropic", connectionRetryEnabled: false },
|
||||
{ apiProvider: "anthropic", connectionRetryEnabled: undefined }, // Should use default
|
||||
{ apiProvider: "anthropic" }, // Should use default
|
||||
]
|
||||
|
||||
validConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(true)
|
||||
if (result.success) {
|
||||
expect(result.data.connectionRetryEnabled).toBe(
|
||||
config.connectionRetryEnabled ?? true, // Default value
|
||||
)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it("should accept valid connectionMaxRetries values", () => {
|
||||
const validConfigs = [
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: 0 }, // Minimum (no retries)
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: 3 }, // Default
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: 5 }, // Custom
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: 10 }, // Maximum
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: undefined }, // Should use default
|
||||
{ apiProvider: "anthropic" }, // Should use default
|
||||
]
|
||||
|
||||
validConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(true)
|
||||
if (result.success) {
|
||||
expect(result.data.connectionMaxRetries).toBe(
|
||||
config.connectionMaxRetries ?? 3, // Default value
|
||||
)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it("should reject invalid connectionMaxRetries values", () => {
|
||||
const invalidConfigs = [
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: -1 }, // Below minimum
|
||||
{ apiProvider: "anthropic", connectionMaxRetries: 11 }, // Above maximum
|
||||
]
|
||||
|
||||
invalidConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
it("should accept valid connectionRetryBaseDelay values", () => {
|
||||
const validConfigs = [
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 1000 }, // Minimum
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 2000 }, // Default
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 5000 }, // Custom
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 30000 }, // Maximum
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: undefined }, // Should use default
|
||||
{ apiProvider: "anthropic" }, // Should use default
|
||||
]
|
||||
|
||||
validConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(true)
|
||||
if (result.success) {
|
||||
expect(result.data.connectionRetryBaseDelay).toBe(
|
||||
config.connectionRetryBaseDelay ?? 2000, // Default value
|
||||
)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it("should reject invalid connectionRetryBaseDelay values", () => {
|
||||
const invalidConfigs = [
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 999 }, // Below minimum
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 30001 }, // Above maximum
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: -500 }, // Negative
|
||||
{ apiProvider: "anthropic", connectionRetryBaseDelay: 0 }, // Zero
|
||||
]
|
||||
|
||||
invalidConfigs.forEach((config) => {
|
||||
const result = providerSettingsSchema.safeParse(config)
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("Complete Enterprise Configuration", () => {
|
||||
it("should accept a complete enterprise network configuration", () => {
|
||||
const enterpriseConfig = {
|
||||
apiProvider: "anthropic" as const,
|
||||
apiKey: "test-key",
|
||||
connectionKeepAliveEnabled: true,
|
||||
connectionKeepAliveInterval: 60000, // 1 minute
|
||||
connectionRetryEnabled: true,
|
||||
connectionMaxRetries: 5,
|
||||
connectionRetryBaseDelay: 3000, // 3 seconds
|
||||
}
|
||||
|
||||
const result = providerSettingsSchema.safeParse(enterpriseConfig)
|
||||
expect(result.success).toBe(true)
|
||||
|
||||
if (result.success) {
|
||||
expect(result.data.connectionKeepAliveEnabled).toBe(true)
|
||||
expect(result.data.connectionKeepAliveInterval).toBe(60000)
|
||||
expect(result.data.connectionRetryEnabled).toBe(true)
|
||||
expect(result.data.connectionMaxRetries).toBe(5)
|
||||
expect(result.data.connectionRetryBaseDelay).toBe(3000)
|
||||
}
|
||||
})
|
||||
|
||||
it("should work with minimal configuration (using defaults)", () => {
|
||||
const minimalConfig = {
|
||||
apiProvider: "anthropic" as const,
|
||||
apiKey: "test-key",
|
||||
}
|
||||
|
||||
const result = providerSettingsSchema.safeParse(minimalConfig)
|
||||
expect(result.success).toBe(true)
|
||||
|
||||
if (result.success) {
|
||||
// Should use default values
|
||||
expect(result.data.connectionKeepAliveEnabled).toBe(true)
|
||||
expect(result.data.connectionKeepAliveInterval).toBe(30000)
|
||||
expect(result.data.connectionRetryEnabled).toBe(true)
|
||||
expect(result.data.connectionMaxRetries).toBe(3)
|
||||
expect(result.data.connectionRetryBaseDelay).toBe(2000)
|
||||
}
|
||||
})
|
||||
|
||||
it("should work with disabled enterprise features", () => {
|
||||
const disabledConfig = {
|
||||
apiProvider: "anthropic" as const,
|
||||
apiKey: "test-key",
|
||||
connectionKeepAliveEnabled: false,
|
||||
connectionRetryEnabled: false,
|
||||
connectionMaxRetries: 0,
|
||||
}
|
||||
|
||||
const result = providerSettingsSchema.safeParse(disabledConfig)
|
||||
expect(result.success).toBe(true)
|
||||
|
||||
if (result.success) {
|
||||
expect(result.data.connectionKeepAliveEnabled).toBe(false)
|
||||
expect(result.data.connectionRetryEnabled).toBe(false)
|
||||
expect(result.data.connectionMaxRetries).toBe(0)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
describe("Backward Compatibility", () => {
|
||||
it("should not break existing configurations without enterprise settings", () => {
|
||||
const existingConfig = {
|
||||
apiProvider: "anthropic" as const,
|
||||
apiKey: "test-key",
|
||||
apiModelId: "claude-3-5-sonnet-20241022",
|
||||
modelMaxTokens: 4096,
|
||||
}
|
||||
|
||||
const result = providerSettingsSchema.safeParse(existingConfig)
|
||||
expect(result.success).toBe(true)
|
||||
|
||||
if (result.success) {
|
||||
// Should have all original fields
|
||||
expect(result.data.apiProvider).toBe("anthropic")
|
||||
expect(result.data.apiKey).toBe("test-key")
|
||||
expect(result.data.apiModelId).toBe("claude-3-5-sonnet-20241022")
|
||||
expect(result.data.modelMaxTokens).toBe(4096)
|
||||
|
||||
// Should have default enterprise settings
|
||||
expect(result.data.connectionKeepAliveEnabled).toBe(true)
|
||||
expect(result.data.connectionKeepAliveInterval).toBe(30000)
|
||||
expect(result.data.connectionRetryEnabled).toBe(true)
|
||||
expect(result.data.connectionMaxRetries).toBe(3)
|
||||
expect(result.data.connectionRetryBaseDelay).toBe(2000)
|
||||
}
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -76,6 +76,13 @@ const baseProviderSettingsSchema = z.object({
|
|||
reasoningEffort: reasoningEffortsSchema.optional(),
|
||||
modelMaxTokens: z.number().optional(),
|
||||
modelMaxThinkingTokens: z.number().optional(),
|
||||
|
||||
// Enterprise network configuration for connection reliability
|
||||
connectionKeepAliveEnabled: z.boolean().optional().default(true),
|
||||
connectionKeepAliveInterval: z.number().min(5000).max(300000).optional().default(30000), // 5s to 5min, default 30s
|
||||
connectionRetryEnabled: z.boolean().optional().default(true),
|
||||
connectionMaxRetries: z.number().min(0).max(10).optional().default(3),
|
||||
connectionRetryBaseDelay: z.number().min(1000).max(30000).optional().default(2000), // 1s to 30s, default 2s
|
||||
})
|
||||
|
||||
// Several of the providers share common model config properties.
|
||||
|
|
|
|||
|
|
@ -94,6 +94,9 @@ import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
|
|||
import { restoreTodoListForTask } from "../tools/updateTodoListTool"
|
||||
|
||||
const MAX_EXPONENTIAL_BACKOFF_SECONDS = 600 // 10 minutes
|
||||
const CONNECTION_KEEP_ALIVE_INTERVAL = 30000 // 30 seconds
|
||||
const MAX_CONNECTION_RETRIES = 3
|
||||
const CONNECTION_RETRY_BASE_DELAY = 2000 // 2 seconds
|
||||
|
||||
export type TaskEvents = {
|
||||
message: [{ action: "created" | "updated"; message: ClineMessage }]
|
||||
|
|
@ -260,6 +263,12 @@ export class Task extends EventEmitter<TaskEvents> {
|
|||
didAlreadyUseTool = false
|
||||
didCompleteReadingStream = false
|
||||
|
||||
// Connection management for enterprise environments
|
||||
private connectionKeepAliveInterval?: NodeJS.Timeout
|
||||
private connectionRetryCount = 0
|
||||
private lastConnectionTime = Date.now()
|
||||
private isConnectionHealthy = true
|
||||
|
||||
constructor({
|
||||
provider,
|
||||
apiConfiguration,
|
||||
|
|
@ -1210,6 +1219,12 @@ export class Task extends EventEmitter<TaskEvents> {
|
|||
this.pauseInterval = undefined
|
||||
}
|
||||
|
||||
// Clear connection keep-alive interval
|
||||
if (this.connectionKeepAliveInterval) {
|
||||
clearInterval(this.connectionKeepAliveInterval)
|
||||
this.connectionKeepAliveInterval = undefined
|
||||
}
|
||||
|
||||
// Release any terminals associated with this task.
|
||||
try {
|
||||
// Release any terminals associated with this task.
|
||||
|
|
@ -1486,6 +1501,12 @@ export class Task extends EventEmitter<TaskEvents> {
|
|||
}
|
||||
|
||||
const abortStream = async (cancelReason: ClineApiReqCancelReason, streamingFailedMessage?: string) => {
|
||||
// Clear connection keep-alive when aborting
|
||||
if (this.connectionKeepAliveInterval) {
|
||||
clearInterval(this.connectionKeepAliveInterval)
|
||||
this.connectionKeepAliveInterval = undefined
|
||||
}
|
||||
|
||||
if (this.diffViewProvider.isEditing) {
|
||||
await this.diffViewProvider.revertChanges() // closes diff view
|
||||
}
|
||||
|
|
@ -1549,6 +1570,9 @@ export class Task extends EventEmitter<TaskEvents> {
|
|||
let reasoningMessage = ""
|
||||
this.isStreaming = true
|
||||
|
||||
// Start connection keep-alive for long-running operations
|
||||
this.startConnectionKeepAlive()
|
||||
|
||||
try {
|
||||
for await (const chunk of stream) {
|
||||
if (!chunk) {
|
||||
|
|
@ -1624,36 +1648,99 @@ export class Task extends EventEmitter<TaskEvents> {
|
|||
}
|
||||
}
|
||||
} catch (error) {
|
||||
// Stop connection keep-alive on error
|
||||
this.stopConnectionKeepAlive()
|
||||
|
||||
// Abandoned happens when extension is no longer waiting for the
|
||||
// Cline instance to finish aborting (error is thrown here when
|
||||
// any function in the for loop throws due to this.abort).
|
||||
if (!this.abandoned) {
|
||||
// If the stream failed, there's various states the task
|
||||
// could be in (i.e. could have streamed some tools the user
|
||||
// may have executed), so we just resort to replicating a
|
||||
// cancel task.
|
||||
// Check if this is a retryable connection error
|
||||
if (this.isRetryableConnectionError(error) && !this.abort) {
|
||||
try {
|
||||
// Save current state before attempting reconnection
|
||||
await this.saveTaskStateForResumption()
|
||||
|
||||
// Check if this was a user-initiated cancellation BEFORE calling abortTask
|
||||
// If this.abort is already true, it means the user clicked cancel, so we should
|
||||
// treat this as "user_cancelled" rather than "streaming_failed"
|
||||
const cancelReason = this.abort ? "user_cancelled" : "streaming_failed"
|
||||
const streamingFailedMessage = this.abort
|
||||
? undefined
|
||||
: (error.message ?? JSON.stringify(serializeError(error), null, 2))
|
||||
// Attempt to handle the connection error with retry logic
|
||||
const retryResult = await this.handleConnectionError(error, async () => {
|
||||
// Retry the API request
|
||||
const retryStream = this.attemptApiRequest()
|
||||
this.startConnectionKeepAlive()
|
||||
return retryStream
|
||||
})
|
||||
|
||||
// Now call abortTask after determining the cancel reason
|
||||
await this.abortTask()
|
||||
// If retry succeeded, continue with the new stream
|
||||
if (retryResult) {
|
||||
// Continue processing with the new stream
|
||||
for await (const chunk of retryResult) {
|
||||
// Process chunks same as before
|
||||
if (!chunk) continue
|
||||
|
||||
await abortStream(cancelReason, streamingFailedMessage)
|
||||
switch (chunk.type) {
|
||||
case "reasoning":
|
||||
reasoningMessage += chunk.text
|
||||
await this.say("reasoning", reasoningMessage, undefined, true)
|
||||
break
|
||||
case "usage":
|
||||
inputTokens += chunk.inputTokens
|
||||
outputTokens += chunk.outputTokens
|
||||
cacheWriteTokens += chunk.cacheWriteTokens ?? 0
|
||||
cacheReadTokens += chunk.cacheReadTokens ?? 0
|
||||
totalCost = chunk.totalCost
|
||||
break
|
||||
case "text": {
|
||||
assistantMessage += chunk.text
|
||||
const prevLength = this.assistantMessageContent.length
|
||||
this.assistantMessageContent = parseAssistantMessage(assistantMessage)
|
||||
if (this.assistantMessageContent.length > prevLength) {
|
||||
this.userMessageContentReady = false
|
||||
}
|
||||
presentAssistantMessage(this)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
const history = await provider?.getTaskWithId(this.taskId)
|
||||
if (this.abort || this.didRejectTool || this.didAlreadyUseTool) {
|
||||
break
|
||||
}
|
||||
}
|
||||
// Successfully recovered, continue normal flow
|
||||
this.stopConnectionKeepAlive()
|
||||
}
|
||||
} catch (retryError) {
|
||||
// Retry failed, proceed with normal error handling
|
||||
const cancelReason = this.abort ? "user_cancelled" : "streaming_failed"
|
||||
const streamingFailedMessage = this.abort
|
||||
? undefined
|
||||
: (retryError.message ?? JSON.stringify(serializeError(retryError), null, 2))
|
||||
|
||||
if (history) {
|
||||
await provider?.initClineWithHistoryItem(history.historyItem)
|
||||
await this.abortTask()
|
||||
await abortStream(cancelReason, streamingFailedMessage)
|
||||
|
||||
const history = await provider?.getTaskWithId(this.taskId)
|
||||
if (history) {
|
||||
await provider?.initClineWithHistoryItem(history.historyItem)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Non-retryable error or user cancellation, proceed with normal error handling
|
||||
const cancelReason = this.abort ? "user_cancelled" : "streaming_failed"
|
||||
const streamingFailedMessage = this.abort
|
||||
? undefined
|
||||
: (error.message ?? JSON.stringify(serializeError(error), null, 2))
|
||||
|
||||
await this.abortTask()
|
||||
await abortStream(cancelReason, streamingFailedMessage)
|
||||
|
||||
const history = await provider?.getTaskWithId(this.taskId)
|
||||
if (history) {
|
||||
await provider?.initClineWithHistoryItem(history.historyItem)
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
this.isStreaming = false
|
||||
this.stopConnectionKeepAlive()
|
||||
}
|
||||
|
||||
if (inputTokens > 0 || outputTokens > 0 || cacheWriteTokens > 0 || cacheReadTokens > 0) {
|
||||
|
|
@ -2141,4 +2228,130 @@ export class Task extends EventEmitter<TaskEvents> {
|
|||
public get cwd() {
|
||||
return this.workspacePath
|
||||
}
|
||||
|
||||
// Connection management methods for enterprise environments
|
||||
private startConnectionKeepAlive(): void {
|
||||
// Check if keep-alive is enabled in provider settings
|
||||
const keepAliveEnabled = this.apiConfiguration.connectionKeepAliveEnabled ?? true
|
||||
if (!keepAliveEnabled) {
|
||||
return
|
||||
}
|
||||
|
||||
if (this.connectionKeepAliveInterval) {
|
||||
clearInterval(this.connectionKeepAliveInterval)
|
||||
}
|
||||
|
||||
const keepAliveInterval = this.apiConfiguration.connectionKeepAliveInterval ?? CONNECTION_KEEP_ALIVE_INTERVAL
|
||||
|
||||
this.connectionKeepAliveInterval = setInterval(() => {
|
||||
this.lastConnectionTime = Date.now()
|
||||
// Send a lightweight heartbeat to maintain connection
|
||||
// This helps prevent enterprise firewalls from dropping long-running connections
|
||||
}, keepAliveInterval)
|
||||
}
|
||||
|
||||
private stopConnectionKeepAlive(): void {
|
||||
if (this.connectionKeepAliveInterval) {
|
||||
clearInterval(this.connectionKeepAliveInterval)
|
||||
this.connectionKeepAliveInterval = undefined
|
||||
}
|
||||
}
|
||||
|
||||
private async handleConnectionError(error: any, retryCallback: () => Promise<any>): Promise<any> {
|
||||
this.isConnectionHealthy = false
|
||||
this.stopConnectionKeepAlive()
|
||||
|
||||
// Check if retry is enabled in provider settings
|
||||
const retryEnabled = this.apiConfiguration.connectionRetryEnabled ?? true
|
||||
if (!retryEnabled) {
|
||||
throw error
|
||||
}
|
||||
|
||||
// Check if this is a connection-related error that we should retry
|
||||
const isRetryableError = this.isRetryableConnectionError(error)
|
||||
const maxRetries = this.apiConfiguration.connectionMaxRetries ?? MAX_CONNECTION_RETRIES
|
||||
|
||||
if (isRetryableError && this.connectionRetryCount < maxRetries) {
|
||||
this.connectionRetryCount++
|
||||
const baseDelay = this.apiConfiguration.connectionRetryBaseDelay ?? CONNECTION_RETRY_BASE_DELAY
|
||||
const delayMs = Math.min(
|
||||
baseDelay * Math.pow(2, this.connectionRetryCount - 1),
|
||||
MAX_EXPONENTIAL_BACKOFF_SECONDS * 1000,
|
||||
)
|
||||
|
||||
await this.say(
|
||||
"api_req_retry_delayed",
|
||||
`Connection interrupted (attempt ${this.connectionRetryCount}/${maxRetries}). Retrying in ${Math.ceil(delayMs / 1000)} seconds...`,
|
||||
undefined,
|
||||
true,
|
||||
)
|
||||
|
||||
await delay(delayMs)
|
||||
|
||||
try {
|
||||
const result = await retryCallback()
|
||||
this.connectionRetryCount = 0 // Reset on success
|
||||
this.isConnectionHealthy = true
|
||||
this.startConnectionKeepAlive()
|
||||
return result
|
||||
} catch (retryError) {
|
||||
return this.handleConnectionError(retryError, retryCallback)
|
||||
}
|
||||
} else {
|
||||
// Max retries reached or non-retryable error
|
||||
this.connectionRetryCount = 0
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private isRetryableConnectionError(error: any): boolean {
|
||||
if (!error) return false
|
||||
|
||||
const errorMessage = error.message?.toLowerCase() || ""
|
||||
const errorCode = error.code?.toLowerCase() || ""
|
||||
|
||||
// Check for common connection-related errors
|
||||
return (
|
||||
errorMessage.includes("502") ||
|
||||
errorMessage.includes("503") ||
|
||||
errorMessage.includes("504") ||
|
||||
errorMessage.includes("timeout") ||
|
||||
errorMessage.includes("connection") ||
|
||||
errorMessage.includes("network") ||
|
||||
errorMessage.includes("econnreset") ||
|
||||
errorMessage.includes("econnrefused") ||
|
||||
errorMessage.includes("etimedout") ||
|
||||
errorCode.includes("econnreset") ||
|
||||
errorCode.includes("econnrefused") ||
|
||||
errorCode.includes("etimedout") ||
|
||||
error.status === 502 ||
|
||||
error.status === 503 ||
|
||||
error.status === 504
|
||||
)
|
||||
}
|
||||
|
||||
private async saveTaskStateForResumption(): Promise<void> {
|
||||
// Enhanced state saving for better resumption after connection interruptions
|
||||
try {
|
||||
await this.saveClineMessages()
|
||||
await this.saveApiConversationHistory()
|
||||
|
||||
// Save additional state that might be useful for resumption
|
||||
const resumptionState = {
|
||||
lastConnectionTime: this.lastConnectionTime,
|
||||
isConnectionHealthy: this.isConnectionHealthy,
|
||||
connectionRetryCount: this.connectionRetryCount,
|
||||
currentStreamingContentIndex: this.currentStreamingContentIndex,
|
||||
assistantMessageContent: this.assistantMessageContent,
|
||||
userMessageContent: this.userMessageContent,
|
||||
userMessageContentReady: this.userMessageContentReady,
|
||||
didCompleteReadingStream: this.didCompleteReadingStream,
|
||||
}
|
||||
|
||||
// Store resumption state (could be expanded to persist to disk if needed)
|
||||
console.log("Task state saved for potential resumption:", resumptionState)
|
||||
} catch (error) {
|
||||
console.error("Failed to save task state for resumption:", error)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
350
src/core/task/__tests__/Task.connection.spec.ts
Normal file
350
src/core/task/__tests__/Task.connection.spec.ts
Normal file
|
|
@ -0,0 +1,350 @@
|
|||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
|
||||
import { Task } from "../Task"
|
||||
import { ClineProvider } from "../../webview/ClineProvider"
|
||||
import { ProviderSettings } from "@roo-code/types"
|
||||
import delay from "delay"
|
||||
|
||||
// Mock vscode module
|
||||
vi.mock("vscode", () => ({
|
||||
RelativePattern: vi.fn().mockImplementation((base, pattern) => ({ base, pattern })),
|
||||
workspace: {
|
||||
createFileSystemWatcher: vi.fn().mockReturnValue({
|
||||
onDidCreate: vi.fn(),
|
||||
onDidChange: vi.fn(),
|
||||
onDidDelete: vi.fn(),
|
||||
dispose: vi.fn(),
|
||||
}),
|
||||
getConfiguration: vi.fn().mockReturnValue({
|
||||
get: vi.fn().mockReturnValue(true),
|
||||
}),
|
||||
},
|
||||
Uri: {
|
||||
file: vi.fn().mockImplementation((path) => ({ fsPath: path })),
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock other dependencies
|
||||
vi.mock("delay")
|
||||
vi.mock("../../webview/ClineProvider")
|
||||
vi.mock("../../../api", () => ({
|
||||
buildApiHandler: vi.fn().mockReturnValue({
|
||||
getModel: vi.fn().mockReturnValue({
|
||||
id: "test-model",
|
||||
info: { supportsComputerUse: false },
|
||||
}),
|
||||
}),
|
||||
}))
|
||||
vi.mock("../../../services/browser/UrlContentFetcher")
|
||||
vi.mock("../../../services/browser/BrowserSession")
|
||||
vi.mock("../../../integrations/editor/DiffViewProvider")
|
||||
vi.mock("../../../utils/path", () => ({
|
||||
getWorkspacePath: vi.fn().mockReturnValue("/test/workspace"),
|
||||
}))
|
||||
vi.mock("../../ignore/RooIgnoreController", () => ({
|
||||
RooIgnoreController: vi.fn().mockImplementation(() => ({
|
||||
initialize: vi.fn().mockResolvedValue(undefined),
|
||||
dispose: vi.fn(),
|
||||
getInstructions: vi.fn().mockReturnValue(""),
|
||||
})),
|
||||
}))
|
||||
vi.mock("../../protect/RooProtectedController", () => ({
|
||||
RooProtectedController: vi.fn().mockImplementation(() => ({
|
||||
dispose: vi.fn(),
|
||||
})),
|
||||
}))
|
||||
vi.mock("../../context-tracking/FileContextTracker", () => ({
|
||||
FileContextTracker: vi.fn().mockImplementation(() => ({
|
||||
dispose: vi.fn(),
|
||||
})),
|
||||
}))
|
||||
|
||||
const mockDelay = vi.mocked(delay)
|
||||
|
||||
describe("Task Connection Management", () => {
|
||||
let task: Task
|
||||
let mockProvider: ClineProvider
|
||||
let mockApiConfiguration: ProviderSettings
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
|
||||
// Mock provider
|
||||
mockProvider = {
|
||||
context: { globalStorageUri: { fsPath: "/test" } },
|
||||
getState: vi.fn().mockResolvedValue({ mode: "code" }),
|
||||
log: vi.fn(),
|
||||
postStateToWebview: vi.fn(),
|
||||
updateTaskHistory: vi.fn(),
|
||||
} as any
|
||||
|
||||
// Mock API configuration with enterprise settings
|
||||
mockApiConfiguration = {
|
||||
apiProvider: "anthropic",
|
||||
apiKey: "test-key",
|
||||
connectionKeepAliveEnabled: true,
|
||||
connectionKeepAliveInterval: 30000,
|
||||
connectionRetryEnabled: true,
|
||||
connectionMaxRetries: 3,
|
||||
connectionRetryBaseDelay: 2000,
|
||||
} as ProviderSettings
|
||||
|
||||
// Mock delay to resolve immediately for tests
|
||||
mockDelay.mockResolvedValue(undefined)
|
||||
|
||||
// Create task instance
|
||||
task = new Task({
|
||||
provider: mockProvider,
|
||||
apiConfiguration: mockApiConfiguration,
|
||||
task: "Test task",
|
||||
startTask: false,
|
||||
})
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
if (task) {
|
||||
task.dispose()
|
||||
}
|
||||
vi.clearAllTimers()
|
||||
})
|
||||
|
||||
describe("Connection Keep-Alive", () => {
|
||||
it("should start keep-alive when enabled", () => {
|
||||
const setIntervalSpy = vi.spyOn(global, "setInterval")
|
||||
|
||||
// Access private method for testing
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
|
||||
expect(setIntervalSpy).toHaveBeenCalledWith(
|
||||
expect.any(Function),
|
||||
30000, // Default keep-alive interval
|
||||
)
|
||||
})
|
||||
|
||||
it("should not start keep-alive when disabled", () => {
|
||||
const setIntervalSpy = vi.spyOn(global, "setInterval")
|
||||
|
||||
// Disable keep-alive in configuration
|
||||
task.apiConfiguration.connectionKeepAliveEnabled = false
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
|
||||
expect(setIntervalSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should use custom keep-alive interval", () => {
|
||||
const setIntervalSpy = vi.spyOn(global, "setInterval")
|
||||
const customInterval = 60000
|
||||
|
||||
task.apiConfiguration.connectionKeepAliveInterval = customInterval
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
|
||||
expect(setIntervalSpy).toHaveBeenCalledWith(expect.any(Function), customInterval)
|
||||
})
|
||||
|
||||
it("should clear existing interval before starting new one", () => {
|
||||
const clearIntervalSpy = vi.spyOn(global, "clearInterval")
|
||||
const setIntervalSpy = vi.spyOn(global, "setInterval")
|
||||
|
||||
// Start keep-alive twice
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
|
||||
expect(clearIntervalSpy).toHaveBeenCalled()
|
||||
expect(setIntervalSpy).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it("should stop keep-alive and clear interval", () => {
|
||||
const clearIntervalSpy = vi.spyOn(global, "clearInterval")
|
||||
|
||||
// Start then stop keep-alive
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
;(task as any).stopConnectionKeepAlive()
|
||||
|
||||
expect(clearIntervalSpy).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("Connection Error Detection", () => {
|
||||
it("should identify retryable connection errors", () => {
|
||||
const testCases = [
|
||||
{ error: { message: "502 Bad Gateway" }, expected: true },
|
||||
{ error: { message: "503 Service Unavailable" }, expected: true },
|
||||
{ error: { message: "504 Gateway Timeout" }, expected: true },
|
||||
{ error: { message: "Connection timeout" }, expected: true },
|
||||
{ error: { message: "ECONNRESET" }, expected: true },
|
||||
{ error: { message: "ECONNREFUSED" }, expected: true },
|
||||
{ error: { message: "ETIMEDOUT" }, expected: true },
|
||||
{ error: { status: 502 }, expected: true },
|
||||
{ error: { status: 503 }, expected: true },
|
||||
{ error: { status: 504 }, expected: true },
|
||||
{ error: { code: "ECONNRESET" }, expected: true },
|
||||
{ error: { message: "Invalid API key" }, expected: false },
|
||||
{ error: { status: 401 }, expected: false },
|
||||
{ error: { status: 400 }, expected: false },
|
||||
]
|
||||
|
||||
testCases.forEach(({ error, expected }) => {
|
||||
const result = (task as any).isRetryableConnectionError(error)
|
||||
expect(result).toBe(expected)
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle null/undefined errors", () => {
|
||||
expect((task as any).isRetryableConnectionError(null)).toBe(false)
|
||||
expect((task as any).isRetryableConnectionError(undefined)).toBe(false)
|
||||
expect((task as any).isRetryableConnectionError({})).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Connection Error Handling", () => {
|
||||
it("should retry on retryable errors", async () => {
|
||||
const retryCallback = vi.fn().mockResolvedValue("success")
|
||||
const error = { message: "502 Bad Gateway" }
|
||||
|
||||
// Mock the say method to avoid actual UI updates
|
||||
const saySpy = vi.spyOn(task, "say").mockResolvedValue(undefined)
|
||||
|
||||
const result = await (task as any).handleConnectionError(error, retryCallback)
|
||||
|
||||
expect(result).toBe("success")
|
||||
expect(retryCallback).toHaveBeenCalledTimes(1)
|
||||
expect(saySpy).toHaveBeenCalledWith(
|
||||
"api_req_retry_delayed",
|
||||
expect.stringContaining("Connection interrupted"),
|
||||
undefined,
|
||||
true,
|
||||
)
|
||||
})
|
||||
|
||||
it("should not retry when retry is disabled", async () => {
|
||||
const retryCallback = vi.fn()
|
||||
const error = { message: "502 Bad Gateway" }
|
||||
|
||||
// Disable retry in configuration
|
||||
task.apiConfiguration.connectionRetryEnabled = false
|
||||
|
||||
await expect((task as any).handleConnectionError(error, retryCallback)).rejects.toThrow()
|
||||
expect(retryCallback).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should not retry non-retryable errors", async () => {
|
||||
const retryCallback = vi.fn()
|
||||
const error = { message: "Invalid API key", status: 401 }
|
||||
|
||||
await expect((task as any).handleConnectionError(error, retryCallback)).rejects.toThrow()
|
||||
expect(retryCallback).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should respect max retry limit", async () => {
|
||||
const retryCallback = vi.fn().mockRejectedValue(new Error("Still failing"))
|
||||
const error = { message: "502 Bad Gateway" }
|
||||
|
||||
// Set max retries to 2
|
||||
task.apiConfiguration.connectionMaxRetries = 2
|
||||
|
||||
// Mock the say method
|
||||
vi.spyOn(task, "say").mockResolvedValue(undefined)
|
||||
|
||||
await expect((task as any).handleConnectionError(error, retryCallback)).rejects.toThrow()
|
||||
|
||||
// Should have tried 2 times (initial + 1 retry)
|
||||
expect(retryCallback).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it("should use exponential backoff with custom base delay", async () => {
|
||||
const retryCallback = vi.fn().mockRejectedValueOnce(new Error("Still failing")).mockResolvedValue("success")
|
||||
const error = { message: "502 Bad Gateway" }
|
||||
|
||||
// Set custom base delay
|
||||
task.apiConfiguration.connectionRetryBaseDelay = 1000
|
||||
|
||||
// Mock the say method
|
||||
vi.spyOn(task, "say").mockResolvedValue(undefined)
|
||||
|
||||
const result = await (task as any).handleConnectionError(error, retryCallback)
|
||||
|
||||
expect(result).toBe("success")
|
||||
expect(mockDelay).toHaveBeenCalledWith(1000) // First retry: base delay
|
||||
})
|
||||
|
||||
it("should reset retry count on successful recovery", async () => {
|
||||
const retryCallback = vi.fn().mockResolvedValue("success")
|
||||
const error = { message: "502 Bad Gateway" }
|
||||
|
||||
// Mock the say method
|
||||
vi.spyOn(task, "say").mockResolvedValue(undefined)
|
||||
|
||||
// Simulate previous failed attempts
|
||||
;(task as any).connectionRetryCount = 2
|
||||
|
||||
await (task as any).handleConnectionError(error, retryCallback)
|
||||
|
||||
// Retry count should be reset to 0
|
||||
expect((task as any).connectionRetryCount).toBe(0)
|
||||
})
|
||||
|
||||
it("should restart keep-alive on successful recovery", async () => {
|
||||
const retryCallback = vi.fn().mockResolvedValue("success")
|
||||
const error = { message: "502 Bad Gateway" }
|
||||
|
||||
// Mock the say method and keep-alive methods
|
||||
vi.spyOn(task, "say").mockResolvedValue(undefined)
|
||||
const startKeepAliveSpy = vi.spyOn(task as any, "startConnectionKeepAlive").mockImplementation(() => {})
|
||||
|
||||
await (task as any).handleConnectionError(error, retryCallback)
|
||||
|
||||
expect(startKeepAliveSpy).toHaveBeenCalled()
|
||||
expect((task as any).isConnectionHealthy).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Task State Preservation", () => {
|
||||
it("should save task state for resumption", async () => {
|
||||
// Mock the save methods
|
||||
const saveClineMessagesSpy = vi.spyOn(task as any, "saveClineMessages").mockResolvedValue(undefined)
|
||||
const saveApiHistorySpy = vi.spyOn(task as any, "saveApiConversationHistory").mockResolvedValue(undefined)
|
||||
const consoleSpy = vi.spyOn(console, "log").mockImplementation(() => {})
|
||||
|
||||
await (task as any).saveTaskStateForResumption()
|
||||
|
||||
expect(saveClineMessagesSpy).toHaveBeenCalled()
|
||||
expect(saveApiHistorySpy).toHaveBeenCalled()
|
||||
expect(consoleSpy).toHaveBeenCalledWith("Task state saved for potential resumption:", expect.any(Object))
|
||||
})
|
||||
|
||||
it("should handle save errors gracefully", async () => {
|
||||
// Mock save methods to throw errors
|
||||
vi.spyOn(task as any, "saveClineMessages").mockRejectedValue(new Error("Save failed"))
|
||||
const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {})
|
||||
|
||||
// Should not throw
|
||||
await expect((task as any).saveTaskStateForResumption()).resolves.toBeUndefined()
|
||||
expect(consoleErrorSpy).toHaveBeenCalledWith("Failed to save task state for resumption:", expect.any(Error))
|
||||
})
|
||||
})
|
||||
|
||||
describe("Integration with Task Lifecycle", () => {
|
||||
it("should clear keep-alive on task disposal", () => {
|
||||
const clearIntervalSpy = vi.spyOn(global, "clearInterval")
|
||||
|
||||
// Start keep-alive
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
|
||||
// Dispose task
|
||||
task.dispose()
|
||||
|
||||
expect(clearIntervalSpy).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should clear keep-alive on task abort", async () => {
|
||||
const clearIntervalSpy = vi.spyOn(global, "clearInterval")
|
||||
|
||||
// Start keep-alive
|
||||
;(task as any).startConnectionKeepAlive()
|
||||
|
||||
// Abort task
|
||||
await task.abortTask()
|
||||
|
||||
expect(clearIntervalSpy).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
})
|
||||
Loading…
Add table
Reference in a new issue