From be1f3eff2d32f72242ffaf69d7af7af22d675d85 Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Tue, 28 Jul 2026 11:12:15 +0530 Subject: [PATCH 1/4] perf(web-shell): throttle Markdown AST parsing during streaming --- .../components/messages/AssistantMessage.tsx | 38 +---- .../components/messages/Markdown.test.ts | 2 + .../client/components/messages/Markdown.tsx | 136 ++++++++++++++++-- 3 files changed, 125 insertions(+), 51 deletions(-) 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 587941fe0f8..7c899d80078 100644 --- a/packages/web-shell/client/components/messages/Markdown.test.ts +++ b/packages/web-shell/client/components/messages/Markdown.test.ts @@ -1335,6 +1335,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;'); diff --git a/packages/web-shell/client/components/messages/Markdown.tsx b/packages/web-shell/client/components/messages/Markdown.tsx index 8d78e47b286..72a46a1e543 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 remarkGfm from 'remark-gfm'; import remarkMath from 'remark-math'; import rehypeKatex from 'rehype-katex'; @@ -708,6 +708,72 @@ 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: T, + isStreaming: boolean | undefined, + intervalMs: number = 80, +): T { + const [throttled, setThrottled] = useState(value); + 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 (throttled !== 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, throttled]); + + 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 @@ -743,6 +809,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, @@ -751,19 +846,27 @@ export const Markdown = memo(function Markdown({ }: MarkdownProps) { const { markdown, markdownTableMode } = useWebShellCustomization(); const sourceMarkdown = source ? markdown : undefined; - const renderedContent = + + const rawRenderedContent = content && source && sourceMarkdown?.transformMarkdown ? sourceMarkdown.transformMarkdown(content, { source }) : content; + + // Throttle the content that actually reaches the parser + const renderedContent = useThrottledValue(rawRenderedContent, isStreaming); + 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; @@ -774,14 +877,20 @@ export const Markdown = memo(function Markdown({ }; }, [components, effectiveTableMode, sourceComponents]); - if (!content) return null; - const remarkPlugins = sourceMarkdown?.remarkPlugins - ? [remarkGfm, remarkMath, ...sourceMarkdown.remarkPlugins] - : [remarkGfm, remarkMath]; - const rehypePlugins = sourceMarkdown?.rehypePlugins - ? [rehypeKatex, ...sourceMarkdown.rehypePlugins] - : [rehypeKatex]; + // 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; return (
- - {renderedContent} - + />
From a4f2e3d64ebbbd0270752dd93a20c2564c478fbe Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Tue, 28 Jul 2026 23:28:00 +0530 Subject: [PATCH 2/4] perf(web-shell): apply review feedback for throttle hook and tests --- .../components/messages/Markdown.test.ts | 55 ++++++++++++++++++- .../client/components/messages/Markdown.tsx | 18 +++--- 2 files changed, 63 insertions(+), 10 deletions(-) diff --git a/packages/web-shell/client/components/messages/Markdown.test.ts b/packages/web-shell/client/components/messages/Markdown.test.ts index 7c899d80078..18e108be3bd 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, @@ -1486,3 +1486,56 @@ 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, + }), + ); + }); + + 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 72a46a1e543..a590a7bdd4f 100644 --- a/packages/web-shell/client/components/messages/Markdown.tsx +++ b/packages/web-shell/client/components/messages/Markdown.tsx @@ -718,6 +718,8 @@ function useThrottledValue( intervalMs: number = 80, ): T { 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); @@ -730,7 +732,7 @@ function useThrottledValue( timeoutRef.current = null; } // Flush immediately when streaming stops - if (throttled !== value) { + if (throttledRef.current !== value) { setThrottled(value); } return; @@ -749,7 +751,7 @@ function useThrottledValue( setThrottled(valueRef.current); }, intervalMs - elapsed); } - }, [value, isStreaming, intervalMs, throttled]); + }, [value, isStreaming, intervalMs]); useEffect(() => { return () => { @@ -847,13 +849,11 @@ export const Markdown = memo(function Markdown({ const { markdown, markdownTableMode } = useWebShellCustomization(); const sourceMarkdown = source ? markdown : undefined; - const rawRenderedContent = - content && source && sourceMarkdown?.transformMarkdown - ? sourceMarkdown.transformMarkdown(content, { source }) - : content; - - // Throttle the content that actually reaches the parser - const renderedContent = useThrottledValue(rawRenderedContent, isStreaming); + const throttledContent = useThrottledValue(content ?? '', isStreaming); + const renderedContent = + throttledContent && source && sourceMarkdown?.transformMarkdown + ? sourceMarkdown.transformMarkdown(throttledContent, { source }) + : throttledContent; const effectiveTableMode = isStreaming ? 'basic' From 83c9be134789d30e39b733493708aed1d5d8b80e Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Wed, 29 Jul 2026 23:42:29 +0530 Subject: [PATCH 3/4] perf: memoize renderedContent and add throttle regression test --- .../client/components/messages/Markdown.test.ts | 3 +++ .../web-shell/client/components/messages/Markdown.tsx | 11 +++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/packages/web-shell/client/components/messages/Markdown.test.ts b/packages/web-shell/client/components/messages/Markdown.test.ts index 18e108be3bd..525e05cd59f 100644 --- a/packages/web-shell/client/components/messages/Markdown.test.ts +++ b/packages/web-shell/client/components/messages/Markdown.test.ts @@ -1527,6 +1527,9 @@ describe('Markdown streaming throttle', () => { ); }); + expect(container.textContent).toContain('Token 1'); + expect(container.textContent).not.toContain('Token 1 Token 2 Token 3'); + act(() => { vi.advanceTimersByTime(80); }); diff --git a/packages/web-shell/client/components/messages/Markdown.tsx b/packages/web-shell/client/components/messages/Markdown.tsx index a590a7bdd4f..299639e1b10 100644 --- a/packages/web-shell/client/components/messages/Markdown.tsx +++ b/packages/web-shell/client/components/messages/Markdown.tsx @@ -850,10 +850,13 @@ export const Markdown = memo(function Markdown({ const sourceMarkdown = source ? markdown : undefined; const throttledContent = useThrottledValue(content ?? '', isStreaming); - const renderedContent = - throttledContent && source && sourceMarkdown?.transformMarkdown - ? sourceMarkdown.transformMarkdown(throttledContent, { source }) - : throttledContent; + const renderedContent = useMemo( + () => + throttledContent && source && sourceMarkdown?.transformMarkdown + ? sourceMarkdown.transformMarkdown(throttledContent, { source }) + : throttledContent, + [throttledContent, source, sourceMarkdown?.transformMarkdown], + ); const effectiveTableMode = isStreaming ? 'basic' From fd14b047a38a28ba77f50ea54e125bea1a79ed18 Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Thu, 30 Jul 2026 12:00:11 +0530 Subject: [PATCH 4/4] fix(web-shell): resolve markdown throttle conflicts and update tests --- .../web-shell/client/components/messages/Markdown.tsx | 8 ++++---- .../components/messages/MarkdownChartRenderer.test.tsx | 6 ++++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/packages/web-shell/client/components/messages/Markdown.tsx b/packages/web-shell/client/components/messages/Markdown.tsx index bfbdea478fe..554e697186f 100644 --- a/packages/web-shell/client/components/messages/Markdown.tsx +++ b/packages/web-shell/client/components/messages/Markdown.tsx @@ -752,11 +752,11 @@ function MarkdownImage({ src, alt }: { src?: string; alt?: string }) { * 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: T, +function useThrottledValue( + value: string, isStreaming: boolean | undefined, intervalMs: number = 80, -): T { +): string { const [throttled, setThrottled] = useState(value); const throttledRef = useRef(throttled); throttledRef.current = throttled; @@ -896,7 +896,7 @@ export const Markdown = memo(function Markdown({ throttledContent && source && sourceMarkdown?.transformMarkdown ? sourceMarkdown.transformMarkdown(throttledContent, { source }) : throttledContent, - [throttledContent, source, sourceMarkdown?.transformMarkdown], + [throttledContent, source, sourceMarkdown], ); const effectiveTableMode = isStreaming diff --git a/packages/web-shell/client/components/messages/MarkdownChartRenderer.test.tsx b/packages/web-shell/client/components/messages/MarkdownChartRenderer.test.tsx index 6ea1a95b661..b75e7082589 100644 --- a/packages/web-shell/client/components/messages/MarkdownChartRenderer.test.tsx +++ b/packages/web-shell/client/components/messages/MarkdownChartRenderer.test.tsx @@ -429,6 +429,12 @@ describe('Web Shell markdown-chart integration', () => { 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();