Add API streaming failed error handling

src/core/ClaudeDev.ts

@@ -21,7 +21,14 @@ import { ApiConfiguration } from "../shared/api"
 import { findLastIndex } from "../shared/array"
 import { combineApiRequests } from "../shared/combineApiRequests"
 import { combineCommandSequences } from "../shared/combineCommandSequences"
-import { ClaudeApiReqInfo, ClaudeAsk, ClaudeMessage, ClaudeSay, ClaudeSayTool } from "../shared/ExtensionMessage"
+import {
+	ClaudeApiReqCancelReason,
+	ClaudeApiReqInfo,
+	ClaudeAsk,
+	ClaudeMessage,
+	ClaudeSay,
+	ClaudeSayTool,
+} from "../shared/ExtensionMessage"
 import { getApiMetrics } from "../shared/getApiMetrics"
 import { HistoryItem } from "../shared/HistoryItem"
 import { ToolName } from "../shared/Tool"
@@ -1600,7 +1607,7 @@ export class ClaudeDev {
 			// update api_req_started. we can't use api_req_finished anymore since it's a unique case where it could come after a streaming message (ie in the middle of being updated or executed)
 			// fortunately api_req_finished was always parsed out for the gui anyways, so it remains solely for legacy purposes to keep track of prices in tasks from history
 			// (it's worth removing a few months from now)
-			const updateApiReqMsg = (cancelled?: boolean) => {
+			const updateApiReqMsg = (cancelReason?: ClaudeApiReqCancelReason, streamingFailedMessage?: string) => {
 				this.claudeMessages[lastApiReqIndex].text = JSON.stringify({
 					...JSON.parse(this.claudeMessages[lastApiReqIndex].text || "{}"),
 					tokensIn: inputTokens,
@@ -1616,10 +1623,51 @@ export class ClaudeDev {
 							cacheWriteTokens,
 							cacheReadTokens
 						),
-					cancelled,
+					cancelReason,
+					streamingFailedMessage,
 				} satisfies ClaudeApiReqInfo)
 			}
 
+			const abortStream = async (cancelReason: ClaudeApiReqCancelReason, streamingFailedMessage?: string) => {
+				if (this.diffViewProvider.isEditing) {
+					await this.diffViewProvider.revertChanges() // closes diff view
+				}
+
+				// if last message is a partial we need to update and save it
+				const lastMessage = this.claudeMessages.at(-1)
+				if (lastMessage && lastMessage.partial) {
+					lastMessage.ts = Date.now()
+					lastMessage.partial = false
+					// instead of streaming partialMessage events, we do a save and post like normal to persist to disk
+					console.log("updating partial message", lastMessage)
+					// await this.saveClaudeMessages()
+				}
+
+				// Let assistant know their response was interrupted for when task is resumed
+				await this.addToApiConversationHistory({
+					role: "assistant",
+					content: [
+						{
+							type: "text",
+							text:
+								assistantMessage +
+								`\n\n[${
+									cancelReason === "streaming_failed"
+										? "Response interrupted by API Error"
+										: "Response interrupted by user"
+								}]`,
+						},
+					],
+				})
+
+				// update api_req_started to have cancelled and cost, so that we can display the cost of the partial stream
+				updateApiReqMsg(cancelReason, streamingFailedMessage)
+				await this.saveClaudeMessages()
+
+				// signals to provider that it can retrieve the saved messages from disk, as abortTask can not be awaited on in nature
+				this.didFinishAborting = true
+			}
+
 			// reset streaming state
 			this.currentStreamingContentIndex = 0
 			this.assistantMessageContent = []
@@ -1632,52 +1680,36 @@ export class ClaudeDev {
 			await this.diffViewProvider.reset()
 
 			let assistantMessage = ""
-			// TODO: handle error being thrown in stream
-			for await (const chunk of stream) {
-				switch (chunk.type) {
-					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
-						this.parseAssistantMessage(assistantMessage)
-						this.presentAssistantMessage()
-						break
-				}
-
-				if (this.abort) {
-					console.log("aborting stream...")
-					if (this.diffViewProvider.isEditing) {
-						await this.diffViewProvider.revertChanges() // closes diff view
+			try {
+				for await (const chunk of stream) {
+					switch (chunk.type) {
+						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
+							this.parseAssistantMessage(assistantMessage)
+							this.presentAssistantMessage()
+							break
 					}
 
-					// if last message is a partial we need to save it
-					const lastMessage = this.claudeMessages.at(-1)
-					if (lastMessage && lastMessage.partial) {
-						lastMessage.ts = Date.now()
-						lastMessage.partial = false
-						// instead of streaming partialMessage events, we do a save and post like normal to persist to disk
-						console.log("saving messages...", lastMessage)
-						// await this.saveClaudeMessages()
+					if (this.abort) {
+						console.log("aborting stream...")
+						await abortStream("user_cancelled")
+						break // aborts the stream
 					}
-
-					//
-					await this.addToApiConversationHistory({
-						role: "assistant",
-						content: [{ type: "text", text: assistantMessage + "\n\n[Response interrupted by user]" }],
-					})
-
-					// update api_req_started to have cancelled and cost, so that we can display the cost of the partial stream
-					updateApiReqMsg(true)
-					await this.saveClaudeMessages()
-
-					// signals to provider that it can retrieve the saved messages from disk, as abortTask can not be awaited on in nature
-					this.didFinishAborting = true
-					break // aborts the stream
+				}
+			} catch (error) {
+				this.abortTask() // 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
+				await abortStream("streaming_failed", error.message ?? JSON.stringify(serializeError(error), null, 2))
+				const history = await this.providerRef.deref()?.getTaskWithId(this.taskId)
+				if (history) {
+					await this.providerRef.deref()?.initClaudeDevWithHistoryItem(history.historyItem)
+					await this.providerRef.deref()?.postStateToWebview()
 				}
 			}
 

src/shared/ExtensionMessage.ts

@@ -95,5 +95,8 @@ export interface ClaudeApiReqInfo {
 	cacheWrites?: number
 	cacheReads?: number
 	cost?: number
-	cancelled?: boolean
+	cancelReason?: ClaudeApiReqCancelReason
+	streamingFailedMessage?: string
 }
+
+export type ClaudeApiReqCancelReason = "streaming_failed" | "user_cancelled"

webview-ui/src/components/chat/ChatRow.tsx

@@ -37,12 +37,12 @@ const ChatRow = memo(
 export default ChatRow
 
 const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessage, isLast }: ChatRowProps) => {
-	const [cost, apiReqCancelled] = useMemo(() => {
+	const [cost, apiReqCancelReason, apiReqStreamingFailedMessage] = useMemo(() => {
 		if (message.text != null && message.say === "api_req_started") {
 			const info: ClaudeApiReqInfo = JSON.parse(message.text)
-			return [info.cost, info.cancelled]
+			return [info.cost, info.cancelReason, info.streamingFailedMessage]
 		}
-		return [undefined, undefined]
+		return [undefined, undefined, undefined]
 	}, [message.text, message.say])
 	const apiRequestFailedMessage =
 		isLast && lastModifiedMessage?.ask === "api_req_failed" // if request is retried then the latest message is a api_req_retried
@@ -96,10 +96,16 @@ const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessa
 			case "api_req_started":
 				return [
 					cost != null ? (
-						apiReqCancelled ? (
-							<span
-								className="codicon codicon-error"
-								style={{ color: cancelledColor, marginBottom: "-1.5px" }}></span>
+						apiReqCancelReason != null ? (
+							apiReqCancelReason === "user_cancelled" ? (
+								<span
+									className="codicon codicon-error"
+									style={{ color: cancelledColor, marginBottom: "-1.5px" }}></span>
+							) : (
+								<span
+									className="codicon codicon-error"
+									style={{ color: errorColor, marginBottom: "-1.5px" }}></span>
+							)
 						) : (
 							<span
 								className="codicon codicon-check"
@@ -113,8 +119,12 @@ const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessa
 						<ProgressIndicator />
 					),
 					cost != null ? (
-						apiReqCancelled ? (
-							<span style={{ color: normalColor, fontWeight: "bold" }}>API Request Cancelled</span>
+						apiReqCancelReason != null ? (
+							apiReqCancelReason === "user_cancelled" ? (
+								<span style={{ color: normalColor, fontWeight: "bold" }}>API Request Cancelled</span>
+							) : (
+								<span style={{ color: errorColor, fontWeight: "bold" }}>API Streaming Failed</span>
+							)
 						) : (
 							<span style={{ color: normalColor, fontWeight: "bold" }}>API Request</span>
 						)
@@ -134,7 +144,7 @@ const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessa
 			default:
 				return [null, null]
 		}
-	}, [type, cost, apiRequestFailedMessage, isCommandExecuting, apiReqCancelled])
+	}, [type, cost, apiRequestFailedMessage, isCommandExecuting, apiReqCancelReason])
 
 	const headerStyle: React.CSSProperties = {
 		display: "flex",
@@ -376,7 +386,10 @@ const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessa
 							<div
 								style={{
 									...headerStyle,
-									marginBottom: cost == null && apiRequestFailedMessage ? 10 : 0,
+									marginBottom:
+										(cost == null && apiRequestFailedMessage) || apiReqStreamingFailedMessage
+											? 10
+											: 0,
 									justifyContent: "space-between",
 									cursor: "pointer",
 									userSelect: "none",
@@ -392,10 +405,10 @@ const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessa
 								</div>
 								<span className={`codicon codicon-chevron-${isExpanded ? "up" : "down"}`}></span>
 							</div>
-							{cost == null && apiRequestFailedMessage && (
+							{((cost == null && apiRequestFailedMessage) || apiReqStreamingFailedMessage) && (
 								<>
 									<p style={{ ...pStyle, color: "var(--vscode-errorForeground)" }}>
-										{apiRequestFailedMessage}
+										{apiRequestFailedMessage || apiReqStreamingFailedMessage}
 										{apiRequestFailedMessage?.toLowerCase().includes("powershell") && (
 											<>
 												<br />