diff --git a/src/hooks/thinking-block-validator/hook.test.ts b/src/hooks/thinking-block-validator/hook.test.ts new file mode 100644 index 000000000..9eab6dff6 --- /dev/null +++ b/src/hooks/thinking-block-validator/hook.test.ts @@ -0,0 +1,108 @@ +const { describe, expect, test } = require("bun:test") + +const { createThinkingBlockValidatorHook } = require("./hook") + +type TestPart = { + type: string + id: string + text?: string + thinking?: string + data?: string + signature?: string +} + +type TestMessage = { + info: { + role: string + id?: string + modelID?: string + } + parts: TestPart[] +} + +function createMessage(info: TestMessage["info"], parts: TestPart[]): TestMessage { + return { info, parts } +} + +function createTextPart(id: string, text: string): TestPart { + return { type: "text", id, text } +} + +function createSignedThinkingPart(id: string, thinking: string, signature: string): TestPart { + return { type: "thinking", id, thinking, signature } +} + +function createRedactedThinkingPart(id: string, signature: string): TestPart { + return { type: "redacted_thinking", id, data: "encrypted", signature } +} + +describe("createThinkingBlockValidatorHook", () => { + test("reuses the previous signed thinking part verbatim when assistant content lacks a leading thinking block", async () => { + const transform = Reflect.get(createThinkingBlockValidatorHook(), "experimental.chat.messages.transform") + expect(typeof transform).toBe("function") + + const previousThinkingPart = createSignedThinkingPart("prt_prev_signed", "prior reasoning", "sig_prev") + const targetTextPart = createTextPart("prt_target_text", "tool result") + const messages: TestMessage[] = [ + createMessage({ role: "user", modelID: "claude-opus-4-6-thinking" }, [createTextPart("prt_user_text", "continue")]), + createMessage({ role: "assistant", id: "msg_prev" }, [previousThinkingPart, createTextPart("prt_prev_text", "done")]), + createMessage({ role: "assistant", id: "msg_target" }, [targetTextPart]), + ] + + await Reflect.apply(transform, undefined, [{}, { messages }]) + + expect(messages[2]?.parts[0]).toBe(previousThinkingPart) + expect(messages[2]?.parts).toEqual([previousThinkingPart, targetTextPart]) + }) + + test("skips injection when no signed Anthropic thinking part exists in history", async () => { + const transform = Reflect.get(createThinkingBlockValidatorHook(), "experimental.chat.messages.transform") + expect(typeof transform).toBe("function") + + const targetTextPart = createTextPart("prt_target_text", "tool result") + const messages: TestMessage[] = [ + createMessage({ role: "user", modelID: "claude-opus-4-6-thinking" }, [createTextPart("prt_user_text", "continue")]), + createMessage({ role: "assistant", id: "msg_prev" }, [{ type: "reasoning", id: "prt_reason", text: "gpt reasoning" }]), + createMessage({ role: "assistant", id: "msg_target" }, [targetTextPart]), + ] + + await Reflect.apply(transform, undefined, [{}, { messages }]) + + expect(messages[2]?.parts).toEqual([targetTextPart]) + }) + + test("does not inject when the assistant message already starts with redacted thinking", async () => { + const transform = Reflect.get(createThinkingBlockValidatorHook(), "experimental.chat.messages.transform") + expect(typeof transform).toBe("function") + + const existingThinkingPart = createRedactedThinkingPart("prt_redacted", "sig_redacted") + const targetTextPart = createTextPart("prt_target_text", "tool result") + const messages: TestMessage[] = [ + createMessage({ role: "user", modelID: "claude-opus-4-6-thinking" }, [createTextPart("prt_user_text", "continue")]), + createMessage({ role: "assistant", id: "msg_target" }, [existingThinkingPart, targetTextPart]), + ] + + await Reflect.apply(transform, undefined, [{}, { messages }]) + + expect(messages[1]?.parts).toEqual([existingThinkingPart, targetTextPart]) + }) + + test("skips processing for models without extended thinking", async () => { + const transform = Reflect.get(createThinkingBlockValidatorHook(), "experimental.chat.messages.transform") + expect(typeof transform).toBe("function") + + const previousThinkingPart = createSignedThinkingPart("prt_prev_signed", "prior reasoning", "sig_prev") + const targetTextPart = createTextPart("prt_target_text", "tool result") + const messages: TestMessage[] = [ + createMessage({ role: "user", modelID: "gpt-5.4" }, [createTextPart("prt_user_text", "continue")]), + createMessage({ role: "assistant", id: "msg_prev" }, [previousThinkingPart]), + createMessage({ role: "assistant", id: "msg_target" }, [targetTextPart]), + ] + + await Reflect.apply(transform, undefined, [{}, { messages }]) + + expect(messages[2]?.parts).toEqual([targetTextPart]) + }) +}) + +export {} diff --git a/src/hooks/thinking-block-validator/hook.ts b/src/hooks/thinking-block-validator/hook.ts index 18a4d3e0b..67b6c12c0 100644 --- a/src/hooks/thinking-block-validator/hook.ts +++ b/src/hooks/thinking-block-validator/hook.ts @@ -21,16 +21,9 @@ interface MessageWithParts { parts: Part[] } -interface ThinkingPart { - thinking?: string - text?: string -} - -interface MessageInfoExtended { - id: string - role: string - sessionID?: string - modelID?: string +type SignedThinkingPart = Part & { + type: "thinking" | "redacted_thinking" + signature: string } type MessagesTransformHook = { @@ -83,57 +76,45 @@ function startsWithThinkingBlock(parts: Part[]): boolean { const firstPart = parts[0] const type = firstPart.type as string - return type === "thinking" || type === "reasoning" + return type === "thinking" || type === "redacted_thinking" || type === "reasoning" } -/** - * Find the most recent thinking content from previous assistant messages - */ -function findPreviousThinkingContent( +function isSignedThinkingPart(part: Part): part is SignedThinkingPart { + const type = part.type as string + if (type !== "thinking" && type !== "redacted_thinking") { + return false + } + + const signature = (part as { signature?: unknown }).signature + return typeof signature === "string" && signature.length > 0 +} + +function findPreviousThinkingPart( messages: MessageWithParts[], currentIndex: number -): string { +): SignedThinkingPart | null { // Search backwards from current message for (let i = currentIndex - 1; i >= 0; i--) { const msg = messages[i] if (msg.info.role !== "assistant") continue - // Look for thinking parts if (!msg.parts) continue for (const part of msg.parts) { - const type = part.type as string - if (type === "thinking" || type === "reasoning") { - const thinking = (part as unknown as ThinkingPart).thinking || (part as unknown as ThinkingPart).text - if (thinking && typeof thinking === "string" && thinking.trim().length > 0) { - return thinking - } + if (isSignedThinkingPart(part)) { + return part } } } - return "" + return null } -/** - * Prepend a thinking block to a message's parts array - */ -function prependThinkingBlock(message: MessageWithParts, thinkingContent: string): void { +function prependThinkingBlock(message: MessageWithParts, thinkingPart: SignedThinkingPart): void { if (!message.parts) { message.parts = [] } - // Create synthetic thinking part - const thinkingPart = { - type: "thinking" as const, - id: `prt_0000000000_synthetic_thinking`, - sessionID: (message.info as unknown as MessageInfoExtended).sessionID || "", - messageID: message.info.id, - thinking: thinkingContent, - synthetic: true, - } - - // Prepend to parts array - message.parts.unshift(thinkingPart as unknown as Part) + message.parts.unshift(thinkingPart) } /** @@ -150,7 +131,8 @@ export function createThinkingBlockValidatorHook(): MessagesTransformHook { // Get the model info from the last user message const lastUserMessage = messages.findLast(m => m.info.role === "user") - const modelID = (lastUserMessage?.info as unknown as MessageInfoExtended)?.modelID || "" + const modelIDValue = (lastUserMessage?.info as { modelID?: unknown } | undefined)?.modelID + const modelID = typeof modelIDValue === "string" ? modelIDValue : "" // Only process if extended thinking might be enabled if (!isExtendedThinkingModel(modelID)) { @@ -166,13 +148,12 @@ export function createThinkingBlockValidatorHook(): MessagesTransformHook { // Check if message has content parts but doesn't start with thinking if (hasContentParts(msg.parts) && !startsWithThinkingBlock(msg.parts)) { - // Find thinking content from previous turns - const previousThinking = findPreviousThinkingContent(messages, i) + const previousThinkingPart = findPreviousThinkingPart(messages, i) + if (!previousThinkingPart) { + continue + } - // Prepend thinking block with content from previous turn or placeholder - const thinkingContent = previousThinking || "[Continuing from previous reasoning]" - - prependThinkingBlock(msg, thinkingContent) + prependThinkingBlock(msg, previousThinkingPart) } } },