chat-one-web/src/components/StreamMessage.tsx

import { useState } from "react";
import { Button, message } from "antd";
import { Check, Copy } from "lucide-react";
import ReactMarkdown from "react-markdown";
import remarkGfm from "remark-gfm";
import remarkMath from "remark-math";
import rehypeKatex from "rehype-katex";
import rehypeHighlight from "rehype-highlight";
import "katex/dist/katex.min.css";
import "highlight.js/styles/github.css";
import type { ReactNode } from "react";

type StreamMessageProps = {
  content: string;
};

function extractText(node: ReactNode): string {
  if (typeof node === "string") return node;
  if (typeof node === "number") return String(node);
  if (!node) return "";
  if (Array.isArray(node)) return node.map((item) => extractText(item)).join("");
  if (typeof node === "object" && "props" in node) {
    const child = (node as { props?: { children?: ReactNode } }).props?.children;
    return extractText(child);
  }
  return "";
}

function MarkdownPreBlock(props: { children: ReactNode }) {
  const [copied, setCopied] = useState(false);
  const codeText = extractText(props.children).replace(/\n$/, "");

  const onCopy = async () => {
    try {
      await navigator.clipboard.writeText(codeText);
      setCopied(true);
      message.success("代码已复制");
      window.setTimeout(() => setCopied(false), 1200);
    } catch {
      message.error("复制失败");
    }
  };

  return (
    <div className="my-2 overflow-hidden rounded-xl bg-neutral-50">
      <div className="flex items-center justify-end px-2 py-1">
        <Button
          type="text"
          size="small"
          icon={copied ? <Check size={16} /> : <Copy size={16} />}
          className="text-neutral-600 hover:bg-black/5! hover:text-neutral-800!"
          onClick={() => {
            void onCopy();
          }}
        >
          {copied ? "已复制" : "复制"}
        </Button>
      </div>
      <pre className="m-0 overflow-x-auto p-3 text-[13px] text-neutral-800">{props.children}</pre>
    </div>
  );
}

export default function StreamMessage(props: StreamMessageProps) {
  return (
    <div className="rounded-2xl px-4 py-2.5 text-[15px] leading-relaxed text-neutral-800">
      <div className="prose prose-sm max-w-none break-words prose-neutral prose-code:before:content-none prose-code:after:content-none">
        <ReactMarkdown
          remarkPlugins={[remarkGfm, remarkMath]}
          rehypePlugins={[rehypeKatex, rehypeHighlight]}
          components={{
            pre: ({ children }) => <MarkdownPreBlock>{children}</MarkdownPreBlock>,
            table: ({ children }) => (
              <div className="my-2 overflow-x-auto rounded-lg border border-[var(--ds-border)]">
                <table className="w-full border-collapse text-sm">{children}</table>
              </div>
            ),
            th: ({ children }) => (
              <th className="border-b border-r border-[var(--ds-border)] bg-neutral-100 px-3 py-2 text-left font-semibold last:border-r-0">
                {children}
              </th>
            ),
            td: ({ children }) => (
              <td className="border-b border-r border-[var(--ds-border)] px-3 py-2 align-top last:border-r-0">
                {children}
              </td>
            ),
            blockquote: ({ children }) => (
              <blockquote className="my-2 border-l-4 border-sky-200 bg-sky-50/60 px-3 py-2 text-neutral-700">
                {children}
              </blockquote>
            ),
          }}
        >
          {props.content}
        </ReactMarkdown>
      </div>
    </div>
  );
}