From 6c95740d9eaed10d214d4a35c2960ee6a08c066a Mon Sep 17 00:00:00 2001 From: Harry Yep Date: Wed, 24 Jun 2026 11:57:48 +0800 Subject: [PATCH] feat: replace custom streaming wrappers with built-in defer and smooth props MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../components/assistant-ui/markdown-text.tsx | 85 +++++-------------- 1 file changed, 20 insertions(+), 65 deletions(-) diff --git a/apps/desktop/src/components/assistant-ui/markdown-text.tsx b/apps/desktop/src/components/assistant-ui/markdown-text.tsx index 694a94e78..841e91fe4 100644 --- a/apps/desktop/src/components/assistant-ui/markdown-text.tsx +++ b/apps/desktop/src/components/assistant-ui/markdown-text.tsx @@ -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 `` 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 ( - - {children} - - ) -} - 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 ( - - - + ) } const MarkdownTextImpl = () => { - return ( - - - - ) + return } export const MarkdownText = memo(MarkdownTextImpl)