diff --git a/components/ai-elements/message.tsx b/components/ai-elements/message.tsx index 97dddf498..80bc10088 100644 --- a/components/ai-elements/message.tsx +++ b/components/ai-elements/message.tsx @@ -4,6 +4,7 @@ import { code } from '@streamdown/code'; import type { ComponentProps, HTMLAttributes } from 'react'; import { memo } from 'react'; import { Streamdown } from 'streamdown'; +import { createSafeCodeHighlighter } from './streamdownCodeHighlighter'; export type MessageProps = HTMLAttributes & { from: 'user' | 'assistant' | 'system' | 'tool'; @@ -60,7 +61,8 @@ export const MessageActions = ({ className, children, ...props }: MessageActions ); -const streamdownPlugins = { cjk, code }; +const safeCode = createSafeCodeHighlighter(code); +const streamdownPlugins = { cjk, code: safeCode }; export type MessageResponseProps = ComponentProps; diff --git a/components/ai-elements/streamdownCodeHighlighter.ts b/components/ai-elements/streamdownCodeHighlighter.ts new file mode 100644 index 000000000..464a8b68b --- /dev/null +++ b/components/ai-elements/streamdownCodeHighlighter.ts @@ -0,0 +1,76 @@ +import type { + CodeHighlighterPlugin, + HighlightOptions, + HighlightResult, +} from 'streamdown'; +import type { BundledLanguage } from 'shiki'; + +const PLAIN_TEXT_LANGUAGES = new Set([ + '', + 'plain', + 'plaintext', + 'text', + 'txt', +]); + +const LANGUAGE_ALIASES: Record = { + 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, + ); + }, +}); diff --git a/components/ai/streamdownCodeHighlighter.test.ts b/components/ai/streamdownCodeHighlighter.test.ts new file mode 100644 index 000000000..7d8fbbee2 --- /dev/null +++ b/components/ai/streamdownCodeHighlighter.test.ts @@ -0,0 +1,90 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import type { + CodeHighlighterPlugin, + HighlightOptions, + HighlightResult, +} from 'streamdown'; +import { + createPlainCodeHighlightResult, + createSafeCodeHighlighter, + resolveSupportedCodeLanguage, +} from '../ai-elements/streamdownCodeHighlighter'; + +const createFakeHighlighter = ( + supportedLanguages: string[], + highlightImpl?: CodeHighlighterPlugin['highlight'], +): CodeHighlighterPlugin => ({ + name: 'shiki', + type: 'code-highlighter', + getSupportedLanguages: () => supportedLanguages as ReturnType, + getThemes: () => ['github-light', 'github-dark'], + supportsLanguage: (language) => supportedLanguages.includes(language), + highlight: highlightImpl ?? ((options: HighlightOptions): HighlightResult => ({ + tokens: [[{ content: options.language, offset: 0 }]], + })), +}); + +test('maps generic conf fences to ini for Streamdown highlighting', () => { + const highlighter = createFakeHighlighter(['ini']); + + assert.equal(resolveSupportedCodeLanguage(highlighter, 'conf'), 'ini'); + assert.equal(resolveSupportedCodeLanguage(highlighter, ' config '), 'ini'); +}); + +test('falls back to plain tokens for unsupported languages', () => { + const highlighter = createSafeCodeHighlighter( + createFakeHighlighter([], () => { + throw new Error('delegate should not be called for unsupported languages'); + }), + ); + + const result = highlighter.highlight({ + code: '*.* action(type="omfwd"\n Target="10.185.3.1")\n', + language: 'conf', + themes: ['github-light', 'github-dark'], + }); + + assert.deepEqual( + result?.tokens.map((line) => line.map((token) => token.content).join('')), + ['*.* action(type="omfwd"', ' Target="10.185.3.1")'], + ); +}); + +test('uses supported aliases when highlighting generic config blocks', () => { + let receivedLanguage: string | null = null; + const highlighter = createSafeCodeHighlighter( + createFakeHighlighter(['ini'], (options: HighlightOptions): HighlightResult => { + receivedLanguage = options.language; + return createPlainCodeHighlightResult(options.code); + }), + ); + + const result = highlighter.highlight({ + code: '*.* action(type="omfwd")', + language: 'conf', + themes: ['github-light', 'github-dark'], + }); + + assert.equal(receivedLanguage, 'ini'); + assert.equal(result?.tokens[0][0].content, '*.* action(type="omfwd")'); +}); + +test('treats text fences as plain code without calling the delegate', () => { + const highlighter = createSafeCodeHighlighter( + createFakeHighlighter(['ini'], () => { + throw new Error('delegate should not be called for text fences'); + }), + ); + + const result = highlighter.highlight({ + code: 'hello\nworld', + language: 'text', + themes: ['github-light', 'github-dark'], + }); + + assert.deepEqual( + result?.tokens.map((line) => line[0].content), + ['hello', 'world'], + ); +});