驯服 AI 聊天流式渲染：RAF 节流、Markdown 分段缓存与虚拟滚动

背景

如果你做过 AI 聊天界面，一定对这种体验不陌生：LLM 的 token 不断流入，UI 需要实时更新，一开始一切正常——直到不正常。包含代码块的长回复开始卡顿，50+ 条消息的对话让滚动变得迟缓，原本觉得够轻量的 Markdown 解析器突然成了瓶颈。

我当时用的是自定义的轻量 Markdown 解析器（刻意避开了沉重的 react-markdown + rehype/remark 插件链），这已经是个不错的起点。但我识别出了三个仍在造成卡顿的具体瓶颈：

1. 每个 SSE chunk 都触发一次完整的 React 重渲染——有时每秒数十次
2. 每次更新都重新解析全部 Markdown 内容——即使 90% 的内容没有变化
3. 所有消息永远留在 DOM 中——长对话导致 DOM 膨胀

让我逐一讲解解决方案。

优化一：基于 requestAnimationFrame 的流式 Chunk 合并

重构前：每个 Chunk = 每次渲染

典型的流式读取循环从 ReadableStream 读取 chunk 并逐个更新状态：

while (true) {
  const { done, value } = await reader.read();
  if (done) break;

  const chunk = decoder.decode(value, { stream: true });
  accumulated += chunk;

  // 每个 chunk 都触发状态更新！
  setMessages((prev) => {
    const updated = [...prev];
    const lastMsg = updated[updated.length - 1];
    if (lastMsg && lastMsg.role === 'assistant') {
      updated[updated.length - 1] = { ...lastMsg, content: accumulated };
    }
    return updated;
  });
}

SSE chunk 可能每 10-50ms 到达一个。每次 setMessages 都触发 React 重渲染，意味着 Markdown 解析器运行、DOM 更新、浏览器重新计算布局——这一切可能每秒发生数十次。其中大部分渲染是浪费的，因为用户根本无法感知那种速度下的差异。

重构后：RAF 节流

浏览器以约 60fps 渲染（每帧约 16ms）。没有比这更快的更新意义。requestAnimationFrame 正是完美的原语——它将更新合并到每帧一次。

const streamingContentRef = useRef('');
const rafIdRef = useRef<number>(0);

// 进入流式循环前：
streamingContentRef.current = '';

const flushStreamUpdate = () => {
  setMessages((prev) => {
    const updated = [...prev];
    const lastMsg = updated[updated.length - 1];
    if (lastMsg && lastMsg.role === 'assistant') {
      updated[updated.length - 1] = {
        ...lastMsg,
        content: streamingContentRef.current,
      };
    }
    return updated;
  });
  rafIdRef.current = 0;
};

while (true) {
  const { done, value } = await reader.read();
  if (done) break;

  const chunk = decoder.decode(value, { stream: true });
  accumulated += chunk;
  streamingContentRef.current = accumulated;

  // 只在没有待执行 RAF 时才调度渲染
  if (rafIdRef.current === 0) {
    rafIdRef.current = requestAnimationFrame(flushStreamUpdate);
  }
}

// 流结束后刷入剩余内容
if (rafIdRef.current !== 0) {
  cancelAnimationFrame(rafIdRef.current);
  rafIdRef.current = 0;
  flushStreamUpdate();
}

关键洞察：streamingContentRef 同步累积最新内容（所以永远不会丢数据），但 setMessages 每帧只调用一次。如果一帧内到达了 5 个 chunk，只有最后一个的内容会被渲染——这正是我们想要的。

别忘了组件卸载时的清理：

useEffect(() => {
  return () => {
    if (rafIdRef.current !== 0) {
      cancelAnimationFrame(rafIdRef.current);
      rafIdRef.current = 0;
    }
  };
}, []);

为什么不用 debounce/throttle？

我首先考虑过 debounce 和 throttle，但 requestAnimationFrame 在这个场景下严格更优：

优化二：Markdown 分段渲染与缓存

重构前：每次更新全量重解析

流式更新期间，Markdown 内容逐步增长。但解析器每次都从头重新解析全部内容。如果 AI 已经写了 500 字并正在添加一个字，我们就在无谓地重新解析那 500 字——实际上只有最后一段在变化。

重构后：按块级边界拆分 + 缓存已完成段

核心思路：在块级边界处将 Markdown 内容拆分为段，缓存已完成段的解析结果，只重新解析最后一段（活跃段）。

将内容拆分为段

function splitContentSegments(content: string): string[] {
  const segments: string[] = [];
  let current = '';
  let inCodeBlock = false;

  for (const line of content.split('\n')) {
    // 尊重代码块边界——绝不在代码块内部拆分
    if (line.startsWith('```')) {
      inCodeBlock = !inCodeBlock;
      current += (current ? '\n' : '') + line;
      if (!inCodeBlock) {
        segments.push(current);
        current = '';
      }
      continue;
    }

    if (inCodeBlock) {
      current += (current ? '\n' : '') + line;
      continue;
    }

    // 空行 = 段落边界
    if (line.trim() === '' && current.trim()) {
      segments.push(current);
      current = '';
      continue;
    }

    current += (current ? '\n' : '') + line;
  }

  if (current.trim()) {
    segments.push(current);
  }

  return segments;
}

拆分逻辑尊重代码块边界（``` 对），所以绝不会把代码块一分为二。空行作为自然的段落分隔符。

带缓存的渲染

const [segmentCache] = useState<Map<string, React.ReactNode>>(() => new Map());

const renderedSegments = useMemo(() => {
  const segments = splitContentSegments(content);
  const result: React.ReactNode[] = [];

  for (let i = 0; i < segments.length; i++) {
    const seg = segments[i];
    const isLast = i === segments.length - 1;

    // 已完成段：优先查缓存
    if (!isLast && segmentCache.has(seg)) {
      result.push(segmentCache.get(seg)!);
    } else {
      // 活跃段（或缓存未命中）：重新解析
      const nodes = (
        <Fragment key={`seg-${i}`}>
          {parseMarkdown(seg, translations, i * 10000)}
        </Fragment>
      );
      // 缓存已完成段以供后续复用
      if (!isLast) {
        segmentCache.set(seg, nodes);
      }
      result.push(nodes);
    }
  }

  // 缓存过大时淘汰旧条目
  if (segmentCache.size > MAX_CACHE_SIZE) {
    const entries = [...segmentCache.entries()];
    segmentCache.clear();
    entries.slice(-MAX_CACHE_SIZE / 2).forEach(([k, v]) => segmentCache.set(k, v));
  }

  return result;
}, [content, translations, segmentCache]);

缓存键就是原始段文本本身。这在流式场景中有效，因为已完成段的文本不会改变——只有最后一段在增长。当一个段"完成"（新段在它之后开始）时，它被缓存。后续渲染时，直接命中缓存。

我使用 useState 的惰性初始化器而非 useRef 来存储缓存 Map。这避免了 React 19 的 react-hooks/refs lint 规则对渲染期间读取 ref 的警告。

效果对比

一段 1000 字的 AI 回复，拆分为约 15 个段：

优化三：消息列表虚拟滚动

重构前：所有消息都在 DOM 中

典型的聊天消息列表用简单的 map 渲染所有消息：

{messages.map((message, index) => (
  <ChatMessage key={message.id} message={message} showCursor={...} />
))}

在 50 条消息的对话中，所有 50 个消息组件（每个可能包含完整的 Markdown 渲染器）都留在 DOM 中。这导致：

• DOM 膨胀：不可见消息产生数百个 DOM 节点
• 滚动缓慢：浏览器必须计算所有节点的布局
• 内存浪费：React 必须维护不可见组件的 fiber 节点

重构后：自定义虚拟滚动

我构建了一个虚拟消息列表，只渲染视口内可见的消息（加一小段缓冲区）。架构如下：

┌─────────────────────────┐
│   顶部占位 (height)      │  ← 不可见消息用占位 div 替代
├─────────────────────────┤
│   消息 #5 (可见)         │  ← 实际渲染
│   消息 #6 (可见)         │  ← 实际渲染
│   消息 #7 (可见)         │  ← 实际渲染
├─────────────────────────┤
│   底部占位 (height)      │  ← 不可见消息用占位 div 替代
└─────────────────────────┘

用 ResizeObserver 追踪高度

聊天消息高度不一（短文本 vs. 长代码块）。我使用 ResizeObserver 测量实际渲染高度，MutationObserver 监听新元素加入 DOM：

useEffect(() => {
  const container = scrollRef.current;
  if (!container) return;

  const resizeObserver = new ResizeObserver((entries) => {
    const updates: Record<string, number> = {};
    for (const entry of entries) {
      const id = (entry.target as HTMLElement).getAttribute('data-msg-id');
      if (id) {
        updates[id] = entry.contentRect.height + MESSAGE_GAP;
      }
    }
    if (Object.keys(updates).length > 0) {
      setHeights((prev) => {
        let changed = false;
        const next = { ...prev };
        for (const [id, height] of Object.entries(updates)) {
          if (prev[id] === undefined || Math.abs(prev[id] - height) > 2) {
            next[id] = height;
            changed = true;
          }
        }
        return changed ? next : prev;
      });
    }
  });

  const observeElements = () => {
    container.querySelectorAll('[data-msg-id]').forEach((el) => {
      resizeObserver.observe(el);
    });
  };

  observeElements();

  const mutationObserver = new MutationObserver(() => {
    observeElements();
  });

  mutationObserver.observe(container, { childList: true, subtree: true });

  return () => {
    resizeObserver.disconnect();
    mutationObserver.disconnect();
  };
}, []);

高度比较中 2px 的容差防止了亚像素舍入差异导致的无限循环。

可见范围计算

有了高度数据，计算哪些消息可见就很直观了：

const visibleRange = useMemo(() => {
  if (messages.length === 0) return { start: 0, end: -1 };

  let start = 0;
  let end = messages.length - 1;
  const bufferHeight = OVERSCAN * (ESTIMATED_HEIGHT + MESSAGE_GAP);
  const viewTop = scrollTop - bufferHeight;
  const viewBottom = scrollTop + viewportHeight + bufferHeight;

  // 扫描第一条可见消息
  for (let i = 0; i < messages.length; i++) {
    const h = heights[messages[i].id] ?? ESTIMATED_HEIGHT + MESSAGE_GAP;
    if (offsets[i] + h >= viewTop) {
      start = Math.max(0, i - OVERSCAN);
      break;
    }
  }

  // 扫描最后一条可见消息
  for (let i = messages.length - 1; i >= 0; i--) {
    if (offsets[i] <= viewBottom) {
      end = Math.min(messages.length - 1, i + OVERSCAN);
      break;
    }
  }

  // 始终包含最后一条消息（保证流式内容可见）
  end = Math.max(end, messages.length - 1);

  return { start, end };
}, [scrollTop, viewportHeight, offsets, heights, messages]);

OVERSCAN = 3 的缓冲区在视口上下各多渲染 3 条消息，确保滚动时不会出现空白闪烁。

智能自动滚动

聊天虚拟列表最棘手的部分之一是自动滚动。行为应该是：

• 用户在底部 → 自动滚动跟随新内容
• 用户已上滑 → 不强制滚动（他们在阅读历史）
• 用户发送新消息 → 始终滚到底部

const isAtBottomRef = useRef(true);

const handleScroll = useCallback(() => {
  const el = scrollRef.current;
  if (!el) return;
  setScrollTop(el.scrollTop);
  // 80px 容差，考虑输入区域和内边距
  isAtBottomRef.current = el.scrollHeight - el.scrollTop - el.clientHeight < 80;
}, []);

// 在底部时自动滚动
useEffect(() => {
  if (!scrollRef.current || !isAtBottomRef.current) return;
  requestAnimationFrame(() => {
    if (scrollRef.current && isAtBottomRef.current) {
      scrollRef.current.scrollTop = scrollRef.current.scrollHeight;
    }
  });
}, [messages, isLoading]);

// 用户发送消息时强制滚到底部
useEffect(() => {
  if (messages.length > prevMsgCountRef.current) {
    const lastMsg = messages[messages.length - 1];
    if (lastMsg?.role === 'user') {
      isAtBottomRef.current = true;
    }
  }
  prevMsgCountRef.current = messages.length;
}, [messages]);

踩坑记录

1. React 19 的 react-hooks/refs 规则

我的第一版实现用 useRef 存储段缓存和高度追踪，在渲染期间读取 ref。React 19 新增的 react-hooks/refs lint 规则会标记这种情况，因为渲染期间读取 ref 可能导致陈旧的 UI（ref 变化不会触发重渲染）。我改用：

• 段缓存：useState<Map> 惰性初始化——Map 跨渲染持久化，但属于"状态"
• 高度追踪：useState<Record<string, number>> 通过 ResizeObserver 回调更新

2. useLayoutEffect 中调用 setState 触发 Lint 错误

React 19 的 react-hooks/set-state-in-effect 规则会标记 effect 体内的同步 setState 调用。我最初用 useLayoutEffect 测量高度并同步调用 setHeights。修复方案是改用 ResizeObserver + MutationObserver——setState 调用发生在观察者回调中（异步的），而非 effect 体内。

3. 虚拟滚动 + 流式 = 始终渲染最后一条消息

一个微妙但关键的细节：流式更新期间，最后一条消息在不断增长。如果用户滚到底部，虚拟列表必须始终将最后一条消息包含在可见范围内，即使滚动位置本应将其排除。否则，当消息增长到足够高将其推出视口时，流式文本就会从 DOM 中消失。

// 始终包含最后一条消息
end = Math.max(end, messages.length - 1);

4. 缓存淘汰策略

段缓存使用简单的"满了清一半"策略。更复杂的 LRU 当然更理想，但对于用户通常单向滚动的聊天场景，这种简单方案已经足够好。200 的 MAX_CACHE_SIZE 足够宽裕，淘汰很少触发。

总结

这三项优化协同工作：RAF 节流降低渲染频率，分段缓存减少每次渲染的解析工作量，虚拟滚动缩减 DOM 体积。最终效果是一个即使面对长流式回复和长对话历史也能保持流畅的聊天界面——而且无需引入 react-window 或 react-virtuoso 等重型依赖。