fix: preserve thinking order around provider tools

ekkoitac · ekkoitac · commit c81fc678e1b4 · 2026-07-14T16:57:56.000+08:00
diff --git a/.changeset/calm-thinkers-wait.md b/.changeset/calm-thinkers-wait.md
@@ -0,0 +1,5 @@
+---
+'@tanstack/ai': patch
+---
+
+Preserve signed thinking blocks relative to tool calls when converting assistant UI messages. A thinking block that follows a provider-executed tool now starts a new assistant segment instead of being replayed before that tool.
diff --git a/packages/ai/src/activities/chat/messages.ts b/packages/ai/src/activities/chat/messages.ts
@@ -1,3 +1,4 @@
+import { isProviderExecutedToolCall } from '../../utilities/provider-executed'
 import { normalizeToolResult } from '../../utilities/tool-result'
 import type { Message as AGUIMessage } from '@ag-ui/core'
 import type {
@@ -297,6 +298,12 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
 
       case 'thinking':
         if (part.content) {
+          // A thinking block after a tool call belongs to the next assistant
+          // segment. Provider-executed tools have no separate tool-result part
+          // to create this boundary for us.
+          if (current.toolCalls.some(isProviderExecutedToolCall)) {
+            flushSegment()
+          }
           pendingThinking.push({
             content: part.content,
             ...(part.signature && { signature: part.signature }),
diff --git a/packages/ai/tests/message-converters.test.ts b/packages/ai/tests/message-converters.test.ts
@@ -201,6 +201,208 @@ describe('Message Converters', () => {
       expect(result[0]?.thinking).toEqual([{ content: 'Let me think...' }])
     })
 
+    it('should preserve thinking order around provider-executed tool calls', () => {
+      const providerToolMetadata = {
+        providerExecuted: true,
+        anthropic: {
+          serverToolType: 'web_search',
+          resultBlockType: 'web_search_tool_result',
+          result: [
+            {
+              type: 'web_search_result',
+              title: 'Example result',
+              url: 'https://example.com',
+              encrypted_content: 'opaque-provider-payload',
+            },
+          ],
+        },
+      }
+      const uiMessage: UIMessage = {
+        id: 'assistant-message',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'thinking',
+            content: 'First signed thinking block',
+            signature: 'signature-1',
+          },
+          {
+            type: 'tool-call',
+            id: 'srvtoolu_search',
+            name: 'web_search',
+            arguments: '{"query":"top AI companies"}',
+            state: 'input-complete',
+            metadata: providerToolMetadata,
+          },
+          {
+            type: 'thinking',
+            content: 'Second signed thinking block',
+            signature: 'signature-2',
+          },
+          {
+            type: 'tool-call',
+            id: 'toolu_create_block',
+            name: 'createBlock',
+            arguments: '{"block":{"type":"entity-grid"}}',
+            state: 'complete',
+            output: { ok: true },
+          },
+          {
+            type: 'tool-result',
+            toolCallId: 'toolu_create_block',
+            content: '{"ok":true}',
+            state: 'complete',
+          },
+        ],
+      }
+
+      const result = convertMessagesToModelMessages([uiMessage])
+
+      expect(result).toEqual([
+        {
+          role: 'assistant',
+          content: null,
+          toolCalls: [
+            {
+              id: 'srvtoolu_search',
+              type: 'function',
+              function: {
+                name: 'web_search',
+                arguments: '{"query":"top AI companies"}',
+              },
+              metadata: providerToolMetadata,
+            },
+          ],
+          thinking: [
+            {
+              content: 'First signed thinking block',
+              signature: 'signature-1',
+            },
+          ],
+        },
+        {
+          role: 'assistant',
+          content: null,
+          toolCalls: [
+            {
+              id: 'toolu_create_block',
+              type: 'function',
+              function: {
+                name: 'createBlock',
+                arguments: '{"block":{"type":"entity-grid"}}',
+              },
+            },
+          ],
+          thinking: [
+            {
+              content: 'Second signed thinking block',
+              signature: 'signature-2',
+            },
+          ],
+        },
+        {
+          role: 'tool',
+          content: '{"ok":true}',
+          toolCallId: 'toolu_create_block',
+        },
+      ])
+    })
+
+    it('should keep local tool calls and results in the same assistant turn', () => {
+      const uiMessage: UIMessage = {
+        id: 'assistant-message',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'tool-call',
+            id: 'tool-call-a',
+            name: 'toolA',
+            arguments: '{"value":"a"}',
+            state: 'input-complete',
+          },
+          {
+            type: 'thinking',
+            content: 'Thinking between local tool calls',
+          },
+          {
+            type: 'tool-call',
+            id: 'tool-call-b',
+            name: 'toolB',
+            arguments: '{"value":"b"}',
+            state: 'input-complete',
+          },
+          {
+            type: 'tool-result',
+            toolCallId: 'tool-call-a',
+            content: '{"result":"a"}',
+            state: 'complete',
+          },
+          {
+            type: 'tool-result',
+            toolCallId: 'tool-call-b',
+            content: '{"result":"b"}',
+            state: 'complete',
+          },
+        ],
+      }
+
+      const modelMessages = uiMessageToModelMessages(uiMessage)
+
+      expect(modelMessages).toEqual([
+        {
+          role: 'assistant',
+          content: null,
+          toolCalls: [
+            {
+              id: 'tool-call-a',
+              type: 'function',
+              function: {
+                name: 'toolA',
+                arguments: '{"value":"a"}',
+              },
+            },
+            {
+              id: 'tool-call-b',
+              type: 'function',
+              function: {
+                name: 'toolB',
+                arguments: '{"value":"b"}',
+              },
+            },
+          ],
+          thinking: [{ content: 'Thinking between local tool calls' }],
+        },
+        {
+          role: 'tool',
+          content: '{"result":"a"}',
+          toolCallId: 'tool-call-a',
+        },
+        {
+          role: 'tool',
+          content: '{"result":"b"}',
+          toolCallId: 'tool-call-b',
+        },
+      ])
+
+      const roundTripped = modelMessagesToUIMessages(modelMessages)
+
+      expect(roundTripped).toHaveLength(1)
+      expect(roundTripped[0]?.parts).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({ type: 'tool-call', id: 'tool-call-a' }),
+          expect.objectContaining({ type: 'tool-call', id: 'tool-call-b' }),
+          expect.objectContaining({
+            type: 'tool-result',
+            toolCallId: 'tool-call-a',
+          }),
+          expect.objectContaining({
+            type: 'tool-result',
+            toolCallId: 'tool-call-b',
+          }),
+        ]),
+      )
+    })
+
     it('should skip system messages', () => {
       const uiMessage: UIMessage = {
         id: 'msg-1',
diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts
@@ -504,11 +504,9 @@ function geminiOmniVideoMount(): Mountable {
 }
 
 /**
- * Mounts a Claude-shaped SSE response that includes a client `tool_use` block
- * followed by a `web_fetch` `server_tool_use` block, plus its
- * `web_fetch_tool_result`. Reproduces the streaming scenario from issue #604
- * — the adapter must not let server-tool `input_json_delta`s leak into the
- * prior client tool's input buffer.
+ * Mounts Claude-shaped responses for server-tool regression coverage. The
+ * streaming response reproduces issue #604, while the request validator
+ * reproduces Anthropic's signed-thinking ordering check from issue #910.
  *
  * The first turn returns the mixed tool_use + server_tool_use response so the
  * adapter can dispatch the client tool. The second turn (after the client
@@ -529,15 +527,25 @@ function anthropicServerToolBugMount(): Mountable {
       }
 
       const bodyText = await readBody(req)
+      type RequestBlock = {
+        type: string
+        text?: string
+        thinking?: string
+        signature?: string
+      }
+      type RequestMessage = {
+        role: string
+        content?: Array<RequestBlock> | string
+      }
+
       let hasToolResult = false
+      let messages: Array<RequestMessage> = []
       try {
         const body = JSON.parse(bodyText) as {
-          messages?: Array<{
-            role: string
-            content?: Array<{ type: string }> | string
-          }>
+          messages?: Array<RequestMessage>
         }
-        hasToolResult = (body.messages ?? []).some(
+        messages = body.messages ?? []
+        hasToolResult = messages.some(
           (m) =>
             Array.isArray(m.content) &&
             m.content.some((c) => c.type === 'tool_result'),
@@ -546,6 +554,74 @@ function anthropicServerToolBugMount(): Mountable {
         // Malformed body — fall through and emit the first-turn stream.
       }
 
+      const thinkingOrderMarker = '[thinking-tool-order] continue'
+      const hasThinkingOrderMarker = messages.some((message) => {
+        if (message.role !== 'user') return false
+        if (typeof message.content === 'string') {
+          return message.content === thinkingOrderMarker
+        }
+        return message.content?.some(
+          (block) =>
+            block.type === 'text' && block.text === thinkingOrderMarker,
+        )
+      })
+      const hasThinkingOrderSignature = messages.some(
+        (message) =>
+          message.role === 'assistant' &&
+          Array.isArray(message.content) &&
+          message.content.some(
+            (block) =>
+              block.type === 'thinking' &&
+              (block.signature === 'signature-1' ||
+                block.signature === 'signature-2'),
+          ),
+      )
+      const validatesThinkingOrder =
+        hasThinkingOrderMarker || hasThinkingOrderSignature
+
+      if (validatesThinkingOrder) {
+        const assistantMessage = messages.find(
+          (message) =>
+            message.role === 'assistant' && Array.isArray(message.content),
+        )
+        const blocks = Array.isArray(assistantMessage?.content)
+          ? assistantMessage.content
+          : []
+        const blockTypes = blocks.map((block) => block.type)
+        const expectedBlockTypes = [
+          'thinking',
+          'server_tool_use',
+          'web_search_tool_result',
+          'thinking',
+          'tool_use',
+        ]
+        const thinkingBlocks = blocks.filter(
+          (block) => block.type === 'thinking',
+        )
+        const preservesOrder =
+          JSON.stringify(blockTypes) === JSON.stringify(expectedBlockTypes) &&
+          thinkingBlocks[0]?.thinking === 'First signed thinking block' &&
+          thinkingBlocks[0]?.signature === 'signature-1' &&
+          thinkingBlocks[1]?.thinking === 'Second signed thinking block' &&
+          thinkingBlocks[1]?.signature === 'signature-2'
+
+        if (!preservesOrder) {
+          res.statusCode = 400
+          res.setHeader('Content-Type', 'application/json')
+          res.end(
+            JSON.stringify({
+              type: 'error',
+              error: {
+                type: 'invalid_request_error',
+                message:
+                  'Signed thinking blocks were reordered relative to tool-use blocks',
+              },
+            }),
+          )
+          return true
+        }
+      }
+
       const events = hasToolResult
         ? buildFollowUpEvents()
         : buildToolPlusServerToolEvents()
diff --git a/testing/e2e/src/routes/api.anthropic-bug-test.ts b/testing/e2e/src/routes/api.anthropic-bug-test.ts
diff --git a/testing/e2e/tests/anthropic-server-tool.spec.ts b/testing/e2e/tests/anthropic-server-tool.spec.ts

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'@tanstack/ai': patch
 +---
++
 +Preserve signed thinking blocks relative to tool calls when converting assistant UI messages. A thinking block that follows a provider-executed tool now starts a new assistant segment instead of being replayed before that tool.