Genarrative/apps/ai-game-creator-shell/src/components/ChatMarkdownMessage/index.tsx

import type { ErrorInfo, ReactNode } from 'react';
import { Component, createContext, useContext } from 'react';
import ReactMarkdown, { type Components } from 'react-markdown';
import remarkGfm from 'remark-gfm';

export type ChatMarkdownMessageProps = {
  text: string;
  role: 'assistant' | 'user';
  streaming?: boolean;
};

type MarkdownErrorBoundaryProps = {
  fallbackText: string;
  children: ReactNode;
};

type MarkdownErrorBoundaryState = {
  hasError: boolean;
};

export class MarkdownErrorBoundary extends Component<
  MarkdownErrorBoundaryProps,
  MarkdownErrorBoundaryState
> {
  state: MarkdownErrorBoundaryState = { hasError: false };

  static getDerivedStateFromError(): MarkdownErrorBoundaryState {
    return { hasError: true };
  }

  componentDidCatch(_error: unknown, _errorInfo: ErrorInfo) {
    // Keep the original message visible without logging its potentially sensitive content.
  }

  componentDidUpdate(prevProps: MarkdownErrorBoundaryProps) {
    if (
      this.state.hasError &&
      prevProps.fallbackText !== this.props.fallbackText
    ) {
      this.setState({ hasError: false });
    }
  }

  render() {
    if (this.state.hasError) {
      return (
        <span className="whitespace-pre-wrap break-words">
          {this.props.fallbackText}
        </span>
      );
    }
    return this.props.children;
  }
}

const ListDepthContext = createContext(0);
const ListKindContext = createContext<'unordered' | 'ordered' | null>(null);
const ListItemContext = createContext(false);

function MarkdownUnorderedList({ children }: { children?: ReactNode }) {
  const depth = useContext(ListDepthContext);
  return (
    <ListDepthContext.Provider value={depth + 1}>
      <ListKindContext.Provider value="unordered">
        <ul
          className={`m-0 mt-2 list-none space-y-1 first:mt-0 ${
            depth > 0 ? 'pl-4' : 'pl-0'
          }`}
        >
          {children}
        </ul>
      </ListKindContext.Provider>
    </ListDepthContext.Provider>
  );
}

function MarkdownOrderedList({ children }: { children?: ReactNode }) {
  const depth = useContext(ListDepthContext);
  return (
    <ListDepthContext.Provider value={depth + 1}>
      <ListKindContext.Provider value="ordered">
        <ol className="m-0 mt-2 list-decimal space-y-1 pl-5 first:mt-0">
          {children}
        </ol>
      </ListKindContext.Provider>
    </ListDepthContext.Provider>
  );
}

function MarkdownParagraph({ children }: { children?: ReactNode }) {
  const inListItem = useContext(ListItemContext);
  return (
    <p
      className={`m-0 break-words ${inListItem ? 'inline' : 'mt-2 first:mt-0'}`}
    >
      {children}
    </p>
  );
}

function MarkdownListItem({ children }: { children?: ReactNode }) {
  const listKind = useContext(ListKindContext);
  return (
    <ListItemContext.Provider value>
      <li className="break-words whitespace-normal">
        {listKind === 'unordered' ? '- ' : null}
        {children}
      </li>
    </ListItemContext.Provider>
  );
}

const markdownComponents: Components = {
  // TODO: 产品确认安全外链策略后，再将链接文本恢复为可点击元素。
  a: ({ children }) => children,
  img: ({ alt }) => (alt?.trim() ? `图片：${alt}` : '图片已省略'),
  h1: ({ children }) => (
    <h1 className="m-0 mt-4 !text-xl font-bold first:mt-0">{children}</h1>
  ),
  h2: ({ children }) => (
    <h2 className="m-0 mt-4 !text-lg font-bold first:mt-0">{children}</h2>
  ),
  h3: ({ children }) => (
    <h3 className="m-0 mt-3 !text-base font-semibold first:mt-0">{children}</h3>
  ),
  h4: ({ children }) => (
    <h4 className="m-0 mt-3 !text-sm font-semibold first:mt-0">{children}</h4>
  ),
  h5: ({ children }) => (
    <h5 className="m-0 mt-2 !text-sm font-medium first:mt-0">{children}</h5>
  ),
  h6: ({ children }) => (
    <h6 className="m-0 mt-2 !text-xs font-medium uppercase tracking-wide first:mt-0">
      {children}
    </h6>
  ),
  p: MarkdownParagraph,
  ul: MarkdownUnorderedList,
  ol: MarkdownOrderedList,
  li: MarkdownListItem,
  blockquote: ({ children }) => (
    <blockquote className="m-0 mt-2 border-l-2 border-(--platform-surface-border) pl-3 text-(--platform-text-soft) first:mt-0">
      {children}
    </blockquote>
  ),
  pre: ({ children }) => (
    <pre className="m-0 mt-2 max-w-full overflow-x-auto rounded-lg bg-black/6 p-3 text-xs leading-5 first:mt-0">
      {children}
    </pre>
  ),
  code: ({ className, children, node: _node, ...props }) => {
    const isBlock =
      Boolean(className?.includes('language-')) ||
      String(children).includes('\n');
    return isBlock ? (
      <code {...props} className="font-mono whitespace-pre">
        {children}
      </code>
    ) : (
      <code
        {...props}
        className="rounded bg-black/6 px-1 py-0.5 font-mono text-[0.9em]"
      >
        {children}
      </code>
    );
  },
  table: ({ children }) => (
    <div className="mt-2 max-w-full overflow-x-auto first:mt-0">
      <table className="min-w-full border-collapse text-left text-sm">
        {children}
      </table>
    </div>
  ),
  th: ({ children }) => (
    <th className="border border-(--platform-surface-border) px-2 py-1 font-semibold">
      {children}
    </th>
  ),
  td: ({ children }) => (
    <td className="border border-(--platform-surface-border) px-2 py-1 align-top">
      {children}
    </td>
  ),
  hr: () => (
    <hr className="mt-2 border-0 border-t border-(--platform-surface-border)" />
  ),
};

const streamingMarkdownComponents: Components = {
  ...markdownComponents,
  p: ({ children }) => (
    <p className="m-0 mt-2 break-words opacity-95 first:mt-0">{children}</p>
  ),
};

function escapeHtmlTagsForDisplay(text: string) {
  return text.replace(/<\/?[A-Za-z][^>]*>/gu, (tag) =>
    tag.replaceAll('<', '&lt;').replaceAll('>', '&gt;'),
  );
}

export function ChatMarkdownMessage({
  text,
  role,
  streaming = false,
}: ChatMarkdownMessageProps) {
  if (role === 'user') {
    return <span className="whitespace-pre-wrap break-words">{text}</span>;
  }

  return (
    <MarkdownErrorBoundary fallbackText={text}>
      <ReactMarkdown
        skipHtml
        remarkPlugins={[remarkGfm]}
        components={
          streaming ? streamingMarkdownComponents : markdownComponents
        }
      >
        {escapeHtmlTagsForDisplay(text)}
      </ReactMarkdown>
    </MarkdownErrorBoundary>
  );
}