188 lines · 5.6 KB · TypeScript
RawDownload
import {  createHighlighterCore,  type HighlighterCore,  type LanguageInput,} from "shiki/core";import { createJavaScriptRegexEngine } from "shiki/engine/javascript"; // ---------------------------------------------------------------------------// Supported languages — explicit imports to avoid bundling all grammars// --------------------------------------------------------------------------- const SUPPORTED_LANGS = [  "json",  "xml",  "yaml",  "shellscript",  "typescript",  "javascript",  "python",  "html",  "css",  "markdown",  "sql",  "graphql",  "go",  "rust",  "java",  "ruby",  "php",  "swift",  "kotlin",  "c",  "cpp",  "csharp",  "tsx",  "jsx",  "toml",  "dockerfile",  "diff",  "http",  "jsonc",  "log",  "proto",] as const; type SupportedLang = (typeof SUPPORTED_LANGS)[number]; const LANG_ALIASES: Record<string, SupportedLang> = {  sh: "shellscript",  shell: "shellscript",  bash: "shellscript",  zsh: "shellscript",  ts: "typescript",  js: "javascript",  py: "python",  rb: "ruby",  rs: "rust",  "c++": "cpp",  "c#": "csharp",  cs: "csharp",  kt: "kotlin",  md: "markdown",  gql: "graphql",  yml: "yaml",}; const LANG_LOADERS: Record<SupportedLang, () => LanguageInput> = {  json: () => import("@shikijs/langs/json"),  xml: () => import("@shikijs/langs/xml"),  yaml: () => import("@shikijs/langs/yaml"),  shellscript: () => import("@shikijs/langs/shellscript"),  typescript: () => import("@shikijs/langs/typescript"),  javascript: () => import("@shikijs/langs/javascript"),  python: () => import("@shikijs/langs/python"),  html: () => import("@shikijs/langs/html"),  css: () => import("@shikijs/langs/css"),  markdown: () => import("@shikijs/langs/markdown"),  sql: () => import("@shikijs/langs/sql"),  graphql: () => import("@shikijs/langs/graphql"),  go: () => import("@shikijs/langs/go"),  rust: () => import("@shikijs/langs/rust"),  java: () => import("@shikijs/langs/java"),  ruby: () => import("@shikijs/langs/ruby"),  php: () => import("@shikijs/langs/php"),  swift: () => import("@shikijs/langs/swift"),  kotlin: () => import("@shikijs/langs/kotlin"),  c: () => import("@shikijs/langs/c"),  cpp: () => import("@shikijs/langs/cpp"),  csharp: () => import("@shikijs/langs/csharp"),  tsx: () => import("@shikijs/langs/tsx"),  jsx: () => import("@shikijs/langs/jsx"),  toml: () => import("@shikijs/langs/toml"),  dockerfile: () => import("@shikijs/langs/dockerfile"),  diff: () => import("@shikijs/langs/diff"),  http: () => import("@shikijs/langs/http"),  jsonc: () => import("@shikijs/langs/jsonc"),  log: () => import("@shikijs/langs/log"),  proto: () => import("@shikijs/langs/proto"),}; const supportedSet = new Set<string>([  ...SUPPORTED_LANGS,  ...Object.keys(LANG_ALIASES),]); export const THEME = "vitesse-dark"; export function resolveLang(lang: string): SupportedLang | null {  const l = lang.trim().toLowerCase();  if (l in LANG_LOADERS) return l as SupportedLang;  if (l in LANG_ALIASES) return LANG_ALIASES[l]!;  return null;} export function isSupportedLang(lang: string): boolean {  return supportedSet.has(lang.trim().toLowerCase());} // ---------------------------------------------------------------------------// Shared singleton highlighter — lazy, created on first use// --------------------------------------------------------------------------- const jsEngine = createJavaScriptRegexEngine({ forgiving: true }); let _promise: Promise<HighlighterCore> | null = null; export function getHighlighter(): Promise<HighlighterCore> {  if (!_promise) {    _promise = createHighlighterCore({      themes: [import("@shikijs/themes/vitesse-dark")],      langs: Object.values(LANG_LOADERS).map((loader) => loader()),      engine: jsEngine,    });  }  return _promise;} // ---------------------------------------------------------------------------// Streamdown code highlighter plugin// --------------------------------------------------------------------------- import type { CodeHighlighterPlugin, ThemeInput } from "streamdown"; const tokensCache = new Map<string, unknown>();const pendingCallbacks = new Map<string, Set<(result: unknown) => void>>(); export function createCodeHighlighterPlugin(): CodeHighlighterPlugin {  return {    name: "shiki" as const,    type: "code-highlighter" as const,    getSupportedLanguages: () => [...SUPPORTED_LANGS] as string[] as never,    getThemes: () => [THEME as ThemeInput, THEME as ThemeInput],    supportsLanguage: (language: string) => isSupportedLang(language),    highlight(options, callback) {      const resolved = resolveLang(options.language);      const lang = resolved ?? "json";      const key = `${lang}:${options.code.length}:${options.code.slice(0, 128)}`;       const cached = tokensCache.get(key);      if (cached) return cached as never;       if (callback) {        if (!pendingCallbacks.has(key)) {          pendingCallbacks.set(key, new Set());        }        pendingCallbacks.get(key)!.add(callback as (result: unknown) => void);      }       void getHighlighter().then((highlighter) => {        if (tokensCache.has(key)) {          const result = tokensCache.get(key);          pendingCallbacks.get(key)?.forEach((cb) => cb(result));          pendingCallbacks.delete(key);          return;        }        const result = highlighter.codeToTokens(options.code, {          lang,          themes: { light: THEME, dark: THEME },        });        tokensCache.set(key, result);        pendingCallbacks.get(key)?.forEach((cb) => cb(result));        pendingCallbacks.delete(key);      });       return null;    },  };}