Streamdown 优化 AI 流式传输时使用客制化的 CodeBlock

Streamdown 简化了流失传输 AI 消息时渲染消息的过程，它比 react-markdown 更适合 AI 场景。同时它也把我带到了 AI-SDK 构建的 AI-Element -- 更适合 AI 原生的组件库。

AI-Element 和 Streamdown

引入：我的业务场景，使用 Response 优化了流式传输效果

我最近正在从零构建一个自己的 agent，基于 Nextjs 和 convex。虽然 convex/agent 已经帮助我完成了后端的绝大部分工作，但是在前端，显示 AI markdown 格式的返回消息仍然是一个需要解决的问题，默认的消息不能渲染 markdown 形式的消息，只是 string。我需要一个漂亮，流畅，可用的消息显示工具。

我去看了我们的老朋友 ai-sdk有没有新的解决方案，这也是 convex/agent 使用的开源库。他们给出了一个完美的解决方案，来自他们最近发布的支持 AI 原生应用构建的组件库 AI-Element 的消息组件 -- Streamdown。

我个人非常建议需要构建 Agent 的朋友们都去看看他们构建的代码。总有一个会对你有帮助。

但本次的课题其实是实现客制化的 blockcode。虽然你可以看到这个其实已经很完善了，但是还有点小缺点，比如字体太小了。我也希望能够探索更更定制化的方案

[TODO markdown 示例组件 - 使用 streamdown，放入 fumadocs card 中]

TLDR： Streamdown 通过精妙的架构设计实现了高性能的 Markdown 流式渲染。核心控制链条为：Streamdown → parseMarkdownIntoBlocks → Block → HardenedMarkdown → components mapping → CodeComponent → CodeBlock → HighlighterManager。采用了 Context Pattern（多层配置传递）、Singleton Pattern（语法高亮管理器）、Factory Pattern（组件映射机制）、Memo Pattern（性能优化策略）和 Composition Pattern（组件功能组合）等设计模式，通过 React.memo、useMemo、useContext 等 API 实现了优雅的状态管理和性能优化。

1. 架构概览：从Markdown到渲染的完整链条

Streamdown 的核心设计围绕着将 Markdown 文本高效转换为 React 组件的完整流程。让我们先看主入口组件的实现：

// https://github.com/vercel/streamdown/blob/main/packages/streamdown/index.tsx
export const Streamdown = memo(
  ({
    children,
    parseIncompleteMarkdown: shouldParseIncompleteMarkdown = true,
    components,
    shikiTheme = ["github-light", "github-dark"],
    mermaidConfig,
    controls = true,
    ...props
  }: StreamdownProps) => {
    const generatedId = useId();
    const blocks = useMemo(
      () =>
        parseMarkdownIntoBlocks(typeof children === "string" ? children : ""),
      [children],
    );

    return (
      <ShikiThemeContext.Provider value={shikiTheme}>
        <MermaidConfigContext.Provider value={mermaidConfig}>
          <ControlsContext.Provider value={controls}>
            <div className={cn("space-y-4", className)} {...props}>
              {blocks.map((block, index) => (
                <Block
                  components={{
                    ...defaultComponents,
                    ...components,
                  }}
                  content={block}
                  key={`${generatedId}-block_${index}`}
                  shouldParseIncompleteMarkdown={shouldParseIncompleteMarkdown}
                />
              ))}
            </div>
          </ControlsContext.Provider>
        </MermaidConfigContext.Provider>
      </ShikiThemeContext.Provider>
    );
  },
  (prevProps, nextProps) =>
    prevProps.children === nextProps.children &&
    prevProps.shikiTheme === nextProps.shikiTheme,
);

React.memo 的作用机制

React.memo 是一个高阶组件，用于优化函数组件的渲染性能：

const Component = memo(ComponentFunction, compareFunction?)

作用原理：

对组件进行浅比较，只有 props 发生变化时才重新渲染
第二个参数是可选的比较函数，返回 true 表示 props 相等（不重新渲染）
在 Streamdown 中，只有当 children 或 shikiTheme 改变时才重新渲染

useMemo 的依赖优化

useMemo 用于缓存计算结果，避免昂贵的重复计算：

const blocks = useMemo(
  () => parseMarkdownIntoBlocks(typeof children === "string" ? children : ""),
  [children], // 依赖数组：只有 children 变化时才重新解析
);

核心优势：

只有依赖项变化时才重新计算
避免了每次渲染都执行 parseMarkdownIntoBlocks 这个较重的操作
保持了引用的稳定性，减少子组件的不必要渲染

分块渲染策略

Streamdown 将整个 Markdown 文本分解为多个 Block 组件：

const Block = memo(
  ({ content, shouldParseIncompleteMarkdown, ...props }: BlockProps) => {
    const parsedContent = useMemo(
      () =>
        typeof content === "string" && shouldParseIncompleteMarkdown
          ? parseIncompleteMarkdown(content.trim())
          : content,
      [content, shouldParseIncompleteMarkdown],
    );

    return <HardenedMarkdown {...props}>{parsedContent}</HardenedMarkdown>;
  },
  (prevProps, nextProps) => prevProps.content === nextProps.content,
);

这种设计的优势：

渐进式渲染：每个 block 可以独立渲染，提升用户体验
精确更新：只有内容变化的 block 才会重新渲染
流式兼容：支持不完整的 Markdown 解析，适合 AI 流式输出

2. Context驱动的配置管理模式

Streamdown 使用 React Context API 实现了优雅的配置层级传递，避免了 props drilling 问题。

createContext 和 Provider 模式

// https://github.com/vercel/streamdown/blob/main/packages/streamdown/index.tsx
export const ShikiThemeContext = createContext<[BundledTheme, BundledTheme]>([
  "github-light" as BundledTheme,
  "github-dark" as BundledTheme,
]);

export const MermaidConfigContext = createContext<MermaidConfig | undefined>(
  undefined,
);

export const ControlsContext = createContext<ControlsConfig>(true);

createContext API 解析：

createContext(defaultValue) 创建一个 Context 对象
defaultValue 是当组件树中没有匹配的 Provider 时使用的值
返回的对象包含 Provider 和 Consumer 组件

多层嵌套的 Provider 模式

return (
  <ShikiThemeContext.Provider value={shikiTheme}>
    <MermaidConfigContext.Provider value={mermaidConfig}>
      <ControlsContext.Provider value={controls}>
        {/* 组件树 */}
      </ControlsContext.Provider>
    </MermaidConfigContext.Provider>
  </ShikiThemeContext.Provider>
);

这种嵌套设计的优势：

配置隔离：每个 Context 负责特定的配置域
类型安全：TypeScript 提供完整的类型检查
按需订阅：组件只消费需要的 Context，减少不必要的重渲染

useContext hook 的使用模式

在子组件中消费 Context：

// https://github.com/vercel/streamdown/blob/main/packages/streamdown/lib/components.tsx
const CodeComponent = ({ node, className, children, ...props }) => {
  const mermaidConfig = useContext(MermaidConfigContext);
  const controlsConfig = useContext(ControlsContext);

  // 根据配置控制功能显示
  const showCodeControls = shouldShowControls(controlsConfig, "code");
  // ...
};

useContext API 特点：

自动订阅 Context 的变化
返回最近的 Provider 的 value
Context 值变化时，所有消费组件都会重新渲染

3. Components映射：从AST到React组件

Components 映射是 Streamdown 的核心机制，将 Markdown AST 节点转换为对应的 React 组件。

工厂模式的 components 对象

// https://github.com/vercel/streamdown/blob/main/packages/streamdown/lib/components.tsx
export const components: Options["components"] = {
  ol: MemoOl,
  li: MemoLi,
  ul: MemoUl,
  hr: MemoHr,
  strong: MemoStrong,
  a: MemoA,
  h1: MemoH1,
  h2: MemoH2,
  h3: MemoH3,
  h4: MemoH4,
  h5: MemoH5,
  h6: MemoH6,
  table: MemoTable,
  thead: MemoThead,
  tbody: MemoTbody,
  tr: MemoTr,
  th: MemoTh,
  td: MemoTd,
  blockquote: MemoBlockquote,
  code: MemoCode,
  img: MemoImg,
  pre: ({ children }) => children,
  sup: MemoSup,
  sub: MemoSub,
};

这个对象充当了工厂模式的注册表，react-markdown 根据 AST 节点类型查找对应的组件进行渲染。

CodeComponent 的条件渲染逻辑

Code 组件是最复杂的组件之一，需要处理行内代码和代码块两种情况：

const CodeComponent = ({ node, className, children, ...props }) => {
  // 判断是否为行内代码
  const inline = node?.position?.start.line === node?.position?.end.line;

  if (inline) {
    return (
      <code
        className={cn(
          "bg-muted rounded px-1.5 py-0.5 font-mono text-sm",
          className,
        )}
      >
        {children}
      </code>
    );
  }

  // 解析语言类型
  const match = className?.match(LANGUAGE_REGEX);
  const language = (match?.at(1) ?? "") as BundledLanguage;

  // 特殊处理 mermaid 图表
  if (language === "mermaid") {
    return <MermaidBlock />;
  }

  // 渲染普通代码块
  return <CodeBlock language={language} code={code} />;
};

React.memo 的精准比较策略

每个组件都使用了精心设计的比较函数：

const MemoCode = memo<ComponentProps>(
  CodeComponent,
  (p, n) => p.className === n.className && sameNodePosition(p.node, n.node),
);

function sameNodePosition(prev?: MarkdownNode, next?: MarkdownNode): boolean {
  if (!(prev?.position || next?.position)) return true;
  if (!(prev?.position && next?.position)) return false;

  const prevStart = prev.position.start;
  const nextStart = next.position.start;
  const prevEnd = prev.position.end;
  const nextEnd = next.position.end;

  return (
    prevStart?.line === nextStart?.line &&
    prevStart?.column === nextStart?.column &&
    prevEnd?.line === nextEnd?.line &&
    prevEnd?.column === nextEnd?.column
  );
}

sameNodePosition 的作用：

比较 AST 节点在源文档中的位置
确保只有内容真正变化的节点才重新渲染
对流式输出场景特别重要，避免了大量无效渲染

4. CodeBlock核心实现解析

CodeBlock 是整个系统最复杂的部分，涉及语法高亮、主题切换、异步加载等多个方面。

HighlighterManager 单例模式

// https://github.com/vercel/streamdown/blob/main/packages/streamdown/lib/code-block.tsx
class HighlighterManager {
  private lightHighlighter: Awaited<
    ReturnType<typeof createHighlighter>
  > | null = null;
  private darkHighlighter: Awaited<
    ReturnType<typeof createHighlighter>
  > | null = null;
  private lightTheme: BundledTheme | null = null;
  private darkTheme: BundledTheme | null = null;
  private readonly loadedLanguages: Set<BundledLanguage> = new Set();
  private initializationPromise: Promise<void> | null = null;

  async highlightCode(
    code: string,
    language: BundledLanguage,
    themes: [BundledTheme, BundledTheme],
    preClassName?: string,
  ): Promise<[string, string]> {
    // 确保只有一个初始化过程
    if (this.initializationPromise) {
      await this.initializationPromise;
    }

    this.initializationPromise = this.ensureHighlightersInitialized(
      themes,
      language,
    );
    await this.initializationPromise;
    this.initializationPromise = null;

    // 返回亮色和暗色主题的高亮结果
    return [lightHtml, darkHtml];
  }
}

const highlighterManager = new HighlighterManager();

单例模式的优势：

全局共享高亮器实例，避免重复创建
缓存已加载的语言和主题，提升性能
统一管理异步初始化过程

useEffect 和 useRef 的异步状态管理

export const CodeBlock = ({
  code,
  language,
  className,
  children,
  preClassName,
  ...rest
}) => {
  const [html, setHtml] = useState<string>("");
  const [darkHtml, setDarkHtml] = useState<string>("");
  const mounted = useRef(false);
  const [lightTheme, darkTheme] = useContext(ShikiThemeContext);

  useEffect(() => {
    mounted.current = true;

    highlighterManager
      .highlightCode(code, language, [lightTheme, darkTheme], preClassName)
      .then(([light, dark]) => {
        if (mounted.current) {
          setHtml(light);
          setDarkHtml(dark);
        }
      });

    return () => {
      mounted.current = false;
    };
  }, [code, language, lightTheme, darkTheme, preClassName]);

  // ...
};

useRef 防止内存泄漏：

mounted.current 标记组件是否仍然挂载
异步操作完成时检查组件状态，避免在卸载的组件上设置状态
这是处理异步操作的标准模式

dangerouslySetInnerHTML 的安全使用

<div
  className={cn("overflow-x-auto dark:hidden", className)}
  dangerouslySetInnerHTML={{ __html: html }}
  data-code-block
  data-language={language}
  {...rest}
/>
<div
  className={cn("hidden overflow-x-auto dark:block", className)}
  dangerouslySetInnerHTML={{ __html: darkHtml }}
  data-code-block
  data-language={language}
  {...rest}
/>

安全考虑：

HTML 来源于 Shiki 库生成，是可信的
使用 data-* 属性提供语义信息
通过 CSS 类控制主题切换

5. 功能扩展的组合模式

Streamdown 通过组合模式实现了代码块的扩展功能，如复制和下载按钮。

Context-based 的数据获取策略

type CodeBlockContextType = {
  code: string;
};

const CodeBlockContext = createContext<CodeBlockContextType>({
  code: "",
});

export const CodeBlock = ({ code, language, className, children, ...rest }) => {
  return (
    <CodeBlockContext.Provider value={{ code }}>
      <div className="my-4 w-full overflow-hidden rounded-xl border">
        <div className="bg-muted/80 flex items-center justify-between p-3">
          <span className="ml-1 font-mono lowercase">{language}</span>
          <div className="flex items-center gap-2">{children}</div>
        </div>
        {/* 代码内容 */}
      </div>
    </CodeBlockContext.Provider>
  );
};

组合模式的优势：

代码块和按钮功能解耦
按钮组件通过 Context 获取代码内容
支持灵活的功能组合

按钮组件的状态管理

export const CodeBlockCopyButton = ({
  onCopy,
  onError,
  timeout = 2000,
  ...props
}) => {
  const [isCopied, setIsCopied] = useState(false);
  const timeoutRef = useRef(0);
  const contextCode = useContext(CodeBlockContext).code;

  const copyToClipboard = async () => {
    try {
      if (!isCopied) {
        await navigator.clipboard.writeText(contextCode);
        setIsCopied(true);
        onCopy?.();
        timeoutRef.current = window.setTimeout(
          () => setIsCopied(false),
          timeout,
        );
      }
    } catch (error) {
      onError?.(error as Error);
    }
  };

  useEffect(() => {
    return () => {
      window.clearTimeout(timeoutRef.current);
    };
  }, []);

  const Icon = isCopied ? CheckIcon : CopyIcon;
  return <button onClick={copyToClipboard}>{Icon}</button>;
};

useState 和 useEffect 的交互：

useState 管理复制状态和图标切换
useRef 存储定时器引用，避免闭包问题
useEffect 清理函数确保组件卸载时清除定时器

错误处理和用户反馈

组件支持自定义的错误处理和成功回调：

<CodeBlockCopyButton
  onCopy={() => console.log("Copied!")}
  onError={(error) => console.error("Copy failed:", error)}
  timeout={3000}
/>

这种设计允许上层应用自定义用户反馈机制。

6. 性能优化的设计哲学

Streamdown 在性能优化方面体现了 React 应用的最佳实践。

React.memo 的分层优化

// 组件级优化
const MemoCode = memo(
  CodeComponent,
  (p, n) => p.className === n.className && sameNodePosition(p.node, n.node),
);

// 通用比较函数
function sameClassAndNode(prev, next) {
  return (
    prev.className === next.className && sameNodePosition(prev.node, next.node)
  );
}

const MemoH1 = memo(H1Component, sameClassAndNode);
const MemoTable = memo(TableComponent, sameClassAndNode);

优化策略：

每个组件都有定制的比较函数
抽取公共比较逻辑避免重复
基于 AST 节点位置的精确比较

useMemo 的缓存机制

const blocks = useMemo(
  () => parseMarkdownIntoBlocks(typeof children === "string" ? children : ""),
  [children],
);

const parsedContent = useMemo(
  () =>
    typeof content === "string" && shouldParseIncompleteMarkdown
      ? parseIncompleteMarkdown(content.trim())
      : content,
  [content, shouldParseIncompleteMarkdown],
);

const rehypeKatexPlugin = useMemo(
  () => () => rehypeKatex({ errorColor: "var(--color-muted-foreground)" }),
  [],
);

缓存层次：

文档级缓存：整个 Markdown 的分块解析
块级缓存：单个块的不完整 Markdown 处理
插件级缓存：rehype 插件的配置对象

HighlighterManager 的资源复用

class HighlighterManager {
  private async ensureHighlightersInitialized(themes, language) {
    // 检查是否需要重新创建高亮器
    const needsLightRecreation =
      !this.lightHighlighter || this.lightTheme !== lightTheme;
    const needsDarkRecreation =
      !this.darkHighlighter || this.darkTheme !== darkTheme;

    // 检查是否需要加载新语言
    const needsLanguageLoad =
      !this.loadedLanguages.has(language) && isLanguageSupported;

    // 智能的资源管理
    if (needsLanguageLoad && !needsLightRecreation) {
      await this.lightHighlighter?.loadLanguage(language);
    }
  }
}

资源复用策略：

高亮器实例复用，避免重复创建
语言按需加载，减少初始化开销
主题变化时的智能重建

异步加载的状态同步

useEffect(() => {
  mounted.current = true;

  highlighterManager
    .highlightCode(code, language, [lightTheme, darkTheme], preClassName)
    .then(([light, dark]) => {
      if (mounted.current) {
        // 关键：检查组件是否仍然挂载
        setHtml(light);
        setDarkHtml(dark);
      }
    });

  return () => {
    mounted.current = false; // 清理标记
  };
}, [code, language, lightTheme, darkTheme, preClassName]);

最佳实践：

使用 useRef 跟踪组件挂载状态
异步操作完成前检查组件状态
避免在卸载的组件上执行状态更新

通过这些精心设计的优化策略，Streamdown 实现了高性能的 Markdown 流式渲染，特别适合 AI 应用场景中的动态内容展示。

// 从 Streamdown 源码复制 CodeComponent 的完整实现
// 保持所有功能：语法高亮、复制按钮、下载功能、Mermaid 支持
const CustomCodeComponentImpl = ({ node, className, children, ...props }) => {
  // ... 完全相同的逻辑 ...

  return (
    <CodeBlock
      // 唯一的修改：字体大小
      preClassName="overflow-x-auto font-mono text-sm p-4 bg-muted/40"
      // 其他参数保持不变
    >
      {/* 完全相同的 children */}
    </CodeBlock>
  );
};

2. 使用 React.memo 保持性能

export const CustomCodeComponent = memo(
  CustomCodeComponentImpl,
  (p, n) => p.className === n.className && sameNodePosition(p.node, n.node),
);

3. 通过 components 参数替换

<Streamdown
  components={{
    code: CustomCodeComponent, // 替换默认的 code 组件
  }}
>
  {markdownContent}
</Streamdown>

字体大小对比

通过上面的交互演示，你可以清楚地看到字体大小的差异：

版本	字体大小	像素值	适用场景
原版	text-xs	12px	紧凑显示，信息密度高
自定义	text-sm	14px	推荐：平衡密度和可读性
大字体	text-base	16px	演示教学，最佳阅读体验

技术实现详解

核心修改只有一行代码

// 原版 Streamdown
preClassName = "overflow-x-auto font-mono text-xs p-4 bg-muted/40";
//                                    ↑
//                                  12px

// 自定义版本
preClassName = "overflow-x-auto font-mono text-sm p-4 bg-muted/40";
//                                    ↑
//                                  14px (提升 17%)

完整的文件结构

src/components/blog/sreamdown-custome-code-block-ai-element/
├── code-block.tsx                 # 从 Streamdown 复制的完整实现
├── custom-code-component.tsx      # 自定义字体大小的 CodeComponent
├── data.ts                        # 演示用的 AI 响应数据
├── custom-response-block-code-demo.tsx  # 交互式演示组件
└── README.md                      # 完整的技术文档

优势总结

我们的方案具有以下优势：

✅ 最小侵入性

只修改 preClassName 中的一个 CSS 类名
保持 Streamdown 的所有原有功能
兼容未来版本更新

✅ 完整功能保留

✅ Shiki 语法高亮 (支持 300+ 语言)
✅ 复制和下载按钮
✅ 明暗主题自动切换
✅ Mermaid 图表支持
✅ 性能优化 (React.memo + 缓存)

✅ 类型安全

修复了 Streamdown 源码的 TypeScript 类型错误
完整的类型定义和 IDE 支持
编译时错误检查

✅ 开发体验

解决了大文件导致的内存溢出问题
完整的 JSDoc 注释和技术文档
支持多种字体大小变体

实际应用

在你的项目中使用只需要三步：

// 1. 导入自定义组件
import { CustomCodeComponent } from "./custom-code-component";

// 2. 替换默认实现
<Streamdown components={{ code: CustomCodeComponent }}>
  {aiResponse}
</Streamdown>;

// 3. 享受更好的代码可读性！

这种方案特别适合：

AI 对话应用：提升代码块的用户体验
技术博客：改善读者的阅读体验
文档网站：平衡信息密度和可读性
代码演示：确保代码清晰可见

通过这个实现，我们证明了有时候最简单的改动也能带来显著的用户体验提升。从 12px 到 14px 的字体调整，让代码块的可读性提升了 17%，同时保持了 Streamdown 的所有强大功能。