fix(ai, ai-anthropic): thinking blocks missing on turn 2+ in tool loops#391
Open
imsherrill wants to merge 2 commits intoTanStack:mainfrom
Open
fix(ai, ai-anthropic): thinking blocks missing on turn 2+ in tool loops#391imsherrill wants to merge 2 commits intoTanStack:mainfrom
imsherrill wants to merge 2 commits intoTanStack:mainfrom
Conversation
- Track thinking per-step via stepId instead of merging into single ThinkingPart - Capture Anthropic signature_delta and preserve through the full stack - Server-side TextEngine accumulates thinking + signatures per iteration - Include thinking blocks in Anthropic message history for multi-turn context - Add interleaved-thinking-2025-05-14 beta header when thinking is enabled - Add tests for multi-step thinking, backward compat, and result aggregation Closes TanStack#340
Contributor
📝 WalkthroughWalkthroughThis PR implements multi-step thinking support for Anthropic's interleaved-thinking API. Changes add step-aware thinking accumulation across tool loops, capture provider signatures from streaming responses, integrate thinking blocks into assistant messages, and refactor stream state to track thinking per step ID rather than a single aggregate field. Changes
Sequence DiagramsequenceDiagram
participant Client
participant AnthropicAPI as Anthropic API
participant StreamProcessor
participant TextEngine
participant Messages as Message Builder
Client->>AnthropicAPI: Request with betas + thinkingBudget
AnthropicAPI-->>StreamProcessor: STEP_STARTED (stepId=step-1)
StreamProcessor->>StreamProcessor: Record pendingThinkingStepId
AnthropicAPI-->>StreamProcessor: STEP_FINISHED (stepId=step-1, thinking content)
StreamProcessor->>StreamProcessor: Accumulate thinkingSteps[step-1]
StreamProcessor->>StreamProcessor: Update currentThinkingStepId
StreamProcessor->>Messages: updateThinkingPart(msgId, step-1, content)
AnthropicAPI-->>StreamProcessor: signature_delta events
StreamProcessor->>StreamProcessor: Accumulate signature
AnthropicAPI-->>StreamProcessor: content_block_stop (thinking)
StreamProcessor->>StreamProcessor: Emit STEP_FINISHED with signature
AnthropicAPI-->>StreamProcessor: Text/Tool-use content
StreamProcessor->>TextEngine: Process stream chunks
TextEngine->>TextEngine: Persist thinking from STEP_FINISHED
TextEngine->>Messages: Build assistant message with thinking array
Messages-->>Client: Return message with thinking blocks + signature
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
View your CI Pipeline Execution ↗ for commit c042c55
☁️ Nx Cloud last updated this comment at |
@tanstack/ai
@tanstack/ai-anthropic
@tanstack/ai-client
@tanstack/ai-devtools-core
@tanstack/ai-elevenlabs
@tanstack/ai-event-client
@tanstack/ai-fal
@tanstack/ai-gemini
@tanstack/ai-grok
@tanstack/ai-groq
@tanstack/ai-ollama
@tanstack/ai-openai
@tanstack/ai-openrouter
@tanstack/ai-preact
@tanstack/ai-react
@tanstack/ai-react-ui
@tanstack/ai-solid
@tanstack/ai-solid-ui
@tanstack/ai-svelte
@tanstack/ai-vue
@tanstack/ai-vue-ui
@tanstack/preact-ai-devtools
@tanstack/react-ai-devtools
@tanstack/solid-ai-devtools
commit: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/typescript/ai/src/activities/chat/messages.ts (1)
181-186:⚠️ Potential issue | 🟠 MajorRound-tripping through
ModelMessagestill strips thinking.Line 185 starts serializing
thinkinginto assistantModelMessages, butmodelMessageToUIMessage()below never materializesmodelMessage.thinkingback intoThinkingParts. Any history hydrated throughnormalizeToUIMessage()/ChatClient.append()will silently drop the thinking blocks and signatures, so a laterconvertMessagesToModelMessages()sends stripped context again.Patch sketch
export function modelMessageToUIMessage( modelMessage: ModelMessage, id?: string, ): UIMessage { const parts: Array<MessagePart> = [] + + if (modelMessage.role === 'assistant' && modelMessage.thinking?.length) { + for (const thinkingPart of modelMessage.thinking) { + parts.push({ + type: 'thinking', + content: thinkingPart.content, + ...(thinkingPart.signature && { + signature: thinkingPart.signature, + }), + }) + } + } // Handle tool results (when role is "tool") - only produce tool-result part, // not a text part (the content IS the tool result, not display text) if (modelMessage.role === 'tool' && modelMessage.toolCallId) {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/ai/src/activities/chat/messages.ts` around lines 181 - 186, The assistant message builder adds a thinking field to ModelMessage (see the push that adds thinking), but modelMessageToUIMessage()/normalizeToUIMessage()/ChatClient.append() never rehydrate modelMessage.thinking back into ThinkingPart objects and signatures, causing round-trip loss; update modelMessageToUIMessage (and any normalizeToUIMessage/ChatClient.append code paths) to detect modelMessage.thinking and reconstruct the original ThinkingPart shape (including reasoning text and signature fields) so that convertMessagesToModelMessages receives intact thinking blocks on subsequent conversions.
🧹 Nitpick comments (3)
packages/typescript/ai/tests/stream-processor.test.ts (1)
778-830: Add one signature propagation regression here.These new cases cover
stepId, but none assert that aSTEP_FINISHED.signaturesurvives into the storedthinkingpart. That field is what Anthropic needs on the follow-up turn, so this branch is still unguarded by tests.Test sketch
+ it('should persist signature on a thinking step', () => { + const processor = new StreamProcessor() + processor.prepareAssistantMessage() + + processor.processChunk( + chunk('STEP_FINISHED', { + stepId: 'step-1', + delta: 'thinking...', + signature: 'sig-1', + }), + ) + + expect(processor.getMessages()[0]?.parts[0]).toEqual({ + type: 'thinking', + content: 'thinking...', + stepId: 'step-1', + signature: 'sig-1', + }) + })🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/ai/tests/stream-processor.test.ts` around lines 778 - 830, Add assertions to the tests that verify STEP_FINISHED.signature is propagated into the stored thinking parts and into the aggregated state.thinking: when using StreamProcessor.processChunk with ev.stepFinished(..., stepId) include a signature value and assert that (from processor.getMessages()[0]!.parts filtered by type 'thinking') each thinking part has the same signature (e.g., for 'step-1' and 'step-2'), and in the getResult()/getState() concatenation case assert that state.thinking retains each step's signature information (or that the per-step signatures are accessible on the stored thinking parts); locate code via StreamProcessor.processChunk, ev.stepFinished, processor.getMessages, and processor.getState to add these checks.packages/typescript/ai-anthropic/src/adapters/text.ts (2)
395-405: Thinking blocks without signatures are silently filtered out.The loop only pushes thinking blocks that have a
signature. This is correct per Anthropic's API requirements (signatures are mandatory for replaying thinking in multi-turn context), but a brief comment explaining this would help maintainability.📝 Suggested comment
if (message.thinking?.length) { for (const thinking of message.thinking) { + // Anthropic requires signatures to replay thinking blocks in multi-turn context; + // blocks without signatures cannot be included. if (thinking.signature) { contentBlocks.push({🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/ai-anthropic/src/adapters/text.ts` around lines 395 - 405, Add a short explanatory comment above the loop that filters thinking blocks to clarify that only thinking entries with a signature are included because Anthropic requires signatures to replay thinking in multi-turn context; reference the variables message.thinking, thinking.signature, the contentBlocks.push of AnthropicContentBlock, and note that omitting unsigned thinking blocks is intentional for API compliance and not a bug.
292-295: Add a comment explaining the beta header requirement and type assertion.The beta header
interleaved-thinking-2025-05-14is required to enable interleaved thinking in Claude models during tool-use conversations. Add a brief comment explaining why this beta is needed and why theas anycast is used (the Anthropic SDK types don't includebetasin this parameter definition).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/ai-anthropic/src/adapters/text.ts` around lines 292 - 295, Add an inline comment above the conditional that sets betas explaining that the beta header "interleaved-thinking-2025-05-14" is required to enable interleaved thinking for Claude during tool-use conversations, and note that the `as any` cast is used because the Anthropic SDK types do not include `betas` on this parameter; update the block around the `thinkingBudget` conditional and the `betas` property to include that short explanatory comment next to `betas` and `as any` so future readers understand both the runtime requirement and the type-workaround.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/typescript/ai/src/activities/chat/messages.ts`:
- Around line 168-169: The issue is that pendingThinking is buffered separately
from current.toolCalls, which loses the original interleaving (e.g., thinking,
tool-call, thinking) when constructing ModelMessage; change the buffering to a
single ordered stream of typed entries (e.g., an array of {type:
'thinking'|'toolCall', payload: ...}) instead of separate pendingThinking and
current.toolCalls so you can preserve insert order when emitting/serializing;
update all places that push to pendingThinking and to current.toolCalls (and the
duplicate buffering logic referenced around the 233-239 region) to push a typed
entry into the unified buffer and update the code that builds the assistant
ModelMessage to iterate this unified buffer and emit thinking/toolCall entries
in original sequence.
In `@packages/typescript/ai/src/activities/chat/stream/processor.ts`:
- Line 146: The field pendingThinkingStepId is not cleared in
resetStreamState(), causing stale step IDs to persist across streams; update the
resetStreamState() method to set this.pendingThinkingStepId = null so that
prepareAssistantMessage() starting a new stream cannot inherit an old id, and
ensure any other stream-reset paths also call resetStreamState() (or explicitly
clear pendingThinkingStepId) to avoid leaking the previous STEP_STARTED
association.
---
Outside diff comments:
In `@packages/typescript/ai/src/activities/chat/messages.ts`:
- Around line 181-186: The assistant message builder adds a thinking field to
ModelMessage (see the push that adds thinking), but
modelMessageToUIMessage()/normalizeToUIMessage()/ChatClient.append() never
rehydrate modelMessage.thinking back into ThinkingPart objects and signatures,
causing round-trip loss; update modelMessageToUIMessage (and any
normalizeToUIMessage/ChatClient.append code paths) to detect
modelMessage.thinking and reconstruct the original ThinkingPart shape (including
reasoning text and signature fields) so that convertMessagesToModelMessages
receives intact thinking blocks on subsequent conversions.
---
Nitpick comments:
In `@packages/typescript/ai-anthropic/src/adapters/text.ts`:
- Around line 395-405: Add a short explanatory comment above the loop that
filters thinking blocks to clarify that only thinking entries with a signature
are included because Anthropic requires signatures to replay thinking in
multi-turn context; reference the variables message.thinking,
thinking.signature, the contentBlocks.push of AnthropicContentBlock, and note
that omitting unsigned thinking blocks is intentional for API compliance and not
a bug.
- Around line 292-295: Add an inline comment above the conditional that sets
betas explaining that the beta header "interleaved-thinking-2025-05-14" is
required to enable interleaved thinking for Claude during tool-use
conversations, and note that the `as any` cast is used because the Anthropic SDK
types do not include `betas` on this parameter; update the block around the
`thinkingBudget` conditional and the `betas` property to include that short
explanatory comment next to `betas` and `as any` so future readers understand
both the runtime requirement and the type-workaround.
In `@packages/typescript/ai/tests/stream-processor.test.ts`:
- Around line 778-830: Add assertions to the tests that verify
STEP_FINISHED.signature is propagated into the stored thinking parts and into
the aggregated state.thinking: when using StreamProcessor.processChunk with
ev.stepFinished(..., stepId) include a signature value and assert that (from
processor.getMessages()[0]!.parts filtered by type 'thinking') each thinking
part has the same signature (e.g., for 'step-1' and 'step-2'), and in the
getResult()/getState() concatenation case assert that state.thinking retains
each step's signature information (or that the per-step signatures are
accessible on the stored thinking parts); locate code via
StreamProcessor.processChunk, ev.stepFinished, processor.getMessages, and
processor.getState to add these checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 19b0cdc6-d9a2-4d83-bb64-7a1aef5b7b29
📒 Files selected for processing (10)
packages/typescript/ai-anthropic/src/adapters/text.tspackages/typescript/ai-client/src/chat-client.tspackages/typescript/ai/src/activities/chat/index.tspackages/typescript/ai/src/activities/chat/messages.tspackages/typescript/ai/src/activities/chat/stream/message-updaters.tspackages/typescript/ai/src/activities/chat/stream/processor.tspackages/typescript/ai/src/activities/chat/stream/types.tspackages/typescript/ai/src/types.tspackages/typescript/ai/tests/message-updaters.test.tspackages/typescript/ai/tests/stream-processor.test.ts
Comment on lines
+168
to
169
| let pendingThinking: Array<{ content: string; signature?: string }> = [] | ||
|
|
Contributor
There was a problem hiding this comment.
This still loses thinking/tool-call order within a single turn.
pendingThinking is buffered independently from current.toolCalls, so a sequence like [thinking(step-1), tool-call(a), thinking(step-2), tool-call(b)] gets flattened into one assistant ModelMessage with thinking: [step-1, step-2] and toolCalls: [a, b]. If a provider emits multiple thinking blocks around multiple tool calls in one turn, the original interleaving is gone and the next request cannot faithfully replay that history.
Also applies to: 233-239
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@packages/typescript/ai/src/activities/chat/messages.ts` around lines 168 -
169, The issue is that pendingThinking is buffered separately from
current.toolCalls, which loses the original interleaving (e.g., thinking,
tool-call, thinking) when constructing ModelMessage; change the buffering to a
single ordered stream of typed entries (e.g., an array of {type:
'thinking'|'toolCall', payload: ...}) instead of separate pendingThinking and
current.toolCalls so you can preserve insert order when emitting/serializing;
update all places that push to pendingThinking and to current.toolCalls (and the
duplicate buffering logic referenced around the 233-239 region) to push a typed
entry into the unified buffer and update the code that builds the assistant
ModelMessage to iterate this unified buffer and emit thinking/toolCall entries
in original sequence.
| private activeMessageIds: Set<string> = new Set() | ||
| private toolCallToMessage: Map<string, string> = new Map() | ||
| private pendingManualMessageId: string | null = null | ||
| private pendingThinkingStepId: string | null = null |
Contributor
There was a problem hiding this comment.
pendingThinkingStepId should be reset in resetStreamState().
If a STEP_STARTED event arrives at the end of one stream but no corresponding STEP_FINISHED follows (e.g., stream aborted), pendingThinkingStepId will retain a stale value. When a new stream starts via prepareAssistantMessage() → resetStreamState(), this stale stepId could incorrectly associate with new thinking content.
🐛 Proposed fix
private resetStreamState(): void {
this.messageStates.clear()
this.activeMessageIds.clear()
this.activeRuns.clear()
this.toolCallToMessage.clear()
this.pendingManualMessageId = null
+ this.pendingThinkingStepId = null
this.finishReason = null
this.hasError = false
this.isDone = false
this.chunkStrategy.reset?.()
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@packages/typescript/ai/src/activities/chat/stream/processor.ts` at line 146,
The field pendingThinkingStepId is not cleared in resetStreamState(), causing
stale step IDs to persist across streams; update the resetStreamState() method
to set this.pendingThinkingStepId = null so that prepareAssistantMessage()
starting a new stream cannot inherit an old id, and ensure any other
stream-reset paths also call resetStreamState() (or explicitly clear
pendingThinkingStepId) to avoid leaking the previous STEP_STARTED association.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ThinkingPartper message instead of one per stepinterleaved-thinking-2025-05-14beta header when thinking is enabled — required by Anthropic API for thinking on tool-result follow-up turnsstepIdthrough the full stack: adapter → engine → processor → UIMessageChanges
@tanstack/ai:ThinkingPartgains optionalstepIdandsignaturefieldsModelMessagegains optionalthinkingarray for multi-turn contextStepFinishedEventgains optionalsignaturefieldStreamProcessorhandlesSTEP_STARTED, tracks thinking per-step viaMap<stepId, content>TextEngineaccumulates thinking + signatures per iteration, includes them in assistant messagesbuildAssistantMessagespreservesThinkingParts intoModelMessage.thinkingupdateThinkingPartkeys onstepIdinstead of replacing a single partonThinkingUpdatecallback gainsstepIdparameter@tanstack/ai-anthropic:signature_deltastream events for thinking block signaturesSTEP_FINISHEDwith signature oncontent_block_stopformatMessagesfor multi-turn historybetas: ['interleaved-thinking-2025-05-14']when thinking is enabled@tanstack/ai-client:onThinkingUpdatecallback updated for newstepIdparameterTest plan
@tanstack/ai(4 new tests for multi-step thinking)@tanstack/ai-clientCloses #340
Summary by CodeRabbit