feat (web-shell): throttle Markdown AST parsing during streaming

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<ReturnType<typeof setTimeout> | 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({
         >
           <div className={styles.contentBody}>
             <Markdown
-              content={markdownContent}
+              content={content}
               source="assistant"
               isStreaming={isStreaming}
             />

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();
+  });
+});

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 <img src={safeSrc} alt={alt || ''} className={styles.image} />;
 }
 
+/**
+ * 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<ReturnType<typeof setTimeout> | 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 (
+    <ReactMarkdown
+      components={components}
+      remarkPlugins={remarkPlugins}
+      rehypePlugins={rehypePlugins}
+      urlTransform={urlTransform}
+    >
+      {content}
+    </ReactMarkdown>
+  );
+});
+
 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 = (
-    <ReactMarkdown
+    <MemoizedMarkdownRenderer
+      content={renderedContent}
+      components={componentsWithCharts}
       remarkPlugins={remarkPlugins}
       rehypePlugins={rehypePlugins}
-      components={componentsWithCharts}
       urlTransform={markdownUrlTransform}
-    >
-      {renderedContent}
-    </ReactMarkdown>
+    />
   );
   const chartAwareMarkdown = chart ? (
     <WebShellMarkdownChartProvider

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();