Use specific languages for shiki (#1938)

drizzle/meta/0017_snapshot.json

@@ -765,3 +765,4 @@
     "indexes": {}
   }
 }
+

e2e-tests/fixtures/context-limit-response.md

@@ -2,3 +2,4 @@ Here is a simple response to test the context limit banner functionality.
 
 This message simulates being close to the model's context window limit.
 
+

package.json

@@ -158,7 +158,7 @@
     "react-dom": "^19.2.1",
     "react-markdown": "^10.1.0",
     "react-resizable-panels": "^2.1.7",
-    "react-shiki": "^0.5.2",
+    "react-shiki": "^0.9.0",
     "recast": "^0.23.11",
     "remark-gfm": "^4.0.1",
     "shell-env": "^4.0.1",

src/__tests__/formatMessagesForSummary.test.ts

 import { formatMessagesForSummary } from "../ipc/handlers/chat_stream_handlers";
+import { describe, it, expect } from "vitest";
 
 describe("formatMessagesForSummary", () => {
   it("should return all messages when there are 8 or fewer messages", () => {

src/components/chat/CodeHighlight.tsx

-import React, {
-  useState,
-  useEffect,
-  useRef,
-  memo,
-  type ReactNode,
-} from "react";
-import { isInlineCode, useShikiHighlighter } from "react-shiki";
-import github from "@shikijs/themes/github-light-default";
-import githubDark from "@shikijs/themes/github-dark-default";
+import React, { useState, useEffect, memo, type ReactNode } from "react";
+import ShikiHighlighter, {
+  isInlineCode,
+  createHighlighterCore,
+  createJavaScriptRegexEngine,
+} from "react-shiki/core";
 import type { Element as HastElement } from "hast";
 import { useTheme } from "../../contexts/ThemeContext";
 import { Copy, Check } from "lucide-react";
+import github from "@shikijs/themes/github-light-default";
+import githubDark from "@shikijs/themes/github-dark-default";
+// common languages
+import astro from "@shikijs/langs/astro";
+import css from "@shikijs/langs/css";
+import graphql from "@shikijs/langs/graphql";
+import html from "@shikijs/langs/html";
+import java from "@shikijs/langs/java";
+import javascript from "@shikijs/langs/javascript";
+import json from "@shikijs/langs/json";
+import jsx from "@shikijs/langs/jsx";
+import less from "@shikijs/langs/less";
+import markdown from "@shikijs/langs/markdown";
+import python from "@shikijs/langs/python";
+import sass from "@shikijs/langs/sass";
+import scss from "@shikijs/langs/scss";
+import shell from "@shikijs/langs/shell";
+import sql from "@shikijs/langs/sql";
+import tsx from "@shikijs/langs/tsx";
+import typescript from "@shikijs/langs/typescript";
+import vue from "@shikijs/langs/vue";
+
+type HighlighterCore = Awaited<ReturnType<typeof createHighlighterCore>>;
+
+// Create a singleton highlighter instance
+let highlighterPromise: Promise<HighlighterCore> | null = null;
+
+function getHighlighter(): Promise<HighlighterCore> {
+  if (!highlighterPromise) {
+    highlighterPromise = createHighlighterCore({
+      themes: [github, githubDark],
+      langs: [
+        astro,
+        css,
+        graphql,
+        html,
+        java,
+        javascript,
+        json,
+        jsx,
+        less,
+        markdown,
+        python,
+        sass,
+        scss,
+        shell,
+        sql,
+        tsx,
+        typescript,
+        vue,
+      ],
+      engine: createJavaScriptRegexEngine(),
+    });
+  }
+  return highlighterPromise as Promise<HighlighterCore>;
+}
+
+function useHighlighter() {
+  const [highlighter, setHighlighter] = useState<HighlighterCore>();
+
+  useEffect(() => {
+    getHighlighter().then(setHighlighter);
+  }, []);
+
+  return highlighter;
+}
 
 interface CodeHighlightProps {
   className?: string | undefined;
@@ -32,29 +94,8 @@ export const CodeHighlight = memo(
     };
 
     const { isDarkMode } = useTheme();
+    const highlighter = useHighlighter();
 
-    // Cache for the highlighted code
-    const highlightedCodeCache = useRef<ReactNode | null>(null);
-
-    // Only update the highlighted code if the inputs change
-    const highlightedCode = useShikiHighlighter(
-      code,
-      language,
-      isDarkMode ? githubDark : github,
-      {
-        delay: 150,
-      },
-    );
-
-    // Update the cache whenever we get a new highlighted code
-    useEffect(() => {
-      if (highlightedCode) {
-        highlightedCodeCache.current = highlightedCode;
-      }
-    }, [highlightedCode]);
-
-    // Use the cached version during transitions to prevent flickering
-    const displayedCode = highlightedCode || highlightedCodeCache.current;
     return !isInline ? (
       <div
         className="shiki not-prose relative [&_pre]:overflow-auto 
@@ -77,7 +118,20 @@ export const CodeHighlight = memo(
             )}
           </div>
         ) : null}
-        {displayedCode}
+        {highlighter ? (
+          <ShikiHighlighter
+            highlighter={highlighter}
+            language={language}
+            theme={isDarkMode ? "github-dark-default" : "github-light-default"}
+            delay={150}
+          >
+            {code}
+          </ShikiHighlighter>
+        ) : (
+          <pre>
+            <code>{code}</code>
+          </pre>
+        )}
       </div>
     ) : (
       <code className={className} {...props}>