[Fix] Billed requests with no response when the provider errors mid-stream

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<Response> {
+		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<Uint8Array>({
+		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)
+		}
+	})
+})

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:

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",

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<string, (state: ModelState) => 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<string>()
+	const reachedLandmarks = new Set<string>()
+	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}`,
+)