Studio: pin why a Streamdown remount keeps its highlighted code by danielhanchen · Pull Request #9048 · unslothai/unsloth · GitHub
Skip to content

Studio: pin why a Streamdown remount keeps its highlighted code - #9048

Merged
danielhanchen merged 4 commits into
mainfrom
pin-shiki-remount
Aug 18, 2026
Merged

Studio: pin why a Streamdown remount keeps its highlighted code#9048
danielhanchen merged 4 commits into
mainfrom
pin-shiki-remount

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

What this is

The incremental Markdown cache bumps renderGeneration when dropping retained blocks cannot be signalled through the Markdown string, and that generation is part of the <Streamdown> React key, so the block subtree unmounts and remounts. That was read as losing Shiki's highlight state and flashing every fence in the reply back to unstyled code.

It does not, and the reason is worth holding in place. This PR adds two tests and changes no behaviour.

Why the remount is safe

None of the highlighting state lives in the component tree.

  • The wrapper in code-plugin.ts is built once at module scope in markdown-text.tsx, so its per-fence slots outlive any remount.
  • @streamdown/code keeps its highlighter instances and its tokenized results in module-scope Maps and answers a repeat request inline.

So a remount re-asks for tokens it already has and gets them back in the same tick.

Measured on this branch's base: with a 40 line TypeScript fence the cold mount answers through its callback, and a fresh plugin standing in for the remount answers synchronously in 0.010 ms with the same 40 token lines.

Upstream already reported and fixed this

  • vercel/streamdown#186, "Code block flickering in virtual lists due to async Shiki highlighting", names the failure precisely: "No caching between mounts: Highlighter state is lost when components unmount."
  • Closed by vercel/streamdown#240, merged 2025-11-21, which rebuilt the code blocks. The pinned @streamdown/code 1.1.1 carries that rebuild.
  • Nothing equivalent is open on shikijs/shiki, where the highlighter is an object the consumer owns and caching it is the consumer's job.

How often the bump actually fires

On this base, driving the cache through the renderer's own stabilizeStreamingMarkdown(preprocessLaTeX(text), true) pipeline over twelve streamed fixtures covering prose, fenced code, inline and display LaTeX, currency dollars, and spans that stay open across thousands of characters:

arrival size updates non-prefix updates retained-block drops generation bumps
1 character 46,181 26 1 0
4 characters 11,550 14 1 0

The generation never moved off zero, so there is no remount to remove and nothing for a rendering benchmark to show. That is why this PR changes no behaviour: the key and the generation counter are left alone.

The tests

Both are in studio/frontend/tests/code-plugin-remount.test.ts, and each was shown to fail on a tree broken for it and only for it.

  1. a remount gets already highlighted code back in the same tick drives a real highlight, then asks a second plugin instance for the same fence and requires it to answer inline. The fixture fence is deliberately past the wrapper's MIN_INCREMENTAL_CHARS, asserted in the test, so it exercises the per-fence slot path rather than the small fence shortcut.
    Broken tree: return null instead of the synchronous cache hit from the wrapper's dispatch. Fails with "a remount had to wait for the highlighter again, so every fence in the reply would flash back to unhighlighted code". The second test still passed.

  2. the chat renderer builds its code plugin once, outside the component walks markdown-text.tsx and requires the single createCodePlugin call to sit at module scope.
    Broken tree: move the plugin construction into a function. Fails with "the code plugin is built inside a component, so its incremental fence slots are discarded on every remount". The first test still passed.

Neither test passes both ways.

Checks

  • npm test: 2938 passed, 0 failed.
  • npm run typecheck: clean.
  • npx biome check on the new file: only the three noNodejsModules warnings that every test file in the repo carries.

The incremental Markdown cache bumps `renderGeneration` when dropping
retained blocks cannot be signalled through the Markdown string, and that
generation is part of the <Streamdown> React key, so the block subtree
unmounts and remounts. That was read as losing Shiki's highlight state and
flashing every fence in the reply back to unstyled code. It does not, and
the reason is worth holding in place.

None of the highlighting state lives in the component tree. The wrapper in
code-plugin.ts is built once at module scope in markdown-text.tsx, so its
per-fence slots outlive any remount, and `@streamdown/code` keeps its
highlighter instances and its tokenized results in module-scope Maps and
answers a repeat request inline. A remount therefore re-asks for tokens it
already has and gets them back in the same tick.

Upstream reported and fixed exactly this. vercel/streamdown issue 186,
"Code block flickering in virtual lists due to async Shiki highlighting",
names "No caching between mounts: Highlighter state is lost when components
unmount", and was closed by PR 240, merged 2025-11-21, which rebuilt the
code blocks. The pinned `@streamdown/code` 1.1.1 carries that rebuild.
Nothing equivalent is open on shikijs/shiki, where the highlighter is an
object the consumer owns and caching it is the consumer's job.

Measured on this branch's base: with a 40 line TypeScript fence the cold
mount answers through its callback, and a fresh plugin standing in for the
remount answers synchronously in 0.010 ms with the same 40 token lines.

Two tests pin the halves of that. One drives a real highlight and then asks
a second plugin instance for the same fence, which has to answer inline.
The other walks markdown-text.tsx and requires the single
`createCodePlugin` call to sit at module scope rather than inside a
component. Both fail on a tree broken for them and only for them: returning
null instead of the synchronous cache hit from the wrapper's dispatch fails
the first and leaves the second passing, and moving the plugin construction
into a function fails the second and leaves the first passing.

No behaviour change. The key and the generation counter are left alone. On
this base, driving the cache through the renderer's own
stabilizeStreamingMarkdown(preprocessLaTeX(text), true) pipeline over
twelve streamed fixtures covering prose, fenced code, inline and display
LaTeX, currency dollars, and spans that stay open across thousands of
characters: at one character per arrival, 46,181 updates produced 26
non-prefix updates, 1 retained-block drop and 0 generation bumps; at four
characters per arrival, 11,550 updates produced 14 non-prefix updates, 1
drop and 0 bumps. The generation never moved off zero, so there is no
remount to remove and nothing for a rendering benchmark to show.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 28d624ab90

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

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".

Comment on lines +83 to +84
const remounted = createCodePlugin({ themes });
const afterRemount = remounted.highlight(options);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reuse the mounted plugin when simulating a remount

A Streamdown remount does not execute the module again, so markdown-text.tsx continues passing the same module-scoped plugin object and its slotsByKey cache. Creating a fresh plugin here instead tests the unrelated implementation detail that @streamdown/code caches tokenization across plugin instances; if upstream safely moves that cache into each plugin instance, actual remounts remain synchronous but this test fails and blocks the upgrade. The remount assertion should call the already mounted plugin, while the separate scope test ensures that object survives.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed, and the check is stronger than it looks. A remount rebuilds the component tree, not the module, so markdown-text.tsx keeps handing Streamdown the same module-scope plugin. Building a fresh wrapper here meant the synchronous answer came from the module-scope caches inside @streamdown/code rather than from our own fence slots.

Measured both directions on the tree at 514d411:

  • Moving createCodePlugin inside MarkdownTextImpl, which is the exact regression this test's failure message describes, left this test green. Only the module-scope walker caught it.
  • Patching @streamdown/code so its caches are per-instance, which is a change upstream could make safely, turned this test red while production stayed synchronous.

So it was failing on a healthy tree and passing on a broken one. Fixed in 1b686c3 by reusing the mounted plugin.

Comment on lines +121 to +124
ts.isFunctionDeclaration(node) ||
ts.isFunctionExpression(node) ||
ts.isArrowFunction(node) ||
ts.isMethodDeclaration(node);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Distinguish module initialization from function syntax

This treats every call lexically contained in a function as runtime component construction. For example, extracting the current setup into function buildCodePlugin() { return createCodePlugin(...); } and invoking buildCodePlugin() once in the module initializer still preserves the plugin across every remount, but this walker reports the inner call as non-module-scoped and fails the test. Check where the factory is invoked rather than merely whether the constructor expression appears inside any function.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The scenario is real and I reproduced it: a helper invoked at module init preserves the plugin across remounts, and the walker still fails it. Keeping the walker as is.

It has no false negative, which is the property that matters here. Moving the call inside MarkdownTextImpl goes red, and adding a second module-scope build goes red on the count assertion. The only cost is a misleading message on a refactor nobody is proposing, and teaching the walker call-graph reachability to cover it is machinery for a case that does not exist.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: db609857d4

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

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".

// through the per-fence slot path a streaming reply uses rather than the small
// fence shortcut straight to the underlying plugin.
const LINES = 90;
const MIN_INCREMENTAL_CHARS = 2000;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Tie the threshold assertion to the production value

When MIN_INCREMENTAL_CHARS in code-plugin.ts is increased beyond this fixture's length, this independent copy remains 2000, so the assertion still passes while highlight takes the small-fence shortcut. The test would then stop exercising the per-fence slot path and could miss the wrapper regression it is intended to pin; derive the boundary from production code or construct the fixture so the exercised path is verified directly.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed, and it was hiding a real gap. The assertion compared against a local copy of the constant, so on the tree at 514d411 I raised MIN_INCREMENTAL_CHARS in code-plugin.ts to 100000 and every one of the four highlight calls took the small-fence shortcut instead of the slot path, with both tests still green. The test silently stopped covering the path its own comment says it covers.

Fixed in 1b686c3: the constant is exported and imported. The same break now fails with "the fixture fence dropped below the incremental threshold, so this test no longer covers the slot path".

pre-commit-ci Bot and others added 2 commits August 18, 2026 03:13
The remount case built a fresh plugin wrapper, so the only thing keeping
its answer synchronous was the module-scope cache inside @streamdown/code.
Production hands Streamdown the same module-scope plugin across a remount,
so reuse the mounted one and the test pins our wrapper instead.

The threshold guard compared against a local copy of MIN_INCREMENTAL_CHARS,
so raising the production value to 100000 sent every fence down the small
fence shortcut with the test still green. Export the constant and import it.

Await the cold ask's callback rather than polling highlight(), so the warm
up loop no longer satisfies the remount assertion on its own.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

@danielhanchen
danielhanchen merged commit 17665d4 into main Aug 18, 2026
42 of 46 checks passed
@danielhanchen
danielhanchen deleted the pin-shiki-remount branch August 18, 2026 05:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant