Create highlighter instance to improve performance

src/components/ChatUI/ShikiHihlighter.tsx

@@ -17,20 +17,20 @@ export const ShikiHighlighter: React.FC<ShikiHighlighterProps> = ({
 }) => {
   const highlight = useShikiHighlighter(theme);
   const [highlightedCode, setHighlightedCode] =
-    useState<React.ReactNode | null>(null);
+    useState<React.ReactNode>(children);
 
   useEffect(() => {
-    highlight(children, language).then(setHighlightedCode);
+    let isMounted = true;
+    highlight(children, language).then((result) => {
+      if (isMounted) {
+        setHighlightedCode(result);
+      }
+    });
+    return () => {
+      isMounted = false;
+    };
   }, [highlight, children, language]);
 
-  if (!highlightedCode) {
-    return React.createElement(
-      PreTag,
-      null,
-      React.createElement('code', null, children)
-    );
-  }
-
   return React.createElement(PreTag, {
     className: 'shiki not-prose',
     children: highlightedCode,

src/hooks/useShikiHighlighter.ts

@@ -1,16 +1,55 @@
-import { useCallback } from 'react';
-import { codeToHtml, type BundledLanguage, type BundledTheme } from 'shiki';
+import { useEffect, useState, useCallback } from 'react';
+import {
+  createHighlighter,
+  type BundledLanguage,
+  type BundledTheme,
+  type Highlighter,
+} from 'shiki';
 import parse from 'html-react-parser';
 
+let highlighterInstance: Highlighter | null = null;
+let highlighterPromise: Promise<Highlighter> | null = null;
+
+const getHighlighterInstance = async (
+  theme: BundledTheme
+): Promise<Highlighter> => {
+  if (!highlighterPromise) {
+    highlighterPromise = createHighlighter({
+      themes: [theme],
+      langs: [],
+    });
+  }
+  highlighterInstance = await highlighterPromise;
+  return highlighterInstance;
+};
+
 export const useShikiHighlighter = (theme: BundledTheme) => {
+  const [highlighter, setHighlighter] = useState<Highlighter | null>(null);
+
+  useEffect(() => {
+    getHighlighterInstance(theme).then(setHighlighter);
+  }, [theme]);
+
   return useCallback(
-    async (code: string, language: BundledLanguage) => {
-      return codeToHtml(code, {
-        lang: language,
-        theme,
-      }).then((html) => parse(html));
+    (code: string, language: BundledLanguage): Promise<React.ReactNode> => {
+      if (!highlighter) {
+        return Promise.resolve(code); // Fallback to plain text if highlighter isn't ready
+      }
+
+      return (async () => {
+        try {
+          if (!highlighter.getLoadedLanguages().includes(language)) {
+            await highlighter.loadLanguage(language);
+          }
+          const html = highlighter.codeToHtml(code, { lang: language, theme });
+          return parse(html);
+        } catch (error) {
+          console.error('Failed to highlight code:', error);
+          return code; // Fallback to plain text on error
+        }
+      })();
     },
-    []
+    [highlighter, theme]
   );
 };