Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit f13db210 by Archer Committed by GitHub
Options

perf(chat): reduce long stream rendering cost (#7297) 

* perf(chat): reduce long stream rendering cost

* fix(chat): close stream render lifecycle gaps

* fix(markdown): skip animation after hidden tail reveal

* chore(chat): remove analysis docs from code PR
parent ea7aa617
Showing with 813 additions and 150 deletions
projects/app/src/components/Markdown/index.module.scss
.waitingAnimation {
:global(.stream-char) {
opacity: 0;
filter: blur(1px);
transform: translateY(1px);
animation: streamFadeIn 420ms cubic-bezier(0.16, 1, 0.3, 1) forwards;
}
:global(.katex-display) :global(.katex-html) span {
animation: none !important;
:global(.stream-tail) {
display: inline;
will-change: opacity, filter, transform;
}
}
......@@ -36,20 +30,6 @@
}
}
@keyframes streamFadeIn {
from {
opacity: 0;
filter: blur(1px);
transform: translateY(1px);
}
to {
opacity: 1;
filter: blur(0);
transform: translateY(0);
}
}
.markdown > *:first-child {
margin-top: 0 !important;
}
......
projects/app/src/components/Markdown/index.tsx
import React, { useCallback, useContext, useMemo } from 'react';
import React, { useCallback, useContext, useEffect, useLayoutEffect, useMemo, useRef } from 'react';
import ReactMarkdown from 'react-markdown';
import 'katex/dist/katex.min.css';
import RemarkMath from 'remark-math'; // Math syntax
......@@ -15,7 +15,7 @@ import { CodeClassNameEnum, hideStreamingIncompleteMarkdownTail, mdTextFormat }
import type { AProps } from './A';
import MarkdownTable from '@fastgpt/web/components/common/Markdown/MarkdownTable';
import { MarkdownRendererRuntimeContext } from './runtimeContext';
import { useStreamAnimatedRehypePlugin } from './rehypeStreamAnimated';
import { getStreamingAppendLength, rehypeStreamAnimated } from './rehypeStreamAnimated';
const CodeLight = dynamic(() => import('./codeBlock/CodeLight'), { ssr: false });
const MermaidCodeBlock = dynamic(() => import('./img/MermaidCodeBlock'), { ssr: false });
......@@ -25,6 +25,7 @@ const IframeCodeBlock = dynamic(() => import('./codeBlock/Iframe'), { ssr: false
const IframeHtmlCodeBlock = dynamic(() => import('./codeBlock/iframe-html'), { ssr: false });
const VideoBlock = dynamic(() => import('./codeBlock/Video'), { ssr: false });
const AudioBlock = dynamic(() => import('./codeBlock/Audio'), { ssr: false });
const useBrowserLayoutEffect = typeof window === 'undefined' ? useEffect : useLayoutEffect;
const ChatGuide = dynamic(() => import('./chat/Guide'), { ssr: false });
const QuestionGuide = dynamic(() => import('./chat/QuestionGuide'), { ssr: false });
......@@ -67,12 +68,43 @@ function MarkdownLinkRenderer(props: any) {
);
}
/** 仅让最新流式文本批次执行淡入；不支持 Web Animations API 时直接展示内容。 */
function MarkdownStreamTailRenderer({ children }: any) {
const ref = useRef<HTMLSpanElement>(null);
useBrowserLayoutEffect(() => {
const element = ref.current;
if (!element?.animate) return;
const animation = element.animate(
[
{ opacity: 0.25, filter: 'blur(1px)', transform: 'translateY(1px)' },
{ opacity: 1, filter: 'blur(0)', transform: 'translateY(0)' }
],
{
duration: 120,
easing: 'cubic-bezier(0.16, 1, 0.3, 1)',
fill: 'both'
}
);
return () => animation.cancel();
}, [children]);
return (
<span ref={ref} className="stream-tail">
{children}
</span>
);
}
const markdownComponents = {
img: MarkdownImgRenderer,
pre: RewritePre,
code: MarkdownCodeRenderer,
table: MarkdownTable as any,
a: MarkdownLinkRenderer
a: MarkdownLinkRenderer,
'stream-tail': MarkdownStreamTailRenderer
};
type Props = {
......@@ -130,7 +162,32 @@ const MarkdownRender = ({
return mdTextFormat(source);
}, [forbidZhFormat, showAnimation, source]);
const streamAnimatedRehypePlugin = useStreamAnimatedRehypePlugin();
const previousSourceRef = useRef('');
const previousFormatSourceRef = useRef('');
// Markdown 尾部未闭合时会暂时隐藏；闭合后的首次 render 可能同时恢复整段文本，
// 此时无法从原始 source 长度准确区分控制符和可见内容，跳过这一帧的尾部动画。
// refs 保存的是上一次已 commit 的 raw/formatted source，只用于计算本次流式 append 的尾部长度。
// eslint-disable-next-line react-hooks/refs
const previousSource = previousSourceRef.current;
// eslint-disable-next-line react-hooks/refs
const previousFormatSource = previousFormatSourceRef.current;
const hasRevealedMarkdownTail = previousSource !== previousFormatSource;
const streamingTailLength = showAnimation
? getStreamingAppendLength({
previousSource: previousFormatSource,
currentSource: formatSource,
previousSourceWasHidden: hasRevealedMarkdownTail
})
: 0;
useBrowserLayoutEffect(() => {
previousSourceRef.current = source;
previousFormatSourceRef.current = formatSource;
}, [formatSource, source]);
const streamAnimatedRehypePlugin = useMemo(
() => [rehypeStreamAnimated, { tailLength: streamingTailLength }],
[streamingTailLength]
);
const rehypePlugins = useMemo(
() =>
showAnimation
......@@ -153,7 +210,7 @@ const MarkdownRender = ({
`}
remarkPlugins={[RemarkMath, [RemarkGfm, { singleTilde: false }], RemarkBreaks]}
rehypePlugins={rehypePlugins as any}
components={markdownComponents}
components={markdownComponents as any}
urlTransform={urlTransform}
>
{formatSource}
......
projects/app/src/components/Markdown/rehypeStreamAnimated.ts
import { useMemo } from 'react';
const STREAM_ANIMATED_BLOCK_TAGS = new Set(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']);
const STREAM_ANIMATED_BLOCK_TAGS = new Set(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'li']);
const STREAM_ANIMATED_SKIP_TAGS = new Set(['pre', 'code', 'table', 'svg']);
const STREAM_ANIMATED_LIST_INLINE_TAGS = new Set([
'a',
'abbr',
'b',
'cite',
'del',
'em',
'i',
'ins',
'kbd',
'mark',
's',
'small',
'span',
'strong',
'sub',
'sup',
'u'
]);
const DEFAULT_STREAMING_TAIL_MAX_LENGTH = 64;
const DEFAULT_STREAMING_TAIL_TAG_NAME = 'stream-tail';
const STREAMING_MARKDOWN_SYNTAX_CHARS = new Set('*_~`[]()>#+-|\\');
type HastElement = {
type: 'element';
......@@ -28,12 +10,13 @@ type HastElement = {
properties?: Record<string, any>;
children?: HastNode[];
};
type HastText = {
type: 'text';
value: string;
};
type HastNode =
| HastElement
| {
type: 'text';
value: string;
}
| HastText
| {
type: string;
[key: string]: any;
......@@ -44,105 +27,137 @@ type HastRoot = {
};
/**
* 流式输出时按 Lobe UI 的方式把新增文本拆成字符 span 做淡入动画。
* 只处理段落、标题和列表里的文本，跳过代码块、表格、公式等高复杂度内容，避免流式阶段 DOM 过重。
* 计算相对上一次已提交 Markdown 新增可见尾部的 Unicode code point 数量。
*
* 只接受纯 append，避免 Markdown 尾部隐藏、内容替换或会话切换时把旧正文误判为新增内容。
* 首尾空白和纯 Markdown 控制标记不产生可见动画。若上一帧存在被隐藏的未闭合 Markdown
* 尾部，闭合时无法仅凭 source 长度区分控制符和可见内容，因此跳过这一帧的尾部动画。
* 返回值有上限，保证单次大 chunk 也只创建固定规模的动画节点。
*/
export const getStreamingAppendLength = ({
previousSource,
currentSource,
previousSourceWasHidden = false,
maxLength = DEFAULT_STREAMING_TAIL_MAX_LENGTH
}: {
previousSource: string;
currentSource: string;
previousSourceWasHidden?: boolean;
maxLength?: number;
}) => {
if (previousSourceWasHidden) return 0;
if (!currentSource.startsWith(previousSource)) return 0;
const appendedSource = currentSource.slice(previousSource.length).trim();
if (!appendedSource) return 0;
const appendedCodePoints = Array.from(appendedSource);
if (appendedCodePoints.every((codePoint) => STREAMING_MARKDOWN_SYNTAX_CHARS.has(codePoint))) {
return 0;
}
return Math.min(appendedCodePoints.length, maxLength);
};
/**
* 只包装最后一个 Markdown 文本块的最新尾部，供流式淡入 renderer 使用。
*
* 与原字符级实现不同，本插件不会重写累计全文。它从最后一个可见 block 的末尾向前消费
* `tailLength` 个 code point，一个原始 text node 最多生成一个 tail element。代码、表格、
* SVG 和 KaTeX 是边界；遇到这些节点后不会继续向前包装旧正文。
*/
export const rehypeStreamAnimated = ({ className = 'stream-char' }: { className?: string }) => {
export const rehypeStreamAnimated = ({
tailLength,
tailTagName = DEFAULT_STREAMING_TAIL_TAG_NAME
}: {
tailLength: number;
tailTagName?: string;
}) => {
return (tree: HastRoot) => {
const isHastElement = (node: HastNode): node is HastElement => {
return node.type === 'element' && typeof (node as HastElement).tagName === 'string';
};
if (tailLength <= 0) return;
const isHastElement = (node: HastNode): node is HastElement =>
node.type === 'element' && typeof (node as HastElement).tagName === 'string';
const isHastText = (node: HastNode): node is HastText =>
node.type === 'text' && typeof (node as HastText).value === 'string';
const hasClass = (node: HastElement, cls: string) => {
const className = node.properties?.className;
if (Array.isArray(className)) return className.some((item) => String(item).includes(cls));
if (typeof className === 'string') return className.includes(cls);
return false;
};
const shouldSkip = (node: HastElement) => {
return STREAM_ANIMATED_SKIP_TAGS.has(node.tagName) || hasClass(node, 'katex');
const shouldSkip = (node: HastElement) =>
STREAM_ANIMATED_SKIP_TAGS.has(node.tagName) || hasClass(node, 'katex');
const hasRenderableText = (node: HastElement): boolean => {
if (shouldSkip(node)) return false;
return !!node.children?.some((child) => {
if (isHastText(child)) return child.value.length > 0;
return isHastElement(child) && hasRenderableText(child);
});
};
let lastTextBlock: HastElement | undefined;
const findLastTextBlock = (node: HastNode, activeTextBlock?: HastElement) => {
if (!isHastElement(node)) return;
if (shouldSkip(node)) {
// 最近候选是当前文本块本身时，跳过节点可能只是行内内容；若候选来自嵌套 block，
// 同级代码、表格或公式则表示尾部已越过候选，不能回退动画更早正文。
if (lastTextBlock !== activeTextBlock) lastTextBlock = undefined;
return;
}
const createStreamCharNode = (char: string): HastElement => {
return {
type: 'element',
tagName: 'span',
properties: { className },
children: [{ type: 'text', value: char }]
};
const renderableText = hasRenderableText(node);
const isTextBlock = STREAM_ANIMATED_BLOCK_TAGS.has(node.tagName) && renderableText;
if (isTextBlock) {
lastTextBlock = node;
}
const nextActiveTextBlock = isTextBlock ? node : activeTextBlock;
node.children?.forEach((child) => findLastTextBlock(child, nextActiveTextBlock));
};
tree.children?.forEach((child) => findLastTextBlock(child));
if (!lastTextBlock) return;
const wrapText = (node: HastElement) => {
const newChildren: HastNode[] = [];
let remainingLength = tailLength;
node.children?.forEach((child) => {
if (child.type === 'text') {
for (const char of child.value) {
newChildren.push(createStreamCharNode(char));
}
return;
}
/** 从尾部反向包装；false 表示遇到不可跨越的渲染边界。 */
const wrapTail = (node: HastElement): boolean => {
if (!node.children) return true;
if (isHastElement(child) && !shouldSkip(child)) {
wrapText(child);
}
newChildren.push(child);
});
for (let index = node.children.length - 1; index >= 0 && remainingLength > 0; index--) {
const child = node.children[index];
node.children = newChildren;
};
if (isHastText(child)) {
const codePoints = Array.from(child.value);
if (codePoints.length === 0) continue;
const wrapListItemText = (node: HastElement) => {
const newChildren: HastNode[] = [];
const animatedLength = Math.min(codePoints.length, remainingLength);
const stableText = codePoints.slice(0, -animatedLength).join('');
const animatedText = codePoints.slice(-animatedLength).join('');
const nextChildren: HastNode[] = [];
node.children?.forEach((child) => {
if (child.type === 'text') {
for (const char of child.value) {
newChildren.push(createStreamCharNode(char));
if (stableText) {
nextChildren.push({ type: 'text', value: stableText });
}
return;
}
if (
isHastElement(child) &&
(child.tagName === 'p' || STREAM_ANIMATED_LIST_INLINE_TAGS.has(child.tagName))
) {
wrapText(child);
nextChildren.push({
type: 'element',
tagName: tailTagName,
properties: {},
children: [{ type: 'text', value: animatedText }]
});
node.children.splice(index, 1, ...nextChildren);
remainingLength -= animatedLength;
continue;
}
newChildren.push(child);
});
node.children = newChildren;
};
const visitElement = (node: HastNode): 'skip' | undefined => {
if (!isHastElement(node)) return;
if (shouldSkip(node)) return 'skip';
if (node.tagName === 'li') {
wrapListItemText(node);
return 'skip';
}
if (STREAM_ANIMATED_BLOCK_TAGS.has(node.tagName)) {
wrapText(node);
return 'skip';
if (!isHastElement(child)) continue;
if (shouldSkip(child)) return false;
if (!wrapTail(child)) return false;
}
node.children?.forEach((child) => {
if (visitElement(child) === 'skip') return;
});
return true;
};
tree.children?.forEach((child) => {
visitElement(child);
});
wrapTail(lastTextBlock);
};
};
/**
* 返回流式 Markdown 字符淡入插件。
* 不在 React render 期间读写 ref；依赖 React 复用已存在的字符 span，仅让新增字符节点触发 CSS animation。
*/
export const useStreamAnimatedRehypePlugin = () => {
return useMemo(() => [rehypeStreamAnimated, { className: 'stream-char' }], []);
};
projects/app/src/components/core/chat/ChatContainer/ChatBox/hooks/useChatGenerate.ts
import { type MutableRefObject, useEffect, useRef } from 'react';
import { type MutableRefObject, useEffect, useMemo, useRef } from 'react';
import { useContextSelector } from 'use-context-selector';
import { useTranslation } from 'next-i18next';
import { useMemoizedFn } from 'ahooks';
......@@ -38,6 +38,7 @@ import {
} from '../utils/interactive';
import { shouldAppendResumeInteractive } from '../utils/resume';
import { formatChatRequestVariables } from '../utils/requestVariables';
import { createStreamRenderScheduler } from '../utils/streamRenderScheduler';
import type { ChatSiteItemType, ChatBoxInputType, SendPromptFnType } from '../type';
import type { StartChatFnProps, generatingMessageProps } from '../../type';
import { cloneDeep } from 'lodash';
......@@ -156,7 +157,6 @@ export const useChatGenerate = ({
const generatingMessageQueueRef = useRef<
Array<generatingMessageProps & { autoTTSResponse?: boolean }>
>([]);
const generatingMessageFrameRef = useRef<number>();
const applyGeneratingMessage = useMemoizedFn(
(
......@@ -524,14 +524,9 @@ export const useChatGenerate = ({
}
);
const flushGeneratingMessageQueue = useMemoizedFn(() => {
if (generatingMessageFrameRef.current !== undefined) {
window.cancelAnimationFrame(generatingMessageFrameRef.current);
generatingMessageFrameRef.current = undefined;
}
const commitGeneratingMessageQueue = useMemoizedFn(() => {
const queue = generatingMessageQueueRef.current;
if (queue.length === 0) return;
if (queue.length === 0) return false;
generatingMessageQueueRef.current = [];
setChatRecords((state) =>
......@@ -539,6 +534,21 @@ export const useChatGenerate = ({
);
generatingScroll(queue.some((message) => message.event === SseResponseEventEnum.interactive));
return true;
});
const streamRenderScheduler = useMemo(
() =>
createStreamRenderScheduler({
onFlush: commitGeneratingMessageQueue
}),
[commitGeneratingMessageQueue]
);
const flushGeneratingMessageQueue = useMemoizedFn(() => {
streamRenderScheduler.flush();
});
const cancelGeneratingMessageQueue = useMemoizedFn(() => {
streamRenderScheduler.cancel();
generatingMessageQueueRef.current = [];
});
const generatingMessage = useMemoizedFn(
......@@ -561,24 +571,20 @@ export const useChatGenerate = ({
}
generatingMessageQueueRef.current.push(message);
if (generatingMessageFrameRef.current !== undefined) return;
generatingMessageFrameRef.current = window.requestAnimationFrame(flushGeneratingMessageQueue);
streamRenderScheduler.schedule();
}
);
useEffect(() => {
return () => {
if (generatingMessageFrameRef.current !== undefined) {
window.cancelAnimationFrame(generatingMessageFrameRef.current);
}
generatingMessageFrameRef.current = undefined;
generatingMessageQueueRef.current = [];
cancelGeneratingMessageQueue();
};
}, []);
}, [cancelGeneratingMessageQueue]);
const abortRequest = useMemoizedFn((reason: string = 'stop') => {
if (reason === 'leave') {
cancelGeneratingMessageQueue();
}
chatControllerRef.current?.abort(new Error(reason));
questionGuideControllerRef.current?.abort(new Error(reason));
pluginControllerRef.current?.abort(new Error(reason));
......@@ -858,6 +864,7 @@ export const useChatGenerate = ({
return {
abortRequest,
flushGeneratingMessages: flushGeneratingMessageQueue,
generatingMessage,
sendPrompt
};
......
projects/app/src/components/core/chat/ChatContainer/ChatBox/hooks/useChatResume.ts
......@@ -49,6 +49,7 @@ type UseChatResumeProps = {
resumedChatTargetRef: MutableRefObject<string | undefined>;
resumeControllerRef: MutableRefObject<AbortController | undefined>;
generatingMessage: (message: generatingMessageProps) => void;
flushGeneratingMessages: () => void;
scrollToBottom: (behavior?: 'smooth' | 'auto', delay?: number) => void;
finishChatGenerateStatus: FinishChatGenerateStatus;
};
......@@ -70,6 +71,8 @@ const isAbortByLeave = (reason: unknown) => {
* 输入约定：
* - `generatingMessage` 仍由 ChatBox 提供，确保普通发送和恢复生成继续共享同一套
* answer/reasoning/tool/plan/interactive 合并逻辑。
* - `flushGeneratingMessages` 在恢复流收尾前提交最后一个 50ms buffer，避免 completedChat、
* finish 或 error 状态先于最后一批 SSE 增量写入。
* - `activeSourceKeyRef/activeChatIdRef` 保存当前页面真实目标，用于防止恢复流异步返回后
* 写入已经切走的会话。
* - `resumedChatTargetRef` 记录本轮已经尝试恢复的 app/chat，避免同一个 generating
......@@ -93,6 +96,7 @@ export const useChatResume = ({
resumedChatTargetRef,
resumeControllerRef,
generatingMessage,
flushGeneratingMessages,
scrollToBottom,
finishChatGenerateStatus
}: UseChatResumeProps) => {
......@@ -256,6 +260,8 @@ export const useChatResume = ({
if (!isActiveResumeTarget({ sourceKey: resumeForSourceKey, chatId: resumeForChatId }))
return;
flushGeneratingMessages();
if (completedChat) {
resumeFinalStatus = completedChat.chatGenerateStatus;
setChatRecords((state) =>
......@@ -351,10 +357,21 @@ export const useChatResume = ({
});
scrollToBottom('auto');
} catch (error) {
if (controller.signal.aborted) return;
if (controller.signal.aborted) {
// 离开页面时不再提交旧会话数据；用户主动停止时仍需立刻落下最后一个 buffer。
if (
!isAbortByLeave(controller.signal.reason) &&
isActiveResumeTarget({ sourceKey: resumeForSourceKey, chatId: resumeForChatId })
) {
flushGeneratingMessages();
}
return;
}
if (!isActiveResumeTarget({ sourceKey: resumeForSourceKey, chatId: resumeForChatId }))
return;
flushGeneratingMessages();
const isStreamError = (error as ResumeStreamErrorType | undefined)?.isStreamError === true;
resumeFinalStatus = isStreamError
? ChatGenerateStatusEnum.error
......@@ -433,6 +450,7 @@ export const useChatResume = ({
chatBoxSourceKey,
chatBoxChatId,
chatGenerateStatus,
flushGeneratingMessages,
generatingMessage,
resumeTargetAiDataId,
scrollToBottom,
......
projects/app/src/components/core/chat/ChatContainer/ChatBox/index.tsx
......@@ -344,7 +344,7 @@ const ChatBox = ({
generatingScroll
});
const { abortRequest, generatingMessage, sendPrompt } = useChatGenerate({
const { abortRequest, flushGeneratingMessages, generatingMessage, sendPrompt } = useChatGenerate({
onStartChat,
isRoundPending,
chatControllerRef: chatController,
......@@ -460,6 +460,7 @@ const ChatBox = ({
resumedChatTargetRef,
resumeControllerRef: resumeController,
generatingMessage,
flushGeneratingMessages,
scrollToBottom,
finishChatGenerateStatus
});
......
projects/app/src/components/core/chat/ChatContainer/ChatBox/utils/streamRenderScheduler.ts 0 → 100644
export const STREAM_RENDER_INTERVAL_MS = 50;
export type StreamRenderSchedulerRuntime = {
now: () => number;
setTimer: (callback: () => void, delay: number) => number;
clearTimer: (id: number) => void;
requestFrame: (callback: () => void) => number;
cancelFrame: (id: number) => void;
};
const browserRuntime: StreamRenderSchedulerRuntime = {
now: () => performance.now(),
setTimer: (callback, delay) => window.setTimeout(callback, delay),
clearTimer: (id) => window.clearTimeout(id),
requestFrame: (callback) => window.requestAnimationFrame(callback),
cancelFrame: (id) => window.cancelAnimationFrame(id)
};
/**
* 创建流式消息 UI 调度器。
*
* 普通增量两次 flush 至少间隔 `intervalMs`，timer 到期后再通过 rAF 对齐 paint。`flush`
* 用于完成和异常收尾，会取消待执行任务并立即提交；`cancel` 用于离开会话，只清理不提交。
*/
export const createStreamRenderScheduler = ({
onFlush,
intervalMs = STREAM_RENDER_INTERVAL_MS,
runtime = browserRuntime
}: {
onFlush: () => boolean | void;
intervalMs?: number;
runtime?: StreamRenderSchedulerRuntime;
}) => {
let lastFlushAt = Number.NEGATIVE_INFINITY;
let timerId: number | undefined;
let frameId: number | undefined;
const cancel = () => {
lastFlushAt = Number.NEGATIVE_INFINITY;
if (timerId !== undefined) {
runtime.clearTimer(timerId);
timerId = undefined;
}
if (frameId !== undefined) {
runtime.cancelFrame(frameId);
frameId = undefined;
}
};
const commit = () => {
// 空 flush 只负责清理待执行任务，不应阻塞下一轮流式输出的首次提交。
if (onFlush() !== false) {
lastFlushAt = runtime.now();
}
};
const schedule = () => {
if (timerId !== undefined || frameId !== undefined) return;
const elapsed = runtime.now() - lastFlushAt;
const delay = Number.isFinite(elapsed) ? Math.max(intervalMs - elapsed, 0) : 0;
timerId = runtime.setTimer(() => {
timerId = undefined;
frameId = runtime.requestFrame(() => {
frameId = undefined;
commit();
});
}, delay);
};
const flush = () => {
cancel();
commit();
};
return {
schedule,
flush,
cancel
};
};
projects/app/test/components/core/chat/ChatContainer/ChatBox/streamRenderScheduler.test.ts 0 → 100644
import { describe, expect, it, vi } from 'vitest';
import {
createStreamRenderScheduler,
STREAM_RENDER_INTERVAL_MS,
type StreamRenderSchedulerRuntime
} from '@/components/core/chat/ChatContainer/ChatBox/utils/streamRenderScheduler';
const createFakeRuntime = () => {
let now = 100;
let nextId = 1;
const timers = new Map<number, { callback: () => void; delay: number }>();
const frames = new Map<number, () => void>();
const runtime: StreamRenderSchedulerRuntime = {
now: () => now,
setTimer: (callback, delay) => {
const id = nextId++;
timers.set(id, { callback, delay });
return id;
},
clearTimer: (id) => {
timers.delete(id);
},
requestFrame: (callback) => {
const id = nextId++;
frames.set(id, callback);
return id;
},
cancelFrame: (id) => {
frames.delete(id);
}
};
return {
runtime,
timers,
frames,
setNow: (value: number) => {
now = value;
},
runNextTimer: () => {
const entry = timers.entries().next().value as
| [number, { callback: () => void; delay: number }]
| undefined;
if (!entry) throw new Error('No timer scheduled');
timers.delete(entry[0]);
entry[1].callback();
},
runNextFrame: () => {
const entry = frames.entries().next().value as [number, () => void] | undefined;
if (!entry) throw new Error('No frame scheduled');
frames.delete(entry[0]);
entry[1]();
}
};
};
describe('createStreamRenderScheduler', () => {
it('should align the first flush to the next frame without waiting a full interval', () => {
const fake = createFakeRuntime();
const onFlush = vi.fn();
const scheduler = createStreamRenderScheduler({ onFlush, runtime: fake.runtime });
scheduler.schedule();
expect([...fake.timers.values()].map((item) => item.delay)).toEqual([0]);
fake.runNextTimer();
expect(fake.frames).toHaveLength(1);
fake.runNextFrame();
expect(onFlush).toHaveBeenCalledTimes(1);
});
it('should coalesce schedules and wait until 50ms after the previous flush', () => {
const fake = createFakeRuntime();
const onFlush = vi.fn();
const scheduler = createStreamRenderScheduler({ onFlush, runtime: fake.runtime });
scheduler.schedule();
scheduler.schedule();
expect(fake.timers).toHaveLength(1);
fake.runNextTimer();
fake.runNextFrame();
fake.setNow(110);
scheduler.schedule();
scheduler.schedule();
expect([...fake.timers.values()].map((item) => item.delay)).toEqual([
STREAM_RENDER_INTERVAL_MS - 10
]);
fake.setNow(150);
fake.runNextTimer();
fake.runNextFrame();
expect(onFlush).toHaveBeenCalledTimes(2);
});
it('should flush immediately and cancel pending work', () => {
const fake = createFakeRuntime();
const onFlush = vi.fn();
const scheduler = createStreamRenderScheduler({ onFlush, runtime: fake.runtime });
scheduler.schedule();
scheduler.flush();
expect(fake.timers).toHaveLength(0);
expect(fake.frames).toHaveLength(0);
expect(onFlush).toHaveBeenCalledTimes(1);
});
it('should cancel pending work without flushing', () => {
const fake = createFakeRuntime();
const onFlush = vi.fn();
const scheduler = createStreamRenderScheduler({ onFlush, runtime: fake.runtime });
scheduler.schedule();
fake.runNextTimer();
scheduler.cancel();
expect(fake.timers).toHaveLength(0);
expect(fake.frames).toHaveLength(0);
expect(onFlush).not.toHaveBeenCalled();
});
it('should reset the interval after canceling a completed stream', () => {
const fake = createFakeRuntime();
const onFlush = vi.fn();
const scheduler = createStreamRenderScheduler({ onFlush, runtime: fake.runtime });
scheduler.schedule();
fake.runNextTimer();
fake.runNextFrame();
fake.setNow(110);
scheduler.cancel();
scheduler.schedule();
expect([...fake.timers.values()].map((item) => item.delay)).toEqual([0]);
});
it('should not throttle after an empty flush', () => {
const fake = createFakeRuntime();
const onFlush = vi.fn(() => false);
const scheduler = createStreamRenderScheduler({ onFlush, runtime: fake.runtime });
scheduler.flush();
scheduler.schedule();
expect([...fake.timers.values()].map((item) => item.delay)).toEqual([0]);
});
it('should not schedule another timer while a frame is pending', () => {
const fake = createFakeRuntime();
const scheduler = createStreamRenderScheduler({ onFlush: vi.fn(), runtime: fake.runtime });
scheduler.schedule();
fake.runNextTimer();
scheduler.schedule();
expect(fake.timers).toHaveLength(0);
expect(fake.frames).toHaveLength(1);
});
it('should use the browser runtime by default', () => {
let timerCallback: (() => void) | undefined;
let frameCallback: (() => void) | undefined;
const setTimeout = vi.fn((callback: () => void) => {
timerCallback = callback;
return 1;
});
const clearTimeout = vi.fn();
const requestAnimationFrame = vi.fn((callback: () => void) => {
frameCallback = callback;
return 2;
});
const cancelAnimationFrame = vi.fn();
vi.stubGlobal('window', {
setTimeout,
clearTimeout,
requestAnimationFrame,
cancelAnimationFrame
});
const onFlush = vi.fn();
const scheduler = createStreamRenderScheduler({ onFlush });
scheduler.schedule();
timerCallback?.();
frameCallback?.();
scheduler.schedule();
scheduler.flush();
scheduler.schedule();
timerCallback?.();
scheduler.cancel();
expect(setTimeout).toHaveBeenCalled();
expect(requestAnimationFrame).toHaveBeenCalled();
expect(clearTimeout).toHaveBeenCalled();
expect(cancelAnimationFrame).toHaveBeenCalled();
expect(onFlush).toHaveBeenCalledTimes(2);
vi.unstubAllGlobals();
});
});
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment