From 5811adfb7f324b5cf039af35fc6962b0f956d1ae Mon Sep 17 00:00:00 2001 From: one Date: Thu, 26 Jun 2025 13:30:49 +0800 Subject: [PATCH] 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 --- .../components/CodeBlockView/CodePreview.tsx | 118 ++++++++---------- .../src/context/CodeStyleProvider.tsx | 13 ++ .../src/services/ShikiStreamService.ts | 83 +++++++++++- 3 files changed, 147 insertions(+), 67 deletions(-) diff --git a/src/renderer/src/components/CodeBlockView/CodePreview.tsx b/src/renderer/src/components/CodeBlockView/CodePreview.tsx index dde163283d..b550cd2467 100644 --- a/src/renderer/src/components/CodeBlockView/CodePreview.tsx +++ b/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([]) - const codeContentRef = useRef(null) - const prevCodeLengthRef = useRef(0) - const safeCodeStringRef = useRef(children) - const highlightQueueRef = useRef>(Promise.resolve()) + const [isInViewport, setIsInViewport] = useState(false) + const codeContainerRef = useRef(null) + const processingRef = useRef(false) + const latestRequestedContentRef = useRef(null) const callerId = useRef(`${Date.now()}-${uuid()}`).current const shikiThemeRef = useRef(activeShikiTheme) @@ -45,7 +46,7 @@ const CodePreview = ({ children, language, setTools }: CodePreviewProps) => { icon: isExpanded ? : , 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 ( Promise + highlightStreamingCode: (code: string, language: string, callerId: string) => Promise cleanupTokenizers: (callerId: string) => void getShikiPreProperties: (language: string) => Promise highlightCode: (code: string, language: string) => Promise @@ -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 = ({ 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 = ({ children }) => const contextValue = useMemo( () => ({ highlightCodeChunk, + highlightStreamingCode, cleanupTokenizers, getShikiPreProperties, highlightCode, @@ -159,6 +171,7 @@ export const CodeStyleProvider: React.FC = ({ children }) => }), [ highlightCodeChunk, + highlightStreamingCode, cleanupTokenizers, getShikiPreProperties, highlightCode, diff --git a/src/renderer/src/services/ShikiStreamService.ts b/src/renderer/src/services/ShikiStreamService.ts index 8ddc2a8be1..8d5c6f6223 100644 --- a/src/renderer/src/services/ShikiStreamService.ts +++ b/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({ + max: 100, // 最大缓存数量 + ttl: 1000 * 60 * 30, // 30分钟过期时间 + updateAgeOnGet: true + }) + // Worker 相关资源 private worker: Worker | null = null private workerInitPromise: Promise | 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 { + 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