feat: support custom shiki themes in code blocks

Add support for custom shiki theme objects (ThemeRegistrationAny-compatible)
in addition to built-in BundledTheme string names.

- Add CustomTheme interface and ThemeInput type alias
- Widen theme types across CodePluginOptions, CodeHighlighterPlugin,
  HighlightOptions, StreamdownProps, and StreamdownContextType
- Extract theme names from objects for cache keys and codeToTokens
- Pass theme objects to createHighlighter for registration
- Add tests for custom theme objects and mixed themes
- Full backward compatibility with existing BundledTheme usage

Closes #409

Author: sleitor

packages/streamdown-code/__tests__/index.test.ts

@@ -1,5 +1,6 @@
 import type { BundledLanguage } from "shiki";
 import { describe, expect, it, vi } from "vitest";
+import type { CustomTheme } from "../index";
 import { code, createCodePlugin } from "../index";
 
 describe("code", () => {
@@ -303,4 +304,87 @@ describe("createCodePlugin", () => {
     expect(typeof plugin.getSupportedLanguages).toBe("function");
     expect(typeof plugin.getThemes).toBe("function");
   });
+
+  it("should create plugin with custom theme objects", () => {
+    const customLight: CustomTheme = {
+      name: "my-light-theme",
+      type: "light",
+      colors: { "editor.background": "#ffffff" },
+      tokenColors: [],
+    };
+    const customDark: CustomTheme = {
+      name: "my-dark-theme",
+      type: "dark",
+      colors: { "editor.background": "#1e1e1e" },
+      tokenColors: [],
+    };
+    const plugin = createCodePlugin({
+      themes: [customLight, customDark],
+    });
+    const themes = plugin.getThemes();
+    expect(themes[0]).toBe(customLight);
+    expect(themes[1]).toBe(customDark);
+  });
+
+  it("should create plugin with mixed built-in and custom themes", () => {
+    const customDark: CustomTheme = {
+      name: "my-dark-theme",
+      type: "dark",
+      colors: { "editor.background": "#1e1e1e" },
+      tokenColors: [],
+    };
+    const plugin = createCodePlugin({
+      themes: ["github-light", customDark],
+    });
+    const themes = plugin.getThemes();
+    expect(themes[0]).toBe("github-light");
+    expect(themes[1]).toBe(customDark);
+  });
+
+  it("should highlight code with custom theme objects", async () => {
+    const customLight: CustomTheme = {
+      name: "custom-light",
+      type: "light",
+      colors: {
+        "editor.background": "#ffffff",
+        "editor.foreground": "#000000",
+      },
+      tokenColors: [],
+    };
+    const customDark: CustomTheme = {
+      name: "custom-dark",
+      type: "dark",
+      colors: {
+        "editor.background": "#1e1e1e",
+        "editor.foreground": "#d4d4d4",
+      },
+      tokenColors: [],
+    };
+    const plugin = createCodePlugin({
+      themes: [customLight, customDark],
+    });
+
+    const callback = vi.fn();
+    const result = plugin.highlight(
+      {
+        code: "const x = 1;",
+        language: "javascript",
+        themes: [customLight, customDark],
+      },
+      callback
+    );
+
+    expect(result).toBeNull();
+
+    await vi.waitFor(
+      () => {
+        expect(callback).toHaveBeenCalled();
+      },
+      { timeout: 5000 }
+    );
+
+    const highlightResult = callback.mock.calls[0][0];
+    expect(highlightResult.tokens).toBeDefined();
+    expect(Array.isArray(highlightResult.tokens)).toBe(true);
+  });
 });

packages/streamdown-code/index.ts

@@ -14,6 +14,20 @@ import { createJavaScriptRegexEngine } from "shiki/engine/javascript";
 
 const jsEngine = createJavaScriptRegexEngine({ forgiving: true });
 
+/**
+ * A custom theme object compatible with shiki's ThemeRegistrationAny.
+ * Must have a `name` property for identification and caching.
+ */
+export interface CustomTheme {
+  name: string;
+  [key: string]: unknown;
+}
+
+/**
+ * Theme input type that accepts either a built-in theme name or a custom theme object.
+ */
+export type ThemeInput = BundledTheme | CustomTheme;
+
 /**
  * Result from code highlighting
  */
@@ -25,7 +39,7 @@ export type HighlightResult = TokensResult;
 export interface HighlightOptions {
   code: string;
   language: BundledLanguage;
-  themes: [string, string];
+  themes: [string, string] | [ThemeInput, ThemeInput];
 }
 
 /**
@@ -54,6 +68,14 @@ export interface CodeHighlighterPlugin {
    * Check if language is supported
    */
   supportsLanguage: (language: BundledLanguage) => boolean;
+  /**
+   * Get list of supported languages
+   */
+  getSupportedLanguages: () => BundledLanguage[];
+  /**
+   * Get the configured themes
+   */
+  getThemes: () => [ThemeInput, ThemeInput];
   type: "code-highlighter";
 }
 
@@ -63,9 +85,10 @@ export interface CodeHighlighterPlugin {
 export interface CodePluginOptions {
   /**
    * Default themes for syntax highlighting [light, dark]
+   * Accepts built-in theme names or custom theme objects.
    * @default ["github-light", "github-dark"]
    */
-  themes?: [BundledTheme, BundledTheme];
+  themes?: [ThemeInput, ThemeInput];
 }
 
 const languageAliases = Object.fromEntries(
@@ -104,10 +127,13 @@ const tokensCache = new Map<string, TokensResult>();
 // Subscribers for async token updates
 const subscribers = new Map<string, Set<(result: TokensResult) => void>>();
 
+const getThemeName = (theme: ThemeInput): string =>
+  typeof theme === "string" ? theme : theme.name;
+
 const getHighlighterCacheKey = (
   language: BundledLanguage,
-  themeNames: [string, string]
-) => `${language}-${themeNames[0]}-${themeNames[1]}`;
+  themes: [ThemeInput, ThemeInput]
+) => `${language}-${getThemeName(themes[0])}-${getThemeName(themes[1])}`;
 
 const getTokensCacheKey = (
   code: string,
@@ -121,9 +147,9 @@ const getTokensCacheKey = (
 
 const getHighlighter = (
   language: BundledLanguage,
-  themeNames: [string, string]
+  themes: [ThemeInput, ThemeInput]
 ): Promise<HighlighterGeneric<BundledLanguage, BundledTheme>> => {
-  const cacheKey = getHighlighterCacheKey(language, themeNames);
+  const cacheKey = getHighlighterCacheKey(language, themes);
 
   if (highlighterCache.has(cacheKey)) {
     return highlighterCache.get(cacheKey) as Promise<
@@ -132,7 +158,7 @@ const getHighlighter = (
   }
 
   const highlighterPromise = createHighlighter({
-    themes: themeNames,
+    themes: themes as (string | Record<string, unknown>)[],
     langs: [language],
     engine: jsEngine,
   });
@@ -147,7 +173,7 @@ const getHighlighter = (
 export function createCodePlugin(
   options: CodePluginOptions = {}
 ): CodeHighlighterPlugin {
-  const defaultThemes: [BundledTheme, BundledTheme] = options.themes ?? [
+  const defaultThemes: [ThemeInput, ThemeInput] = options.themes ?? [
     "github-light",
     "github-dark",
   ];
@@ -165,19 +191,23 @@ export function createCodePlugin(
       return Array.from(languageNames);
     },
 
-    getThemes(): [BundledTheme, BundledTheme] {
+    getThemes(): [ThemeInput, ThemeInput] {
       return defaultThemes;
     },
 
     highlight(
-      { code, language, themes: themeNames }: HighlightOptions,
+      { code, language, themes }: HighlightOptions,
       callback?: (result: HighlightResult) => void
     ): HighlightResult | null {
       const resolvedLanguage = normalizeLanguage(language);
+      const themeNames: [string, string] = [
+        getThemeName(themes[0]),
+        getThemeName(themes[1]),
+      ];
       const tokensCacheKey = getTokensCacheKey(
         code,
         resolvedLanguage,
-        themeNames as [string, string]
+        themeNames
       );
 
       // Return cached result if available
@@ -197,10 +227,7 @@ export function createCodePlugin(
       }
 
       // Start highlighting in background
-      getHighlighter(
-        resolvedLanguage as BundledLanguage,
-        themeNames as [string, string]
-      )
+      getHighlighter(resolvedLanguage as BundledLanguage, defaultThemes)
         .then((highlighter) => {
           const availableLangs = highlighter.getLoadedLanguages();
           const langToUse = (

packages/streamdown/index.tsx

@@ -16,7 +16,6 @@ import rehypeRaw from "rehype-raw";
 import rehypeSanitize, { defaultSchema } from "rehype-sanitize";
 import remarkGfm from "remark-gfm";
 import remend, { type RemendOptions } from "remend";
-import type { BundledTheme } from "shiki";
 import type { Pluggable } from "unified";
 import { type AnimateOptions, createAnimatePlugin } from "./lib/animate";
 import { BlockIncompleteContext } from "./lib/block-incomplete-context";
@@ -25,7 +24,7 @@ import { hasIncompleteCodeFence, hasTable } from "./lib/incomplete-code-utils";
 import { Markdown, type Options } from "./lib/markdown";
 import { parseMarkdownIntoBlocks } from "./lib/parse-blocks";
 import { PluginContext } from "./lib/plugin-context";
-import type { PluginConfig } from "./lib/plugin-types";
+import type { PluginConfig, ThemeInput } from "./lib/plugin-types";
 import { preprocessCustomTags } from "./lib/preprocess-custom-tags";
 import { cn } from "./lib/utils";
 
@@ -45,10 +44,12 @@ export { parseMarkdownIntoBlocks } from "./lib/parse-blocks";
 export type {
   CjkPlugin,
   CodeHighlighterPlugin,
+  CustomTheme,
   DiagramPlugin,
   HighlightOptions,
   MathPlugin,
   PluginConfig,
+  ThemeInput,
 } from "./lib/plugin-types";
 
 // Patterns for HTML indentation normalization
@@ -128,7 +129,7 @@ export type StreamdownProps = Options & {
   /** Normalize HTML block indentation to prevent 4+ spaces being treated as code blocks. @default false */
   normalizeHtmlIndentation?: boolean;
   className?: string;
-  shikiTheme?: [BundledTheme, BundledTheme];
+  shikiTheme?: [ThemeInput, ThemeInput];
   mermaid?: MermaidOptions;
   controls?: ControlsConfig;
   isAnimating?: boolean;
@@ -179,12 +180,12 @@ const carets = {
 
 // Combined context for better performance - reduces React tree depth from 5 nested providers to 1
 export interface StreamdownContextType {
+  shikiTheme: [ThemeInput, ThemeInput];
   controls: ControlsConfig;
   isAnimating: boolean;
   linkSafety?: LinkSafetyConfig;
   mermaid?: MermaidOptions;
   mode: "static" | "streaming";
-  shikiTheme: [BundledTheme, BundledTheme];
 }
 
 const defaultStreamdownContext: StreamdownContextType = {
@@ -285,7 +286,7 @@ export const Block = memo(
 
 Block.displayName = "Block";
 
-const defaultShikiTheme: [BundledTheme, BundledTheme] = [
+const defaultShikiTheme: [ThemeInput, ThemeInput] = [
   "github-light",
   "github-dark",
 ];

packages/streamdown/lib/plugin-types.ts

@@ -1,5 +1,20 @@
 import type { MermaidConfig } from "mermaid";
 import type { BundledLanguage, BundledTheme } from "shiki";
+
+/**
+ * A custom theme object compatible with shiki's ThemeRegistrationAny.
+ * Must have a `name` property for identification and caching.
+ */
+export interface CustomTheme {
+  name: string;
+  [key: string]: unknown;
+}
+
+/**
+ * Theme input type that accepts either a built-in theme name or a custom theme object.
+ */
+export type ThemeInput = BundledTheme | CustomTheme;
+
 import type { Pluggable } from "unified";
 
 /**
@@ -30,7 +45,7 @@ export interface HighlightResult {
 export interface HighlightOptions {
   code: string;
   language: BundledLanguage;
-  themes: [string, string];
+  themes: [string, string] | [ThemeInput, ThemeInput];
 }
 
 /**
@@ -59,6 +74,14 @@ export interface CodeHighlighterPlugin {
    * Check if language is supported
    */
   supportsLanguage: (language: BundledLanguage) => boolean;
+  /**
+   * Get list of supported languages
+   */
+  getSupportedLanguages: () => BundledLanguage[];
+  /**
+   * Get the configured themes
+   */
+  getThemes: () => [ThemeInput, ThemeInput];
   type: "code-highlighter";
 }