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>
This commit is contained in:
oobabooga 2026-06-02 08:27:04 -03:00 committed by GitHub
commit de21723a0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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]} />;
}