Studio: incrementally tokenize streaming code fences by oobabooga · Pull Request #8935 · unslothai/unsloth · GitHub
Skip to content

Studio: incrementally tokenize streaming code fences - #8935

Merged
oobabooga merged 16 commits into
unslothai:mainfrom
oobabooga:worktree-codeblock-highlight-perf
Aug 19, 2026
Merged

Studio: incrementally tokenize streaming code fences#8935
oobabooga merged 16 commits into
unslothai:mainfrom
oobabooga:worktree-codeblock-highlight-perf

Conversation

@oobabooga

Copy link
Copy Markdown
Member

Main already renders large streaming fences incrementally. Above 2,000 characters, #7537 reuses highlighted lines, displays the new tail as plain text, and refreshes every 250 ms. However, each refresh still asks Shiki to tokenize the entire block. If that takes longer than 250 ms, the throttle stops helping. Shorter fences are tokenized in full on every frame.

This PR makes the tokenization incremental too. It caches completed-line tokens and the grammar state after them, so each refresh processes only new text. The existing plain-tail cadence remains unchanged, so user-visible rendering does not change.

What changed

  • Continue tokenization from Shiki's saved grammar state instead of processing the full fence again.
  • Preserve the existing plain-tail rendering between refreshes for fences over 2,000 characters.
  • Cache completed fences for thread re-mounts, bounded at 512 fences and 512,000 source characters.
  • Match whole-document Shiki output for LF and CRLF input.
  • Keep every callback waiting on a grammar load, without merging prefix-related sibling blocks or evicting pending fences.
  • Release pending entries after a failed grammar load so they do not pin the cache or fire during a later retry.

Measurements

Production builds were run alternately on the same machine, with three paired runs per row. The table reports medians from identical streamed updates.

Fixture CPU throttle Main thread CPU, main to PR Change Long tasks, main to PR Frames taking >33 ms, main to PR
TypeScript fence, 7.4 KB 1x 7,082 to 3,428 ms 52% lower 3,183 to 698 ms 32 to 7
TypeScript fence, 7.4 KB 4x 105,505 to 13,954 ms 87% lower 93,816 to 3,288 ms 418 to 31
HTML fence, 7.3 KB 1x 3,270 to 2,955 ms 10% lower 458 to 347 ms 6 to 4
Single-line JSON, 8.6 KB 1x 3,341 to 3,212 ms 4% lower 167 to 209 ms 10 to 13
Single-line JSON, 8.6 KB 4x 21,148 to 17,144 ms 19% lower 10,105 to 6,828 ms 127 to 75

Absolute totals vary between sessions, but the paired TypeScript reduction repeated at 40% to 52% at 1x and 86% to 87% at 4x.

Tokenizer instrumentation confirms the source of the improvement:

Fixture Characters passed to Shiki, main to PR Reduction
Two Python fences, 10 KB 374,170 to 17,530 21.3x
TypeScript fence, 7.4 KB 196,104 to 10,880 18.0x
HTML fence, 7.3 KB 234,526 to 12,245 19.2x
Eight short fences, 6.5 KB 128,049 to 18,825 6.8x
Single-line JSON, 8.6 KB 283,182 to 269,130 No meaningful change

A replay through the real chat UI showed the same pattern. TypeScript main-thread CPU fell from 6,420 to 4,498 ms at 1x and from 81,348 to 19,646 ms at 4x. HTML moved from 4,532 to 4,309 ms. The gain depends on grammar cost and completed lines; one long line still relies on the existing throttle.

Rendering parity

  • Settled DOM output was byte-identical between builds and to mounting the completed reply in one pass for Python, TypeScript, HTML, eight short fences, and single-line JSON.
  • Settled plugin tokens matched whole-document codeToTokens output.
  • Mid-stream rendering also matched. TypeScript showed a plain tail on all 106 measured frames in both builds, while single-line JSON remained fully plain between refreshes in both builds.

Tradeoff

The cache is bounded, unlike the @streamdown/code cache. Mounted blocks are unaffected; only evicted blocks that later remount are tokenized again. Re-mounting 250 fences took 0.7 ms on both builds, while 400 fences beyond the budget took 974 ms here versus 1.9 ms on main. This is the cost of avoiding unbounded retention of streamed prefixes.

Out of scope

preprocessLaTeX can still rewrite LaTeX-like text inside an open streaming fence until its closing delimiter arrives. That behavior already exists on main and is separate from highlighting.

Checks

  • npm test: 2,717 passed, including 15 focused highlighting tests
  • npm run typecheck
  • npx eslint on the changed files
  • npm run build
  • Production renderer, tokenizer instrumentation, plugin-only driver, and real chat replay against fd6ab9a

chatgpt-codex-connector[bot]

This comment was marked as resolved.

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
chatgpt-codex-connector[bot]

This comment was marked as resolved.

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
chatgpt-codex-connector[bot]

This comment was marked as resolved.

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
chatgpt-codex-connector[bot]

This comment was marked as resolved.

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@oobabooga

Copy link
Copy Markdown
Member Author

@codex review

@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: 99a34e7f71

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

// tildes, then spaces. It also starts a line, so the body it leaves behind ends
// at a newline. The run can be short, because the cached code is the last text
// tokenized and the run may still have been arriving then.
const CLOSING_FENCE = /^ {0,3}(?:`+|~+)[ \t]*$/;

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 Match closing suffixes to the opening fence marker

When a large tilde-opened fence ends with a backtick-only content line (or vice versa) and a later same-language sibling contains only the preceding prefix, this pattern accepts the wrong marker family as a closing delimiter. findFence then shares the cache entry, and the shorter sibling can clear the longer sibling's pending refresh, leaving the longer block with a plain approximate tail. Markdown requires the closing fence to use the same marker as the opener, so preserve that marker for this check or keep these siblings separate.

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 plugin never sees the opening marker. HighlightOptions is {code, language, themes}, and the info string is identical for a tilde fence and a backtick fence, so there is nothing to match the suffix against. Keeping the siblings separate is not available either: a body of P plus a bare marker line shortening to P is byte-identical whether one fence just closed or two fences share a prefix, so no content-only rule tells them apart.

I did reproduce the mechanism, and it needs the same repeated render of unchanged code as the identical-fence item. One document-order pass over both fences returns exact tokens for each on this branch, and that pass is what the consumer performs, since its effect is keyed on the code. The identical-fence fix removed the cancellation for a caller whose code already equals the entry's; here the caller's code is shorter, so it does not reach. Leaving this rather than adding a check that cannot be correct.

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. You're on a roll.

Reviewed commit: 99a34e7f71

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

The incremental tokenization here is correct, including for the cases the
existing tests do not reach. These add coverage for the three gaps.

Embedded grammars. Every existing fixture is python, json or typescript,
where one grammar spans the whole fence, so resuming from a saved
GrammarState is never asked to restore a stack more than one level deep.
Five fixtures that are: HTML with embedded script and style, TSX with JSX
children, markdown with a multi-line HTML comment, shell heredocs with
runtime terminators, and nested template literals.

The engine. The desktop app's CSP is default-src 'self' with no
wasm-unsafe-eval, and the packaged frontend is this same bundle, so
switching to shiki's Oniguruma engine would pass every test and every
browser and fail only in the shipped desktop app. Nothing in CI builds
the desktop bundle and drives its webview, so this pins the JS engine.

The cache budget. MAX_CACHED_CHARACTERS is 512,000 and the largest
existing fixture is about 36 KB, so the character-budget branch, its
never-evict-the-only-fence carve-out and the running character count
were all unreachable.

Each test was checked against a deliberately broken tree rather than
assumed to discriminate, and two did not until they were rebuilt:

- The evict-then-resume test spent 22 of 23 comparisons on the throttled
  approximation, which renders the uncommitted tail plain and never
  reads the grammar state, so it could not fail however the resume broke.
  It now waits out REFRESH_MS between updates.
- The same test used a python triple-quoted string, whose tokens are
  byte-identical with the resumed state dropped. It now uses HTML with an
  embedded script, where the grammar is several levels deep at the resume
  point, and fails at character 2907 of 4414 when the state is dropped.

Tests only, no source changes. 32 passed, 0 failed on this head.
@danielhanchen

Copy link
Copy Markdown
Member

Reviewed this closely and it holds up. The mechanism is sound: the object being retained across updates is shiki's public GrammarState, and keeping it without copying is safe because vscode-textmate's StateStack is a persistent immutable linked list, so structural sharing needs no clone. The CRLF trim is necessary, since splitLines only drops \r when a \n follows. Offset reconstruction via shiftLine is consistent.

I have pushed tests to the branch rather than only describing them, f961eb48a. Tests only, no source changes. 32 passed, 0 failed on this head.

What was not covered

Embedded grammars. Every existing fixture is python, json or typescript, where a single grammar spans the whole fence, so resuming from a saved state is never asked to restore a stack more than one level deep. That is the case most likely to diverge, and a divergence there produces plausible tokens with wrong scopes rather than an obvious break. Added five: HTML with embedded script and style, TSX with JSX children and an expression container, markdown with a multi-line HTML comment, shell heredocs with runtime terminators, and nested template literals. All match whole-document codeToTokens at every prefix, with extra cut points immediately before and after each delimiter that pushes or pops a grammar.

The regex engine. This is the one desktop-specific hazard in the file and it is worth stating explicitly: tauri.conf.json sets default-src 'self' with no wasm-unsafe-eval, and frontendDist is the same frontend/dist the browser gets. Switching to shiki's Oniguruma engine would pass every test here and in every browser, and fail only inside the packaged desktop app. No CI job builds the desktop bundle and drives its webview, so nothing would catch it. The new test pins createJavaScriptRegexEngine.

The cache budget. MAX_CACHED_CHARACTERS is 512,000 and the largest existing fixture is about 36 KB, so the character-budget branch, the never-evict-the-only-fence carve-out and the running character count were all unreachable. They are correct; they were just untested. Driving 600 KB through confirms eviction fires under MAX_FENCES, and that a single oversized fence is retained rather than evicting itself.

Two things worth knowing about testing this file

Each test was checked against a deliberately broken tree rather than assumed to discriminate, and two did not until rebuilt. Both traps are easy to fall into again, so they are written into the test comments with the measured numbers.

  1. A tight await loop never leaves the approximate path. Once a fence clears MIN_INCREMENTAL_CHARS, two updates inside REFRESH_MS return approximateResult, which reuses committed lines and renders the uncommitted tail plain. It never reads fence.state. Measured 22 of 23 comparisons taking that path, so the test could not fail however the resume was broken. It now waits out the throttle.

  2. A python triple-quoted string cannot detect a lost grammar state. Its tokens are byte-identical with the state dropped. The rebuilt test uses HTML with an embedded script, where the grammar is several levels deep at the resume point, and it diverges at character 2907 of 4414 when the state is dropped.

The existing fixtures all sit under MIN_INCREMENTAL_CHARS, so they always take the exact path and never exercise the throttle at all. Neither side of that boundary is currently covered by the suite.

Not changed, raised for your call

The React tree is still rebuilt every frame. Tokenization is incremental now, but tokenize returns a fresh result object each call and Streamdown's body memo compares by identity, so the whole fence's spans are recreated per frame regardless. On a 600 line fence that is several thousand elements. This is the natural follow-up to the work here and probably the larger remaining win on long code cells, but it is a change to the boundary with Streamdown rather than to this file, so I have not touched it.

Fence identity remains a heuristic. The plugin never receives the opening marker, so it cannot apply CommonMark's rule that a closer matches the opener's family and is at least as long. shedsClosingRun stands in for information the renderer has and does not pass down. This is a design question rather than a defect, and worth a separate look.

The unstyled tail is visible. During fast streaming of a large fence the tail renders plain for up to 250 ms. That is the intended tradeoff of the throttle, not a bug, but it is the behaviour users will notice on exactly the long code cells this targets, and it is currently undocumented.

@danielhanchen

Copy link
Copy Markdown
Member

Picking this back up now that main has moved a long way underneath it. Three things: a current answer on the red check, the conflict, and two effects of this change that I do not think either of us had noticed.

The failing check was stale, and the current answer is green

Repo tests (CPU) and Chat UI Tests were both from Aug 16 and their logs have since expired, so I rebuilt the merge locally against current main and ran the real gates rather than quote them.

  • Chat UI Tests was never a failure. The job's own conclusion is cancelled, superseded by a newer push. gh pr checks renders that as fail.
  • Repo tests (CPU) is a real red, but not from this PR.

Full frontend suite on the merged tree: 3790 tests, 3788 pass, 2 fail. Both failures are unrelated to this change and neither test imports code-plugin.

test on this PR on pristine main (01cf3728c) verdict
queued-model-capabilities fail fail pre-existing on main
openai-codex-connect fail in the parallel suite, passes in isolation same test isolation flake

npm run typecheck clean, npm run build clean, all 35 code-plugin tests pass. Staging CI is running the cross-platform matrix separately.

Worth noting for anyone else checking this: typecheck here has to be the npm script. tsc --noEmit -p tsconfig.json checks nothing at all, because that tsconfig is "files": [] with project references only.

The conflict, and one trap in it

Resolved and pushed. The conflict itself was small, both sides had edited the comment above MIN_INCREMENTAL_CHARS, but the resolution has a trap worth stating: main now exports that constant, because tests/code-plugin-remount.test.ts (new on main) imports it. Taking this branch's side of the hunk verbatim drops the export and breaks that test. I kept the export and merged the two comments. mergeable is now MERGEABLE and staging CI merged it clean.

I also fixed 5 no-useless-escape eslint errors that I introduced in the test file I pushed on Aug 16. My mistake, not yours.

This change fixes a real correctness bug, which the description undersells

The plugin this replaces delegated to @streamdown/code, whose result cache keys on language : theme : theme : length : first 100 chars : last 100 chars and then serves the hit without confirming it against the actual code. That key is a lossy digest, so two different fences can collide, and the second one is rendered with the first one's contents. No error, no warning, just the wrong code on screen.

It reproduces on current main in a few lines. Two blocks of equal length sharing an opening and a closing but differing in the middle:

const rendered_A contains "AAAA": true
const rendered_B contains "BBBB": false
const rendered_B contains "AAAA": true   <- B was served A's tokens

This branch is immune, because findFence treats the compact key as a candidate and confirms it with exact.code === code before serving. I checked it both with the fixture above and with one tailored to this branch's own 32-character key shape, so it is the confirmation doing the work and not a wider digest:

fixture main this PR
shared first/last 100 chars collision no collision
shared first/last 32 chars n/a no collision

Two similar config or import blocks in one reply is a realistic shape for this, so I would call it reachable rather than theoretical. I have pushed a regression test for it (39d5174b5). It fails on a tree with the exact.code === code confirmation removed and passes with it, checked individually rather than assumed.

It also closes an unbounded cache

The same upstream cache is module scope and never evicts. Every intermediate prefix seen while a fence streams becomes a permanent entry holding a full tokenization of that prefix. Measured with --expose-gc, forcing a collection before and after, so this is retained heap and not collection lag:

streamed main, retained this PR, retained
one 1.9 KB fence 12.87 MB 0.37 MB
five 1.9 KB fences 65.42 MB 1.09 MB

Linear in fence count on main (12.87 x 5 = 64.4, measured 65.4), bounded here by MAX_FENCES and MAX_CACHED_CHARACTERS. Extrapolating to the 77 fences a 300 KB thread carries, that is on the order of a gigabyte of permanently retained heap on main for a single thread. Small fences are the worse case, not the better one: below MIN_INCREMENTAL_CHARS the old wrapper passed straight through on every frame, so each of ~60 frames per second cached its own full tokenization.

So the honest framing of this PR is broader than the title. The CPU reduction is the headline and I have measured it separately, but the change also removes a wrong-content rendering bug and an unbounded memory leak that exist in the tree today.

Does it break anything

The old pathways still work, checked rather than assumed. The behaviour this replaces is preserved on every axis I could find a way to test:

  • Engine. I want to be explicit here since I raised it before as a hazard. @streamdown/code builds its highlighter with createJavaScriptRegexEngine({ forgiving: true }), and so does this branch. The engine choice is preserved, not changed. That matters because studio/src-tauri/tauri.conf.json sets default-src 'self' with no wasm-unsafe-eval, so Shiki's default Oniguruma WASM engine would be blocked in the packaged desktop app while passing every browser test and every CI job we have. The test I added pins it with that reason in the comment.
  • Language aliases. This branch keeps the hand-written overrides and additionally derives Shiki's own aliases from bundledLanguagesInfo, which is what the upstream plugin did. Nothing that resolved before stops resolving.
  • Remount. tests/code-plugin-remount.test.ts from main passes unchanged, so a remount is still answered inline from cache with no unstyled flash.
  • Existing threads. Nothing here is persisted. The cache is in-memory and rebuilt per session, there is no schema, no on-disk format and no config key, so an existing ~/.unsloth/studio home with threads full of code fences is unaffected in both directions. Downgrading is equally safe.

Still open from my earlier review

The fence-identity question is the one structural item I have not closed. This branch cannot see the opening fence marker, so shedsClosingRun stands in for CommonMark's actual rule that a closer matches the opener's family and length. I am checking now whether a wrong findFence match can produce wrong tokens or only ever a delayed plain tail, since that is the difference between a blocker and a follow-up, and I will post that separately rather than hold this up. My current expectation, from how tokenize always re-derives from the real code, is the latter.

Nice piece of work. The mechanism is sound and it turns out to be doing more good than it claims.

@danielhanchen

Copy link
Copy Markdown
Member

Browser matrix, since this is a rendering path and the node suite only proves it in node.

I bundled the plugin from this branch with esbuild and drove it in Chromium, Firefox and Playwright WebKit, streaming an HTML fence with embedded <script> and <style> (3,958 characters, so well past MIN_INCREMENTAL_CHARS) and comparing the emitted tokens against a whole-document codeToTokens from a separately constructed highlighter. Embedded grammars are the interesting case here because the grammar switches mid document, so a resumed GrammarState has a real stack to restore rather than a single level.

Six cut points, all strictly increasing and all inside the source, with a 300 ms wait before each so every one of them outlasts REFRESH_MS and takes the resume path rather than the approximate plain-tail path.

tree Chromium Firefox WebKit first divergence
this branch 6/6 identical 6/6 identical 6/6 identical none
resume deliberately disabled 5 of 6 differ 5 of 6 differ 5 of 6 differ at 2,500 chars

The second row is the part that makes the first row mean anything. I built a second bundle from a tree with one line changed, grammarState: fence.state to grammarState: undefined, so the tokenizer never resumes. It diverges at the same cut point in all three engines. Without that control the top row would only prove the harness runs, and a parity test that passes on a correct tree and a broken one measures nothing.

The one cut that still matches on the broken tree is the first, at 2,100 characters, which is expected: committedLength is still 0 there, so there is no saved state to resume from and both trees do the same full pass.

Two honesty notes on this. WebKit here is Playwright's WebKit, which I am treating as a proxy for Safari and for the WebKitGTK that Tauri embeds on Linux. It is not real Safari and I am not claiming it is. And the engine question I raised earlier is now settled in the reassuring direction: @streamdown/code already used createJavaScriptRegexEngine({ forgiving: true }), so this branch preserves the engine rather than changing it, and the desktop CSP hazard is unchanged by this PR. Per Shiki's own docs the JS engine wants the RegExp v flag (Chrome 117+, Firefox 119+, Safari 17.4+) and falls back to u otherwise, so all three engines above are on the v path.

@danielhanchen

Copy link
Copy Markdown
Member

Closing out the fence-identity question I left open, since it was the one structural item that could have blocked this.

It does not block. But shedsClosingRun is dormant, and that is worth knowing.

The heuristic cannot produce wrong tokens

CommonMark 0.31.2 section 4.5 requires a closer to use the same marker family as the opener, with a run at least as long, at most three spaces of indentation and nothing after but spaces or tabs. CLOSING_FENCE diverges from that in three ways, all of them false accepts: it accepts runs shorter than three, it accepts a tilde run under a backtick opener and the reverse, and it accepts a run shorter than the opener's. So Codex's item is describing something real about the regex.

It cannot corrupt tokens, though, and the reason is structural rather than lucky. findFence only ever matches an entry whose committed prefix is a genuine prefix of the incoming code: the startsWith branch extends the anchor, and the shedsClosingRun branch additionally requires longer.startsWith(shorter). Since committedLength never exceeds the last newline of the cached code, the resumed grammar state always describes text the incoming code really begins with. A wrong match costs a shared cache entry, not wrong scopes.

The part I did not expect

shedsClosingRun requires shorter === "" || shorter.endsWith("\n"), and the code it is tested against never ends in a newline.

Streamdown's CodeBlock strips every trailing newline before highlighting. In node_modules/streamdown/dist/chunk-BO2N2NFS.js:

lr = e => { let t = e.length; for (; t > 0 && e[t-1] === "\n";) t--; return e.slice(0, t) }

and it is applied as d = useMemo(() => lr(e), [e]), with d passed straight down as the code prop that reaches highlight(). So for any non-empty fence the precondition is false and the predicate never fires, in either of its two call sites.

That has a visible consequence, because the renderer's callback has no cancellation:

let r = o.highlight({code: s, language: e, themes: l}, c => { i(c); });

useEffect re-runs on a new code, but nothing unregisters the callback the previous run left with the plugin, so a refresh queued for an older code value still lands and overwrites the current result. Replaying the exact string sequence Streamdown produces for a large backtick fence whose closing run straddles a chunk boundary:

body length: 6499 | ends with newline: false
final rendered last line: "``"
matches the real source : false

The block permanently renders a stray `` line that is not in the document. Reproduced on both trees:

tree final last line matches source
main (01cf3728c) `` no
this PR `` no

So this is pre-existing on main, not a regression, and not a reason to hold this PR. main has no equivalent guard at all; this branch added one that happens not to reach.

Two things that follow

First, this is further support for the 👎 already on the marker-family item. Not only is the information unavailable at this layer, as you said, the predicate that would use it does not currently run. Matching the marker family would not change any observable behaviour today.

Second, the fix is close but it is a deliberate trade rather than free, which is why I am flagging it rather than pushing it. The shape to match is \n+ then the run, since the strip removes as many newlines as the body ends with, and the endsWith("\n") precondition has to go. Doing that makes the predicate reachable for ordinary bodies for the first time, which is exactly what makes the marker-family question real instead of theoretical. I would rather that landed as its own change, with its own tests, than get folded into this one. Happy to open it if you would prefer.

The cheap thing worth doing here, if you want it, is a one-line comment above CLOSING_FENCE recording that CodeBlock strips trailing newlines, so the next reader does not have to rediscover why the guard is quiet.

For completeness, on the question this PR does own: a 35 KB Python fence streamed over 300 frames measured 845/989/848 ms of process CPU on main against 269/221/259 ms here, so roughly 3.4x to 4.5x, consistent with the numbers I posted earlier at a different fixture size.

Where the opening marker actually lives

For the record, since it will come up again: Streamdown does know it. CodeComponent receives the hast node and already reads node.position.start.line/column for its memo comparator, and that position spans the opener. It just is not forwarded, and HighlightOptions is {code, language, themes} in both streamdown and @streamdown/code. Adding an optional blockKey would be source-compatible for existing plugins but is a change in vercel/streamdown plus a version bump, so it is not something this PR could carry. startLine is already plumbed to HighlightedCodeBlockBody and would be a usable approximation without any upstream change, if a stable identity is ever needed.

@danielhanchen

Copy link
Copy Markdown
Member

Correcting one number in my previous comment, and adding the thing that makes it useful.

I quoted "845/989/848 ms on main against 269/221/259 ms here, so roughly 3.4x to 4.5x" for a 35 KB Python fence. That figure came from a measurement I had not re-run myself before posting it, which is not the standard I have held the rest of this review to. So I measured it, three interleaved repetitions per cell, process.cpuUsage() over the whole streamed run so the trailing refresh is counted, grammar warmed first.

At that fixture size I reproduce the main side closely and get a smaller ratio than I quoted:

35,709-char Python fence, 300 frames main this PR ratio
rep 1 917 ms 327 ms 2.8x
rep 2 868 ms 321 ms 2.7x
rep 3 945 ms 323 ms 2.9x

So 2.7x to 2.9x, not 3.4x to 4.5x. The main column matches what I quoted (868 to 945 against 845 to 989); the difference is entirely on this branch's side, where I measure 321 to 327 rather than 221 to 269. Same direction, same conclusion, but I would rather the record carry the number I can stand behind.

The more useful finding is that the ratio is strongly a function of fence size, which a single multiplier hides. Same harness, same pacing, only the fence doubled:

fence main this PR ratio
35,709 chars 868 to 945 ms 321 to 327 ms 2.7x to 2.9x
73,869 chars 4,080 to 6,984 ms 469 to 513 ms 8.7x to 13.6x

Doubling the fence roughly quintuples main's cost while this branch's grows by about half. That is exactly the shape the design predicts, since main re-tokenizes the whole fence on every refresh and this branch tokenizes only what is new, so the gap should widen with length rather than hold constant.

This is the same caution I raised earlier about the character-reduction figure being cadence-dependent, and it applies to the CPU figures too: quoting one multiplier invites someone to measure a different fence, get 2.7x, and think the work regressed. A range with the fixture attached ages better. The numbers in the PR description are fine on that front since they name their fixtures.

Nothing here changes the verdict. It makes the case stronger at the sizes that motivated the PR.

@danielhanchen

Copy link
Copy Markdown
Member

Rendering evidence, since this touches a rendering path and the premise is that the output does not change.

Two separate Studio installs, one per tree, each with its own home and port, driven by the same scripted SSE bytes with no model loaded so both sides receive identical reply text. BEFORE is 01cf3728c, which is the exact merge base of the AFTER tree, so the pair isolates this PR rather than a fortnight of unrelated main. AFTER is e89e4e693; that is one commit behind the current head and the difference is 61 lines in a test file, no source change.

Two fences per run, each streamed 100 characters every 45 ms with three deliberate pauses mid-fence so the same prefix could be photographed on both builds fully refreshed:

  • a 10,305-character HTML fence carrying real <style> and <script> blocks, which is the embedded-grammar case,
  • an 8,789-character TypeScript fence including a regex literal, where divide-versus-regex is precisely the token whose scope depends on state carried into the line.

Every checkpoint sits above MIN_INCREMENTAL_CHARS, so the incremental path is the one being photographed rather than the short-fence path both builds share.

The painted fence markup is byte-identical at every sample. Hashes over the live DOM of the fence, so this is token boundaries and colours rather than a visual impression:

sample BEFORE AFTER
html mid-fence, 3,011 / 6,091 / 8,632 chars a8070855d9f9 / ffd5610b1bb8 / eab8a7c74195 identical
html settled 053ded6e93aa, 2,755 coloured tokens identical
ts settled 108ba5594771, 2,645 coloured tokens identical

Inside <style>: grid-template-columns #383A42, linear-gradient #0184BC. Inside <script>: async and await #A626A4. Same on both sides. The plain-tail cadence is unchanged too, longest plain tail 167 and 69 characters on both.

The control, which is the part that makes a parity pair mean anything. A pair of identical screenshots is also what you get from photographing one build twice, so the scene's expect was written before the run to require that something must differ:

BEFORE AFTER
main-thread long tasks (html / ts) 14 / 14 5 / 2
long-task total (html) 1,101 ms 356 ms, -68%
long-task total (ts) 1,179 ms 113 ms, -90%
bundle chat-CrAeM29I.js chat-DSiSoSnw.js, +2,669 bytes, grammarState 1 -> 4

So the two halves are demonstrably two different builds producing the same pixels at a third of the main-thread cost.

One divergence, named rather than smoothed over: in the mid-stream pair the images differ in a single 16x16 pixel box, the sidebar "generating" spinner caught at a different point of its CSS rotation. Wall clock, not content. Four of the six screenshot pairs are byte-identical PNGs.

On the trap that a settled screenshot could mask a broken resume by re-tokenizing everything at the end: tokenize() only ever tokenizes code.slice(fence.committedLength) and concatenates the cached fence.lines, so a streamed fence never gets a final full pass, and the scene never reloads the thread. The settled DOM is the incremental path's own output. The mid-fence samples were taken anyway rather than relying on that argument.

@danielhanchen

Copy link
Copy Markdown
Member

@codex review

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 18, 2026
@danielhanchen

Copy link
Copy Markdown
Member

Follow-up review at head fbf4a93db. Two things worth putting on the record: a content-loss question that was hanging over this PR, and the state of the gaps I listed earlier.

1. The streaming content-loss report does not belong to this PR

A streaming failure has been reported where text streams in, generation finishes, and then a large part of the rendered reply disappears. The build it was seen on had this PR and another in-flight streaming change applied together, so it could not tell the two apart. I ran the three arms separately.

Method: same frozen main (c675be1ad) in all arms, so the only variable is the build. No model is loaded; a fetch shim serves an identical scripted SSE reply (13,065 chars, sha256 57cfd8be2ae7) from an out-of-process server, so the bytes are the same in every arm and every trial. Sampled every 100 ms during the stream, then again after it settles. Every sample records the last assistant node, the sum over all assistant nodes, and the max over them, so a node-selection artifact cannot be mistaken for real loss. Chromium, 11 trials per arm.

arm build peak rendered final rendered lost after completion
A frozen main 12,273 12,273 0 in 11/11
B this PR alone 12,273 12,273 0 in 11/11
C this PR + the other streaming change 4,026 to 4,631 3,863 to 4,032 damaged in 11/11

This PR alone loses nothing. It renders 12,273 of 12,273 characters and matches frozen main exactly, in 11 out of 11 trials. Arm C additionally never renders more than about a third of the reply in the first place, which is the larger of its two problems and one that a "lost after completion" number understates. I am not drawing a conclusion here about the other change on its own, since arm C is not a clean superset of arm B; the point I can support is the one about this PR.

Two honest caveats. Chromium here is a proxy, not the packaged desktop webview. And the useClientLookup: Clamped stale index warnings that accompanied the original report did not reproduce in any arm, including C. That identifier is not in this repository at all; it comes from @assistant-ui/core. I could not reproduce those warnings and I am not claiming to have explained them.

2. The gaps I raised earlier are closed, and the new tests are proven non-vacuous

A test that passes is not evidence until it can fail. I cut six worktrees from head, broke one thing in each, and re-ran the whole suite. Baseline at head is 37 pass, 0 fail, 0 skipped.

what I broke result
re-tokenize the whole fence every update (tokens stay identical) 36 pass, 1 fail: a streaming fence is tokenized once, not once per update
swap to the WASM Oniguruma engine fails the engine test, and hangs the rest of the suite
drop the character-budget clause from evict() 36 pass, 1 fail: the character budget evicts even while the fence count is under its limit
remove the never-evict carve-out 36 pass, 1 fail: a fence larger than the whole budget is kept rather than evicting itself
never release a failed grammar key 36 pass, 1 fail: a transiently failed grammar load is retried under the same cache key
never resume grammar state 26 pass, 11 fail, including all five embedded-grammar tests

Four of those fail exactly one test while 36 others pass, which is the useful property: each guard is specific and nothing else in the suite can detect that regression. In particular the first one confirms what the performance test claims about itself, that re-tokenizing whole would pass the entire rest of the suite, so that test is the only thing standing between this PR and a silent loss of the property it exists for.

On the fence-identity item: I agree with your reasoning and I am not asking for a change. The plugin genuinely cannot see the opening marker, and a body plus a bare marker line is ambiguous by construction. That is a limitation of the interface, not a defect in this PR.

The engine assertion is worth one note. It pins source text rather than behaviour, which is unusual, but tauri.conf.json sets default-src 'self' with no script-src and no wasm-unsafe-eval, and per the CSP spec WebAssembly compilation is blocked in exactly that configuration. Shiki's Oniguruma engine is WebAssembly. Nothing in CI drives the packaged webview, so a future switch would ship broken and no job would catch it. I pushed one small commit there so a failure prints the message instead of 10 KB of source.

3. Does merging break anything

The PR touches one source file plus its tests. No storage, schema, migration or thread-persistence path is touched, so an existing ~/.unsloth/studio home with threads containing code fences is unaffected: this is render-time only and nothing on disk changes shape. The old pathway is covered rather than assumed, since every incremental and embedded-grammar test asserts the streamed result is byte-identical to whole-document tokenization, which is the pre-PR behaviour, and the pre-existing code-plugin-remount.test.ts still passes.

Gates at head: 37/37 tests pass with 0 skipped, npm run typecheck clean, eslint clean on the changed files, npm run build succeeds. On a staging replica the only red job is tests/queued-model-capabilities.test.ts, which this PR does not touch and which fails the same way on plain upstream main with ERR_MODULE_NOT_FOUND for mmproj-fallback. That one is pre-existing and not yours.

One small thing: the body still says "15 focused highlighting tests" and head now has 37.

4. Scope note on the reasoning pane

Separately measured, so it does not get misattributed to this PR later: incremental tokenization does not fix the slow streaming reasoning pane. Mounted span counts come out essentially identical with and without this PR (peaks within a handful of spans of each other). The pane renders through the same path, so this PR is structurally on it, but this PR changes the cost of producing tokens, not the number of spans mounted, and the pane's problem is the mount count. That is a separate piece of work and it takes nothing away from the CPU result here.

Nice work on this one. The incremental design holds up under everything I could throw at it.

@danielhanchen

Copy link
Copy Markdown
Member

@codex review

@danielhanchen

Copy link
Copy Markdown
Member

A measurement that bears directly on what this PR has left to win

I have been measuring the parse-side cost of a streaming reasoning reply for a separate change
(#9228, which bounds the Shiki token cache). Two results here are relevant to incremental
tokenisation, and one of them may reduce the headroom this PR is aiming at.

Shiki costs ~0.00 ms per streamed chunk, and the reason matters

Per chunk, on a reasoning-pane-shaped fixture at 90,000 characters (3,750 chunks of 24 characters,
prose plus repeating 1,800-character python fences):

build        task  script  string  react  parse  shiki  layout  style
dev server   5.23    4.02    0.18   2.53   0.07   0.00    0.11   0.02
production   2.63    1.51    0.18   0.78   0.06   0.00    0.10   0.02

Shiki reads 0.00 on both builds. That is because every fence is tokenised once and then served
from cache, not because tokenisation is cheap.
The distinction decides what incremental
tokenisation can recover: if the repeated work has already been eliminated by caching, making the
remaining single tokenisation incremental saves a fraction of something that is already near zero
per chunk.

I want to be careful about the scope of that. It says nothing about the FIRST tokenisation of a
long fence, or about a fence that is still growing while it is on screen, which is exactly the case
this PR targets. It does say that summed over a stream, the per-chunk cost of highlighting is not
where the time goes on this fixture.

The number was reached by ablation, not by a wrapper, which is why I trust it

A wrapper around highlight sees far less than the whole cost. I injected a known delay per call
and watched where it landed:

parse   injected 7500 ms -> parse rose 7348 ms (98%), React rose 7896 ms (105%)
Shiki   injected 3344 ms -> Shiki rose 3335 ms (99.7%), React rose  219 ms (7%)

So Shiki runs about 7% inside React's render phase and 93% outside it, in the trailing
setTimeout that code-plugin.ts schedules. Any instrument that measures Shiki as render-phase
work sees roughly a fourteenth of it, and a split built that way would under-report tokenisation
badly.

The 0.00 above is therefore not from that wrapper. It is from ABLATION: replacing the real
highlighter with a stub that splits lines on whitespace instead of running the grammar, keeping the
token count within tolerance so the DOM does not change materially, and comparing total main-thread
task time from CDP. Full 19,602 ms against stub 20,015 ms at 90 K, with the stub marginally
slower. Ablation catches the asynchronous grammar work that a wrapper cannot see, and it still
finds nothing to remove.

What I would suggest measuring before landing

Since the caching already removes the repeated work, the case for this PR probably rests on the
single-fence latency it improves rather than on throughput across a stream. A measurement of
time-to-first-highlighted-token for one long growing fence would make that case directly, and it
would be immune to the caching effect above.

Method and caveats: Chromium only; one run per cell; production numbers taken against a bundle
built with react-dom/profiling so <Profiler> still reports; and on the production runs 36-38%
of main-thread task time is neither script, layout nor style and is unattributed, so treat the
production column as indicative rather than exact. The harnesses are on #9228 if you want to rerun
any of it.

The markdown fixture already here holds an HTML comment, which pushes one
shallow level and leaves the nested fence's body doing no work. This one
opens a python fence and a bash fence inside a markdown body and closes
both, so the grammar switches in and back out twice.

That matters because a `#` line inside a fence is body text and the same
line outside one is a heading. A resume that stays a level too deep, or
returns a level too shallow, colours them the other way round and drags
the prose after each fence with it.

Checked against a tree with `grammarState: fence.state` replaced by
`undefined` rather than assumed to discriminate: it fails there at
character 100, the `# Collect the rows` line inside the python fence,
which the broken tree renders as an H1.

The shared helper now waits out REFRESH_MS once a fixture is past
MIN_INCREMENTAL_CHARS. Without it this test fails on correct code at the
first comparison past the threshold, because the throttled approximation
renders the uncommitted tail plain and never reads the grammar state. All
119 comparisons here take the resume path, 15 of them past the threshold.
The five existing fixtures are all below it and pay nothing.

Tests only, no source changes. 38 passed, 0 failed.
@danielhanchen

Copy link
Copy Markdown
Member

Came back to this with a specific question, and the answer turned out to be worth writing up properly because it nearly went the other way.

The question

A separate piece of profiling on the streaming pane measured Shiki at approximately 0.00 ms per chunk, and attributed it to every fence being tokenized once and cached. Taken at face value that would undercut this PR: if fences are already tokenized once in the shipped path, making tokenization incremental would win much less than the numbers here suggest.

I did not want to raise that with you as an inference, so I measured it. It does not undercut this PR. Two of the three parts of that claim are right, and the part that would have mattered is not.

What the fixture behind that measurement actually is

The ablation ran on plain main, without this branch, on a vite dev server, with a fixture of 21 fences built from a repeating 1,250-character prose paragraph plus an 1,800-character fence. Only one arm was ever run, so it is not a comparison of the two trees. It also cannot be one: on main, code-plugin.ts:133-135 short-circuits every fence under MIN_INCREMENTAL_CHARS straight into @streamdown/code, so on that fixture almost none of the region this PR rewrites executes.

That is a fact about the fixture, not a criticism of the instrument. The ablation is sound for what it does measure, and I say below what that is.

Measured, on both trees, same probe

Both trees import createHighlighter from the same shiki, so I wrapped that one entry point and counted characters and calls into codeToTokens, in an isolated install so nothing else in the tree was perturbed. Frame boundaries come from marked's Lexer, the splitter Streamdown 2.5 uses, cross-validated against remark-parse on all 501 frames of a 12 K run with 0 mismatches. Merge base fd6ab9aca against 99a34e7f7, which is your last source commit, so the pair isolates the source change. Grammars warmed first, three interleaved repetitions per cell, arm order alternating.

1. The win does not start at the threshold, and it is not flat

Single TypeScript fence, streamed 24 characters every 73 ms, which is a real token rate. Medians of three.

fence body main chars into Shiki this PR ratio main CPU this PR CPU
477 5,237 1,023 5.1x 254 ms 176 ms
1,081 25,290 2,368 10.7x 499 ms 212 ms
1,842 72,835 4,060 17.9x 1,029 ms 262 ms
2,612 103,336 5,181 19.9x 1,347 ms 288 ms
7,540 412,367 11,958 34.5x 4,377 ms 418 ms
20,104 2,537,353 28,918 87.7x 24,553 ms 764 ms

I expected the two trees to be indistinguishable below 2,000 and predicted that in writing before running it. That was wrong, and wrong in your favour. There is no discontinuity at the threshold at all; the curve is smooth from 477 characters upward.

The reason is worth stating in your description, because I had to read the source twice to see it: MIN_INCREMENTAL_CHARS gates the plain-tail throttle, not the incrementality. Below it, main re-tokenizes the entire fence on every single frame, visible as its call count equalling the frame count exactly (21/21, 46/46, 78/78). Above it, main re-tokenizes at most every 250 ms. So below the threshold is main's most wasteful regime, not its least, and your change covers it while the description implies it does not.

2. On the profiled fixture itself, this PR is 3.9x cheaper

Same 90,000-character fixture as the ablation, real 73 ms gap, three interleaved repetitions per arm.

arm chars into Shiki calls CPU
main 1,617,982 1,649 7,153 ms
this PR 421,903 2,364 1,850 ms
ratio 3.83x 0.70x 3.87x

My main-arm count of 1,649 calls over 21 fences lines up with the other instrument's 1,672 over 22 on the same fixture shape, the difference being one truncated fence at the slice boundary. Two independently built probes within 1.4% is the best evidence either of us has that both are on the real path.

3. Why 3.83x and not 41x, which is the honest limit of this design

Attributing every call to a fence body:

fence class count main this PR ratio
python and typescript, 59 to 67 lines 16 1,255,183 30,714 40.9x
json, one very long line 5 358,851 358,851 1.00x

The single-line fences cost both trees 358,851 characters, identical to the character, and they are 85% of everything this PR still spends. That follows directly from committing completed lines: a fence containing no newline commits nothing, so its live line is re-tokenized every frame on both trees. Your description already says one long line still relies on the existing throttle, and this is that, quantified. It means the aggregate ratio on any real stream is a property of the content mix, so an all-multi-line stream reads about 41x and an all-minified-JSON one reads about 1x.

4. Closed fences are not re-tokenized on either tree

This is the part of the caching claim that is true, and I want to be precise about who it belongs to. Driving highlight() on every fence on every frame (28,656 calls) versus only on changed fences (1,652) produces byte-identical Shiki totals. Every fence pays exactly one tokenization of its final body when it closes and zero thereafter, across thousands of subsequent frames.

That caching is @streamdown/code's module-global cache and it exists on main today. This PR neither introduces it nor removes it. So "every fence is tokenized once and cached" is true of each fence's final body and false of the roughly 77 intermediate bodies that precede it, and those intermediate bodies are where main's 1.6 million characters go. That is the whole gap this PR closes.

So what does the 0.00 ms reading mean

It means Shiki is not the streaming pane's bottleneck, and that is correct. On that fixture main spends 1.54 ms of Shiki per 73 ms chunk, about a 2% duty cycle, against a pane holding 18,536 elements. An ablation whose resolution floor is roughly plus or minus 400 ms of 19,600 cannot see a 2% component, and its measured delta was 413 ms and flipped sign at 45,000 characters. Cutting a 2% component by 3.9x is worth having but it is about 1.5% of that pane's time, and I would not want either of us quoting the 3.9x as if it fixed the pane. It does not, and I said as much separately.

None of that reaches your headline claims, which are a single large fence in the answer, measured on production builds. It also does not reach mine: my earlier browser A/B here was two Studio installs each serving frontend/dist with the hashed bundles named, not a dev server.

One wrinkle that is yours to judge

When a chunk boundary splits the closing backtick run, the body transiently grows to include a partial delimiter and then shrinks, and this branch pays one extra near-full pass. Confirmed causally rather than inferred: at 1,800-character fences exactly one fence sheds a partial closing run and exactly that fence pays 1,864 characters, while at 2,500 none sheds and none pays over 142. Your worst case at close equals main's every case at close, so it is a note rather than a defect, and I would leave it alone.

The metric I would add

For a single long growing fence the number that shows what this design buys is time to first highlighted token, not steady-state cost per chunk. On main at 20,104 characters the tokenizer falls up to 1.2 seconds behind the stream at a fast cadence, while this branch stays within about 90 ms. Steady-state CPU understates that, and latency-to-paint is what a user actually experiences on a long code cell.

Housekeeping

I pushed one more test, 7c63fb497. Tests only, no source changes, 38 passed and 0 failed. It covers a markdown body whose nested python and bash fences both close again, so the grammar switches in and back out twice. The markdown fixture already there holds an HTML comment, which pushes one shallow level; this one makes the pop back out load-bearing, because a # line inside a fence is body text and the same line outside one is a heading. Checked against a tree with grammarState: fence.state replaced by undefined rather than assumed to discriminate: it fails there at character 100, on the # Collect the rows line, which the broken tree renders as an H1. The shared helper now waits out REFRESH_MS past the threshold, because without it the test fails on correct code at the first comparison past 2,000, where the throttled approximation renders the tail plain and never reads the grammar state. The five existing fixtures are below the threshold and pay nothing.

The red checks are not yours. All six were checked. Frontend build + bundle sanity, Frontend unit tests (Windows), Repo tests (CPU) and (Python 3.13) fail identically on plain main and were introduced by #9173, which added mmproj-fallback.ts; #9220, #9192 and #9189 are all open against that breakage. The other two, Chat UI Tests (chat) and Loaded-models indicator, were cancelled inside apt-get update before any repo code ran, with the same signature on unrelated PRs in the same window. Neither code-plugin nor shiki appears in any failure log, and your own tests passed on the Windows job.

Where I land

I went looking for a reason this PR wins less than it claims and did not find one. The mechanism is sound, the parity evidence holds across three browser engines and on production Studio builds, the win is real at every fence size I could measure, and it extends below the threshold in a way the description undersells. The single-line case is a genuine limit and you had already named it. Two small things for the body when you next touch it: it still says "15 focused highlighting tests" and head now has 38, and it is worth saying that the improvement covers short fences too.

@oobabooga

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Already looking forward to the next diff.

Reviewed commit: 8d07fa46ff

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@oobabooga

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

@oobabooga
oobabooga merged commit 531bc48 into unslothai:main Aug 19, 2026
27 of 31 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants