Merge pull request #582 from luoxk/fix/assistant-tool-use-text-interleave

fix: reorder assistant tool_use blocks before sending history

src/utils/messages.ts

@@ -2310,6 +2310,21 @@ export function normalizeMessagesForAPI(
   // mismatched thinking block signatures cause API 400 errors.
   const withFilteredOrphans = filterOrphanedThinkingOnlyMessages(relocated)
 
+  // Reorder assistant content so any tool_use blocks form a contiguous run.
+  // mergeAssistantMessages also reorders, but this pass additionally protects
+  // single-message cases (e.g. sessions resumed from disk that were persisted
+  // before this fix landed). See reorderAssistantToolUseBlocks for the
+  // Bedrock validation rationale.
+  const withReorderedToolUse = withFilteredOrphans.map(msg => {
+    if (msg.type !== 'assistant') return msg
+    const reordered = reorderAssistantToolUseBlocks(msg.message.content)
+    if (reordered === msg.message.content) return msg
+    return {
+      ...msg,
+      message: { ...msg.message, content: reordered },
+    }
+  })
+
   // Order matters: strip trailing thinking first, THEN filter whitespace-only
   // messages. The reverse order has a bug: a message like [text("\n\n"), thinking("...")]
   // survives the whitespace filter (has a non-text block), then thinking stripping
@@ -2319,7 +2334,7 @@ export function normalizeMessagesForAPI(
   // conditions a prior pass was meant to handle. Consider unifying into a single
   // pass that cleans content, then validates in one shot.
   const withFilteredThinking =
-    filterTrailingThinkingFromLastAssistant(withFilteredOrphans)
+    filterTrailingThinkingFromLastAssistant(withReorderedToolUse)
   const withFilteredWhitespace =
     filterWhitespaceOnlyAssistantMessages(withFilteredThinking)
   const withNonEmpty = ensureNonEmptyAssistantContent(withFilteredWhitespace)
@@ -2394,11 +2409,85 @@ export function mergeAssistantMessages(
     ...a,
     message: {
       ...a.message,
-      content: [...a.message.content, ...b.message.content],
+      // Reorder so that any tool_use blocks introduced by `b` don't end up
+      // separated from `a`'s tool_use blocks by intervening text. Without
+      // this, Bedrock's strict history validation rejects "tool_use ids …
+      // without tool_result blocks immediately after" because text inside
+      // the tool_use cluster makes the earlier tool_use blocks no longer
+      // count as "trailing" — only the final tool_use is paired with the
+      // next message's tool_results.
+      content: reorderAssistantToolUseBlocks([
+        ...a.message.content,
+        ...b.message.content,
+      ]),
     },
   }
 }
 
+/**
+ * Reorder an assistant message's content so that all `tool_use` blocks form
+ * a contiguous run. Any non-`tool_use` blocks that the model emitted in the
+ * middle of that run (typically `text`) are pushed to the position right
+ * after the last `tool_use`.
+ *
+ * Why: Anthropic's history validation (and Bedrock's stricter copy of it)
+ * requires every `tool_use` block to be paired with a matching `tool_result`
+ * in the next message. The validator only treats the *trailing* run of
+ * `tool_use` blocks as "needing tool_results next", so when the model
+ * streams `text` between tool calls — e.g. `tu1, tu2, tu3, tu4, text, tu5` —
+ * only `tu5` is considered trailing on the next request, and `tu1..tu4` are
+ * reported as missing tool_results, producing a 400 on the *next* turn even
+ * though the previous turn returned all 5 tool_results correctly.
+ *
+ * Block-type policy:
+ * - `thinking` and `redacted_thinking` keep their relative positions
+ *   (signatures are position-sensitive within a turn).
+ * - `tool_use` blocks become contiguous, in their original id order
+ *   (preserves any caller logic that pairs tool_results by index).
+ * - Non-`tool_use`, non-thinking blocks that were interleaved between
+ *   tool_use blocks are moved to immediately after the tool_use cluster.
+ *
+ * No blocks are dropped. The function is a no-op when there are fewer than
+ * two `tool_use` blocks or when the existing tool_use run is already
+ * contiguous.
+ */
+export function reorderAssistantToolUseBlocks<T extends { type: string }>(
+  content: T[],
+): T[] {
+  if (content.length < 2) return content
+
+  const toolUseIndices: number[] = []
+  for (let i = 0; i < content.length; i++) {
+    if (content[i]!.type === 'tool_use') toolUseIndices.push(i)
+  }
+  if (toolUseIndices.length < 2) return content
+
+  const first = toolUseIndices[0]!
+  const last = toolUseIndices[toolUseIndices.length - 1]!
+
+  let hasInterleaved = false
+  for (let i = first; i <= last; i++) {
+    if (content[i]!.type !== 'tool_use') {
+      hasInterleaved = true
+      break
+    }
+  }
+  if (!hasInterleaved) return content
+
+  const head = content.slice(0, first)
+  const window = content.slice(first, last + 1)
+  const tail = content.slice(last + 1)
+
+  const tools: T[] = []
+  const displaced: T[] = []
+  for (const block of window) {
+    if (block.type === 'tool_use') tools.push(block)
+    else displaced.push(block)
+  }
+
+  return [...head, ...tools, ...displaced, ...tail]
+}
+
 function isToolResultMessage(msg: Message): boolean {
   if (msg.type !== 'user') {
     return false

tests/reorderAssistantToolUseBlocks.test.ts

@@ -0,0 +1,181 @@
+import { describe, expect, test } from "bun:test"
+import {
+  mergeAssistantMessages,
+  reorderAssistantToolUseBlocks,
+} from "../src/utils/messages.js"
+import type { AssistantMessage } from "../src/types/message.js"
+
+const tu = (id: string) => ({
+  type: "tool_use" as const,
+  id,
+  name: "TaskCreate",
+  input: {},
+})
+const text = (s: string) => ({
+  type: "text" as const,
+  text: s,
+  citations: [],
+})
+const thinking = (s: string) => ({
+  type: "thinking" as const,
+  thinking: s,
+  signature: "",
+})
+const redactedThinking = (data: string) => ({
+  type: "redacted_thinking" as const,
+  data,
+})
+
+describe("reorderAssistantToolUseBlocks", () => {
+  test("no-op for content with 0 tool_use", () => {
+    const content = [thinking("t"), text("hello")]
+    expect(reorderAssistantToolUseBlocks(content)).toBe(content)
+  })
+
+  test("no-op for content with 1 tool_use", () => {
+    const content = [thinking("t"), text("a"), tu("1"), text("b")]
+    expect(reorderAssistantToolUseBlocks(content)).toBe(content)
+  })
+
+  test("no-op when tool_use blocks are already contiguous", () => {
+    const content = [thinking("t"), tu("1"), tu("2"), tu("3"), text("after")]
+    expect(reorderAssistantToolUseBlocks(content)).toBe(content)
+  })
+
+  test("hoists interleaved text out of the tool_use cluster (Bedrock bug case)", () => {
+    // The exact pattern from CLIENT_TOOL_USE_BUG_REPORT.md messages[31].
+    const content = [
+      thinking("planning"),
+      tu("a"),
+      tu("b"),
+      tu("c"),
+      tu("d"),
+      text("brief explanation"),
+      tu("e"),
+      redactedThinking("zzz"),
+    ]
+    const out = reorderAssistantToolUseBlocks(content)
+    expect(out.map((b) => b.type)).toEqual([
+      "thinking",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "text",
+      "redacted_thinking",
+    ])
+    // tool_use ids preserved in original order — downstream tool_result
+    // pairing relies on this.
+    expect(
+      out
+        .filter((b): b is ReturnType<typeof tu> => b.type === "tool_use")
+        .map((b) => b.id),
+    ).toEqual(["a", "b", "c", "d", "e"])
+  })
+
+  test("preserves head/tail blocks outside the tool_use window", () => {
+    const content = [
+      thinking("head1"),
+      text("head2"),
+      tu("a"),
+      text("middle"),
+      tu("b"),
+      text("tail1"),
+      redactedThinking("tail2"),
+    ]
+    const out = reorderAssistantToolUseBlocks(content)
+    // head stays in place
+    expect(out[0]).toEqual(thinking("head1"))
+    expect(out[1]).toEqual(text("head2"))
+    // tool_uses become contiguous
+    expect(out[2]!.type).toBe("tool_use")
+    expect(out[3]!.type).toBe("tool_use")
+    // displaced text follows the tool_use cluster
+    expect(out[4]).toEqual(text("middle"))
+    // tail preserved
+    expect(out[5]).toEqual(text("tail1"))
+    expect(out[6]).toEqual(redactedThinking("tail2"))
+  })
+
+  test("does not drop any blocks", () => {
+    const content = [
+      tu("1"),
+      thinking("mid-thought"),
+      tu("2"),
+      text("explain"),
+      tu("3"),
+    ]
+    const out = reorderAssistantToolUseBlocks(content)
+    expect(out.length).toBe(content.length)
+    // every original block still present (by reference equality where applicable)
+    for (const block of content) {
+      expect(out).toContain(block)
+    }
+  })
+
+  test("is idempotent", () => {
+    const content = [
+      tu("a"),
+      tu("b"),
+      text("explain"),
+      tu("c"),
+    ]
+    const once = reorderAssistantToolUseBlocks(content)
+    const twice = reorderAssistantToolUseBlocks(once)
+    expect(twice).toBe(once)
+  })
+})
+
+describe("mergeAssistantMessages reordering", () => {
+  const makeAsst = (
+    id: string,
+    content: AssistantMessage["message"]["content"],
+  ): AssistantMessage =>
+    ({
+      type: "assistant",
+      uuid: `uuid-${id}` as AssistantMessage["uuid"],
+      timestamp: "2026-05-23T00:00:00.000Z",
+      message: {
+        id: `msg-${id}`,
+        role: "assistant",
+        type: "message",
+        model: "claude-opus-4-7",
+        content,
+        stop_reason: null,
+        stop_sequence: null,
+        usage: { input_tokens: 0, output_tokens: 0 } as never,
+      },
+    }) as unknown as AssistantMessage
+
+  test("reorders interleaved text after concat", () => {
+    // Simulates the streaming case: each content_block_stop produces its own
+    // AssistantMessage, then normalizeMessagesForAPI merges them by message.id.
+    const a = makeAsst("x", [thinking("t"), tu("a"), tu("b"), tu("c"), tu("d"), text("aside")])
+    const b = makeAsst("x", [tu("e"), redactedThinking("zzz")])
+    const merged = mergeAssistantMessages(a, b)
+    expect(merged.message.content.map((c) => c.type)).toEqual([
+      "thinking",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "text",
+      "redacted_thinking",
+    ])
+  })
+
+  test("is a no-op when concat result is already valid", () => {
+    const a = makeAsst("x", [thinking("t"), tu("a"), tu("b")])
+    const b = makeAsst("x", [tu("c"), text("after")])
+    const merged = mergeAssistantMessages(a, b)
+    expect(merged.message.content.map((c) => c.type)).toEqual([
+      "thinking",
+      "tool_use",
+      "tool_use",
+      "tool_use",
+      "text",
+    ])
+  })
+})