fix(ai, ai-anthropic): thinking blocks missing on turn 2+ in tool loops#391
Open
imsherrill wants to merge 11 commits intoTanStack:mainfrom
Open
fix(ai, ai-anthropic): thinking blocks missing on turn 2+ in tool loops#391imsherrill wants to merge 11 commits intoTanStack:mainfrom
imsherrill wants to merge 11 commits intoTanStack:mainfrom
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds per-step tracking and optional provider signatures for "thinking" blocks across streaming, processor, engine, adapter, and client wiring so thinking content/signatures are emitted, accumulated, and replayed per stepId instead of being merged across steps. ChangesPer-Step Thinking Tracking & Signature Preservation
Sequence DiagramsequenceDiagram
participant User
participant Model
participant StreamProcessor
participant TextEngine
participant Context
User->>Model: user query / turn n
Model->>StreamProcessor: STEP_STARTED (stepId)
Model->>StreamProcessor: STEP_FINISHED (stepId, thinking_delta, signature_delta)
StreamProcessor->>TextEngine: onThinkingUpdate(msgId, stepId, thinking_content)
TextEngine->>Context: accumulate thinking[stepId] + signature[stepId]
StreamProcessor->>TextEngine: emit tool_use / assistant message (attach thinking if present)
TextEngine->>Context: store assistant message with thinking array
User->>Model: next turn (with prior tool_result)
Context->>Model: include prior thinking entries + signatures
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
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 c53d501
☁️ Nx Cloud last updated this comment at |
@tanstack/ai
@tanstack/ai-anthropic
@tanstack/ai-client
@tanstack/ai-code-mode
@tanstack/ai-code-mode-skills
@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-isolate-cloudflare
@tanstack/ai-isolate-node
@tanstack/ai-isolate-quickjs
@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.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
packages/typescript/smoke-tests/e2e/tests/thinking.spec.ts (2)
48-51: Tighten types ongetMessagesand filter callbacks.The
Array<any>return type and the(p: any)predicates throughout the tests disable type-checking on the assertions (e.g.assistantMessage.parts,firstStep.content,firstStep.signature). SinceUIMessageis already exported from@tanstack/ai-clientand used inmock.tsx, importing it here would catch shape regressions (e.g. ifThinkingPart.signatureis renamed) at compile time rather than flakily at runtime.Suggested change
-import { test, expect, type Page } from '@playwright/test' +import { expect, test } from '@playwright/test' +import type { Page } from '@playwright/test' +import type { UIMessage } from '@tanstack/ai-client' @@ -async function getMessages(page: Page): Promise<Array<any>> { +async function getMessages(page: Page): Promise<Array<UIMessage>> { const jsonContent = await page.locator('#messages-json-content').textContent() return JSON.parse(jsonContent || '[]') }This also addresses the ESLint
sort-importsandimport/consistent-type-specifier-stylehints on line 1.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/smoke-tests/e2e/tests/thinking.spec.ts` around lines 48 - 51, Tighten the types by importing UIMessage from "@tanstack/ai-client" and update getMessages to return Promise<UIMessage[]> (change the function signature and JSON parse typing), then replace loose any usages in filter/map callbacks (e.g., predicates like (p: any)) with the concrete UIMessage type and update assertions to use the typed properties (e.g., assistantMessage.parts, firstStep.content, firstStep.signature) so TypeScript catches shape/regression changes; ensure imports follow the project's import style to address sort-imports and import/consistent-type-specifier-style hints.
115-122:afterEachis outside thedescribeblock.The
test.afterEachhook at line 115 is declared after thedescribecloses on line 113, so it attaches to the root file scope rather than theChat E2E Tests - Multi-Step Thinking (Mock API)group. That's functionally fine today (only one describe block in this file), but if additional describes are added later the screenshot hook will run for all of them, which may or may not be intended. Consider moving it inside the describe for clearer scoping.Suggested change
test.describe('Chat E2E Tests - Multi-Step Thinking (Mock API)', () => { ... + + test.afterEach(async ({ page }, testInfo) => { + if (testInfo.status !== testInfo.expectedStatus) { + await page.screenshot({ + path: `test-results/failure-${testInfo.title.replace(/\s+/g, '-')}.png`, + fullPage: true, + }) + } + }) }) - -test.afterEach(async ({ page }, testInfo) => { - if (testInfo.status !== testInfo.expectedStatus) { - await page.screenshot({ - path: `test-results/failure-${testInfo.title.replace(/\s+/g, '-')}.png`, - fullPage: true, - }) - } -})🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/smoke-tests/e2e/tests/thinking.spec.ts` around lines 115 - 122, The screenshot cleanup hook test.afterEach is currently declared outside the describe "Chat E2E Tests - Multi-Step Thinking (Mock API)" so it attaches to the root scope; move the test.afterEach block inside that describe (before its closing brace) so the hook is scoped to that describe only, preserving the same screenshot logic and parameters (page, testInfo) and keeping the filename generation code unchanged.packages/typescript/smoke-tests/e2e/src/routes/api.mock-chat.ts (1)
310-394: TwoSTEP_FINISHEDevents per step — intentional but worth a comment.Each thinking step emits
STEP_FINISHEDtwice (first with adelta, then withcontent+signature). This appears to intentionally simulate streaming deltas followed by a final summary event, butSTEP_FINISHEDconventionally fires once per step. Since the PR's tests pass and rely on both events being processed (finalsignaturestored on the step), consider adding a brief inline comment documenting that the firstSTEP_FINISHEDcarries the incrementaldeltaand the second carries the finalcontent/signature, so future readers don't mistake it for a duplicate event bug.Also: the
Date.now()calls in the chunk literals are evaluated at module load time;createMockStreamalready overridestimestampon yield (line 407), so these initial values are dead — not a defect, just noise. Consider using a placeholder liketimestamp: 0for the static chunks to make that clearer.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/typescript/smoke-tests/e2e/src/routes/api.mock-chat.ts` around lines 310 - 394, Add a short inline comment above the "thinking-multi-step" mock describing that each thinking step intentionally emits two STEP_FINISHED events (first carrying the streaming delta, second carrying final content+signature) so readers don't treat it as a duplicate; also replace the Date.now() timestamps in the chunk literals with a placeholder like 0 since createMockStream overrides timestamps at yield (see "thinking-multi-step", STEP_FINISHED entries and createMockStream).
🤖 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-anthropic/src/adapters/text.ts`:
- Around line 407-417: The code drops signed thinking for assistant messages
that have no toolCalls; ensure any message.thinking items with a signature are
always appended to contentBlocks as AnthropicContentBlock regardless of the
assistant/toolCalls branch. Update the logic around the handling of
message.thinking (the loop that pushes {type:'thinking', thinking:
thinking.content, signature: thinking.signature} into contentBlocks) so it runs
for both the assistant+toolCalls branch and the later assistant-only branch (the
code paths around the existing message.thinking loop and the block at lines
~468-478); you can either hoist the loop above the branch so it always executes
for message.thinking or add the same signed-thinking push into the
assistant-only branch, referencing message.thinking, contentBlocks, and
AnthropicContentBlock.
In `@packages/typescript/ai/src/activities/chat/stream/processor.ts`:
- Around line 1119-1136: The code currently prefers state.currentThinkingStepId
over chunk.stepId which causes distinct chunk.stepId values to be merged into an
existing thinking step; change the selection to prefer the explicit chunk.stepId
when present: compute stepId = chunk.stepId ?? state.currentThinkingStepId, then
if the resolved stepId is not present in state.thinkingSteps, initialize it
(state.thinkingSteps.set(stepId, ''), state.thinkingStepOrder.push(stepId)) and
update state.currentThinkingStepId = stepId; keep the existing
pendingThinkingStepId consumption logic but ensure it does not override an
explicit chunk.stepId.
---
Nitpick comments:
In `@packages/typescript/smoke-tests/e2e/src/routes/api.mock-chat.ts`:
- Around line 310-394: Add a short inline comment above the
"thinking-multi-step" mock describing that each thinking step intentionally
emits two STEP_FINISHED events (first carrying the streaming delta, second
carrying final content+signature) so readers don't treat it as a duplicate; also
replace the Date.now() timestamps in the chunk literals with a placeholder like
0 since createMockStream overrides timestamps at yield (see
"thinking-multi-step", STEP_FINISHED entries and createMockStream).
In `@packages/typescript/smoke-tests/e2e/tests/thinking.spec.ts`:
- Around line 48-51: Tighten the types by importing UIMessage from
"@tanstack/ai-client" and update getMessages to return Promise<UIMessage[]>
(change the function signature and JSON parse typing), then replace loose any
usages in filter/map callbacks (e.g., predicates like (p: any)) with the
concrete UIMessage type and update assertions to use the typed properties (e.g.,
assistantMessage.parts, firstStep.content, firstStep.signature) so TypeScript
catches shape/regression changes; ensure imports follow the project's import
style to address sort-imports and import/consistent-type-specifier-style hints.
- Around line 115-122: The screenshot cleanup hook test.afterEach is currently
declared outside the describe "Chat E2E Tests - Multi-Step Thinking (Mock API)"
so it attaches to the root scope; move the test.afterEach block inside that
describe (before its closing brace) so the hook is scoped to that describe only,
preserving the same screenshot logic and parameters (page, testInfo) and keeping
the filename generation code unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 52fe2bf7-eb5c-4660-a208-3e1292bede10
📒 Files selected for processing (7)
.changeset/thinking-blocks-per-step.mdpackages/typescript/ai-anthropic/src/adapters/text.tspackages/typescript/ai/src/activities/chat/stream/processor.tspackages/typescript/ai/tests/stream-processor.test.tspackages/typescript/smoke-tests/e2e/src/routes/api.mock-chat.tspackages/typescript/smoke-tests/e2e/src/routes/mock.tsxpackages/typescript/smoke-tests/e2e/tests/thinking.spec.ts
✅ Files skipped from review due to trivial changes (1)
- .changeset/thinking-blocks-per-step.md
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/typescript/ai/tests/stream-processor.test.ts
Comment on lines
+407
to
+417
| 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) | ||
| } | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
Replay signed thinking for non-tool assistant messages too.
Line 407 only preserves thinking inside the assistant && toolCalls branch. Assistant responses that contain signed thinking but no tool calls fall through to lines 468-478, so their message.thinking is dropped from future Anthropic context.
🐛 Proposed fix
+ const appendThinkingBlocks = (
+ contentBlocks: AnthropicContentBlocks,
+ thinkingParts: ModelMessage['thinking'],
+ ) => {
+ if (!thinkingParts?.length) return
+
+ for (const thinking of thinkingParts) {
+ if (thinking.signature) {
+ contentBlocks.push({
+ type: 'thinking',
+ thinking: thinking.content,
+ signature: thinking.signature,
+ } as unknown as AnthropicContentBlock)
+ }
+ }
+ }
+
for (const message of messages) {
const role = message.role
@@
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)
- }
- }
- }
+ appendThinkingBlocks(contentBlocks, message.thinking)
@@
+ if (role === 'assistant' && message.thinking?.length) {
+ const contentBlocks: AnthropicContentBlocks = []
+ appendThinkingBlocks(contentBlocks, message.thinking)
+
+ if (message.content) {
+ contentBlocks.push({
+ type: 'text',
+ text: typeof message.content === 'string' ? message.content : '',
+ } as AnthropicContentBlock)
+ }
+
+ formattedMessages.push({
+ role: 'assistant',
+ content: contentBlocks,
+ })
+ continue
+ }
+
formattedMessages.push({Also applies to: 468-478
🤖 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 407 -
417, The code drops signed thinking for assistant messages that have no
toolCalls; ensure any message.thinking items with a signature are always
appended to contentBlocks as AnthropicContentBlock regardless of the
assistant/toolCalls branch. Update the logic around the handling of
message.thinking (the loop that pushes {type:'thinking', thinking:
thinking.content, signature: thinking.signature} into contentBlocks) so it runs
for both the assistant+toolCalls branch and the later assistant-only branch (the
code paths around the existing message.thinking loop and the block at lines
~468-478); you can either hoist the loop above the branch so it always executes
for message.thinking or add the same signed-thinking push into the
assistant-only branch, referencing message.thinking, contentBlocks, and
AnthropicContentBlock.
Comment on lines
+1119
to
+1136
| // Consume pending stepId from STEP_STARTED that arrived before the message existed | ||
| if (this.pendingThinkingStepId) { | ||
| state.currentThinkingStepId = this.pendingThinkingStepId | ||
| if (!state.thinkingSteps.has(this.pendingThinkingStepId)) { | ||
| state.thinkingSteps.set(this.pendingThinkingStepId, '') | ||
| state.thinkingStepOrder.push(this.pendingThinkingStepId) | ||
| } | ||
| this.pendingThinkingStepId = null | ||
| } | ||
|
|
||
| const stepId = state.currentThinkingStepId ?? chunk.stepId | ||
|
|
||
| // Auto-initialize if no prior STEP_STARTED (backward compat) | ||
| if (!state.thinkingSteps.has(stepId)) { | ||
| state.thinkingSteps.set(stepId, '') | ||
| state.thinkingStepOrder.push(stepId) | ||
| state.currentThinkingStepId = stepId | ||
| } |
Contributor
There was a problem hiding this comment.
Prefer the chunk’s explicit stepId to avoid merging distinct thinking steps.
Line 1129 ignores chunk.stepId whenever currentThinkingStepId is already set. If a replay/provider emits STEP_FINISHED events with distinct stepIds but without STEP_STARTED, later steps are accumulated into the previous ThinkingPart.
🐛 Proposed fix
- const stepId = state.currentThinkingStepId ?? chunk.stepId
+ const stepId = chunk.stepId ?? state.currentThinkingStepId
// Auto-initialize if no prior STEP_STARTED (backward compat)
if (!state.thinkingSteps.has(stepId)) {
state.thinkingSteps.set(stepId, '')
state.thinkingStepOrder.push(stepId)
- state.currentThinkingStepId = stepId
}
+ state.currentThinkingStepId = stepId🤖 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` around lines
1119 - 1136, The code currently prefers state.currentThinkingStepId over
chunk.stepId which causes distinct chunk.stepId values to be merged into an
existing thinking step; change the selection to prefer the explicit chunk.stepId
when present: compute stepId = chunk.stepId ?? state.currentThinkingStepId, then
if the resolved stepId is not present in state.thinkingSteps, initialize it
(state.thinkingSteps.set(stepId, ''), state.thinkingStepOrder.push(stepId)) and
update state.currentThinkingStepId = stepId; keep the existing
pendingThinkingStepId consumption logic but ensure it does not override an
explicit chunk.stepId.
- 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
…int; clear stale pendingThinkingStepId - Move `betas: ['interleaved-thinking-2025-05-14']` out of the shared mapper and onto the `beta.messages.create` call site in chatStream. Prevents the non-beta structuredOutput endpoint from receiving an invalid `betas` field when a thinking budget is configured. - Clear `pendingThinkingStepId` when a later STEP_STARTED takes the active-message branch, and also in `resetStreamState`, so a stale pending id can't misattribute a later STEP_FINISHED's delta to an earlier step. - Add covering test for the pendingThinkingStepId leak (red-green verified).
- Add `thinking-multi-step` mock scenario to the e2e harness emitting STEP_STARTED/STEP_FINISHED pairs for two distinct stepIds with provider signatures, followed by a text message and RUN_FINISHED. - Expose thinkingPartCount / thinkingStepIds on the mock chat page via data-* attributes for assertion. - Add tests/thinking.spec.ts asserting two ThinkingParts with distinct stepIds and matching signatures are produced (pre-PR behavior merged them into a single part). - Add .changeset/thinking-blocks-per-step.md bumping @tanstack/ai, @tanstack/ai-anthropic, and @tanstack/ai-client.
When STEP_STARTED arrives before the assistant message exists, its stepId is stashed in pendingThinkingStepId. handleStepFinishedEvent already consumes it, but handleReasoningMessageContentEvent did not, so reasoning deltas were keyed by the reasoning messageId and the matching signature from STEP_FINISHED landed on a different ThinkingPart. With Anthropic's interleaved thinking around tool calls this produced two ThinkingParts per block (one unsigned content, one signed empty). Consume the pending stepId here too so both event paths attribute to the same step.
Both handleStepFinishedEvent and handleReasoningMessageContentEvent need to promote a pending stepId from a STEP_STARTED that arrived before the assistant message existed. Pull the shared logic into a small private helper instead of duplicating it.
The bare yield bypassed the file's existing StreamChunk cast helper, so after the merge from main StepFinishedEvent now extends the stricter AG-UI base (requires stepName, no signature) and CI failed to build. Use asChunk like every other yield in this generator and add stepName.
imsherrill
force-pushed
the
fix/thinking-blocks-per-step
branch
from
962e396 to
c53d501
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.changeset/thinking-blocks-per-step.md:
- Around line 2-4: The changeset currently lists patch bumps for packages that
introduce a breaking API change: update the version bumps to minor for the
affected packages; specifically change the entries in
.changeset/thinking-blocks-per-step.md from patch to minor for '@tanstack/ai'
and '@tanstack/ai-client' (and any other package exposing the new
onThinkingUpdate signature), and ensure the changeset message mentions the
breaking change to the onThinkingUpdate callback (old signature (messageId,
content) -> new signature (messageId, stepId, content)) and references the
public API types StreamProcessorEvents and ChatClient so the release tooling
applies a minor bump.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ef6014b3-a4ed-43d1-bc2b-e065a9dec715
📒 Files selected for processing (1)
.changeset/thinking-blocks-per-step.md
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
packages/typescript/ai/src/activities/chat/tools/tool-calls.ts (1)
96-98: ⚡ Quick winAdd
toolCallName?: stringoverride toToolCallStartEventinterface and remove the cast workaround
ToolCallEndEventalready declares bothtoolCallName?: stringandtoolName?: stringas optional fields (lines 930–934).ToolCallStartEventshould follow the same pattern by explicitly overridingtoolCallNameto be optional, rather than using thePartial<>cast at the call site. This makes the optionality explicit in the type contract and aligns with howToolCallEndEventhandles backward compatibility.Suggested changes
In
packages/typescript/ai/src/types.ts, updateToolCallStartEvent:export interface ToolCallStartEvent extends AGUIToolCallStartEvent { /** Model identifier for multi-model support */ model?: string /** * `@deprecated` Use `toolCallName` instead (from `@ag-ui/core` spec). * Kept for backward compatibility. */ toolName: string + /** Name of the tool being called (from `@ag-ui/core`) */ + toolCallName?: string /** Index for parallel tool calls */ index?: number /** Provider-specific metadata to carry into the ToolCall */ providerMetadata?: Record<string, unknown> }Then in
packages/typescript/ai/src/activities/chat/tools/tool-calls.ts, simplify lines 96–98:- const runtimeEvent = event as Partial<ToolCallStartEvent> & - Pick<ToolCallStartEvent, 'toolName'> - const name = runtimeEvent.toolCallName ?? runtimeEvent.toolName + const name = event.toolCallName ?? event.toolName🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/typescript/ai/src/activities/chat/tools/tool-calls.ts` around lines 96 - 98, Update the ToolCallStartEvent interface to include an explicit optional toolCallName?: string field (mirroring ToolCallEndEvent's toolCallName?: string and toolName?: string) in the types definition (ToolCallStartEvent), then remove the Partial<> cast workaround in the tool call handler and directly treat event as ToolCallStartEvent so the code in tool-calls.ts can simply read const name = event.toolCallName ?? event.toolName; ensure the change preserves backward compatibility and compiles against existing uses of ToolCallStartEvent and tool-calls.ts.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@packages/typescript/ai/src/activities/chat/tools/tool-calls.ts`:
- Around line 96-98: Update the ToolCallStartEvent interface to include an
explicit optional toolCallName?: string field (mirroring ToolCallEndEvent's
toolCallName?: string and toolName?: string) in the types definition
(ToolCallStartEvent), then remove the Partial<> cast workaround in the tool call
handler and directly treat event as ToolCallStartEvent so the code in
tool-calls.ts can simply read const name = event.toolCallName ?? event.toolName;
ensure the change preserves backward compatibility and compiles against existing
uses of ToolCallStartEvent and tool-calls.ts.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9bd6764e-e99e-4415-89ac-f8e667b04a7e
📒 Files selected for processing (1)
packages/typescript/ai/src/activities/chat/tools/tool-calls.ts
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
2 participants
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
New Features
Documentation
Tests