Studio: stop every message part re-rendering on each streaming chunk - #9014
Conversation
A streaming assistant reply rebuilt all of its already-finished parts on every chunk, so the per-chunk cost grew with the length of the reply. MessagePrimitivePartByIndex is memoized, and its comparator checks the components fields one at a time rather than comparing the object as a whole: prev.components?.Text === next.components?.Text && ... prev.components?.tools === next.components?.tools && Every field thread.tsx passes is a module-level component, so all of them compare equal across renders. The exception was tools, an object literal built inline in the JSX, which meant a fresh identity on every render. That single mismatch failed the comparator and re-rendered every part of the message. Both maps move to module scope. THREAD_MESSAGE_COMPONENTS is declared after the three components it names, because a module-scope initializer runs at import time and would otherwise read them in their temporal dead zone. The guard test covers the two ways this regresses. It fails on the tree before this change, and it also pins the upstream assumption: if an assistant-ui upgrade stops comparing components.tools by identity, the test fails and says to re-measure rather than leaving behind a hoist and a comment that no longer describe what the library does. npm test 2939 passed, 0 failed. typecheck clean. biome adds no new errors on either file.
The ThreadPrimitive.Messages half of this change is reverted. Hoisting the map there cannot help, and it collides with #9042 which fixes that call properly. Reading the primitive settles it. Given a components map, assistant-ui builds: children: () => <ThreadMessageComponent components={components} /> so the per-message element always carries a props object and never reaches the propless bail-out in RenderChildrenWithAccessor. A stable map only lets the outer memo bail out, which is the cheap part. #9042 moves the call to the children form returning one shared propless element, which does reach the bail-out, and measures 7.44x on a delete at 300K characters. Its test asserts ThreadPrimitive.Messages carries no components prop at all, which is the exact opposite of what a hoist looks like, so the two could not both land. MessagePrimitive.Parts is unaffected and keeps the fix. MessagePrimitivePartByIndex compares components field by field, checking components.tools by identity, so the inline tools literal really did defeat it on every render. That part is unchanged and still measured. The test is scoped to MessagePrimitive.Parts and carries the reason ThreadPrimitive.Messages is excluded, so nobody re-adds the hoist there on the strength of the same reasoning. Discrimination: re-inlining the Parts literal fails 2 of the 3 tests. The third pins the upstream comparator and passes on both trees by design, as before. npm test 2,939 passed, 0 failed. typecheck clean.
|
Narrowed this after #9042 landed, and the reason is worth recording rather than leaving in the diff. The children: () => _jsx(ThreadMessageComponent, { components: components })so the per-message element always carries a props object and never reaches the propless bail-out in The two were also mutually exclusive: #9042 asserts
The test is now scoped to Discrimination unchanged in kind: re-inlining the Parts literal fails 2 of the 3 tests. The third pins the upstream comparator and passes on both trees by design, as stated before.
|
The upstream comparator is read with readFileSync, so a half-installed or relocated node_modules surfaces as a bare ENOENT stack inside this test. That happened once and read as a real regression in the code under test until it was disbelieved. Wrap the read and say plainly that it is an install problem. Also make the inline-literal assertion message one template literal, which drops the useTemplate error the file was adding to biome.
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
Before and after, one scene driven against both buildsBEFORE is the merge base Settled
Both halves show the tool card in its completed state, with the tick, the argument row That earlier part is the point of this shot. Hoisting the components map makes an upstream memo start bailing out, so the failure mode to look for is not a visible redesign but a finished part going stale. It has not. Mid-stream
The same tool card while the run is still going, showing the running spinner and the collapsed chevron on both sides, with The two settled screenshots are byte-identical. Facts agree: |
1 participant
A streaming assistant reply rebuilt all of its already-finished parts on every chunk, so the per-chunk cost grew with the length of the reply. This is one of the paths behind reports of Studio getting sluggish during long generations with code cells and text.
Cause
MessagePrimitivePartByIndexis memoized, and its comparator checks thecomponentsfields one at a time rather than comparing the object as a whole. From@assistant-ui/core/dist/react/primitives/message/MessageParts.js:Every field
thread.tsxpassed is a module-level component, so all of them compare equal across renders. The exception wastools, an object literal built inline in the JSX:toolsgot a fresh identity on every render. That single mismatch failed the comparator, so every part of the message re-rendered on every render ofAssistantMessage, which during streaming is every chunk.There is a second effect in the same file, and it is roughly half of what this fixes.
MessagePrimitivePartsCompatbuilds the whole element array inside auseMemowhose dependency list is[messageRanges, components, contentLength]. An inlinecomponentsinvalidated that memo on every render too, rebuilding the entire array independently of the per-part comparator.Note that hoisting only helps because of how that comparator is written. The
MessagePrimitivePartswrapper builds a newmergedobject each render regardless, so an identity comparison on the map as a whole would never have held. It is the per-field check that makes stable fields worth having.Change
One map moves to module scope:
ASSISTANT_PART_COMPONENTS, the map handed to<MessagePrimitive.Parts>inAssistantMessage.Every field in it is already a module-level binding. The five non-tool entries are ES module imports; the eight tool entries are the existing module-scope
*Confirmableconsts.withToolConfirmationcloses over nothing but itsComponentargument and calls no hooks at wrap time, so the literal already had identical contents on every render and this changes its identity only. Nothing here needed auseMemo.The other two call sites are deliberately left alone:
<MessagePrimitive.Parts />inUserMessagepasses nocomponentsprop at all.undefined === undefinedsatisfies the comparator, so it never had the problem.componentsliteral onThreadPrimitive.Messagesis left inline.ThreadPrimitiveMessageByIndexis memoized withisComponentsSame, a field-by-field structural compare, so a fresh literal whose fields are stable already compares equal and the memo already bails out. Hoisting there buys nothing, and it would collide with Chat: stop a message delete from re-rendering the whole thread #9042.Testing
tests/thread-part-components-stable.test.ts, three tests. Each was checked against the tree before this change rather than assumed to discriminate:MessagePrimitive.Partscomponents map is an inline object literalcomponents.toolsby identityThe third passes on both sides deliberately. It is not a regression guard for this change; it pins the upstream behaviour this fix depends on, so that an assistant-ui upgrade which stops comparing
components.toolsby identity fails the test and says to re-measure, rather than leaving behind a hoist and a comment that no longer describe what the library does.The first two were also checked against three further broken trees, since a source-text assertion is easy to write so that it passes on the broken tree too: the literal restored inline, everything hoisted except
tools, and the hoisted const spread back into the JSX with{...ASSISTANT_PART_COMPONENTS}. Two of the three tests fail on each.npm test3,482 passed, 0 failed.npm run typecheckclean.biome checkonthread.tsxis unchanged; the new test file carries 6 diagnostics, all of the kind the rest oftests/already carries (noNodejsModules,useTopLevelRegex).Measured
Live generation is not usable for a before-and-after here: the GGUF loads with
--spec-type ngram-mod --parallel 4, so two runs of the same prompt at temperature 0 produced 770 and 814 characters. Instead the real SSE stream fromunsloth/Qwen3.5-2B-MTP-GGUF:UD-Q4_K_XLat 4096 context was recorded (600 chunks, 145 KB, reasoning plus three python tool calls plus interleaved text) and replayed byte-for-byte, with its original pacing, into a build of each tree. Counters are non-memoized module-scope wrappers identical on both sides.AssistantMessagerendersReproduced three times on Chromium, once each on Firefox and WebKit.
One qualification worth stating: the large ratio needs a message with several parts mounted at once, which in practice means tool calls. For a plain reasoning-plus-text reply the reasoning card collapses when it finishes, so only one part is mounted and the before side already does about as little work as the after side. The gain is real for the "code cells and text" case this started from, and close to nothing for a short prose answer.
Rendered output is unchanged. Comparing the assistant message body across the same replay,
textContentis identical at 2,506 characters andinnerHTMLis identical at 74,510 characters once Radix auto-ids are normalized (those also differ between two runs of the same tree). That holds on Chromium, Firefox and WebKit, and under a mid-stream interrupt at chunk 300 (930 characters, identical).Because the risk in a change like this is over-memoizing rather than under-memoizing, the cases where a part must still update were driven explicitly on the built frontend: a tool call going running to complete, parts appended mid-stream, regenerate, assistant branch switch both ways, user branch switch both ways, editing a user message, editing an assistant response in place, and the tool confirmation gate. All render; none goes stale. A thread written by the previous build also reopens byte-identically on this one.
WebKit here is Playwright's, which is a proxy for the webviews Desktop embeds rather than those webviews. macOS and Windows runners were not used; the change is a JavaScript object hoist with no platform-dependent semantics.