From de21723a0daaed3268bbbe78d781802afb577831 Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Tue, 2 Jun 2026 08:27:04 -0300
Subject: [PATCH] Studio: optimize chat streaming by batching renders to one
 per animation frame (#5788)

* Studio: optimize chat streaming by batching renders to one per animation frame

* Fix dev-mode streaming edge case

* Studio: harden streaming markdown coalescing

---------

Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
---
 .../components/assistant-ui/markdown-text.tsx | 53 ++++++++++++++++++-
 1 file changed, 51 insertions(+), 2 deletions(-)

diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
index 8e1d0f03b7..8cb33d0ed2 100644
--- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx
+++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
@@ -382,11 +382,60 @@ function StreamdownBlock(props: BlockProps) {
 }
 const AUDIO_PLAYER_RE = /<audio-player\s+src="([^"]+)"\s*\/>/;
 
+// Coalesce markdown re-parses to one per animation frame while streaming: the
+// runtime notifies on every token (hundreds/sec) and the monitor can't paint
+// that fast. When not streaming we return live text rather than the throttled
+// state, so the final text never lags and a reused instance (parts are keyed by
+// index) shows a completed message's text immediately instead of a stale frame.
+function useRafCoalescedText(text: string, isStreaming: boolean): string {
+  const [displayed, setDisplayed] = useState(text);
+  const pendingRef = useRef(text);
+  const rafRef = useRef<number | null>(null);
+
+  useEffect(() => {
+    pendingRef.current = text;
+    if (!isStreaming) {
+      if (rafRef.current !== null) {
+        cancelAnimationFrame(rafRef.current);
+        rafRef.current = null;
+      }
+      return;
+    }
+    if (rafRef.current === null) {
+      rafRef.current = requestAnimationFrame(() => {
+        rafRef.current = null;
+        setDisplayed(pendingRef.current);
+      });
+    }
+  }, [text, isStreaming]);
+
+  // Unmount cleanup. Cancel the in-flight rAF and null the handle so a
+  // StrictMode remount isn't gated out by a stale id. Kept separate from the
+  // scheduling effect so it doesn't cancel mid-stream and defeat the throttle.
+  useEffect(() => {
+    return () => {
+      if (rafRef.current !== null) {
+        cancelAnimationFrame(rafRef.current);
+        rafRef.current = null;
+      }
+    };
+  }, []);
+
+  if (isStreaming && text.startsWith(displayed)) {
+    return displayed;
+  }
+  return text;
+}
+
 const MarkdownTextImpl = () => {
   const { text, status } = useMessagePartText();
-  const processedText = useMemo(() => preprocessLaTeX(text), [text]);
+  const displayText = useRafCoalescedText(text, status.type === "running");
+  const processedText = useMemo(
+    () => preprocessLaTeX(displayText),
+    [displayText],
+  );
 
-  const audioMatch = text.match(AUDIO_PLAYER_RE);
+  const audioMatch = displayText.match(AUDIO_PLAYER_RE);
   if (audioMatch) {
     return <AudioPlayer src={audioMatch[1]} />;
   }