YeonGyu-Kim
296 lines
11 KiB
TypeScript
296 lines
11 KiB
TypeScript
import type { PluginInput } from "@opencode-ai/plugin"
|
|
import type { ExperimentalConfig } from "../../config"
|
|
import { log } from "../../shared/logger"
|
|
import { detectErrorType } from "./detect-error-type"
|
|
import type { RecoveryErrorType } from "./detect-error-type"
|
|
import type { MessageData } from "./types"
|
|
import { normalizeSDKResponse } from "../../shared"
|
|
import {
|
|
getInterruptedIdleMessagesFetchTimeoutMs,
|
|
withInterruptedIdleMessagesFetchTimeout,
|
|
} from "./interrupted-idle-message-fetch-timeout"
|
|
import { recoverToolResultMissing } from "./recover-tool-result-missing"
|
|
import { recoverUnavailableTool } from "./recover-unavailable-tool"
|
|
import { recoverThinkingBlockOrder } from "./recover-thinking-block-order"
|
|
import { recoverThinkingDisabledViolation } from "./recover-thinking-disabled-violation"
|
|
import { extractResumeConfig, findLastUserMessage, resumeSession } from "./resume"
|
|
|
|
interface MessageInfo {
|
|
id?: string
|
|
role?: string
|
|
sessionID?: string
|
|
parentID?: string
|
|
error?: unknown
|
|
}
|
|
|
|
export interface SessionRecoveryOptions {
|
|
experimental?: ExperimentalConfig
|
|
}
|
|
|
|
export interface SessionRecoveryHook {
|
|
handleSessionRecovery: (info: MessageInfo) => Promise<boolean>
|
|
handleInterruptedToolResultsOnIdle: (sessionID: string) => Promise<boolean>
|
|
isRecoverableError: (error: unknown) => boolean
|
|
setOnAbortCallback: (callback: (sessionID: string) => void) => void
|
|
setOnRecoveryCompleteCallback: (callback: (sessionID: string) => void) => void
|
|
}
|
|
|
|
export function createSessionRecoveryHook(ctx: PluginInput, options?: SessionRecoveryOptions): SessionRecoveryHook {
|
|
const processingErrors = new Set<string>()
|
|
const processingInterruptedToolMessages = new Set<string>()
|
|
const experimental = options?.experimental
|
|
let onAbortCallback: ((sessionID: string) => void) | null = null
|
|
let onRecoveryCompleteCallback: ((sessionID: string) => void) | null = null
|
|
|
|
const setOnAbortCallback = (callback: (sessionID: string) => void): void => {
|
|
onAbortCallback = callback
|
|
}
|
|
|
|
const setOnRecoveryCompleteCallback = (callback: (sessionID: string) => void): void => {
|
|
onRecoveryCompleteCallback = callback
|
|
}
|
|
|
|
const isRecoverableError = (error: unknown): boolean => {
|
|
return detectErrorType(error) !== null
|
|
}
|
|
|
|
const assistantMessageIsFinished = (message: MessageData): boolean => {
|
|
if (message.info?.error) {
|
|
return true
|
|
}
|
|
|
|
const finish = message.info?.finish
|
|
if (finish === "tool-calls") {
|
|
return false
|
|
}
|
|
if ((typeof finish === "string" && finish.length > 0) || finish === true) {
|
|
return true
|
|
}
|
|
|
|
const completed = message.info?.time?.completed
|
|
if (typeof completed === "number" && Number.isFinite(completed)) {
|
|
return true
|
|
}
|
|
return typeof completed === "string" && completed.length > 0
|
|
}
|
|
|
|
const partHasValidToolUseID = (part: NonNullable<MessageData["parts"]>[number]): boolean => {
|
|
const callID = part.callID
|
|
if (typeof callID === "string" && /^(toolu_|call_)/.test(callID)) {
|
|
return true
|
|
}
|
|
|
|
const id = part.id
|
|
return typeof id === "string" && /^(toolu_|call_)/.test(id)
|
|
}
|
|
|
|
const messageHasInterruptedToolResults = (message: MessageData): boolean => {
|
|
return message.parts?.some((part) =>
|
|
(part.type === "tool" || part.type === "tool_use")
|
|
&& (part.state?.status === "pending" || part.state?.status === "running")
|
|
&& partHasValidToolUseID(part)
|
|
) === true
|
|
}
|
|
|
|
const findLatestAssistantMessage = (messages: MessageData[]): MessageData | undefined => {
|
|
for (let index = messages.length - 1; index >= 0; index--) {
|
|
const message = messages[index]
|
|
if (message?.info?.role === "assistant") {
|
|
return message
|
|
}
|
|
}
|
|
return undefined
|
|
}
|
|
|
|
const handleInterruptedToolResultsOnIdle = async (sessionID: string): Promise<boolean> => {
|
|
let recoveryStarted = false
|
|
let assistantMessageIDForRecovery: string | undefined
|
|
try {
|
|
const messagesResp = await withInterruptedIdleMessagesFetchTimeout(
|
|
ctx.client.session.messages({
|
|
path: { id: sessionID },
|
|
query: { directory: ctx.directory },
|
|
}),
|
|
getInterruptedIdleMessagesFetchTimeoutMs(),
|
|
)
|
|
const messages = normalizeSDKResponse(messagesResp, [] as MessageData[])
|
|
const latestAssistant = findLatestAssistantMessage(messages)
|
|
if (!latestAssistant?.info?.id) {
|
|
return false
|
|
}
|
|
|
|
if (assistantMessageIsFinished(latestAssistant) || !messageHasInterruptedToolResults(latestAssistant)) {
|
|
return false
|
|
}
|
|
|
|
const assistantMessageID = latestAssistant.info.id
|
|
if (processingInterruptedToolMessages.has(assistantMessageID)) {
|
|
return false
|
|
}
|
|
processingInterruptedToolMessages.add(assistantMessageID)
|
|
assistantMessageIDForRecovery = assistantMessageID
|
|
|
|
if (onAbortCallback) {
|
|
onAbortCallback(sessionID)
|
|
}
|
|
recoveryStarted = true
|
|
|
|
const lastUser = findLastUserMessage(messages)
|
|
const resumeConfig = extractResumeConfig(lastUser, sessionID)
|
|
const success = await recoverToolResultMissing(ctx.client, sessionID, latestAssistant, resumeConfig, {
|
|
recoverStatuses: new Set(["pending", "running"]),
|
|
resultText: "Tool execution was interrupted before producing a result.",
|
|
source: "session-recovery-interrupted-tool-results",
|
|
})
|
|
if (!success) {
|
|
processingInterruptedToolMessages.delete(assistantMessageID)
|
|
}
|
|
return success
|
|
} catch (err) {
|
|
if (assistantMessageIDForRecovery) {
|
|
processingInterruptedToolMessages.delete(assistantMessageIDForRecovery)
|
|
}
|
|
log("[session-recovery] Interrupted tool result recovery failed:", { sessionID, error: err })
|
|
return false
|
|
} finally {
|
|
if (recoveryStarted && onRecoveryCompleteCallback) {
|
|
onRecoveryCompleteCallback(sessionID)
|
|
}
|
|
}
|
|
}
|
|
|
|
const handleSessionRecovery = async (info: MessageInfo): Promise<boolean> => {
|
|
if (!info || info.role !== "assistant" || !info.error) return false
|
|
|
|
const errorType = detectErrorType(info.error)
|
|
if (!errorType) return false
|
|
|
|
const sessionID = info.sessionID
|
|
let assistantMsgID = info.id
|
|
|
|
if (!sessionID) return false
|
|
|
|
if (!assistantMsgID) {
|
|
try {
|
|
const messagesResp = await ctx.client.session.messages({
|
|
path: { id: sessionID },
|
|
query: { directory: ctx.directory },
|
|
})
|
|
const msgs = (messagesResp as { data?: MessageData[] }).data
|
|
const lastAssistant = msgs?.findLast((m) => m.info?.role === "assistant" && m.info?.error)
|
|
assistantMsgID = lastAssistant?.info?.id
|
|
} catch {
|
|
log("[session-recovery] Failed to fetch messages for messageID fallback", { sessionID })
|
|
}
|
|
}
|
|
|
|
if (!assistantMsgID) return false
|
|
if (processingErrors.has(assistantMsgID)) return false
|
|
processingErrors.add(assistantMsgID)
|
|
|
|
try {
|
|
if (onAbortCallback) {
|
|
onAbortCallback(sessionID)
|
|
}
|
|
|
|
await ctx.client.session.abort({ path: { id: sessionID } }).catch(() => {})
|
|
|
|
const messagesResp = await ctx.client.session.messages({
|
|
path: { id: sessionID },
|
|
query: { directory: ctx.directory },
|
|
})
|
|
const msgs = (messagesResp as { data?: MessageData[] }).data
|
|
|
|
const failedMsg = msgs?.find((m) => m.info?.id === assistantMsgID)
|
|
if (!failedMsg) {
|
|
return false
|
|
}
|
|
|
|
const toastTitles: Record<RecoveryErrorType & string, string> = {
|
|
tool_result_missing: "Tool Crash Recovery",
|
|
unavailable_tool: "Tool Recovery",
|
|
thinking_block_order: "Thinking Block Recovery",
|
|
thinking_disabled_violation: "Thinking Strip Recovery",
|
|
thinking_block_modified: "Thinking Block Recovery",
|
|
"assistant_prefill_unsupported": "Prefill Unsupported",
|
|
}
|
|
const toastMessages: Record<RecoveryErrorType & string, string> = {
|
|
tool_result_missing: "Injecting cancelled tool results...",
|
|
unavailable_tool: "Recovering from unavailable tool call...",
|
|
thinking_block_order: "Fixing message structure...",
|
|
thinking_disabled_violation: "Stripping thinking blocks...",
|
|
thinking_block_modified: "Stripping corrupted thinking blocks...",
|
|
"assistant_prefill_unsupported": "Prefill not supported; continuing without recovery.",
|
|
}
|
|
|
|
await ctx.client.tui
|
|
.showToast({
|
|
body: {
|
|
title: toastTitles[errorType],
|
|
message: toastMessages[errorType],
|
|
variant: "warning",
|
|
duration: 3000,
|
|
},
|
|
})
|
|
.catch(() => {})
|
|
|
|
let success = false
|
|
|
|
if (errorType === "tool_result_missing") {
|
|
const lastUser = findLastUserMessage(msgs ?? [])
|
|
const resumeConfig = extractResumeConfig(lastUser, sessionID)
|
|
success = await recoverToolResultMissing(ctx.client, sessionID, failedMsg, resumeConfig)
|
|
} else if (errorType === "unavailable_tool") {
|
|
success = await recoverUnavailableTool(ctx.client, sessionID, failedMsg)
|
|
} else if (errorType === "thinking_block_order") {
|
|
success = await recoverThinkingBlockOrder(ctx.client, sessionID, failedMsg, ctx.directory, info.error)
|
|
if (success && experimental?.auto_resume) {
|
|
const lastUser = findLastUserMessage(msgs ?? [])
|
|
const resumeConfig = extractResumeConfig(lastUser, sessionID)
|
|
await resumeSession(ctx.client, resumeConfig)
|
|
}
|
|
} else if (errorType === "thinking_disabled_violation") {
|
|
success = await recoverThinkingDisabledViolation(ctx.client, sessionID, failedMsg)
|
|
if (success && experimental?.auto_resume) {
|
|
const lastUser = findLastUserMessage(msgs ?? [])
|
|
const resumeConfig = extractResumeConfig(lastUser, sessionID)
|
|
await resumeSession(ctx.client, resumeConfig)
|
|
}
|
|
} else if (errorType === "thinking_block_modified") {
|
|
success = await recoverThinkingDisabledViolation(ctx.client, sessionID, failedMsg)
|
|
if (success && experimental?.auto_resume) {
|
|
const lastUser = findLastUserMessage(msgs ?? [])
|
|
const resumeConfig = extractResumeConfig(lastUser, sessionID)
|
|
await resumeSession(ctx.client, resumeConfig)
|
|
}
|
|
} else if (errorType === "assistant_prefill_unsupported") {
|
|
success = false
|
|
}
|
|
|
|
return success
|
|
} catch (err) {
|
|
log("[session-recovery] Recovery failed:", err)
|
|
return false
|
|
} finally {
|
|
// Keep assistantMsgID in processingErrors permanently so that a
|
|
// stale duplicate session.error for the SAME assistant message
|
|
// does not retrigger recovery (and a second resumeSession
|
|
// promptAsync injection) after the first attempt resolves.
|
|
// Successful recovery starts a new assistant message on the next
|
|
// turn with a different id, so this dedupe never blocks future
|
|
// legitimate errors.
|
|
if (sessionID && onRecoveryCompleteCallback) {
|
|
onRecoveryCompleteCallback(sessionID)
|
|
}
|
|
}
|
|
}
|
|
|
|
return {
|
|
handleSessionRecovery,
|
|
handleInterruptedToolResultsOnIdle,
|
|
isRecoverableError,
|
|
setOnAbortCallback,
|
|
setOnRecoveryCompleteCallback,
|
|
}
|
|
}
|