From 9f51f9f0345362bfe9d63ec76878b0053a715305 Mon Sep 17 00:00:00 2001 From: Bentlybro Date: Fri, 13 Feb 2026 09:14:48 +0000 Subject: [PATCH] fix(frontend): use singleton shiki highlighter for code blocks Fixes SENTRY-1051: shiki warning about too many instances. @streamdown/code creates a new highlighter per language, causing memory bloat and warnings when chats have multiple code blocks. This adds a custom code plugin that: - Uses a single shiki highlighter instance - Dynamically loads languages on demand - Pre-loads common languages (js, ts, python, json, etc.) - Caches results like the original plugin Drop-in replacement - same interface as @streamdown/code. --- .../src/components/ai-elements/message.tsx | 2 +- .../src/lib/streamdown-code-singleton.ts | 230 ++++++++++++++++++ 2 files changed, 231 insertions(+), 1 deletion(-) create mode 100644 autogpt_platform/frontend/src/lib/streamdown-code-singleton.ts diff --git a/autogpt_platform/frontend/src/components/ai-elements/message.tsx b/autogpt_platform/frontend/src/components/ai-elements/message.tsx index 5cc330e57c..d83266dc0c 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-singleton"; import { math } from "@streamdown/math"; import { mermaid } from "@streamdown/mermaid"; import type { UIMessage } from "ai"; diff --git a/autogpt_platform/frontend/src/lib/streamdown-code-singleton.ts b/autogpt_platform/frontend/src/lib/streamdown-code-singleton.ts new file mode 100644 index 0000000000..14c4fa417d --- /dev/null +++ b/autogpt_platform/frontend/src/lib/streamdown-code-singleton.ts @@ -0,0 +1,230 @@ +/** + * Custom Streamdown code plugin with proper shiki singleton. + * + * Fixes SENTRY-1051: "@streamdown/code creates a new shiki highlighter per language, + * causing "10 instances created" warnings and memory bloat. + * + * This plugin creates ONE highlighter and loads languages dynamically. + */ + +import { + createHighlighter, + bundledLanguages, + type Highlighter, + type BundledLanguage, + type BundledTheme, +} from "shiki"; + +// Types matching streamdown's expected interface +interface HighlightToken { + content: string; + color?: string; + bgColor?: string; + htmlStyle?: Record; + htmlAttrs?: Record; + offset?: number; +} + +interface HighlightResult { + tokens: HighlightToken[][]; + fg?: string; + bg?: string; +} + +interface HighlightOptions { + code: string; + language: BundledLanguage; + themes: [string, string]; +} + +interface CodeHighlighterPlugin { + name: "shiki"; + type: "code-highlighter"; + highlight: ( + options: HighlightOptions, + callback?: (result: HighlightResult) => void + ) => HighlightResult | null; + supportsLanguage: (language: BundledLanguage) => boolean; + getSupportedLanguages: () => BundledLanguage[]; + getThemes: () => [BundledTheme, BundledTheme]; +} + +// Singleton state +let highlighterPromise: Promise | null = null; +let highlighterInstance: Highlighter | null = null; +const loadedLanguages = new Set(); +const pendingLanguages = new Map>(); + +// Result cache (same as @streamdown/code) +const resultCache = new Map(); +const pendingCallbacks = new Map void>>(); + +// All supported languages +const supportedLanguages = new Set(Object.keys(bundledLanguages)); + +// Cache key for results +function getCacheKey(code: string, language: string, themes: [string, string]): string { + const prefix = code.slice(0, 100); + const suffix = code.length > 100 ? code.slice(-100) : ""; + return `${language}:${themes[0]}:${themes[1]}:${code.length}:${prefix}:${suffix}`; +} + +// Get or create the singleton highlighter +async function getHighlighter(themes: [string, string]): Promise { + if (highlighterInstance) { + return highlighterInstance; + } + + if (!highlighterPromise) { + highlighterPromise = createHighlighter({ + themes: themes as [BundledTheme, BundledTheme], + // Start with common languages pre-loaded for faster first render + langs: ["javascript", "typescript", "python", "json", "html", "css", "bash", "markdown"], + }).then((h) => { + highlighterInstance = h; + ["javascript", "typescript", "python", "json", "html", "css", "bash", "markdown"].forEach( + (l) => loadedLanguages.add(l) + ); + return h; + }); + } + + return highlighterPromise; +} + +// Load a language dynamically +async function ensureLanguageLoaded( + highlighter: Highlighter, + language: string +): Promise { + if (loadedLanguages.has(language)) { + return; + } + + if (pendingLanguages.has(language)) { + return pendingLanguages.get(language); + } + + const loadPromise = highlighter + .loadLanguage(language as BundledLanguage) + .then(() => { + loadedLanguages.add(language); + pendingLanguages.delete(language); + }) + .catch((err) => { + console.warn(`[streamdown-code-singleton] Failed to load language: ${language}`, err); + pendingLanguages.delete(language); + }); + + pendingLanguages.set(language, loadPromise); + return loadPromise; +} + +// Convert shiki tokens to streamdown format +function convertTokens( + shikiResult: ReturnType +): HighlightResult { + return { + tokens: shikiResult.tokens.map((line) => + line.map((token) => ({ + content: token.content, + color: token.color, + htmlStyle: token.htmlStyle, + })) + ), + fg: shikiResult.fg, + bg: shikiResult.bg, + }; +} + +export interface CodePluginOptions { + themes?: [BundledTheme, BundledTheme]; +} + +export function createCodePlugin( + options: CodePluginOptions = {} +): CodeHighlighterPlugin { + const themes = options.themes ?? ["github-light", "github-dark"]; + + return { + name: "shiki", + type: "code-highlighter", + + supportsLanguage(language: BundledLanguage): boolean { + return supportedLanguages.has(language); + }, + + getSupportedLanguages(): BundledLanguage[] { + return Array.from(supportedLanguages) as BundledLanguage[]; + }, + + getThemes(): [BundledTheme, BundledTheme] { + return themes as [BundledTheme, BundledTheme]; + }, + + highlight( + { code, language, themes: highlightThemes }: HighlightOptions, + callback?: (result: HighlightResult) => void + ): HighlightResult | null { + const cacheKey = getCacheKey(code, language, highlightThemes); + + // Return cached result if available + if (resultCache.has(cacheKey)) { + return resultCache.get(cacheKey)!; + } + + // Register callback for async result + if (callback) { + if (!pendingCallbacks.has(cacheKey)) { + pendingCallbacks.set(cacheKey, new Set()); + } + pendingCallbacks.get(cacheKey)!.add(callback); + } + + // Start async highlighting + getHighlighter(highlightThemes) + .then(async (highlighter) => { + // Ensure language is loaded + const lang = supportedLanguages.has(language) ? language : "text"; + if (lang !== "text") { + await ensureLanguageLoaded(highlighter, lang); + } + + // Highlight the code + const effectiveLang = highlighter.getLoadedLanguages().includes(lang) + ? lang + : "text"; + + const shikiResult = highlighter.codeToTokens(code, { + lang: effectiveLang, + themes: { + light: highlightThemes[0] as BundledTheme, + dark: highlightThemes[1] as BundledTheme, + }, + }); + + const result = convertTokens(shikiResult); + resultCache.set(cacheKey, result); + + // Notify all pending callbacks + const callbacks = pendingCallbacks.get(cacheKey); + if (callbacks) { + for (const cb of callbacks) { + cb(result); + } + pendingCallbacks.delete(cacheKey); + } + }) + .catch((err) => { + console.error("[streamdown-code-singleton] Failed to highlight code:", err); + pendingCallbacks.delete(cacheKey); + }); + + // Return null while async loading + return null; + }, + }; +} + +// Pre-configured plugin with default settings (drop-in replacement for @streamdown/code) +export const code = createCodePlugin();