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/lib/hooks/useShiki.ts

@@ -1,11 +1,13 @@
 import { useState, useEffect, type ReactNode } from 'react';
 import {
-  codeToHtml,
+  createHighlighter,
   type BundledLanguage,
   type BundledTheme,
+  type Highlighter,
 } from 'shiki';
 import parse from 'html-react-parser';
 import { removeTabIndexFromPre } from '@/lib/utils';
+import customTheme from '@styles/rainglow-azure-constrast.mjs';
 
 export const useShikiHighlighter = (
   code: string,
@@ -15,28 +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(() => {
-    const highlightCode = async () => {
-      try {
-        const html = await codeToHtml(code, {
-          lang: lang as BundledLanguage,
-          theme,
-          transformers: [removeTabIndexFromPre],
+    const initializeHighlighter = async () => {
+      if (!highlighter) {
+        const newHighlighter = await createHighlighter({
+          langs: [lang as BundledLanguage],
+          themes: [customTheme as unknown as BundledTheme],
         });
-        setHighlightedCode(parse(html));
-      } catch (error) {
-        const fallbackHtml = await codeToHtml(code, {
-          lang: 'plaintext',
-          theme,
-          transformers: [removeTabIndexFromPre],
-        });
-        setHighlightedCode(parse(fallbackHtml));
+        setHighlighter(newHighlighter);
+        setIsReady(true);
       }
     };
 
-    highlightCode();
-  }, [code]);
+    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));
+    }
+    console.timeEnd('shiki');
+  }, [code, lang, highlighter, isReady]);
 
-  return highlightedCode;
-};
+  return isReady ? highlightedCode : null;
+};