oh-my-opencode-free-fork/src/hooks/session-recovery.ts

/**
 * Session Recovery - Message State Error Recovery
 *
 * Handles FOUR specific scenarios:
 * 1. tool_use block exists without tool_result
 *    - Recovery: inject tool_result with "cancelled" content
 *
 * 2. Thinking block order violation (first block must be thinking)
 *    - Recovery: prepend empty thinking block
 *
 * 3. Thinking disabled but message contains thinking blocks
 *    - Recovery: strip thinking/redacted_thinking blocks
 *
 * 4. Empty content message (non-empty content required)
 *    - Recovery: inject text part directly via filesystem
 */

import { existsSync, mkdirSync, readdirSync, readFileSync, writeFileSync } from "node:fs"
import { join } from "node:path"
import { xdgData } from "xdg-basedir"
import type { PluginInput } from "@opencode-ai/plugin"
import type { createOpencodeClient } from "@opencode-ai/sdk"

type Client = ReturnType<typeof createOpencodeClient>

const OPENCODE_STORAGE = join(xdgData ?? "", "opencode", "storage")
const MESSAGE_STORAGE = join(OPENCODE_STORAGE, "message")
const PART_STORAGE = join(OPENCODE_STORAGE, "part")

type RecoveryErrorType = "tool_result_missing" | "thinking_block_order" | "thinking_disabled_violation" | "empty_content_message" | null

interface MessageInfo {
  id?: string
  role?: string
  sessionID?: string
  parentID?: string
  error?: unknown
}

interface ToolUsePart {
  type: "tool_use"
  id: string
  name: string
  input: Record<string, unknown>
}

interface ThinkingPart {
  type: "thinking"
  thinking: string
}

interface MessagePart {
  type: string
  id?: string
  text?: string
  thinking?: string
  name?: string
  input?: Record<string, unknown>
}

interface MessageData {
  info?: MessageInfo
  parts?: MessagePart[]
}

function getErrorMessage(error: unknown): string {
  if (!error) return ""
  if (typeof error === "string") return error.toLowerCase()
  const errorObj = error as { data?: { message?: string }; message?: string }
  return (errorObj.data?.message || errorObj.message || "").toLowerCase()
}

function detectErrorType(error: unknown): RecoveryErrorType {
  const message = getErrorMessage(error)

  if (message.includes("tool_use") && message.includes("tool_result")) {
    return "tool_result_missing"
  }

  if (
    message.includes("thinking") &&
    (message.includes("first block") || message.includes("must start with") || message.includes("preceeding"))
  ) {
    return "thinking_block_order"
  }

  if (message.includes("thinking is disabled") && message.includes("cannot contain")) {
    return "thinking_disabled_violation"
  }

  if (message.includes("non-empty content") || message.includes("must have non-empty content")) {
    return "empty_content_message"
  }

  return null
}

function extractToolUseIds(parts: MessagePart[]): string[] {
  return parts.filter((p): p is ToolUsePart => p.type === "tool_use" && !!p.id).map((p) => p.id)
}

async function recoverToolResultMissing(
  client: Client,
  sessionID: string,
  failedAssistantMsg: MessageData
): Promise<boolean> {
  const parts = failedAssistantMsg.parts || []
  const toolUseIds = extractToolUseIds(parts)

  if (toolUseIds.length === 0) {
    return false
  }

  const toolResultParts = toolUseIds.map((id) => ({
    type: "tool_result" as const,
    tool_use_id: id,
    content: "Operation cancelled by user (ESC pressed)",
  }))

  try {
    await client.session.prompt({
      path: { id: sessionID },
      // @ts-expect-error - SDK types may not include tool_result parts, but runtime accepts it
      body: { parts: toolResultParts },
    })

    return true
  } catch {
    return false
  }
}

async function recoverThinkingBlockOrder(
  client: Client,
  sessionID: string,
  failedAssistantMsg: MessageData,
  directory: string
): Promise<boolean> {
  const messageID = failedAssistantMsg.info?.id
  if (!messageID) {
    return false
  }

  const existingParts = failedAssistantMsg.parts || []
  const patchedParts: MessagePart[] = [{ type: "thinking", thinking: "" } as ThinkingPart, ...existingParts]

  try {
    // @ts-expect-error - Experimental API
    await client.message?.update?.({
      path: { id: messageID },
      body: { parts: patchedParts },
    })

    return true
  } catch {
    // message.update not available
  }

  try {
    // @ts-expect-error - Experimental API
    await client.session.patch?.({
      path: { id: sessionID },
      body: {
        messageID,
        parts: patchedParts,
      },
    })

    return true
  } catch {
    // session.patch not available
  }

  return await fallbackRevertStrategy(client, sessionID, failedAssistantMsg, directory)
}

async function recoverThinkingDisabledViolation(
  client: Client,
  sessionID: string,
  failedAssistantMsg: MessageData
): Promise<boolean> {
  const messageID = failedAssistantMsg.info?.id
  if (!messageID) {
    return false
  }

  const existingParts = failedAssistantMsg.parts || []
  const strippedParts = existingParts.filter((p) => p.type !== "thinking" && p.type !== "redacted_thinking")

  if (strippedParts.length === 0) {
    return false
  }

  try {
    // @ts-expect-error - Experimental API
    await client.message?.update?.({
      path: { id: messageID },
      body: { parts: strippedParts },
    })

    return true
  } catch {
    // message.update not available
  }

  try {
    // @ts-expect-error - Experimental API
    await client.session.patch?.({
      path: { id: sessionID },
      body: {
        messageID,
        parts: strippedParts,
      },
    })

    return true
  } catch {
    // session.patch not available
  }

  return false
}

const THINKING_TYPES = new Set(["thinking", "redacted_thinking", "reasoning"])
const META_TYPES = new Set(["step-start", "step-finish"])

interface StoredMessageMeta {
  id: string
  sessionID: string
  role: string
  parentID?: string
}

interface StoredPart {
  id: string
  sessionID: string
  messageID: string
  type: string
  text?: string
}

function generatePartId(): string {
  const timestamp = Date.now().toString(16)
  const random = Math.random().toString(36).substring(2, 10)
  return `prt_${timestamp}${random}`
}

function getMessageDir(sessionID: string): string {
  const projectHash = readdirSync(MESSAGE_STORAGE).find((dir) => {
    const sessionDir = join(MESSAGE_STORAGE, dir)
    try {
      return readdirSync(sessionDir).some((f) => f.includes(sessionID.replace("ses_", "")))
    } catch {
      return false
    }
  })

  if (projectHash) {
    return join(MESSAGE_STORAGE, projectHash, sessionID)
  }

  for (const dir of readdirSync(MESSAGE_STORAGE)) {
    const sessionPath = join(MESSAGE_STORAGE, dir, sessionID)
    if (existsSync(sessionPath)) {
      return sessionPath
    }
  }

  return ""
}

function readMessagesFromStorage(sessionID: string): StoredMessageMeta[] {
  const messageDir = getMessageDir(sessionID)
  if (!messageDir || !existsSync(messageDir)) return []

  const messages: StoredMessageMeta[] = []
  for (const file of readdirSync(messageDir)) {
    if (!file.endsWith(".json")) continue
    try {
      const content = readFileSync(join(messageDir, file), "utf-8")
      messages.push(JSON.parse(content))
    } catch {
      continue
    }
  }

  return messages.sort((a, b) => a.id.localeCompare(b.id))
}

function readPartsFromStorage(messageID: string): StoredPart[] {
  const partDir = join(PART_STORAGE, messageID)
  if (!existsSync(partDir)) return []

  const parts: StoredPart[] = []
  for (const file of readdirSync(partDir)) {
    if (!file.endsWith(".json")) continue
    try {
      const content = readFileSync(join(partDir, file), "utf-8")
      parts.push(JSON.parse(content))
    } catch {
      continue
    }
  }

  return parts
}

function injectTextPartToStorage(sessionID: string, messageID: string, text: string): boolean {
  const partDir = join(PART_STORAGE, messageID)

  if (!existsSync(partDir)) {
    mkdirSync(partDir, { recursive: true })
  }

  const partId = generatePartId()
  const part: StoredPart = {
    id: partId,
    sessionID,
    messageID,
    type: "text",
    text,
  }

  try {
    writeFileSync(join(partDir, `${partId}.json`), JSON.stringify(part, null, 2))
    return true
  } catch {
    return false
  }
}

function findEmptyContentMessageFromStorage(sessionID: string): string | null {
  const messages = readMessagesFromStorage(sessionID)

  for (let i = 0; i < messages.length; i++) {
    const msg = messages[i]
    if (msg.role !== "assistant") continue

    const isLastMessage = i === messages.length - 1
    if (isLastMessage) continue

    const parts = readPartsFromStorage(msg.id)
    const hasContent = parts.some((p) => {
      if (THINKING_TYPES.has(p.type)) return false
      if (META_TYPES.has(p.type)) return false
      if (p.type === "text" && p.text?.trim()) return true
      if (p.type === "tool_use") return true
      if (p.type === "tool_result") return true
      return false
    })

    if (!hasContent && parts.length > 0) {
      return msg.id
    }
  }

  return null
}

function hasNonEmptyOutput(msg: MessageData): boolean {
  const parts = msg.parts
  if (!parts || parts.length === 0) return false

  return parts.some((p) => {
    if (THINKING_TYPES.has(p.type)) return false
    if (p.type === "step-start" || p.type === "step-finish") return false
    if (p.type === "text" && p.text && p.text.trim()) return true
    if (p.type === "tool_use" && p.id) return true
    if (p.type === "tool_result") return true
    return false
  })
}

function findEmptyContentMessage(msgs: MessageData[]): MessageData | null {
  for (let i = 0; i < msgs.length; i++) {
    const msg = msgs[i]
    const isLastMessage = i === msgs.length - 1
    const isAssistant = msg.info?.role === "assistant"

    if (isLastMessage && isAssistant) continue

    if (!hasNonEmptyOutput(msg)) {
      return msg
    }
  }
  return null
}

async function recoverEmptyContentMessage(
  _client: Client,
  sessionID: string,
  failedAssistantMsg: MessageData,
  _directory: string
): Promise<boolean> {
  const emptyMessageID = findEmptyContentMessageFromStorage(sessionID) || failedAssistantMsg.info?.id
  if (!emptyMessageID) return false

  return injectTextPartToStorage(sessionID, emptyMessageID, "(interrupted)")
}

async function fallbackRevertStrategy(
  client: Client,
  sessionID: string,
  failedAssistantMsg: MessageData,
  directory: string
): Promise<boolean> {
  const parentMsgID = failedAssistantMsg.info?.parentID

  const messagesResp = await client.session.messages({
    path: { id: sessionID },
    query: { directory },
  })
  const msgs = (messagesResp as { data?: MessageData[] }).data
  if (!msgs || msgs.length === 0) {
    return false
  }

  let targetUserMsg: MessageData | null = null
  if (parentMsgID) {
    targetUserMsg = msgs.find((m) => m.info?.id === parentMsgID) ?? null
  }
  if (!targetUserMsg) {
    for (let i = msgs.length - 1; i >= 0; i--) {
      if (msgs[i].info?.role === "user") {
        targetUserMsg = msgs[i]
        break
      }
    }
  }

  if (!targetUserMsg?.parts?.length) {
    return false
  }

  await client.session.revert({
    path: { id: sessionID },
    body: { messageID: targetUserMsg.info?.id ?? "" },
    query: { directory },
  })

  const textParts = targetUserMsg.parts
    .filter((p) => p.type === "text" && p.text)
    .map((p) => ({ type: "text" as const, text: p.text ?? "" }))

  if (textParts.length === 0) {
    return false
  }

  await client.session.prompt({
    path: { id: sessionID },
    body: { parts: textParts },
    query: { directory },
  })

  return true
}

export function createSessionRecoveryHook(ctx: PluginInput) {
  const processingErrors = new Set<string>()
  let onAbortCallback: ((sessionID: string) => void) | null = null

  const setOnAbortCallback = (callback: (sessionID: string) => void): void => {
    onAbortCallback = callback
  }

  const isRecoverableError = (error: unknown): boolean => {
    return detectErrorType(error) !== null
  }

  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
    const assistantMsgID = info.id

    if (!sessionID || !assistantMsgID) return false
    if (processingErrors.has(assistantMsgID)) return false
    processingErrors.add(assistantMsgID)

    try {
      await ctx.client.session.abort({ path: { id: sessionID } }).catch(() => {})

      if (onAbortCallback) {
        onAbortCallback(sessionID)
      }

      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",
        thinking_block_order: "Thinking Block Recovery",
        thinking_disabled_violation: "Thinking Strip Recovery",
        empty_content_message: "Empty Message Recovery",
      }
      const toastMessages: Record<RecoveryErrorType & string, string> = {
        tool_result_missing: "Injecting cancelled tool results...",
        thinking_block_order: "Fixing message structure...",
        thinking_disabled_violation: "Stripping thinking blocks...",
        empty_content_message: "Deleting empty message...",
      }
      const toastTitle = toastTitles[errorType]
      const toastMessage = toastMessages[errorType]

      await ctx.client.tui
        .showToast({
          body: {
            title: toastTitle,
            message: toastMessage,
            variant: "warning",
            duration: 3000,
          },
        })
        .catch(() => {})

      let success = false

      if (errorType === "tool_result_missing") {
        success = await recoverToolResultMissing(ctx.client, sessionID, failedMsg)
      } else if (errorType === "thinking_block_order") {
        success = await recoverThinkingBlockOrder(ctx.client, sessionID, failedMsg, ctx.directory)
      } else if (errorType === "thinking_disabled_violation") {
        success = await recoverThinkingDisabledViolation(ctx.client, sessionID, failedMsg)
      } else if (errorType === "empty_content_message") {
        success = await recoverEmptyContentMessage(ctx.client, sessionID, failedMsg, ctx.directory)
      }

      return success
    } catch {
      return false
    } finally {
      processingErrors.delete(assistantMsgID)
    }
  }

  return {
    handleSessionRecovery,
    isRecoverableError,
    setOnAbortCallback,
  }
}