From 41629b526284575b28b142696b17d8015f96394e Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 10 May 2026 15:54:29 +0900 Subject: [PATCH] fix(plugin): guard assistant prefill message tails --- src/plugin/messages-transform.test.ts | 24 +++++++++- src/plugin/messages-transform.ts | 65 +++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 1 deletion(-) diff --git a/src/plugin/messages-transform.test.ts b/src/plugin/messages-transform.test.ts index d3c0d6315..d4cb0637d 100644 --- a/src/plugin/messages-transform.test.ts +++ b/src/plugin/messages-transform.test.ts @@ -7,10 +7,13 @@ import type { CreatedHooks } from "../create-hooks" type TestPart = { type: string id?: string + sessionID?: string + messageID?: string callID?: string tool_use_id?: string content?: string text?: string + synthetic?: boolean } type TestMessage = { @@ -38,7 +41,7 @@ function makeHooks(overrides: { contextInjectorMessagesTransform: overrides.contextInjector ? makeHook(overrides.contextInjector) : undefined, thinkingBlockValidator: overrides.thinkingBlock ? makeHook(overrides.thinkingBlock) : undefined, toolPairValidator: overrides.toolPair ? makeHook(overrides.toolPair) : undefined, - } as unknown as CreatedHooks + } as CreatedHooks } async function runHandler( @@ -157,6 +160,25 @@ describe("createMessagesTransformHandler", () => { //#when / #then await runHandler(hooks, []) }) + + it("appends a synthetic user turn when transformed messages end with assistant prefill", async () => { + //#given + const messages: TestMessage[] = [ + { info: { role: "user" }, parts: [{ type: "text", text: "work on this" }] }, + { info: { role: "assistant" }, parts: [{ type: "text", text: "partial assistant tail" }] }, + ] + + //#when + await runHandler(makeHooks({}), messages) + + //#then + expect(messages.at(-1)?.info).toMatchObject({ role: "user" }) + expect(messages.at(-1)?.parts[0]).toMatchObject({ + type: "text", + text: "[internal] Continue from the previous assistant state.", + synthetic: true, + }) + }) }) function createRealToolPairValidator(): TransformHook { diff --git a/src/plugin/messages-transform.ts b/src/plugin/messages-transform.ts index 1eaef3eba..ab3999456 100644 --- a/src/plugin/messages-transform.ts +++ b/src/plugin/messages-transform.ts @@ -3,12 +3,75 @@ 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 + +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( hookName: string, @@ -79,5 +142,7 @@ export function createMessagesTransformHandler(args: { input, output, ) + + ensureUserTurnAfterAssistantTail(output) } }