diff --git a/apps/vscode-e2e/src/suite/mid-stream-retry.test.ts b/apps/vscode-e2e/src/suite/mid-stream-retry.test.ts new file mode 100644 index 0000000000..8ba6e03fdf --- /dev/null +++ b/apps/vscode-e2e/src/suite/mid-stream-retry.test.ts @@ -0,0 +1,106 @@ +import { providerIdentifiers, RooCodeEventName, type ClineMessage } from "@roo-code/types" +import * as assert from "assert" + +import { setDefaultSuiteTimeout } from "./test-utils" +import { isCompletedAsk, waitFor } from "./utils" + +const PROBE = "mid-stream-retry-e2e:" + +function installMidStreamFailureInterceptor(requests: { count: number }): () => void { + const originalFetch = globalThis.fetch + + globalThis.fetch = async function (input: RequestInfo | URL, init?: RequestInit): Promise { + const body = typeof init?.body === "string" ? init.body : "" + if (body.includes(PROBE)) { + requests.count++ + return makePartialFailureResponse() + } + return originalFetch.call(globalThis, input, init as RequestInit) + } as typeof globalThis.fetch + + return () => { + globalThis.fetch = originalFetch + } +} + +function makePartialFailureResponse(): Response { + const encoder = new TextEncoder() + let emitted = false + const stream = new ReadableStream({ + pull(controller) { + if (!emitted) { + emitted = true + controller.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + id: "mid-stream-failure", + object: "chat.completion.chunk", + model: "openai/gpt-4.1", + choices: [ + { index: 0, delta: { role: "assistant", content: "partial" }, finish_reason: null }, + ], + })}\n\n`, + ), + ) + return + } + controller.error(new Error("mid-stream provider failure")) + }, + }) + + return new Response(stream, { status: 200, headers: { "content-type": "text/event-stream" } }) +} + +suite("Mid-stream retry", function () { + setDefaultSuiteTimeout(this) + + let restoreFetch: (() => void) | undefined + const requests = { count: 0 } + + suiteSetup(async () => { + restoreFetch = installMidStreamFailureInterceptor(requests) + const aimockUrl = process.env.AIMOCK_URL + await globalThis.api.setConfiguration({ + apiProvider: providerIdentifiers.openrouter, + openRouterApiKey: "mock-key", + openRouterModelId: "openai/gpt-4.1", + requestDelaySeconds: 1, + ...(aimockUrl && { openRouterBaseUrl: `${aimockUrl}/v1` }), + }) + }) + + suiteTeardown(async () => { + restoreFetch?.() + restoreFetch = undefined + await globalThis.api.clearCurrentTask() + }) + + test("bounds partial-stream retries and surfaces the failure prompt", async () => { + requests.count = 0 + const messages: ClineMessage[] = [] + const handler = ({ message }: { message: ClineMessage }) => messages.push(message) + globalThis.api.on(RooCodeEventName.Message, handler) + + try { + await globalThis.api.startNewTask({ + configuration: { mode: "ask", autoApprovalEnabled: false }, + text: `${PROBE} fail after a partial provider response`, + }) + + await waitFor( + () => messages.some((message) => isCompletedAsk(message) && message.ask === "api_req_failed"), + { timeout: 60_000 }, + ) + assert.strictEqual(requests.count, 4, "Should make one initial request and three automatic retries") + assert.ok( + messages.some((message) => message.type === "say" && message.say === "api_req_retry_delayed"), + "Should expose automatic retry backoff to the user", + ) + + await new Promise((resolve) => setTimeout(resolve, 250)) + assert.strictEqual(requests.count, 4, "Waiting at the retry prompt must not issue another billed request") + } finally { + globalThis.api.off(RooCodeEventName.Message, handler) + } + }) +}) diff --git a/docs/architecture/task-lifecycle-model.md b/docs/architecture/task-lifecycle-model.md index 9266d49987..dff3aacd6f 100644 --- a/docs/architecture/task-lifecycle-model.md +++ b/docs/architecture/task-lifecycle-model.md @@ -6,14 +6,15 @@ Zoo Code checks task lifecycle protocols through one compositional verification pnpm lifecycle:model-check ``` -The command runs six independent bounded submodels in sequence: +The command runs seven independent bounded submodels in sequence: 1. the persisted task delegation lifecycle; 2. shared-store concurrency across task-history hosts; 3. production-backed provider handoff and scheduler ordering; 4. the task cleanup protocol; 5. request-stream parser scoping; and -6. completion persistence. +6. completion persistence; and +7. mid-stream provider retry budgeting and user handoff. This umbrella command is the single model-check entry point in the `compile` CI job after type checking. Command-level composition does not merge the submodels' state spaces: each checker retains its own bounds, transitions, invariant ownership, reachability requirements, and counterexample format. In particular, parser state is not part of the persisted lifecycle graph. The focused parser checker remains directly runnable with `pnpm parser-scope:model-check` for debugging. @@ -104,6 +105,12 @@ The model abstracts restart visibility as the `durable` history phase. It allows Seven semantic landmarks keep the intended positive and negative paths reachable: delayed completion remains pending, failed completion remains pending, exhausted retries settle without completion, cancellation can win after retry delay but before persistence, delegated reopen failure emits no delegated completion, and both standalone and delegated tasks can complete after durable history. The checker explores all reachable states through depth 10 and fails rather than reporting a truncated pass if an unseen successor remains. +## Mid-stream provider retry model + +`scripts/check-mid-stream-retry.ts` models the independent request-level protocol used when a provider fails after yielding at least one stream chunk. It imports the production `decideMidStreamFailure` decision, exhaustively explores success, failure, backoff cancellation, prompt cancellation, user approval, and user decline through one bounded user-approved retry round, and requires every automatic retry to have one visible announcement. Its invariants prevent automatic requests beyond the configured budget, require the user prompt exactly at exhaustion, require approval to reset the budget, and make decline or either cancellation path terminal. + +The model does not claim provider transport liveness or token-billing accuracy. Focused `Task` tests cover conversation-history bookkeeping and both prompt responses; the VS Code E2E suite injects a real partial SSE response followed by a transport failure to verify the extension-host boundary, request count, visible retry handoff, and stable waiting at the failure prompt. + ## Invariants The task delegation checker currently enforces: diff --git a/package.json b/package.json index 1fd9ddc8fe..b5032b4cfb 100644 --- a/package.json +++ b/package.json @@ -13,7 +13,7 @@ "check-types": "turbo check-types --log-order grouped --output-logs new-only", "test": "turbo test --log-order grouped --output-logs new-only", "test:mutation-ci": "node --test scripts/stryker-diff.test.mjs", - "lifecycle:model-check": "tsx scripts/check-task-lifecycle.ts && tsx scripts/check-task-store-concurrency.ts && tsx scripts/check-provider-handoff-scheduler.ts && pnpm cleanup-protocol:model-check && pnpm parser-scope:model-check && tsx scripts/check-completion-persistence.ts", + "lifecycle:model-check": "tsx scripts/check-task-lifecycle.ts && tsx scripts/check-task-store-concurrency.ts && tsx scripts/check-provider-handoff-scheduler.ts && pnpm cleanup-protocol:model-check && pnpm parser-scope:model-check && tsx scripts/check-completion-persistence.ts && tsx scripts/check-mid-stream-retry.ts", "cleanup-protocol:model-check": "tsx scripts/check-task-cleanup-protocol.ts", "parser-scope:model-check": "node scripts/run-native-tool-call-parser-scoping.mjs", "mcp:integration-check": "tsx scripts/check-mcp-oauth-integration.ts", diff --git a/scripts/check-mid-stream-retry.ts b/scripts/check-mid-stream-retry.ts new file mode 100644 index 0000000000..a24ed2f93b --- /dev/null +++ b/scripts/check-mid-stream-retry.ts @@ -0,0 +1,162 @@ +import assert from "node:assert/strict" + +import { decideMidStreamFailure, MAX_MID_STREAM_RETRIES } from "../src/core/task/midStreamRetry" + +type Phase = "requesting" | "backoff" | "awaiting-user" | "stopped" | "succeeded" + +interface ModelState { + phase: Phase + stopReason?: "decline" | "backoff-abort" | "prompt-abort" + retryAttempt: number + requests: number + announcements: number + roundsApproved: number +} + +interface Transition { + name: string + next: ModelState +} + +interface TraceStep { + action: string + state: ModelState +} + +const MAX_APPROVED_ROUNDS = 1 +const MAX_DEPTH = 16 +const MAX_STATES = 100 +const expectedActions = ["fail", "retry", "approve", "decline", "abort-backoff", "abort-prompt", "succeed"] as const +const landmarks = { + "automatic-budget-exhausted": (state: ModelState) => + state.phase === "awaiting-user" && state.requests === MAX_MID_STREAM_RETRIES + 1, + "approved-round-reset": (state: ModelState) => + state.roundsApproved === 1 && state.phase === "requesting" && state.retryAttempt === 0, + "declined-after-approved-round": (state: ModelState) => + state.roundsApproved === 1 && state.stopReason === "decline", + "backoff-cancelled": (state: ModelState) => state.stopReason === "backoff-abort", + "prompt-cancelled": (state: ModelState) => state.stopReason === "prompt-abort", +} satisfies Record boolean> + +function initialState(): ModelState { + return { phase: "requesting", retryAttempt: 0, requests: 1, announcements: 0, roundsApproved: 0 } +} + +function transitions(state: ModelState): Transition[] { + if (state.phase === "requesting") { + const decision = decideMidStreamFailure(state.retryAttempt) + return [ + { name: "fail", next: { ...state, phase: decision === "retry" ? "backoff" : "awaiting-user" } }, + { name: "succeed", next: { ...state, phase: "succeeded" } }, + ] + } + if (state.phase === "backoff") { + return [ + { + name: "retry", + next: { + ...state, + phase: "requesting", + retryAttempt: state.retryAttempt + 1, + requests: state.requests + 1, + announcements: state.announcements + 1, + }, + }, + { name: "abort-backoff", next: { ...state, phase: "stopped", stopReason: "backoff-abort" } }, + ] + } + if (state.phase === "awaiting-user") { + const result: Transition[] = [ + { name: "decline", next: { ...state, phase: "stopped", stopReason: "decline" } }, + { name: "abort-prompt", next: { ...state, phase: "stopped", stopReason: "prompt-abort" } }, + ] + if (state.roundsApproved < MAX_APPROVED_ROUNDS) { + result.push({ + name: "approve", + next: { + ...state, + phase: "requesting", + retryAttempt: 0, + requests: state.requests + 1, + roundsApproved: state.roundsApproved + 1, + }, + }) + } + return result + } + return [] +} + +function invariantViolations(state: ModelState): string[] { + const violations: string[] = [] + if (state.announcements !== state.requests - 1 - state.roundsApproved) { + violations.push("every automatic retry must have exactly one user-visible announcement") + } + if (state.retryAttempt > MAX_MID_STREAM_RETRIES) { + violations.push("an automatic retry exceeded the configured retry budget") + } + if (state.phase === "awaiting-user" && state.retryAttempt !== MAX_MID_STREAM_RETRIES) { + violations.push("user input must be requested exactly when the automatic retry budget is exhausted") + } + return violations +} + +function canonical(state: ModelState): string { + return JSON.stringify(state) +} + +function formatCounterexample(message: string, trace: TraceStep[]): string { + return [ + `Mid-stream retry invariant failed: ${message}`, + `Bounds: depth=${MAX_DEPTH}, states=${MAX_STATES}, approvedRounds=${MAX_APPROVED_ROUNDS}`, + ...trace.map((step, index) => `${index}. ${step.action} ${JSON.stringify(step.state)}`), + ].join("\n") +} + +function runModelCheck(): number { + const start = initialState() + const queue: Array<{ state: ModelState; trace: TraceStep[] }> = [ + { state: start, trace: [{ action: "initial", state: start }] }, + ] + const visited = new Set([canonical(start)]) + const reachedActions = new Set() + const reachedLandmarks = new Set() + const frontier: ModelState[] = [] + + for (let index = 0; index < queue.length; index++) { + const node = queue[index]! + for (const [name, predicate] of Object.entries(landmarks)) { + if (predicate(node.state)) reachedLandmarks.add(name) + } + const violations = invariantViolations(node.state) + if (violations.length) throw new Error(formatCounterexample(violations.join("; "), node.trace)) + if (node.trace.length - 1 === MAX_DEPTH) { + frontier.push(node.state) + continue + } + for (const transition of transitions(node.state)) { + reachedActions.add(transition.name) + const key = canonical(transition.next) + if (visited.has(key)) continue + visited.add(key) + queue.push({ + state: transition.next, + trace: [...node.trace, { action: transition.name, state: transition.next }], + }) + if (visited.size > MAX_STATES) throw new Error(`Mid-stream retry model exceeded ${MAX_STATES} states`) + } + } + + const unreachableActions = expectedActions.filter((action) => !reachedActions.has(action)) + assert.deepEqual(unreachableActions, [], `Unreachable actions: ${unreachableActions.join(", ")}`) + const missingLandmarks = Object.keys(landmarks).filter((name) => !reachedLandmarks.has(name)) + assert.deepEqual(missingLandmarks, [], `Unreachable landmarks: ${missingLandmarks.join(", ")}`) + const unexploredSuccessor = frontier.flatMap(transitions).find((next) => !visited.has(canonical(next.next))) + assert.equal(unexploredSuccessor, undefined, "Increase MAX_DEPTH to cover unseen successors") + return visited.size +} + +const checkedStates = runModelCheck() +console.log( + `Mid-stream retry model check passed: ${checkedStates} reachable states, ${expectedActions.length}/${expectedActions.length} actions reachable, ${Object.keys(landmarks).length}/${Object.keys(landmarks).length} landmarks reached, depth <= ${MAX_DEPTH}`, +) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 92ee8184d6..47b47c6cf7 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -138,6 +138,7 @@ import { validateAndFixToolResultIds } from "./validateToolResultIds" import { mergeConsecutiveApiMessages } from "./mergeConsecutiveApiMessages" import { prepareApiConversationMessage } from "./apiConversationHistory" import { shouldAddUserMessageToHistory } from "./messageCounting" +import { decideMidStreamFailure, findRetryRequestMessageIndex, MAX_MID_STREAM_RETRIES } from "./midStreamRetry" import { type TaskExecutionContext } from "./providerHandoff" const MAX_EXPONENTIAL_BACKOFF_SECONDS = 600 // 10 minutes @@ -2696,6 +2697,7 @@ export class Task extends EventEmitter implements TaskLike { private async disposeOnce(): Promise { console.log(`[Task#dispose] disposing task ${this.taskId}.${this.instanceId}`) + this.abort = true this.cancelAssistantMessagePersistence() // Stop the idle telemetry check and report any unflushed activity as a @@ -2924,6 +2926,7 @@ export class Task extends EventEmitter implements TaskLike { includeFileDetails: boolean retryAttempt?: number userMessageWasRemoved?: boolean // Track if user message was removed due to empty response + requestMessageId?: string } const stack: StackItem[] = [{ userContent, includeFileDetails, retryAttempt: 0 }] @@ -3065,9 +3068,11 @@ export class Task extends EventEmitter implements TaskLike { isEmptyUserContent, userMessageWasRemoved: currentItem.userMessageWasRemoved, }) + let requestMessageId = currentItem.requestMessageId if (shouldAddUserMessage) { await this.addToApiConversationHistory({ role: "user", content: finalUserContent }) this.messageCounts.user++ + requestMessageId = this.apiConversationHistory.at(-1)?.messageId } // Since we sent off a placeholder api_req_started message to update the @@ -3655,34 +3660,124 @@ export class Task extends EventEmitter implements TaskLike { this.abortReason = cancelReason await this.abortTask() } else { + // Stream failed - retry with the same content, but only up to + // MAX_MID_STREAM_RETRIES automatic attempts. Every attempt + // re-bills the full input context, so retries must be both + // bounded and visible to the user. + const midStreamRetryAttempt = currentItem.retryAttempt ?? 0 + + if (decideMidStreamFailure(midStreamRetryAttempt) === "ask") { + // Automatic retries exhausted - surface the failure instead of + // retrying (and re-billing the request) silently forever. + // Stryker disable next-line CallExpression: console output has no retry-protocol side effect. + console.error( + // Stryker disable next-line StringLiteral: diagnostic-only task identity and retry-limit text. + `[Task#${this.taskId}.${this.instanceId}] Stream failed, automatic retry limit (${MAX_MID_STREAM_RETRIES}) reached: ${streamingFailedMessage}`, + ) + + const { response } = await this.ask( + "api_req_failed", + // Stryker disable next-line LogicalOperator: both fallbacks describe the same provider failure. + streamingFailedMessage ?? rawErrorMessage, + ) + + if (response === "yesButtonClicked") { + await this.say("api_req_retried") + + // The user approved another round of retries, so reset the + // automatic retry budget. Remove the user message this request + // added first so it is not duplicated in history on retry. + if (requestMessageId) { + const requestMessageIndex = findRetryRequestMessageIndex( + this.apiConversationHistory, + requestMessageId, + ) + if (requestMessageIndex === -1) { + await this.say( + "error", + "Failed to locate the API request in conversation history.", + ) + return false + } + const [requestMessage] = this.apiConversationHistory.splice( + requestMessageIndex, + 1, + ) + this.messageCounts.user-- + if (!(await this.saveApiConversationHistory(false))) { + this.apiConversationHistory.splice(requestMessageIndex, 0, requestMessage!) + this.messageCounts.user++ + await this.say( + "error", + "Failed to persist conversation history before retrying.", + ) + return false + } + } + + stack.push({ + userContent: currentUserContent, + // Stryker disable next-line BooleanLiteral: file details are already materialized in the persisted request being retried. + includeFileDetails: false, + retryAttempt: 0, + }) + + // Continue to retry the request + continue + } + + // User declined to retry - record the failure visibly and stop. + // Stryker disable next-line LogicalOperator: both fallbacks describe the same provider failure. + await this.say("error", streamingFailedMessage ?? rawErrorMessage) + + // Synthetic assistant message recording the failure -- increment + // messageCounts.assistant to match, same as the normal + // assistant-message-saved path. + await this.addToApiConversationHistory({ + role: "assistant", + content: [ + { + type: "text", + text: "Failure: The API request failed mid-stream and the retry was declined.", + }, + ], + }) + this.messageCounts.assistant++ + + return false + } + // Stream failed - log the error and retry with the same content - // The existing rate limiting will prevent rapid retries console.error( - `[Task#${this.taskId}.${this.instanceId}] Stream failed, will retry: ${streamingFailedMessage}`, + // Stryker disable next-line StringLiteral,ArithmeticOperator: diagnostic-only attempt metadata. + `[Task#${this.taskId}.${this.instanceId}] Stream failed, will retry (attempt ${midStreamRetryAttempt + 1}/${MAX_MID_STREAM_RETRIES}): ${streamingFailedMessage}`, ) - // Apply exponential backoff similar to first-chunk errors when auto-resubmit is enabled - const stateForBackoff = await this.providerRef.deref()?.getState() - if (stateForBackoff?.autoApprovalEnabled) { - await this.backoffAndAnnounce(currentItem.retryAttempt ?? 0, error) - - // Check if task was aborted during the backoff - if (this.abort) { - console.log( - `[Task#${this.taskId}.${this.instanceId}] Task aborted during mid-stream retry backoff`, - ) - // Abort the entire task - this.abortReason = "user_cancelled" - await this.abortTask() - break - } + // Announce every automatic retry with the shared exponential + // backoff countdown (api_req_retry_delayed) so no retry - and its + // associated token cost - happens silently. + await this.backoffAndAnnounce(midStreamRetryAttempt, error) + + // Check if task was aborted during the backoff + if (this.abort) { + // Stryker disable next-line CallExpression: console output has no cancellation side effect. + console.log( + // Stryker disable next-line StringLiteral: diagnostic-only task identity. + `[Task#${this.taskId}.${this.instanceId}] Task aborted during mid-stream retry backoff`, + ) + // Abort the entire task + // Stryker disable next-line StringLiteral: abort reason is existing diagnostic metadata. + this.abortReason = "user_cancelled" + await this.abortTask() + break } // Push the same content back onto the stack to retry, incrementing the retry attempt counter stack.push({ userContent: currentUserContent, includeFileDetails: false, - retryAttempt: (currentItem.retryAttempt ?? 0) + 1, + retryAttempt: midStreamRetryAttempt + 1, + requestMessageId, }) // Continue to retry the request diff --git a/src/core/task/__tests__/Task.dispose.test.ts b/src/core/task/__tests__/Task.dispose.test.ts index 472218fce5..cb8ce0ed6a 100644 --- a/src/core/task/__tests__/Task.dispose.test.ts +++ b/src/core/task/__tests__/Task.dispose.test.ts @@ -118,6 +118,18 @@ describe("Task dispose method", () => { expect(disposalComplete).toBe(true) }) + test("should cancel a pending ask when disposed directly", async () => { + const pendingAsk = task.ask("api_req_failed", "provider failed") + await vi.waitFor(() => + expect(task.clineMessages).toContainEqual(expect.objectContaining({ type: "ask", ask: "api_req_failed" })), + ) + + await task.dispose() + + await expect(pendingAsk).rejects.toThrow(/aborted/) + expect(task.abort).toBe(true) + }) + test("should reject the memoized completion promise when disposal cannot start", async () => { const disposalError = new Error("disposal failed") skipCleanup = true diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 1bcacd459c..35c9cc0b33 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -647,6 +647,199 @@ describe("Cline", () => { }) }) + describe("mid-stream retries", () => { + function midStreamFailingRequest(error: Error) { + return (async function* () { + yield { type: "text", text: "partial response" } as ApiStreamChunk + throw error + })() + } + + async function createMidStreamRetryTask() { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + const state = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...state, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: false, + }) + vi.spyOn(task.diffViewProvider, "reset").mockResolvedValue(undefined) + return task + } + + it("stops auto-retrying and surfaces the failure after the mid-stream retry limit", async () => { + const task = await createMidStreamRetryTask() + const streamError = new Error("Overloaded") + const askSpy = vi + .spyOn(task, "ask") + .mockResolvedValue({ response: "noButtonClicked" } satisfies TaskAskResult) + const saySpy = vi.spyOn(task, "say") + const attemptSpy = vi + .spyOn(task, "attemptApiRequest") + .mockImplementation(() => midStreamFailingRequest(streamError)) + + const result = await task.recursivelyMakeClineRequests([{ type: "text", text: "original user request" }]) + + // 1 initial attempt + MAX_MID_STREAM_RETRIES (3) automatic retries, then it stops + // instead of looping (and re-billing the request) forever. + expect(attemptSpy).toHaveBeenCalledTimes(4) + + // Every automatic retry is announced through the shared backoff countdown so + // no retry happens silently. + const completedRetryAnnouncements = saySpy.mock.calls.filter( + (call) => call[0] === "api_req_retry_delayed" && call[3] === false, + ) + expect(completedRetryAnnouncements).toHaveLength(3) + + // Once the cap is exhausted the failure is surfaced to the user. + expect(askSpy).toHaveBeenCalledTimes(1) + expect(askSpy).toHaveBeenCalledWith("api_req_failed", expect.stringContaining("Overloaded")) + expect(saySpy).toHaveBeenCalledWith("error", expect.stringContaining("Overloaded")) + + expect(result).toBe(false) + expect(task.apiConversationHistory).toMatchObject([ + { + role: "user", + content: expect.arrayContaining([expect.objectContaining({ text: "original user request" })]), + }, + { role: "assistant", content: [{ type: "text", text: expect.stringContaining("Failure") }] }, + ]) + expect(task.messageCounts).toEqual({ user: 1, assistant: 1 }) + }) + + it("resets the retry budget when the user approves another retry round", async () => { + const task = await createMidStreamRetryTask() + task["saveApiConversationHistory"] = vi.fn().mockResolvedValue(true) + const streamError = new Error("Overloaded") + const summaryMessage = { + role: "user" as const, + content: [{ type: "text" as const, text: "context summary" }], + messageId: "summary-after-request", + isSummary: true, + } + const askSpy = vi + .spyOn(task, "ask") + .mockImplementationOnce(async () => { + task.apiConversationHistory.push(summaryMessage) + return { response: "yesButtonClicked" } satisfies TaskAskResult + }) + .mockResolvedValueOnce({ response: "noButtonClicked" } satisfies TaskAskResult) + const saySpy = vi.spyOn(task, "say") + const attemptSpy = vi + .spyOn(task, "attemptApiRequest") + .mockImplementation(() => midStreamFailingRequest(streamError)) + + await task.recursivelyMakeClineRequests([{ type: "text", text: "original user request" }]) + + expect(askSpy).toHaveBeenCalledTimes(2) + expect(saySpy).toHaveBeenCalledWith("api_req_retried") + // Each round gets one initial request and three automatic retries. + expect(attemptSpy).toHaveBeenCalledTimes(8) + const completedRetryAnnouncements = saySpy.mock.calls.filter( + (call) => call[0] === "api_req_retry_delayed" && call[3] === false, + ) + expect(completedRetryAnnouncements).toHaveLength(6) + expect(task.messageCounts).toEqual({ user: 1, assistant: 1 }) + expect(task["saveApiConversationHistory"]).toHaveBeenCalledWith(false) + expect(task.apiConversationHistory).toContainEqual(summaryMessage) + expect( + task.apiConversationHistory.findIndex((message) => message.messageId === summaryMessage.messageId), + ).toBe(0) + }) + + it("does not remove an earlier user turn when approving an empty continuation retry", async () => { + const task = await createMidStreamRetryTask() + task["saveApiConversationHistory"] = vi.fn().mockResolvedValue(true) + const earlierUserMessage = { role: "user" as const, content: [{ type: "text" as const, text: "earlier" }] } + task.apiConversationHistory.push(earlierUserMessage) + task.messageCounts.user++ + vi.spyOn(task, "ask") + .mockResolvedValueOnce({ response: "yesButtonClicked" } satisfies TaskAskResult) + .mockResolvedValueOnce({ response: "noButtonClicked" } satisfies TaskAskResult) + vi.spyOn(task, "say") + const attemptSpy = vi + .spyOn(task, "attemptApiRequest") + .mockImplementation(() => midStreamFailingRequest(new Error("Overloaded"))) + + await task.recursivelyMakeClineRequests([]) + + expect(task.ask).toHaveBeenCalledTimes(2) + expect(attemptSpy).toHaveBeenCalledTimes(8) + expect(task.apiConversationHistory).toContainEqual(earlierUserMessage) + expect(task.messageCounts).toEqual({ user: 1, assistant: 1 }) + }) + + it("stops and restores history when the approved retry deletion cannot be persisted", async () => { + const task = await createMidStreamRetryTask() + task["saveApiConversationHistory"] = vi.fn().mockImplementation(async (merge = true) => merge) + const summaryMessage = { + role: "user" as const, + content: [{ type: "text" as const, text: "context summary" }], + messageId: "summary-after-request", + isSummary: true, + } + vi.spyOn(task, "ask").mockImplementation(async () => { + task.apiConversationHistory.push(summaryMessage) + return { response: "yesButtonClicked" } satisfies TaskAskResult + }) + const saySpy = vi.spyOn(task, "say") + const attemptSpy = vi + .spyOn(task, "attemptApiRequest") + .mockImplementation(() => midStreamFailingRequest(new Error("Overloaded"))) + + const result = await task.recursivelyMakeClineRequests([{ type: "text", text: "original user request" }]) + + expect(result).toBe(false) + expect(attemptSpy).toHaveBeenCalledTimes(4) + expect(task["saveApiConversationHistory"]).toHaveBeenCalledWith(false) + expect(saySpy).toHaveBeenCalledWith("error", "Failed to persist conversation history before retrying.") + expect(task.apiConversationHistory).toHaveLength(2) + expect(task.apiConversationHistory[0]?.role).toBe("user") + expect(task.apiConversationHistory[1]).toEqual(summaryMessage) + expect(task.messageCounts.user).toBe(1) + }) + + it("stops when the owned request message cannot be located by id", async () => { + const task = await createMidStreamRetryTask() + vi.spyOn(task, "ask").mockImplementation(async () => { + task.apiConversationHistory[0]!.messageId = "replaced-request-id" + return { response: "yesButtonClicked" } satisfies TaskAskResult + }) + const saySpy = vi.spyOn(task, "say") + const attemptSpy = vi + .spyOn(task, "attemptApiRequest") + .mockImplementation(() => midStreamFailingRequest(new Error("Overloaded"))) + + const result = await task.recursivelyMakeClineRequests([{ type: "text", text: "original user request" }]) + + expect(result).toBe(false) + expect(attemptSpy).toHaveBeenCalledTimes(4) + expect(saySpy).toHaveBeenCalledWith("error", "Failed to locate the API request in conversation history.") + expect(task.messageCounts.user).toBe(1) + }) + + it("does not enqueue another request when aborted during retry backoff", async () => { + const task = await createMidStreamRetryTask() + const abortSpy = vi.spyOn(task, "abortTask").mockResolvedValue(undefined) + task["backoffAndAnnounce"] = vi.fn().mockImplementation(async () => { + task.abort = true + }) + const attemptSpy = vi + .spyOn(task, "attemptApiRequest") + .mockImplementation(() => midStreamFailingRequest(new Error("Overloaded"))) + + await task.recursivelyMakeClineRequests([{ type: "text", text: "original user request" }]) + + expect(attemptSpy).toHaveBeenCalledTimes(1) + expect(abortSpy).toHaveBeenCalledOnce() + }) + }) + describe("constructor", () => { it.each([{ apiConfigName: "parent-local-profile" }, { apiConfigName: undefined }])( "uses an explicit delegated-child context without shared state or startup persistence", diff --git a/src/core/task/__tests__/midStreamRetry.spec.ts b/src/core/task/__tests__/midStreamRetry.spec.ts new file mode 100644 index 0000000000..eafaa49001 --- /dev/null +++ b/src/core/task/__tests__/midStreamRetry.spec.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from "vitest" + +import { decideMidStreamFailure, findRetryRequestMessageIndex, MAX_MID_STREAM_RETRIES } from "../midStreamRetry" + +describe("decideMidStreamFailure", () => { + it.each([0, 1, 2])("retries attempt %i", (attempt) => { + expect(decideMidStreamFailure(attempt)).toBe("retry") + }) + + it("asks after the automatic retry budget is exhausted", () => { + expect(decideMidStreamFailure(MAX_MID_STREAM_RETRIES)).toBe("ask") + }) +}) + +describe("findRetryRequestMessageIndex", () => { + const messages = [ + { messageId: "other-user", role: "user" }, + { messageId: "request", role: "user" }, + { messageId: "request", role: "assistant" }, + ] + + it("finds the exact user request", () => { + expect(findRetryRequestMessageIndex(messages, "request")).toBe(1) + }) + + it("does not match a different id or a non-user message", () => { + expect(findRetryRequestMessageIndex(messages, "missing")).toBe(-1) + expect(findRetryRequestMessageIndex([{ messageId: "request", role: "assistant" }], "request")).toBe(-1) + }) +}) diff --git a/src/core/task/midStreamRetry.ts b/src/core/task/midStreamRetry.ts new file mode 100644 index 0000000000..c467d5d928 --- /dev/null +++ b/src/core/task/midStreamRetry.ts @@ -0,0 +1,14 @@ +export const MAX_MID_STREAM_RETRIES = 3 + +export type MidStreamFailureDecision = "retry" | "ask" + +export function decideMidStreamFailure(retryAttempt: number): MidStreamFailureDecision { + return retryAttempt < MAX_MID_STREAM_RETRIES ? "retry" : "ask" +} + +export function findRetryRequestMessageIndex( + messages: Array<{ messageId?: string; role?: string }>, + requestMessageId: string, +): number { + return messages.findIndex((message) => message.messageId === requestMessageId && message.role === "user") +}