diff --git a/.changeset/calm-thinkers-wait.md b/.changeset/calm-thinkers-wait.md new file mode 100644 index 000000000..253d7c7cd --- /dev/null +++ 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 index 8d47ba320..5dbc6b268 100644 --- 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 { 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 index f037a396d..92dd1e54d 100644 --- 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 index 27204a5cd..aa695e745 100644 --- 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 | string + } + let hasToolResult = false + let messages: Array = [] try { const body = JSON.parse(bodyText) as { - messages?: Array<{ - role: string - content?: Array<{ type: string }> | string - }> + messages?: Array } - 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 index 4359138e9..3d8aa4506 100644 --- a/testing/e2e/src/routes/api.anthropic-bug-test.ts +++ b/testing/e2e/src/routes/api.anthropic-bug-test.ts @@ -8,6 +8,7 @@ import { import { createAnthropicChat } from '@tanstack/ai-anthropic' import { webFetchTool } from '@tanstack/ai-anthropic/tools' import { z } from 'zod' +import type { UIMessage } from '@tanstack/ai' const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' const DUMMY_KEY = 'sk-e2e-test-dummy-key' @@ -28,11 +29,108 @@ const DUMMY_KEY = 'sk-e2e-test-dummy-key' export const Route = createFileRoute('/api/anthropic-bug-test')({ server: { handlers: { - POST: async () => { + POST: async ({ request }) => { const adapter = createAnthropicChat('claude-sonnet-4-5', DUMMY_KEY, { baseURL: `${LLMOCK_DEFAULT_BASE}/anthropic-bug-test`, }) + if ( + new URL(request.url).searchParams.get('case') === 'thinking-order' + ) { + const messages: Array = [ + { + id: 'initial-user-message', + role: 'user', + parts: [{ type: 'text', content: 'Research top AI companies' }], + }, + { + 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: { + 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', + }, + ], + }, + }, + }, + { + 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', + }, + ], + }, + { + id: 'follow-up-user-message', + role: 'user', + parts: [ + { + type: 'text', + content: '[thinking-tool-order] continue', + }, + ], + }, + ] + + const chunks: Array = [] + try { + for await (const chunk of chat({ + ...createChatOptions({ adapter }), + messages, + })) { + chunks.push(chunk) + } + } catch (error) { + return new Response( + JSON.stringify({ + chunks, + error: error instanceof Error ? error.message : String(error), + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ) + } + + return new Response(JSON.stringify({ chunks, error: null }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }) + } + // A trivial client tool the user might pair with webFetchTool(). The // server lookup just echoes back so the agent loop can settle into // the follow-up text response. diff --git a/testing/e2e/tests/anthropic-server-tool.spec.ts b/testing/e2e/tests/anthropic-server-tool.spec.ts index bd85451a1..bf4c516da 100644 --- a/testing/e2e/tests/anthropic-server-tool.spec.ts +++ b/testing/e2e/tests/anthropic-server-tool.spec.ts @@ -86,3 +86,24 @@ test.describe('anthropic — webFetchTool() streaming (#604)', () => { expect(chunks.some((c) => c.type === 'RUN_FINISHED')).toBe(true) }) }) + +test.describe('anthropic — signed thinking replay (#910)', () => { + test('keeps thinking blocks on the correct side of server tool calls', async ({ + request, + }) => { + const res = await request.post( + '/api/anthropic-bug-test?case=thinking-order', + ) + expect(res.ok()).toBe(true) + const { chunks, error } = (await res.json()) as { + chunks: Array> + error: string | null + } + + // The mounted Anthropic endpoint rejects the request unless the replayed + // content blocks are thinking(A) → server tool/result → thinking(B) → + // local tool. This mirrors Anthropic's signed-thinking validation. + expect(error).toBeNull() + expect(chunks.some((chunk) => chunk.type === 'RUN_FINISHED')).toBe(true) + }) +})