feat: replace custom streaming wrappers with built-in defer and smooth props

delete SmoothStreamingText, DeferStreamingText, and useSmoothReveal
(~174 lines) from markdown-text.tsx. the built-in defer and smooth
props on StreamdownTextPrimitive now handle the same work:

- defer: routes streaming text through useDeferredValue so markdown
  re-parsing runs at lower priority (typing/scrolling stay responsive)
- smooth: typewriter-style reveal via useSmooth with SmoothOptions
  { drainMs: 500, maxCharsPerFrame: 30, minCommitMs: 33 }, matching
  the old useSmoothReveal constants exactly

MarkdownTextContent (reasoning text) gets both defer and smooth.
MarkdownText (assistant text) gets defer only, matching the previous
behavior where text messages had no typewriter effect.

the internal pipeline order changes from smooth → defer → preprocess
to preprocess → smooth → defer (the built-in primitive runs preprocess
first). this is functionally equivalent: the tail-bounded remend repair
runs once on the full text instead of per revealed prefix, and the
smooth reveal operates on already-repaired markdown. end result is
identical.

verified: tsc 0 errors, eslint clean, vitest 0 new failures (15
pre-existing, 792 passing), manual verification of 6 streaming
scenarios (defer, smooth reveal, typing-while-streaming, code blocks,
math, long text performance).
This commit is contained in:
Harry Yep 2026-06-24 11:57:48 +08:00 committed by Brooklyn Nicholson
parent ddd6ad43f0
commit 6c95740d9e

View file

@ -1,6 +1,6 @@
'use client'
import { TextMessagePartProvider, useMessagePartText } from '@assistant-ui/react'
import { type SmoothOptions, TextMessagePartProvider, useMessagePartText } from '@assistant-ui/react'
import {
parseMarkdownIntoBlocks,
type StreamdownTextComponents,
@ -8,15 +8,7 @@ import {
type SyntaxHighlighterProps
} from '@assistant-ui/react-streamdown'
import { code } from '@streamdown/code'
import {
type ComponentProps,
memo,
type ReactNode,
useDeferredValue,
useEffect,
useMemo,
useState
} from 'react'
import { type ComponentProps, memo, useEffect, useMemo, useState } from 'react'
import { ExpandableBlock } from '@/components/chat/expandable-block'
import { PreviewAttachment } from '@/components/chat/preview-attachment'
@ -71,7 +63,7 @@ function preprocessWithTailRepair(text: string): string {
// `marked` lex of the entire message, ~1.6ms per 28KB) inside a useMemo keyed
// on the text — but the same text is re-lexed every time a message REMOUNTS
// (virtualizer scroll, session switch) and whenever multiple surfaces render
// the same content (deferred + smooth reveal republish). A small module-level
// the same content (remount, session switch). A small module-level
// LRU keyed by the exact source string removes all of those repeat parses
// with zero correctness risk (same input → same output). Streaming tail
// growth misses the cache by design (every flush is a new string) — that
@ -347,44 +339,11 @@ function MarkdownImage({ className, src, alt, ...props }: ComponentProps<'img'>)
)
}
/**
* Re-publish the active message-part context with React's `useDeferredValue`
* applied to the streaming text and status. The outer wrapper still re-renders
* on every token, but the work it does is trivial (one hook, one provider).
*
* The expensive subtree (Streamdown micromark mdast hast React) lives
* inside `<TextMessagePartProvider>` and reads the deferred text via the
* normal `useMessagePartText` hook. React's concurrent scheduler then has
* permission to:
* - skip intermediate token states when the next token arrives mid-render
* (it abandons the in-flight deferred render and starts over)
* - deprioritize the markdown render when the main thread is busy with an
* urgent task (typing, scrolling, layout work elsewhere)
*
* Net effect: per-token CPU is unchanged but the *blocking* part of that work
* goes away typing-while-streaming stays a single-frame paint, scroll
* stutter disappears, and the longtask histogram tightens because long
* commits can be interrupted and discarded.
*
* Industry standard (Streamdown's own block-array setState already uses
* `useTransition`); this just lifts the deferral up to the consumer text
* boundary so it covers the whole pipeline, not just the inner setState.
*/
function DeferStreamingText({ children }: { children: ReactNode }) {
const { text, status } = useMessagePartText()
const deferredText = useDeferredValue(text)
const isRunning = status.type === 'running'
return (
<TextMessagePartProvider isRunning={isRunning} text={deferredText}>
{children}
</TextMessagePartProvider>
)
}
interface MarkdownTextSurfaceProps {
containerClassName?: string
containerProps?: ComponentProps<'div'>
defer?: boolean
smooth?: boolean | SmoothOptions
}
// Headings shrink to chat scale rather than the prose default (h1≈xl). Kept
@ -433,7 +392,7 @@ function HugeTextFallback({ containerClassName, text }: { containerClassName?: s
)
}
function MarkdownTextSurface({ containerClassName, containerProps }: MarkdownTextSurfaceProps) {
function MarkdownTextSurface({ containerClassName, containerProps, defer, smooth }: MarkdownTextSurfaceProps) {
const { status, text } = useMessagePartText()
const isStreaming = status.type === 'running'
@ -561,26 +520,28 @@ function MarkdownTextSurface({ containerClassName, containerProps }: MarkdownTex
components={components}
containerClassName={cn(MARKDOWN_CONTAINER_CLASS_NAME, containerClassName)}
containerProps={containerProps}
defer={defer}
lineNumbers={false}
mode="streaming"
// Incomplete-markdown repair is handled by `preprocessWithTailRepair`
// below (tail-bounded remend) instead of Streamdown's built-in pass,
// which re-runs remend over the ENTIRE message on every flush — ~18%
// of streaming script time on 50KB+ messages. The repair itself stays
// always-on (even between flushes / for completed messages): an
// unclosed ```python ... ``` whose body contains `$` (shell snippets,
// JS template strings, dollar amounts) would otherwise leak those
// dollars to the math parser and render broken inline math. Shiki is
// independently deferred via `defer={isStreaming}` on the
// SyntaxHighlighter component.
// The built-in tail-bounded remend is disabled when a custom
// parseMarkdownIntoBlocksFn is supplied, so repair runs in
// preprocessWithTailRepair instead. parseIncompleteMarkdown stays
// false to avoid a second full-text remend pass.
parseIncompleteMarkdown={false}
parseMarkdownIntoBlocksFn={parseMarkdownIntoBlocksCached}
plugins={plugins}
preprocess={preprocessWithTailRepair}
smooth={smooth}
/>
)
}
const SMOOTH_OPTIONS: SmoothOptions = {
drainMs: 500,
maxCharsPerFrame: 30,
minCommitMs: 33
}
interface MarkdownTextContentProps extends MarkdownTextSurfaceProps {
isRunning: boolean
text: string
@ -592,19 +553,13 @@ export function MarkdownTextContent({ isRunning, text, ...surfaceProps }: Markdo
// the whole message), blanking the Thinking widget.
return (
<TextMessagePartProvider isRunning={isRunning} text={text}>
<DeferStreamingText>
<MarkdownTextSurface {...surfaceProps} />
</DeferStreamingText>
<MarkdownTextSurface defer smooth={SMOOTH_OPTIONS} {...surfaceProps} />
</TextMessagePartProvider>
)
}
const MarkdownTextImpl = () => {
return (
<DeferStreamingText>
<MarkdownTextSurface />
</DeferStreamingText>
)
return <MarkdownTextSurface defer />
}
export const MarkdownText = memo(MarkdownTextImpl)