feat(ai-claude-code): stream tool-call arguments (+ fix processor id-rename orphaning tool results)#904
Conversation
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`, `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. The partial TOOL_CALL_START also carries parentMessageId (mirroring the partial text path and every other adapter) so a tool-first stream groups the call under the correct message, and an abort mid-tool-args flushes the in-flight partial tool call so its START is still paired with an END and an interrupted RESULT.
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.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (7)
📝 WalkthroughWalkthroughThis PR fixes two issues: StreamProcessor now uses a centralized ChangesStreamProcessor Message-ID Rename Fix
Claude Code Partial Tool-Call Argument Streaming
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant SDK as Claude Agent SDK
participant Translator as translateSdkStream
participant Consumer as AG-UI Event Stream
SDK->>Translator: content_block_start (tool_use)
Translator->>Translator: set partial tool-use state
Translator-->>Consumer: TOOL_CALL_START (parentMessageId)
SDK->>Translator: content_block_delta (input_json_delta)
Translator->>Translator: accumulate partial args
Translator-->>Consumer: TOOL_CALL_ARGS (delta)
SDK->>Translator: content_block_stop
Translator->>Translator: closePartialToolUse() parses JSON
Translator-->>Consumer: TOOL_CALL_END (input)
SDK->>Translator: complete assistant message (tool_use block)
Translator->>Translator: skip duplicate via streamedToolCallIds
sequenceDiagram
participant Stream as AG-UI Stream
participant Processor as StreamProcessor
Stream->>Processor: TOOL_CALL_START (no parentMessageId)
Processor->>Processor: create placeholder assistant message
Stream->>Processor: TEXT_MESSAGE_START (real messageId)
Processor->>Processor: renameMessageId(placeholder, real id)
Processor->>Processor: update messages, toolCallToMessage, structured maps
Stream->>Processor: TOOL_CALL_RESULT
Processor-->>Stream: attach result to renamed message
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ 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 |
What
Two related changes, surfaced while wiring structured output through a tool call on the
@tanstack/ai-claude-codeharness adapter.1.
@tanstack/ai-claude-code: stream tool-call arguments (feat, patch)The stream translator only forwarded
text_deltaandthinking_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_startfor tool_use intoTOOL_CALL_START,input_json_deltainto incrementalTOOL_CALL_ARGS,content_block_stopintoTOOL_CALL_END), matching how the@tanstack/ai-anthropicand@tanstack/ai-openaiadapters already 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 lets structured output routed through a tool call paint incrementally as the model writes it, instead of arriving all at once 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_ENDand an 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, andactiveMessageIdsbut nottoolCallToMessage,structuredMessageIds, orstructuredOutputUpdateBatches, 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 this harness (native
structured_outputonly appears in the final result, 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 covering partial tool-arg streaming plus dedup, abort mid-tool-args pairing an interrupted result,parentMessageIdgrouping, sequential tool blocks, and zero-argSTART/END. 6 files, 39 tests pass.@tanstack/ai: newstream-processorregression test (tool-first withoutparentMessageId, then rename, then the tool result attaches, which was silently dropped); the existing provided-parentMessageIdtest still passes becauserenameMessageIdno-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
New Features
Bug Fixes
Tests