oh-my-opencode/src/plugin/messages-transform.ts

import type { Message, Part } from "@opencode-ai/sdk"

import { log } from "../shared/logger"
import type { CreatedHooks } from "../create-hooks"

const ASSISTANT_PREFILL_RECOVERY_TEXT = "[internal] Continue from the previous assistant state."

type MessageWithParts = {
  info: Message
  parts: Part[]
}

type MessagesTransformOutput = { messages: MessageWithParts[] }
type UserMessageInfo = Extract<Message, { role: "user" }>

function getSessionID(message: MessageWithParts): string | undefined {
  return message.info.sessionID
}

function findLastUserMessage(messages: MessageWithParts[]): UserMessageInfo | undefined {
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const message = messages[index]
    if (message?.info.role === "user") {
      return message.info
    }
  }

  return undefined
}

function createAssistantPrefillRecoveryMessage(
  lastAssistantMessage: MessageWithParts,
  messages: MessageWithParts[],
): MessageWithParts {
  const lastUserMessage = findLastUserMessage(messages)
  const sessionID = getSessionID(lastAssistantMessage) ?? lastUserMessage?.sessionID ?? ""
  const messageID = `${lastAssistantMessage.info.id}_prefill_recovery`
  const model = lastUserMessage?.model ?? {
    providerID: "internal",
    modelID: "assistant-prefill-guard",
  }

  return {
    info: {
      id: messageID,
      sessionID,
      role: "user",
      time: { created: Date.now() },
      agent: lastUserMessage?.agent ?? "internal",
      model,
      ...(lastUserMessage?.system ? { system: lastUserMessage.system } : {}),
      ...(lastUserMessage?.tools ? { tools: lastUserMessage.tools } : {}),
    },
    parts: [
      {
        id: `${messageID}_text`,
        sessionID,
        messageID,
        type: "text",
        text: ASSISTANT_PREFILL_RECOVERY_TEXT,
        synthetic: true,
      },
    ],
  }
}

function ensureUserTurnAfterAssistantTail(output: MessagesTransformOutput): void {
  const lastMessage = output.messages.at(-1)
  if (!lastMessage || lastMessage.info.role !== "assistant") {
    return
  }

  output.messages.push(createAssistantPrefillRecoveryMessage(lastMessage, output.messages))
}

async function runMessagesTransformHookSafely<I, O>(
  hookName: string,
  handler: ((input: I, output: O) => unknown | Promise<unknown>) | null | undefined,
  input: I,
  output: O,
): Promise<void> {
  if (!handler) return
  try {
    await Promise.resolve(handler(input, output))
  } catch (error) {
    // Isolate per-handler failures so later handlers (notably toolPairValidator)
    // always run. A throw here used to leave orphaned tool_use blocks in the
    // post-compaction payload, producing API 400s like
    // "tool_use ids were found without tool_result blocks immediately after".
    log("[messages-transform] hook execution failed", {
      hook: hookName,
      error,
    })
  }
}

export function createMessagesTransformHandler(args: {
  hooks: CreatedHooks
}): (input: Record<string, never>, output: MessagesTransformOutput) => Promise<void> {
  return async (input, output): Promise<void> => {
    await runMessagesTransformHookSafely(
      "contextInjectorMessagesTransform",
      args.hooks.contextInjectorMessagesTransform?.[
        "experimental.chat.messages.transform"
      ],
      input,
      output,
    )

    await runMessagesTransformHookSafely(
      "teamModeStatusInjector",
      args.hooks.teamModeStatusInjector?.[
        "experimental.chat.messages.transform"
      ],
      input,
      output,
    )

    await runMessagesTransformHookSafely(
      "teamMailboxInjector",
      args.hooks.teamMailboxInjector?.[
        "experimental.chat.messages.transform"
      ],
      input,
      output,
    )

    await runMessagesTransformHookSafely(
      "thinkingBlockValidator",
      args.hooks.thinkingBlockValidator?.[
        "experimental.chat.messages.transform"
      ],
      input,
      output,
    )

    await runMessagesTransformHookSafely(
      "toolPairValidator",
      args.hooks.toolPairValidator?.[
        "experimental.chat.messages.transform"
      ],
      input,
      output,
    )

    ensureUserTurnAfterAssistantTail(output)
  }
}