feat: Add codehighlight for the Code'

packages/ui/src/components/Code/Code.stories.tsx

@@ -90,6 +90,24 @@ export const LongWithScroll: StoryObj<typeof CodeBlock> = {
   },
 };
 
+export const ShikiHighlight: StoryObj<typeof CodeBlock> = {
+  render: (args) => (
+    <CodeBlock {...args}>
+      <Code>
+        {`import React from 'react';
+
+export const MyComponent = () => {
+  return <div>Hello World</div>;
+};`}
+      </Code>
+    </CodeBlock>
+  ),
+  args: {
+    label: 'ShikiHighlight.tsx',
+    language: 'tsx',
+  },
+};
+
 export const PlainText: StoryObj<typeof CodeBlock> = {
   render: (args) => (
     <CodeBlock {...args}>

packages/ui/src/components/Code/Code.test.tsx

@@ -1,9 +1,17 @@
 import { render, screen, fireEvent, act } from '@testing-library/react';
 import React from 'react';
+import { vi, describe, it, expect, beforeEach } from 'vitest';
 
 import { Code } from './Code';
 import { CodeBlock } from './CodeBlock/CodeBlock';
 
+// Mock shiki
+vi.mock('shiki', () => ({
+  createHighlighter: vi.fn().mockResolvedValue({
+    codeToHtml: vi.fn().mockReturnValue('<pre><code><span class="line">highlighted code</span></code></pre>'),
+  }),
+}));
+
 // Mock clipboard API
 const mockWriteText = vi.fn().mockImplementation(() => Promise.resolve());
 Object.assign(navigator, {
@@ -15,6 +23,7 @@ Object.assign(navigator, {
 describe('Code Component', () => {
   beforeEach(() => {
     mockWriteText.mockClear();
+    vi.clearAllMocks();
   });
 
   it('renders inline code correctly', () => {
@@ -46,7 +55,6 @@ describe('Code Component', () => {
 
     await act(async () => {
       fireEvent.click(codeElement);
-      await Promise.resolve();
     });
 
     expect(mockWriteText).toHaveBeenCalledWith('copied!');
@@ -64,7 +72,6 @@ describe('Code Component', () => {
     const copyButton = screen.getByLabelText('Copy code');
     await act(async () => {
       fireEvent.click(copyButton);
-      await Promise.resolve();
     });
 
     expect(mockWriteText).toHaveBeenCalledWith('block copy');
@@ -76,9 +83,22 @@ describe('Code Component', () => {
     const codeElement = screen.getByText('extracted');
     await act(async () => {
       fireEvent.click(codeElement);
-      await Promise.resolve();
     });
 
     expect(mockWriteText).toHaveBeenCalledWith('extracted');
   });
+
+  it('uses shiki for highlighting when language is provided in block mode', async () => {
+    await act(async () => {
+      render(
+        <CodeBlock language="typescript">
+          <Code>const x: number = 1;</Code>
+        </CodeBlock>
+      );
+    });
+
+    // Wait for shiki async highlighting
+    const highlighted = await screen.findByText('highlighted code');
+    expect(highlighted).toBeInTheDocument();
+  });
 });

packages/ui/src/components/Code/Code.tsx

@@ -3,8 +3,12 @@ import React, { memo, useCallback, useEffect, useRef, useState } from 'react';
 import { usePress } from 'react-aria';
 
 import { useCodeContext } from './contexts/CodeContext/CodeContext';
+import { codeHighlighter } from './lib/highlighter/highlighter';
 
 import type { ClassNameProps, TestIdProps } from '@ermnvldmr/stl';
+import type { CodeLanguage } from './lib/highlighter/highlighter';
+
+export type { CodeLanguage };
 
 /**
  * Props for the Code component.
@@ -18,6 +22,12 @@ export interface CodeProps extends ClassNameProps, TestIdProps {
    * If omitted in inline mode, the component will attempt to extract text from `children`.
    */
   copyValue?: string;
+
+  /**
+   * The language of the code. Only used when rendered inside a CodeBlock.
+   * If provided, shiki will be used for syntax highlighting.
+   */
+  language?: CodeLanguage;
 }
 
 /**
@@ -30,12 +40,16 @@ export interface CodeProps extends ClassNameProps, TestIdProps {
 export const Code = memo(function Code({
   children,
   copyValue,
+  language: propsLanguage,
   className,
   'data-testid': testId,
 }: CodeProps) {
   const context = useCodeContext();
   const [isCopied, setIsCopied] = useState(false);
   const codeRef = useRef<HTMLElement>(null);
+  const [highlightedHtml, setHighlightedHtml] = useState<string | null>(null);
+
+  const language = propsLanguage ?? context?.language;
 
   const handleCopy = useCallback(async () => {
     try {
@@ -67,6 +81,22 @@ export const Code = memo(function Code({
     }
   }, [isCopied]);
 
+  useEffect(() => {
+    if (context?.isInCodeBlock && language && typeof children === 'string') {
+      const highlight = async () => {
+        try {
+          const html = await codeHighlighter.highlight(children, language);
+          setHighlightedHtml(html);
+        } catch (err) {
+          console.error('Failed to highlight code:', err);
+        }
+      };
+      void highlight();
+    } else {
+      setHighlightedHtml(null);
+    }
+  }, [children, language, context?.isInCodeBlock]);
+
   const { pressProps } = usePress({
     onPress: onCopyPress,
   });
@@ -74,7 +104,39 @@ export const Code = memo(function Code({
   if (context?.isInCodeBlock) {
     const { showLineNumbers } = context;
 
-    let content = children;
+    let content: React.ReactNode = children;
+
+    if (highlightedHtml) {
+      return (
+        <div
+          ref={codeRef as any}
+          className={cn(
+            'block min-w-full font-mono text-sm',
+            '[&>pre]:!bg-transparent [&>pre]:!p-0 [&>pre]:!m-0',
+            showLineNumbers && [
+              '[&>pre>code]:[counter-reset:line]',
+              '[&>pre>code>.line]:relative [&>pre>code>.line]:pl-10',
+              '[&>pre>code>.line]:[counter-increment:line]',
+              '[&>pre>code>.line]:before:content-[counter(line)]',
+              '[&>pre>code>.line]:before:absolute [&>pre>code>.line]:before:left-0 [&>pre>code>.line]:before:top-0',
+              '[&>pre>code>.line]:before:inline-block',
+              '[&>pre>code>.line]:before:w-8',
+              '[&>pre>code>.line]:before:text-right',
+              '[&>pre>code>.line]:before:text-[var(--rb-muted-text)]/50',
+              '[&>pre>code>.line]:before:select-none',
+              '[&>pre>code>.line]:before:border-r',
+              '[&>pre>code>.line]:before:border-[var(--rb-outline)]/10',
+              '[&>pre>code>.line]:before:pr-3',
+              '[&>pre>code>.line]:before:font-mono [&>pre>code>.line]:before:text-[10px]',
+            ],
+            className
+          )}
+          dangerouslySetInnerHTML={{ __html: highlightedHtml }}
+          data-testid={testId}
+        />
+      );
+    }
+
     if (typeof children === 'string') {
       content = children
         .trimEnd()

packages/ui/src/components/Code/CodeBlock/CodeBlock.tsx

@@ -7,6 +7,7 @@ import { Container } from '../../Container/Container';
 import { CodeContext, type CodeContextValue } from '../contexts/CodeContext/CodeContext';
 
 import type { ClassNameProps, TestIdProps } from '@ermnvldmr/stl';
+import type { CodeLanguage } from '../lib/highlighter/highlighter';
 
 /**
  * Props for the CodeBlock component.
@@ -21,6 +22,11 @@ export interface CodeBlockProps extends ClassNameProps, TestIdProps {
    */
   label?: string;
 
+  /**
+   * The language of the code block.
+   */
+  language?: CodeLanguage;
+
   /**
    * Whether to render line number gutters.
    * Requires children (Code) to be rendered as block-level elements.
@@ -49,6 +55,7 @@ export interface CodeBlockProps extends ClassNameProps, TestIdProps {
 export const CodeBlock = memo(function CodeBlock({
   children,
   label,
+  language,
   showLineNumbers = true,
   copyValue,
   maxHeight,
@@ -92,8 +99,9 @@ export const CodeBlock = memo(function CodeBlock({
     () => ({
       isInCodeBlock: true,
       showLineNumbers,
+      language,
     }),
-    [showLineNumbers]
+    [showLineNumbers, language]
   );
 
   return (

packages/ui/src/components/Code/contexts/CodeContext/CodeContext.ts

@@ -1,5 +1,7 @@
 import { createContext, useContext } from 'react';
 
+import type { CodeLanguage } from '../../lib/highlighter/highlighter';
+
 /**
  * Value for the Code context.
  */
@@ -8,6 +10,8 @@ export interface CodeContextValue {
   isInCodeBlock: boolean;
   /** Whether to show line numbers in the block. */
   showLineNumbers?: boolean;
+  /** The language of the code in the block. */
+  language?: CodeLanguage;
 }
 
 /**

packages/ui/src/components/Code/lib/highlighter/highlighter.ts

@@ -0,0 +1,113 @@
+import { createHighlighter } from 'shiki';
+
+import type { Highlighter } from 'shiki';
+
+/**
+ * Supported languages for syntax highlighting.
+ */
+export type CodeLanguage =
+  | 'typescript'
+  | 'javascript'
+  | 'tsx'
+  | 'jsx'
+  | 'bash'
+  | 'shell'
+  | 'json'
+  | 'css'
+  | 'html'
+  | 'markdown'
+  | 'yaml'
+  | 'dockerfile'
+  | 'toml'
+  | 'rust'
+  | 'go'
+  | 'python';
+
+export const SUPPORTED_LANGUAGES: CodeLanguage[] = [
+  'typescript',
+  'javascript',
+  'tsx',
+  'jsx',
+  'bash',
+  'shell',
+  'json',
+  'css',
+  'html',
+  'markdown',
+  'yaml',
+  'dockerfile',
+  'toml',
+  'rust',
+  'go',
+  'python',
+];
+
+/**
+ * A singleton highlighter service for code blocks.
+ */
+class CodeHighlighter {
+  private static instance: CodeHighlighter;
+  private highlighter: Highlighter | null = null;
+  private initializingPromise: Promise<Highlighter> | null = null;
+
+  private constructor() {}
+
+  /**
+   * Gets the singleton instance of the highlighter.
+   */
+  public static getInstance(): CodeHighlighter {
+    if (!CodeHighlighter.instance) {
+      CodeHighlighter.instance = new CodeHighlighter();
+    }
+    return CodeHighlighter.instance;
+  }
+
+  /**
+   * Initializes the highlighter if it hasn't been already.
+   */
+  private async getHighlighter(): Promise<Highlighter> {
+    if (this.highlighter) {
+      return this.highlighter;
+    }
+
+    if (this.initializingPromise) {
+      return this.initializingPromise;
+    }
+
+    this.initializingPromise = createHighlighter({
+      themes: ['github-dark', 'github-light'],
+      langs: SUPPORTED_LANGUAGES,
+    }).then((highlighter) => {
+      this.highlighter = highlighter;
+      this.initializingPromise = null;
+      return highlighter;
+    });
+
+    return this.initializingPromise;
+  }
+
+  /**
+   * Highlights the given code.
+   *
+   * @param code - The raw source code to highlight
+   * @param lang - The language of the code
+   * @returns A promise resolving to the highlighted HTML
+   */
+  public async highlight(code: string, lang: CodeLanguage): Promise<string> {
+    try {
+      const highlighter = await this.getHighlighter();
+      return highlighter.codeToHtml(code, {
+        lang,
+        themes: {
+          light: 'github-light',
+          dark: 'github-dark',
+        },
+      });
+    } catch (err) {
+      console.error('Failed to highlight code:', err);
+      throw err;
+    }
+  }
+}
+
+export const codeHighlighter = CodeHighlighter.getInstance();

packages/ui/src/components/Markdown/MarkdownCode/MarkdownCode.tsx

@@ -1,12 +1,25 @@
-import React, { memo } from 'react';
+import React, { memo, useMemo } from 'react';
 
 import { Code } from '../../Code/Code';
 
 import type { CodeProps } from '../../Code/Code';
+import type { CodeLanguage } from '../../Code/lib/highlighter/highlighter';
 
 /**
  * A Markdown code component for inline snippets.
  */
 export const MarkdownCode = memo(function MarkdownCode(props: CodeProps) {
-  return <Code {...props} />;
+  const { className, ...rest } = props;
+
+  const language = useMemo(() => {
+    if (className) {
+      const match = /language-(\w+)/.exec(className);
+      if (match) {
+        return match[1] as CodeLanguage;
+      }
+    }
+    return undefined;
+  }, [className]);
+
+  return <Code language={language} {...rest} />;
 });

packages/ui/src/components/Markdown/MarkdownPre/MarkdownPre.tsx

@@ -1,8 +1,9 @@
-import React, { memo } from 'react';
+import React, { memo, useMemo } from 'react';
 
 import { CodeBlock } from '../../Code/CodeBlock/CodeBlock';
 
 import type { CodeBlockProps } from '../../Code/CodeBlock/CodeBlock';
+import type { CodeLanguage } from '../../Code/lib/highlighter/highlighter';
 
 /**
  * A Markdown pre component for code blocks.
@@ -10,5 +11,26 @@ import type { CodeBlockProps } from '../../Code/CodeBlock/CodeBlock';
  * Wraps content in the design system CodeBlock.
  */
 export const MarkdownPre = memo(function MarkdownPre(props: CodeBlockProps) {
-  return <CodeBlock {...props} />;
+  const { children, ...rest } = props;
+
+  // MDX often passes a <code> element as the only child of <pre>.
+  // That <code> element usually has the language class (e.g., "language-typescript").
+  const language = useMemo(() => {
+    if (React.isValidElement<{ className?: string }>(children) && children.props.className) {
+      const className = children.props.className;
+      const match = /language-(\w+)/.exec(className);
+      if (match) {
+        // We cast here because we know it comes from markdown,
+        // shiki will handle it or fallback to plaintext if unknown at runtime.
+        return match[1] as CodeLanguage;
+      }
+    }
+    return undefined;
+  }, [children]);
+
+  return (
+    <CodeBlock language={language} {...rest}>
+      {children}
+    </CodeBlock>
+  );
 });