Experimenting with custom themes in Shiki, first working version

src/components/ChatUI/CodeHighlight.tsx

@@ -24,7 +24,6 @@ export const CodeHighlight = ({
   return !isInline ? (
     <ShikiHighlighter
       language={language as BundledLanguage}
-      theme={'houston'}
       {...props}>
       {String(children)}
     </ShikiHighlighter>

src/components/ChatUI/ShikiHihlighter.tsx

@@ -5,7 +5,7 @@ import type { BundledTheme, BundledLanguage } from 'shiki';
 interface ShikiHighlighterProps {
   language: BundledLanguage;
   children: string;
-  theme: BundledTheme;
+  theme?: BundledTheme;
   as?: React.ElementType;
 }
 
@@ -15,9 +15,9 @@ export const ShikiHighlighter = ({
   children: code,
   as: Element = 'pre',
 }: ShikiHighlighterProps) => {
-  const highlightedCode = useShikiHighlighter(code, language, theme);
+  const highlightedCode = useShikiHighlighter(code, language, theme as BundledTheme);
 
-  return (
+  return highlightedCode && (
     <Element className="shiki not-prose relative [&_pre]:overflow-auto [&_pre]:rounded-lg [&_pre]:px-6 [&_pre]:py-5">
       {language ? (
         <span className="absolute right-3 top-2 text-xs tracking-tighter text-muted-foreground/85">

src/hooks/useShiki.ts

@@ -1,7 +1,13 @@
 import { useState, useEffect, type ReactNode } from 'react';
-import { codeToHtml, type BundledLanguage, type BundledTheme } from 'shiki';
+import {
+  createHighlighter,
+  type BundledLanguage,
+  type BundledTheme,
+  type Highlighter,
+} from 'shiki';
 import parse from 'html-react-parser';
 import { removeTabIndexFromPre } from '@utils/shikiTransformers';
+import customTheme from '@styles/rainglow-azure-constrast.mjs';
 
 export const useShikiHighlighter = (
   code: string,
@@ -11,17 +17,36 @@ export const useShikiHighlighter = (
   const [highlightedCode, setHighlightedCode] = useState<ReactNode | null>(
     null
   );
+  const [highlighter, setHighlighter] = useState<Highlighter | null>(null);
+  const [isReady, setIsReady] = useState(false);
 
   useEffect(() => {
-    codeToHtml(code,
-      {
-        lang: lang as BundledLanguage,
-        theme,
-        transformers: [removeTabIndexFromPre]
-      }).then((html) => {
+    const initializeHighlighter = async () => {
+      if (!highlighter) {
+        const newHighlighter = await createHighlighter({
+          langs: [lang as BundledLanguage],
+          themes: [customTheme as unknown as BundledTheme],
+        });
+        setHighlighter(newHighlighter);
+        setIsReady(true);
+      }
+    };
+
+    initializeHighlighter();
+  }, [lang]);
+
+  useEffect(() => {
+    if (isReady && highlighter && lang) {
+      console.time('shiki');
+      const html = highlighter.codeToHtml(code, {
+        lang: lang,
+        theme: customTheme.name as BundledTheme,
+        transformers: [removeTabIndexFromPre],
+      });
       setHighlightedCode(parse(html));
-    });
-  }, [code]);
+    }
+    console.timeEnd('shiki');
+  }, [code, lang, highlighter, isReady]);
 
-  return highlightedCode;
-};
+  return isReady ? highlightedCode : null;
+};