Files
Netcatty/components/ai-elements/streamdownCodeHighlighter.ts

78 lines
1.9 KiB
TypeScript

import type {
CodeHighlighterPlugin,
HighlightOptions,
} from 'streamdown';
import type { BundledLanguage } from 'shiki';
type HighlightResult = NonNullable<ReturnType<CodeHighlighterPlugin['highlight']>>;
const PLAIN_TEXT_LANGUAGES = new Set([
'',
'plain',
'plaintext',
'text',
'txt',
]);
const LANGUAGE_ALIASES: Record<string, BundledLanguage> = {
cfg: 'ini',
conf: 'ini',
config: 'ini',
};
export const createPlainCodeHighlightResult = (source: string): HighlightResult => {
const code = source.replace(/\n+$/, '');
return {
bg: 'transparent',
fg: 'inherit',
tokens: code.split('\n').map((line) => [
{
content: line,
color: 'inherit',
bgColor: 'transparent',
htmlStyle: {},
offset: 0,
},
]),
};
};
const normalizeLanguageKey = (language: string): string =>
language.trim().toLowerCase();
export const resolveSupportedCodeLanguage = (
highlighter: CodeHighlighterPlugin,
language: string,
): BundledLanguage | null => {
const key = normalizeLanguageKey(language);
if (PLAIN_TEXT_LANGUAGES.has(key)) return null;
const direct = key as BundledLanguage;
if (highlighter.supportsLanguage(direct)) return direct;
const alias = LANGUAGE_ALIASES[key];
if (alias && highlighter.supportsLanguage(alias)) return alias;
return null;
};
export const createSafeCodeHighlighter = (
highlighter: CodeHighlighterPlugin,
): CodeHighlighterPlugin => ({
...highlighter,
supportsLanguage(language) {
return resolveSupportedCodeLanguage(highlighter, language) !== null;
},
highlight(options: HighlightOptions, callback?: (result: HighlightResult) => void) {
const supportedLanguage = resolveSupportedCodeLanguage(highlighter, options.language);
if (!supportedLanguage) {
return createPlainCodeHighlightResult(options.code);
}
return highlighter.highlight(
{ ...options, language: supportedLanguage },
callback,
);
},
});