Add API streaming failed error handling

This commit is contained in:
Saoud Rizwan
2024-09-30 19:52:00 -04:00
parent 9b1b9c10a1
commit 42bcc4420d
3 changed files with 110 additions and 62 deletions

View File

@@ -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,7 +1680,7 @@ export class ClaudeDev {
await this.diffViewProvider.reset()
let assistantMessage = ""
// TODO: handle error being thrown in stream
try {
for await (const chunk of stream) {
switch (chunk.type) {
case "usage":
@@ -1651,35 +1699,19 @@ export class ClaudeDev {
if (this.abort) {
console.log("aborting stream...")
if (this.diffViewProvider.isEditing) {
await this.diffViewProvider.revertChanges() // closes diff view
}
// 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()
}
//
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
await abortStream("user_cancelled")
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()
}
}
// need to call here in case the stream was aborted
if (this.abort) {

View File

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

View File

@@ -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 ? (
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 ? (
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 />