feat(ai-claude-code): stream tool-call arguments (+ fix processor id-rename orphaning tool results)#903
Closed
ozolcar wants to merge 2 commits into
Closed
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR adds incremental Claude Code tool-call argument streaming and centralizes message-id renaming in the stream processor so tool results stay attached after placeholder renames. ChangesClaude Code Partial Tool-Call Streaming
Processor Message ID Rename Fix
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ClaudeAgentSDK
participant translateSdkStream
participant StreamConsumer
ClaudeAgentSDK->>translateSdkStream: content_block_start (tool_use)
translateSdkStream->>StreamConsumer: TOOL_CALL_START
ClaudeAgentSDK->>translateSdkStream: content_block_delta (input_json_delta)
translateSdkStream->>StreamConsumer: TOOL_CALL_ARGS (delta)
ClaudeAgentSDK->>translateSdkStream: content_block_stop
translateSdkStream->>translateSdkStream: closePartialToolUse (parse args)
translateSdkStream->>StreamConsumer: TOOL_CALL_END (input)
sequenceDiagram
participant StreamProcessor
participant ProcessorState
participant ToolCallRouting
StreamProcessor->>ProcessorState: renameMessageId(oldId, newId)
StreamProcessor->>ToolCallRouting: remap toolCallToMessage values
StreamProcessor->>ProcessorState: update messages, messageStates, activeMessageIds, structured batches
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
The stream translator only handled `text_delta` and `thinking_delta` partial events, so a `tool_use` block surfaced as a single completed TOOL_CALL_START/ARGS/END at the end of the message. It now also translates the SDK's partial tool input — `content_block_start` (tool_use), `content_block_delta` with `input_json_delta`, and `content_block_stop` — into incremental TOOL_CALL_ARGS deltas, matching how the model adapters (`@tanstack/ai-anthropic`, `@tanstack/ai-openai`) already stream tool args. The complete `assistant` message dedupes any tool call already streamed via partials (tracked by tool-call id, mirroring the existing text/thinking message-id dedup), so nothing is emitted twice. With `streamPartials: false`, tool calls still emit whole from the complete message.
When a tool call streams before any text and without parentMessageId (optional per the AG-UI spec), the StreamProcessor creates a placeholder assistant message and renames it to the real provider id on TEXT_MESSAGE_START. The rename updated messages/messageStates/activeMessageIds but not toolCallToMessage/structuredMessageIds/structuredOutputUpdateBatches, so a TOOL_CALL_RESULT arriving after the rename resolved to the vanished placeholder id and was silently dropped. Route the rename through a single renameMessageId() seam that remaps every id-keyed structure — the same set enumerated by pruneToMessages() and reset() — so no adapter that legitimately omits parentMessageId can orphan a tool result.
ozolcar
force-pushed
the
feat/claude-code-stream-tool-args
branch
from
8b25a36 to
760cbb6
Compare
![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: 1
🧹 Nitpick comments (1)
packages/ai-claude-code/tests/translate.test.ts (1)
636-701: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding a test for tool_use streaming with no preceding
message_start.Existing grouping test covers
parentMessageIdwhen a message_start/text block precedes the tool call. Since translate.ts conditionally omitsparentMessageIdwhenpartialMessageId === null(Line 449-451 in translate.ts), a test asserting the omission in that case would close the gap and guard the branch this layer is responsible for (the actual orphan-prevention itself is fixed at the processor layer per the PR stack).🤖 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/ai-claude-code/tests/translate.test.ts` around lines 636 - 701, Add a focused translate test for streamed tool_use events when no preceding message_start exists, since the current sibling-grouping case only covers the partialMessageId path. In translate.test.ts, extend the collect-based stream_event coverage around the existing tool call assertions to simulate a tool_use block without an earlier message_start, and verify that TOOL_CALL_START does not include parentMessageId when translate() has no partialMessageId to attach. Use the translate.ts branch that conditionally omits parentMessageId and the existing TOOLS_CALL_START/TEXT_MESSAGE_START matching pattern to locate the behavior.
🤖 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 `@packages/ai/src/activities/chat/stream/processor.ts`:
- Around line 764-773: The docstring on the message-id renaming helper refers to
a nonexistent pruneToMessages() method, which will mislead readers. Update that
reference in the comment for the id-renaming logic around the processor’s
message-keyed map cleanup so it points to the actual enumerating method,
removeMessagesAfter(), and keep the wording aligned with the mirrored cleanup
behavior described by reset().
---
Nitpick comments:
In `@packages/ai-claude-code/tests/translate.test.ts`:
- Around line 636-701: Add a focused translate test for streamed tool_use events
when no preceding message_start exists, since the current sibling-grouping case
only covers the partialMessageId path. In translate.test.ts, extend the
collect-based stream_event coverage around the existing tool call assertions to
simulate a tool_use block without an earlier message_start, and verify that
TOOL_CALL_START does not include parentMessageId when translate() has no
partialMessageId to attach. Use the translate.ts branch that conditionally omits
parentMessageId and the existing TOOLS_CALL_START/TEXT_MESSAGE_START matching
pattern to locate the behavior.
🪄 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: 1d9d7458-6358-4a9f-bd17-c8de624e08ec
📒 Files selected for processing (7)
.changeset/claude-code-stream-tool-args.md.changeset/processor-message-id-rename.mdpackages/ai-claude-code/src/stream/sdk-types.tspackages/ai-claude-code/src/stream/translate.tspackages/ai-claude-code/tests/translate.test.tspackages/ai/src/activities/chat/stream/processor.tspackages/ai/tests/stream-processor.test.ts
Comment on lines
+764
to
+773
| /** | ||
| * Rename a message id across every structure keyed by it. | ||
| * | ||
| * A placeholder id (from ensureAssistantMessage / startAssistantMessage) is | ||
| * renamed to the real provider id on TEXT_MESSAGE_START. Every structure that | ||
| * holds that id must move together — miss one and the data it holds is | ||
| * silently orphaned onto the vanished id (e.g. a tool result whose | ||
| * toolCallToMessage entry still points at the placeholder). This is the same | ||
| * id-keyed set enumerated by pruneToMessages() and reset(). | ||
| */ |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Docstring references a nonexistent method.
pruneToMessages() doesn't exist in this class — the method that enumerates the same id-keyed maps is removeMessagesAfter() (see its comment at Line 424-431). Fix the reference so future readers can actually find the mirrored cleanup logic.
📝 Proposed fix
- * silently orphaned onto the vanished id (e.g. a tool result whose
- * toolCallToMessage entry still points at the placeholder). This is the same
- * id-keyed set enumerated by pruneToMessages() and reset().
+ * silently orphaned onto the vanished id (e.g. a tool result whose
+ * toolCallToMessage entry still points at the placeholder). This is the same
+ * id-keyed set enumerated by removeMessagesAfter() and reset().📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| /** | |
| * Rename a message id across every structure keyed by it. | |
| * | |
| * A placeholder id (from ensureAssistantMessage / startAssistantMessage) is | |
| * renamed to the real provider id on TEXT_MESSAGE_START. Every structure that | |
| * holds that id must move together — miss one and the data it holds is | |
| * silently orphaned onto the vanished id (e.g. a tool result whose | |
| * toolCallToMessage entry still points at the placeholder). This is the same | |
| * id-keyed set enumerated by pruneToMessages() and reset(). | |
| */ | |
| /** | |
| * Rename a message id across every structure keyed by it. | |
| * | |
| * A placeholder id (from ensureAssistantMessage / startAssistantMessage) is | |
| * renamed to the real provider id on TEXT_MESSAGE_START. Every structure that | |
| * holds that id must move together — miss one and the data it holds is | |
| * silently orphaned onto the vanished id (e.g. a tool result whose | |
| * toolCallToMessage entry still points at the placeholder). This is the same | |
| * id-keyed set enumerated by removeMessagesAfter() and reset(). | |
| */ |
🤖 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/ai/src/activities/chat/stream/processor.ts` around lines 764 - 773,
The docstring on the message-id renaming helper refers to a nonexistent
pruneToMessages() method, which will mislead readers. Update that reference in
the comment for the id-renaming logic around the processor’s message-keyed map
cleanup so it points to the actual enumerating method, removeMessagesAfter(),
and keep the wording aligned with the mirrored cleanup behavior described by
reset().
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.
What
Two related changes, surfaced by wiring structured output through the Claude Code harness:
1.
@tanstack/ai-claude-code— stream tool-call arguments (feat, patch)The stream translator only forwarded
text_delta/thinking_deltapartial events, so atool_useblock surfaced as a single completedTOOL_CALL_START/ARGS/ENDat the end of the message. It now also translates the SDK's partial tool input —content_block_start(tool_use) →TOOL_CALL_START,input_json_delta→ incrementalTOOL_CALL_ARGS,content_block_stop→TOOL_CALL_END— matching how@tanstack/ai-anthropic/@tanstack/ai-openaialready stream tool args. The completeassistantmessage dedupes anything already streamed via partials (tracked by tool-call id, mirroring the existing text/thinking message-id dedup). WithstreamPartials: false, tool calls still emit whole from the complete message.This unblocks streaming structured output on the Claude Code harness: routing structured data through a tool call now paints incrementally as the model writes it, instead of arriving atomically at run end.
Also in this commit:
TOOL_CALL_STARTcarriesparentMessageId(mirroring the partial text path and every other adapter), so a tool-first stream groups the call under the correct message.catchflushes the in-flight partial tool call (mirroringhandleResult) so itsTOOL_CALL_STARTis still paired with aTOOL_CALL_END+ interruptedTOOL_CALL_RESULT, upholding the module's documented invariant.2.
@tanstack/ai— atomic message-id rename (fix, patch)Independently reachable from any adapter that omits
parentMessageId(optional per the AG-UI spec — Mistral omits it today). When a tool call streams before any text, theStreamProcessorcreates a placeholder assistant message and renames it to the real provider id onTEXT_MESSAGE_START. That rename updatedmessages/messageStates/activeMessageIdsbut nottoolCallToMessage/structuredMessageIds/structuredOutputUpdateBatches, so aTOOL_CALL_RESULTarriving after the rename resolved to the vanished placeholder id and was silently dropped. The rename now goes through a singlerenameMessageId()seam that remaps every id-keyed structure — the same set already enumerated bypruneToMessages()andreset(). Complements the adapter-side fix in #480.Why
Routing structured output through a tool call is the sanctioned way to stream it on the Claude Code harness (native
structured_outputonly appears in the finalResultMessage, not as deltas). Getting there surfaced (1) the translator droppinginput_json_delta, and (2) the processor's rename orphaning a tool result whenparentMessageIdis absent.Test plan
@tanstack/ai-claude-code— newtranslatetests: partial tool-arg streaming + dedup, abort-mid-tool-args pairs an interrupted result,parentMessageIdgrouping, sequential tool blocks, zero-argSTART→END. 6 files / 39 tests pass.@tanstack/ai— newstream-processorregression test: tool-first withoutparentMessageId→ rename → tool result attaches (was silently dropped); the existing provided-parentMessageIdtest still passes (renameMessageIdno-ops when ids match). 63 files / 1133 tests pass.test:types,test:eslint(0 errors),buildclean.@tanstack/ai-claude-codepatch,@tanstack/aipatch).Closes #901.
Summary by CodeRabbit