diff --git a/packages/web-shell/client/components/messages/AssistantMessage.tsx b/packages/web-shell/client/components/messages/AssistantMessage.tsx index 8f681ce620c..2f9c0f6c5f7 100644 --- a/packages/web-shell/client/components/messages/AssistantMessage.tsx +++ b/packages/web-shell/client/components/messages/AssistantMessage.tsx @@ -32,41 +32,6 @@ interface AssistantMessageProps { customFooterInfo?: WebShellAssistantTurnFooterRenderInfo; } -const STREAMING_MARKDOWN_UPDATE_MS = 80; - -function useStreamingMarkdownContent(content: string, isStreaming?: boolean) { - const [streamingContent, setStreamingContent] = useState(content); - const latestContentRef = useRef(content); - const timerRef = useRef | undefined>(undefined); - latestContentRef.current = content; - - useEffect(() => { - if (!isStreaming) { - if (timerRef.current !== undefined) { - clearTimeout(timerRef.current); - timerRef.current = undefined; - } - if (streamingContent !== content) setStreamingContent(content); - return; - } - if (timerRef.current !== undefined || streamingContent === content) return; - timerRef.current = setTimeout(() => { - timerRef.current = undefined; - setStreamingContent(latestContentRef.current); - }, STREAMING_MARKDOWN_UPDATE_MS); - }, [content, isStreaming, streamingContent]); - - useEffect( - () => () => { - if (timerRef.current !== undefined) clearTimeout(timerRef.current); - }, - [], - ); - - if (!isStreaming) return content; - return content.startsWith(streamingContent) ? streamingContent : content; -} - export const AssistantMessage = memo(function AssistantMessage({ content, isStreaming, @@ -79,7 +44,6 @@ export const AssistantMessage = memo(function AssistantMessage({ }: AssistantMessageProps) { const { t } = useI18n(); const { renderAssistantTurnFooter } = useWebShellCustomization(); - const markdownContent = useStreamingMarkdownContent(content, isStreaming); const [copied, setCopied] = useState(false); const showFooter = !!content && !isStreaming && showFooterActions; const customFooter = useMemo( @@ -111,7 +75,7 @@ export const AssistantMessage = memo(function AssistantMessage({ >
diff --git a/packages/web-shell/client/components/messages/Markdown.test.ts b/packages/web-shell/client/components/messages/Markdown.test.ts index 361e51903f8..d575fdab455 100644 --- a/packages/web-shell/client/components/messages/Markdown.test.ts +++ b/packages/web-shell/client/components/messages/Markdown.test.ts @@ -3,7 +3,7 @@ */ import { act, createElement, type ReactNode } from 'react'; import { createRoot } from 'react-dom/client'; -import { afterEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { WebShellCustomizationProvider, type WebShellCodeBlockRenderInfo, @@ -1372,6 +1372,8 @@ describe('Markdown code highlighting while streaming', () => { isStreaming: true, }), ); + // Wait for the 80ms streaming throttle to flush the new content + await new Promise((resolve) => setTimeout(resolve, 100)); }); expect(container.querySelector('.shiki')).toBeNull(); expect(container.textContent).toContain('const b = 2;'); @@ -1521,3 +1523,59 @@ describe('Markdown code highlighting while streaming', () => { container.remove(); }); }); + +describe('Markdown streaming throttle', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('flushes the latest content when multiple tokens arrive in one throttle window', () => { + const container = document.createElement('div'); + document.body.appendChild(container); + const root = createRoot(container); + + act(() => { + root.render( + createElement(Markdown, { + content: 'Token 1', + isStreaming: true, + }), + ); + }); + + act(() => { + root.render( + createElement(Markdown, { + content: 'Token 1 Token 2', + isStreaming: true, + }), + ); + }); + act(() => { + root.render( + createElement(Markdown, { + content: 'Token 1 Token 2 Token 3', + isStreaming: true, + }), + ); + }); + + expect(container.textContent).toContain('Token 1'); + expect(container.textContent).not.toContain('Token 1 Token 2 Token 3'); + + act(() => { + vi.advanceTimersByTime(80); + }); + + expect(container.textContent).toContain('Token 1 Token 2 Token 3'); + + act(() => { + root.unmount(); + }); + container.remove(); + }); +}); diff --git a/packages/web-shell/client/components/messages/Markdown.tsx b/packages/web-shell/client/components/messages/Markdown.tsx index 846652a384b..554e697186f 100644 --- a/packages/web-shell/client/components/messages/Markdown.tsx +++ b/packages/web-shell/client/components/messages/Markdown.tsx @@ -12,7 +12,7 @@ import { import { useTheme } from '../../themeContext'; import { useTranscriptRenderMode } from '../../transcriptRenderMode'; import ReactMarkdown, { defaultUrlTransform } from 'react-markdown'; -import type { Components } from 'react-markdown'; +import type { Components, Options } from 'react-markdown'; import { isMarkdownFenceClosed } from '@datafe-open/markdown-chart'; import remarkGfm from 'remark-gfm'; import remarkMath from 'remark-math'; @@ -748,6 +748,74 @@ function MarkdownImage({ src, alt }: { src?: string; alt?: string }) { return {alt; } +/** + * Throttles a rapidly changing value (like a streaming string) to prevent + * O(n²) re-parsing of the entire Markdown AST on every token. + */ +function useThrottledValue( + value: string, + isStreaming: boolean | undefined, + intervalMs: number = 80, +): string { + const [throttled, setThrottled] = useState(value); + const throttledRef = useRef(throttled); + throttledRef.current = throttled; + const lastRunRef = useRef(0); + const timeoutRef = useRef | null>(null); + const valueRef = useRef(value); + valueRef.current = value; + + useEffect(() => { + if (!isStreaming) { + if (timeoutRef.current) { + clearTimeout(timeoutRef.current); + timeoutRef.current = null; + } + // Flush immediately when streaming stops + if (throttledRef.current !== value) { + setThrottled(value); + } + return; + } + + const now = Date.now(); + const elapsed = now - lastRunRef.current; + + if (elapsed >= intervalMs) { + lastRunRef.current = now; + setThrottled(valueRef.current); + } else if (!timeoutRef.current) { + timeoutRef.current = setTimeout(() => { + lastRunRef.current = Date.now(); + timeoutRef.current = null; + setThrottled(valueRef.current); + }, intervalMs - elapsed); + } + }, [value, isStreaming, intervalMs]); + + useEffect(() => { + return () => { + if (timeoutRef.current) { + clearTimeout(timeoutRef.current); + timeoutRef.current = null; + } + }; + }, []); + + if (!isStreaming) return value; + + // Bypass throttle for non-monotonic changes + if ( + typeof value === 'string' && + typeof throttled === 'string' && + !value.startsWith(throttled) + ) { + return value; + } + + return throttled; +} + // `code`/`pre`/`a`/`img` are stable references; only `table` is created per // call (it closes over tableMode/tableResetKey). Recreating the components // object for a table reset therefore never changes the `code` element type, so @@ -783,6 +851,35 @@ function createComponents( const COMPONENTS_DEFAULT = createComponents(); +/** + * Isolated memoized renderer. This ensures react-markdown ONLY re-parses + * when the throttled content or plugin references actually change. + */ +const MemoizedMarkdownRenderer = memo(function MemoizedMarkdownRenderer({ + content, + components, + remarkPlugins, + rehypePlugins, + urlTransform, +}: { + content: string; + components: Options['components']; + remarkPlugins: Options['remarkPlugins']; + rehypePlugins: Options['rehypePlugins']; + urlTransform: Options['urlTransform']; +}) { + return ( + + {content} + + ); +}); + export const Markdown = memo(function Markdown({ content, source, @@ -792,19 +889,28 @@ export const Markdown = memo(function Markdown({ const { markdown, markdownTableMode } = useWebShellCustomization(); const theme = useTheme(); const sourceMarkdown = source ? markdown : undefined; - const renderedContent = - content && source && sourceMarkdown?.transformMarkdown - ? sourceMarkdown.transformMarkdown(content, { source }) - : content; + + const throttledContent = useThrottledValue(content ?? '', isStreaming); + const renderedContent = useMemo( + () => + throttledContent && source && sourceMarkdown?.transformMarkdown + ? sourceMarkdown.transformMarkdown(throttledContent, { source }) + : throttledContent, + [throttledContent, source, sourceMarkdown], + ); + const effectiveTableMode = isStreaming ? 'basic' : (tableMode ?? markdownTableMode ?? 'basic'); + + // Memoize components so references stay stable during throttle window const components = useMemo(() => { if (effectiveTableMode === 'advanced') { return createComponents('advanced', renderedContent); } return COMPONENTS_DEFAULT; }, [effectiveTableMode, renderedContent]); + const sourceComponents = sourceMarkdown?.components; const renderedComponents = useMemo(() => { if (!sourceComponents) return components; @@ -842,23 +948,29 @@ export const Markdown = memo(function Markdown({ [chartPre, renderedComponents], ); + // Memoize plugins so their array references remain stable. + const remarkPlugins = useMemo(() => { + return sourceMarkdown?.remarkPlugins + ? [remarkGfm, remarkMath, ...sourceMarkdown.remarkPlugins] + : [remarkGfm, remarkMath]; + }, [sourceMarkdown?.remarkPlugins]); + + const rehypePlugins = useMemo(() => { + return sourceMarkdown?.rehypePlugins + ? [rehypeKatex, ...sourceMarkdown.rehypePlugins] + : [rehypeKatex]; + }, [sourceMarkdown?.rehypePlugins]); + if (!content) return null; - const remarkPlugins = sourceMarkdown?.remarkPlugins - ? [remarkGfm, remarkMath, ...sourceMarkdown.remarkPlugins] - : [remarkGfm, remarkMath]; - const rehypePlugins = sourceMarkdown?.rehypePlugins - ? [rehypeKatex, ...sourceMarkdown.rehypePlugins] - : [rehypeKatex]; const renderedMarkdown = ( - - {renderedContent} - + /> ); const chartAwareMarkdown = chart ? ( { expect(mountedChart.container.textContent).toContain('Rendering chart'); await mountedChart.rerender(tree(`\`\`\`markdown-chart\n${chart}\n\`\`\``)); + + // Wait for the 80ms streaming throttle to flush the new content + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 100)); + }); + await flushChart(); expect(runtime.init).toHaveBeenCalledOnce();