Merge pull request #4477 from code-yeongyu/fix-4374-dual-thinking-streams

fix(thinking-block-validator): prevent dual thinking streams after question tool response (#4374)

src/hooks/thinking-block-validator/hook.ts

@@ -78,14 +78,15 @@ function hasContentParts(parts: Part[]): boolean {
 }
 
 /**
- * Check if a message starts with a thinking/reasoning block
+ * Check if a message already carries a thinking/reasoning block anywhere.
  */
-function startsWithThinkingBlock(parts: Part[]): boolean {
+function hasThinkingBlock(parts: Part[]): boolean {
   if (!parts || parts.length === 0) return false
 
-  const firstPart = parts[0]
+  return parts.some((part) => {
-  const type = firstPart.type as string
+    const type = part.type as string
-  return type === "thinking" || type === "redacted_thinking" || type === "reasoning"
+    return type === "thinking" || type === "redacted_thinking" || type === "reasoning"
+  })
 }
 
 /**
@@ -160,8 +161,8 @@ export function createThinkingBlockValidatorHook(): MessagesTransformHook {
         // Only check assistant messages
         if (msg.info.role !== "assistant") continue
 
-        // Check if message has content parts but doesn't start with thinking
+        // Check if message has content parts but no thinking block yet.
-        if (hasContentParts(msg.parts) && !startsWithThinkingBlock(msg.parts)) {
+        if (hasContentParts(msg.parts) && !hasThinkingBlock(msg.parts)) {
           // Find the most recent real thinking part (with valid signature) from
           // previous turns.  If none exists we cannot safely inject a thinking
           // block - a synthetic block without a signature would cause the API

src/plugin/messages-transform-thinking-block.test.ts

@@ -0,0 +1,100 @@
+declare const describe: (name: string, fn: () => void) => void
+declare const it: (name: string, fn: () => void | Promise<void>) => void
+declare const expect: <T>(value: T) => {
+  toBe(expected: T): void
+}
+
+import type { CreatedHooks } from "../create-hooks"
+import { createThinkingBlockValidatorHook } from "../hooks/thinking-block-validator/hook"
+import { createToolPairValidatorHook } from "../hooks/tool-pair-validator/hook"
+import { createMessagesTransformHandler } from "./messages-transform"
+
+type TestPart = {
+  type: string
+  id?: string
+  toolUseId?: string
+  tool_use_id?: string
+  name?: string
+  content?: Array<{ type: "text"; text: string }>
+  text?: string
+  thinking?: string
+  signature?: string
+}
+
+type TestMessage = {
+  info: {
+    role: "assistant" | "user"
+    id?: string
+    sessionID?: string
+  }
+  parts: TestPart[]
+}
+
+function createTestHooks(): CreatedHooks {
+  return {
+    thinkingBlockValidator: createThinkingBlockValidatorHook(),
+    toolPairValidator: createToolPairValidatorHook(),
+  } as CreatedHooks
+}
+
+async function runMessagesTransform(messages: TestMessage[]): Promise<void> {
+  const handler = createMessagesTransformHandler({ hooks: createTestHooks() })
+  await handler({}, { messages: messages as never })
+}
+
+function countThinkingParts(parts: TestPart[]): number {
+  return parts.filter((part) => part.type === "thinking" || part.type === "redacted_thinking").length
+}
+
+describe("messages transform thinking block integration", () => {
+  it("#given a question tool answer and a resumed assistant turn with existing thinking #when messages transform runs #then it keeps one thinking block in that assistant turn", async () => {
+    //#given
+    const thinkingBeforeQuestion: TestPart = {
+      type: "thinking",
+      thinking: "ask a clarifying question",
+      signature: "sig-before-question",
+    }
+    const thinkingAfterAnswer: TestPart = {
+      type: "thinking",
+      thinking: "continue after answer",
+      signature: "sig-after-answer",
+    }
+    const messages = [
+      {
+        info: { id: "msg_user_prompt", role: "user", sessionID: "ses_question_thinking" },
+        parts: [{ type: "text", text: "think, then ask a question" }],
+      },
+      {
+        info: { id: "msg_question", role: "assistant", sessionID: "ses_question_thinking" },
+        parts: [thinkingBeforeQuestion, { type: "tool_use", id: "toolu_question", name: "question" }],
+      },
+      {
+        info: { id: "msg_question_answer", role: "user", sessionID: "ses_question_thinking" },
+        parts: [
+          {
+            type: "tool_result",
+            toolUseId: "toolu_question",
+            tool_use_id: "toolu_question",
+            content: [{ type: "text", text: "answer" }],
+          },
+        ],
+      },
+      {
+        info: { id: "msg_resumed", role: "assistant", sessionID: "ses_question_thinking" },
+        parts: [
+          { type: "text", text: "resuming" },
+          thinkingAfterAnswer,
+          { type: "tool_use", id: "toolu_after_answer", name: "bash" },
+        ],
+      },
+    ] satisfies TestMessage[]
+
+    //#when
+    await runMessagesTransform(messages)
+
+    //#then
+    const resumedMessage = messages.find((message) => message.info.id === "msg_resumed")
+    expect(resumedMessage?.parts[1]).toBe(thinkingAfterAnswer)
+    expect(countThinkingParts(resumedMessage?.parts ?? [])).toBe(1)
+  })
+})