fix(ai, ai-anthropic): thinking blocks missing on turn 2+ in tool loops

.changeset/thinking-blocks-per-step.md

@@ -0,0 +1,28 @@
+---
+'@tanstack/ai': patch
+'@tanstack/ai-anthropic': patch
+'@tanstack/ai-client': patch
+---
+
+**Fix thinking blocks getting merged across steps and lost on turn 2+ of Anthropic tool loops.**
+
+Each thinking step emitted by the adapter now produces its own `ThinkingPart` on the `UIMessage` instead of being merged into a single part, and thinking content + Anthropic signatures are preserved in server-side message history so multi-turn tool flows with extended thinking work correctly.
+
+`@tanstack/ai`:
+
+- `ThinkingPart` gains optional `stepId` and `signature` fields.
+- `ModelMessage` gains an optional `thinking?: Array<{ content; signature? }>` field so prior thinking can be replayed in subsequent turns.
+- `StepFinishedEvent` gains an optional `signature` field for provider-supplied thinking signatures.
+- `StreamProcessor` tracks thinking per-step via `stepId` and keeps step ordering. `getState().thinking` / `getResult().thinking` concatenate step contents in order.
+- The `onThinkingUpdate` callback on `StreamProcessorEvents` now receives `(messageId, stepId, content)` — consumers implementing it directly must add the `stepId` parameter.
+- `TextEngine` accumulates thinking + signatures per iteration and includes them in assistant messages with tool calls so the next turn can replay them.
+
+`@tanstack/ai-anthropic`:
+
+- Captures `signature_delta` stream events and emits the final `STEP_FINISHED` with the signature on `content_block_stop`.
+- Includes thinking blocks with signatures in `formatMessages` for multi-turn history.
+- Passes `betas: ['interleaved-thinking-2025-05-14']` to the `beta.messages.create` call site when a thinking budget is configured. The beta flag is scoped to the streaming path only, so `structuredOutput` (which uses the non-beta `messages.create` endpoint) is unaffected.
+
+`@tanstack/ai-client`:
+
+- `ChatClient`'s internal `onThinkingUpdate` wiring is updated for the new `stepId` parameter.

packages/typescript/ai-anthropic/src/adapters/text.ts

@@ -26,6 +26,7 @@ import type {
   URLPDFSource,
 } from '@anthropic-ai/sdk/resources/messages'
 import type Anthropic_SDK from '@anthropic-ai/sdk'
+import type { AnthropicBeta } from '@anthropic-ai/sdk/resources/beta/beta'
 import type {
   ContentPart,
   Modality,
@@ -121,8 +122,22 @@ export class AnthropicTextAdapter<
     try {
       const requestParams = this.mapCommonOptionsToAnthropic(options)
 
+      // Interleaved thinking is only supported on the beta messages endpoint,
+      // so the `betas` flag is attached here rather than in the shared mapper
+      // (structuredOutput uses the non-beta endpoint which rejects `betas`).
+      const modelOptions = options.modelOptions as
+        | InternalTextProviderOptions
+        | undefined
+      const useInterleavedThinking =
+        modelOptions?.thinking?.type === 'enabled' &&
+        typeof modelOptions.thinking.budget_tokens === 'number' &&
+        modelOptions.thinking.budget_tokens > 0
+      const betas: Array<AnthropicBeta> | undefined = useInterleavedThinking
+        ? ['interleaved-thinking-2025-05-14']
+        : undefined
+
       const stream = await this.client.beta.messages.create(
-        { ...requestParams, stream: true },
+        { ...requestParams, stream: true, ...(betas && { betas }) },
         {
           signal: options.request?.signal,
           headers: options.request?.headers,
@@ -389,6 +404,18 @@ export class AnthropicTextAdapter<
       if (role === 'assistant' && message.toolCalls?.length) {
         const contentBlocks: AnthropicContentBlocks = []
 
+        if (message.thinking?.length) {
+          for (const thinking of message.thinking) {
+            if (thinking.signature) {
+              contentBlocks.push({
+                type: 'thinking',
+                thinking: thinking.content,
+                signature: thinking.signature,
+              } as unknown as AnthropicContentBlock)
+            }
+          }
+        }
+
         if (message.content) {
           const content =
             typeof message.content === 'string' ? message.content : ''
@@ -528,6 +555,7 @@ export class AnthropicTextAdapter<
   ): AsyncIterable<StreamChunk> {
     let accumulatedContent = ''
     let accumulatedThinking = ''
+    let accumulatedSignature = ''
     const timestamp = Date.now()
     const toolCallsMap = new Map<
       number,
@@ -570,6 +598,7 @@ export class AnthropicTextAdapter<
             })
           } else if (event.content_block.type === 'thinking') {
             accumulatedThinking = ''
+            accumulatedSignature = ''
             // Emit STEP_STARTED for thinking
             stepId = genId()
             yield {
@@ -615,6 +644,11 @@ export class AnthropicTextAdapter<
               delta,
               content: accumulatedThinking,
             }
+          } else if (
+            (event.delta as { type: string }).type === 'signature_delta'
+          ) {
+            accumulatedSignature +=
+              (event.delta as { signature: string }).signature || ''
           } else if (event.delta.type === 'input_json_delta') {
             const existing = toolCallsMap.get(currentToolIndex)
             if (existing) {
@@ -644,7 +678,20 @@ export class AnthropicTextAdapter<
             }
           }
         } else if (event.type === 'content_block_stop') {
-          if (currentBlockType === 'tool_use') {
+          if (currentBlockType === 'thinking') {
+            // Emit signature so it can be replayed in multi-turn context
+            if (accumulatedSignature && stepId) {
+              yield {
+                type: 'STEP_FINISHED',
+                stepId,
+                model,
+                timestamp,
+                delta: '',
+                content: accumulatedThinking,
+                signature: accumulatedSignature,
+              }
+            }
+          } else if (currentBlockType === 'tool_use') {
             const existing = toolCallsMap.get(currentToolIndex)
             if (existing) {
               // If tool call wasn't started yet (no args), start it now

packages/typescript/ai-client/src/chat-client.ts

@@ -153,7 +153,11 @@ export class ChatClient {
             this.events.textUpdated(this.currentStreamId, messageId, content)
           }
         },
-        onThinkingUpdate: (messageId: string, content: string) => {
+        onThinkingUpdate: (
+          messageId: string,
+          _stepId: string,
+          content: string,
+        ) => {
           // Emit thinking update to devtools
           if (this.currentStreamId) {
             this.events.thinkingUpdated(

packages/typescript/ai/src/activities/chat/index.ts

@@ -253,6 +253,10 @@ class TextEngine<
   private totalChunkCount = 0
   private currentMessageId: string | null = null
   private accumulatedContent = ''
+  private accumulatedThinking: Array<{ content: string; signature?: string }> =
+    []
+  private currentThinkingContent = ''
+  private currentThinkingSignature = ''
   private eventOptions?: Record<string, unknown>
   private eventToolNames?: Array<string>
   private finishedEvent: RunFinishedEvent | null = null
@@ -494,6 +498,9 @@ class TextEngine<
   private async beginIteration(): Promise<void> {
     this.currentMessageId = this.createId('msg')
     this.accumulatedContent = ''
+    this.accumulatedThinking = []
+    this.currentThinkingContent = ''
+    this.currentThinkingSignature = ''
     this.finishedEvent = null
 
     // Update mutable context fields
@@ -585,12 +592,15 @@ class TextEngine<
       case 'RUN_ERROR':
         this.handleRunErrorEvent(chunk)
         break
+      case 'STEP_STARTED':
+        this.handleStepStartedEvent()
+        break
       case 'STEP_FINISHED':
         this.handleStepFinishedEvent(chunk)
         break
 
       default:
-        // RUN_STARTED, TEXT_MESSAGE_START, TEXT_MESSAGE_END, STEP_STARTED,
+        // RUN_STARTED, TEXT_MESSAGE_START, TEXT_MESSAGE_END,
         // STATE_SNAPSHOT, STATE_DELTA, CUSTOM
         // - no special handling needed in chat activity
         break
@@ -633,10 +643,32 @@ class TextEngine<
     this.earlyTermination = true
   }
 
+  private finalizeCurrentThinkingStep(): void {
+    if (this.currentThinkingContent) {
+      this.accumulatedThinking.push({
+        content: this.currentThinkingContent,
+        ...(this.currentThinkingSignature && {
+          signature: this.currentThinkingSignature,
+        }),
+      })
+      this.currentThinkingContent = ''
+      this.currentThinkingSignature = ''
+    }
+  }
+
+  private handleStepStartedEvent(): void {
+    this.finalizeCurrentThinkingStep()
+  }
+
   private handleStepFinishedEvent(
-    _chunk: Extract<StreamChunk, { type: 'STEP_FINISHED' }>,
+    chunk: Extract<StreamChunk, { type: 'STEP_FINISHED' }>,
   ): void {
-    // State tracking for STEP_FINISHED is handled by middleware
+    if (chunk.delta) {
+      this.currentThinkingContent += chunk.delta
+    }
+    if (chunk.signature) {
+      this.currentThinkingSignature = chunk.signature
+    }
   }
 
   private async *checkForPendingToolCalls(): AsyncGenerator<
@@ -939,12 +971,17 @@ class TextEngine<
   }
 
   private addAssistantToolCallMessage(toolCalls: Array<ToolCall>): void {
+    this.finalizeCurrentThinkingStep()
+
     this.messages = [
       ...this.messages,
       {
         role: 'assistant',
         content: this.accumulatedContent || null,
         toolCalls,
+        ...(this.accumulatedThinking.length > 0 && {
+          thinking: this.accumulatedThinking,
+        }),
       },
     ]
   }

packages/typescript/ai/src/activities/chat/messages.ts

@@ -165,6 +165,7 @@ function isToolCallIncluded(part: ToolCallPart): boolean {
 function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
   const messageList: Array<ModelMessage> = []
   let current = createSegment()
+  let pendingThinking: Array<{ content: string; signature?: string }> = []
 
   // Track emitted tool result IDs to avoid duplicates.
   // A tool call can have BOTH an explicit tool-result part AND an output
@@ -181,7 +182,9 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
         role: 'assistant',
         content,
         ...(hasToolCalls && { toolCalls: current.toolCalls }),
+        ...(pendingThinking.length > 0 && { thinking: pendingThinking }),
       })
+      pendingThinking = []
     }
     current = createSegment()
   }
@@ -227,7 +230,15 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
         }
         break
 
-      // thinking parts are skipped - they're UI-only
+      case 'thinking':
+        if (part.content) {
+          pendingThinking.push({
+            content: part.content,
+            ...(part.signature && { signature: part.signature }),
+          })
+        }
+        break
+
       default:
         break
     }

packages/typescript/ai/src/activities/chat/stream/message-updaters.ts

@@ -244,28 +244,35 @@ export function updateToolCallApprovalResponse(
 }
 
 /**
- * Update or add a thinking part to a message.
+ * Update or add a thinking part to a message, keyed by stepId.
+ * Each distinct stepId produces its own ThinkingPart.
  */
 export function updateThinkingPart(
   messages: Array<UIMessage>,
   messageId: string,
+  stepId: string,
   content: string,
+  signature?: string,
 ): Array<UIMessage> {
   return messages.map((msg) => {
     if (msg.id !== messageId) {
       return msg
     }
 
     const parts = [...msg.parts]
-    const thinkingPartIndex = parts.findIndex((p) => p.type === 'thinking')
+    const thinkingPartIndex = parts.findIndex(
+      (p) => p.type === 'thinking' && p.stepId === stepId,
+    )
 
     const thinkingPart: ThinkingPart = {
       type: 'thinking',
       content,
+      stepId,
+      ...(signature && { signature }),
     }
 
     if (thinkingPartIndex >= 0) {
-      // Update existing thinking part
+      // Update existing thinking part for this step
       parts[thinkingPartIndex] = thinkingPart
     } else {
       // Add new thinking part at the end (preserve natural streaming order)