refactor: replace react-syntax-highlighter with shiki in CodeBlock (#149)

* refactor: replace react-syntax-highlighter with shiki in CodeBlock

* Create perky-wolves-grow.md

---------

Co-authored-by: Hayden Bleasel <hello@haydenbleasel.com>

Author: tobiasbueschel

.changeset/perky-wolves-grow.md

@@ -0,0 +1,5 @@
+---
+"ai-elements": patch
+---
+
+refactor: replace react-syntax-highlighter with shiki in CodeBlock

packages/elements/__tests__/code-block.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 import { CodeBlock, CodeBlockCopyButton } from "../src/code-block";
@@ -11,22 +11,26 @@ Object.assign(navigator, {
 });
 
 describe("CodeBlock", () => {
-  it("renders code content", () => {
+  it("renders code content", async () => {
     const { container } = render(
       <CodeBlock code="const foo = 'bar';" language="javascript" />
     );
-    expect(container.textContent).toContain("const foo");
+    await waitFor(() => {
+      expect(container.textContent).toContain("const foo");
+    });
   });
 
-  it("renders with line numbers", () => {
+  it("renders with line numbers", async () => {
     const { container } = render(
       <CodeBlock
         code="line1\nline2"
         language="javascript"
         showLineNumbers={true}
       />
     );
-    expect(container.textContent).toContain("line1");
+    await waitFor(() => {
+      expect(container.textContent).toContain("line1");
+    });
   });
 
   it("renders children actions", () => {
@@ -43,6 +47,8 @@ describe("CodeBlock", () => {
       <CodeBlock className="custom-class" code="code" language="javascript" />
     );
     expect(container.firstChild).toHaveClass("custom-class");
+    expect(container.firstChild).toHaveClass("group");
+    expect(container.firstChild).toHaveClass("relative");
   });
 });
 

packages/elements/__tests__/tool.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import { describe, expect, it } from "vitest";
 import {
   Tool,
@@ -137,7 +137,7 @@ describe("ToolContent", () => {
 });
 
 describe("ToolInput", () => {
-  it("renders input parameters", () => {
+  it("renders input parameters", async () => {
     const input = { query: "test search" };
     render(
       <Tool defaultOpen>
@@ -148,7 +148,9 @@ describe("ToolInput", () => {
       </Tool>
     );
     expect(screen.getByText("Parameters")).toBeInTheDocument();
-    expect(screen.getAllByText(/"query"/)[0]).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getAllByText(/"query"/)[0]).toBeInTheDocument();
+    });
   });
 
   it("renders JSON formatted input", () => {

packages/elements/package.json

@@ -21,7 +21,7 @@
     "motion": "^12.23.24",
     "nanoid": "^5.1.6",
     "react": "19.2.0",
-    "react-syntax-highlighter": "^15.6.6",
+    "shiki": "3.13.0",
     "streamdown": "^1.4.0",
     "tokenlens": "^1.3.1",
     "use-stick-to-bottom": "^1.1.1"
@@ -31,8 +31,8 @@
     "@testing-library/jest-dom": "^6.9.1",
     "@testing-library/react": "^16.3.0",
     "@testing-library/user-event": "^14.6.1",
+    "@types/hast": "^3.0.4",
     "@types/react": "19.2.2",
-    "@types/react-syntax-highlighter": "^15.5.13",
     "@vitejs/plugin-react": "^5.0.4",
     "@vitest/coverage-v8": "^4.0.1",
     "jsdom": "^27.0.1",

packages/elements/src/code-block.tsx

@@ -2,14 +2,24 @@
 
 import { Button } from "@repo/shadcn-ui/components/ui/button";
 import { cn } from "@repo/shadcn-ui/lib/utils";
+import type { Element } from "hast";
 import { CheckIcon, CopyIcon } from "lucide-react";
-import type { ComponentProps, HTMLAttributes, ReactNode } from "react";
-import { createContext, useContext, useState } from "react";
-import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import {
-  oneDark,
-  oneLight,
-} from "react-syntax-highlighter/dist/esm/styles/prism";
+  type ComponentProps,
+  createContext,
+  type HTMLAttributes,
+  useContext,
+  useEffect,
+  useRef,
+  useState,
+} from "react";
+import { type BundledLanguage, codeToHtml, type ShikiTransformer } from "shiki";
+
+type CodeBlockProps = HTMLAttributes<HTMLDivElement> & {
+  code: string;
+  language: BundledLanguage;
+  showLineNumbers?: boolean;
+};
 
 type CodeBlockContextType = {
   code: string;
@@ -19,85 +29,106 @@ const CodeBlockContext = createContext<CodeBlockContextType>({
   code: "",
 });
 
-export type CodeBlockProps = HTMLAttributes<HTMLDivElement> & {
-  code: string;
-  language: string;
-  showLineNumbers?: boolean;
-  children?: ReactNode;
+const lineNumberTransformer: ShikiTransformer = {
+  name: "line-numbers",
+  line(node: Element, line: number) {
+    node.children.unshift({
+      type: "element",
+      tagName: "span",
+      properties: {
+        className: [
+          "inline-block",
+          "min-w-10",
+          "mr-4",
+          "text-right",
+          "select-none",
+          "text-muted-foreground",
+        ],
+      },
+      children: [{ type: "text", value: String(line) }],
+    });
+  },
 };
 
+export async function highlightCode(
+  code: string,
+  language: BundledLanguage,
+  showLineNumbers = false
+) {
+  const transformers: ShikiTransformer[] = showLineNumbers
+    ? [lineNumberTransformer]
+    : [];
+
+  return await Promise.all([
+    codeToHtml(code, {
+      lang: language,
+      theme: "one-light",
+      transformers,
+    }),
+    codeToHtml(code, {
+      lang: language,
+      theme: "one-dark-pro",
+      transformers,
+    }),
+  ]);
+}
+
 export const CodeBlock = ({
   code,
   language,
   showLineNumbers = false,
   className,
   children,
   ...props
-}: CodeBlockProps) => (
-  <CodeBlockContext.Provider value={{ code }}>
-    <div
-      className={cn(
-        "relative w-full overflow-hidden rounded-md border bg-background text-foreground",
-        className
-      )}
-      {...props}
-    >
-      <div className="relative">
-        <SyntaxHighlighter
-          className="overflow-hidden dark:hidden"
-          codeTagProps={{
-            className: "font-mono text-sm",
-          }}
-          customStyle={{
-            margin: 0,
-            padding: "1rem",
-            fontSize: "0.875rem",
-            background: "hsl(var(--background))",
-            color: "hsl(var(--foreground))",
-          }}
-          language={language}
-          lineNumberStyle={{
-            color: "hsl(var(--muted-foreground))",
-            paddingRight: "1rem",
-            minWidth: "2.5rem",
-          }}
-          showLineNumbers={showLineNumbers}
-          style={oneLight}
-        >
-          {code}
-        </SyntaxHighlighter>
-        <SyntaxHighlighter
-          className="hidden overflow-hidden dark:block"
-          codeTagProps={{
-            className: "font-mono text-sm",
-          }}
-          customStyle={{
-            margin: 0,
-            padding: "1rem",
-            fontSize: "0.875rem",
-            background: "hsl(var(--background))",
-            color: "hsl(var(--foreground))",
-          }}
-          language={language}
-          lineNumberStyle={{
-            color: "hsl(var(--muted-foreground))",
-            paddingRight: "1rem",
-            minWidth: "2.5rem",
-          }}
-          showLineNumbers={showLineNumbers}
-          style={oneDark}
-        >
-          {code}
-        </SyntaxHighlighter>
-        {children && (
-          <div className="absolute top-2 right-2 flex items-center gap-2">
-            {children}
-          </div>
+}: CodeBlockProps) => {
+  const [html, setHtml] = useState<string>("");
+  const [darkHtml, setDarkHtml] = useState<string>("");
+  const mounted = useRef(false);
+
+  useEffect(() => {
+    highlightCode(code, language, showLineNumbers).then(([light, dark]) => {
+      if (!mounted.current) {
+        setHtml(light);
+        setDarkHtml(dark);
+        mounted.current = true;
+      }
+    });
+
+    return () => {
+      mounted.current = false;
+    };
+  }, [code, language, showLineNumbers]);
+
+  return (
+    <CodeBlockContext.Provider value={{ code }}>
+      <div
+        className={cn(
+          "group relative w-full overflow-hidden rounded-md border bg-background text-foreground",
+          className
         )}
+        {...props}
+      >
+        <div className="relative">
+          <div
+            className="overflow-hidden dark:hidden [&>pre]:m-0 [&>pre]:bg-background! [&>pre]:p-4 [&>pre]:text-foreground! [&>pre]:text-sm [&_code]:font-mono [&_code]:text-sm"
+            // biome-ignore lint/security/noDangerouslySetInnerHtml: "this is needed."
+            dangerouslySetInnerHTML={{ __html: html }}
+          />
+          <div
+            className="hidden overflow-hidden dark:block [&>pre]:m-0 [&>pre]:bg-background! [&>pre]:p-4 [&>pre]:text-foreground! [&>pre]:text-sm [&_code]:font-mono [&_code]:text-sm"
+            // biome-ignore lint/security/noDangerouslySetInnerHtml: "this is needed."
+            dangerouslySetInnerHTML={{ __html: darkHtml }}
+          />
+          {children && (
+            <div className="absolute top-2 right-2 flex items-center gap-2">
+              {children}
+            </div>
+          )}
+        </div>
       </div>
-    </div>
-  </CodeBlockContext.Provider>
-);
+    </CodeBlockContext.Provider>
+  );
+};
 
 export type CodeBlockCopyButtonProps = ComponentProps<typeof Button> & {
   onCopy?: () => void;
@@ -117,7 +148,7 @@ export const CodeBlockCopyButton = ({
   const { code } = useContext(CodeBlockContext);
 
   const copyToClipboard = async () => {
-    if (typeof window === "undefined" || !navigator.clipboard.writeText) {
+    if (typeof window === "undefined" || !navigator?.clipboard?.writeText) {
       onError?.(new Error("Clipboard API not available"));
       return;
     }