Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/calm-thinkers-wait.md
Original file line number Diff 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.
7 changes: 7 additions & 0 deletions packages/ai/src/activities/chat/messages.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down Expand Up @@ -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 }),
Expand Down
202 changes: 202 additions & 0 deletions packages/ai/tests/message-converters.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
96 changes: 86 additions & 10 deletions testing/e2e/global-setup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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'),
Expand All @@ -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()
Expand Down
Loading