[codex] Fix AI code block rendering fallback

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<HTMLDivElement> & {
   from: 'user' | 'assistant' | 'system' | 'tool';
@@ -60,7 +61,8 @@ export const MessageActions = ({ className, children, ...props }: MessageActions
   </div>
 );
 
-const streamdownPlugins = { cjk, code };
+const safeCode = createSafeCodeHighlighter(code);
+const streamdownPlugins = { cjk, code: safeCode };
 
 export type MessageResponseProps = ComponentProps<typeof Streamdown>;
 

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<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,
+    );
+  },
+});

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<CodeHighlighterPlugin['getSupportedLanguages']>,
+  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'],
+  );
+});