refactor(CodePreview): handle chunking in ShikiStreamService, make the algorithm more robust (#7409)

* refactor(ShikiStreamService, CodePreview): handle chunking in ShikiStreamService, make the algorithm more robust

- Add highlightStreamingCode with improved robustness
- Improve viewport detection

* perf: improve checks for appending

* chore: update comments

src/renderer/src/components/CodeBlockView/CodePreview.tsx

@@ -4,7 +4,7 @@ import { useSettings } from '@renderer/hooks/useSettings'
 import { uuid } from '@renderer/utils'
 import { getReactStyleFromToken } from '@renderer/utils/shiki'
 import { ChevronsDownUp, ChevronsUpDown, Text as UnWrapIcon, WrapText as WrapIcon } from 'lucide-react'
-import React, { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react'
+import React, { memo, useCallback, useEffect, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { ThemedToken } from 'shiki/core'
 import styled from 'styled-components'
@@ -18,19 +18,20 @@ interface CodePreviewProps {
 /**
  * Shiki 流式代码高亮组件
  *
- * - 通过 shiki tokenizer 处理流式响应
- * - 为了正确执行语法高亮，必须保证流式响应都依次到达 tokenizer，不能跳过
+ * - 通过 shiki tokenizer 处理流式响应，高性能
+ * - 进入视口后触发高亮，改善页面内有大量长代码块时的响应
+ * - 并发安全
  */
 const CodePreview = ({ children, language, setTools }: CodePreviewProps) => {
   const { codeShowLineNumbers, fontSize, codeCollapsible, codeWrappable } = useSettings()
-  const { activeShikiTheme, highlightCodeChunk, cleanupTokenizers } = useCodeStyle()
+  const { activeShikiTheme, highlightStreamingCode, cleanupTokenizers } = useCodeStyle()
   const [isExpanded, setIsExpanded] = useState(!codeCollapsible)
   const [isUnwrapped, setIsUnwrapped] = useState(!codeWrappable)
   const [tokenLines, setTokenLines] = useState<ThemedToken[][]>([])
-  const codeContentRef = useRef<HTMLDivElement>(null)
-  const prevCodeLengthRef = useRef(0)
-  const safeCodeStringRef = useRef(children)
-  const highlightQueueRef = useRef<Promise<void>>(Promise.resolve())
+  const [isInViewport, setIsInViewport] = useState(false)
+  const codeContainerRef = useRef<HTMLDivElement>(null)
+  const processingRef = useRef(false)
+  const latestRequestedContentRef = useRef<string | null>(null)
   const callerId = useRef(`${Date.now()}-${uuid()}`).current
   const shikiThemeRef = useRef(activeShikiTheme)
 
@@ -45,7 +46,7 @@ const CodePreview = ({ children, language, setTools }: CodePreviewProps) => {
       icon: isExpanded ? <ChevronsDownUp className="icon" /> : <ChevronsUpDown className="icon" />,
       tooltip: isExpanded ? t('code_block.collapse') : t('code_block.expand'),
       visible: () => {
-        const scrollHeight = codeContentRef.current?.scrollHeight
+        const scrollHeight = codeContainerRef.current?.scrollHeight
         return codeCollapsible && (scrollHeight ?? 0) > 350
       },
       onClick: () => setIsExpanded((prev) => !prev)
@@ -77,81 +78,63 @@ const CodePreview = ({ children, language, setTools }: CodePreviewProps) => {
     setIsUnwrapped(!codeWrappable)
   }, [codeWrappable])
 
-  // 处理尾部空白字符
-  const safeCodeString = useMemo(() => {
-    return typeof children === 'string' ? children.trimEnd() : ''
-  }, [children])
-
   const highlightCode = useCallback(async () => {
-    if (!safeCodeString) return
+    const currentContent = typeof children === 'string' ? children.trimEnd() : ''
 
-    if (prevCodeLengthRef.current === safeCodeString.length) return
+    // 记录最新要处理的内容，为了保证最终状态正确
+    latestRequestedContentRef.current = currentContent
 
-    // 捕获当前状态
-    const startPos = prevCodeLengthRef.current
-    const endPos = safeCodeString.length
+    // 如果正在处理，先跳出，等到完成后会检查是否有新内容
+    if (processingRef.current) return
 
-    // 添加到处理队列，确保按顺序处理
-    highlightQueueRef.current = highlightQueueRef.current.then(async () => {
-      // FIXME: 长度有问题，或者破坏了流式内容，需要清理 tokenizer 并使用完整代码重新高亮
-      if (prevCodeLengthRef.current > safeCodeString.length || !safeCodeString.startsWith(safeCodeStringRef.current)) {
-        cleanupTokenizers(callerId)
-        prevCodeLengthRef.current = 0
-        safeCodeStringRef.current = ''
+    processingRef.current = true
 
-        const result = await highlightCodeChunk(safeCodeString, language, callerId)
-        setTokenLines(result.lines)
+    try {
+      // 循环处理，确保会处理最新内容
+      while (latestRequestedContentRef.current !== null) {
+        const contentToProcess = latestRequestedContentRef.current
+        latestRequestedContentRef.current = null // 标记开始处理
 
-        prevCodeLengthRef.current = safeCodeString.length
-        safeCodeStringRef.current = safeCodeString
+        // 传入完整内容，让 ShikiStreamService 检测变化并处理增量高亮
+        const result = await highlightStreamingCode(contentToProcess, language, callerId)
 
-        return
+        // 如有结果，更新 tokenLines
+        if (result.lines.length > 0 || result.recall !== 0) {
+          setTokenLines((prev) => {
+            return result.recall === -1
+              ? result.lines
+              : [...prev.slice(0, Math.max(0, prev.length - result.recall)), ...result.lines]
+          })
+        }
       }
-
-      // 跳过 race condition，延迟到后续任务
-      if (prevCodeLengthRef.current !== startPos) {
-        return
-      }
-
-      const incrementalCode = safeCodeString.slice(startPos, endPos)
-      const result = await highlightCodeChunk(incrementalCode, language, callerId)
-      setTokenLines((lines) => [...lines.slice(0, Math.max(0, lines.length - result.recall)), ...result.lines])
-      prevCodeLengthRef.current = endPos
-      safeCodeStringRef.current = safeCodeString
-    })
-  }, [callerId, cleanupTokenizers, highlightCodeChunk, language, safeCodeString])
+    } finally {
+      processingRef.current = false
+    }
+  }, [highlightStreamingCode, language, callerId, children])
 
   // 主题变化时强制重新高亮
   useEffect(() => {
     if (shikiThemeRef.current !== activeShikiTheme) {
-      prevCodeLengthRef.current++
       shikiThemeRef.current = activeShikiTheme
+      cleanupTokenizers(callerId)
+      setTokenLines([])
     }
-  }, [activeShikiTheme])
+  }, [activeShikiTheme, callerId, cleanupTokenizers])
 
   // 组件卸载时清理资源
   useEffect(() => {
     return () => cleanupTokenizers(callerId)
   }, [callerId, cleanupTokenizers])
 
-  // 触发代码高亮
-  // - 进入视口后触发第一次高亮
-  // - 内容变化后触发之后的高亮
+  // 视口检测逻辑，进入视口后触发第一次代码高亮
   useEffect(() => {
-    let isMounted = true
-
-    if (prevCodeLengthRef.current > 0) {
-      setTimeout(highlightCode, 0)
-      return
-    }
-
-    const codeElement = codeContentRef.current
+    const codeElement = codeContainerRef.current
     if (!codeElement) return
 
     const observer = new IntersectionObserver(
       (entries) => {
-        if (entries[0].intersectionRatio > 0 && isMounted) {
-          setTimeout(highlightCode, 0)
+        if (entries[0].intersectionRatio > 0) {
+          setIsInViewport(true)
           observer.disconnect()
         }
       },
@@ -161,15 +144,18 @@ const CodePreview = ({ children, language, setTools }: CodePreviewProps) => {
     )
 
     observer.observe(codeElement)
+    return () => observer.disconnect()
+  }, []) // 只执行一次
 
-    return () => {
-      isMounted = false
-      observer.disconnect()
-    }
-  }, [highlightCode])
+  // 触发代码高亮
+  useEffect(() => {
+    if (!isInViewport) return
+
+    setTimeout(highlightCode, 0)
+  }, [isInViewport, highlightCode])
 
   useEffect(() => {
-    const container = codeContentRef.current
+    const container = codeContainerRef.current
     if (!container || !codeShowLineNumbers) return
 
     const digits = Math.max(tokenLines.length.toString().length, 1)
@@ -180,7 +166,7 @@ const CodePreview = ({ children, language, setTools }: CodePreviewProps) => {
 
   return (
     <ContentContainer
-      ref={codeContentRef}
+      ref={codeContainerRef}
       $lineNumbers={codeShowLineNumbers}
       $wrap={codeWrappable && !isUnwrapped}
       $fadeIn={hasHighlightedCode}

src/renderer/src/context/CodeStyleProvider.tsx

@@ -10,6 +10,7 @@ import { createContext, type PropsWithChildren, use, useCallback, useEffect, use
 
 interface CodeStyleContextType {
   highlightCodeChunk: (trunk: string, language: string, callerId: string) => Promise<HighlightChunkResult>
+  highlightStreamingCode: (code: string, language: string, callerId: string) => Promise<HighlightChunkResult>
   cleanupTokenizers: (callerId: string) => void
   getShikiPreProperties: (language: string) => Promise<ShikiPreProperties>
   highlightCode: (code: string, language: string) => Promise<string>
@@ -22,6 +23,7 @@ interface CodeStyleContextType {
 
 const defaultCodeStyleContext: CodeStyleContextType = {
   highlightCodeChunk: async () => ({ lines: [], recall: 0 }),
+  highlightStreamingCode: async () => ({ lines: [], recall: 0 }),
   cleanupTokenizers: () => {},
   getShikiPreProperties: async () => ({ class: '', style: '', tabindex: 0 }),
   highlightCode: async () => '',
@@ -114,6 +116,15 @@ export const CodeStyleProvider: React.FC<PropsWithChildren> = ({ children }) =>
     shikiStreamService.cleanupTokenizers(callerId)
   }, [])
 
+  // 高亮流式输出的代码
+  const highlightStreamingCode = useCallback(
+    async (fullContent: string, language: string, callerId: string) => {
+      const normalizedLang = languageMap[language as keyof typeof languageMap] || language.toLowerCase()
+      return shikiStreamService.highlightStreamingCode(fullContent, normalizedLang, activeShikiTheme, callerId)
+    },
+    [activeShikiTheme, languageMap]
+  )
+
   // 获取 Shiki pre 标签属性
   const getShikiPreProperties = useCallback(
     async (language: string) => {
@@ -148,6 +159,7 @@ export const CodeStyleProvider: React.FC<PropsWithChildren> = ({ children }) =>
   const contextValue = useMemo(
     () => ({
       highlightCodeChunk,
+      highlightStreamingCode,
       cleanupTokenizers,
       getShikiPreProperties,
       highlightCode,
@@ -159,6 +171,7 @@ export const CodeStyleProvider: React.FC<PropsWithChildren> = ({ children }) =>
     }),
     [
       highlightCodeChunk,
+      highlightStreamingCode,
       cleanupTokenizers,
       getShikiPreProperties,
       highlightCode,

src/renderer/src/services/ShikiStreamService.ts

@@ -20,7 +20,7 @@ export type ShikiPreProperties = {
  * 代码 chunk 高亮结果
  *
  * @param lines 所有高亮行（包括稳定和不稳定）
- * @param recall 需要撤回的行数
+ * @param recall 需要撤回的行数，-1 表示撤回所有行
  */
 export interface HighlightChunkResult {
   lines: ThemedToken[][]
@@ -47,6 +47,13 @@ class ShikiStreamService {
     }
   })
 
+  // 缓存每个 callerId 对应的已处理内容
+  private codeCache = new LRUCache<string, string>({
+    max: 100, // 最大缓存数量
+    ttl: 1000 * 60 * 30, // 30分钟过期时间
+    updateAgeOnGet: true
+  })
+
   // Worker 相关资源
   private worker: Worker | null = null
   private workerInitPromise: Promise<void> | null = null
@@ -261,6 +268,72 @@ class ShikiStreamService {
     return hast.children[0].properties as ShikiPreProperties
   }
 
+  /**
+   * 高亮流式输出的代码，调用方传入完整代码内容，得到增量高亮结果。
+   *
+   * - 检测当前内容与上次处理内容的差异。
+   * - 如果是末尾追加，只传输增量部分（此时性能最好，如遇性能问题，考虑检查这里的逻辑）。
+   * - 如果不是追加，重置 tokenizer 并处理完整内容。
+   *
+   * 调用者需要自行处理撤回。
+   * @param code 完整代码内容
+   * @param language 语言
+   * @param theme 主题
+   * @param callerId 调用者ID
+   * @returns 高亮结果，recall 为 -1 表示撤回所有行
+   */
+  async highlightStreamingCode(
+    code: string,
+    language: string,
+    theme: string,
+    callerId: string
+  ): Promise<HighlightChunkResult> {
+    const cacheKey = `${callerId}-${language}-${theme}`
+    const lastContent = this.codeCache.get(cacheKey) || ''
+
+    let isAppend = false
+
+    if (code.length === lastContent.length) {
+      // 内容没有变化，返回空结果
+      if (code === lastContent) {
+        return { lines: [], recall: 0 }
+      }
+    } else if (code.length > lastContent.length) {
+      // 长度增加，可能是追加
+      isAppend = code.startsWith(lastContent)
+    }
+
+    try {
+      let result: HighlightChunkResult
+
+      if (isAppend) {
+        // 流式追加，只传输增量
+        const chunk = code.slice(lastContent.length)
+        result = await this.highlightCodeChunk(chunk, language, theme, callerId)
+      } else {
+        // 非追加变化，重置并处理完整内容
+        this.cleanupTokenizers(callerId)
+        this.codeCache.delete(cacheKey) // 清除缓存
+
+        result = await this.highlightCodeChunk(code, language, theme, callerId)
+
+        // 撤回所有行
+        result = {
+          ...result,
+          recall: -1
+        }
+      }
+
+      // 成功处理后更新缓存
+      this.codeCache.set(cacheKey, code)
+      return result
+    } catch (error) {
+      // 处理失败时不更新缓存，保持之前的状态
+      console.error('Failed to highlight streaming code:', error)
+      throw error
+    }
+  }
+
   /**
    * 高亮代码 chunk，返回本次高亮的所有 ThemedToken 行
    *
@@ -405,6 +478,13 @@ class ShikiStreamService {
       })
     }
 
+    // 清理对应的内容缓存
+    for (const key of this.codeCache.keys()) {
+      if (key.startsWith(`${callerId}-`)) {
+        this.codeCache.delete(key)
+      }
+    }
+
     // 再清理主线程中的 tokenizers，移除所有以 callerId 开头的缓存项
     for (const key of this.tokenizerCache.keys()) {
       if (key.startsWith(`${callerId}-`)) {
@@ -429,6 +509,7 @@ class ShikiStreamService {
 
     this.workerDegradationCache.clear()
     this.tokenizerCache.clear()
+    this.codeCache.clear()
     this.highlighter = null
     this.workerInitPromise = null
     this.workerInitRetryCount = 0