From 905373a712b7e88917016388a299de588e4db2a3 Mon Sep 17 00:00:00 2001 From: Bently Date: Tue, 17 Feb 2026 12:15:53 +0000 Subject: [PATCH] fix(frontend): use singleton Shiki highlighter for code syntax highlighting (#12144) ## Summary Addresses SENTRY-1051: Shiki warning about multiple highlighter instances. ## Problem The `@streamdown/code` package creates a **new Shiki highlighter for each language** encountered. When users view AI chat responses with code blocks in multiple languages (JavaScript, Python, JSON, YAML, etc.), this creates 10+ highlighter instances, triggering Shiki's warning: > "10 instances have been created. Shiki is supposed to be used as a singleton, consider refactoring your code to cache your highlighter instance" This causes memory bloat and performance degradation. ## Solution Introduced a custom code highlighting plugin that properly implements the singleton pattern: ### New files: - `src/lib/shiki-highlighter.ts` - Singleton highlighter management - `src/lib/streamdown-code-plugin.ts` - Drop-in replacement for `@streamdown/code` ### Key features: - **Single shared highlighter** - One instance serves all code blocks - **Preloaded common languages** - JS, TS, Python, JSON, Bash, YAML, etc. - **Lazy loading** - Additional languages loaded on demand - **Result caching** - Avoids re-highlighting identical code blocks ### Changes: - Added `shiki` as direct dependency - Updated `message.tsx` to use the new plugin ## Testing - [ ] Verify code blocks render correctly in AI chat - [ ] Confirm no Shiki singleton warnings in console - [ ] Test with multiple languages in same conversation ## Related - Linear: SENTRY-1051 - Sentry: Multiple Shiki instances warning

Greptile Summary

 Replaced `@streamdown/code` with a custom singleton-based Shiki highlighter implementation to resolve memory bloat from creating multiple highlighter instances per language. The new implementation creates a single shared highlighter with preloaded common languages (JS, TS, Python, JSON, etc.) and lazy-loads additional languages on demand. Results are cached to avoid re-highlighting identical code blocks. **Key changes:** - Added `shiki` v3.21.0 as a direct dependency - Created `shiki-highlighter.ts` with singleton pattern and language management utilities - Created `streamdown-code-plugin.ts` as a drop-in replacement for `@streamdown/code` - Updated `message.tsx` to import from the new plugin instead of `@streamdown/code` The implementation follows React best practices with async highlighting and callback-based notifications. The cache key uses code length + prefix/suffix for efficient lookups on large code blocks.

Confidence Score: 4/5

 - Safe to merge with minor considerations for edge cases - The implementation is solid with proper singleton pattern, caching, and async handling. The code is well-structured and addresses the stated problem. However, there's a subtle potential race condition in the callback handling where multiple concurrent requests for the same cache key could trigger duplicate highlight operations before the first completes. The cache key generation using prefix/suffix could theoretically cause false cache hits for large files with identical prefixes and suffixes. Despite these edge cases, the implementation should work correctly for the vast majority of use cases. - No files require special attention

Sequence Diagram

 ```mermaid sequenceDiagram participant UI as Streamdown Component participant Plugin as Custom Code Plugin participant Cache as Token Cache participant Singleton as Shiki Highlighter (Singleton) participant Callbacks as Pending Callbacks UI->>Plugin: highlight(code, lang) Plugin->>Cache: Check cache key alt Cache hit Cache-->>Plugin: Return cached result Plugin-->>UI: Return highlighted tokens else Cache miss Plugin->>Callbacks: Register callback Plugin->>Singleton: Get highlighter instance alt First call Singleton->>Singleton: Create highlighter with preloaded languages end Singleton-->>Plugin: Return highlighter alt Language not loaded Plugin->>Singleton: Load language dynamically end Plugin->>Singleton: codeToTokens(code, lang, themes) Singleton-->>Plugin: Return tokens Plugin->>Cache: Store result Plugin->>Callbacks: Notify all waiting callbacks Callbacks-->>UI: Async callback with result end ```
Last reviewed commit: 96c793b --- autogpt_platform/frontend/package.json | 2 +- autogpt_platform/frontend/pnpm-lock.yaml | 16 +- .../src/components/ai-elements/message.tsx | 2 +- .../frontend/src/lib/shiki-highlighter.ts | 70 ++++++++ .../src/lib/streamdown-code-plugin.ts | 159 ++++++++++++++++++ 5 files changed, 234 insertions(+), 15 deletions(-) create mode 100644 autogpt_platform/frontend/src/lib/shiki-highlighter.ts create mode 100644 autogpt_platform/frontend/src/lib/streamdown-code-plugin.ts diff --git a/autogpt_platform/frontend/package.json b/autogpt_platform/frontend/package.json index 5988e59c90..80fb46beff 100644 --- a/autogpt_platform/frontend/package.json +++ b/autogpt_platform/frontend/package.json @@ -62,7 +62,6 @@ "@rjsf/validator-ajv8": "6.1.2", "@sentry/nextjs": "10.27.0", "@streamdown/cjk": "1.0.1", - "@streamdown/code": "1.0.1", "@streamdown/math": "1.0.1", "@streamdown/mermaid": "1.0.1", "@supabase/ssr": "0.7.0", @@ -116,6 +115,7 @@ "remark-gfm": "4.0.1", "remark-math": "6.0.0", "shepherd.js": "14.5.1", + "shiki": "^3.21.0", "sonner": "2.0.7", "streamdown": "2.1.0", "tailwind-merge": "2.6.0", diff --git a/autogpt_platform/frontend/pnpm-lock.yaml b/autogpt_platform/frontend/pnpm-lock.yaml index 468e2f312d..4d76cd0cb9 100644 --- a/autogpt_platform/frontend/pnpm-lock.yaml +++ b/autogpt_platform/frontend/pnpm-lock.yaml @@ -108,9 +108,6 @@ importers: '@streamdown/cjk': specifier: 1.0.1 version: 1.0.1(@types/mdast@4.0.4)(micromark-util-types@2.0.2)(micromark@4.0.2)(react@18.3.1)(unified@11.0.5) - '@streamdown/code': - specifier: 1.0.1 - version: 1.0.1(react@18.3.1) '@streamdown/math': specifier: 1.0.1 version: 1.0.1(react@18.3.1) @@ -270,6 +267,9 @@ importers: shepherd.js: specifier: 14.5.1 version: 14.5.1 + shiki: + specifier: ^3.21.0 + version: 3.21.0 sonner: specifier: 2.0.7 version: 2.0.7(react-dom@18.3.1(react@18.3.1))(react@18.3.1) @@ -3307,11 +3307,6 @@ packages: peerDependencies: react: ^18.0.0 || ^19.0.0 - '@streamdown/code@1.0.1': - resolution: {integrity: sha512-U9LITfQ28tZYAoY922jdtw1ryg4kgRBdURopqK9hph7G2fBUwPeHthjH7SvaV0fvFv7EqjqCzARJuWUljLe9Ag==} - peerDependencies: - react: ^18.0.0 || ^19.0.0 - '@streamdown/math@1.0.1': resolution: {integrity: sha512-R9WdHbpERiRU7WeO7oT1aIbnLJ/jraDr89F7X9x2OM//Y8G8UMATRnLD/RUwg4VLr8Nu7QSIJ0Pa8lXd2meM4Q==} peerDependencies: @@ -11907,11 +11902,6 @@ snapshots: - micromark-util-types - unified - '@streamdown/code@1.0.1(react@18.3.1)': - dependencies: - react: 18.3.1 - shiki: 3.21.0 - '@streamdown/math@1.0.1(react@18.3.1)': dependencies: katex: 0.16.28 diff --git a/autogpt_platform/frontend/src/components/ai-elements/message.tsx b/autogpt_platform/frontend/src/components/ai-elements/message.tsx index 5cc330e57c..cc11c2f01b 100644 --- a/autogpt_platform/frontend/src/components/ai-elements/message.tsx +++ b/autogpt_platform/frontend/src/components/ai-elements/message.tsx @@ -10,7 +10,7 @@ import { } from "@/components/ui/tooltip"; import { cn } from "@/lib/utils"; import { cjk } from "@streamdown/cjk"; -import { code } from "@streamdown/code"; +import { code } from "@/lib/streamdown-code-plugin"; import { math } from "@streamdown/math"; import { mermaid } from "@streamdown/mermaid"; import type { UIMessage } from "ai"; diff --git a/autogpt_platform/frontend/src/lib/shiki-highlighter.ts b/autogpt_platform/frontend/src/lib/shiki-highlighter.ts new file mode 100644 index 0000000000..f6bbcb9b94 --- /dev/null +++ b/autogpt_platform/frontend/src/lib/shiki-highlighter.ts @@ -0,0 +1,70 @@ +import { + bundledLanguages, + bundledLanguagesInfo, + createHighlighter, + type BundledLanguage, + type BundledTheme, + type HighlighterGeneric, +} from "shiki"; + +export type { BundledLanguage, BundledTheme }; + +const LANGUAGE_ALIASES: Record = Object.fromEntries( + bundledLanguagesInfo.flatMap((lang) => + (lang.aliases ?? []).map((alias) => [alias, lang.id]), + ), +); + +const SUPPORTED_LANGUAGES = new Set(Object.keys(bundledLanguages)); + +const PRELOAD_LANGUAGES: BundledLanguage[] = [ + "javascript", + "typescript", + "python", + "json", + "bash", + "yaml", + "markdown", + "html", + "css", + "sql", + "tsx", + "jsx", +]; + +export const SHIKI_THEMES: [BundledTheme, BundledTheme] = [ + "github-light", + "github-dark", +]; + +let highlighterPromise: Promise< + HighlighterGeneric +> | null = null; + +export function getShikiHighlighter(): Promise< + HighlighterGeneric +> { + if (!highlighterPromise) { + highlighterPromise = createHighlighter({ + themes: SHIKI_THEMES, + langs: PRELOAD_LANGUAGES, + }).catch((err) => { + highlighterPromise = null; + throw err; + }); + } + return highlighterPromise; +} + +export function resolveLanguage(lang: string): string { + const normalized = lang.trim().toLowerCase(); + return LANGUAGE_ALIASES[normalized] ?? normalized; +} + +export function isLanguageSupported(lang: string): boolean { + return SUPPORTED_LANGUAGES.has(resolveLanguage(lang)); +} + +export function getSupportedLanguages(): BundledLanguage[] { + return Array.from(SUPPORTED_LANGUAGES) as BundledLanguage[]; +} diff --git a/autogpt_platform/frontend/src/lib/streamdown-code-plugin.ts b/autogpt_platform/frontend/src/lib/streamdown-code-plugin.ts new file mode 100644 index 0000000000..52ccc5afed --- /dev/null +++ b/autogpt_platform/frontend/src/lib/streamdown-code-plugin.ts @@ -0,0 +1,159 @@ +import type { CodeHighlighterPlugin } from "streamdown"; + +import { + type BundledLanguage, + type BundledTheme, + getShikiHighlighter, + getSupportedLanguages, + isLanguageSupported, + resolveLanguage, + SHIKI_THEMES, +} from "./shiki-highlighter"; + +interface HighlightResult { + tokens: { + content: string; + color?: string; + htmlStyle?: Record; + }[][]; + fg?: string; + bg?: string; +} + +type HighlightCallback = (result: HighlightResult) => void; + +const MAX_CACHE_SIZE = 500; +const tokenCache = new Map(); +const pendingCallbacks = new Map>(); +const inFlightLanguageLoads = new Map>(); + +function simpleHash(str: string): string { + let hash = 0; + for (let i = 0; i < str.length; i++) { + const char = str.charCodeAt(i); + hash = (hash << 5) - hash + char; + hash = hash & hash; + } + return hash.toString(36); +} + +function getCacheKey( + code: string, + lang: string, + themes: readonly string[], +): string { + return `${lang}:${themes.join(",")}:${simpleHash(code)}`; +} + +function evictOldestIfNeeded(): void { + if (tokenCache.size > MAX_CACHE_SIZE) { + const oldestKey = tokenCache.keys().next().value; + if (oldestKey) { + tokenCache.delete(oldestKey); + } + } +} + +export function createSingletonCodePlugin(): CodeHighlighterPlugin { + return { + name: "shiki", + type: "code-highlighter", + + supportsLanguage(lang: BundledLanguage): boolean { + return isLanguageSupported(lang); + }, + + getSupportedLanguages(): BundledLanguage[] { + return getSupportedLanguages(); + }, + + getThemes(): [BundledTheme, BundledTheme] { + return SHIKI_THEMES; + }, + + highlight({ code, language, themes }, callback) { + const lang = resolveLanguage(language); + const cacheKey = getCacheKey(code, lang, themes); + + if (tokenCache.has(cacheKey)) { + return tokenCache.get(cacheKey)!; + } + + if (callback) { + if (!pendingCallbacks.has(cacheKey)) { + pendingCallbacks.set(cacheKey, new Set()); + } + pendingCallbacks.get(cacheKey)!.add(callback); + } + + getShikiHighlighter() + .then(async (highlighter) => { + const loadedLanguages = highlighter.getLoadedLanguages(); + + if (!loadedLanguages.includes(lang) && isLanguageSupported(lang)) { + let loadPromise = inFlightLanguageLoads.get(lang); + if (!loadPromise) { + loadPromise = highlighter + .loadLanguage(lang as BundledLanguage) + .finally(() => { + inFlightLanguageLoads.delete(lang); + }); + inFlightLanguageLoads.set(lang, loadPromise); + } + await loadPromise; + } + + const finalLang = ( + highlighter.getLoadedLanguages().includes(lang) ? lang : "text" + ) as BundledLanguage; + + const shikiResult = highlighter.codeToTokens(code, { + lang: finalLang, + themes: { light: themes[0], dark: themes[1] }, + }); + + const result: HighlightResult = { + tokens: shikiResult.tokens.map((line) => + line.map((token) => ({ + content: token.content, + color: token.color, + htmlStyle: token.htmlStyle, + })), + ), + fg: shikiResult.fg, + bg: shikiResult.bg, + }; + + evictOldestIfNeeded(); + tokenCache.set(cacheKey, result); + + const callbacks = pendingCallbacks.get(cacheKey); + if (callbacks) { + callbacks.forEach((cb) => { + cb(result); + }); + pendingCallbacks.delete(cacheKey); + } + }) + .catch((error) => { + console.error("[Shiki] Failed to highlight code:", error); + + const fallback: HighlightResult = { + tokens: code.split("\n").map((line) => [{ content: line }]), + }; + + const callbacks = pendingCallbacks.get(cacheKey); + if (callbacks) { + callbacks.forEach((cb) => { + cb(fallback); + }); + pendingCallbacks.delete(cacheKey); + } + }); + + return null; + }, + }; +} + +export const code = createSingletonCodePlugin();